Switch to Clima-compatible data structures

[ekdejong]

It seems like we need to take a several steps to make Cloudy compatible with ClimaAtmos...

• Data structures
  • Remove mutable structs (must be immutable)
  • Replace Arrays and Vectors with Tuples or SVectors
• Function signatures
  • Ensure function calls return values, rather than updating structs
  • They also should not allocate (where possible)
• Unit testing
  • Add tests for @allocated as necessary

Next steps will include:

• Updating CloudMicrophysics.jl to use the appropriate new structures and function calls in a way that is still compatible with ClimaAtmos.jl
• (Re)integrate Cloudy.jl and/or CM.jl into ClimaAtmos. (Note the scaffolding for this already exists in branch ej/add_cloudy)

[charleskawczynski]

Some notes from todays discussion:

Compatibility todos

Top items

• Split up data structures by spatially varying and spatially invariant
• Convert mutable -> immutable structs
• Arrays -> Tuples/SArrays
• Fixing allocations (in a stack-allocated way). Example:

function foo(a)
	b = zeros(size(a))
	for i = 1:length(a)
		b[i]+=rand()
	end
	return b
end

# non-allocating, but still mutating
function foo!(a, b)
	for i = 1:length(a)
		b[i]+=rand()
	end
	return b
end

# stack-allocated, much better
foo(a) = map(a_i->rand(), a)

Additional + details

• @assert is not allowed -> use error, and don't capture variables in the error message
• Vector -> Tuple / StaticArrays.SArray
• Use map over initializing / + loops, when applicable
• Add JET & @allocated tests to functions we expect to be in tendency kernels
• If a function is only used during initialization, then we can leave these for a later time?

    norm = zeros(FT, sum(NProgMoms))
    n_dist = length(NProgMoms)
    for (i, n_mom) in enumerate(NProgMoms)
        for j in 1:n_mom
            norm[get_dist_moment_ind(NProgMoms, i, j)] = norms[1] * norms[2]^(j - 1)
        end
    end

• Use concrete eltypes in structs (e.g., in CoalescenceTensor):

struct CoalescenceTensor{FT} <: KernelTensor{FT}
    c::Array{FT}

to

struct CoalescenceTensor{FT, AFT <: AbstractArray{FT}} <: KernelTensor{FT}
    c::AFT

• (maybe) remove keyword arguments, which are not always well supported inside gpu kernels
• Remove runtime allocations:
• collect - explicit allocation
• reshape - explicit allocation
• [getproperty(dist, p) for p in params] - remove use of list comprehensions
• Array{Symbol, 1}/Matrix{}(...) - remove explicit array construction ->
• broadcasting: y = floor.(x) can allocate or cause issues, maybe just use map
• findall
• ones(...) / zeros(...)
• [C_1_1] makes a vector (assumes Vector type)
• intersect
• @views are likely problematic
• Dicts are not allowed (on CPU or GPU kernels) -> NamedTuples, or custom structs
  • NamedTuple keeps track of all separate types

julia> typeof((;a=1, b=2.0))
@NamedTuple{a::Int64, b::Float64}
```julia
 - Remove `n, θ = get_params(dist)[2]` -> `(;n, θ) = dist`
 - Remove
```julia
function get_params(dist::PrimitiveParticleDistribution{FT}) where {FT <: Real}
    params = Array{Symbol, 1}(collect(propertynames(dist)))
    values = Array{FT, 1}([getproperty(dist, p) for p in params])
    return params, values
end

entirely

[ekdejong]

Re: making Coalescence.jl GPU compatible...

1. update_coal_ints should return an SVector

• Perhaps we don't need to store Q, R, S
• Remove update_moments!, maybe pass the current moments in instead
• We should store finite_2d_integrals as a local map, rather than updating coal_data
• As a first pass, try (Q, R, S) = get_coalescence_integral_qrs...
• Expand out the arguments so we don't have coal_data, and keep local and global variables separate

2. Get rid of intialize_coalescence_data