Merge pull request #1680 from CliMA/ck/local_geometry_type

Define `local_geometry_type`

Author: charleskawczynski

src/Fields/Fields.jl

@@ -7,7 +7,7 @@ import ..DataLayouts: DataLayouts, AbstractData, DataStyle
 import ..Domains
 import ..Topologies
 import ..Quadratures
-import ..Grids: ColumnIndex
+import ..Grids: ColumnIndex, local_geometry_type
 import ..Spaces: Spaces, AbstractSpace, AbstractPointSpace
 import ..Geometry: Geometry, Cartesian12Vector
 import ..Utilities: PlusHalf, half, UnrolledFunctions
@@ -39,6 +39,8 @@ Field(values::V, space::S) where {V <: AbstractData, S <: AbstractSpace} =
 Field(::Type{T}, space::S) where {T, S <: AbstractSpace} =
     Field(similar(Spaces.coordinates_data(space), T), space)
 
+local_geometry_type(::Field{V, S}) where {V, S} = local_geometry_type(S)
+
 ClimaComms.context(field::Field) = ClimaComms.context(axes(field))
 
 ClimaComms.context(topology::Topologies.Topology2D) = topology.context

src/Grids/Grids.jl

@@ -49,6 +49,16 @@ If the grid is not staggered, `staggering` should be `nothing`.
 """
 function local_geometry_data end
 
+"""
+    Grids.local_geometry_type(::Type)
+
+Get the `LocalGeometry` type.
+"""
+function local_geometry_type end
+
+# Fallback, but this requires user error-handling
+local_geometry_type(::Type{T}) where {T} = Union{}
+
 function local_dss_weights end
 function quadrature_style end
 function vertical_topology end

src/Grids/column.jl

@@ -36,6 +36,8 @@ struct ColumnGrid{
     colidx::C
 end
 
+local_geometry_type(::Type{ColumnGrid{G, C}}) where {G, C} =
+    local_geometry_type(G)
 
 column(grid::AbstractExtrudedFiniteDifferenceGrid, colidx::ColumnIndex) =
     ColumnGrid(grid, colidx)

src/Grids/extruded.jl

@@ -37,6 +37,10 @@ mutable struct ExtrudedFiniteDifferenceGrid{
     face_local_geometry::LG
 end
 
+local_geometry_type(
+    ::Type{ExtrudedFiniteDifferenceGrid{H, V, A, GG, LG}},
+) where {H, V, A, GG, LG} = eltype(LG) # calls eltype from DataLayouts
+
 function ExtrudedFiniteDifferenceGrid(
     horizontal_grid::Union{SpectralElementGrid1D, SpectralElementGrid2D},
     vertical_grid::FiniteDifferenceGrid,
@@ -125,7 +129,6 @@ vertical_topology(grid::ExtrudedFiniteDifferenceGrid) =
 local_dss_weights(grid::ExtrudedFiniteDifferenceGrid) =
     local_dss_weights(grid.horizontal_grid)
 
-
 local_geometry_data(grid::AbstractExtrudedFiniteDifferenceGrid, ::CellCenter) =
     grid.center_local_geometry
 local_geometry_data(grid::AbstractExtrudedFiniteDifferenceGrid, ::CellFace) =
@@ -147,6 +150,10 @@ struct DeviceExtrudedFiniteDifferenceGrid{VT, Q, GG, LG} <:
     face_local_geometry::LG
 end
 
+local_geometry_type(
+    ::Type{DeviceExtrudedFiniteDifferenceGrid{VT, Q, GG, LG}},
+) where {VT, Q, GG, LG} = eltype(LG) # calls eltype from DataLayouts
+
 Adapt.adapt_structure(to, grid::ExtrudedFiniteDifferenceGrid) =
     DeviceExtrudedFiniteDifferenceGrid(
         Adapt.adapt(to, vertical_topology(grid)),

src/Grids/finitedifference.jl

@@ -168,6 +168,9 @@ FiniteDifferenceGrid(mesh::Meshes.IntervalMesh) =
 topology(grid::FiniteDifferenceGrid) = grid.topology
 vertical_topology(grid::FiniteDifferenceGrid) = grid.topology
 
+local_geometry_type(::Type{FiniteDifferenceGrid{T, GG, LG}}) where {T, GG, LG} =
+    eltype(LG) # calls eltype from DataLayouts
+
 local_geometry_data(grid::FiniteDifferenceGrid, ::CellCenter) =
     grid.center_local_geometry
 local_geometry_data(grid::FiniteDifferenceGrid, ::CellFace) =
@@ -182,6 +185,10 @@ struct DeviceFiniteDifferenceGrid{T, GG, LG} <: AbstractFiniteDifferenceGrid
     face_local_geometry::LG
 end
 
+local_geometry_type(
+    ::Type{DeviceFiniteDifferenceGrid{T, GG, LG}},
+) where {T, GG, LG} = eltype(LG) # calls eltype from DataLayouts
+
 Adapt.adapt_structure(to, grid::FiniteDifferenceGrid) =
     DeviceFiniteDifferenceGrid(
         Adapt.adapt(to, grid.topology),

src/Grids/level.jl

@@ -18,6 +18,8 @@ topology(levelgrid::LevelGrid) = topology(levelgrid.full_grid)
 
 local_dss_weights(grid::LevelGrid) = local_dss_weights(grid.full_grid)
 
+local_geometry_type(::Type{LevelGrid{G, L}}) where {G, L} =
+    local_geometry_type(G)
 local_geometry_data(levelgrid::LevelGrid{<:Any, Int}, ::Nothing) = level(
     local_geometry_data(levelgrid.full_grid, CellCenter()),
     levelgrid.level,

src/Grids/spectralelement.jl

@@ -21,6 +21,10 @@ mutable struct SpectralElementGrid1D{
     dss_weights::D
 end
 
+local_geometry_type(
+    ::Type{SpectralElementGrid1D{T, Q, GG, LG}},
+) where {T, Q, GG, LG} = eltype(LG) # calls eltype from DataLayouts
+
 # non-view grids are cached based on their input arguments
 # this means that if data is saved in two different files, reloading will give fields which live on the same grid
 function SpectralElementGrid1D(
@@ -118,6 +122,10 @@ mutable struct SpectralElementGrid2D{
     boundary_surface_geometries::BS
 end
 
+local_geometry_type(
+    ::Type{SpectralElementGrid2D{T, Q, GG, LG, D, IS, BS}},
+) where {T, Q, GG, LG, D, IS, BS} = eltype(LG) # calls eltype from DataLayouts
+
 """
     SpectralElementSpace2D(topology, quadrature_style; enable_bubble)
 

src/MatrixFields/MatrixFields.jl

@@ -60,6 +60,7 @@ import ..DataLayouts: AbstractData
 import ..Geometry
 import ..Topologies
 import ..Spaces
+import ..Spaces: local_geometry_type
 import ..Fields
 import ..Operators
 

src/MatrixFields/matrix_multiplication.jl

@@ -270,11 +270,29 @@ function Operators.right_interior_idx(
     end
 end
 
+pick_inferred_type(
+    ::Type{Union{}},
+    ::Type{Y},
+) where {Y <: Geometry.LocalGeometry} = Y
+pick_inferred_type(
+    ::Type{X},
+    ::Type{Union{}},
+) where {X <: Geometry.LocalGeometry} = X
+pick_inferred_type(::Type{T}, ::Type{T}) where {T <: Geometry.LocalGeometry} = T
+pick_inferred_type(::Type{Union{}}, ::Type{Union{}}) =
+    error("Both LGs are not inferred")
+pick_inferred_type(::Type{X}, ::Type{Y}) where {X, Y} =
+    error("LGs do not match: X=$X, Y=$Y")
+
 function Operators.return_eltype(
     ::MultiplyColumnwiseBandMatrixField,
     matrix1,
     arg,
 )
+    # LG1 = local_geometry_type(typeof(axes(matrix1)))
+    # LG2 = local_geometry_type(typeof(axes(arg)))
+    # LG = pick_inferred_type(LG1, LG2)
+    # return Operators.return_eltype(op, matrix1, arg, LG)
     eltype(matrix1) <: BandMatrixRow || error(
         "The first argument of ⋅ must have elements of type BandMatrixRow, but \
          the given argument has elements of type $(eltype(matrix1))",
@@ -293,6 +311,39 @@ function Operators.return_eltype(
     end
 end
 
+function Operators.return_eltype(
+    ::MultiplyColumnwiseBandMatrixField,
+    matrix1,
+    arg,
+    ::Type{LG},
+) where {LG}
+    eltype(matrix1) <: BandMatrixRow || error(
+        "The first argument of ⋅ must have elements of type BandMatrixRow, but \
+         the given argument has elements of type $(eltype(matrix1))",
+    )
+    if eltype(arg) <: BandMatrixRow # matrix-matrix multiplication
+        matrix2 = arg
+        ld1, ud1 = outer_diagonals(eltype(matrix1))
+        ld2, ud2 = outer_diagonals(eltype(matrix2))
+        prod_ld, prod_ud = ld1 + ld2, ud1 + ud2
+        prod_value_type = Base.promote_op(
+            rmul_with_projection,
+            eltype(eltype(matrix1)),
+            eltype(eltype(matrix2)),
+            LG,
+        )
+        return band_matrix_row_type(prod_ld, prod_ud, prod_value_type)
+    else # matrix-vector multiplication
+        vector = arg
+        prod_value_type = Base.promote_op(
+            rmul_with_projection,
+            eltype(eltype(matrix1)),
+            eltype(vector),
+            LG,
+        )
+    end
+end
+
 Operators.return_space(::MultiplyColumnwiseBandMatrixField, space1, space2) =
     space1
 
@@ -314,6 +365,8 @@ boundary_modified_ud(::BottomRightMatrixCorner, ud, column_space, i) =
 # matrix field broadcast expressions to take roughly 3 or 4 times longer to
 # evaluate, but this is less significant than the decrease in compilation time.
 function multiply_matrix_at_index(loc, space, idx, hidx, matrix1, arg, bc)
+    # lg = Geometry.LocalGeometry(space, idx, hidx)
+    # prod_type = Operators.return_eltype(⋅, matrix1, arg, typeof(lg))
     prod_type = Operators.return_eltype(⋅, matrix1, arg)
 
     column_space1 = column_axes(matrix1, space)

src/MatrixFields/single_field_solver.jl

@@ -13,11 +13,14 @@ x_eltype(A::ColumnwiseBandMatrixField, b) =
     x_eltype(eltype(eltype(A)), eltype(b))
 x_eltype(::Type{T_A}, ::Type{T_b}) where {T_A, T_b} =
     rmul_return_type(inv_return_type(T_A), T_b)
+# Base.promote_op(rmul_with_projection, inv_return_type(T_A), T_b, LG)
 
-unit_eltype(A::UniformScaling) = unit_eltype(eltype(A))
-unit_eltype(A::ColumnwiseBandMatrixField) = unit_eltype(eltype(eltype(A)))
-unit_eltype(::Type{T_A}) where {T_A} =
+unit_eltype(A::UniformScaling) = eltype(A)
+unit_eltype(A::ColumnwiseBandMatrixField) =
+    unit_eltype(eltype(eltype(A)), local_geometry_type(A))
+unit_eltype(::Type{T_A}, ::Type{LG}) where {T_A, LG} =
     rmul_return_type(inv_return_type(T_A), T_A)
+# Base.promote_op(rmul_with_projection, inv_return_type(T_A), T_A, LG)
 
 ################################################################################