API for individual stages of HPS

Local Solve Stage

jaxhps.local_solve.local_solve_stage_adaptive_2D_DtN(pde_problem: PDEProblem, device: Device = jax.devices()[0], host_device: Device = jax.devices('cpu')[0]) → Tuple[Array, Array, Array, Array]

This function performs the local solve stage for 2D adaptive discretization problems, creating DtN matrices.

The 2D adaptive version of the local solve stage is similar to the 2D uniform version. The major difference appears when dealing with the pre-computed operator Q, which must be scaled by the side length of each leaf. This means Q must be re-computed for each leaf solve.

Parameters:

pde_problem (PDEProblem) – Specifies the discretization, differential operator, source function, and keeps track of the pre-computed differentiation and interpolation matrices.
device (jax.Device) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device) – Where to place the output. Defaults to jax.devices("cpu")[0].

Returns:

Y (jax.Array) – Solution operators mapping from Dirichlet boundary data to homogeneous solutions on the leaf interiors. Has shape (n_leaves, p^2, 4q)
T (jax.Array) – Dirichlet-to-Neumann matrices for each leaf. Has shape (n_leaves, 4q, 4q)
v (jax.Array) – Leaf-level particular solutions. Has shape (n_leaves, p^2)
h (jax.Array) – Outgoing boundary data. This is the outward-pointing normal derivative of the particular solution. Has shape (n_leaves, 4q)

jaxhps.local_solve.local_solve_stage_adaptive_3D_DtN(pde_problem: PDEProblem, device: Device = jax.devices()[0], host_device: Device = jax.devices('cpu')[0]) → Tuple[Array, Array, Array, Array]

This function performs the local solve stage for 3D adaptive discretization problems, creating DtN matrices.

The 3D adaptive version of the local solve stage is similar to the 3D uniform version. The major difference appears when dealing with the pre-computed operator Q, which must be scaled by the side length of each leaf. This means Q must be re-computed for each leaf solve.

Parameters:

pde_problem (PDEProblem) – Specifies the discretization, differential operator, source function, and keeps track of the pre-computed differentiation and interpolation matrices.
device (jax.Device) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device) – Where to place the output. Defaults to jax.devices("cpu")[0].

Returns:

Y (jax.Array) – Solution operators mapping from Dirichlet boundary data to homogeneous solutions on the leaf interiors. Has shape (n_leaves, p^3, 6q^2)
T (jax.Array) – Dirichlet-to-Neumann matrices for each leaf. Has shape (n_leaves, 6q^2, 6q^2)
v (jax.Array) – Leaf-level particular solutions. Has shape (n_leaves, p^3)
h (jax.Array) – Outgoing boundary data. This is the outward-pointing normal derivative of the particular solution. Has shape (n_leaves, 6q^2)

jaxhps.local_solve.local_solve_stage_uniform_2D_DtN(pde_problem: PDEProblem, host_device: Device = jax.devices('cpu')[0], device: Device = jax.devices()[0]) → Tuple[Array, Array, Array, Array]

This function performs the local solve stage for 2D problems with a uniform quadtree, creating DtN matrices.

Parameters:

pde_problem (PDEProblem) – Specifies the discretization, differential operator, source function, and keeps track of the pre-computed differentiation and interpolation matrices.
device (jax.Device) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device) – Where to place the output. Defaults to jax.devices("cpu")[0].

Returns:

Y (jax.Array) – Solution operators mapping from Dirichlet boundary data to homogeneous solutions on the leaf interiors. Has shape (n_leaves, p^2, 4q)
T (jax.Array) – Dirichlet-to-Neumann matrices for each leaf. Has shape (n_leaves, 4q, 4q)
v (jax.Array) – Leaf-level particular solutions. Has shape (n_leaves, p^2) or (n_leaves, p**2, n_src) if multi-source.
h (jax.Array) – Outgoing boundary data. This is the outward-pointing normal derivative of the particular solution. Has shape (n_leaves, 4q) or (n_leaves, 4q, n_src) if multi-source.

jaxhps.local_solve.local_solve_stage_uniform_2D_ItI(pde_problem: PDEProblem, device: Device = jax.devices()[0], host_device: Device = jax.devices('cpu')[0]) → Tuple[Array, Array, Array, Array]

This function performs the local solve stage for 2D problems with a uniform quadtree, creating ItI matrices.

Parameters:

pde_problem (PDEProblem) – Specifies the discretization, differential operator, source function, and keeps track of the pre-computed differentiation and interpolation matrices.
device (jax.Device) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device) – Where to place the output. Defaults to jax.devices("cpu")[0].

Returns:

Y (jax.Array) – Solution operators mapping from Impedance boundary data to homogeneous solutions on the leaf interiors. Has shape (n_leaves, p^2, 4q)
T (jax.Array) – Impedance-to-Impedance matrices for each leaf. Has shape (n_leaves, 4q, 4q)
v (jax.Array) – Leaf-level particular solutions. Has shape (n_leaves, p^2) if there is a single source term, or (n_leaves, p^2, nsrc) if there are multiple source terms.
h (jax.Array) – Outgoing boundary data. This is the outgoing impedance data of the particular solution \(v_n - i \eta v\). Has shape (n_leaves, 4q) if there is a single source term, or (n_leaves, 4q, nsrc) if there are multiple source terms.

jaxhps.local_solve.local_solve_stage_uniform_3D_DtN(pde_problem: PDEProblem, device: Device = jax.devices()[0], host_device: Device = jax.devices('cpu')[0]) → Tuple[Array, Array, Array, Array]

This function performs the local solve stage for 3D problems with a uniform quadtree, creating DtN matrices.

The 3D adaptive version of the local solve stage is similar to the 3D uniform version. The major difference appears when dealing with the pre-computed operator Q, which must be scaled by the side length of each leaf. This means Q must be re-computed for each leaf solve.

Parameters:

pde_problem (PDEProblem) – Specifies the discretization, differential operator, source function, and keeps track of the pre-computed differentiation and interpolation matrices.
device (jax.Device) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device) – Where to place the output. Defaults to jax.devices("cpu")[0].

Returns:

Y (jax.Array) – Solution operators mapping from Dirichlet boundary data to homogeneous solutions on the leaf interiors. Has shape (n_leaves, p^3, 6q^2)
T (jax.Array) – Dirichlet-to-Neumann matrices for each leaf. Has shape (n_leaves, 6q^2, 6q^2)
v (jax.Array) – Leaf-level particular solutions. Has shape (n_leaves, p^3) or (n_leaves, p^3, n_src) if multi-source.
h (jax.Array) – Outgoing boundary data. This is the outward-pointing normal derivative of the particular solution. Has shape (n_leaves, 6q^2) or (n_leaves, 6q^2, n_src) if multi-source.

jaxhps.local_solve.nosource_local_solve_stage_uniform_2D_DtN(pde_problem: PDEProblem, device: Device = jax.devices()[0], host_device: Device = jax.devices('cpu')[0]) → Tuple[Array, Array, Array]

This function performs the local solve stage for 2D problems with a uniform quadtree, creating DtN matrices.

Parameters:

pde_problem (PDEProblem) – Specifies the discretization, differential operator, source function, and keeps track of the pre-computed differentiation and interpolation matrices.
device (jax.Device) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device) – Where to place the output. Defaults to jax.devices("cpu")[0].

Returns:

Y (jax.Array) – Solution operators mapping from Dirichlet boundary data to homogeneous solutions on the leaf interiors. Has shape (n_leaves, p^2, 4q)
T (jax.Array) – Dirichlet-to-Neumann matrices for each leaf. Has shape (n_leaves, 4q, 4q)
Phi (jax.Array) – Particular solution operators mapping from the source term evaluated on the interior Cheby nodes to the particular solution on all of the Cheby nodes. Has shape (n_leaves, p^2, (p-1)^2).

jaxhps.local_solve.nosource_local_solve_stage_uniform_2D_ItI(pde_problem: PDEProblem, device: Device = jax.devices()[0], host_device: Device = jax.devices('cpu')[0]) → Tuple[Array, Array, Array]

This function performs the local solve stage for 2D problems with a uniform quadtree, creating ItI matrices.

Parameters:

pde_problem (PDEProblem) – Specifies the discretization, differential operator, source function, and keeps track of the pre-computed differentiation and interpolation matrices.
device (jax.Device) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device) – Where to place the output. Defaults to jax.devices("cpu")[0].

Returns:

Y (jax.Array) – Solution operators mapping from Impedance boundary data to homogeneous solutions on the leaf interiors. Has shape (n_leaves, p^2, 4q)
T (jax.Array) – Impedance-to-Impedance matrices for each leaf. Has shape (n_leaves, 4q, 4q)
Phi (jax.Array) – Particular solution operators mapping from the source term evaluated on the interior Cheby nodes to the particular solution on all of the Cheby nodes. Has shape (n_leaves, p^2, (p-1)^2).

Merge Stage

jaxhps.merge.merge_stage_adaptive_2D_DtN(pde_problem: PDEProblem) → None

Implements adaptive 2D merges of DtN matrices. Merges the nodes in the quadtree four at a time, projecting the rows and columns of DtN matrices as necessary when neighboring nodes have different boundary discretizations.

This function saves the output of each merge step inside the DiscretizationNode2D.data objects that define the quadtree. After this function is called, the top-level DtN matrix can be accessed at pde_problem.domain.root.data.T.

Parameters:: pde_problem (PDEProblem) – Specifies the discretization, differential operator, source function, and keeps track of the pre-computed differentiation and interpolation matrices.
Return type:: None. All of the solution operators are stored inside the pde_problem object.

jaxhps.merge.merge_stage_adaptive_3D_DtN(pde_problem: PDEProblem, T_arr: Array, h_arr: Array, device: Device = jax.devices()[0], host_device: Device = jax.devices('cpu')[0]) → None

Implements adaptive 3D merges of DtN matrices. Merges the nodes in the quadtree eight at a time, projecting the rows and columns of DtN matrices as necessary when neighboring nodes have different boundary discretizations.

This function saves the output of each merge step inside the DiscretizationNode3D.data objects that define the quadtree. After this function is called, the top-level DtN matrix can be accessed at pde_problem.domain.root.data.T.

This function takes as arguments parts of the local solve stage output. These arrays are indexed at the locations of the lowest-level leaves; the lowest-level leaves are then merged using the vectorized code used in the uniform version of this algorithm. This gives us a slight performance boost. At higher levels, the merges are not vectorized.

Parameters:

pde_problem (PDEProblem) – Specifies the discretization, differential operator, source function, and keeps track of the pre-computed differentiation and interpolation matrices.
T_arr (jax.Array) – Array of DtN matrices from the local solve stage. Has shape (n_leaves, 6q^2, 6q^2)
h_arr (jax.Array) – Array of outgoing boundary data from the local solve stage. Has shape (n_leaves, 6q^2)
device (jax.Device) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device) – Where to place the output. Defaults to jax.devices(“cpu”)[0].

Return type:

None. All of the solution operators are stored inside the pde_problem object.

jaxhps.merge.merge_stage_uniform_2D_DtN(T_arr: array, h_arr: array, l: int, device: Device = jax.devices()[0], host_device: Device = jax.devices('cpu')[0], subtree_recomp: bool = False, return_T: bool = False) → Tuple[List[array], List[array], List[array]]

Implements uniform 2D merges of DtN matrices. Merges the nodes in the quadtree four at a time. This function returns lists containing \(S\) and \(\tilde{g}\), giving enough information to propagate boundary data back down the tree in a later part of the algorithm.

If this function is called with the argument return_T=True, the top-level DtN matrix is also returned.

Parameters:

pde_problem (PDEProblem) – Specifies the discretization, differential operator, source function, and keeps track of the pre-computed differentiation and interpolation matrices.
T_arr (jax.Array) – Array of DtN matrices from the local solve stage. Has shape (n_leaves, 4q, 4q)
h_arr (jax.Array) – Array of outgoing boundary data from the local solve stage. Has shape (n_leaves, 4q)
l (int) – Number of levels to merge.
device (jax.Device) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device) – Where to place the output. Defaults to jax.devices(“cpu”)[0].
subtree_recomp (bool) – A flag for used by the subtree recomputation methods, which triggers an early return with just the top-level T and h.
return_T (bool) – A flag to return the top-level T matrix. Defaults to False.

Returns:

S_lst (List[jax.Array]) – A list of propagation operators. The first element of the list are the propagation operators for the nodes just above the leaves, and the last element of the list is the propagation operator for the root of the quadtree.
g_tilde_lst (List[jax.Array]) – A list of incoming particular solution data along the merge interfaces. The first element of the list corresponds to the nodes just above the leaves, and the last element of the list corresponds to the root of the quadtree.
T_last (jax.Array) – The top-level DtN matrix, which is only returned if return_T=True. Has shape (4q, 4q).

jaxhps.merge.merge_stage_uniform_2D_ItI(T_arr: array, h_arr: array, l: int, device: Device = jax.devices()[0], host_device: Device = jax.devices('cpu')[0], subtree_recomp: bool = False, return_T: bool = False, return_h: bool = False) → Tuple[List[Array], List[Array], List[Array]]

Implements uniform 2D merges of ItI matrices. Merges the nodes in the quadtree four at a time. This function uses a Schur complement strategy to reduce the size of matrix inverted in each merge operation. This function returns lists containing \(S\) and \(\tilde{g}\), giving enough information to propagate boundary data back down the tree in a later part of the algorithm.

If this function is called with the argument return_T=True, the top-level DtN matrix is also returned.

Parameters:

pde_problem (PDEProblem) – Specifies the discretization, differential operator, source function, and keeps track of the pre-computed differentiation and interpolation matrices.
T_arr (jax.Array) – Array of ItI matrices from the local solve stage. Has shape (n_leaves, 4q, 4q)
h_arr (jax.Array) – Array of outgoing boundary data from the local solve stage. Has shape (n_leaves, 4q) or (n_leaves, 4q, nsrc)
l (int) – Number of levels to merge
device (jax.Device) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device) – Where to place the output. Defaults to jax.devices(“cpu”)[0].
subtree_recomp (bool) – A flag for used by the subtree recomputation methods, which triggers an early return with just the top-level T and h.
return_T (bool) – A flag to return the top-level T matrix. Defaults to False.

Returns:

S_lst (List[jax.Array]) – A list of propagation operators. The first element of the list are the propagation operators for the nodes just above the leaves, and the last element of the list is the propagation operator for the root of the quadtree.
g_tilde_lst (List[jax.Array]) – A list of incoming particular solution data along the merge interfaces. The first element of the list corresponds to the nodes just above the leaves, and the last element of the list corresponds to the root of the quadtree.
T_last (jax.Array) – The top-level DtN matrix, which is only returned if return_T=True. Has shape (4q, 4q).

jaxhps.merge.merge_stage_uniform_3D_DtN(T_arr: Array, h_arr: Array, l: int, device: Device = jax.devices()[0], host_device: Device = jax.devices('cpu')[0], return_T: bool = False) → Tuple[List[Array], List[Array], List[Array]]

Implements uniform 3D merges of DtN matrices. Merges the nodes in the quadtree eight at a time. This function returns lists containing \(S\) and \(\tilde{g}\), giving enough information to propagate boundary data back down the tree in a later part of the algorithm.

If this function is called with the argument return_T=True, the top-level DtN matrix is also returned.

Parameters:

pde_problem (PDEProblem) – Specifies the discretization, differential operator, source function, and keeps track of the pre-computed differentiation and interpolation matrices.
T_arr (jax.Array) – Array of DtN matrices from the local solve stage. Has shape (n_leaves, 6q^2, 6q^2)
h_arr (jax.Array) – Array of outgoing boundary data from the local solve stage. Has shape (n_leaves, 6q^2)
l (int) – The number of levels in the quadtree.
device (jax.Device) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device) – Where to place the output. Defaults to jax.devices(“cpu”)[0].
return_T (bool) – If True, the top-level DtN matrix is returned. Defaults to False.

Returns:

S_lst (List[jax.Array]) – A list of propagation operators. The first element of the list are the propagation operators for the nodes just above the leaves, and the last element of the list is the propagation operator for the root of the quadtree.
g_tilde_lst (List[jax.Array]) – A list of incoming particular solution data along the merge interfaces. The first element of the list corresponds to the nodes just above the leaves, and the last element of the list corresponds to the root of the quadtree.
T_last (jax.Array) – The top-level DtN matrix, which is only returned if return_T=True. Has shape (4q, 4q).

jaxhps.merge.nosource_merge_stage_uniform_2D_DtN(T_arr: Array, l: int, device: Device = jax.devices()[0], host_device: Device = jax.devices('cpu')[0], return_T: bool = False) → Tuple[List[Array], List[Array], List[Array]]

Implements uniform 2D merges of DtN matrices. Merges the nodes in the quadtree four at a time. This function uses a Schur complement strategy to reduce the size of matrix inverted in each merge operation. This function returns lists containing \(S\) and \(\tilde{g}\), giving enough information to propagate boundary data back down the tree in a later part of the algorithm.

If this function is called with the argument return_T=True, the top-level DtN matrix is also returned.

Parameters:

T_arr (jax.Array) – Array of DtN matrices from the local solve stage. Has shape (n_leaves, 4q, 4q)
l (int) – Number of levels to merge
device (jax.Device) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device) – Where to place the output. Defaults to jax.devices(“cpu”)[0].
return_T (bool) – A flag to return the top-level T matrix. Defaults to False.

Returns:

S_lst (List[jax.Array]) – A list of propagation operators. The first element of the list are the propagation operators for the nodes just above the leaves, and the last element of the list is the propagation operator for the root of the quadtree.
D_inv_lst (List[jax.Array]) – List of pre-computed D^{-1} matrices for each level of the quadtree.
BD_inverse_lst (List[jax.Array]) – List of pre-computed BD^{-1} matrices for each level of the quadtree.
T_last (jax.Array) – The top-level DtN matrix, which is only returned if return_T=True. Has shape (4q, 4q).

jaxhps.merge.nosource_merge_stage_uniform_2D_ItI(T_arr: array, l: int, device: Device = jax.devices()[0], host_device: Device = jax.devices('cpu')[0], return_T: bool = False) → Tuple[List[Array], List[Array], List[Array]]

Implements uniform 2D merges of ItI matrices. Merges the nodes in the quadtree four at a time. This function uses a Schur complement strategy to reduce the size of matrix inverted in each merge operation. This function returns lists containing \(S\) and \(\tilde{g}\), giving enough information to propagate boundary data back down the tree in a later part of the algorithm.

If this function is called with the argument return_T=True, the top-level DtN matrix is also returned.

Parameters:

T_arr (jax.Array) – Array of ItI matrices from the local solve stage. Has shape (n_leaves, 4q, 4q)
l (int) – Number of levels to merge
device (jax.Device) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device) – Where to place the output. Defaults to jax.devices(“cpu”)[0].
return_T (bool) – A flag to return the top-level T matrix. Defaults to False.

Returns:

S_lst (List[jax.Array]) – A list of propagation operators. The first element of the list are the propagation operators for the nodes just above the leaves, and the last element of the list is the propagation operator for the root of the quadtree.
D_inv_lst (List[jax.Array]) – List of pre-computed D^{-1} matrices for each level of the quadtree.
BD_inverse_lst (List[jax.Array]) – List of pre-computed BD^{-1} matrices for each level of the quadtree.
T_last (jax.Array) – The top-level DtN matrix, which is only returned if return_T=True. Has shape (4q, 4q).

Upward Pass

jaxhps.up_pass.up_pass_uniform_2D_DtN(source: Array, pde_problem: PDEProblem, device: Device = jax.devices()[0], host_device: Device = jax.devices('cpu')[0], return_h_last: bool = False) → Tuple[Array, List[Array], Array]

This function performs the upward pass for 2D DtN problems. It recomputes the local solve stage to get outgoing impedance data from the particular solution, which is now known because the source is specified.

Parameters:

source (jax.Array) – Specifies the source term(s). Can have shape (nleaves, p^2) or (nleaves, p^2, nsrc).
pde_probem (PDEProblem) – Specifies the discretization, differential operator, source function, and keeps track of the pre-computed differentiation and interpolation matrices.
device (jax.Device, optional) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device, optional) – Where to place the output. Defaults to jax.devices("cpu")[0].
return_h_last (bool, optional) – If True, return the last h vector, which gives the outward pointing normal derivative data for the particular solution evaluated at the domain boundary.

Returns:

v (jax.Array) – Leaf-level particular solutions. Has shape (n_leaves, p^2)
g_tilde_lst (List[jax.Array]) – List of pre-computed g_tilde matrices for each level of the quadtree.
h_last (jax.Array) – Outward pointing normal derivative data for the particular solution evaluated at the domain boundary. Has shape (nbdry, nsrc). Is only returned if return_h_last=True.

jaxhps.up_pass.up_pass_uniform_2D_ItI(source: Array, pde_problem: PDEProblem, device: Device = jax.devices()[0], host_device: Device = jax.devices('cpu')[0], return_h_last: bool = False) → Tuple[Array, List[Array], Array]

This function performs the upward pass for 2D ItI problems. It recomputes the local solve stage to get outgoing impedance data from the particular solution, which is now known because the source is specified.

Parameters:

source (jax.Array) – Specifies the source term(s). Can have shape (nleaves, p^2) or (nleaves, p^2, nsrc).
pde_probem (PDEProblem) – Specifies the discretization, differential operator, source function, and keeps track of the pre-computed differentiation and interpolation matrices.
device (jax.Device, optional) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device, optional) – Where to place the output. Defaults to jax.devices("cpu")[0].
return_h_last (bool, optional) – If True, return the last h vector, which gives the outgoing impedance data for the particular solution evaluated at the domain boundary.

Returns:

v (jax.Array) – Leaf-level particular solutions. Has shape (n_leaves, p^2)
g_tilde_lst (List[jax.Array]) – List of pre-computed g_tilde matrices for each level of the quadtree.
h_last (jax.Array) – Outgoing impedance data for the particular solution evaluated at the domain boundary. Has shape (nbdry, nsrc). Is only returned if return_h_last=True.

Downward Pass

jaxhps.down_pass.down_pass_adaptive_2D_DtN(pde_problem: PDEProblem, boundary_data: List[Array]) → Array

This function performs the downward pass for 2D adaptive discretizations using DtN matrices.

The boundary data is assumed to be Dirichlet data

Parameters:

pde_problem (PDEProblem) – Specifies the discretization, differential operator, source function, and keeps track of the pre-computed differentiation and interpolation matrices.
boundary_data (List[jax.Array]) – A length-4 list of arrays, specifying samples of the boundary data. jaxhps.Domain.get_adaptive_boundary_data_lst() is a utility for constructing this list.

Returns:

solns – The solutions on the interior points. Has shape (n_leaves, p^2)

Return type:

jax.Array

jaxhps.down_pass.down_pass_adaptive_3D_DtN(pde_problem: PDEProblem, boundary_data: List[Array]) → Array

This function performs the downward pass for 2D adaptive discretizations using DtN matrices.

The boundary data is assumed to be Dirichlet data

Parameters:

pde_problem (PDEProblem) – Specifies the discretization, differential operator, source function, and keeps track of the pre-computed differentiation and interpolation matrices.
boundary_data (List[jax.Array]) – A length-6 list of arrays, specifying samples of the boundary data. jaxhps.Domain.get_adaptive_boundary_data_lst() is a utility for constructing this list.

Returns:

solns – The solutions on the interior points. Has shape (n_leaves, p^3)

Return type:

jax.Array

jaxhps.down_pass.down_pass_uniform_2D_DtN(boundary_data: Array, S_lst: List[Array], g_tilde_lst: List[Array], Y_arr: Array, v_arr: Array, device: Device = jax.devices()[0], host_device: Device = jax.devices('cpu')[0]) → Array

Computes the downward pass of the HPS algorithm. This function takes the Dirichlet data at the boundary of the domain and propagates it down the tree to the leaf nodes.

If Y_arr is None, the function will exit early after doing all of the downward propagation operations.

Parameters:

boundary_data (jax.Array) – An array specifying Dirichlet data on the boundary of the domain. Has shape (n_bdry,) or (n_bdry, nsrc) for multi-source problems.
S_lst (List[jax.Array]) – A list of propagation operators. The first element of the list are the propagation operators for the nodes just above the leaves, and the last element of the list is the propagation operator for the root of the quadtree.
g_tilde_lst (List[jax.Array]) – A list of incoming particular solution data along the merge interfaces. The first element of the list corresponds to the nodes just above the leaves, and the last element of the list corresponds to the root of the quadtree.
Y_arr (jax.Array) – Matrices mapping the solution to the interior of the leaf nodes. Has shape (n_leaf, p^2, 4q).
v_arr (jax.Array) – Particular solution data at the interior of the leaves. Has shape (n_leaf, p^2).
device (jax.Device) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device) – Where to place the output. Defaults to jax.devices(“cpu”)[0].

Returns:

solns – Has shape (n_leaves, p^2). Interior solution on all of the leaf nodes.

Return type:

jax.Array

jaxhps.down_pass.down_pass_uniform_2D_ItI(boundary_data: Array, S_lst: List[Array], g_tilde_lst: List[Array], Y_arr: Array, v_arr: Array, device: Device = jax.devices()[0], host_device: Device = jax.devices('cpu')[0]) → None

Computes the downward pass of the HPS algorithm. This function takes the incoming impedance data at the boundary of the domain:

\[u_n(x) + i \eta u(x) = g(x)\]

and propagates it down to the leaves.

If Y_arr is None, the function will exit early after doing all of the downward propagation operations.

Parameters:

boundary_data (jax.Array) – An array specifying Dirichlet data on the boundary of the domain. Has shape (n_bdry,) or (n_bdry, nsrc).
S_lst (List[jax.Array]) – A list of propagation operators. The first element of the list are the propagation operators for the nodes just above the leaves, and the last element of the list is the propagation operator for the root of the quadtree.
g_tilde_lst (List[jax.Array]) – A list of incoming particular solution data along the merge interfaces. The first element of the list corresponds to the nodes just above the leaves, and the last element of the list corresponds to the root of the quadtree.
Y_arr (jax.Array) – Matrices mapping the solution to the interior of the leaf nodes. Has shape (n_leaf, p^2, 4q).
v_arr (jax.Array) – Particular solution data at the interior of the leaves. Has shape (n_leaf, p^2) or (n_leaf, p^2, nsrc).
device (jax.Device) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device) – Where to place the output. Defaults to jax.devices(“cpu”)[0].

Returns:

solns – Has shape (n_leaves, p^2). Interior solution on all of the leaf nodes.

Return type:

jax.Array

jaxhps.down_pass.down_pass_uniform_3D_DtN(boundary_data: Array, S_lst: List[Array], g_tilde_lst: List[Array], Y_arr: Array, v_arr: Array, device: Device = jax.devices()[0], host_device: Device = jax.devices('cpu')[0]) → None

Computes the downward pass of the HPS algorithm. This function takes the Dirichlet data at the boundary of the domain and propagates it down the tree to the leaf nodes.

If Y_arr is None, the function will exit early after doing all of the downward propagation operations.

Parameters:

boundary_data (jax.Array) – An array specifying Dirichlet data on the boundary of the domain. Has shape (n_bdry,) or (n_bdry, nsrc) for multi-source problems.
S_lst (List[jax.Array]) – A list of propagation operators. The first element of the list are the propagation operators for the nodes just above the leaves, and the last element of the list is the propagation operator for the root of the quadtree.
g_tilde_lst (List[jax.Array]) – A list of incoming particular solution data along the merge interfaces. The first element of the list corresponds to the nodes just above the leaves, and the last element of the list corresponds to the root of the quadtree.
Y_arr (jax.Array) – Matrices mapping the solution to the interior of the leaf nodes. Has shape (n_leaf, p^3, 6q^2).
v_arr (jax.Array) – Particular solution data at the interior of the leaves. Has shape (n_leaf, p^3).
device (jax.Device) – Where to perform the computation. Defaults to jax.devices()[0].
host_device (jax.Device) – Where to place the output. Defaults to jax.devices(“cpu”)[0].

Returns:

solns – Has shape (n_leaves, p^3). Interior solution on all of the leaf nodes.

Return type:

jax.Array