More DMStag: local_to_local!, with_product_coordinates, set_matrix_preallocate_only!

CHANGELOG.md

@@ -12,6 +12,8 @@
 - `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.
+- `LibPETSc.DMStagRestoreProductCoordinateArraysRead` restored the arrays through the writable restore, which does not match the read-only get.
+- The `KSP` and `SNES` docstrings did not say when the options given to the constructor are applied: once at construction for `KSP`, at every `solve!` for `SNES` (and `TS`).
 
 ### Added
 
@@ -26,6 +28,10 @@
 - `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!`.
+- `local_to_local!(dst, dm, src)` and the in-place `local_to_local!(v, dm)`, to refresh ghost points.
+- `with_product_coordinates(f, dm)`: read-only access to a DMStag's per-axis coordinates, handed back when `f` returns.
+- `set_matrix_preallocate_only!(dm, flag)`, so a matrix from `dm` gets its nonzero pattern from the first assembly.
+- `with_field_views!(f, dm, vecs...; fields, read, write)`: concretely typed views of DMStag local vectors by `location => dof`, indexed like `stencil`, handed back when `f` returns. About 9 allocations per vector, whatever the grid size.
 
 ### Changed
 

docs/src/man/dmstag.md

@@ -82,6 +82,9 @@ lvec = PETSc.local_vec(dm)
 # The written vector comes first, the DM follows it (naming.md §8)
 PETSc.global_to_local!(lvec, dm, gvec, PETSc.INSERT_VALUES)
 PETSc.local_to_global!(gvec, dm, lvec, PETSc.ADD_VALUES)
+
+# Refresh the ghost points of a local vector from its neighbours, in place
+PETSc.local_to_local!(lvec, dm)
 ```
 
 ### Getting Location Indices
@@ -96,6 +99,23 @@ indices = PETSc.local_indices(dm)
 indices = PETSc.global_indices(dm)
 ```
 
+### Views by field
+
+[`with_field_views!`](@ref) checks out local vectors and hands the block one view per field. Views are indexed by the element index, ghosts included, like `stencil`, and their types are concrete, so a loop over them compiles to plain array code:
+
+```julia
+flow = (PETSc.face_location(dm, 1) => 0, PETSc.face_location(dm, 2) => 0,
+        PETSc.element_location(dm) => 0)
+c = PETSc.corners(dm)
+PETSc.with_field_views!(dm, x_local, r_local; fields = flow, write = (false, true)) do (Vx, Vy, P), (Rx, Ry, Rp)
+    for I in c.lower:c.upper
+        Rp[I] = Vx[I + CartesianIndex(1, 0)] - Vx[I] + Vy[I + CartesianIndex(0, 1)] - Vy[I]
+    end
+end
+```
+
+Without `fields`, the block gets each whole array, indexed `[I..., slot]` with `slot` from `dof_slot`. `write = false` checks out read-only.
+
 ### 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:
@@ -150,6 +170,11 @@ PETSc.set_uniform_coordinates!(dm, xmin, xmax, ymin, ymax, zmin, zmax)   # 3D
 
 # Get local coordinate array
 coords = PETSc.local_coordinate_array(dm)
+
+# Read the per-axis coordinates: x[i, 1] is the lower face of element i, x[i, 2] its centre
+PETSc.with_product_coordinates(dm) do x, y
+    x[i, 2], y[j, 2]
+end
 ```
 
 ## Stencil Types
@@ -176,9 +201,11 @@ PETSc.set_uniform_coordinates!(dm, 0.0, 1.0, 0.0, 1.0)
 # Create vectors and matrix
 x = PETSc.global_vec(dm)
 b = PETSc.global_vec(dm)
-A = LibPETSc.DMCreateMatrix(petsclib, dm)
+A = PETSc.PetscMat(dm)
 ```
 
+`PetscMat(dm)` stores every coupling the stencil allows, as explicit zeros. To let the first assembly define the pattern instead, call `PETSc.set_matrix_preallocate_only!(dm, true)` before creating the matrix.
+
 ## Functions
 
 ```@autodocs

scripts/renames.jl

@@ -398,6 +398,9 @@ const INTERNAL = Set{Symbol}([
     :check_axis,
     :stencil_vector,
     :value_vector,
+    :field_slots,
+    :call_with_views,
+    :field_views,
     :TSStepHookFn,
     :_unsafe_local_array,
     :_local_arrays,
@@ -484,6 +487,10 @@ const UNCHANGED_PUBLIC = Symbol[
     :face_location,          # new in v0.5.2
     :element_location,       # new in v0.5.2
     :stencil,                # new in v0.5.2
+    :local_to_local!,        # new in v0.5.2
+    :set_matrix_preallocate_only!,  # new in v0.5.2
+    :with_product_coordinates,      # new in v0.5.2
+    :with_field_views!,             # new in v0.5.2
     :PetscDS,
     :AbstractPetscDS,
     :add_boundary!,

src/autowrapped/DM_wrappers.jl

@@ -43681,7 +43681,7 @@ function DMStagRestoreProductCoordinateArraysRead(petsclib::PetscLibType, dm::Ab
     refZ = arrZ === nothing ? C_NULL : Ref(arrZ.ptr)
 
     @chk ccall(
-               (:DMStagRestoreProductCoordinateArrays, $petsc_library),
+               (:DMStagRestoreProductCoordinateArraysRead, $petsc_library),
                PetscErrorCode,
                (CDM, Ptr{Ptr{Ptr{$PetscScalar}}}, Ptr{Ptr{Ptr{$PetscScalar}}}, Ptr{Ptr{Ptr{$PetscScalar}}}),
                dm, refX, refY, refZ,

src/dm.jl

@@ -408,6 +408,35 @@ function global_to_local!(
     return lvec
 end
 
+"""
+    local_to_local!(dst::AbstractPetscVec, dm::AbstractPetscDM, src::AbstractPetscVec, mode = INSERT_VALUES)
+    local_to_local!(v::AbstractPetscVec, dm::AbstractPetscDM, mode = INSERT_VALUES)
+
+Fill the ghost points of the local vector `dst` from the owned values of the local
+vector `src` on the neighbouring ranks, and return `dst`. `dst` may be `src`: the
+second form refreshes the ghost points of `v` in place.
+
+# External Links
+$(doc_external("DM/DMLocalToLocalBegin"))
+$(doc_external("DM/DMLocalToLocalEnd"))
+"""
+function local_to_local!(
+    dst::AbstractPetscVec{PetscLib},
+    dm::AbstractPetscDM{PetscLib},
+    src::AbstractPetscVec{PetscLib},
+    mode::InsertMode = INSERT_VALUES,
+) where {PetscLib}
+    LibPETSc.DMLocalToLocalBegin(getlib(PetscLib), dm, src, mode, dst)
+    LibPETSc.DMLocalToLocalEnd(getlib(PetscLib), dm, src, mode, dst)
+    return dst
+end
+
+local_to_local!(
+    v::AbstractPetscVec{PetscLib},
+    dm::AbstractPetscDM{PetscLib},
+    mode::InsertMode = INSERT_VALUES,
+) where {PetscLib} = local_to_local!(v, dm, v, mode)
+
 
 """
     set_uniform_coordinates!(
@@ -594,3 +623,22 @@ function LibPETSc.PetscMat(da::AbstractPetscDM{PetscLib}) where {PetscLib}
     return J
 end
 
+"""
+    set_matrix_preallocate_only!(dm::AbstractPetscDM, flag::Bool)
+
+With `flag` set, a matrix made from `dm` by [`PetscMat`](@ref) is preallocated but
+its nonzero pattern is not inserted, so the first assembly defines the pattern
+instead of every coupling the stencil allows being stored as an explicit zero.
+The setting stays on `dm` for every later matrix. Returns `dm`.
+
+# External Links
+$(doc_external("DM/DMSetMatrixPreallocateOnly"))
+"""
+function set_matrix_preallocate_only!(
+    dm::AbstractPetscDM{PetscLib},
+    flag::Bool,
+) where {PetscLib}
+    LibPETSc.DMSetMatrixPreallocateOnly(getlib(PetscLib), dm, LibPETSc.PetscBool(flag))
+    return dm
+end
+

src/dmstag.jl

@@ -669,3 +669,132 @@ 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)
+
+# ============================================================================
+#   Coordinates
+# ============================================================================
+
+"""
+    with_product_coordinates(f, dm::DMStag)
+
+Call `f(x)`, `f(x, y)` or `f(x, y, z)` with the local coordinate arrays of `dm`,
+one per axis, and return what `f` returns. It needs product coordinates, which
+[`set_uniform_coordinates!`](@ref) sets. The arrays are PETSc's own and are read
+only: they are handed back when `f` returns or throws, and writing to them, or
+using them after `f`, is undefined.
+
+Each array is indexed `[i, slot]`: `i` is the 1-based element index, ghost
+elements included, as [`ghost_corners`](@ref) gives it, and `slot` is 1 for the
+coordinate of the element's lower face and 2 for its centre.
+
+```julia
+with_product_coordinates(dm) do x, y
+    x[i, 1], x[i, 2]   # x of the lower face and of the centre of element i
+end
+```
+
+# External Links
+$(doc_external("DMStag/DMStagGetProductCoordinateArraysRead"))
+$(doc_external("DMStag/DMStagRestoreProductCoordinateArraysRead"))
+"""
+function with_product_coordinates(f, dm::DMStag{PetscLib, N}) where {PetscLib, N}
+    lib = getlib(PetscLib)
+    x, y, z = LibPETSc.DMStagGetProductCoordinateArraysRead(lib, dm)
+    try
+        N == 1 && return f(x)
+        N == 2 && return f(x, y)
+        return f(x, y, z)
+    finally
+        LibPETSc.DMStagRestoreProductCoordinateArraysRead(lib, dm, x, y, z)
+    end
+end
+
+# ============================================================================
+#   Views of a local vector by field
+# ============================================================================
+
+"""
+    with_field_views!(f, dm::DMStag, vecs::AbstractPetscVec...; fields = nothing, read = true, write = true)
+
+Check out the local vectors `vecs` of `dm`, call `f` with one argument per vector,
+hand the vectors back when `f` returns or throws, and return what `f` returns.
+
+`fields` is a tuple of `location => dof` pairs, as [`IS`](@ref LibPETSc.IS) takes
+them, with the 0-based component of [`stencil`](@ref). For each vector `f` gets a
+tuple of views, one per field and in that order. Without `fields`, `f` gets the
+whole array per vector, indexed `[I..., slot]` with `slot` from [`dof_slot`](@ref).
+Either way the element index `I` is 1-based and covers the ghost elements, as in
+[`ghost_corners`](@ref), so a view and [`stencil`](@ref) address the same point.
+
+`read` and `write` work as in [`with_local_array!`](@ref): a `Bool`, or one per
+vector. `write = false` checks out read-only; the views have the same types either
+way. The vectors must be local vectors of `dm` ([`local_vec`](@ref)); anything else
+throws an `ArgumentError` before any vector is checked out.
+
+```julia
+flow = (face_location(dm, 1) => 0, face_location(dm, 2) => 0, element_location(dm) => 0)
+c = corners(dm)
+with_field_views!(dm, x, r; fields = flow, write = (false, true)) do (Vx, Vy, P), (Rx, Ry, Rp)
+    for I in c.lower:c.upper      # the owned elements
+        Rp[I] = Vx[I + CartesianIndex(1, 0)] - Vx[I] + Vy[I + CartesianIndex(0, 1)] - Vy[I]
+    end
+end
+```
+
+# External Links
+$(doc_external("DMStag/DMStagVecGetArray"))
+$(doc_external("DMStag/DMStagGetLocationSlot"))
+"""
+function with_field_views!(
+    f,
+    dm::DMStag{PetscLib, N},
+    vecs::Vararg{AbstractPetscVec{PetscLib}, M};
+    fields = nothing,
+    read::Union{Bool, NTuple{M, Bool}} = true,
+    write::Union{Bool, NTuple{M, Bool}} = true,
+) where {PetscLib, N, M}
+    read  isa Bool && (read  = ntuple(_ -> read,  Val(M)))
+    write isa Bool && (write = ntuple(_ -> write, Val(M)))
+    lib = getlib(PetscLib)
+    gc = ghost_corners(dm)
+    q = Int(LibPETSc.DMStagGetEntriesPerElement(lib, dm))
+    dims = (q, gc.size...)
+    offsets = ntuple(d -> gc.lower[d] - 1, Val(N))
+    slots = field_slots(dm, fields)
+    for v in vecs
+        LibPETSc.VecGetLocalSize(lib, v) == prod(dims) || throw(
+            ArgumentError("with_field_views! takes local vectors of the DMStag (local_vec(dm))"),
+        )
+    end
+    acquired = map(vecs, read, write) do v, r, w
+        acquire_local_array(v; read = r, write = w)
+    end
+    try
+        return call_with_views(f, map(first, acquired), dims, offsets, slots)
+    finally
+        foreach(vecs, acquired, read, write) do v, (_, cpu_arr, backend), r, w
+            release_local_array(cpu_arr, backend, v; read = r, write = w)
+        end
+    end
+end
+
+# The 1-based slot of each field, or `nothing` for the whole array
+field_slots(dm, ::Nothing) = nothing
+field_slots(dm::DMStag{PetscLib}, fields::Tuple) where {PetscLib} = map(fields) do (loc, dof)
+    dof_slot(dm, loc, Int(dof))
+end
+
+# A function barrier: the array types come from a runtime memory type, and
+# everything below, `f` included, compiles for the concrete ones
+function call_with_views(f, arrays, dims, offsets, slots)
+    return f(map(a -> field_views(reshape(a, dims), offsets, slots), arrays)...)
+end
+
+# The array is stored with the component first; the views put it last, or drop it
+function field_views(A::AbstractArray{T, D}, offsets::NTuple{N, Int}, ::Nothing) where {T, D, N}
+    whole = PermutedDimsArray(A, (ntuple(d -> d + 1, Val(N))..., 1))
+    return OffsetArray(whole, offsets..., 0)
+end
+
+field_views(A::AbstractArray{T, D}, offsets::NTuple{N, Int}, slots::Tuple) where {T, D, N} =
+    map(slot -> OffsetArray(view(A, slot, ntuple(_ -> Colon(), Val(N))...), offsets...), slots)

src/ksp.jl

@@ -61,6 +61,9 @@ end
 Create a `KSP` using the matrix `A` and preconditioner construction matrix `P`
 with optional `prefix` and `options`.
 
+The options are applied here, once, through `KSPSetFromOptions`. 
+A setting made in code afterwards, such as [`set_tolerances!`](@ref), overrides them.
+
 The communicator is obtained from `A` and if it has size `1` then the garbage
 collector is set, otherwise the user is responsible for calling
 [`destroy!`](@ref).
@@ -110,6 +113,9 @@ end
 
 Create a `KSP` associated with the `dm` with optional `prefix` and `options`.
 
+The options are applied here, once, through `KSPSetFromOptions`.
+A setting made in code afterwards overrides them.
+
 The communicator is obtained from `dm`. The KSP can be used with geometric
 multigrid when the DM provides grid hierarchy information.
 

src/public_names.jl

@@ -20,24 +20,25 @@ public @bd_fn, @jacobian_fn, @residual_fn, @simple_fn, AbstractPetscDS,
        inttype, isassembled, isfinalized, isinitialized, issimplex, iteration_number, ksp,
        ksp_iterations, l2diff, label, library_info, local_arrays, local_coordinate_array,
        local_coordinates, local_indices, local_interior_linear_index, local_to_global!,
-       local_vec, mat_nullspace_create, max_steps, max_time, memtype, memtype_backend,
-       narrow, ndofs, options_prefix, ownership_range, owns, parse_option, parse_options,
-       pc, prev_time, project_field!, project_function!, release_local_array, reset!,
-       reshape_local_array, restore_local_arrays!, save_vtk, scalartype, set_adapt_type!,
-       set_auxiliary_vec!, set_compute_operators!, set_compute_rhs!, set_constants!,
-       set_convergence_test!, set_dm!, set_dm_active!, set_equation_type!,
+       local_to_local!, local_vec, mat_nullspace_create, max_steps, max_time, memtype,
+       memtype_backend, narrow, ndofs, options_prefix, ownership_range, owns, parse_option,
+       parse_options, pc, prev_time, project_field!, project_function!, release_local_array,
+       reset!, reshape_local_array, restore_local_arrays!, save_vtk, scalartype,
+       set_adapt_type!, set_auxiliary_vec!, set_compute_operators!, set_compute_rhs!,
+       set_constants!, set_convergence_test!, set_dm!, set_dm_active!, set_equation_type!,
        set_error_if_step_fails!, set_exact_final_time!, set_exact_solution!, set_field!,
        set_fieldsplit_is!, set_from_options!, set_function!, set_function_domain_error!,
        set_ifunction!, set_ijacobian!, set_jacobian!, set_jacobian_nullspace!,
-       set_jacobian_preconditioner!, set_library!, set_max_snes_failures!,
-       set_max_step_rejections!, set_max_steps!, set_max_time!, set_monitor!, set_name!,
-       set_nullspace!, set_operators!, set_option!, set_options_prefix!, set_petscint!,
-       set_post_step!, set_pre_step!, set_problem_type!, set_residual!, set_rhs_function!,
-       set_rhs_jacobian!, set_shell_apply!, set_shell_setup!, set_snes_jacobian!,
-       set_snes_local_fem!, set_solution!, set_step_number!, set_time!, set_timestep!,
-       set_tolerances!, set_type!, set_uniform_coordinates!, set_user_ctx!, set_values!,
-       setup!, snes, snes_failures, snes_iterations, solution, solve!, solve_time,
-       star_fd_coloring, stencil, step!, step_number, step_rejections,
-       tao_usable_after_reinitialize, timestep, tolerances, type_name, unsafe_local_array,
-       unset_library!, user_ctx, vertex_location, vtk_merge_tensor, with_local_array!,
+       set_jacobian_preconditioner!, set_library!, set_matrix_preallocate_only!,
+       set_max_snes_failures!, set_max_step_rejections!, set_max_steps!, set_max_time!,
+       set_monitor!, set_name!, set_nullspace!, set_operators!, set_option!,
+       set_options_prefix!, set_petscint!, set_post_step!, set_pre_step!, set_problem_type!,
+       set_residual!, set_rhs_function!, set_rhs_jacobian!, set_shell_apply!,
+       set_shell_setup!, set_snes_jacobian!, set_snes_local_fem!, set_solution!,
+       set_step_number!, set_time!, set_timestep!, set_tolerances!, set_type!,
+       set_uniform_coordinates!, set_user_ctx!, set_values!, setup!, snes, snes_failures,
+       snes_iterations, solution, solve!, solve_time, star_fd_coloring, stencil, step!,
+       step_number, step_rejections, tao_usable_after_reinitialize, timestep, tolerances,
+       type_name, unsafe_local_array, unset_library!, user_ctx, vertex_location,
+       vtk_merge_tensor, with_field_views!, with_local_array!, with_product_coordinates,
        wrap_local_array, zero_rows!, zero_rows_local!

src/snes.jl

@@ -67,6 +67,11 @@ end
 
 Create a PETSc nonlinear solver (SNES) context on the communicator `comm`.
 
+The options are stored and applied at the start of every [`solve!`](@ref), so that a
+DM and callbacks attached after construction are visible to `SNESSetFromOptions`.
+Because they are applied on every solve, they override a setting made in code for
+the same option.
+
 # Arguments
 - `petsclib`: The PETSc library instance
 - `comm::MPI.Comm`: MPI communicator

src/ts.jl

@@ -248,6 +248,8 @@ Create a PETSc time stepper on the communicator `comm`.
 
 Options are stored and applied in [`solve!`](@ref) rather than here, so that a
 DM and callbacks attached after construction are visible to `TSSetFromOptions`.
+They are applied on every `solve!`, and override a setting made in code for the
+same option. [`step!`](@ref) does not apply them.
 
 `exact_final_time` defaults to `TS_EXACTFINALTIME_MATCHSTEP`, so the last step
 lands on the time set by [`set_max_time!`](@ref).