Refactor FieldMatrix indexing for readability

Also add support for tensor slicing and
improved docs

wip

finish docs

minor formatting fix

another minor docs formatting fix

Author: imreddyTeja

.buildkite/pipeline.yml

@@ -874,13 +874,13 @@ steps:
         agents:
           slurm_gpus: 1
 
-      - label: "Unit: scalar_field_matrix (CPU)"
-        key: cpu_scalar_field_matrix
-        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/MatrixFields/scalar_field_matrix.jl"
+      - label: "Unit: field_matrix_indexing (CPU)"
+        key: cpu_field_matrix_indexing
+        command: "julia --color=yes --check-bounds=yes --project=.buildkite test/MatrixFields/field_matrix_indexing.jl"
 
-      - label: "Unit: scalar_field_matrix (GPU)"
-        key: gpu_scalar_field_matrix
-        command: "julia --color=yes --project=.buildkite test/MatrixFields/scalar_field_matrix.jl"
+      - label: "Unit: field_matrix_indexing (GPU)"
+        key: gpu_field_matrix_indexing
+        command: "julia --color=yes --project=.buildkite test/MatrixFields/field_matrix_indexing.jl"
         env:
           CLIMACOMMS_DEVICE: "CUDA"
         agents:

docs/src/matrix_fields.md

@@ -108,8 +108,8 @@ scalar_field_matrix
 A FieldMatrix entry can be:
 
 - An `UniformScaling`, which contains a `Number`
-- A `DiagonalMatrixRow`, which can contain aything
-- A `ColumnwiseBandMatrixField`, where each row is a [`BandMatrixRow`](@ref) where the band element type is representable with the space's base number type.
+- A `DiagonalMatrixRow`, which can contain either a `Number` or a tensor (represented as a `Geometry.Axis2Tensor`)
+- 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.
 This is also true for nested composite types.
@@ -164,10 +164,7 @@ An example of this is:
 
 Now consider what happens indexing `A` with the key `(@name(name1.foo.bar.buz), @name(name2.biz.bop.fud))`.
 
-First, a function searches the keys of `A` for a key that `(@name(foo.bar.buz), @name(biz.bop.fud))`
-is a child of. In this example, `(@name(foo.bar.buz), @name(biz.bop.fud))` is a child of
-the key `(@name(name1), @name(name2))`, and
-`(@name(foo.bar.buz), @name(biz.bop.fud))` is referred to as the internal key.
+First, `getindex` finds a key in `A` that contains the key being indexed. In this example, `(@name(name1.foo.bar.buz), @name(name2.biz.bop.fud))` is contained within `(@name(name1), @name(name2))`, so `(@name(name1), @name(name2))` is called the "parent key" and `(@name(foo.bar.buz), @name(biz.bop.fud))` is referred to as the "internal key".
 
 Next, the entry that `(@name(name1), @name(name2))` is paired with is recursively indexed
 by the internal key.
@@ -181,9 +178,9 @@ works as follows:
 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.
 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 on the 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 row of `entry`,
-extract that field from `entry`, and recurse on the it with all of `internal_name_pair[1]` and the remaining names of `internal_name_pair[2]`
+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.
@@ -194,3 +191,161 @@ the following situations:
 
 1. The internal key indexes to a type different than the basetype of the entry
 2. The internal key indexes to a zero-ed value
+
+```@setup 2
+using ClimaCore.CommonSpaces
+using ClimaCore.Geometry
+using ClimaCore.Fields
+import ClimaCore: MatrixFields
+import ClimaCore.MatrixFields: @name
+FT = Float64
+space = ColumnSpace(FT ;
+           z_elem = 6,
+           z_min = 0,
+           z_max = 10,
+           staggering = CellCenter()
+       )
+f = map(x -> rand(Geometry.Covariant12Vector{Float64}), Fields.local_geometry_field(space))
+g = map(x -> rand(Geometry.Covariant12Vector{Float64}), Fields.local_geometry_field(space))
+identity_axis2tensor = Geometry.Covariant12Vector(FT(1), FT(0)) *
+                   Geometry.Contravariant12Vector(FT(1), FT(0))' +
+                   Geometry.Covariant12Vector(FT(0), FT(1)) *
+                   Geometry.Contravariant12Vector(FT(0), FT(1))'
+∂f_∂g = fill(MatrixFields.TridiagonalMatrixRow(-0.5 * identity_axis2tensor, identity_axis2tensor, -0.5 * identity_axis2tensor), space)
+J = MatrixFields.FieldMatrix((@name(f), @name(g))=> ∂f_∂g)
+```
+
+## Optimizations
+
+Each entry of a `FieldMatrix` can be a `ColumnwiseBandMatrixField`, a `DiagonalMatrixRow`, or an
+`UniformScaling`.
+
+A `ColumnwiseBandMatrixField` is a `Field` with a `BandMatrixRow` at each point. It is intended
+to represent a collection of banded matrices, where there is one band matrix for each column
+of the space the `Field` is on. Beyond only storing the diagonals of the band matrix, an `entry`
+can be optimized to use less memory. Each optimized representation can be indexed equivalently to
+non optimized representation, and used in addition, subtraction, Matrix-vector multiplication,
+Matrix-matrix multiplication, `RecursiveApply`, and `FieldMatrixSolver`.
+
+For the following sections, `space` is a column space with $N_v$ levels. A column space is
+used for simplicity in this example, but the optimizations work with any space with columns.
+
+Let $f$ and $g$ be `Fields` on `space` with elements of type with elements of type
+`T_f` and `T_g`. $f_i$ and $g_i$ refers to the values of $f$ and $g$ at the $ 0 < i \leq N_v$ level.
+
+Let $M$ be a $N_v \times N_v$ banded matrix with lower and upper bandwidth of $b_1$ and $b_2$.
+$M$ represents $\frac{\partial f}{\partial g}$, so $M_{i,j} = \frac{\partial f_i}{\partial g_j}$
+
+### `ScalingFieldMatrixEntry` Optimization
+
+Consider the case where $b_1 = 0$ and $b_2 = 0$, i.e $M$ is a diagonal matrix, and
+where $M = k * I$, and $k$ is of type `T_k`. This would happen if
+$\frac{\partial f_i}{\partial g_j} = [i=j] * k$. Instead of storing
+each element on the diagonal, less memory can be used by storing a single value that represents a scaling of the identity. This reduces memory usage by a factor of $N_v$.
+
+```julia
+entry = fill(DiagonalMatrixRow(k), space)
+```
+
+can also be represented by
+
+```julia
+entry = DiagonalMatrixRow(k)
+```
+
+or, if `T_k` is a scalar, then
+
+```julia
+entry = I * k
+```
+
+### Implicit Tensor Structure Optimization
+
+The functions that index an entry with an internal key assume the implicit tensor structure optimization is being used
+when all of the following are true for `entry` where `T_k` is the element type of each band, and
+`(internal_key_1, internal_key_2)` is the internal key indexing `entry`.
+
+- the first key in the `internal_key_1` name chain is not a fieldname of `T_k`
+- the first key in the `internal_key_2` name chain is not a fieldname of `T_k`
+- `internal_key_1` and `internal_key_2` are both not empty
+
+For most use cases, `T_k` is a scalar.
+
+If the above conditions are met, the optimization assumes that the user intends the
+entry to have an implicit tensor structure, with the values of type `T_k` representing a scaling of the
+tensor identity. If both the first and second name in the name pair are equivalent, then they index onto the diagonal,
+and the value is returned. Otherwise, they index off the diagonal, and a zero value
+is returned.
+
+This optimization is intended to be used when `T_f = T_g`, and they are both `AxisVectors`
+The notation $f_{n}[i]$ where $0 < n \leq N_v$  refers to the $i$ component of the element
+at the $n$ vertical level of $f$. In the following example, `T_f=T_g=Covariant12Vector`, and
+$b_1 = b_2 = 1$, and
+
+```math
+\frac{\partial f_n[i]}{\partial g_m[j]} = \begin{cases}
+  -0.5, & \text{if } i = j \text{ and }  m = n-1 \text{ or } m = n+1 \\
+  1, & \text{if } i = j \text{ and } m = n \\
+  0, & \text{if } i \neq j \text{ or } m < n -1 \text{ or } m > n +1
+\end{cases}
+```
+
+The non-zero values of each row of `M` are equivalent in this example, but they can also vary in value.
+
+```julia
+∂f_∂g = fill(MatrixFields.TridiagonalMatrixRow(-0.5 * identity_axis2tensor, identity_axis2tensor, -0.5 * identity_axis2tensor), space)
+J = MatrixFields.FieldMatrix((@name(f), @name(g))=> ∂f_∂g)
+```
+
+`∂f_∂g` can be indexed into to get the partial derrivatives of individual components.
+
+
+```@example 2
+J[(@name(f.components.data.:(1)), @name(g.components.data.:(1)))]
+```
+
+```@example 2
+J[(@name(f.components.data.:(2)), @name(g.components.data.:(1)))]
+```
+
+This can be more optimally stored with the implicit tensor structure optimization:
+
+```@example 2
+∂f_∂g = fill(MatrixFields.TridiagonalMatrixRow(-0.5, 1.0, -0.5), space) # hide
+J = MatrixFields.FieldMatrix((@name(f), @name(g))=> ∂f_∂g) # hide
+```
+
+```julia
+∂f_∂g = fill(MatrixFields.TridiagonalMatrixRow(-0.5, 1.0, -0.5), space)
+J = MatrixFields.FieldMatrix((@name(f), @name(g))=> ∂f_∂g)
+```
+
+```@example 2
+J[(@name(f.components.data.:(1)), @name(g.components.data.:(1)))]
+```
+
+```@example 2
+Base.Broadcast.materialize(J[(@name(f.components.data.:(2)), @name(g.components.data.:(1)))])
+```
+
+If it is the case that
+
+```math
+\frac{\partial f_n[i]}{\partial g_m[j]} = \begin{cases}
+  k, & \text{if } i = j \text{ and } m = n \\
+  0, & \text{if } i \neq j \text{ or } m \neq n
+\end{cases}
+```
+
+where $k$ is a constant scalar, both the implicit tensor structure optimization and
+`ScalingFieldMatrixEntry` optimization can both be applied:
+
+```julia
+∂f_∂g = fill(MatrixFields.DiagonalMatrixRow(k), space)
+```
+
+can equivalently be represented with
+
+```julia
+∂f_∂g = MatrixFields.DiagonalMatrixRow(k)
+```