wip

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,141 @@ 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
+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()
+       )
+```
+
+If using a `FieldMatrix` to represent a jacobian, entries with certain structures
+can be stored in an optimized manner.
+
+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.
+
+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
+
+```@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))
+```
+
+#### Uniform Scaling Case
+
+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
+
+```@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)
+```
+
+Individual components can be indexed into:
+
+```@example 2
+J = MatrixFields.FieldMatrix((@name(f), @name(g))=> ∂f_∂g)
+J[[(@name(f.components.data.:(1)), @name(g.components.data.:(1)))]]
+```
+
+```@example 2
+J[[(@name(f.components.data.:(2)), @name(g.components.data.:(1)))]]
+```
+
+The above example indexes into $\frac{\partial f_n[1]}{\partial g_n[1]}$ where $ 0 < n \leq N_v$
+
+The entry can
+also be represeted with a single `DiagonalMatrixRow`, as follows:
+
+```@example 2
+∂f_∂g_optimized = MatrixFields.DiagonalMatrixRow(k * identity_axis2tensor)
+```
+
+`∂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
+
+```@example 2
+∂f_∂g_more_optimized = MatrixFields.DiagonalMatrixRow(k * identity_axis2tensor)
+```
+
+Individual components of `∂f_∂g_optimized` can be indexed in the same way as `∂f_∂g`.
+
+```@example 2
+J_unoptimized = MatrixFields.FieldMatrix((@name(f), @name(g)) => ∂f_∂g)
+J_unoptimized[(@name(f.components.data.:(1)), @name(g.components.data.:(1)))]
+```
+
+```@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)))]
+```
+
+```@example 2
+J_more_optimized[(@name(f.components.data.:(1)), @name(g.components.data.:(2)))]
+```
+
+`∂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.
+
+```@example 2
+∂f_∂g_optimized = map(x -> MatrixFields.DiagonalMatrixRow(rand(Float64)), ∂f_∂g)
+```
+
+```@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)))]
+```
+
+```@example 2
+Base.Broadcast.materialize(J_optimized[(@name(f.components.data.:(2)), @name(g.components.data.:(1)))])
+```
+
+#### bandwidth > 1 case
+
+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
+
+```@example 2
+∂f_∂g_optimized = map(x -> MatrixFields.TridiagonalMatrixRow(rand(Float64), rand(Float64), rand(Float64)), ∂f_∂g)
+```
+
+```@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)))]
+```
+
+```@example 2
+Base.Broadcast.materialize(J_optimized[(@name(f.components.data.:(2)), @name(g.components.data.:(1)))])
+```

src/MatrixFields/field_name_dict.jl

@@ -213,11 +213,11 @@ function get_internal_entry(
     name_pair == (@name(), @name()) && return entry
     S = eltype(eltype(entry))
     T = eltype(parent(entry))
-    (start_offset, target_type, apply_zero) =
+    (start_offset, target_type, index_method) =
         field_offset_and_type(name_pair, T, S, name_pair)
-    if target_type <: eltype(parent(entry)) && !apply_zero
-        band_element_size =
-            DataLayouts.typesize(eltype(parent(entry)), eltype(eltype(entry)))
+    if isa(index_method, Val{:view})
+        @assert target_type <: T
+        band_element_size = DataLayouts.typesize(T, S)
         singleton_datalayout = DataLayouts.singleton(Fields.field_values(entry))
         scalar_band_type =
             band_matrix_row_type(outer_diagonals(eltype(entry))..., target_type)
@@ -234,14 +234,15 @@ function get_internal_entry(
             scalar_data,
         )
         return Fields.Field(values, axes(entry))
-    elseif apply_zero && start_offset == 0
+    elseif isa(index_method, Val{:broadcasted_zero})
+        # implicit tensor structure optimization, off diagonal
         zero_value = zero(target_type)
         return Base.broadcasted(entry) do matrix_row
             map(x -> zero_value, matrix_row)
         end
-    elseif target_type == S && start_offset == 0
+    elseif target_type == S && isa(index_method, Val{:view_of_blocks})
         return entry
-    else
+    else # fallback to broadcasted indexing on each element, currently no support for view_of_blocks
         return Base.broadcasted(entry) do matrix_row
             map(matrix_row) do matrix_row_entry
                 get_internal_entry(matrix_row_entry, name_pair)
@@ -306,35 +307,49 @@ end
     field_offset_and_type(name_pair::FieldNamePair, ::Type{T}, ::Type{S}, full_key::FieldNamePair)
 
 Returns the offset of the field with name `name_pair` in an object of type `S` in
-multiples of `sizeof(T)` and the type of the field with name `name_pair`.
+multiples of `sizeof(T)`, the type of the field with name `name_pair`, and a `Val` indicating
+what method can index a ClimaCore `Field` of `S` with `name_pair`.
+
+The third return value is one of the following:
+- `Val(:view)`: indexing with a view is possible
+-  `Val(:view_of_blocks)`: indexing with a view of blocks is possible (this is not implemented)
+- `Val(:broadcasted_fallback)`: indexing with a view is not possible
+- `Val(:broadcasted_zero)`: indexing with a view is not possible, and the `name_pair` indexes
+off diagonal with implicit tensor structure optimization (see MatrixFields docs)
 
-When `S` is a `Geometry.Axis2Tensor`, the name pair must index into a scalar of
-the tensor or be empty. In other words, the name pair cannot index into a slice.
+When `S` is a `Geometry.Axis2Tensor`, and the name pair indexes to a slice of
+the tensor, an offset of `-1` is returned . In other words, the name pair cannot index into a slice.
 
 If neither element of `name_pair` is `@name()`, the first name in the pair is indexed with
 first, and then the second name is used to index the result of the first.
+
+This is an internal funtion designed to be used with `get_internal_entry(::ColumnwiseBandMatrixField)`
 """
 function field_offset_and_type(
     name_pair::FieldNamePair,
     ::Type{T},
     ::Type{S},
     full_key::FieldNamePair,
 ) where {S, T}
-    name_pair == (@name(), @name()) && return (0, S, false) # base case
-    if S <: Geometry.Axis2Tensor &&
-       all(n -> is_child_name(n, @name(components.data)), name_pair)# special case to calculate index
-        (name_pair[1] == @name() || name_pair[2] == @name()) &&
-            throw(KeyError(full_key))
+    if name_pair == (@name(), @name()) # recursion base case
+        # 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
         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)) # slicing not supported
+            throw(KeyError(full_key))
         (n_rows, n_cols) = map(length, axes(S))
-        (remaining_offset, end_type, apply_zero) = field_offset_and_type(
+        (remaining_offset, end_type, index_method) = field_offset_and_type(
             (drop_first(internal_row_name), drop_first(internal_col_name)),
             T,
             eltype(S),
@@ -345,20 +360,19 @@ function field_offset_and_type(
         return (
             (n_rows * (col_index - 1) + row_index - 1) + remaining_offset,
             end_type,
-            apply_zero,
+            index_method,
         )
-    elseif S <: Geometry.AdjointAxisVector
+    elseif S <: Geometry.AdjointAxisVector # bypass adjoint because indexing parent is equivalent
         return field_offset_and_type(name_pair, T, fieldtype(S, 1), full_key)
     elseif name_pair[1] != @name() &&
-           extract_first(name_pair[1]) in fieldnames(S)
-
+           extract_first(name_pair[1]) in fieldnames(S) # index with first part of name_pair[1]
         remaining_field_chain = (drop_first(name_pair[1]), name_pair[2])
         child_type = fieldtype(S, extract_first(name_pair[1]))
         field_index = unrolled_filter(
             i -> fieldname(S, i) == extract_first(name_pair[1]),
             1:fieldcount(S),
         )[1]
-        (remaining_offset, end_type, apply_zero) = field_offset_and_type(
+        (remaining_offset, end_type, index_method) = field_offset_and_type(
             remaining_field_chain,
             T,
             child_type,
@@ -367,18 +381,17 @@ function field_offset_and_type(
         return (
             DataLayouts.fieldtypeoffset(T, S, field_index) + remaining_offset,
             end_type,
-            apply_zero,
+            index_method,
         )
     elseif name_pair[2] != @name() &&
-           extract_first(name_pair[2]) in fieldnames(S)
-
+           extract_first(name_pair[2]) in fieldnames(S) # index with first part of name_pair[2]
         remaining_field_chain = name_pair[1], drop_first(name_pair[2])
         child_type = fieldtype(S, extract_first(name_pair[2]))
         field_index = unrolled_filter(
             i -> fieldname(S, i) == extract_first(name_pair[2]),
             1:fieldcount(S),
         )[1]
-        (remaining_offset, end_type, apply_zero) = field_offset_and_type(
+        (remaining_offset, end_type, index_method) = field_offset_and_type(
             remaining_field_chain,
             T,
             child_type,
@@ -387,10 +400,10 @@ function field_offset_and_type(
         return (
             DataLayouts.fieldtypeoffset(T, S, field_index) + remaining_offset,
             end_type,
-            apply_zero,
+            index_method,
         )
-    elseif !any(isequal(@name()), name_pair) # implicit tensor structure
-        (remaining_offset, end_type, apply_zero) = field_offset_and_type(
+    elseif !any(isequal(@name()), name_pair) # implicit tensor structure optimization
+        (remaining_offset, end_type, index_method) = field_offset_and_type(
             (drop_first(name_pair[1]), drop_first(name_pair[2])),
             T,
             S,
@@ -400,7 +413,7 @@ function field_offset_and_type(
             remaining_offset,
             end_type,
             extract_first(name_pair[1]) == extract_first(name_pair[2]) ?
-            apply_zero : true,
+            index_method : Val(:broadcasted_zero), # zero if off diagonal
         )
     else
         throw(KeyError(full_key))