Skip to main content

feax.DCboundary

Dirichlet boundary condition implementation for FEAX finite element framework.

This module provides the core DirichletBC class for boundary condition application and dataclass-based BC specification classes for type-safe definition.

Key Classes:

  • DirichletBC: JAX-compatible BC class with apply methods
  • DirichletBCSpec: Dataclass for specifying individual boundary conditions
  • DirichletBCConfig: Container for multiple BC specifications with convenience methods

DirichletBC Objects​

@dataclass(frozen=True)
class DirichletBC()

JAX-compatible dataclass for Dirichlet boundary conditions.

This class pre-computes and stores all BC information as static JAX arrays, making it suitable for JIT compilation.

bc_rows​

All boundary condition row indices

bc_mask​

Boolean mask for BC rows (size: total_dofs)

bc_vals​

Boundary condition values for each BC row

replace_vals​

def replace_vals(new_bc_vals: np.ndarray) -> 'DirichletBC'

Return a new DirichletBC with updated prescribed values.

The DOF locations (bc_rows, bc_mask) are preserved; only the values at those DOFs change. This is useful for incremental loading where the same DOFs are constrained but the prescribed displacement increases each step.

Parameters​

  • new_bc_vals (np.ndarray): New boundary condition values, same shape as self.bc_vals.

Returns​

DirichletBC A new instance with the updated values.

from_specs​

@staticmethod
def from_specs(problem: 'Problem',
               specs: List['DirichletBCSpec']) -> 'DirichletBC'

Create DirichletBC directly from a list of DirichletBCSpec objects.

This is a convenient factory method that creates a DirichletBC without needing to create an intermediate DirichletBCConfig object.

Parameters​

  • problem (Problem): The finite element problem instance
  • specs (List[DirichletBCSpec]): List of boundary condition specifications

Returns​

DirichletBC The compiled boundary condition object

Examples​

>>> bc = DirichletBC.from_specs(problem, [
...     DirichletBCSpec(left_boundary, 'all', 0.0),
...     DirichletBCSpec(right_boundary, 'x', 0.1)
... ])

apply_boundary_to_J_csr​

def apply_boundary_to_J_csr(bc: DirichletBC,
                            problem: 'Problem',
                            csr_data: np.ndarray,
                            symmetric: bool = True,
                            ts=None) -> np.ndarray

Apply Dirichlet BCs directly to a CSR data array (no BCOO).

The CSR structure (problem.csr_indptr / csr_indices) is fixed, and the canonical pattern already contains every diagonal slot (d, d), so BC enforcement is a pure value transform on csr_data:

  1. Zero every entry whose row (and, when symmetric, column) is a BC DOF.
  2. Set the diagonal slot of each BC row to 1.

This matches textbook symmetric/non-symmetric Dirichlet elimination: the single canonical diagonal slot is zeroed then set to 1. The transform depends only on bc_rows / bc_mask (DOF locations), never on bc_vals, and is fully JIT/trace compatible.

Parameters​

  • bc (DirichletBC): Boundary condition (only bc_mask / bc_rows are used).
  • problem (Problem): Provides the precomputed CSR structure (csr_row_of_slot, csr_indices, csr_diag_slots).
  • csr_data (np.ndarray): CSR value array of length problem.csr_nse (e.g. from :func:``2).
  • symmetric (bool, default True): Symmetric elimination (zero BC rows and columns) vs non-symmetric (zero BC rows only, keep columns / K10 coupling).

Returns​

np.ndarray BC-applied CSR value array (same shape as csr_data).

apply_boundary_to_res​

def apply_boundary_to_res(bc: DirichletBC,
                          res_vec: np.ndarray,
                          sol_vec: np.ndarray,
                          scale: float = 1.0) -> np.ndarray

Apply Dirichlet boundary conditions to residual vector using row elimination.

This function modifies the residual vector to enforce Dirichlet boundary conditions by setting residual entries at BC DOFs to: res[bc_dof] = sol[bc_dof] - bc_val * scale This ensures that the Newton step will drive the solution towards the prescribed values.

The modified residual enforces the constraint that after the Newton update: sol_new[bc_dof] = bc_val * scale

Parameters​

  • bc (DirichletBC): Pre-computed boundary condition information containing: - bc_rows: DOF indices where BCs are applied - bc_vals: Prescribed values at boundary DOFs - bc_mask: Boolean mask (not used in this function) - total_dofs: Total number of DOFs (for validation)
  • res_vec (np.ndarray): The residual vector (flattened) with shape (total_dofs,)
  • sol_vec (np.ndarray): The current solution vector (flattened) with shape (total_dofs,)
  • scale (float, optional): Scaling factor for boundary condition values, by default 1.0 Useful for ramping up BCs or unit conversion

Returns​

np.ndarray The residual vector with boundary conditions applied, same shape as input

Notes​

This function is JAX-JIT compatible and creates a copy of the input residual to avoid modifying the original array. The boundary condition enforcement follows the standard penalty method approach for constraint enforcement in Newton-Raphson solvers.

DirichletBCSpec Objects​

@dataclass
class DirichletBCSpec()

Specification for a single Dirichlet boundary condition.

This dataclass provides a clear, type-safe way to specify boundary conditions.

Parameters​

  • location (Callable[[np.ndarray], bool]): Function that takes a point (x, y, z) and returns True if the point is on the boundary where this BC should be applied
  • component (Union[int, str]): Which component to constrain: - For scalar problems: must be 0 or 'all' - For vector problems: 0='x', 1='y', 2='z', or 'all' for all components
  • value (Union[float, Callable[[np.ndarray], float]]): The prescribed value, either: - A constant float value - A function that takes a point and returns the value at that point
  • variable_index (Optional[int]): For multi-variable problems, specifies which variable this BC applies to. If None (default), applies to all variables (backward compatible behavior). For single-variable problems, this parameter is ignored.

Examples​

>>> # Fix left boundary in x-direction to zero
>>> bc1 = DirichletBCSpec(
...     location=lambda pt: np.isclose(pt[0], 0.0),
...     component='x',  # or component=0
...     value=0.0
... )

__post_init__​

def __post_init__() -> None

Validate and normalize the component specification.

  1. Convert string component names ('x', 'y', 'z', 'all') to integers
  2. Validate integer component indices are non-negative
  3. Convert constant values to functions for uniform interface

Raises​

ValueError If component string is invalid or integer component is negative

DirichletBCConfig Objects​

@dataclass
class DirichletBCConfig()

Configuration for all Dirichlet boundary conditions in a problem.

This dataclass holds a collection of DirichletBCSpec objects and provides methods to convert to the format expected by DirichletBC.from_bc_info.

Parameters​

  • specs (List[DirichletBCSpec]): List of boundary condition specifications

Examples​

>>> # Create BC configuration for elasticity problem
>>> bc_config = DirichletBCConfig([
...     DirichletBCSpec(
...         location=lambda pt: np.isclose(pt[0], 0.0),
...         component='all',
...         value=0.0
...     ),
...     DirichletBCSpec(
...         location=lambda pt: np.isclose(pt[0], 1.0),
...         component='x',
...         value=0.1
...     )
... ])
>>>
>>> # Create DirichletBC from config
>>> bc = bc_config.create_bc(problem)

add​

def add(location: Callable[[np.ndarray], bool], component: Union[int, str],
        value: Union[float, Callable[[np.ndarray],
                                     float]]) -> 'DirichletBCConfig'

Add a boundary condition specification to the configuration.

This method allows for fluent-style chaining when building BC configurations.

Parameters​

  • location (Callable[[np.ndarray], bool]): Function that takes a point coordinate array and returns True if the point is on the boundary where this BC should be applied
  • component (Union[int, str]): Which component to constrain: - For scalar problems: 0 or 'all' - For vector problems: 0/'x', 1/'y', 2/'z', or 'all'
  • value (Union[float, Callable[[np.ndarray], float]]): The prescribed value, either a constant or a spatial function

Returns​

  • self (DirichletBCConfig): Returns self for method chaining, allowing: config.add(...).add(...)

Examples​

>>> config = DirichletBCConfig()
>>> config.add(lambda pt: np.isclose(pt[0], 0.0), 'all', 0.0)
>>> config.add(lambda pt: np.isclose(pt[0], 1.0), 'x', 0.1)

create_bc​

def create_bc(problem: 'Problem') -> 'DirichletBC'

Create a DirichletBC object from this configuration.

This method directly processes the BC specification without intermediate format conversion.

Parameters​

  • problem (Problem): The finite element problem instance containing mesh information and vector dimension specifications

Returns​

  • bc (DirichletBC): The compiled boundary condition object ready for use in solvers

Notes​

This method automatically detects the vector dimension from the problem and handles both single-variable and multi-variable problems.

dirichlet_bc_config​

def dirichlet_bc_config(*specs: DirichletBCSpec) -> DirichletBCConfig

Convenience function to create a DirichletBCConfig from multiple specs.

This function provides a concise way to create BC configurations without explicitly instantiating the DirichletBCConfig class.

Parameters​

*specs : DirichletBCSpec Variable number of boundary condition specifications

Returns​

  • config (DirichletBCConfig): The BC configuration containing all provided specifications

Examples​

>>> # Create BC config with multiple specifications
>>> config = dirichlet_bc_config(
...     DirichletBCSpec(left_boundary, 'all', 0.0),
...     DirichletBCSpec(right_boundary, 'x', 0.1),
...     DirichletBCSpec(top_boundary, 'y', lambda pt: 0.01 * pt[0])
... )
>>> bc = config.create_bc(problem)

See Also​

DirichletBCConfig : The main configuration class DirichletBCSpec : Individual BC specification