Add support for more grid operators (#81)

This PR adds support for second-order derivatives, including Laplacian and general `divg_grad` operators.

---------

Co-authored-by: Ivan Utkin <iutkin@ethz.ch>

Author: luraess

Project.toml

@@ -1,7 +1,7 @@
 name = "Chmy"
 uuid = "33a72cf0-4690-46d7-b987-06506c2248b9"
 authors = ["Ivan Utkin <iutkin@ethz.ch>, Ludovic Raess <ludovic.rass@gmail.com>, and contributors"]
-version = "0.1.23"
+version = "0.1.24"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

docs/src/concepts/grid_operators.md

@@ -7,32 +7,42 @@ Chmy.jl currently supports various finite difference operators for fields define
 | $\frac{\partial}{\partial x} P$ | ` ∂x(P, g, I)` |
 | $\frac{\partial}{\partial y} P$ | ` ∂y(P, g, I)` |
 | $\frac{\partial}{\partial z} P$ | ` ∂z(P, g, I)` |
-| $\nabla P$ | ` divg(P, g, I)` |
+| $\frac{\partial^2}{\partial x^2} P$ | ` ∂²x(P, g, I)` |
+| $\frac{\partial^2}{\partial y^2} P$ | ` ∂²y(P, g, I)` |
+| $\frac{\partial^2}{\partial z^2} P$ | ` ∂²z(P, g, I)` |
+| $\nabla P$ | `divg(P, g, I)` |
+| $\Delta P$ | `lapl(P, g, I)` |
+| $\nabla \cdot χ \nabla P$ | `divg_grad(P, χ, g, I)` |
 
 ## Computing the Divergence of a Vector Field
 
-To illustrate the usage of grid operators, we compute the divergence of an vector field $V$ using the `divg` function. We first allocate memory for required fields.
+To illustrate the usage of grid operators, we compute the divergence of a vector field $V$, the Laplacian of a scalar field $P$ and the divergence-gradient of a scalar field $C$ weighted by the coefficient $χ$ using the `divg`, `lapl` and `divg_grad` functions, respectively. We first allocate memory for required fields.
 
 ```julia
-V  = VectorField(backend, grid)
+P  = Field(backend, grid, Center())
+C  = Field(backend, grid, Center())
+χ  = Field(backend, grid, Center()) # works also for Vertex()
 ∇V = Field(backend, grid, Center())
+V  = VectorField(backend, grid)
 # use set! to set up the initial vector field...
 ```
 
 The kernel that computes the divergence needs to have the grid information passed as for other finite difference operators.
 
 ```julia
-@kernel inbounds = true function update_∇!(V, ∇V, g::StructuredGrid, O)
+@kernel inbounds = true function update!(∇V, V, P, C, χ, g::StructuredGrid, O)
     I = @index(Global, Cartesian)
     I = I + O
     ∇V[I] = divg(V, g, I)
+    P[I]  = lapl(P, g, I)
+    C[I]  = divg_grad(C, χ, g, I)
 end
 ```
 
 The kernel can then be launched when required as we detailed in section [Kernels](./kernels.md).
 
 ```julia
-launch(arch, grid, update_∇! => (V, ∇V, grid))
+launch(arch, grid, update! => (∇V, V, P, C, χ, grid))
 ```
 
 ## Masking
@@ -125,4 +135,3 @@ and can be launched with some launcher defined using `launch = Launcher(arch, gr
 ```julia
 launch(arch, grid, interpolate_ρ! => (ρ, ρx, ρy, grid))
 ```
-

src/Chmy.jl

@@ -51,19 +51,19 @@ export
     xcenters, ycenters, zcenters,
 
     # GridOperators
-    left, right, δ, ∂,
+    left, right, δ, ∂, ∂², ∂k∂,
     InterpolationRule, Linear, HarmonicLinear,
     itp, lerp, hlerp,
-    divg, vmag,
+    divg, divg_grad, lapl, vmag,
     AbstractMask, FieldMask, FieldMask1D, FieldMask2D, FieldMask3D, at,
 
-    leftx, rightx, δx, ∂x,
-    lefty, righty, δy, ∂y,
-    leftz, rightz, δz, ∂z,
+    leftx, rightx, δx, ∂x, ∂²x, ∂k∂x,
+    lefty, righty, δy, ∂y, ∂²y, ∂k∂y,
+    leftz, rightz, δz, ∂z, ∂²z, ∂k∂z,
 
-    leftx_masked, rightx_masked, δx_masked, ∂x_masked,
-    lefty_masked, righty_masked, δy_masked, ∂y_masked,
-    leftz_masked, rightz_masked, δz_masked, ∂z_masked,
+    leftx_masked, rightx_masked, δx_masked, ∂x_masked, ∂²x_masked, ∂k∂x_masked,
+    lefty_masked, righty_masked, δy_masked, ∂y_masked, ∂²y_masked, ∂k∂y_masked,
+    leftz_masked, rightz_masked, δz_masked, ∂z_masked, ∂²z_masked, ∂k∂z_masked,
 
     # KernelLaunch
     Launcher,

src/GridOperators/GridOperators.jl

@@ -1,11 +1,11 @@
 module GridOperators
 
-export left, right, δ, ∂
+export left, right, δ, ∂, ∂², ∂k∂
 
 export InterpolationRule, Linear, HarmonicLinear
 export itp, lerp, hlerp
 
-export divg, vmag
+export divg, divg_grad, lapl, vmag
 
 export AbstractMask, FieldMask, FieldMask1D, FieldMask2D, FieldMask3D, at
 

src/GridOperators/cartesian_field_operators.jl

@@ -3,9 +3,11 @@ for (dim, coord) in enumerate((:x, :y, :z))
     _r = Symbol(:right, coord)
     _δ = Symbol(:δ, coord)
     _∂ = Symbol(:∂, coord)
+    _∂² = Symbol(:∂², coord)
+    _∂k∂ = Symbol(:∂k∂, coord)
 
     @eval begin
-        export $_δ, $_∂, $_l, $_r
+        export $_δ, $_∂, $_∂², $_∂k∂, $_l, $_r
 
         """
             $($_l)(f, I)
@@ -42,5 +44,23 @@ for (dim, coord) in enumerate((:x, :y, :z))
         @add_cartesian function $_∂(f::AbstractField, grid, I::Vararg{Integer,N}) where {N}
             ∂(f, grid, Dim($dim), I...)
         end
+
+        """
+            $($_∂²)(f, grid, I)
+
+        Directional partial second derivative in $($(string(coord))) direction.
+        """
+        @add_cartesian function $_∂²(f::AbstractField, grid, I::Vararg{Integer,N}) where {N}
+            ∂²(f, grid, Dim($dim), I...)
+        end
+
+        """
+            $($_∂k∂)(f, k, grid, I)
+
+        Directional divergence of gradient times coefficient `k` in $($(string(coord))) direction.
+        """
+        @add_cartesian function $_∂k∂(f::AbstractField, k::AbstractField, grid, I::Vararg{Integer,N}) where {N}
+            ∂k∂(f, k, grid, Dim($dim), I...)
+        end
     end
 end

src/GridOperators/field_operators.jl

@@ -2,13 +2,13 @@
 @add_cartesian function left(f::AbstractField, dim, I::Vararg{Integer,N}) where {N}
     loc  = location(f)
     from = flipped(loc, dim)
-    left(f, loc, from, dim, I...)
+    return left(f, loc, from, dim, I...)
 end
 
 @add_cartesian function right(f::AbstractField, dim, I::Vararg{Integer,N}) where {N}
     loc  = location(f)
     from = flipped(loc, dim)
-    right(f, loc, from, dim, I...)
+    return right(f, loc, from, dim, I...)
 end
 
 @add_cartesian function δ(f::AbstractField, dim, I::Vararg{Integer,N}) where {N}
@@ -20,10 +20,33 @@ end
 @add_cartesian function ∂(f::AbstractField, grid, dim, I::Vararg{Integer,N}) where {N}
     loc  = location(f)
     from = flipped(loc, dim)
-    ∂(f, loc, from, grid, dim, I...)
+    return ∂(f, loc, from, grid, dim, I...)
+end
+
+@add_cartesian function ∂²(f::AbstractField, grid, dim, I::Vararg{Integer,N}) where {N}
+    loc  = location(f)
+    from = location(f)
+    return ∂²(f, loc, from, grid, dim, I...)
+end
+
+@add_cartesian function ∂k∂(f::AbstractField, k::AbstractField, grid, dim, I::Vararg{Integer,N}) where {N}
+    loc  = location(f)
+    from = location(f)
+    return ∂k∂(f, k, loc, from, grid, dim, I...)
 end
 
 # covariant derivatives
+"""
+    divg(V, grid, I...)
+
+Compute the divergence of a vector field `V` on a structured grid `grid`.
+This operation is performed along all dimensions of the grid.
+
+## Arguments
+- `V`: The vector field represented as a named tuple of fields.
+- `grid`: The structured grid on which the operation is performed.
+- `I...`: The indices specifying the location on the grid
+"""
 @propagate_inbounds @generated function divg(V::NamedTuple{names,<:NTuple{N,AbstractField}}, grid::StructuredGrid{N}, I::Vararg{Integer,N}) where {names,N}
     quote
         @inline
@@ -35,6 +58,61 @@ end
     return divg(V, grid, Tuple(I)...)
 end
 
+"""
+    lapl(F, grid, I...)
+
+Compute the Laplacian of a field `F` on a structured grid `grid`.
+This operation is performed along all dimensions of the grid.
+
+## Arguments
+- `F`: The field whose Laplacian is to be computed.
+- `grid`: The structured grid on which the operation is performed.
+- `I...`: The indices specifying the location on the grid
+"""
+@propagate_inbounds @generated function lapl(F::AbstractField, grid::StructuredGrid{N}, I::Vararg{Integer,N}) where {N}
+    quote
+        @inline
+        Base.Cartesian.@ncall $N (+) D -> ∂²(F, grid, Dim(D), I...)
+    end
+end
+
+@propagate_inbounds function lapl(F::AbstractField, grid::StructuredGrid{N}, I::CartesianIndex{N}) where {N}
+    return lapl(F, grid, Tuple(I)...)
+end
+
+"""
+    divg_grad(F, K, grid, I...)
+
+Compute the divergence of the gradient of a field `F` weighted by a coefficient `K` on a structured grid `grid`.
+This operation is performed along all dimensions of the grid.
+
+## Arguments
+- `F`: The field whose gradient is to be computed.
+- `K`: The weighting field for the gradient.
+- `grid`: The structured grid on which the operation is performed.
+- `I`: The indices specifying the location on the grid.
+"""
+@propagate_inbounds @generated function divg_grad(F::AbstractField, K::AbstractField, grid::StructuredGrid{N}, I::Vararg{Integer,N}) where {N}
+    quote
+        @inline
+        Base.Cartesian.@ncall $N (+) D -> ∂k∂(F, K, grid, Dim(D), I...)
+    end
+end
+
+@propagate_inbounds function divg_grad(F::AbstractField, K::AbstractField, grid::StructuredGrid{N}, I::CartesianIndex{N}) where {N}
+    return divg_grad(F, K, grid, Tuple(I)...)
+end
+
+"""
+    vmag(V, grid, I...)
+
+Compute the magnitude of a vector field `V` at a given grid location `I` in a structured grid `grid`.
+
+## Arguments
+- `V`: A named tuple representing the vector field, where each component is an `AbstractField`.
+- `grid`: The structured grid on which the vector field is defined.
+- `I...`: The indices specifying the location on the grid.
+"""
 @propagate_inbounds @generated function vmag(V::NamedTuple{names,<:NTuple{N,AbstractField}}, grid::StructuredGrid{N}, I::Vararg{Integer,N}) where {names,N}
     quote
         @inline

src/GridOperators/masked_cartesian_field_operators.jl

@@ -3,9 +3,11 @@ for (dim, coord) in enumerate((:x, :y, :z))
     _r = Symbol(:right, coord)
     _δ = Symbol(:δ, coord)
     _∂ = Symbol(:∂, coord)
+    _∂² = Symbol(:∂², coord)
+    _∂k∂ = Symbol(:∂k∂, coord)
 
     @eval begin
-        export $_δ, $_∂, $_l, $_r
+        export $_δ, $_∂, $_∂², $_∂k∂, $_l, $_r
 
         """
             $($_l)(f, ω, I)
@@ -42,5 +44,23 @@ for (dim, coord) in enumerate((:x, :y, :z))
         @add_cartesian function $_∂(f::AbstractField, ω::AbstractMask, grid, I::Vararg{Integer,N}) where {N}
             ∂(f, ω, grid, Dim($dim), I...)
         end
+
+        """
+        $($_∂²)(f, ω, grid, I)
+
+        Directional partial second derivative in $($(string(coord))) direction, masked with `ω`.
+        """
+        @add_cartesian function $_∂²(f::AbstractField, ω::AbstractMask, grid, I::Vararg{Integer,N}) where {N}
+            ∂²(f, ω, grid, Dim($dim), I...)
+        end
+
+        """
+        $($_∂k∂)(f, k, ω, grid, I)
+
+        Directional divergence of gradient times coefficient `k` in $($(string(coord))) direction, masked with `ω`.
+        """
+        @add_cartesian function $_∂k∂(f::AbstractField, k::AbstractField, ω::AbstractMask, grid, I::Vararg{Integer,N}) where {N}
+            ∂k∂(f, k, ω, grid, Dim($dim), I...)
+        end
     end
 end

src/GridOperators/masked_field_operators.jl

@@ -14,7 +14,7 @@ end
 @add_cartesian function δ(f::AbstractField, ω::AbstractMask, dim, I::Vararg{Integer,N}) where {N}
     loc  = location(f)
     from = flipped(loc, dim)
-    δ(f, loc, from, ω, dim, I...)
+    return δ(f, loc, from, ω, dim, I...)
 end
 
 @add_cartesian function ∂(f::AbstractField, ω::AbstractMask, grid, dim, I::Vararg{Integer,N}) where {N}
@@ -23,7 +23,31 @@ end
     return ∂(f, loc, from, ω, grid, dim, I...)
 end
 
+@add_cartesian function ∂²(f::AbstractField, ω::AbstractMask, grid, dim, I::Vararg{Integer,N}) where {N}
+    loc  = location(f)
+    from = location(f)
+    return ∂²(f, loc, from, ω, grid, dim, I...)
+end
+
+@add_cartesian function ∂k∂(f::AbstractField, k::AbstractField, ω::AbstractMask, grid, dim, I::Vararg{Integer,N}) where {N}
+    loc  = location(f)
+    from = location(f)
+    return ∂k∂(f, k, loc, from, ω, grid, dim, I...)
+end
+
 # covariant derivatives
+"""
+    divg(V::NamedTuple{names,<:NTuple{N,AbstractField}}, ω::AbstractMask{T,N}, grid::StructuredGrid{N}, I::Vararg{Integer,N}) where {names,T,N}
+
+Compute the divergence of a vector field `V` on a structured grid `grid`, using the mask `ω` to handle masked regions.
+This operation is performed along all dimensions of the grid.
+
+## Arguments:
+- `V::NamedTuple{names,<:NTuple{N,AbstractField}}`: The vector field represented as a named tuple of fields.
+- `ω::AbstractMask`: The mask for the grid.
+- `grid::StructuredGrid{N}`: The structured grid on which the operation is performed.
+- `I`: The indices specifying the location on the grid (Tuple or Cartesian indices).
+"""
 @propagate_inbounds @generated function divg(V::NamedTuple{names,<:NTuple{N,AbstractField}},
                                              ω::AbstractMask{T,N},
                                              grid::StructuredGrid{N},
@@ -40,3 +64,64 @@ end
                                   I::CartesianIndex{N}) where {names,T,N}
     return divg(V, ω, grid, Tuple(I)...)
 end
+
+"""
+    lapl(F::AbstractField, ω::AbstractMask{T,N}, grid::StructuredGrid{N}, I::Vararg{Integer,N}) where {T,N}
+
+Compute the Laplacian of a field `F` on a structured grid `grid`.
+This operation is performed along all dimensions of the grid.
+
+## Arguments
+- `F::AbstractField`: The field whose Laplacian is to be computed.
+- `ω::AbstractMask`: The mask for the grid.
+- `grid::StructuredGrid{N}`: The structured grid on which the operation is performed.
+- `I`: The indices specifying the location on the grid (Tuple or Cartesian indices).
+"""
+@propagate_inbounds @generated function lapl(F::AbstractField,
+                                             ω::AbstractMask{T,N},
+                                             grid::StructuredGrid{N},
+                                             I::Vararg{Integer,N}) where {T,N}
+    quote
+        @inline
+        Base.Cartesian.@ncall $N (+) D -> ∂²(F, ω, grid, Dim(D), I...)
+    end
+end
+
+@propagate_inbounds function lapl(F::AbstractField,
+                                  ω::AbstractMask{T,N},
+                                  grid::StructuredGrid{N},
+                                  I::CartesianIndex{N}) where {T,N}
+    return lapl(F, ω, grid, Tuple(I)...)
+end
+
+"""
+    divg_grad(F::AbstractField, K::AbstractField, ω::AbstractMask{T,N}, grid::StructuredGrid{N}, I::Vararg{Integer,N}) where {T,N}
+
+Compute the divergence of the diffusion flux of a field `F` weighted by a coefficient `K` on a structured grid `grid`.
+This operation is performed along all dimensions of the grid.
+
+## Arguments
+- `F::AbstractField`: The field whose gradient is to be computed.
+- `K::AbstractField`: The weighting field for the gradient.
+- `ω::AbstractMask`: The mask for the grid.
+- `grid::StructuredGrid{N}`: The structured grid on which the operation is performed.
+- `I`: The indices specifying the location on the grid (Tuple or Cartesian indices).
+"""
+@propagate_inbounds @generated function divg_grad(F::AbstractField,
+                                                  K::AbstractField,
+                                                  ω::AbstractMask{T,N},
+                                                  grid::StructuredGrid{N},
+                                                  I::Vararg{Integer,N}) where {T,N}
+    quote
+        @inline
+        Base.Cartesian.@ncall $N (+) D -> ∂k∂(F, K, ω, grid, Dim(D), I...)
+    end
+end
+
+@propagate_inbounds function divg_grad(F::AbstractField,
+                                       K::AbstractField,
+                                       ω::AbstractMask{T,N},
+                                       grid::StructuredGrid{N},
+                                       I::CartesianIndex{N}) where {T,N}
+    return divg_grad(F, K, ω, grid, Tuple(I)...)
+end