add transform function

examples/column/hydrostatic_ekman.jl

@@ -124,14 +124,15 @@ function tendency!(dY, Y, _, t)
 
     uv_1 = Operators.getidx(uv, Operators.Interior(), 1)
     u_wind = LinearAlgebra.norm(uv_1)
+
     A = Operators.AdvectionC2C(
         bottom = Operators.SetValue(Geometry.Cartesian12Vector(0.0, 0.0)),
         top = Operators.SetValue(Geometry.Cartesian12Vector(0.0, 0.0)),
     )
 
     # uv
     bcs_bottom =
-        Operators.SetValue(Geometry.Cartesian3Vector(Cd * u_wind * uv_1))
+        Operators.SetValue(Geometry.Cartesian3Vector(Cd * u_wind) ⊗ uv_1)
     bcs_top = Operators.SetValue(uvg)
     ∂c = Operators.DivergenceF2C(bottom = bcs_bottom)
     ∂f = Operators.GradientC2F(top = bcs_top)

src/Geometry/axistensors.jl

@@ -176,13 +176,17 @@ import Base: +, -, *, /, \, ==
     axes(a) == axes(b) && components(a) == components(b)
 
 # vectors
-
 const AxisVector{T, A1, S} = AxisTensor{T, 1, Tuple{A1}, S}
+
 AxisVector(ax::A1, v::SVector{N, T}) where {A1 <: AbstractAxis, N, T} =
     AxisVector{T, A1, SVector{N, T}}((ax,), v)
 
-(AxisVector{T, A, SVector{N, T}} where {T})(args...) where {A, N} =
-    AxisVector(A.instance, SVector(args))
+(AxisVector{T, A, SVector{1, T}} where {T})(arg1) where {A} =
+    AxisVector(A.instance, SVector(arg1))
+(AxisVector{T, A, SVector{2, T}} where {T})(arg1, arg2) where {A} =
+    AxisVector(A.instance, SVector(arg1, arg2))
+(AxisVector{T, A, SVector{3, T}} where {T})(arg1, arg2, arg3) where {A} =
+    AxisVector(A.instance, SVector(arg1, arg2, arg3))
 
 const CovariantVector{T, A1 <: CovariantAxis, S} = AxisVector{T, A1, S}
 const ContravariantVector{T, A1 <: ContravariantAxis, S} = AxisVector{T, A1, S}
@@ -215,7 +219,17 @@ components(va::AdjointAxis2Tensor) = components(parent(va))'
 const Axis2TensorOrAdj{T, A, S} =
     Union{Axis2Tensor{T, A, S}, AdjointAxis2Tensor{T, A, S}}
 
+# based on 1st dimension
+const Covariant2Tensor{T, A, S} =
+    Axis2Tensor{T, A, S} where {A <: Tuple{CovariantAxis, AbstractAxis}}
+const Contravariant2Tensor{T, A, S} =
+    Axis2Tensor{T, A, S} where {A <: Tuple{ContravariantAxis, AbstractAxis}}
+const Cartesian2Tensor{T, A, S} =
+    Axis2Tensor{T, A, S} where {A <: Tuple{CartesianAxis, AbstractAxis}}
 
+const CovariantTensor = Union{CovariantVector, Covariant2Tensor}
+const ContravariantTensor = Union{ContravariantVector, Contravariant2Tensor}
+const CartesianTensor = Union{CartesianVector, Cartesian2Tensor}
 
 for I in [(1,), (2,), (3,), (1, 2), (1, 3), (2, 3), (1, 2, 3)]
     strI = join(I)
@@ -349,3 +363,103 @@ end
 function outer(x::AbstractVector, y)
     RecursiveApply.rmap(y -> x ⊗ y, y)
 end
+
+function _transform(
+    ato::Ato,
+    x::AxisVector{T, Afrom, SVector{N, T}},
+) where {Ato <: AbstractAxis{I}, Afrom <: AbstractAxis{I}} where {I, T, N}
+    x
+end
+@generated function _transform(
+    ato::Ato,
+    x::AxisVector{T, Afrom, SVector{N, T}},
+) where {
+    Ato <: AbstractAxis{Ito},
+    Afrom <: AbstractAxis{Ifrom},
+} where {Ito, Ifrom, T, N}
+    errcond = false
+    for n in 1:N
+        i = Ifrom[n]
+        if i ∉ Ito
+            errcond = :($errcond || x[$n] != 0)
+        end
+    end
+    vals = []
+    for i in Ito
+        val = :(zero(T))
+        for n in 1:N
+            if i == Ifrom[n]
+                val = :(x[$n])
+                break
+            end
+        end
+        push!(vals, val)
+    end
+    quote
+        if $errcond
+            throw(InexactError(:transform, Ato, x))
+        end
+        AxisVector(ato, SVector($(vals...)))
+    end
+end
+
+function _transform(
+    ato::Ato,
+    x::Axis2Tensor{T, Tuple{Afrom, A2}},
+) where {
+    Ato <: AbstractAxis{I},
+    Afrom <: AbstractAxis{I},
+    A2 <: AbstractAxis{J},
+} where {I, J, T}
+    x
+end
+
+@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)
+    errcond = false
+    for n in 1:N
+        i = Ifrom[n]
+        if i ∉ Ito
+            for m in 1:M
+                errcond = :($errcond || x[$n, $m] != 0)
+            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
+        if $errcond
+            throw(InexactError(:transform, Ato, x))
+        end
+        Axis2Tensor((ato, axes(x, 2)), SMatrix{$(length(Ito)), $M}($(vals...)))
+    end
+end
+
+function transform(ato::CovariantAxis, v::CovariantTensor)
+    _transform(ato, v)
+end
+function transform(ato::ContravariantAxis, v::ContravariantTensor)
+    _transform(ato, v)
+end
+function transform(ato::CartesianAxis, v::CartesianTensor)
+    _transform(ato, v)
+end

src/Geometry/conversions.jl

@@ -44,7 +44,7 @@ Covariant12Vector(u::CartesianVector, local_geometry::LocalGeometry) =
     local_geometry.∂x∂ξ' * u
 Covariant12Vector(u::ContravariantVector, local_geometry::LocalGeometry) =
     local_geometry.∂x∂ξ' * local_geometry.∂x∂ξ * u
-
+3
 Cartesian12Vector(u::CartesianVector, local_geometry::LocalGeometry) = u
 Cartesian12Vector(u::CovariantVector, local_geometry::LocalGeometry) =
     local_geometry.∂ξ∂x' * u
@@ -59,30 +59,23 @@ covariant2(u::AxisVector, local_geometry::LocalGeometry) =
 covariant3(u::AxisVector, local_geometry::LocalGeometry) =
     CovariantVector(u, local_geometry).u₃
 
+# TODO: specialize?
 contravariant1(u::AxisVector, local_geometry::LocalGeometry) =
-    ContravariantVector(u, local_geometry).u¹
+    transform(Contravariant123Axis(), u, local_geometry).u¹
 contravariant2(u::AxisVector, local_geometry::LocalGeometry) =
-    ContravariantVector(u, local_geometry).u²
+    transform(Contravariant123Axis(), u, local_geometry).u²
 contravariant3(u::AxisVector, local_geometry::LocalGeometry) =
-    ContravariantVector(u, local_geometry).u³
-
-Jcontravariant3(u::AxisVector, local_geometry::LocalGeometry) =
-    local_geometry.J * contravariant3(u, local_geometry)
+    transform(Contravariant123Axis(), u, local_geometry).u³
 
+contravariant1(u::Axis2Tensor, local_geometry::LocalGeometry) =
+    transform(Contravariant123Axis(), u, local_geometry)[1, :]
+contravariant2(u::Axis2Tensor, local_geometry::LocalGeometry) =
+    transform(Contravariant123Axis(), u, local_geometry)[2, :]
+contravariant3(u::Axis2Tensor, local_geometry::LocalGeometry) =
+    transform(Contravariant123Axis(), u, local_geometry)[3, :]
 
-contravariant1(
-    A::Axis2Tensor{<:Any, Tuple{Cartesian1Axis, Cartesian1Axis}},
-    local_geometry::LocalGeometry,
-) = (local_geometry.∂ξ∂x * A)[1, :]
-contravariant1(
-    A::Axis2Tensor{<:Any, Tuple{Cartesian12Axis, Cartesian12Axis}},
-    local_geometry::LocalGeometry,
-) = (local_geometry.∂ξ∂x * A)[1, :]
-contravariant2(
-    A::Axis2Tensor{<:Any, Tuple{Cartesian12Axis, Cartesian12Axis}},
-    local_geometry::LocalGeometry,
-) = (local_geometry.∂ξ∂x * A)[2, :]
-
+Jcontravariant3(u::AxisTensor, local_geometry::LocalGeometry) =
+    local_geometry.J * contravariant3(u, local_geometry)
 
 
 # conversions
@@ -94,6 +87,91 @@ function Covariant3Vector(
     Covariant3Vector(uⁱ.u³)
 end
 
+"""
+    transform(axis, V[, local_geometry])
+
+Transform the first axis of the vector or tensor `V` to `axis`.
+"""
+function transform end
+
+# Covariant <-> Cartesian
+function transform(
+    ax::CartesianAxis,
+    v::CovariantTensor,
+    local_geometry::LocalGeometry,
+)
+    transform(
+        ax,
+        local_geometry.∂ξ∂x' * transform(dual(axes(local_geometry.∂ξ∂x, 1)), v),
+    )
+end
+function transform(
+    ax::CovariantAxis,
+    v::CartesianTensor,
+    local_geometry::LocalGeometry,
+)
+    transform(
+        ax,
+        local_geometry.∂x∂ξ' * transform(dual(axes(local_geometry.∂x∂ξ, 1)), v),
+    )
+end
+
+# Contravariant <-> Cartesian
+function transform(
+    ax::ContravariantAxis,
+    v::CartesianTensor,
+    local_geometry::LocalGeometry,
+)
+    transform(
+        ax,
+        local_geometry.∂ξ∂x * transform(dual(axes(local_geometry.∂ξ∂x, 2)), v),
+    )
+end
+function transform(
+    ax::CartesianAxis,
+    v::ContravariantTensor,
+    local_geometry::LocalGeometry,
+)
+    transform(
+        ax,
+        local_geometry.∂x∂ξ * transform(dual(axes(local_geometry.∂x∂ξ, 2)), v),
+    )
+end
+
+# Covariant <-> Contravariant
+function transform(
+    ax::ContravariantAxis,
+    v::CovariantTensor,
+    local_geometry::LocalGeometry,
+)
+    transform(
+        ax,
+        local_geometry.∂ξ∂x *
+        local_geometry.∂ξ∂x' *
+        transform(dual(axes(local_geometry.∂ξ∂x, 1)), v),
+    )
+end
+function transform(
+    ax::CovariantAxis,
+    v::ContravariantTensor,
+    local_geometry::LocalGeometry,
+)
+    transform(
+        ax,
+        local_geometry.∂x∂ξ' *
+        local_geometry.∂x∂ξ *
+        transform(dual(axes(local_geometry.∂x∂ξ, 2)), v),
+    )
+end
+
+transform(ato::CovariantAxis, v::CovariantTensor, ::LocalGeometry) =
+    transform(ato, v)
+transform(ato::ContravariantAxis, v::ContravariantTensor, ::LocalGeometry) =
+    transform(ato, v)
+transform(ato::CartesianAxis, v::CartesianTensor, ::LocalGeometry) =
+    transform(ato, v)
+
+
 
 """
     divergence_result_type(V)
@@ -107,10 +185,10 @@ divergence_result_type(
     ::Type{Axis2Tensor{FT, Tuple{A1, A2}, S}},
 ) where {
     FT,
-    A1 <: CartesianAxis,
+    A1,
     A2 <: CartesianAxis,
     S <: StaticMatrix{S1, S2},
-} where {S1, S2} = AxisVector{FT, A1, SVector{S1, FT}}
+} where {S1, S2} = AxisVector{FT, A2, SVector{S2, FT}}
 
 curl_result_type(::Type{V}) where {V <: Covariant12Vector{FT}} where {FT} =
     Contravariant3Vector{FT}