Domains¶

A domain is what spec.domain(name, domain) registers, and its discretization decides how the operator bank evaluates each operator. See domains and discretizations.

Recommended defaults

Structured problems use the grid factories with the default method="fvm". For finite elements, build a TaggedMesh so boundary conditions can select boundaries by name. The mesh factories return one with with_boundary_tags=True.

Grid factories¶

diff_soup.domains.cartesian_grid_1d ¶

cartesian_grid_1d(n: int | None = None, x_min: float | None = None, x_max: float | None = None, *, faces: ArrayLike | None = None, half_cell_left: bool = False, half_cell_right: bool = False, method: str = 'fvm') -> Grid1D

Build a one-dimensional Cartesian grid on [x_min, x_max].

Pass n, x_min and x_max for a uniform grid, or faces= alone for explicit face positions. Each node sits at the center of its cell. A half cell, on the uniform grid only, puts its node on the boundary.

Parameters:

Name	Type	Description	Default
`n`	`int \| None`	The number of cells.	`None`
`x_min`	`float \| None`	The left end.	`None`
`x_max`	`float \| None`	The right end, greater than `x_min`.	`None`
`faces`	`ArrayLike \| None`	Strictly increasing face positions, at least two, in place of `n`, `x_min` and `x_max`.	`None`
`half_cell_left`	`bool`	If `True`, the first cell is a half cell with its node at `x_min`.	`False`
`half_cell_right`	`bool`	If `True`, the last cell is a half cell with its node at `x_max`.	`False`
`method`	`str`	The discretization, `"fvm"` or `"fdm"`.	`'fvm'`

Returns:

Type	Description
`Grid1D`	A Grid1D.

diff_soup.domains.cartesian_grid_2d ¶

cartesian_grid_2d(nx: int, ny: int, lx: float, ly: float, *, method: str = 'fvm') -> GridND

Build a two-dimensional Cartesian grid on [0, lx] x [0, ly].

Parameters:

Name	Type	Description	Default
`nx`	`int`	The number of cells along x.	required
`ny`	`int`	The number of cells along y.	required
`lx`	`float`	The length along x.	required
`ly`	`float`	The length along y.	required
`method`	`str`	The discretization, `"fvm"` or `"fdm"`.	`'fvm'`

Returns:

Type	Description
`GridND`	A GridND.

diff_soup.domains.cartesian_grid_3d ¶

cartesian_grid_3d(nx: int, ny: int, nz: int, lx: float, ly: float, lz: float, *, method: str = 'fvm') -> GridND

Build a three-dimensional Cartesian grid on [0, lx] x [0, ly] x [0, lz].

Parameters:

Name	Type	Description	Default
`nx`	`int`	The number of cells along x.	required
`ny`	`int`	The number of cells along y.	required
`nz`	`int`	The number of cells along z.	required
`lx`	`float`	The length along x.	required
`ly`	`float`	The length along y.	required
`lz`	`float`	The length along z.	required
`method`	`str`	The discretization, `"fvm"` or `"fdm"`.	`'fvm'`

Returns:

Type	Description
`GridND`	A GridND.

diff_soup.domains.cylindrical_grid_1d ¶

cylindrical_grid_1d(n: int, r_max: float, *, r_min: float = 0.0, half_cell_right: bool = True, method: str = 'fvm') -> Grid1D

Build a radial grid in cylindrical coordinates on [r_min, r_max].

Cell volumes and face areas carry the metric, per radian, so the axis at r = 0 needs no special treatment. With r_min > 0 the first node sits at r_min.

Parameters:

Name	Type	Description	Default
`n`	`int`	The number of cells.	required
`r_max`	`float`	The outer radius.	required
`r_min`	`float`	The inner radius.	`0.0`
`half_cell_right`	`bool`	If `True`, the last cell is a half cell with its node at `r_max`, the surface.	`True`
`method`	`str`	The discretization, `"fvm"` or `"fdm"`.	`'fvm'`

Returns:

Type	Description
`Grid1D`	A Grid1D.

diff_soup.domains.spherical_grid_1d ¶

spherical_grid_1d(n: int, r_max: float, *, r_min: float = 0.0, half_cell_right: bool = True, method: str = 'fvm') -> Grid1D

Build a radial grid in spherical coordinates on [r_min, r_max].

Cell volumes and face areas carry the metric, per steradian, so the center at r = 0 needs no special treatment. With r_min > 0 the first node sits at r_min.

Parameters:

Name	Type	Description	Default
`n`	`int`	The number of cells.	required
`r_max`	`float`	The outer radius.	required
`r_min`	`float`	The inner radius.	`0.0`
`half_cell_right`	`bool`	If `True`, the last cell is a half cell with its node at `r_max`, the surface.	`True`
`method`	`str`	The discretization, `"fvm"` or `"fdm"`.	`'fvm'`

Returns:

Type	Description
`Grid1D`	A Grid1D.

diff_soup.domains.staggered_grid_2d ¶

staggered_grid_2d(nx: int, ny: int, lx: float, ly: float, *, method: str = 'fvm') -> StaggeredGrid

Build a two-dimensional staggered (MAC) grid on [0, lx] x [0, ly].

Parameters:

Name	Type	Description	Default
`nx`	`int`	The number of cells along x.	required
`ny`	`int`	The number of cells along y.	required
`lx`	`float`	The length along x.	required
`ly`	`float`	The length along y.	required
`method`	`str`	The discretization of the underlying `cell_grid`, `"fvm"` or `"fdm"`.	`'fvm'`

Returns:

Type	Description
`StaggeredGrid`	A StaggeredGrid.

diff_soup.domains.staggered_grid_3d ¶

staggered_grid_3d(nx: int, ny: int, nz: int, lx: float, ly: float, lz: float, *, method: str = 'fvm') -> StaggeredGrid

Build a three-dimensional staggered (MAC) grid on [0, lx] x [0, ly] x [0, lz].

Parameters:

Name	Type	Description	Default
`nx`	`int`	The number of cells along x.	required
`ny`	`int`	The number of cells along y.	required
`nz`	`int`	The number of cells along z.	required
`lx`	`float`	The length along x.	required
`ly`	`float`	The length along y.	required
`lz`	`float`	The length along z.	required
`method`	`str`	The discretization of the underlying `cell_grid`, `"fvm"` or `"fdm"`.	`'fvm'`

Returns:

Type	Description
`StaggeredGrid`	A StaggeredGrid.

Mesh factories¶

diff_soup.domains.interval_mesh ¶

interval_mesh(n_cells: int, x_min: float, x_max: float, *, with_boundary_tags: bool = True) -> Mesh

Build a one-dimensional mesh of two-node line elements on [x_min, x_max].

It is the finite-element counterpart of cartesian_grid_1d. With tags, its end nodes are named "left" and "right".

Parameters:

Name	Type	Description	Default
`n_cells`	`int`	The number of elements.	required
`x_min`	`float`	The left end.	required
`x_max`	`float`	The right end, greater than `x_min`.	required
`with_boundary_tags`	`bool`	If `True`, return a TaggedMesh with both ends tagged.	`True`

Returns:

Type	Description
`Mesh`	A Mesh, or a `TaggedMesh` when tagged.

diff_soup.domains.cylindrical_interval_mesh ¶

cylindrical_interval_mesh(n_cells: int, r_max: float, *, r_min: float = 0.0, with_boundary_tags: bool = True) -> Mesh

Build a radial line mesh in cylindrical coordinates on [r_min, r_max].

Its element measure carries the metric r dr per radian, the same metric as cylindrical_grid_1d. With tags, the inner and outer ends are named "center" and "surface".

Parameters:

Name	Type	Description	Default
`n_cells`	`int`	The number of elements.	required
`r_max`	`float`	The outer radius, greater than `r_min`.	required
`r_min`	`float`	The inner radius, zero or greater.	`0.0`
`with_boundary_tags`	`bool`	If `True`, return a TaggedMesh with both ends tagged.	`True`

Returns:

Type	Description
`Mesh`	A Mesh, or a `TaggedMesh` when tagged.

diff_soup.domains.spherical_interval_mesh ¶

spherical_interval_mesh(n_cells: int, r_max: float, *, r_min: float = 0.0, with_boundary_tags: bool = True) -> Mesh

Build a radial line mesh in spherical coordinates on [r_min, r_max].

Its element measure carries the metric r**2 dr per steradian, the same metric as spherical_grid_1d. With tags, the inner and outer ends are named "center" and "surface".

Parameters:

Name	Type	Description	Default
`n_cells`	`int`	The number of elements.	required
`r_max`	`float`	The outer radius, greater than `r_min`.	required
`r_min`	`float`	The inner radius, zero or greater.	`0.0`
`with_boundary_tags`	`bool`	If `True`, return a TaggedMesh with both ends tagged.	`True`

Returns:

Type	Description
`Mesh`	A Mesh, or a `TaggedMesh` when tagged.

diff_soup.domains.rectangle_mesh ¶

rectangle_mesh(nx: int, ny: int, lx: float, ly: float, *, periodic_axes: tuple[int, ...] = (), ele_type: str = 'QUAD4', with_boundary_tags: bool = False) -> Mesh

Build a structured mesh of the rectangle [0, lx] x [0, ly].

ele_type="TRI3" splits each cell into two triangles. With tags, the sides are named "left", "right", "bottom" and "top".

Parameters:

Name	Type	Description	Default
`nx`	`int`	The number of cells along x.	required
`ny`	`int`	The number of cells along y.	required
`lx`	`float`	The length along x.	required
`ly`	`float`	The length along y.	required
`periodic_axes`	`tuple[int, ...]`	Axes whose opposite sides are identified as periodic.	`()`
`ele_type`	`str`	The element type, `"QUAD4"` or `"TRI3"`.	`'QUAD4'`
`with_boundary_tags`	`bool`	If `True`, return a TaggedMesh with the four sides tagged.	`False`

Returns:

Type	Description
`Mesh`	A Mesh, or a `TaggedMesh` when tagged.

diff_soup.domains.box_mesh ¶

box_mesh(nx: int, ny: int, nz: int, lx: float, ly: float, lz: float, *, periodic_axes: tuple[int, ...] = (), ele_type: str = 'HEX8', with_boundary_tags: bool = False) -> Mesh

Build a structured mesh of the box [0, lx] x [0, ly] x [0, lz].

ele_type="TET4" splits each cell into six tetrahedra. With tags, the faces are named "x_min", "x_max", "y_min", "y_max", "z_min" and "z_max".

Parameters:

Name	Type	Description	Default
`nx`	`int`	The number of cells along x.	required
`ny`	`int`	The number of cells along y.	required
`nz`	`int`	The number of cells along z.	required
`lx`	`float`	The length along x.	required
`ly`	`float`	The length along y.	required
`lz`	`float`	The length along z.	required
`periodic_axes`	`tuple[int, ...]`	Axes whose opposite faces are identified as periodic.	`()`
`ele_type`	`str`	The element type, `"HEX8"` or `"TET4"`.	`'HEX8'`
`with_boundary_tags`	`bool`	If `True`, return a TaggedMesh with the six faces tagged.	`False`

Returns:

Type	Description
`Mesh`	A Mesh, or a `TaggedMesh` when tagged.

diff_soup.domains.geometry.polygons_to_tagged_mesh ¶

polygons_to_tagged_mesh(regions: Dict[str, List[Tuple[float, float]]], *, mesh_size: float = 1.0, segment_tags: Optional[Dict[str, List[int]]] = None, segment_names: Optional[dict] = None, order: int = 1, verbose: bool = False) -> TaggedMesh

Triangulate several named polygons into one mesh whose regions share their interfaces.

Each cell is tagged with its region, and each region's outer edges with its name. The edges two regions share are tagged with both names, sorted and joined by |, such as "bowl|soup". The tag group "boundary" holds the whole outer boundary.

segment_tags splits a region's outer edges as polygon_to_tagged_mesh does: one integer per polygon edge, with ids repeated to group edges, and segment_names naming each id (segment_tags={"bowl": [1, 1, 2, 1]} with segment_names={1: "bowl", 2: "table"}). Interface edges keep their "a|b" tag.

Parameters:

Name	Type	Description	Default
`regions`	`Dict[str, List[Tuple[float, float]]]`	The polygon of each region by name, where a later region takes any overlap with an earlier one.	required
`mesh_size`	`float`	The target element size.	`1.0`
`segment_tags`	`Optional[Dict[str, List[int]]]`	The tag of each polygon edge, by region name. A region left out keeps one tag, its name.	`None`
`segment_names`	`Optional[dict]`	The name of each id in `segment_tags`, such as `{2: "table"}`. Defaults to `"<region>_seg_<i>"`, after the first polygon edge `i` carrying the id.	`None`
`order`	`int`	`1` for TRI3 elements, `2` for TRI6.	`1`
`verbose`	`bool`	If `True`, show gmsh's output.	`False`

Returns:

Type	Description
`TaggedMesh`	A TaggedMesh with boundary and cell tags.

Raises:

Type	Description
`ValueError`	If there are no regions, a polygon has fewer than three vertices, a name contains `\|`, a region ends up with no cells, or `segment_tags` and `segment_names` do not match the regions, their edges and each other.

Mesh types¶

diff_soup.domains.Mesh ¶

A finite-element mesh given by its node coordinates and element connectivity.

Build one from your own points and cells, or with a mesh factory. A mesh always discretizes by finite elements.

Parameters:

Name	Type	Description	Default
`points`		The node coordinates, shape `(num_nodes, dim)`.	required
`cells`		The node indices of each element, one row per element. A row may run clockwise or counter-clockwise.	required
`ele_type`	`str`	The element type, such as `"LINE2"`, `"TRI3"`, `"QUAD4"`, `"TET4"` or `"HEX8"`.	required
`periodic_axes`		Axes whose opposite boundaries are identified as periodic, in two or more dimensions.	`()`
`coord_system`	`CoordSystem`	The CoordSystem, which must be Cartesian unless the mesh is one-dimensional.	`CARTESIAN`

diff_soup.domains.TaggedMesh ¶

A Mesh whose boundary facets, and optionally cells, carry integer tags.

Boundary conditions select facets by name, as tag="rim" or with ds.tags. Facet integrals, such as flux conditions and boundary operators, support LINE2, TRI3, TRI6, QUAD4, QUAD8, TET4, TET10, HEX8, HEX20 and HEX27 elements. Cell tags name regions, and cell_indicator turns one into a piecewise coefficient.

Parameters:

Name	Type	Description	Default
`points`		The node coordinates, shape `(num_nodes, dim)`.	required
`cells`		The node indices of each element, one row per element.	required
`ele_type`	`str`	The element type.	required
`boundary_edges`		The boundary edges of a two-dimensional mesh, each as its two end nodes, then its midside node on a quadratic mesh.	`None`
`boundary_faces`		The boundary faces of a three-dimensional mesh by their corners in cyclic order, or the boundary nodes of a one-dimensional mesh.	`None`
`boundary_tags`		One integer tag per boundary edge or face, in the same order.	`None`
`tag_names`		The name of each tag, such as `{1: "left"}`.	`None`
`tag_groups`		Group names that each select several tags, such as `{"boundary": (1, 2)}`.	`None`
`periodic_axes`		As for Mesh.	`()`
`coord_system`	`CoordSystem`	As for Mesh.	`CARTESIAN`
`cell_tags`		One integer tag per cell, shape `(num_cells,)`, or `None` (the default) for no regions.	`None`
`cell_tag_names`		The tag of each region name, such as `{"soup": 1, "bowl": 2}`.	`None`

facets_for_tag ¶

facets_for_tag(tag: int) -> np.ndarray

Return the boundary facets that carry an integer tag.

Parameters:

Name	Type	Description	Default
`tag`	`int`	The integer tag.	required

Returns:

Type	Description
`ndarray`	The matching edges, or faces if no edge matches, as rows of node indices, possibly empty.

facets_for_name ¶

facets_for_name(name: str) -> np.ndarray

Return the boundary facets of a named tag or tag group.

Parameters:

Name	Type	Description	Default
`name`	`str`	A name in `tag_names`, or a key of `tag_groups`.	required

Returns:

Type	Description
`ndarray`	The facets of every tag with that name or in that group, as rows of node indices.

Raises:

Type	Description
`KeyError`	If the mesh names no tags, or no tag or group is called `name`.

cell_indicator ¶

cell_indicator(name: str) -> Float[Array, ' num_cells']

Return one per cell of a named region and zero elsewhere.

Pass the result, or a weighted sum of several, wherever an operator takes a value per cell, such as D in q.ops.div_flux.

Parameters:

Name	Type	Description	Default
`name`	`str`	A name in `cell_tag_names`.	required

Returns:

Type	Description
`Float[Array, ' num_cells']`	A JAX array of shape `(num_cells,)` in the default float dtype.

Raises:

Type	Description
`KeyError`	If the mesh names no region `name`.

Domain types

The details matter only when extending the library.

diff_soup.domains.Domain ¶

The base class of every grid and mesh that spec.domain registers.

Subclass it to add a domain type. The operator bank chooses how to evaluate each operator from discretization_type.

dim `abstractmethod` `property` ¶

dim: int

The number of spatial dimensions.

num_nodes `abstractmethod` `property` ¶

num_nodes: int

The number of nodes, which on a grid is the number of cells.

discretization_type `abstractmethod` `property` ¶

discretization_type: str

The discretization, "fvm", "fdm" or "fem" for the built-in domains.

node_coordinates `abstractmethod` ¶

node_coordinates() -> jax.Array

Return the coordinates of every node.

Returns:

Type	Description
`Array`	An array of shape `(num_nodes, dim)`.

diff_soup.domains.Grid1D ¶

A one-dimensional grid of n cells, each holding one node.

Build it with cartesian_grid_1d, cylindrical_grid_1d or spherical_grid_1d.

Attributes:

Name	Type	Description
`nodes`	`Float[Array, 'n']`	The node positions, one per cell.
`faces`	`Float[Array, 'n_plus_1']`	The face positions, `n + 1` of them.
`cell_widths`	`Float[Array, 'n']`	The width of each cell.
`node_distances`	`Float[Array, 'n_minus_1']`	The distances between neighbouring nodes, `n - 1` of them.
`cell_volumes`	`Float[Array, 'n']`	The volume of each cell, per radian or per steradian on a curvilinear grid.
`face_areas`	`Float[Array, 'n_plus_1']`	The area of each face, per radian or per steradian on a curvilinear grid.
`coord_system`	`CoordSystem`	The CoordSystem of the grid.
`n`	`int`	The number of cells.

nodes `instance-attribute` ¶

nodes: Float[Array, 'n']

faces `instance-attribute` ¶

faces: Float[Array, 'n_plus_1']

cell_volumes `instance-attribute` ¶

cell_volumes: Float[Array, 'n']

coordinates `property` ¶

coordinates: tuple[Array]

The cell-centre positions along the axis, as a 1-tuple of nodes.

meshgrid ¶

meshgrid() -> tuple[Array]

Return (nodes,), the one-axis case of GridND.meshgrid.

scaled ¶

scaled(factor: ArrayLike) -> 'Grid1D'

Return the grid with every length multiplied by factor.

Face areas and cell volumes scale with the metric, so spherical_grid_1d(n, 1.0).scaled(R) matches spherical_grid_1d(n, R). factor may be a traced value, which is how spec.geometry makes the size of a grid a parameter.

Parameters:

Name	Type	Description	Default
`factor`	`ArrayLike`	The length scale.	required

Returns:

Type	Description
`'Grid1D'`	A Grid1D with the same cells, coordinate system and discretization.

diff_soup.domains.GridND ¶

A tensor-product Cartesian grid in two or three dimensions.

Build it with cartesian_grid_2d or cartesian_grid_3d. A field on it is a flat array of num_nodes values in row-major order over shape.

Attributes:

Name	Type	Description
`axes`	`tuple`	One Grid1D per axis.

axes `instance-attribute` ¶

axes: tuple

shape `property` ¶

shape: tuple[int, ...]

The number of cells along each axis.

coordinates `property` ¶

coordinates: tuple[Array, ...]

The cell-centre positions along each axis, one 1D array per axis.

meshgrid ¶

meshgrid() -> tuple[Array, ...]

Return one coordinate array of shape shape per axis.

The arrays are in field storage order, so grid.meshgrid()[k].ravel() equals grid.node_coordinates()[:, k].

Example

x, y = grid.meshgrid()
u0 = (jnp.sin(jnp.pi * x) * jnp.sin(jnp.pi * y)).ravel()

node_coordinates ¶

node_coordinates() -> Array

Return the coordinates of every node, in row-major order over shape.

Returns:

Type	Description
`Array`	An array of shape `(num_nodes, dim)`.

diff_soup.domains.StaggeredGrid ¶

A marker-and-cell (MAC) grid, with scalars at cell centers and velocities on cell faces.

Each velocity component lives on the faces normal to its axis, so its array has one extra entry along that axis. The built-in operator bank has no staggered stencils, so apply its operators on cell_grid or register a bank that implements them.

Attributes:

Name	Type	Description
`cell_grid`	`GridND`	The GridND of cell centers.

cell_grid `instance-attribute` ¶

cell_grid: GridND

shape `property` ¶

shape: tuple[int, ...]

The number of cells along each axis.

face_shape ¶

face_shape(location: Location) -> tuple[int, ...]

Return the array shape of the values stored at location.

Parameters:

Name	Type	Description	Default
`location`	`Location`	A `Location` from `diff_soup.domains.fields`, either `CELL_CENTER` or a face location `FACE_X`, `FACE_Y` or `FACE_Z` within the grid's dimension.	required

Returns:

Type	Description
`tuple[int, ...]`	`shape` at cell centers, or `shape` with one more entry along the face's axis.

diff_soup.domains.CoordSystem ¶

The coordinate system of a grid or mesh.

Its members CARTESIAN, CYLINDRICAL and SPHERICAL compare equal to their lowercase names. Cylindrical and spherical domains are one-dimensional, with the radius as their coordinate.

diff_soup.domains.DiscreteField ¶

An array of values paired with the location they live at.

The location is a Location from diff_soup.domains.fields: NODE, CELL_CENTER, FACE_X, FACE_Y, FACE_Z or QUAD_POINT. The field is a pytree whose one leaf is the array. replace(array) returns a copy with new values, and at(location) a copy at a new location.

Parameters:

Name	Type	Description	Default
`array`		The values.	required
`location`	`Location`	A `Location`, or its string value such as `"face_x"`.	required

Domains¶

Grid factories¶

diff_soup.domains.cartesian_grid_1d ¶

diff_soup.domains.cartesian_grid_2d ¶

diff_soup.domains.cartesian_grid_3d ¶

diff_soup.domains.cylindrical_grid_1d ¶

diff_soup.domains.spherical_grid_1d ¶

diff_soup.domains.staggered_grid_2d ¶

diff_soup.domains.staggered_grid_3d ¶

Mesh factories¶

diff_soup.domains.interval_mesh ¶

diff_soup.domains.cylindrical_interval_mesh ¶

diff_soup.domains.spherical_interval_mesh ¶

diff_soup.domains.rectangle_mesh ¶

diff_soup.domains.box_mesh ¶

diff_soup.domains.geometry.polygons_to_tagged_mesh ¶

Mesh types¶

diff_soup.domains.Mesh ¶

diff_soup.domains.TaggedMesh ¶

facets_for_tag ¶

facets_for_name ¶

cell_indicator ¶

diff_soup.domains.Domain ¶

dim abstractmethod property ¶

num_nodes abstractmethod property ¶

discretization_type abstractmethod property ¶

node_coordinates abstractmethod ¶

diff_soup.domains.Grid1D ¶

nodes instance-attribute ¶

faces instance-attribute ¶

cell_volumes instance-attribute ¶

coordinates property ¶

meshgrid ¶

scaled ¶

diff_soup.domains.GridND ¶

axes instance-attribute ¶

shape property ¶

coordinates property ¶

meshgrid ¶

node_coordinates ¶

diff_soup.domains.StaggeredGrid ¶

cell_grid instance-attribute ¶

shape property ¶

face_shape ¶

diff_soup.domains.CoordSystem ¶

diff_soup.domains.DiscreteField ¶

dim `abstractmethod` `property` ¶

num_nodes `abstractmethod` `property` ¶

discretization_type `abstractmethod` `property` ¶

node_coordinates `abstractmethod` ¶

nodes `instance-attribute` ¶

faces `instance-attribute` ¶

cell_volumes `instance-attribute` ¶

coordinates `property` ¶

axes `instance-attribute` ¶

shape `property` ¶

coordinates `property` ¶

cell_grid `instance-attribute` ¶

shape `property` ¶