beamax.geometry

Computational domains and sensor geometries.

beamax.geometry

Domain(N: Tuple[int, ...], dx: Tuple[float, ...], periodic: Tuple[bool, ...], cfl: float = 0.3, c: Param = 1500.0, density: Param = 1.0, alpha_coeff: Param = None, lam: float = 0.0, alpha_power: Param = None)

Bases: Module

Axis-aligned rectangular domain with physical spacing and medium fields.

Attributes:

Name	Type	Description
N	Tuple[int, ...]	Grid shape per axis.
dx	Tuple[float, ...]	Physical spacing per axis (same length as N).
periodic	Tuple[bool, ...]	Per-axis periodicity flags.
cfl	float	CFL number used to pick dt in generate_time_domain (default 0.3).
c	Callable | float	Speed of sound \(c(\mathbf{x})\) or constant (default 1500.0).
density	Callable | float | None	Density \(\rho(\mathbf{x})\) or constant.
alpha_coeff	Callable | float | None	Absorption prefactor \(\alpha_0(\mathbf{x})\).
lam	float	Absorption coefficient for GB ODEs (default 0.0).
alpha_power	Callable | float | None	Absorption exponent \(y(\mathbf{x})\) in \(\alpha(f, \mathbf{x}) = \alpha_0(\mathbf{x})f^{y(\mathbf{x})}\).

Notes

• All callable fields are evaluated lazily on accessors (pure, JAX-friendly).
• N, periodic are static; others can be traced.
• Derived arrays (grid, sound_speed_array, density_array, alpha_coeff_array, alpha_power_array) and ndim are available as properties.

Construct and validate a computational domain.

c_fn: FieldFn property

Callable sound speed, wrapping scalars into a JAX-friendly function.

Returns:

Type	Description
Callable[[ndarray], ndarray]	If c is already callable, returns it. Otherwise returns a function that broadcasts the constant c over the leading shape of x.

grid_size: Float[Array, ' d'] property

Physical size of the domain per axis.

Returns:

Type	Description
(ndarray, shape(d))	Component \(i\) is \(L_i = N_i\,\Delta x_i\).

xmax: Float[Array, ''] property

Maximum extent \(x_{\max} = \lVert\mathbf{L}\rVert_2\), where \(\mathbf{L}\) is grid_size.

Returns:

Type	Description
float

k_max: Float[Array, ''] property

Largest axis-aligned Nyquist wavenumber.

Returns:

Type	Description
float	\(k_{\max} = \pi / \min_i \Delta x_i\).

ndim: int property

Number of spatial dimensions.

Returns:

Type	Description
int

grid: Float[Array, '*N d'] property

Stacked spatial coordinates.

Returns:

Type	Description
(ndarray, shape(*N, d))	Meshgrid stacked along last axis.

sound_speed_array: Float[Array, '*N'] property

Speed of sound evaluated on grid.

Returns:

Type	Description
(ndarray, shape(*N))

density_array: Optional[Float[Array, '*N']] property

Density evaluated on grid.

Returns:

Type	Description
ndarray | None

alpha_coeff_array: Optional[Float[Array, '*N']] property

Absorption prefactor evaluated on grid.

Returns:

Type	Description
ndarray | None

alpha_power_array: Optional[Float[Array, '*N']] property

Absorption exponent evaluated on grid.

Returns:

Type	Description
ndarray | None

compute_max_speed() -> Float[Array, '']

Maximum sound speed on grid.

Returns:

Type	Description
ndarray	Scalar (0-D array) equal to \(\max_{\mathbf{x}}c(\mathbf{x})\).

compute_min_speed() -> Float[Array, '']

Minimum sound speed on grid.

Returns:

Type	Description
ndarray	Scalar (0-D array) equal to \(\min_{\mathbf{x}}c(\mathbf{x})\).

generate_meshgrid() -> Tuple[tuple[Float[Array, '*N'], ...], tuple[Int[Array, '*N'], ...]]

Spatial and Fourier meshgrids.

Returns:

Type	Description
(spatial_meshgrid, fourier_meshgrid)	Each is a tuple of length d whose arrays have shape N. Spatial coordinates: \(x_{i,n} = n\,\Delta x_i\) for \(n = 0, \ldots, N_i - 1\). Fourier indices: \(-\lceil N_i / 2 \rceil, \ldots, \lfloor N_i / 2 \rfloor - 1\).

compute_max_freq() -> Float[Array, '']

Max frequency allowed by grid / CFL proxy.

Returns:

Type	Description
ndarray	Scalar \[ f_{\max} \approx \frac{\max_{\mathbf{x}}c(\mathbf{x})}{2\min_i \Delta x_i}. \]

generate_time_domain() -> Float[Array, ' Nt']

CFL-based uniform time grid covering one diameter-crossing.

Returns:

Type	Description
(ndarray, shape(Nt))	With \[ \Delta t = \mathrm{CFL}\, \frac{\min_i \Delta x_i}{\max_{\mathbf{x}}c(\mathbf{x})}, \qquad t_{\max} = \frac{\lVert\mathbf{L}\rVert_2}{\min_{\mathbf{x}}c(\mathbf{x})}, \] where \(\mathbf{L}\) is grid_size, this returns arange(0, tmax, dt).

Sensor(domain: Domain, positions: Optional[Float[Array, 'Ns d']] = None, binary_mask: Optional[Num[Array, '*N']] = None)

Bases: Module

Sampling geometry for receivers or sources.

Construct with exactly one of positions or binary_mask. The other representation is derived deterministically from the supplied form and is available via the :attr:positions / :attr:binary_mask properties.

Parameters:

Name	Type	Description	Default
domain	Domain	The spatial domain the sensor lives on.	required
positions	(ndarray, shape(Ns, d))	Cartesian positions in physical units.	None
binary_mask	(ndarray, shape(*N))	Mask with 1 at sensor voxels, 0 elsewhere.	None

Raises:

Type	Description
ValueError	If neither or both of positions and binary_mask are provided.

Notes

positions are converted to the nearest integer grid index when the mask is derived. Sub-pixel positions are quantised to grid voxels.

Construct from positions or mask.

Parameters:

Name	Type	Description	Default
domain	Domain	Spatial domain the sensor belongs to.	required
positions	(ndarray, shape(Ns, d))	Physical sensor positions.	None
binary_mask	(ndarray, shape(*N))	Binary sensor mask.	None

Raises:

Type	Description
ValueError	If neither or both of positions and binary_mask are provided.

positions: Float[Array, 'Ns d'] property

Sensor positions (physical units).

Returns:

Type	Description
(ndarray, shape(Ns, d))

binary_mask: Num[Array, '*N'] property

Sensor mask aligned to domain.N.

Returns:

Type	Description
jnp.ndarray, shape (*N,), dtype=int