implicit flat fields

Author: kmdeck

src/MatrixFields/MatrixFields.jl

@@ -88,8 +88,7 @@ const ColumnwiseBandMatrixField{V, S} = Fields.Field{
 } where {
     V <: AbstractData{<:BandMatrixRow},
     S <: Union{
-        Spaces.FiniteDifferenceSpace,
-        Spaces.ExtrudedFiniteDifferenceSpace,
+        Spaces.AbstractSpace,
         Operators.PlaceholderSpace, # so that this can exist inside cuda kernels
     },
 }

src/MatrixFields/matrix_shape.jl

@@ -3,19 +3,34 @@ struct Square <: AbstractMatrixShape end
 struct FaceToCenter <: AbstractMatrixShape end
 struct CenterToFace <: AbstractMatrixShape end
 
+matrix_shape(matrix_field) = matrix_shape(matrix_field, axes(matrix_field))
+
 """
     matrix_shape(matrix_field, [matrix_space])
 
-Returns either `Square()`, `FaceToCenter()`, or `CenterToFace()`, depending on
-whether the diagonal indices of `matrix_field` are `Int`s or `PlusHalf`s and
-whether `matrix_space` is on cell centers or cell faces. By default,
+Returns the matrix shape for a matrix field defined on the `matrix_space`. By default,
 `matrix_space` is set to `axes(matrix_field)`.
+
+When the matrix_space is a finite difference space (extruded or otherwise): the shape is
+either `Square()`, `FaceToCenter()`, or `CenterToFace()`, depending on
+whether the diagonal indices of `matrix_field` are `Int`s or `PlusHalf`s and
+whether `matrix_space` is on cell centers or cell faces. 
+
+When the matrix_space is a spectral element or point space: only a Square() shape is supported.
 """
-matrix_shape(matrix_field, matrix_space = axes(matrix_field)) = _matrix_shape(
+matrix_shape(matrix_field, matrix_space) = _matrix_shape(
     eltype(outer_diagonals(eltype(matrix_field))),
     matrix_space.staggering,
 )
 
+function matrix_shape(
+    matrix_field,
+    matrix_space::Union{Spaces.AbstractSpectralElementSpace, Spaces.PointSpace},
+)
+    @assert eltype(matrix_field) <: DiagonalMatrixRow
+    Square()
+end
+
 _matrix_shape(::Type{Int}, _) = Square()
 _matrix_shape(::Type{PlusHalf{Int}}, ::Spaces.CellCenter) = FaceToCenter()
 _matrix_shape(::Type{PlusHalf{Int}}, ::Spaces.CellFace) = CenterToFace()

src/MatrixFields/single_field_solver.jl

@@ -98,10 +98,10 @@ end
 
 function _single_field_solve_col!(
     ::ClimaComms.AbstractCPUDevice,
-    cache::Fields.ColumnField,
-    x::Fields.ColumnField,
+    cache,
+    x,
     A,
-    b::Fields.ColumnField,
+    b,
 )
     if A isa Fields.ColumnField
         band_matrix_solve!(

test/MatrixFields/flat_spaces.jl

@@ -0,0 +1,37 @@
+import ClimaCore
+include(
+    joinpath(pkgdir(ClimaCore), "test", "TestUtilities", "TestUtilities.jl"),
+)
+import .TestUtilities as TU
+
+include("matrix_field_test_utils.jl")
+import ClimaCore.MatrixFields: @name, ⋅
+
+@testset "Matrix Fields with Spectral Element and Point Spaces" begin
+    get_j_field(space, FT) = fill(MatrixFields.DiagonalMatrixRow(FT(1)), space)
+
+    implicit_vars = (@name(tmp.v1), @name(tmp.v2))
+    for FT in (Float32, Float64)
+        comms_ctx = ClimaComms.SingletonCommsContext(comms_device)
+        ps = TU.PointSpace(FT; context = comms_ctx)
+        ses = TU.SpectralElementSpace2D(FT; context = comms_ctx)
+        v1 = Fields.zeros(ps)
+        v2 = Fields.zeros(ses)
+        Y = Fields.FieldVector(; :tmp => (; :v1 => v1, :v2 => v2))
+        implicit_blocks = MatrixFields.unrolled_map(
+            var ->
+                (var, var) =>
+                    get_j_field(axes(MatrixFields.get_field(Y, var)), FT),
+            implicit_vars,
+        )
+        matrix = MatrixFields.FieldMatrix(implicit_blocks...)
+        alg = MatrixFields.BlockDiagonalSolve()
+        solver = MatrixFields.FieldMatrixSolver(alg, matrix, Y)
+        b1 = random_field(FT, ps)
+        b2 = random_field(FT, ses)
+        x = similar(Y)
+        b = Fields.FieldVector(; :tmp => (; :v1 => b1, :v2 => b2))
+        MatrixFields.field_matrix_solve!(solver, x, matrix, b)
+        @test x == b
+    end
+end

test/runtests.jl

@@ -81,6 +81,8 @@ UnitTest("MatrixFields - non-scalar broadcasting (1)" ,"MatrixFields/matrix_fiel
 UnitTest("MatrixFields - non-scalar broadcasting (2)" ,"MatrixFields/matrix_fields_broadcasting/test_non_scalar_2.jl"),
 UnitTest("MatrixFields - non-scalar broadcasting (3)" ,"MatrixFields/matrix_fields_broadcasting/test_non_scalar_3.jl"),
 UnitTest("MatrixFields - non-scalar broadcasting (4)" ,"MatrixFields/matrix_fields_broadcasting/test_non_scalar_4.jl"),
+UnitTest("MatrixFields - flat spaces" ,"MatrixFields/flat_spaces.jl"),
+    
 # UnitTest("MatrixFields - matrix field broadcast"   ,"MatrixFields/matrix_field_broadcasting.jl"), # too long
 # UnitTest("MatrixFields - operator matrices"        ,"MatrixFields/operator_matrices.jl"), # too long
 # UnitTest("MatrixFields - field matrix solvers"     ,"MatrixFields/field_matrix_solvers.jl"), # too long