Refactor Geometry module

Author: nefrathenrici

.buildkite/Manifest.toml

@@ -24,11 +24,11 @@ steps:
       - "echo $$JULIA_DEPOT_PATH"
       - 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("MPI")'''
-      - "julia --project=$COUPLER_PATH/experiments/AMIP/ -e 'using Pkg; Pkg.precompile()'"
-      - "julia --project=$COUPLER_PATH/experiments/AMIP/ -e 'using Pkg; Pkg.status()'"
+      - 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("MPI")'
+      - julia --project=$COUPLER_PATH/experiments/AMIP/ -e 'using Pkg; Pkg.precompile()'
+      - julia --project=$COUPLER_PATH/experiments/AMIP/ -e 'using Pkg; Pkg.status()'
     agents:
       slurm_gpus: 1
       slurm_cpus_per_task: 8

.buildkite/perf/pipeline.yml

@@ -252,10 +252,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
@@ -1090,12 +1090,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"
@@ -1621,14 +1621,6 @@ steps:
 
   # TODO: improve performance label: [inference, allocs, flops, latency, benchmark, boundscheck]
   # TODO: use perf env for perf jobs, or merge test/perf envs
-  - group: "Perf: Geometry"
-    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"
-
   - group: "Perf: DataLayouts"
     steps:
 
@@ -2500,9 +2492,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"

.buildkite/pipeline.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/ClimaCoreVTK.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

.github/workflows/docs.yml

@@ -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"

Project.toml

@@ -20,7 +20,6 @@ Geometry.LocalGeometry
 ## Internals
 
 ```@docs
-Geometry.Δz_metric_component
 Geometry.:⊗
 ```
 

docs/src/APIs/geometry_api.md

@@ -108,7 +108,9 @@ 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` containing a `Number` or a `Geometry.Tensor{2}`. The
+  tensor's basis is whatever the user supplies; there's no padding convention
+  imposed here.
 - 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 +175,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.

docs/src/matrix_fields.md

@@ -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/box/bubble_3d_rhotheta.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)),