Compute/output value calculated from tracers as an array

[j-ostiguy]

I would like to calculate pCO2 as a function of my model tracers (alkalinity, dic, S and T). I am using the CO2System package, which will calculate pCO2 from these 4 values but only accepts scalars or vectors. I can't seem to find a way to either build an array at each time step or compute the pCO2 field.

I defined this function:

function co2sys(alk,dic,S,T)
    Pin = 0
    Pout = 0
    in1type = 1
    in2type = 2
    SI = 0
    PO4 = 0
    pHscalein = 1 
    K1K2 = 4
    KSO4 = 1
  
    vals, _, _ = CO2SYS(alk,dic,in1type,in2type,S,T,T,Pin,Pout,SI,PO4,pHscalein,K1K2,KSO4)
  
    return vals[19] #return pco2 value
  end

and I've tried several things along these lines:

#pCO2w_calc =  co2sys(model.tracers.alk,model.tracers.dic,model.tracers.S,model.tracers.T)

#pCO2w = compute!(Field(pCO2w_calc))

#=
#calculate pCO2
pCO2w = zeros(grid.Nx,grid.Ny,grid.Nz) #initialize pCO2w
for xx in 1:grid.Nx
    for yy in 1:grid.Ny
        for zz in 1:grid.Nz
                pCO2w[xx,yy,zz] = co2sys(model.tracers.alk[xx,yy,zz],model.tracers.dic[xx,yy,zz],
                                                model.tracers.S[xx,yy,zz],model.tracers.T[xx,yy,zz])
        end
    end
end
=#

#=
@inline function pCO2w(i,j,k,grid,clock,model_fields) 
    for i in 1:grid.Nx
        for j in 1:grid.Ny
            for k in 1:grid.Nz
                pCO2w_val[i,j,k] = co2sys(model_fields.alk[i,j,k],model_fields.dic[i,j,k],model_fields.S[i,j,k],model_fields.T[i,j,k])
            end
        end
    end
    return pCO2w_val
end
=#

[simone-silvestri]

Hello @j-ostiguy. You are trying to do a local operation, so there are two ways you can go about it, one more "high level" (1) approach and one more "low level" approach (2)

1. Use KernelFunctionOperations. This allows you to write a function with the following signature: (i, j, k, grid, args...) that is applied locally to all i, j, k in the grid. A way to do that would be:

@inline co2sys_op(i, j, k, grid, alk, dic, S, T) = @inbounds co2sys(alk[i, j, k], dic[i, j, k], S[i, j, k], T[i, j, k])

alk = model.tracers.alk
dic = model.tracers.alk
T = model.tracers.T
S = model.tracers.S

pCO2w_comp = KernelFunctionOperation{Center, Center, Center}(co2sys_op, grid, computed_dependencies = (alk, dic, S, T))
pCO2w_val = compute!(Field(pCO2w_comp))

2. To explain a little bit of what happens behind the scenes, KernelFunctionOperation constructs a lazy operation that (when compute!d in a field) launches a kernel that (not exactly but in essence) looks something like this

@kernel function compute_kernel!(pCO2w_val, pCO2w_comp::KernelFunctionOperation)
    i, j, k = @index(Global, NTuple)
    pCO2w_val[i, j, k] = pCO2w_comp.function(i, j, k, pCO2w_comp.grid, pCO2w_comp.args....)
end

where in this case pCO2w_comp.function is your function co2sys_op and pCO2w_comp.args is your compute_dependencies.
This function (compute_kernel!) is a kernel implemented using the KernelAbstractions.jl package. It basically abstracts the notion of for loops, to threaded loops on CPU and kernel functions on the GPU and allows you to write local computations, "looping" over a set of i, j, k, indices. You have to specify the sizes of these loops as launch parameters. On how to call and interact with @kernel functions I recommend you to look at the KernelAbstractions.jl package itself. You can write your own kernels and attach them to the simulation through callbacks.

[j-ostiguy]

Thanks for your help!