Revise docstrings across mpi-igg branch

  Add missing docstrings for exported and internal functions (set_backend,
  global_dot/sum/length/min, scalar_halo!, velocity_halo!, global_perdot,
  L₂, conv_diff!, BDIM!, project!, CFL, _apply_offset, up/down/restrict/
  restrictL/restrictML/restrictL!, Vcycle!, MultiLevelPoisson solver!).

  Fix inaccuracies and typos and update existing docstrings with MPI-relevant details

Author: b-fg

ext/WaterLilyMPIExt.jl

@@ -10,9 +10,9 @@ hooks through `par_mode[]` (defaults to `Serial()`).  This extension defines
 `Parallel <: AbstractParMode` and adds new dispatch methods — no method
 overwriting, so precompilation works normally.
 
-Functions with MPI-specific behavior (via dispatch or subtype specialization):
+Functions with MPI-specific behavior (via dispatch on `::Parallel`):
   _wallBC_L!  — zero L at physical walls only (skip MPI-internal) + halo on L
-  _exitBC!    — global reductions for inflow/outflow mass flux (dispatches on ::Parallel)
+  _exitBC!    — global reductions for inflow/outflow mass flux
   _divisible  — stricter coarsening threshold (N>8) for multigrid
 """
 module WaterLilyMPIExt

src/AutoBody.jl

@@ -12,6 +12,13 @@ struct AutoBody{F1<:Function,F2<:Function} <: AbstractBody
     map::F2
 end
 AutoBody(sdf, map=(x,t)->x) = AutoBody(sdf, map)
+"""
+    _apply_offset(body::AutoBody, offset)
+
+Wrap the body's `map` function so that rank-local coordinates are shifted
+by `offset` before evaluation.  This lets users write SDFs in global
+coordinates while each MPI rank indexes with local coordinates.
+"""
 _apply_offset(body::AutoBody, offset) =
     AutoBody(body.sdf, let m=body.map, o=offset; (x,t)->m(x .+ o, t); end)
 

src/Body.jl

@@ -97,7 +97,13 @@ struct SetBody{O<:Function,Ta<:AbstractBody,Tb<:AbstractBody} <: AbstractBody
     b::Tb
 end
 
-# Default _apply_offset: identity for generic bodies
+"""
+    _apply_offset(body::AbstractBody, offset)
+
+Wrap the body's coordinate mapping so rank-local coordinates are shifted by
+`offset` before evaluation.  Default is identity (no-op).  `SetBody` recurses
+into its children; `AutoBody` wraps the `map` function (see `AutoBody.jl`).
+"""
 _apply_offset(body::AbstractBody, offset) = body
 _apply_offset(body::SetBody, offset) = SetBody(body.op, _apply_offset(body.a, offset), _apply_offset(body.b, offset))
 

src/Flow.jl

@@ -32,6 +32,13 @@ function median(a,b,c)
     return a
 end
 
+"""
+    conv_diff!(r, u, Φ, λ; ν=0.1, perdir=())
+
+Compute the convective and diffusive fluxes of velocity `u` using the
+convection scheme `λ(u,c,d)` and kinematic viscosity `ν`.  The result is
+accumulated into `r` (force per unit volume) and `Φ` is a working scalar.
+"""
 function conv_diff!(r,u,Φ,λ::F;ν=0.1,perdir=()) where {F}
     r .= zero(eltype(r))
     N,n = size_u(u)
@@ -53,14 +60,17 @@ accelerate!(r,t,g::Function,::Union{Nothing,Tuple}) = accelerate!(r,t,g)
 accelerate!(r,t,::Nothing,U::Function) = accelerate!(r,t,(i,x,t)->ForwardDiff.derivative(τ->U(i,x,τ),t))
 accelerate!(r,t,g::Function,U::Function) = accelerate!(r,t,(i,x,t)->g(i,x,t)+ForwardDiff.derivative(τ->U(i,x,τ),t))
 """
-    Flow{D::Int, T::Float, Sf<:AbstractArray{T,D}, Vf<:AbstractArray{T,D+1}, Tf<:AbstractArray{T,D+2}}
+    Flow{D, T, Sf, Vf, Tf}
 
 Composite type for a multidimensional immersed boundary flow simulation.
 
 Flow solves the unsteady incompressible [Navier-Stokes equations](https://en.wikipedia.org/wiki/Navier%E2%80%93Stokes_equations) on a Cartesian grid.
 Solid boundaries are modelled using the [Boundary Data Immersion Method](https://eprints.soton.ac.uk/369635/).
 The primary variables are the scalar pressure `p` (an array of dimension `D`)
 and the velocity vector field `u` (an array of dimension `D+1`).
+
+All arrays use the N+4 staggered layout: `N` interior cells plus 2 ghost/boundary
+cells per side, giving total size `M = N + 4` per spatial dimension.
 """
 struct Flow{D, T, Sf<:AbstractArray{T}, Vf<:AbstractArray{T}, Tf<:AbstractArray{T}}
     # Fluid fields
@@ -77,7 +87,7 @@ struct Flow{D, T, Sf<:AbstractArray{T}, Vf<:AbstractArray{T}, Tf<:AbstractArray{
     uBC :: Union{NTuple{D,Number},Function} # domain boundary values/function
     Δt:: Vector{T} # time step (stored in CPU memory)
     ν :: T # kinematic viscosity
-    g :: Union{Function,Nothing} # acceleration field funciton
+    g :: Union{Function,Nothing} # acceleration field function
     exitBC :: Bool # Convection exit
     perdir :: NTuple # tuple of periodic direction
     function Flow(N::NTuple{D}, uBC; f=Array, Δt=0.25, ν=0., g=nothing,
@@ -103,12 +113,26 @@ Current flow time.
 """
 time(a::Flow) = sum(@view(a.Δt[1:end-1]))
 
+"""
+    BDIM!(a::Flow)
+
+Apply the Boundary Data Immersion Method to enforce the body velocity.
+Uses the zeroth and first kernel moments (`μ₀`, `μ₁`) to blend the
+body velocity `V` with the predicted fluid velocity.
+"""
 function BDIM!(a::Flow)
     dt = a.Δt[end]
     @loop a.f[Ii] = a.u⁰[Ii]+dt*a.f[Ii]-a.V[Ii] over Ii in CartesianIndices(a.f)
     @loop a.u[Ii] += μddn(Ii,a.μ₁,a.f)+a.V[Ii]+a.μ₀[Ii]*a.f[Ii] over Ii ∈ inside_u(size(a.p))
 end
 
+"""
+    project!(a::Flow, b::AbstractPoisson, w=1)
+
+Project the velocity field onto a divergence-free field by solving
+the pressure Poisson equation `∇⋅(L∇p) = ∇⋅u/Δt` and correcting
+the velocity: `u -= L∇p`.  The weight `w` scales `Δt` (0.5 for the corrector).
+"""
 function project!(a::Flow{n},b::AbstractPoisson,w=1) where n
     dt = w*a.Δt[end]
     @inside b.z[I] = div(I,a.u); b.x .*= dt # set source term & solution IC
@@ -146,6 +170,12 @@ and the `AbstractPoisson` pressure solver to project the velocity onto an incomp
 end
 scale_u!(a,scale) = @loop a.u[Ii] *= scale over Ii ∈ inside_u(size(a.p))
 
+"""
+    CFL(a::Flow; Δt_max=10)
+
+Compute the CFL-limited time step from the maximum outward flux at each cell.
+Uses `global_min` for MPI-safe reduction across all ranks.
+"""
 function CFL(a::Flow;Δt_max=10)
     @inside a.σ[I] = flux_out(I,a.u)
     global_min(Δt_max,inv(maximum(a.σ)+5a.ν))

src/Metrics.jl

@@ -72,10 +72,12 @@ Compute ``∥𝛚∥`` at the center of cell `I`.
 """
 ω_mag(I::CartesianIndex{3},u) = norm2(ω(I,u))
 """
-    ω_θ(I::CartesianIndex{3},z,center,u)
+    ω_θ(I::CartesianIndex{3}, z, center, u; offset=zero(SVector{3,eltype(u)}))
 
 Compute ``𝛚⋅𝛉`` at the center of cell `I` where ``𝛉`` is the azimuth
-direction around vector `z` passing through `center`.
+direction around vector `z` passing through `center`.  The optional
+`offset` shifts the cell location — used in MPI parallel to map
+rank-local indices to global coordinates.
 """
 function ω_θ(I::CartesianIndex{3},z,center,u;offset=zero(SVector{3,eltype(u)}))
     θ = z × (loc(0,I,eltype(u))+offset-SVector{3}(center))
@@ -138,9 +140,11 @@ total_force(sim) = pressure_force(sim) .+ viscous_force(sim)
 
 using LinearAlgebra: cross
 """
-    pressure_moment(x₀,sim::Simulation)
+    pressure_moment(x₀, sim::Simulation; offset=zero(x₀))
 
-Computes the pressure moment on an immersed body relative to point x₀.
+Compute the pressure moment on an immersed body relative to point `x₀`.
+The optional `offset` shifts cell locations — used in MPI parallel to map
+rank-local indices to global coordinates.
 """
 pressure_moment(x₀,sim;kwargs...) = pressure_moment(x₀,sim.flow,sim.body;kwargs...)
 pressure_moment(x₀,flow,body;kwargs...) = pressure_moment(x₀,flow.p,flow.f,body,time(flow);kwargs...)

src/MultiLevelPoisson.jl

@@ -1,12 +1,34 @@
+"""
+    up(I, a=0)
+
+Return the fine-grid `CartesianIndices` range that maps to coarse cell `I`.
+When `a≠0`, the range is shifted by `-δ(a,I)` for staggered (face) restriction.
+"""
 @inline up(I::CartesianIndex,a=0) = (2I-3oneunit(I)):(2I-2oneunit(I)-δ(a,I))
+"""
+    down(I)
+
+Map fine-grid index `I` to the corresponding coarse-grid index.
+"""
 @inline down(I::CartesianIndex) = CI((I.I .- 1) .÷ 2 .+ 2)
+"""
+    restrict(I, b)
+
+Sum the fine-grid scalar values in `b` that map to coarse cell `I`.
+"""
 @fastmath @inline function restrict(I::CartesianIndex,b)
     s = zero(eltype(b))
     for J ∈ up(I)
      s += @inbounds(b[J])
     end
     return s
 end
+"""
+    restrictL(I, i, b)
+
+Restrict the Poisson conductivity `b` in dimension `i` from the fine grid
+to coarse cell `I`, averaging the two fine-grid face values.
+"""
 @fastmath @inline function restrictL(I::CartesianIndex,i,b)
     s = zero(eltype(b))
     for J ∈ up(I,i)
@@ -15,6 +37,12 @@ end
     return 0.5s
 end
 
+"""
+    restrictML(b::Poisson)
+
+Build a new coarse-level `Poisson` from fine-level `b` by restricting the
+conductivity `L` and allocating matching solution/residual arrays.
+"""
 function restrictML(b::Poisson)
     N,n = size_u(b.L)
     Na = map(i->i÷2+2,N)
@@ -23,6 +51,12 @@ function restrictML(b::Poisson)
     restrictL!(aL,b.L,perdir=b.perdir)
     Poisson(ax,aL,copy(ax);b.perdir)
 end
+"""
+    restrictL!(a, b; perdir=())
+
+Restrict the fine-grid conductivity `b` into coarse-grid `a`, then apply
+boundary conditions (`BC!` and `wallBC_L!`) on the coarse level.
+"""
 function restrictL!(a::AbstractArray{T,M},b;perdir=()) where {T,M}
     Na,n = size_u(a)
     for i ∈ 1:n
@@ -36,10 +70,10 @@ prolongate!(a,b) = @inside a[I] = b[down(I)]
 
 @inline divisible(l::Poisson) = all(size(l.x) .|> divisible)
 """
-    MultiLevelPoisson{N,M}
+    MultiLevelPoisson{T,S,V}
 
 Composite type used to solve the pressure Poisson equation with a [geometric multigrid](https://en.wikipedia.org/wiki/Multigrid_method) method.
-The only variable is `levels`, a vector of nested `Poisson` systems.
+The main field is `levels`, a vector of nested `Poisson` systems from fine to coarse.
 """
 struct MultiLevelPoisson{T,S<:AbstractArray{T},V<:AbstractArray{T}} <: AbstractPoisson{T,S,V}
     x::S
@@ -67,6 +101,13 @@ function update!(ml::MultiLevelPoisson)
     end
 end
 
+"""
+    Vcycle!(ml::MultiLevelPoisson; l=1, ω=1)
+
+Perform one multigrid V-cycle starting at level `l`: smooth on the fine grid,
+restrict the residual, solve (or recurse) on the coarse grid, then prolongate
+and correct the fine-grid solution.
+"""
 function Vcycle!(ml::MultiLevelPoisson;l=1,ω=1)
     fine,coarse = ml.levels[l],ml.levels[l+1]
     # set up coarse level
@@ -86,6 +127,13 @@ residual!(ml::MultiLevelPoisson,x) = residual!(ml.levels[1],x)
 
 smooth! = GaussSeidelRB!
 
+"""
+    solver!(ml::MultiLevelPoisson; tol=1e-4, itmx=32)
+
+Multigrid solver: iterates V-cycles with adaptive relaxation `ω` until the
+`L₂`-norm residual drops below `tol`.  Ends with `pin_pressure!` + `comm!`
+to remove the null-space mode and synchronize halos.
+"""
 function solver!(ml::MultiLevelPoisson{T};tol=1e-4,itmx=32) where T
     p = ml.levels[1]
     residual!(p); r₂ = L₂(p); ω = T(1)

src/Poisson.jl

@@ -1,7 +1,7 @@
 abstract type AbstractPoisson{T,S,V} end
 
 """
-    Poisson{N,M}
+    Poisson{T, S, V}
 
 Composite type for conservative variable coefficient Poisson equations:
 
@@ -13,11 +13,18 @@ The resulting linear system is
 
 where A is symmetric, block-tridiagonal and extremely sparse. Moreover,
 `D[I]=-∑ᵢ(L[I,i]+L'[I,i])`. This means matrix storage, multiplication,
-ect can be easily implemented and optimized without external libraries.
-
-To help iteratively solve the system above, the Poisson structure holds
-helper arrays for `inv(D)`, the error `ϵ`, and residual `r=z-Ax`. An iterative
-solution method then estimates the error `ϵ=̃A⁻¹r` and increments `x+=ϵ`, `r-=Aϵ`.
+etc. can be easily implemented and optimized without external libraries.
+
+The conductivity `L` is initialized from `flow.μ₀` and modified by
+`wallBC_L!` (zero at physical walls for Neumann BCs). Ghost-cell
+synchronization uses `comm!` (= `perBC!` + `scalar_halo!`) rather than
+bare `perBC!`, so that MPI rank-internal boundaries are handled correctly.
+
+To help iteratively solve the system, the structure holds helper arrays
+for `inv(D)`, the error `ϵ`, and residual `r=z-Ax`. An iterative solution
+method estimates the error `ϵ≈A⁻¹r` and increments `x+=ϵ`, `r-=Aϵ`.
+The solver ends with `pin_pressure!` + `comm!` to remove the null-space
+mode and synchronize halos.
 """
 struct Poisson{T,S<:AbstractArray{T},V<:AbstractArray{T}} <: AbstractPoisson{T,S,V}
     L :: V # Lower diagonal coefficients
@@ -58,7 +65,7 @@ end
     mult!(p::Poisson,x)
 
 Efficient function for Poisson matrix-vector multiplication.
-Fills `p.z = p.A x` with 0 in the ghost cells.
+Fills `p.z = Ax` with 0 in the ghost cells, where `A` is the Poisson matrix implied by `L` and `D`.
 """
 function mult!(p::Poisson,x)
     @assert axes(p.z)==axes(x)
@@ -135,7 +142,7 @@ end
     GaussSeidelRB!(p::Poisson;it=4, ω=1)
 
 Red-black Gauss-Seidel smoother. Runs `it` iterations; a complete red-black cycle requires `it` to be even.
-`ω` under-/over-relaxs the solution through scaling the deferred corrections in `increment!`.
+`ω` under-/over-relaxes the solution through scaling the deferred corrections in `increment!`.
 Note: This performs best on GPU configurations and is the default smoother.
 """
 function GaussSeidelRB!(p::Poisson{T};it=4, ω=1) where {T}
@@ -150,7 +157,7 @@ end
 """
     pcg!(p::Poisson; it=6)
 
-Conjugate-Gradient smoother with Jacobi predictioning. Runs at most `it` iterations,
+Conjugate-Gradient smoother with Jacobi preconditioning. Runs at most `it` iterations,
 but will exit early if the Gram-Schmidt update parameter `|α| < 1%` or `|r D⁻¹ r| < 1e-8`.
 Note: This runs for general backends.
 """
@@ -180,16 +187,19 @@ L₂(p::Poisson) = global_dot(p.r, p.r) # special method since outside(p.r)≡0
 L∞(p::Poisson) = maximum(abs,p.r)
 
 """
-    solver!(A::Poisson;tol=1e-4,itmx=1e3)
+    solver!(A::Poisson; tol=1e-4, itmx=1e3)
 
-Approximate iterative solver for the Poisson matrix equation `Ax=b`.
+Iterative solver for the Poisson matrix equation `Ax=b` using
+preconditioned conjugate gradients (`pcg!`).
 
-  - `A`: Poisson matrix with working arrays.
-  - `A.x`: Solution vector. Can start with an initial guess.
-  - `A.z`: Right-Hand-Side vector. Will be overwritten!
-  - `A.n[end]`: stores the number of iterations performed.
+  - `A.x`: Solution vector (can start with an initial guess).
+  - `A.z`: Right-hand-side vector (overwritten).
+  - `A.n[end]`: Number of iterations performed.
   - `tol`: Convergence tolerance on the `L₂`-norm residual.
   - `itmx`: Maximum number of iterations.
+
+Ends with `pin_pressure!` (remove null-space mean) and `comm!`
+(halo sync) so the solution is ready for use in `project!`.
 """
 function solver!(p::Poisson;tol=1e-4,itmx=1e3)
     residual!(p); r₂ = L₂(p)

src/WaterLily.jl

@@ -60,7 +60,7 @@ export RigidMap,setmap
 
 """
     Simulation(dims::NTuple, uBC::Union{NTuple,Function}, L::Number;
-               U=norm2(Uλ), Δt=0.25, ν=0., ϵ=1, g=nothing,
+               U=√sum(abs2,uBC), Δt=0.25, ν=0., ϵ=1, g=nothing,
                perdir=(), exitBC=false,
                body::AbstractBody=NoBody(),
                T=Float32, mem=Array)