Switch to Clima-compatible data structures

[edejong-caltech]

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

  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

[edejong-caltech]

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