Add sketch of new Tensor interface

.buildkite/Manifest.toml

@@ -1,6 +1,6 @@
 # This file is machine-generated - editing it directly is not advised
 
-julia_version = "1.10.8"
+julia_version = "1.10.10"
 manifest_format = "2.0"
 project_hash = "1df235e0c8e5246d3a9cdc32ad82313e2fd7f343"
 
@@ -373,7 +373,7 @@ weakdeps = ["CUDA", "MPI"]
     ClimaCommsMPIExt = "MPI"
 
 [[deps.ClimaCore]]
-deps = ["Adapt", "BandedMatrices", "BlockArrays", "ClimaComms", "ClimaInterpolations", "CubedSphere", "DataStructures", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LLVM", "LazyBroadcast", "LinearAlgebra", "MultiBroadcastFusion", "NVTX", "PkgVersion", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "StaticArrays", "Statistics", "UnrolledUtilities"]
+deps = ["Adapt", "BandedMatrices", "BlockArrays", "ClimaComms", "ClimaInterpolations", "CubedSphere", "DataStructures", "ForwardDiff", "GaussQuadrature", "GilbertCurves", "HDF5", "InteractiveUtils", "IntervalSets", "KrylovKit", "LLVM", "LazyBroadcast", "LinearAlgebra", "MultiBroadcastFusion", "NVTX", "PkgVersion", "Random", "RecursiveArrayTools", "RootSolvers", "SparseArrays", "StaticArrays", "Statistics", "UnrolledUtilities"]
 path = ".."
 uuid = "d414da3d-4745-48bb-8d80-42e94e092884"
 version = "0.14.51"
@@ -1707,7 +1707,7 @@ version = "0.3.23+4"
 [[deps.OpenLibm_jll]]
 deps = ["Artifacts", "Libdl"]
 uuid = "05823500-19ac-5b8b-9628-191a04bc5112"
-version = "0.8.1+2"
+version = "0.8.5+0"
 
 [[deps.OpenMPI_jll]]
 deps = ["Artifacts", "CompilerSupportLibraries_jll", "Hwloc_jll", "JLLWrappers", "LazyArtifacts", "Libdl", "MPIPreferences", "TOML", "Zlib_jll"]

.buildkite/perf/pipeline.yml

@@ -10,7 +10,7 @@ env:
   GKSwstype: 100
   SLURM_KILL_BAD_EXIT: 1
   COUPLER_PATH: ".buildkite/perf/ClimaCoupler.jl"
-  COUPLER_COMMIT: "7f4e6f02ce990bdd6d42c9db3a11abf4b158b610"
+  COUPLER_COMMIT: "4508dca4d2fb8380b49effdd2864c86f1e87d291"
   CONFIG_PATH: "$COUPLER_PATH/config/nightly_configs"
   BENCHMARK_CONFIG_PATH: "$COUPLER_PATH/config/benchmark_configs"
 
@@ -25,7 +25,8 @@ steps:
       - echo "--- Instantiate AMIP env"
       # Instantiate and dev install ClimaCore
       - "julia --project=$COUPLER_PATH/experiments/AMIP/ -e 'using Pkg; Pkg.instantiate(;verbose=true)'"
-      - 'julia --project=$COUPLER_PATH/experiments/AMIP/ -e ''using Pkg; Pkg.develop(path=".")'''
+      - "julia --project=$COUPLER_PATH/experiments/AMIP/ -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name=\"ClimaCore\", rev=\"main\"))'"
+      # - "julia --project=$COUPLER_PATH/experiments/AMIP/ -e 'using Pkg; Pkg.add(Pkg.PackageSpec(;name=\"ClimaAtmos\", rev=\"ne/compat\"))'"
       - 'julia --project=$COUPLER_PATH/experiments/AMIP/ -e ''using Pkg; Pkg.add("MPI")'''
       - "julia --project=$COUPLER_PATH/experiments/AMIP/ -e 'using Pkg; Pkg.precompile()'"
       - "julia --project=$COUPLER_PATH/experiments/AMIP/ -e 'using Pkg; Pkg.status()'"

.buildkite/pipeline.yml

@@ -251,10 +251,10 @@ steps:
         retry: *retry_policy
         command: "julia --color=yes --check-bounds=yes --project=.buildkite test/Geometry/geometry.jl"
 
-      - label: "Unit: axistensors"
-        key: unit_axistensors
+      - label: "Unit: tensors"
+        key: unit_tensors
         retry: *retry_policy
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/Geometry/axistensors.jl"
+        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/Geometry/tensors.jl"
 
       - label: "Unit: mul_with_projection"
         key: unit_mul_with_projection
@@ -1079,12 +1079,12 @@ steps:
         agents:
           slurm_gpus: 1
 
-      - label: "Unit: field_matrix_indexing (CPU)"
+      - label: "Unit: field matrix indexing (CPU)"
         key: cpu_field_matrix_indexing
         retry: *retry_policy
         command: "julia --color=yes --check-bounds=yes --project=.buildkite test/MatrixFields/field_matrix_indexing.jl"
 
-      - label: "Unit: field_matrix_indexing (GPU)"
+      - label: "Unit: field matrix indexing (GPU)"
         key: gpu_field_matrix_indexing
         retry: *retry_policy
         command: "julia --color=yes --project=.buildkite test/MatrixFields/field_matrix_indexing.jl"
@@ -1616,7 +1616,7 @@ steps:
       - label: "Perf: Axis tensor conversion benchmarks"
         key: "cpu_axis_tensor_conversion_perf_bm"
         retry: *retry_policy
-        command: "julia --color=yes --project=.buildkite test/Geometry/axistensor_conversion_benchmarks.jl"
+        command: "julia --color=yes --project=.buildkite test/Geometry/tensor_conversion_benchmarks.jl"
 
   - group: "Perf: DataLayouts"
     steps:
@@ -2489,9 +2489,10 @@ steps:
         command:
           - "julia --color=yes --project=.buildkite examples/hybrid/driver.jl"
         artifact_paths:
-          - "examples/hybrid/sphere/output/held_suarez_rhoe/Float32/*"
+          - "examples/hybrid/sphere/output/held_suarez_rhoe/Float64/*"
         env:
           TEST_NAME: "sphere/held_suarez_rhoe"
+          FLOAT_TYPE: "Float64"
 
       - label: ":computer: Float64 3D sphere dry Held-Suarez (ρθ)"
         key: "cpu_held_suarez_rho_theta_float64"

.github/workflows/ClimaCoreVTK.yml

@@ -27,7 +27,8 @@ jobs:
           sudo apt-get update && sudo apt-get -y install paraview python3-paraview
       - name: Install Julia dependencies
         run: |
-          julia --project=monorepo -e 'using Pkg; Pkg.develop(path="$(pwd())/lib/ClimaCoreVTK")'
+          julia --project=monorepo -e 'using Pkg; Pkg.develop(path="lib/ClimaCoreVTK")'
+          julia --project=monorepo -e 'using Pkg; Pkg.develop(path=".")'
       - name: Run the tests
         env:
           CI_OUTPUT_DIR: output

.github/workflows/docs.yml

@@ -16,7 +16,10 @@ jobs:
         with:
           version: '1.10'
       - name: Install dependencies
-        run: julia --project=docs -e 'using Pkg; Pkg.develop(path="."); Pkg.instantiate(;verbose=true)'
+        run: >
+          julia --project=docs -e 'using Pkg; Pkg.develop(path=".");
+          Pkg.develop(path="lib/ClimaCoreVTK"); 
+          Pkg.instantiate(;verbose=true)'
       - name: Build and deploy
         env:
           GKSwstype: "100" # headless GR: https://discourse.julialang.org/t/generation-of-documentation-fails-qt-qpa-xcb-could-not-connect-to-display/60988/2

Project.toml

@@ -24,6 +24,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MultiBroadcastFusion = "c3c07f87-98de-43f2-a76f-835b330b2cbb"
 NVTX = "5da4648a-3479-48b8-97b9-01cb529c0a1f"
 PkgVersion = "eebad327-c553-4316-9ea0-9fa01ccd7688"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 RootSolvers = "7181ea78-2dcb-4de3-ab41-2b8ab5a31e74"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"

docs/src/matrix_fields.md

@@ -108,7 +108,7 @@ scalar_field_matrix
 A FieldMatrix entry can be:
 
 - A `UniformScaling`, which contains a `Number`
-- A `DiagonalMatrixRow`, which can contain either a `Number` or a tensor (represented as a `Geometry.Axis2Tensor`)
+- A `DiagonalMatrixRow`, which can contain either a `Number` or a rank-2 tensor (represented as a `Geometry.Tensor{2}`)
 - A `ColumnwiseBandMatrixField`, where each value is a [`BandMatrixRow`](@ref) with entries of any type that can be represented using the field's base number type.
 
 If an entry contains a composite type, the fields of that type can be extracted.
@@ -173,8 +173,8 @@ The recursive indexing of an internal entry given some entry `entry` and interna
 works as follows:
 
 1. If the  `internal_name_pair` is blank, return `entry`
-2. If the element type of each band of `entry` is an `Axis2Tensor`, and `internal_name_pair` is of the form `(@name(components.data.1...), @name(components.data.2...))` (potentially with different numbers), then extract the specified component, and recurse on it with the remaining `internal_name_pair`.
-3. If the element type of each band of `entry` is a `Geometry.AdjointAxisVector`, then recurse on the parent of the adjoint.
+2. If the element type of each band of `entry` is a `Geometry.Tensor{2}`, and `internal_name_pair` is of the form `(@name(components.data.1...), @name(components.data.2...))` (potentially with different numbers), then extract the specified component, and recurse on it with the remaining `internal_name_pair`.
+3. If the element type of each band of `entry` is an `Adjoint` of a `Geometry.AbstractTensor{1}` (an adjoint rank-1 tensor), then recurse on the parent of the adjoint.
 4. If `internal_name_pair[1]` is not empty, and the first name in it is a field of the element type of each band of `entry`, extract that field from `entry`, and recurse into it with the remaining names of `internal_name_pair[1]` and all of `internal_name_pair[2]`
 5. If `internal_name_pair[2]` is not empty, and the first name in it is a field of the element type of each band of `entry`, extract that field from `entry`, and recurse into it with all of `internal_name_pair[1]` and the remaining names of `internal_name_pair[2]`
 6. At this point, if none of the previous cases are true, both `internal_name_pair[1]` and `internal_name_pair[2]` should be non-empty, and it is assumed that `entry` is being used to implicitly represent some tensor structure. If the first name in `internal_name_pair[1]` is equivalent to `internal_name_pair[2]`, then both the first names are dropped, and entry is recursed onto. If the first names are different, both the first names are dropped, and the zero of entry is recursed onto.

examples/hybrid/box/bubble_3d_rhotheta.jl

@@ -253,9 +253,9 @@ function rhs!(dY, Y, _, t)
     # Horizontal momentum
     @. dρuₕ += -uvdivf2c(ρw ⊗ If(uₕ))
     Ih = Ref(
-        Geometry.Axis2Tensor(
+        Geometry.Tensor(
+            I,
             (Geometry.UVAxis(), Geometry.UVAxis()),
-            @SMatrix [1.0 0.0; 0.0 1.0]
         ),
     )
     @. dρuₕ -= hdiv(ρuₕ ⊗ uₕ + p * Ih)

examples/hybrid/implicit_equation_jacobian.jl

@@ -62,14 +62,14 @@ function ImplicitEquationJacobian(Y, transform, flags = (;))
     ᶠ𝕄_name = @name(f.w)
 
     BidiagonalRow_C3 = BidiagonalMatrixRow{C3{FT}}
-    BidiagonalRow_ACT3 = BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}}
-    QuaddiagonalRow_ACT3 = QuaddiagonalMatrixRow{Adjoint{FT, CT3{FT}}}
+    BidiagonalRow_ACT3 = BidiagonalMatrixRow{typeof(CT3(FT(0))')}
+    QuaddiagonalRow_ACT3 = QuaddiagonalMatrixRow{typeof(CT3(FT(0))')}
     TridiagonalRow_C3xACT3 =
         TridiagonalMatrixRow{typeof(C3(FT(0)) * CT3(FT(0))')}
     ∂ᶜ𝔼ₜ∂ᶠ𝕄_Row_ACT3 =
         flags.∂ᶜ𝔼ₜ∂ᶠ𝕄_mode == :exact && :ρe in propertynames(Y.c) ?
         QuaddiagonalRow_ACT3 : BidiagonalRow_ACT3
-    ∂Yₜ∂Y = FieldMatrix(
+    ∂Yₜ∂Y = MatrixFields.FieldMatrix(
         (ᶜρ_name, ᶠ𝕄_name) => zeros(BidiagonalRow_ACT3, axes(Y.c)),
         (ᶜ𝔼_name, ᶠ𝕄_name) => zeros(∂ᶜ𝔼ₜ∂ᶠ𝕄_Row_ACT3, axes(Y.c)),
         (ᶠ𝕄_name, ᶜρ_name) => zeros(BidiagonalRow_C3, axes(Y.f)),

examples/hybrid/plane/density_current_2d_flux_form.jl

@@ -241,10 +241,7 @@ function rhs!(dY, Y, _, t)
 
     # horizontal momentum
     Ih = Ref(
-        Geometry.Axis2Tensor(
-            (Geometry.UAxis(), Geometry.UAxis()),
-            @SMatrix [1.0]
-        ),
+        Geometry.Tensor(I, (Geometry.UAxis(), Geometry.UAxis())),
     )
     @. dρuₕ += -uvdivf2c(ρw ⊗ If(uₕ))
     @. dρuₕ -= hdiv(ρuₕ ⊗ uₕ + p * Ih)

examples/hybrid/plane/inertial_gravity_wave.jl

@@ -109,8 +109,8 @@ function center_initial_condition(ᶜlocal_geometry)
     if is_discrete_hydrostatic_balance
         face_space =
             Spaces.FaceExtrudedFiniteDifferenceSpace(axes(ᶜlocal_geometry))
-        ᶠΔz = Fields.local_geometry_field(face_space).∂x∂ξ.components.data.:4
-        ᶜΔz = ᶜlocal_geometry.∂x∂ξ.components.data.:4
+        ᶠΔz = Fields.local_geometry_field(face_space).metric.tensor.components.data.:4
+        ᶜΔz = ᶜlocal_geometry.metric.tensor.components.data.:4
         ᶜp = discrete_hydrostatic_balance!(ᶠΔz, ᶜΔz, grav)
     else
         ᶜp = @. p₀(ᶜz)

examples/hybrid/plane/isothermal_channel.jl

@@ -121,7 +121,7 @@ gⁱʲ =
     ).gⁱʲ
 g13 = gⁱʲ.components.data.:3
 g11 = gⁱʲ.components.data.:1
-g33 = gⁱʲ.components.data.:4
+g33 = gⁱʲ.components.data.:9
 u₃_bc = Geometry.Covariant3Vector.(-1 .* g13 .* u₁_bc.components.data.:1 ./ g33)
 apply_boundary_w =
     Operators.SetBoundaryOperator(bottom = Operators.SetValue(u₃_bc))
@@ -175,7 +175,7 @@ function rhs_invariant!(dY, Y, _, t)
         ).gⁱʲ
     g13 = gⁱʲ.components.data.:3
     g11 = gⁱʲ.components.data.:1
-    g33 = gⁱʲ.components.data.:4
+    g33 = gⁱʲ.components.data.:9
     u₃_bc =
         Geometry.Covariant3Vector.(-1 .* g13 .* u₁_bc.components.data.:1 ./ g33)
     apply_boundary_w =

examples/hybrid/plane/topo_schar_nh.jl

@@ -258,7 +258,7 @@ function rhs_invariant!(dY, Y, _, t)
             Fields.local_geometry_field(hv_face_space),
             ClimaCore.Utilities.half,
         ).gⁱʲ
-    g33 = gⁱʲ.components.data.:4
+    g33 = gⁱʲ.components.data.:9
     u₃_bc = Geometry.Covariant3Vector.(-1 .* u₁_bc ./ g33)  # fw = -g^31 cuₕ/ g^33
     apply_boundary_w =
         Operators.SetBoundaryOperator(bottom = Operators.SetValue(u₃_bc))

examples/hybrid/staggered_nonhydrostatic_model.jl

@@ -1,4 +1,4 @@
-using LinearAlgebra: ×, norm, norm_sqr, dot, Adjoint
+using LinearAlgebra: ×, norm, norm_sqr, dot
 using ClimaCore: Operators, Fields
 
 include("implicit_equation_jacobian.jl")
@@ -120,7 +120,7 @@ function default_cache(ᶜlocal_geometry, ᶠlocal_geometry, Y, upwinding_mode)
         ᶜf,
         ∂ᶜK∂ᶠw = similar(
             ᶜlocal_geometry,
-            BidiagonalMatrixRow{Adjoint{FT, CT3{FT}}},
+            BidiagonalMatrixRow{typeof(CT3(FT(0))')},
         ),
         ᶠupwind_product,
         ᶠupwind_product_matrix,
@@ -310,11 +310,11 @@ function implicit_equation_jacobian!(j, Y, p, δtγ, t)
     ∂ᶠ𝕄ₜ∂ᶠ𝕄 = ∂Yₜ∂Y[ᶠ𝕄_name, ᶠ𝕄_name]
 
     ᶠgⁱʲ = Fields.local_geometry_field(ᶠw).gⁱʲ
-    g³³(gⁱʲ) = Geometry.AxisTensor(
-        (Geometry.Contravariant3Axis(), Geometry.Contravariant3Axis()),
-        Geometry.components(gⁱʲ)[end],
+    g³³(gⁱʲ) = reshape(
+        gⁱʲ,
+        Geometry.Contravariant3Axis(),
+        Geometry.Contravariant3Axis(),
     )
-
     # If ∂(ᶜχ)/∂(ᶠw) = 0, then
     # ∂(ᶠupwind_product(ᶠw, ᶜχ))/∂(ᶠw) =
     #     ∂(ᶠupwind_product(ᶠw, ᶜχ))/∂(CT3(ᶠw)) * ∂(CT3(ᶠw))/∂(ᶠw) =

ext/ClimaCoreCUDAExt.jl

@@ -5,7 +5,7 @@ import ClimaCore.Limiters
 import ClimaComms
 import ClimaCore: DataLayouts, Grids, Spaces, Fields
 import ClimaCore: Geometry
-import ClimaCore.Geometry: AxisTensor
+import ClimaCore.Geometry: AbstractTensor
 import CUDA
 using CUDA
 using CUDA: threadIdx, blockIdx, blockDim

ext/cuda/column_matrix_helpers.jl

@@ -427,4 +427,4 @@ Base.@propagate_inbounds outer_or_mul(x::T1, y::T2) where {T1 <: AbstractVector,
 Base.@propagate_inbounds outer_or_mul(
     x::T1,
     y::T2,
-) where {T1 <: Geometry.AdjointAxisVector, T2 <: Geometry.Axis2Tensor} = (x * y)'
+) where {T1 <: Geometry.AbstractCovector, T2 <: Geometry.Tensor{2}} = (x * y)'

ext/cuda/operators_fd_eager.jl

@@ -507,7 +507,7 @@ Base.@propagate_inbounds recursively_project(projection_tuple::T, y::Y) where {T
 Base.@propagate_inbounds recursively_project(
     projection_tuple::T,
     y::Y,
-) where {T, Y <: AxisTensor} =
+) where {T, Y <: AbstractTensor} =
     @inbounds @inline project(projection_tuple[1], y, projection_tuple[2])
 
 if hasfield(Method, :recursion_relation)

ext/cuda/operators_spectral_element.jl

@@ -363,10 +363,10 @@ Base.@propagate_inbounds function operator_evaluate(
     end
     if eltype(input) <: Number
         return Geometry.Covariant1Vector(∂f∂ξ₁)
-    elseif eltype(input) <: Geometry.AxisVector
-        tensor_axes = (Geometry.Covariant1Axis(), axes(eltype(input))[1])
-        tensor_components = hcat(Geometry.components(∂f∂ξ₁))'
-        return Geometry.AxisTensor(tensor_axes, tensor_components)
+    elseif eltype(input) <: Geometry.AbstractTensor{1}
+        tensor_axes = (Geometry.Covariant1Axis(), Geometry.tensor_bases(eltype(input))[1])
+        tensor_components = hcat(parent(∂f∂ξ₁))'
+        return Geometry.Tensor(tensor_components, tensor_axes)
     else
         error("Unsupported input type for gradient operator: $(eltype(input))")
     end
@@ -396,11 +396,11 @@ Base.@propagate_inbounds function operator_evaluate(
     end
     if eltype(input) <: Number
         return Geometry.Covariant12Vector(∂f∂ξ₁, ∂f∂ξ₂)
-    elseif eltype(input) <: Geometry.AxisVector
-        tensor_axes = (Geometry.Covariant12Axis(), axes(eltype(input))[1])
+    elseif eltype(input) <: Geometry.AbstractTensor{1}
+        tensor_axes = (Geometry.Covariant12Axis(), Geometry.tensor_bases(eltype(input))[1])
         tensor_components =
-            hcat(Geometry.components(∂f∂ξ₁), Geometry.components(∂f∂ξ₂))'
-        return Geometry.AxisTensor(tensor_axes, tensor_components)
+            hcat(parent(∂f∂ξ₁), parent(∂f∂ξ₂))'
+        return Geometry.Tensor(tensor_components, tensor_axes)
     else
         error("Unsupported input type for gradient operator: $(eltype(input))")
     end

lib/ClimaCoreVTK/src/addfield.jl

@@ -53,7 +53,7 @@ function addfield!(
     field,
     dataspace,
     ::Type{T},
-) where {T <: Geometry.AxisVector}
+) where {T <: Geometry.AbstractTensor{1}}
     interp = Operators.Interpolate(dataspace)
     # should we convert then interpolate, or vice versa?
     ifield = interp.(Geometry.Cartesian123Vector.(field))

src/Fields/broadcast.jl

@@ -423,28 +423,12 @@ function Base.Broadcast.broadcasted(
     fs::AbstractFieldStyle,
     ::Type{V},
     arg,
-) where {V <: Geometry.AxisVector}
+) where {V <: Geometry.AbstractTensor{1}}
     space = Fields.axes(arg)
     # wrap in a Field so that the axes line up correctly (it just get's unwraped so effectively a no-op)
     Base.Broadcast.broadcasted(fs, V, arg, local_geometry_field(space))
 end
 
-function Base.Broadcast.broadcasted(
-    fs::AbstractFieldStyle,
-    ::Type{V},
-    arg,
-) where {V <: Geometry.CartesianVector}
-    space = Fields.axes(arg)
-    # wrap in a Field so that the axes line up correctly (it just get's unwraped so effectively a no-op)
-    Base.Broadcast.broadcasted(
-        fs,
-        V,
-        arg,
-        tuple(Spaces.global_geometry(space)),
-        local_geometry_field(space),
-    )
-end
-
 function Base.copyto!(
     field::Field,
     bc::Base.Broadcast.Broadcasted{Base.Broadcast.DefaultArrayStyle{0}},

src/Fields/field_iterator.jl

@@ -47,10 +47,10 @@ function isa_12_covariant_field(
     f::Type{CF},
 ) where {
     FT,
-    CF <: Geometry.AxisTensor{
-        FT,
+    CF <: Geometry.Tensor{
         1,
-        Tuple{Geometry.CovariantAxis{(1, 2)}},
+        FT,
+        Tuple{Geometry.Covariant12Axis},
         StaticArrays.SVector{2, FT},
     },
 }
@@ -61,10 +61,10 @@ function isa_3_covariant_field(
     f::Type{CF},
 ) where {
     FT,
-    CF <: Geometry.AxisTensor{
-        FT,
+    CF <: Geometry.Tensor{
         1,
-        Tuple{Geometry.CovariantAxis{(3,)}},
+        FT,
+        Tuple{Geometry.Covariant3Axis},
         StaticArrays.SVector{1, FT},
     },
 }