finish docs

Author: imreddyTeja

docs/src/matrix_fields.md

@@ -191,9 +191,6 @@ 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
-3. The internal key slices an `AxisTensor`
-
-### Implicit Tensor Structure Optimization
 
 ```@setup 2
 using ClimaCore.CommonSpaces
@@ -208,124 +205,147 @@ space = ColumnSpace(FT ;
            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)
 ```
 
-If using a `FieldMatrix` to represent a jacobian, entries with certain structures
-can be stored in an optimized manner.
+## Optimizations
 
-The optimization assumes that if indexing into an entry of scalars, the user intends the
-entry to have an implicit tensor structure, with the scalar values 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 scalar value is returned. Otherwise, they index off the diagonal, and a zero value
-is returned.
+Each entry of a `FieldMatrix` can be a `ColumnwiseBandMatrixField`, a `DiagonalMatrixRow`, or an
+`UniformScaling`.
 
-The following sections refer the `Field`s
-$f$ and $g$, which both have values of type `Covariant12Vector` and are defined on a column domain, which is discretized with $N_v$ layers.
-The notation $f_{n}[i]$ where $ 0 < n \leq N_v$ and $i \in (1,2)$ refers to the $i$ component of the element of $f$
-at the $i$ vertical level. $g$ is indexed similarly. Although $f$ and $g$ have values of type
-`Covariant12Vector`, this optimization works for any two `Field`s of `AxisVector`s
+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`.
 
-```@example 2
-f = map(x -> rand(Geometry.Covariant12Vector{Float64}), Fields.local_geometry_field(space))
-g = map(x -> rand(Geometry.Covariant12Vector{Float64}), Fields.local_geometry_field(space))
-```
+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.
 
-#### Uniform Scaling Case
+$f$ and $g$ are both `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.
 
-If $\frac{\partial f_n[i]}{\partial g_n[j]} = [i = j]$ for some scalar $k$, then the
-non-optimized entry would be represented by a diagonal matrix with values of an identity 2d tensor. If $k=2$, then
+$M$ is 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}$
 
-```@example 2
-identity_axis2tensor = Geometry.Covariant12Vector(FT(1), FT(0)) * # hide
-                   Geometry.Contravariant12Vector(FT(1), FT(0))' + # hide
-                   Geometry.Covariant12Vector(FT(0), FT(1)) * # hide
-                   Geometry.Contravariant12Vector(FT(0), FT(1))' # hide
-k = 2
-∂f_∂g = fill(MatrixFields.DiagonalMatrixRow(k * identity_axis2tensor), space)
+### `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)
 ```
 
-Individual components can be indexed into:
+can also be represented by
 
-```@example 2
-J = MatrixFields.FieldMatrix((@name(f), @name(g))=> ∂f_∂g)
-J[[(@name(f.components.data.:(1)), @name(g.components.data.:(1)))]]
+```julia
+entry = DiagonalMatrixRow(k)
 ```
 
-```@example 2
-J[[(@name(f.components.data.:(2)), @name(g.components.data.:(1)))]]
+or, if `T_k` is a scalar, then
+
+```julia
+entry = I * k
 ```
 
-The above example indexes into $\frac{\partial f_n[1]}{\partial g_n[1]}$ where $ 0 < n \leq N_v$
+### Implicit Tensor Structure Optimization
 
-The entry can
-also be represeted with a single `DiagonalMatrixRow`, as follows:
+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`.
 
-```@example 2
-∂f_∂g_optimized = MatrixFields.DiagonalMatrixRow(k * identity_axis2tensor)
-```
+- 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
 
-`∂f_∂g_optimized` is a single `DiagonalMatrixRow`, which represents a diagonal matrix with the
-same tensor along the diagonal. In this case, that tensor is $k$ multiplied by the identity matrix, and that can be
-represented with `k * I` as follows
+For most use cases, `T_k` is a scalar.
 
-```@example 2
-∂f_∂g_more_optimized = MatrixFields.DiagonalMatrixRow(k * identity_axis2tensor)
+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}
 ```
 
-Individual components of `∂f_∂g_optimized` can be indexed in the same way as `∂f_∂g`.
+The non-zero values of each row of `M` are equivalent in this example, but they can also vary in value.
 
-```@example 2
-J_unoptimized = MatrixFields.FieldMatrix((@name(f), @name(g)) => ∂f_∂g)
-J_unoptimized[(@name(f.components.data.:(1)), @name(g.components.data.:(1)))]
+```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_more_optimized = MatrixFields.FieldMatrix((@name(f), @name(g)) => ∂f_∂g_optimized)
-J_more_optimized[(@name(f.components.data.:(1)), @name(g.components.data.:(1)))]
+J[(@name(f.components.data.:(1)), @name(g.components.data.:(1)))]
 ```
 
 ```@example 2
-J_more_optimized[(@name(f.components.data.:(1)), @name(g.components.data.:(2)))]
+J[(@name(f.components.data.:(2)), @name(g.components.data.:(1)))]
 ```
 
-`∂f_∂g` stores $2 * 2 * N_v$ floats in memory, `∂f_∂g_optimized` stores `$2*2$ floats, and
-`∂f_∂g_more_optimized` stores only one float.
-
-#### Vertically Varying Case
-
-The implicit tensor optimization can also be used when
-$\frac{\partial f_n[i]}{\partial g_n[j]} = [i = j] * h(f_n, g_n)$.
-
-In this case, a full `ColumnWiseBandMatrixField` must be used.
+This can be more optimally with the implicit tensor structure:
 
 ```@example 2
-∂f_∂g_optimized = map(x -> MatrixFields.DiagonalMatrixRow(rand(Float64)), ∂f_∂g)
+∂f_∂g = fill(MatrixFields.TridiagonalMatrixRow(-0.5, 1.0, -0.5), space) # hide
+J = MatrixFields.FieldMatrix((@name(f), @name(g))=> ∂f_∂g) # hide
 ```
 
-```@example 2
-J_optimized = MatrixFields.FieldMatrix((@name(f), @name(g)) => ∂f_∂g_optimized)
-J_optimized[(@name(f.components.data.:(1)), @name(g.components.data.:(1)))]
+```julia
+∂f_∂g = fill(MatrixFields.TridiagonalMatrixRow(-0.5, 1.0, -0.5), space)
+J = MatrixFields.FieldMatrix((@name(f), @name(g))=> ∂f_∂g)
 ```
 
 ```@example 2
-Base.Broadcast.materialize(J_optimized[(@name(f.components.data.:(2)), @name(g.components.data.:(1)))])
+J[(@name(f.components.data.:(1)), @name(g.components.data.:(1)))]
 ```
 
-#### bandwidth > 1 case
+```@example 2
+Base.Broadcast.materialize(J[(@name(f.components.data.:(2)), @name(g.components.data.:(1)))])
+```
 
-The implicit tensor optimization can also be used when
-$\frac{\partial f_n[i]}{\partial g[j]} = [i = j] * h(f_n, g_{n-k_1}, ..., g_{n+k_2})$ where
-$b_1$ and $b_2$ are the lower and upper bandwidth. Say $b_1 = b_2 = 1$. Then
+If it is the case that
 
-```@example 2
-∂f_∂g_optimized = map(x -> MatrixFields.TridiagonalMatrixRow(rand(Float64), rand(Float64), rand(Float64)), ∂f_∂g)
+```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}
 ```
 
-```@example 2
-J_optimized = MatrixFields.FieldMatrix((@name(f), @name(g)) => ∂f_∂g_optimized)
-J_optimized[(@name(f.components.data.:(1)), @name(g.components.data.:(1)))]
+where $k$ is a constant scalar, both the implicit tensor structure optimization and
+`ScalingFieldMatrixEntry` optimization can be applied:
+
+```julia
+∂f_∂g = fill(MatrixFields.DiagonalMatrixRow(k), space)
 ```
 
-```@example 2
-Base.Broadcast.materialize(J_optimized[(@name(f.components.data.:(2)), @name(g.components.data.:(1)))])
+can equivalently be represented with
+
+```julia
+∂f_∂g = MatrixFields.DiagonalMatrixRow(k)
 ```

src/MatrixFields/field_name_dict.jl

@@ -148,9 +148,22 @@ get_internal_key(child_name_pair::FieldNamePair, name_pair::FieldNamePair) = (
     extract_internal_name(child_name_pair[1], name_pair[1]),
     extract_internal_name(child_name_pair[2], name_pair[2]),
 )
+"""
+    get_internal_entry(entry, name::FieldName)
 
+Returns the field indexed to by `name` from `entry`
+"""
 get_internal_entry(entry, name::FieldName) = get_field(entry, name)
 # call get_internal_entry on scaling value, and rebuild entry container
+"""
+    get_internal_entry(entry, name_pair::FieldNamePair)
+
+Returns the field indexed to by `name_pair` from `entry`. Indexing behavior is described
+in the MatrixFields section of the documentation. If `entry` is a `ColumnwiseBandMatrixField`,
+and the field indexed to by `name_pair` is not a field of scalars, a broadcasted object
+is returned. This also happens when indexing off diagonal with the implicit tensor structure
+optimization (see MatrixFields documentation).
+"""
 get_internal_entry(entry::UniformScaling, name_pair::FieldNamePair) =
     UniformScaling(get_internal_entry(scaling_value(entry), name_pair))
 get_internal_entry(entry::DiagonalMatrixRow, name_pair::FieldNamePair) =
@@ -179,6 +192,24 @@ function get_internal_entry(
             (drop_first(internal_row_name), drop_first(internal_col_name)),
             full_key,
         )
+    elseif T <: Geometry.Axis2Tensor && # slicing a 2d tensor
+           is_child_name(name_pair[1], @name(components.data))
+        internal_row_name =
+            extract_internal_name(name_pair[1], @name(components.data))
+        return get_internal_entry(
+            entry[extract_first(internal_row_name), :],
+            (drop_first(internal_row_name), name_pair[2]),
+            full_key,
+        )
+    elseif T <: Geometry.Axis2Tensor && # slicing a 2d tensor
+           is_child_name(name_pair[2], @name(components.data))
+        internal_col_name =
+            extract_internal_name(name_pair[2], @name(components.data))
+        return get_internal_entry(
+            entry[:, extract_first(internal_col_name)],
+            (name_pair[1], drop_first(internal_col_name)),
+            full_key,
+        )
     elseif T <: Geometry.AdjointAxisVector # bypass parent for adjoint vectors
         return get_internal_entry(getfield(entry, :parent), name_pair, full_key)
     elseif name_pair[1] != @name() &&
@@ -335,17 +366,15 @@ function field_offset_and_type(
         # if S <: T, then its possible to construct a strided view in the indexing function
         return (0, S, S <: T ? Val(:view) : Val(:view_of_blocks))
     elseif S <: Geometry.Axis2Tensor &&
-           all(n -> is_child_name(n, @name(components.data)), name_pair) # special case to calculate index
+           any(n -> is_child_name(n, @name(components.data)), name_pair) # special case to calculate index
+        all(n -> is_child_name(n, @name(components.data)), name_pair) ||
+            return (0, S, Val{:broadcasted_fallback}())
         internal_row_name =
             extract_internal_name(name_pair[1], @name(components.data))
         internal_col_name =
             extract_internal_name(name_pair[2], @name(components.data))
         row_index = extract_first(internal_row_name)
         col_index = extract_first(internal_col_name)
-        if ((row_index isa Number) && (col_index isa Colon)) ||
-           ((row_index isa Colon) && (col_index isa Number))
-            return (0, S, Val{:broadcasted_fallback}()) # slice case, return trigger fallback
-        end
         ((row_index isa Number) && (col_index isa Number)) ||
             throw(KeyError(full_key))
         (n_rows, n_cols) = map(length, axes(S))

test/MatrixFields/field_matrix_indexing.jl

@@ -184,19 +184,17 @@ end
     A, _ = dycore_prognostic_EDMF_FieldMatrix(FT)
     A_scaling, _ = scaling_only_dycore_prognostic_EDMF_FieldMatrix(FT)
 
-    colon_field_name =
-        append_internal_name(@name(c.uₕ.components.data), FieldName(Colon()))
     index_name_1 = @name(c.uₕ.components.data.:(1))
     index_name_2 = @name(c.uₕ.components.data.:(2))
 
     @test eltype(
-        eltype(Base.Broadcast.materialize(A[(colon_field_name, index_name_1)])),
+        eltype(Base.Broadcast.materialize(A[(@name(c.uₕ), index_name_1)])),
     ) <: Geometry.CovariantVector
     @test eltype(
-        eltype(Base.Broadcast.materialize(A[(index_name_2, colon_field_name)])),
+        eltype(Base.Broadcast.materialize(A[(index_name_2, @name(c.uₕ))])),
     ) <: Geometry.ContravariantVector
-    @test eltype(A_scaling[(colon_field_name, index_name_1)]) <:
+    @test eltype(A_scaling[(@name(c.uₕ), index_name_1)]) <:
           Geometry.CovariantVector
-    @test eltype(A_scaling[(index_name_2, colon_field_name)]) <:
+    @test eltype(A_scaling[(index_name_2, @name(c.uₕ))]) <:
           Geometry.ContravariantVector
 end