Domains & meshes — Scimba Jax

1 · No mesh

Mesh-less domains

These domains are sampled on the fly, at every optimization step. A domain is defined by a level set and a mapping, and a solver only uses the points drawn from it by a sampler.

Every VolumetricDomain is defined by a level set (a signed distance function, tested as sdf(x) < 0) and reshaped by a mapping. The bottom panel of the figure shows a main domain built this way, with a subdomain and a hole, sampled by a DomainSampler (interior points, boundary points and normals).

A level set (signed distance contours) reshaped by a mapping into a curved domain, and below it a composite meshless domain made of an outer square, an inner subdomain and a hole, sampled by random interior and boundary points with normals.

1. A volumetric domain, with a subdomain and a hole

A VolumetricDomain that is a main domain can be extended with subdomains, holes and a mapping. Each piece carries its own label, which shows up as a key in the dictionaries returned by samplers.

Square2D and Disk2D come from domains/meshless_domains/; arbitrary-dimension equivalents (HypercubeND, BallND) are in domains_nd.py. Each of them is defined by a level set and a mapping, as explained below.

from scimba_jax.domains.meshless_domains.domains_2d import (
    Square2D,
    Disk2D,
)
bounds_outer = [(-1.0, 1.0), (-1.0, 1.0)]
bounds_inner = [(-0.7, 0.3), (-0.25, 0.25)]
outer = Square2D(
    bounds_outer, is_main_domain=True, label_str="outer"
)
inner = Square2D(
    bounds_inner, is_main_domain=False, label_str="inner"
)
hole = Disk2D(
    (0.55, 0.0), 0.2, is_main_domain=False, label_str="hole"
)
outer.add_subdomain(inner)
outer.add_hole(hole)

2. Defining a domain from a level set and a mapping

A VolumetricDomain is fundamentally a SignedDistance (the level set: < 0 inside, > 0 outside) plus bounds. Square2D, Disk2D and similar classes are thin wrappers around this pattern. The same domain can be written directly, and then reshaped with a Mapping.

import jax.numpy as jnp
from scimba_jax.domains.sdf import SignedDistance
from scimba_jax.domains.meshless_domains.base import VolumetricDomain
from scimba_jax.domains.domain_mapping import DomainMapping as Mapping
class Disk2DSignedDistance(SignedDistance):
    def __init__(self, center, radius):
        super().__init__(dim=2, threshold=0.0)
        self.center = center
        self.radius = radius
    def _sdf_pointwise(self, x):
        return (jnp.linalg.norm(x - self.center, axis=0) - self.radius)[None]
disk = VolumetricDomain(
    domain_type="Disk2D",
    dim=2,
    sdf=Disk2DSignedDistance(jnp.array([0.0, 0.0]), 1.0),
    bounds=[(-1.0, 1.0), (-1.0, 1.0)],
    is_main_domain=True,
)
disk.set_mapping(
    Mapping.rot_2d(angle=jnp.pi / 6.0, center=jnp.array([0.0, 0.0])),
    [(-1.5, 1.5), (-1.5, 1.5)],
)

3. Mappings of subdomains and holes

Calling set_mapping on a composite main domain (like outer from step 1) reshapes it, and automatically applies the same mapping to its subdomains and holes. Built-in mappings cover common cases such as rotations. Custom mappings are JAX-traceable functions, so vmap and automatic differentiation (e.g. for normals) can be used directly.

import jax.numpy as jnp
from scimba_jax.domains.domain_mapping import (
    DomainMapping as Mapping,
)
outer.set_mapping(
    Mapping.rot_2d(
        angle=jnp.pi / 3.0, center=jnp.array([0.0, 0.0])
    ),
    [(-2.0, 2.0), (-2.0, 2.0)],
)

4. Boundary domains as parameterized curves

A boundary piece is a SurfacicDomain: a 1D parametric_domain (an interval) mapped into physical space by a surface (a DomainMapping). Segment2D, Circle2D and ArcCircle2D are built this way. For instance, ArcCircle2D maps an angle interval through Mapping.circle.

The boundary of hole can be split into two named arcs, registered with add_bc_domain. Each arc can then be treated independently, e.g. with different boundary conditions on the north and south arcs.

The figure: a boundary curve is a parametric domain (here an angle interval) mapped through a surface mapping. Splitting [0, 2π] into [0, π] and [π, 2π] and mapping each half with Mapping.circle gives two independent boundary pieces. This is the construction used by ArcCircle2D above, and it allows different boundary conditions on the north and south parts of hole.

import jax.numpy as jnp
from scimba_jax.domains.meshless_domains.domains_2d import (
    ArcCircle2D,
)
# split the circular boundary of the hole into two named arcs
disk_up = ArcCircle2D(
    (0.55, 0.0), 0.2, (0.0, jnp.pi), label_str="bc hole north"
)
disk_down = ArcCircle2D(
    (0.55, 0.0),
    0.2,
    (jnp.pi, 2 * jnp.pi),
    label_str="bc hole south",
)
hole.add_bc_domain(disk_up)
hole.add_bc_domain(disk_down)

A parametric angle interval [0, 2π], split into [0, π] and [π, 2π], mapped via Mapping.circle onto a physical circle whose top half is the 'bc hole north' arc and bottom half is the 'bc hole south' arc, with arrows showing the parameterization direction.

2 · Built once

Meshes

There are four kinds of meshes, and the choice depends on the geometry: a single mapped grid, an arbitrary boundary imported through Gmsh, a few smooth patches, or the product of two existing meshes.

1. A structured mesh for the DG/FEM solver

Classical solvers (linear_approximation/dg/) use a Mesh, which is a cartesian grid of n_cells reference cells, each with its own Gauss quadrature (ref_quad), mapped to physical space through the same Mapping abstraction as the mesh-less domains above.

Adapted from the diffusion-advection DG example: view source ↗.

A structured reference cell mapped through a Mapping object onto a curved physical mesh, with Gauss quadrature points.

One reference cell, [0, 1]², mapped onto a physical Mesh by the same Mapping abstraction (here onto a curved sector). Every cell has its own Gauss quadrature points.

import jax.numpy as jnp
from scimba_jax.linear_approximation.meshes.mesh import Mesh
from scimba_jax.linear_approximation.quad.gauss_quad import (
    UnitSquareTensorized,
)
from scimba_jax.mapping.mapping import (
    InvertibleFunction,
    Mapping,
)
physical_dim = 2
quad_order = 3
n_cells = 40
# a quarter annulus, theta in [0, pi / 2]
r_min, r_max = 0.25, 0.8
# [0, 1]^2 -> (r, theta) -> (x, y), and its inverse
def to_sector(x):
    r = r_min + x[0] * (r_max - r_min)
    theta = x[1] * jnp.pi / 2.0
    return jnp.array([r * jnp.cos(theta), r * jnp.sin(theta)])
def from_sector(y):
    r = jnp.sqrt(y[0] ** 2 + y[1] ** 2)
    theta = jnp.arctan2(y[1], y[0])
    return jnp.array(
        [(r - r_min) / (r_max - r_min), theta / (jnp.pi / 2.0)]
    )
mesh = Mesh(
    dim=physical_dim,
    n_cells=(n_cells, n_cells),
    ref_quad=UnitSquareTensorized(
        dim=physical_dim, order=quad_order
    ),
    mapping=Mapping(
        mappings=[InvertibleFunction(to_sector, from_sector)]
    ),
    # curved cells: faces use the mapping
    is_identity_mapping=False,
)

2. Unstructured quads, imported through Gmsh

An unstructured mesh is useful when a curved or arbitrary boundary is not naturally described by a single mapping. Gmsh returns a MacroMesh with complete tensor-product quadrilaterals; UnstructuredMesh.from_macro_mesh turns it into the mesh consumed by a solver.

Gmsh is an optional dependency, installed with scimba[mesh]. The requested geometry order is kept, including the nodes on curved cell edges.

A mesh of quadrilateral cells following a smooth curved boundary, with one cell K highlighted and its four corner nodes marked, and a box listing what a mesh file holds: nodes (physical coordinates), cells (each quad's 4 nodes) and faces (the edges shared by two neighbouring cells).

A mesh is made of points (nodes), the quads they form (cells), and the edges shared by neighbouring quads (faces). Cells are smaller near a tight curve and larger where the boundary is flat. A curved cell is described in the same way, with moved corner nodes.

import numpy as np
from scimba_jax.linear_approximation.meshes.unstructured_mesh import (
    UnstructuredMesh,
)
from scimba_jax.linear_approximation.quad.gauss_quad import (
    UnitSquareTensorized,
)
from scimba_jax.mapping.macro_mesh import macro_mesh_from_points
# Requires: pip install scimba[mesh]
# Gmsh closes this ordered boundary and fills it with curved
# Q_p cells.
boundary = np.array(
    [
        [-1.0, -0.7],
        [0.8, -0.9],
        [1.2, 0.1],
        [0.2, 1.0],
        [-0.9, 0.6],
    ]
)
macro = macro_mesh_from_points(
    boundary, order=3, mesh_size=0.18, smooth=True
)
mesh = UnstructuredMesh.from_macro_mesh(
    macro,
    UnitSquareTensorized(dim=2, order=5),
)

3. Block-structured meshes for patchwise geometry

This representation is useful when a domain is made of a small number of smooth patches, for example an O-grid around a disk or a hole. Each patch has its own mapped Cartesian grid, and the container stores the interface pairings. Patches may use different resolutions when an interface is non-conforming.

For DG, an interface uses the same numerical flux as an interior face. For continuous FEM, an explicit Nitsche/SIPG interface flux has to be given, since no default penalty is provided.

Five patches around a curved boundary: a square grid in the middle and four curved grids around it, meeting exactly at their shared edges, with a box listing three properties: one grid per patch, edges match exactly, and curved patches are allowed.

Each patch is a regular grid. Here, five patches with matching edges form a single mesh, whose outer boundary is round while each patch stays simple.

from scimba_jax.linear_approximation.meshes.block_structured_mesh import (
    BlockStructuredMesh,
)
from scimba_jax.linear_approximation.quad.gauss_quad import (
    UnitSquareTensorized,
)
from scimba_jax.mapping.macro_mesh import macro_mesh_ogrid_disk
# Five macro-patches: a square centre plus four curved outer
# patches.
macro = macro_mesh_ogrid_disk(
    radius=1.0, inner=0.45, n=1, order=3
)
quad = UnitSquareTensorized(dim=2, order=5)
# One structured 12 x 12 mesh in every patch.
mesh = BlockStructuredMesh(
    macro, n_cells=(12, 12), ref_quad=quad
)
# Per-patch counts are also valid, e.g. [(8, 8), (16, 16), ...].

4. Tensor meshes: product of two existing meshes

tensor_mesh(mesh_a, mesh_b) does not introduce a new mesh class. A cell of the product is cell_a × cell_b, its quadrature is quad_a ⊗ quad_b, and the result is a Mesh, BlockStructuredMesh or UnstructuredMesh, depending on which of the two factors is the least structured. All the existing features of that class (faces, fluxes, fast Cartesian paths) can then be used.

A space-time slab, a phase-space mesh or a toroidal tokamak cross-section are built in this way, as the product of two meshes. The argument names lets the physics be written as f(x, t) instead of slices of a flat array, and periodic_directions makes an axis of the product periodic (for instance, the toroidal angle).

mesh_a, a 1-D structured mesh in x, and mesh_b, a 1-D structured mesh in t, combined by tensor_mesh into mesh_a times mesh_b: a 2-D Mesh named (x, t), with one cell K equal to cell_a times cell_b highlighted. A box lists three properties: kind (least structured of the two), periodic (glue an axis, e.g. tokamak angle) and names (so physics reads f of x and t).

The product of two 1D meshes gives a 2D mesh. Here, the product is a Mesh since both factors are. The basis and the physics are the same as on each factor, and the tensor structure is only contained in the mesh and its connectivity.

from scimba_jax.linear_approximation.meshes.mesh import Mesh
from scimba_jax.linear_approximation.meshes.tensor_mesh import (
    tensor_mesh,
)
from scimba_jax.linear_approximation.quad.gauss_quad import (
    UnitSquareTensorized,
)
from scimba_jax.mapping.mapping import (
    InvertibleFunction,
    Mapping,
)
def _line(n_cells, length, order=3):
    physical = InvertibleFunction(
        lambda x: length * x, lambda y: y / length
    )
    return Mesh(
        dim=1,
        n_cells=(n_cells,),
        ref_quad=UnitSquareTensorized(dim=1, order=order),
        mapping=Mapping(mappings=[physical]),
    )
# 1D mesh in x times 1D mesh in t: a 2D space-time slab,
# without time stepping.
# The product is a `Mesh` (the least structured factor),
# usable by DG/FEM schemes.
mesh_x, mesh_t = _line(32, length=1.0), _line(12, length=0.1)
slab = tensor_mesh(mesh_x, mesh_t, names=("x", "t"))