More DMStag stuff: locations by axis, stencils from Julia indices, stencil assembly

CHANGELOG.md

@@ -11,6 +11,7 @@
 - `LibPETSc.MatCreateSubMatrices` and `MatCreateSubMatricesMPI` did not work with `MAT_REUSE_MATRIX`, and `MatDestroySubMatrices` could not take their result. Pass the returned vector back to either.
 - `LibPETSc.VecNestRestoreSubVecsRead` accepts the vector `VecNestGetSubVecsRead` returns; before, it only took a raw pointer that no wrapper handed out, so the read lock could not be released.
 - `LibPETSc.PCASMDestroySubdomains` and `PCGASMDestroySubdomains` accept the vectors the subdomain creators return.
+- On an Int32 PETSc build (the `PetscInt = "Int32"` preference), `LibPETSc.DMStagStencil`, `MatStencil` and the other C structs had `Int64` fields, so PETSc read them wrong. They now take the integer width of the loaded library, and the two stencil constructors convert their indices to it.
 
 ### Added
 
@@ -21,6 +22,10 @@
 - Long `TS` runs: `set_pre_step!(f!, ts)` and `set_post_step!(f!, ts)` hooks that may change the run, `set_max_snes_failures!`, `set_max_step_rejections!`, `set_error_if_step_fails!`, `set_step_number!` for restarts, and `equation_type`/`set_equation_type!`.
 - `set_solution!(snes, x)`, as on `TS`.
 - `LibPETSc.PETSC_UNLIMITED` and `LibPETSc.PETSC_CURRENT`, which PETSc's manual pages ask for.
+- DMStag locations by axis: `vertex_location(dm)`, `edge_location(dm, a, b)`, `face_location(dm, axis)` and `element_location(dm)`, which mean the same in 2D and 3D, unlike `DMSTAG_DOWN`.
+- `stencil(dm, loc, I; dof = 0)` builds a `DMStagStencil` from a 1-based element index, allocation free.
+- Assembly with DMStag stencils: `set_values!(J, dm, rows, cols, vals, mode)`, `set_values!(v, dm, positions, vals, mode)` and `zero_rows_local!(J, dm, rows, diag)`. They take any `AbstractVector`.
+- `LibPETSc.IS(dm, loc => dof, ...)`, the index set of whole DMStag fields, for `set_fieldsplit_is!`.
 
 ### Changed
 

docs/src/man/dmstag.md

@@ -96,6 +96,50 @@ indices = PETSc.local_indices(dm)
 indices = PETSc.global_indices(dm)
 ```
 
+### Locations and stencils
+
+DMStag names a point by where it sits on an element: `DMSTAG_LEFT`, `DMSTAG_DOWN`, `DMSTAG_BACK_DOWN_LEFT`, and so on. Those names shift meaning between dimensions (`DOWN` is the second axis, whatever it's called in your model), so the location functions count axes instead:
+
+```julia
+PETSc.vertex_location(dm)        # the element's lower corner
+PETSc.face_location(dm, axis)    # the face normal to `axis`, on the lower side
+PETSc.edge_location(dm, a, b)    # the edge touching the lower faces of `a` and `b` (3D; the vertex in 2D)
+PETSc.element_location(dm)       # the element interior
+```
+
+`face_location(dm, ndims(dm))` is the last axis in any dimension. An axis outside `1:ndims(dm)` throws an `ArgumentError`.
+
+A stencil addresses one unknown: a location, the element index and a component. [`stencil`](@ref) takes the 1-based element index that `corners` and `ghost_corners` use, and PETSc's 0-based component, as [`dof_slot`](@ref) does. Indices are not bounds-checked, so ghost elements work, and the call allocates nothing:
+
+```julia
+I = corners(dm).lower
+row  = PETSc.stencil(dm, PETSc.face_location(dm, 1), I)
+cols = [PETSc.stencil(dm, PETSc.element_location(dm), I),
+        PETSc.stencil(dm, PETSc.element_location(dm), I + CartesianIndex(1, 0))]
+```
+
+### Assembling with stencils
+
+```julia
+J = PETSc.PetscMat(dm)
+PETSc.set_values!(J, dm, [row], cols, [-1.0, 1.0], PETSc.ADD_VALUES)   # row-major block
+PETSc.assemble!(J)
+
+PETSc.zero_rows_local!(J, dm, wall_rows, 1.0)   # Dirichlet rows, given as stencils
+PETSc.set_values!(b, dm, wall_rows, wall_values)
+```
+
+Under `ADD_VALUES`, repeated (row, column) entries in one call are summed. `zero_rows_local!` is collective: a rank that owns no wall passes an empty vector.
+
+An index set of whole fields, for a field split, takes (location, component) pairs:
+
+```julia
+flow = LibPETSc.IS(dm, PETSc.face_location(dm, 1) => 0,
+                       PETSc.face_location(dm, 2) => 0,
+                       PETSc.element_location(dm) => 0)
+PETSc.set_fieldsplit_is!(PETSc.pc(ksp), "flow", flow)
+```
+
 ## Setting Coordinates
 
 ```julia

scripts/renames.jl

@@ -395,6 +395,9 @@ const INTERNAL = Set{Symbol}([
     # the underscore stays: inner worker beside a wrapper of the same name
     :_mul!,
     :_zero_rows_vecs,
+    :check_axis,
+    :stencil_vector,
+    :value_vector,
     :TSStepHookFn,
     :_unsafe_local_array,
     :_local_arrays,
@@ -476,6 +479,11 @@ const UNCHANGED_PUBLIC = Symbol[
     :DMPlex,
     :narrow,
     :ndofs,
+    :vertex_location,        # new in v0.5.2
+    :edge_location,          # new in v0.5.2
+    :face_location,          # new in v0.5.2
+    :element_location,       # new in v0.5.2
+    :stencil,                # new in v0.5.2
     :PetscDS,
     :AbstractPetscDS,
     :add_boundary!,

src/autowrapped/petsc_library.jl

@@ -37,7 +37,9 @@ const PETSC_CURRENT = -2     # keep the value already set
 const PETSC_UNLIMITED = -3   # no limit, where a function takes a count or a tolerance
 const PETSC_COMM_SELF = MPI.COMM_SELF
 
-PetscInt = Int64
+# The one integer width this process loads (LibPETSc_startup.jl): 
+# C structs such as DMStagStencil and MatStencil take their field types from it
+PetscInt = libs[1][3]
 PetscInt64 = Int64
 PetscInt32 = Int32
 PetscScalar = Float64

src/autowrapped/struct_wrappers.jl

@@ -73,7 +73,8 @@ struct DMStagStencil
 
     DMStagStencil() = new()
 
-    DMStagStencil(loc,i,j,k,c) = new(loc::DMStagStencilLocation,i::PetscInt ,    j::PetscInt ,    k::PetscInt ,    c::PetscInt)
+    # `new` converts the indices to the PetscInt width this process loads
+    DMStagStencil(loc,i,j,k,c) = new(loc::DMStagStencilLocation, i, j, k, c)
 
 end 
 
@@ -87,7 +88,8 @@ struct MatStencil
     MatStencil() = new()
 
     # Convenience constructor accepting integers in i,j,k,c order
-    MatStencil(i,j,k,c) = new(i::PetscInt, j::PetscInt, k::PetscInt, c::PetscInt)
+    # `new` converts the indices to the PetscInt width this process loads
+    MatStencil(i,j,k,c) = new(i, j, k, c)
 end 
 
 struct MatInfo

src/dmstag.jl

@@ -412,9 +412,260 @@ end
     slot::Int = dof_slot(dm::DMStag, loc::LibPETSc.DMStagStencilLocation, dof::Int) 
 
 Returns the location `slot` for a degree of freedom `dof` at a given stencil location `loc` in the DMStag `dm`.
-Note that the returned `slot` is 1-based for Julia compatibility.    
+Note that the returned `slot` is 1-based for Julia compatibility.
+`dof` is PETSc's 0-based component number at that location, as in [`stencil`](@ref).
 """
 function dof_slot(dm::DMStag{PetscLib}, loc::LibPETSc.DMStagStencilLocation, dof::Int) where {PetscLib} 
     slot = LibPETSc.DMStagGetLocationSlot(getlib(PetscLib), dm, loc, PetscLib.PetscInt(dof))
     return slot+1
 end
+
+
+# ============================================================================
+#   Locations and stencils
+# ============================================================================
+
+"""
+    vertex_location(dm::DMStag)
+
+The location of the vertex DMStag stores with each element, its lower corner:
+`DMSTAG_LEFT` in 1D, `DMSTAG_DOWN_LEFT` in 2D and `DMSTAG_BACK_DOWN_LEFT` in 3D.
+
+See also [`face_location`](@ref), [`edge_location`](@ref), [`element_location`](@ref)
+and [`stencil`](@ref).
+
+# External Links
+$(doc_external("DMStag/DMStagStencilLocation"))
+"""
+vertex_location(::DMStag{PetscLib, 1}) where {PetscLib} = LibPETSc.DMSTAG_LEFT
+vertex_location(::DMStag{PetscLib, 2}) where {PetscLib} = LibPETSc.DMSTAG_DOWN_LEFT
+vertex_location(::DMStag{PetscLib, 3}) where {PetscLib} = LibPETSc.DMSTAG_BACK_DOWN_LEFT
+
+"""
+    edge_location(dm::DMStag, a::Integer, b::Integer)
+
+The location of the edge touching the lower faces of axes `a` and `b`, in either
+order: the edge a shear component τ_ab lives on. In 3D, `(1, 2)` is
+`DMSTAG_DOWN_LEFT`, `(1, 3)` is `DMSTAG_BACK_LEFT` and `(2, 3)` is
+`DMSTAG_BACK_DOWN`. In 2D the vertex plays that role, so `(1, 2)` returns
+`DMSTAG_DOWN_LEFT`. A 1D DM has no edges, and an axis outside `1:ndims(dm)`, or
+`a == b`, throws an `ArgumentError`.
+
+See also [`vertex_location`](@ref), [`face_location`](@ref).
+
+# External Links
+$(doc_external("DMStag/DMStagStencilLocation"))
+"""
+function edge_location(dm::DMStag{PetscLib, N}, a::Integer, b::Integer) where {PetscLib, N}
+    N == 1 && throw(ArgumentError("a 1D DMStag has no edges"))
+    check_axis(dm, a)
+    check_axis(dm, b)
+    a == b && throw(ArgumentError("an edge needs two different axes, got ($a, $b)"))
+    lo, hi = minmax(a, b)
+    lo == 1 && hi == 2 && return LibPETSc.DMSTAG_DOWN_LEFT
+    lo == 1 && return LibPETSc.DMSTAG_BACK_LEFT
+    return LibPETSc.DMSTAG_BACK_DOWN
+end
+
+"""
+    face_location(dm::DMStag, axis::Integer)
+
+The location of the face normal to `axis` on the lower side of each element:
+`DMSTAG_LEFT` for axis 1, `DMSTAG_DOWN` for axis 2 and `DMSTAG_BACK` for axis 3.
+Axes count the DM's own axes, so `face_location(dm, ndims(dm))` is the last axis in
+any dimension. An axis outside `1:ndims(dm)` throws an `ArgumentError`.
+
+See also [`vertex_location`](@ref), [`edge_location`](@ref), [`element_location`](@ref).
+
+# External Links
+$(doc_external("DMStag/DMStagStencilLocation"))
+"""
+function face_location(dm::DMStag, axis::Integer)
+    check_axis(dm, axis)
+    axis == 1 && return LibPETSc.DMSTAG_LEFT
+    axis == 2 && return LibPETSc.DMSTAG_DOWN
+    return LibPETSc.DMSTAG_BACK
+end
+
+"""
+    element_location(dm::DMStag)
+
+The location of the element interior, `DMSTAG_ELEMENT` in every dimension.
+
+See also [`vertex_location`](@ref), [`face_location`](@ref), [`edge_location`](@ref).
+
+# External Links
+$(doc_external("DMStag/DMStagStencilLocation"))
+"""
+element_location(::DMStag) = LibPETSc.DMSTAG_ELEMENT
+
+function check_axis(::DMStag{PetscLib, N}, axis::Integer) where {PetscLib, N}
+    1 <= axis <= N || throw(ArgumentError("axis $axis is not an axis of a $(N)D DMStag"))
+    return nothing
+end
+
+"""
+    stencil(dm::DMStag, loc::LibPETSc.DMStagStencilLocation, I; dof = 0)
+
+The `LibPETSc.DMStagStencil` addressing component `dof` at location `loc` of
+element `I`. `I` is a `CartesianIndex{N}` or an `NTuple{N, Integer}` with the
+1-based element indices [`corners`](@ref) and [`ghost_corners`](@ref) use, and the
+stencil holds them 0-based, as PETSc reads them. Indices are not checked, so a ghost
+element (0 or `N + 1` on a periodic axis) is valid. `dof` is PETSc's 0-based
+component number, as in [`dof_slot`](@ref).
+
+Allocation free, for use inside assembly loops.
+
+```julia
+lower = corners(dm).lower
+row = stencil(dm, face_location(dm, 1), lower)             # the first owned x face
+col = stencil(dm, element_location(dm), lower; dof = 1)    # second element component
+```
+
+See also [`set_values!`](@ref), [`zero_rows_local!`](@ref).
+
+# External Links
+$(doc_external("DMStag/DMStagStencil"))
+"""
+@inline function stencil(
+    dm::DMStag{PetscLib, N},
+    loc::LibPETSc.DMStagStencilLocation,
+    I::NTuple{N, Integer};
+    dof::Integer = 0,
+) where {PetscLib, N}
+    T = inttype(PetscLib)
+    i = ntuple(d -> d <= N ? T(I[d] - 1) : zero(T), Val(3))
+    return LibPETSc.DMStagStencil(loc, i[1], i[2], i[3], T(dof))
+end
+
+@inline stencil(
+    dm::DMStag{PetscLib, N},
+    loc::LibPETSc.DMStagStencilLocation,
+    I::CartesianIndex{N};
+    dof::Integer = 0,
+) where {PetscLib, N} = stencil(dm, loc, Tuple(I); dof)
+
+# ============================================================================
+#   Operations on stencils
+# ============================================================================
+
+"""
+    set_values!(J::AbstractPetscMat, dm::DMStag, rows, cols, vals, mode = INSERT_VALUES)
+
+Write the dense block `vals` into `J` at the stencils `rows` × `cols` and return
+`J`. `vals` is row-major, with `length(rows) * length(cols)` entries. Under
+`ADD_VALUES`, entries whose row and column repeat are summed. `rows`, `cols` and
+`vals` can be any `AbstractVector`; one that is not a `Vector` is copied first.
+
+# External Links
+$(doc_external("DMStag/DMStagMatSetValuesStencil"))
+"""
+function set_values!(
+    J::AbstractPetscMat{PetscLib},
+    dm::DMStag{PetscLib},
+    rows::AbstractVector{LibPETSc.DMStagStencil},
+    cols::AbstractVector{LibPETSc.DMStagStencil},
+    vals::AbstractVector,
+    mode::InsertMode = INSERT_VALUES,
+) where {PetscLib}
+    length(vals) == length(rows) * length(cols) || throw(
+        DimensionMismatch(
+            "a $(length(rows))x$(length(cols)) block needs " *
+            "$(length(rows) * length(cols)) values, got $(length(vals))",
+        ),
+    )
+    T = inttype(PetscLib)
+    LibPETSc.DMStagMatSetValuesStencil(
+        getlib(PetscLib), dm, J,
+        T(length(rows)), stencil_vector(rows),
+        T(length(cols)), stencil_vector(cols),
+        value_vector(PetscLib, vals), mode,
+    )
+    return J
+end
+
+"""
+    set_values!(v::AbstractPetscVec, dm::DMStag, positions, vals, mode = INSERT_VALUES)
+
+Write `vals` into `v` at the stencils `positions` and return `v`. Both can be any
+`AbstractVector` of the same length; one that is not a `Vector` is copied first.
+
+# External Links
+$(doc_external("DMStag/DMStagVecSetValuesStencil"))
+"""
+function set_values!(
+    v::AbstractPetscVec{PetscLib},
+    dm::DMStag{PetscLib},
+    positions::AbstractVector{LibPETSc.DMStagStencil},
+    vals::AbstractVector,
+    mode::InsertMode = INSERT_VALUES,
+) where {PetscLib}
+    length(vals) == length(positions) || throw(
+        DimensionMismatch("$(length(positions)) positions, but $(length(vals)) values"),
+    )
+    LibPETSc.DMStagVecSetValuesStencil(
+        getlib(PetscLib), dm, v, inttype(PetscLib)(length(positions)),
+        stencil_vector(positions), value_vector(PetscLib, vals), mode,
+    )
+    return v
+end
+
+"""
+    zero_rows_local!(J::AbstractPetscMat, dm::DMStag, rows, diag = 1; x = nothing, b = nothing)
+
+[`zero_rows_local!`](@ref) with the rows given as stencils of `dm`, which `J` was
+created from. Collective: a rank that owns none of the rows passes an empty vector.
+
+# External Links
+$(doc_external("DMStag/DMStagStencilToIndexLocal"))
+"""
+function zero_rows_local!(
+    J::AbstractPetscMat{PetscLib},
+    dm::DMStag{PetscLib, N},
+    rows::AbstractVector{LibPETSc.DMStagStencil},
+    diag = 1;
+    x = nothing,
+    b = nothing,
+) where {PetscLib, N}
+    T = inttype(PetscLib)
+    rows_0b = LibPETSc.DMStagStencilToIndexLocal(
+        getlib(PetscLib), dm, T(N), T(length(rows)), stencil_vector(rows),
+    )
+    return zero_rows_local!(J, rows_0b, diag; x, b)
+end
+
+"""
+    LibPETSc.IS(dm::DMStag, loc => dof, ...)
+    LibPETSc.IS(dm::DMStag, pairs::AbstractVector{<:Pair})
+
+The index set, in the global numbering, of every point of `dm` at the given
+locations and components: each pair is a `LibPETSc.DMStagStencilLocation` and a
+0-based component, as [`stencil`](@ref) takes them. The caller owns the result. It
+suits [`set_fieldsplit_is!`](@ref):
+
+```julia
+flow = LibPETSc.IS(dm, face_location(dm, 1) => 0, face_location(dm, 2) => 0,
+                   element_location(dm) => 0)
+set_fieldsplit_is!(pc(ksp), "flow", flow)
+```
+
+# External Links
+$(doc_external("DMStag/DMStagCreateISFromStencils"))
+"""
+LibPETSc.IS(dm::DMStag, first::Pair, rest::Pair...) = LibPETSc.IS(dm, [first, rest...])
+
+function LibPETSc.IS(dm::DMStag{PetscLib}, pairs::AbstractVector{<:Pair}) where {PetscLib}
+    T = inttype(PetscLib)
+    stencils = LibPETSc.DMStagStencil[
+        LibPETSc.DMStagStencil(loc, zero(T), zero(T), zero(T), T(dof)) for (loc, dof) in pairs
+    ]
+    return LibPETSc.DMStagCreateISFromStencils(getlib(PetscLib), dm, T(length(stencils)), stencils)
+end
+
+# The C calls take a `Vector`; any other vector is copied into one
+stencil_vector(v::Vector{LibPETSc.DMStagStencil}) = v
+stencil_vector(v::AbstractVector{LibPETSc.DMStagStencil}) = collect(v)
+value_vector(::Type{PetscLib}, v::Vector) where {PetscLib} =
+    eltype(v) === scalartype(PetscLib) ? v : Vector{scalartype(PetscLib)}(v)
+value_vector(::Type{PetscLib}, v::AbstractVector) where {PetscLib} =
+    Vector{scalartype(PetscLib)}(v)