Merge pull request #2378 from CliMA/he/ft-add-left-and-right-projection

ft: add (left and right) projection of 2-tensors

Author: haakon-e

NEWS.md

@@ -4,6 +4,11 @@ ClimaCore.jl Release Notes
 main
 -------
 
+v0.14.41
+-------
+- Add `project` for both sides of 2-tensors (previously only left side was supported). 
+  Simplify implementation. [2379](https://github.com/Clima/ClimaCore.jl/pull/2379)
+
 v0.14.40
 -------
 - Store `reverse_mode` in the `IntervalMesh` struct, so that we can access it when writing

Project.toml

@@ -1,7 +1,7 @@
 name = "ClimaCore"
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
 authors = ["CliMA Contributors <[email protected]>"]
-version = "0.14.40"
+version = "0.14.41"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/Geometry/axistensors.jl

@@ -1,4 +1,4 @@
-using StaticArrays, LinearAlgebra
+using StaticArrays, LinearAlgebra, UnrolledUtilities
 
 """
     AbstractAxis
@@ -8,6 +8,7 @@ An axis of a [`AxisTensor`](@ref).
 abstract type AbstractAxis{I} end
 
 Base.Broadcast.broadcastable(a::AbstractAxis) = a
+axis_indices(::AbstractAxis{I}) where {I} = I
 
 """
     dual(ax::AbstractAxis)
@@ -521,101 +522,41 @@ function _project(
     x
 end
 
-#= Set `assert_exact_transform() = true` for debugging=#
-assert_exact_transform() = false
-
-@generated function _transform(
-    ato::Ato,
-    x::Axis2Tensor{T, Tuple{Afrom, A2}},
-) where {
-    Ato <: AbstractAxis{Ito},
-    Afrom <: AbstractAxis{Ifrom},
-    A2 <: AbstractAxis{J},
-} where {Ito, Ifrom, J, T}
-    N = length(Ifrom)
-    M = length(J)
-    if assert_exact_transform()
-        errcond = false
-        for n in 1:N
-            i = Ifrom[n]
-            if i ∉ Ito
-                for m in 1:M
-                    errcond = :($errcond || x[$n, $m] != zero(T))
-                end
-            end
-        end
-    end
-    vals = []
-    for m in 1:M
-        for i in Ito
-            val = :(zero(T))
-            for n in 1:N
-                if i == Ifrom[n]
-                    val = :(x[$n, $m])
-                    break
-                end
-            end
-            push!(vals, val)
-        end
-    end
-    quote
-        Base.@_propagate_inbounds_meta
-        if assert_exact_transform()
-            if $errcond
-                throw(InexactError(:transform, Ato, x))
-            end
-        end
-        @inbounds Axis2Tensor(
-            (ato, axes(x, 2)),
-            SMatrix{$(length(Ito)), $M}($(vals...)),
+@inline transform(ato::CovariantAxis, v::CovariantTensor) = project(ato, v)
+@inline transform(ato::ContravariantAxis, v::ContravariantTensor) =
+    project(ato, v)
+@inline transform(ato::CartesianAxis, v::CartesianTensor) = project(ato, v)
+@inline transform(ato::LocalAxis, v::LocalTensor) = project(ato, v)
+
+@inline function project(
+    ato_l::AbstractAxis,
+    v::Axis2Tensor,
+    ato_r::AbstractAxis,
+)
+    @assert symbols.(axes(v)) == symbols.((ato_l, ato_r)) "Axes do not match"
+    T = eltype(v)
+    Ifrom_l, Ifrom_r = axis_indices.(axes(v))
+    product_to = unrolled_product(axis_indices(ato_l), axis_indices(ato_r))
+    vals = unrolled_map(product_to) do (m_l, m_r)
+        n_l = unrolled_findfirst(
+            ((n_l, ifrom),) -> m_l == ifrom,
+            enumerate(Ifrom_l),
         )
-    end
-end
-
-@generated function _project(
-    ato::Ato,
-    x::Axis2Tensor{T, Tuple{Afrom, A2}},
-) where {
-    Ato <: AbstractAxis{Ito},
-    Afrom <: AbstractAxis{Ifrom},
-    A2 <: AbstractAxis{J},
-} where {Ito, Ifrom, J, T}
-    N = length(Ifrom)
-    M = length(J)
-    vals = []
-    for m in 1:M
-        for i in Ito
-            val = :(zero(T))
-            for n in 1:N
-                if i == Ifrom[n]
-                    val = :(x[$n, $m])
-                    break
-                end
-            end
-            push!(vals, val)
-        end
-    end
-    quote
-        Base.@_propagate_inbounds_meta
-        @inbounds Axis2Tensor(
-            (ato, axes(x, 2)),
-            SMatrix{$(length(Ito)), $M}($(vals...)),
+        n_r = unrolled_findfirst(
+            ((n_r, ifrom),) -> m_r == ifrom,
+            enumerate(Ifrom_r),
         )
+        isnothing(n_l) || isnothing(n_r) ? zero(T) : v[n_l, n_r]
     end
+    S = SMatrix{length(ato_l), length(ato_r)}(vals...)
+    @inbounds Axis2Tensor((ato_l, ato_r), S)
+end
+@inline project(ato::AbstractAxis, v::Axis2Tensor) = project(ato, v, axes(v, 2))
+@inline project(v::Axis2Tensor, ato::AbstractAxis) = project(axes(v, 1), v, ato)
+@inline function project(ato_l, v::AxisVector)
+    @assert symbols(axes(v, 1)) == symbols(ato_l)
+    _project(ato_l, v)
 end
-
-@inline transform(ato::CovariantAxis, v::CovariantTensor) = _project(ato, v)
-@inline transform(ato::ContravariantAxis, v::ContravariantTensor) =
-    _project(ato, v)
-@inline transform(ato::CartesianAxis, v::CartesianTensor) = _project(ato, v)
-@inline transform(ato::LocalAxis, v::LocalTensor) = _project(ato, v)
-
-@inline project(ato::CovariantAxis, v::CovariantTensor) = _project(ato, v)
-@inline project(ato::ContravariantAxis, v::ContravariantTensor) =
-    _project(ato, v)
-@inline project(ato::CartesianAxis, v::CartesianTensor) = _project(ato, v)
-@inline project(ato::LocalAxis, v::LocalTensor) = _project(ato, v)
-
 
 """
     outer(x, y)

test/Geometry/axistensors.jl

@@ -2,7 +2,6 @@ using Test, JET
 using ClimaCore.Geometry, ClimaCore.DataLayouts
 using LinearAlgebra, StaticArrays
 import ClimaCore
-ClimaCore.Geometry.assert_exact_transform() = true
 
 @testset "AxisTensors" begin
     x = Geometry.Covariant12Vector(1.0, 2.0)
@@ -169,6 +168,26 @@ end
         Geometry.Covariant12Axis(),
         Geometry.Covariant13Vector(2.0, 2.0) * Geometry.Cartesian1Vector(1.0)',
     ) == Geometry.Covariant12Vector(2.0, 0.0) * Geometry.Cartesian1Vector(1.0)'
+
+    # Test projection over rightmost axis
+    x_C12 = Geometry.Covariant12Vector(2.0, 2.0)
+    x_Cart123 = Geometry.Cartesian123Vector(1.0, 1.0, 1.0)
+    @test Geometry.project(x_C12 * x_Cart123', Geometry.Cartesian3Axis()) ==
+          x_C12 * Geometry.Cartesian3Vector(1.0)'
+    @test Geometry.project(x_C12 * x_Cart123', Geometry.Cartesian23Axis()) ==
+          x_C12 * Geometry.Cartesian23Vector(1.0, 1.0)'
+
+    # Test projection over both axes
+    @test Geometry.project(
+        Geometry.Covariant12Axis(),
+        x_C12 * x_Cart123',
+        Geometry.Cartesian123Axis(),
+    ) == x_C12 * x_Cart123'
+    @test Geometry.project(
+        Geometry.Covariant2Axis(),
+        x_C12 * x_Cart123',
+        Geometry.Cartesian13Axis(),
+    ) == Geometry.Covariant2Vector(2.0) * Geometry.Cartesian13Vector(1.0, 1.0)'
 end