ClimaCore.jl
ClimaCore.jl copied to clipboard
Spectral element operator patterns (possible improvements)
Following a review comment by charleskawczynski, I'll leave this suggestion in an issue for future visibility - we currently use @simd in the spectral operator blocks, but an alternative is as follows:
e1 = length(input) ≥ 1 ? (mapreduce(k -> D[i, k] ⊠ input[1][k, j, vt], ⊞, 1:Nq), mapreduce(k -> D[j, k] ⊠ input[1][i, k, vt], ⊞, 1:Nq)) : ()
e2 = length(input) ≥ 2 ? (mapreduce(k -> D[i, k] ⊠ input[2][k, j, vt], ⊞, 1:Nq), mapreduce(k -> D[j, k] ⊠ input[2][i, k, vt], ⊞, 1:Nq)) : ()
e3 = length(input) ≥ 3 ? (mapreduce(k -> D[i, k] ⊠ input[3][k, j, vt], ⊞, 1:Nq), mapreduce(k -> D[j, k] ⊠ input[3][i, k, vt], ⊞, 1:Nq)) : ()
t = (e1...,e2...,e3...)
if length(input) == 1 # check types
return Geometry.Covariant12Vector(t...)
elseif length(input) == 2
return Geometry.AxisTensor((Geometry.Covariant12Axis(), Geometry.UVAxis()), t)
else
return Geometry.AxisTensor((Geometry.Covariant12Axis(), Geometry.UVWAxis()), t)
end
Originally posted by @charleskawczynski in https://github.com/CliMA/ClimaCore.jl/pull/1736#discussion_r1710014985