Spatial discretization

At a high level, we represent the spatial domain with a tree of jaxhps.DiscretizationNode2D or jaxhps.DiscretizationNode3D objects. Generally, the user will specify a root Node, and then pass the root along with other information to a constructor of a jaxhps.Domain object. The jaxhps.Domain object will build the discretization tree and populate all of the discretization points.

For instance, in the jaxhps documentation, we saw how to use the jaxhps.Domain constructor to build a uniform 2D discretization. We can also use a constructor to build an adaptive discretization. The adaptive discretization is built by recursively subdividing nodes until a specified function can be represented to a desired accuracy. We will use a simulated 2D permittivity field as an example.

import jaxhps

# Not distributed in the jaxhps package,
# see https://github.com/meliao/jaxhps/blob/main/examples/poisson_boltzmann_utils.py
from poisson_boltzmann_utils import permittivity_2D

# Create a root node for the domain
root = jaxhps.DiscretizationNode2D(xmin=-1.0, xmax=1.0, ymin=-1.0, ymax=1.0)

# Create a Domain object with adaptive discretization
domain = jaxhps.Domain.from_adaptive_discretization(p=16,
                                                   q=14,
                                                   root=root,
                                                   f=permittivity_2D, # Can also specify a list of functions
                                                   tol=1e-06)

print(domain.n_leaves)

DiscretizationNode objects

class jaxhps.DiscretizationNode2D(xmin: float, xmax: float, ymin: float, ymax: float, depth: int = 0, children: Tuple[DiscretizationNode2D] = ())

xmin: float

xmax: float

ymin: float

ymax: float

depth: int

n_0: int: Number of quadrature points on the boundary of side 0.

n_1: int: Number of quadrature points on the boundary of side 1.

n_2: int: Number of quadrature points on the boundary of side 2.

n_3: int: Number of quadrature points on the boundary of side 3.

children: Tuple[DiscretizationNode2D]: Child discretization nodes. Initialized to an empty tuple.

class jaxhps.DiscretizationNode3D(xmin: float, xmax: float, ymin: float, ymax: float, zmin: float, zmax: float, depth: int = 0, children: Tuple[DiscretizationNode3D] = ())

xmin: float

xmax: float

ymin: float

ymax: float

zmin: float

zmax: float

depth: int

n_0: int: Number of quadrature points on the boundary of face 0.

n_1: int: Number of quadrature points on the boundary of face 1.

n_2: int: Number of quadrature points on the boundary of face 2.

n_3: int: Number of quadrature points on the boundary of face 3.

n_4: int: Number of quadrature points on the boundary of face 4.

n_5: int: Number of quadrature points on the boundary of face 5.

children: Child discretization nodes. Initialized to an empty tuple.

jaxhps.get_all_leaves(node: DiscretizationNode2D | DiscretizationNode3D) → Tuple[DiscretizationNode2D | DiscretizationNode3D]

Returns all of the leaf descendants of the given node.

Parameters:: node (DiscretizationNode2D | DiscretizationNode3D) – Root of the tree.
Returns:: All of the leaves.
Return type:: Tuple[DiscretizationNode2D | DiscretizationNode3D]

Domain object

class jaxhps.Domain(p: int, q: int, root: DiscretizationNode2D | DiscretizationNode3D, L: int | None = None)

p: int: Polynomial order for Chebyshev points.

q: int: Polynomial order for Gauss-Legendre points.

root: DiscretizationNode2D | DiscretizationNode3D: Root node of the discretization tree

L: int | None: Number of levels of uniform refinement, or None for adaptive refinement.

n_leaves: int: The number of leaves in the discretization tree.

interior_points: Array: Interior Chebyshev points with shape (n_leaves, p^d, d)

boundary_points: Array: Boundary Gauss points with shape (x q^{d-1}, d), where x is the number of leaves touching the boundary.

interp_to_interior_points(values: Array, sample_points_x: Array, sample_points_y: Array, sample_points_z: Array = None) → Array

This is a utility for interpolating from values on a rectangular grid to values on the HPS grid. The interpolation method builds a separate barycentric Lagrange interpolation matrix for each leaf on the HPS grid and maps the values to the leaf discretization points.

The values are assumed to be defined on sample_points:

X, Y, Z = jnp.meshgrid(sample_points_x,
                     sample_points_y,
                     sample_points_z,
                     indexing="ij")
sample_points = jnp.concatenate((jnp.expand_dims(X, 3),
                               jnp.expand_dims(Y, 3),
                               jnp.expand_dims(Z, 3)),
                              axis=3
                             )

Parameters:

values ((jax.Array)) – Has shape (n_x, n_y) or (n_x, n_y, n_z), and specifies the values of the function on the rectangular grid.
sample_points_x ((jax.Array)) – Has shape (n_x,). Specifies the x-coordinates of the rectangular grid.
sample_points_y ((jax.Array)) – Has shape (n_y,). Specifies the y-coordinates of the rectangular grid.
sample_points_z ((optional, jax.Array)) – Has shape (n_z,). Specifies the z-coordinates of the rectangular grid. This parameter is optional and should be provided only for 3D cases. Defaults to None.

Returns:

Samples on the HPS grid. Has shape (n_leaves, p^d), where d is the dimension of the problem.

Return type:

jax.Array

interp_from_interior_points(samples: Array, eval_points_x: Array, eval_points_y: Array, eval_points_z: Array = None) → Tuple[Array]

This is a method for interpolating from the HPS grid to a rectangular grid. For each point in the rectangular grid, a barycentric Lagrange interpolation matrix is built from the containing leaf to the point, and the function is interpolated from the leaf to that point.

Parameters:

samples (jax.Array) – Function sampled on the HPS grid. Has shape (n_leaves, p^d).
eval_points_x (jax.Array) – Evaluation points in the x dimension. Has shape (n_x,).
eval_points_y (jax.Array) – Evaluation points in the y dimension. Has shape (n_y,).
eval_points_z (optional, Jax.Array) – Evaluation points in the z dimension. Has shape (n_z,). Defaults to None.

Returns:

vals (jax.Array) – Function sampled on a rectangular grid of shape (n_x, n_y) or (n_x, n_y, n_z)
target_pts (jax.Array) – The target points that the vals are sampled on. Has shape (n_x, n_y, 2) or (n_x, n_y, n_z, 3).

get_adaptive_boundary_data_lst(f: Callable[[Array], Array]) → List[Array]

Given a callable object f, this function evaluates the function at the boundary discretization points, and organizes the results into a list. This is helpful for specifying boundary conditions in the adaptive case, where it is not clear a priori how many points will be on each part of the boundary.

Parameters:: f (Callable[[jax.Array], jax.Array]) – Must have signature […, d] -> […].
Returns:: Each element of the list corresponds to a side (2D) or face (3D) of the boundary.
Return type:: List[jax.Array]

classmethod from_adaptive_discretization(p: int, q: int, root: DiscretizationNode2D | DiscretizationNode3D, f: Callable[[Array], Array] | List[Callable[[Array], Array]], tol: float, use_level_restriction: bool = True, use_l_2_norm: bool = False) → Domain

This is a constructor for creating a Domain when using an adaptive discretization. Given the root of the domain and a function f to be evaluated on the domain, this method will adaptively refine the HPS grid until reaching a specified tolerance tol. Multiple functions for adaptive refinement can be specified in a list.

The tolerance is enforced in the \(\ell_\infty\) norm by default, but can also be enforced in \(\ell_2\).

Parameters:

p (int) – Polynomial order for Chebyshev points.
q (int) – Polynomial order for Gauss-Legendre points.
root (DiscretizationNode2D | DiscretizationNode3D) – Specifies the root of the discretization tree.
f (Callable[[jax.Array], jax.Array] | List[Callable[[jax.Array], jax.Array]]) – Function(s) to be used in the adaptive refinement method.
tol (float) – Tolerance parameter.
use_level_restriction (bool, optional) – Whether to enforce the level restriction criterion. Defaults to True. If set to False, could cause errors in the merge stage.
use_l_2_norm (bool, optional) – If set to True, the refinement uses a relative L_2 criterion instead of L_infinity.

Returns:

The domain object with an adaptively-refined discretization tree.

Return type:

Domain