pde_opt.numerics
Numerical methods for PDEs.
- class pde_opt.numerics.Domain(points: Tuple[int, ...], box: Tuple[Tuple[float, float], ...], units: str, geometry: Shape | None = None)[source]
Sets up a simulation domain for the model.
The following information is stored in a Domain: – points[i] is the number of collocation points in the i’th dimension – dx[i] is the spacing between each collocation point in the i’th dimension – box[i] is the bounds of the simulation box in the i’th dimension – units are the length units these. values are stored in
- points: Tuple[int, ...]
- box: Tuple[Tuple[float, float], ...]
- units: str
- class pde_opt.numerics.Shape(binary: jax.Array, refine_factor: float | None = None, refine_edge: float | None = None, dx: Tuple[float, float] | None = (1.0, 1.0), smooth_epsilon: float = 1.0, smooth_curvature: float = 0.0, smooth_dt: float = 0.1, smooth_tf: float = 1.0)[source]
Sets up a geometry/shape for solving PDE on with smoothed boundary method.
The user creates a shape by providing a binary representation and an optional smoothing parameter.
- binary: jax.Array
- refine_factor: float | None = None
- refine_edge: float | None = None
- dx: Tuple[float, float] | None = (1.0, 1.0)
- smooth_epsilon: float = 1.0
- smooth_curvature: float = 0.0
- smooth_dt: float = 0.1
- smooth_tf: float = 1.0
- refine_binary_mask() jax.Array[source]
Refine the binary mask by upsampling and adding edge padding.
- Parameters:
refine_factor – Factor by which to upsample the binary mask (e.g., 2.0 for 2x upsampling)
refine_edge – Percentage to increase the edges (e.g., 0.5 for 50% increase)
- Returns:
Refined binary mask with upsampling and edge padding applied
- smooth_shape() jax.Array[source]
Smooths the shape using the Allen-Cahn equation with curvature minimization.
- laplacian_from_mask(periodic: bool = False)[source]
Unnormalized graph Laplacian (4-neighbour) from a 0/1 mask. Nodes are entries where mask==1. Two nodes connect if they are up/down/left/right neighbours and both are 1.
- Returns:
(n_nodes, n_nodes) CSR Laplacian ids: (H, W) array, node index in [0, n_nodes) or -1 if not a node
- Return type:
L
- get_shape_modes(N: int | None = None)[source]
Get the first N eigenvectors of the graph Laplacian of the binary mask.
Creates a graph where nodes are the 1-valued pixels, with edges between adjacent pixels (left, right, top, bottom neighbors).
- Parameters:
N – Number of eigenvectors to return. If None, returns all eigenvectors.
downsampling_factor – If provided, downsample binary by this factor before computing modes, then upsample results back to original size. This can significantly reduce memory usage and computation time for large binary masks.
- Returns:
Array of shape (num_nodes, N) containing the first N eigenvectors
- __init__(binary: jax.Array, refine_factor: float | None = None, refine_edge: float | None = None, dx: Tuple[float, float] | None = (1.0, 1.0), smooth_epsilon: float = 1.0, smooth_curvature: float = 0.0, smooth_dt: float = 0.1, smooth_tf: float = 1.0) None
- class pde_opt.numerics.AllenCahn2DPeriodic(domain: Domain, kappa: float, mu: Callable | equinox.Module, R: Callable | equinox.Module, derivs: str = 'fd')[source]
Allen-Cahn equation in 2D with periodic boundary conditions.
The Allen-Cahn equation describes phase transitions and interface dynamics. The equation is:
\[\frac{\partial u}{\partial t} = -R(u) \mu\]where u is the concentration, R(u) is the reaction term, μ is the chemical potential, and κ is a parameter (the gradient energy coefficient). The chemical potential is given by:
\[\mu = \mu_h(u) - \kappa \nabla^2 u\]- kappa: float
Gradient energy coefficient
- mu: Callable | equinox.Module
Function for the chemical potential
- R: Callable | equinox.Module
Function for the reaction term
- derivs: str = 'fd'
Type of derivative computation
- class pde_opt.numerics.AllenCahn2DSmoothedBoundary(domain: Domain, kappa: float, f: Callable | equinox.Module, mu: Callable | equinox.Module, R: Callable | equinox.Module, theta: Callable | equinox.Module, derivs: str = 'fd')[source]
Allen-Cahn equation with smoothed boundary method for arbitrary geometries.
This class implements the Allen-Cahn equation using the smoothed boundary method, which allows for complex domain geometries through a smooth level-set function ψ.
The equation is:
\[\frac{\partial u}{\partial t} = -R(u) \mu\]where the chemical potential includes boundary effects:
\[\mu = \mu_h(u) - \frac{\kappa}{\psi} \nabla \cdot (\psi \nabla u) - \sqrt{\kappa} \frac{|\nabla \psi|}{\psi} \sqrt{2f} \cos(\theta)\]- kappa: float
Gradient energy coefficient
- f: Callable | equinox.Module
Function for the free energy density
- mu: Callable | equinox.Module
Function for the chemical potential
- R: Callable | equinox.Module
Function for the reaction term
- theta: Callable | equinox.Module
Function for the contact angle
- derivs: str = 'fd'
Type of derivative computation
- class pde_opt.numerics.CahnHilliard2DPeriodic(domain: Domain, kappa: float, mu: Callable | equinox.Module, D: Callable | equinox.Module, derivs: str = 'fd')[source]
Cahn–Hilliard equation in 2D with periodic boundary conditions.
The Cahn-Hilliard equation describes phase separation and coarsening dynamics. The equation is:
\[\frac{\partial u}{\partial t} = \nabla \cdot (D(u) \nabla \mu)\]where u is the concentration, D(u) is the mobility, and μ is the chemical potential. The chemical potential is given by:
\[\mu = \mu_h(u) - \kappa \nabla^2 u\]- kappa: float
Gradient energy coefficient
- mu: Callable | equinox.Module
Function for the chemical potential
- D: Callable | equinox.Module
Function for the mobility
- derivs: str = 'fd'
Type of derivative computation
- fft = None
- ifft = None
- fourier_symbol = None
- class pde_opt.numerics.CahnHilliard3DPeriodic(domain: Domain, kappa: float, mu: Callable | equinox.Module, D: Callable | equinox.Module, derivs: str = 'fd')[source]
Cahn–Hilliard equation in 3D with periodic boundary conditions.
The Cahn-Hilliard equation describes phase separation and coarsening dynamics. The equation is:
\[\frac{\partial u}{\partial t} = \nabla \cdot (D(u) \nabla \mu)\]where u is the concentration, D(u) is the mobility, and μ is the chemical potential. The chemical potential is given by:
\[\mu = \mu_h(u) - \kappa \nabla^2 u\]- kappa: float
Gradient energy coefficient
- mu: Callable | equinox.Module
Function for the chemical potential
- D: Callable | equinox.Module
Function for the mobility
- derivs: str = 'fd'
Type of derivative computation
- fft = None
- ifft = None
- fourier_symbol = None
- class pde_opt.numerics.CahnHilliard2DSmoothedBoundary(domain: Domain, kappa: float, f: Callable | equinox.Module, mu: Callable | equinox.Module, D: Callable | equinox.Module, theta: Callable | equinox.Module, flux: Callable | equinox.Module, derivs: str = 'fd')[source]
Cahn–Hilliard equation with smoothed boundary method for arbitrary geometries.
This class implements the Cahn-Hilliard equation using the smoothed boundary method, which allows for complex domain geometries through a smooth level-set function ψ.
The equation is:
\[\frac{\partial u}{\partial t} = \frac{1}{\psi} \nabla \cdot (\psi D(u) \nabla \mu) + \frac{|\nabla \psi|}{\psi} J_n\]where the chemical potential includes boundary effects:
\[\mu = \mu_h(u) - \frac{\kappa}{\psi} \nabla \cdot (\psi \nabla u) - \sqrt{\kappa} \frac{|\nabla \psi|}{\psi} \sqrt{2f} \cos(\theta)\]- __init__(domain: Domain, kappa: float, f: Callable | equinox.Module, mu: Callable | equinox.Module, D: Callable | equinox.Module, theta: Callable | equinox.Module, flux: Callable | equinox.Module, derivs: str = 'fd') None
- kappa: float
Gradient energy coefficient
- f: Callable | equinox.Module
Function for the free energy density
- mu: Callable | equinox.Module
Function for the chemical potential
- D: Callable | equinox.Module
Function for the mobility
- theta: Callable | equinox.Module
Function for the contact angle
- flux: Callable | equinox.Module
Function for the normal flux
- derivs: str = 'fd'
Type of derivative computation
- class pde_opt.numerics.GPE2DTSControl(domain: Domain, k: float, e: float, lights: Callable, trap_factor: float = 1.0)[source]
Gross-Pitaevskii equation in 2D with time-splitting and control.
The Gross-Pitaevskii equation describes the dynamics of Bose-Einstein condensates. The equation is:
\[i\hbar \frac{\partial \psi}{\partial t} = \left[-\frac{\hbar^2}{2m}\nabla^2 + V(\mathbf{r}, t) + g|\psi|^2\right]\psi\]where ψ is the wave function, V is the external potential, and g is the interaction strength. The external potential includes a harmonic trap and control field:
\[V(\mathbf{r}, t) = \frac{1}{2}m\omega^2\left[(1+\epsilon)x^2 + (1-\epsilon)y^2\right] + V_{control}(\mathbf{r}, t)\]- k: float
Interaction strength parameter
- e: float
Trap ellipticity parameter
- lights: Callable
Function for the control field
- trap_factor: float = 1.0
Scaling factor for the harmonic trap
- fft = None
- ifft = None
- A_term = None
- dx = None
- class pde_opt.numerics.GPE2DTSRot(domain: Domain, k: float, e: float, omega: float)[source]
Gross-Pitaevskii equation in 2D with time-splitting and rotation.
The Gross-Pitaevskii equation describes the dynamics of Bose-Einstein condensates. The equation is:
\[i\hbar \frac{\partial \psi}{\partial t} = \left[-\frac{\hbar^2}{2m}\nabla^2 + V(\mathbf{r}) + g|\psi|^2 - \Omega L_z\right]\psi\]where ψ is the wave function, V is the external potential, g is the interaction strength, and Ω is the rotation frequency with L_z being the angular momentum operator. The external potential includes a harmonic trap:
\[V(\mathbf{r}) = \frac{1}{2}m\omega^2\left[(1+\epsilon)x^2 + (1-\epsilon)y^2\right]\]- k: float
Interaction strength parameter
- e: float
Trap ellipticity parameter
- omega: float
Rotation frequency
- class pde_opt.numerics.PeriodicCNN(*args: Any, **kwargs: Any)[source]
Stack of periodic conv blocks; final conv returns requested channels.
Translation-equivariant on a torus so long as stride=1 and only pointwise nonlinearities are used. Accepts (C,H,W) or (B,C,H,W); returns same spatial size.
- __init__(in_channels: int, hidden_channels: Sequence[int] = (32, 64, 64), out_channels: int | None = None, kernel_size: int = 3, act: Callable[[jax.Array], jax.Array] = jax.nn.gelu, *, key: jax.Array)[source]
- layers: Tuple[equinox.Module, ...]
- class pde_opt.numerics.LegendrePolynomialExpansion(*args: Any, **kwargs: Any)[source]
Legendre polynomial expansion.
- params: jax.Array
- max_degree: int
- class pde_opt.numerics.DiffusionLegendrePolynomials(*args: Any, **kwargs: Any)[source]
Diffusion Legendre polynomials.
Uses exp to ensure positivity.
- expansion: LegendrePolynomialExpansion
- class pde_opt.numerics.ChemicalPotentialLegendrePolynomials(*args: Any, **kwargs: Any)[source]
Chemical potential Legendre polynomials.
- expansion: LegendrePolynomialExpansion
- prior_fn: Callable
- class pde_opt.numerics.Mixer2d(*args: Any, **kwargs: Any)[source]
- __init__(img_size, patch_size, hidden_size, mix_patch_size, mix_hidden_size, num_blocks, *, key)[source]
- conv_in: equinox.nn.Conv2d
- conv_out: equinox.nn.ConvTranspose2d
- blocks: list
- norm: equinox.nn.LayerNorm
- class pde_opt.numerics.SemiImplicitFourierSpectral(*args: Any, **kwargs: Any)[source]
Semi-implicit Fourier spectral method.
This solver implements a semi-implicit Fourier spectral method for phase-field simulations with variable mobility.
- Required Equation Attributes:
fourier_symbol: Fourier space representation of the highest order differential operator. fft: Forward Fourier transform function. ifft: Inverse Fourier transform function.
- Parameters:
A (float) – Constant for splitting the mobility term.
References
Zhu, Jingzhi, et al. “Coarsening kinetics from a variable-mobility Cahn-Hilliard equation: Application of a semi-implicit Fourier spectral method.” Physical Review E 60.4 (1999): 3564.
- required_equation_attrs = ['fourier_symbol', 'fft', 'ifft']
- A: float
- fourier_symbol: jax.Array
- fft: Callable
- ifft: Callable
- class pde_opt.numerics.StrangSplitting(*args: Any, **kwargs: Any)[source]
Strang splitting method for time-dependent PDEs with separable operators.
References
Bao, Weizhu, and Yongyong Cai. “Mathematical theory and numerical methods for Bose-Einstein condensation.” arXiv preprint arXiv:1212.5341 (2012).
- required_equation_attrs = ['A_term', 'dx', 'fft', 'ifft']
- A_term: jax.Array
- dx: float
- fft: Callable
- ifft: Callable
- time_scale: float
Modules
This module contains the Domain class, which is used to set up a simulation domain for the model. |
|
PDE equation classes. |
|
Function representations for PDEs. |
|
This module contains the Shape class, which is used to set up a geometry/shape for solving PDE on with smoothed boundary method. |
|
Custom numerical solvers for partial differential equations. |
|
This module contains symbolic equation classes for the PDEs for testing purposes. |
|