Skip to content

beamax.geometry

Domain and sensor geometry primitives used by forward and inverse solvers.

Key Objects

  • Domain: grid size, spacing, wave-speed model, CFL/time grid helpers, and mesh generation.
  • Sensor: sensor masks/positions tied to a domain.

API Reference

beamax.geometry

Domain

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(x) or constant (default 1500.0).
density Callable | float | None

Density rho(x) or constant.
alpha_coeff Callable | float | None

Absorption prefactor alpha0(x).
lam float

Absorption coefficient for GB ODEs (default 0.0).
alpha_power Callable | float | None

Absorption exponent y(x) in alpha = alpha0 * f**y.
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.
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))

N * dx per axis.
xmax: Float[Array, ''] property

Max extent (Euclidean norm of grid_size).

Returns:

Type Description
float
k_max: Float[Array, ''] property

Max wavenumber given sampling.

Returns:

Type Description
float

π * min(1/dx).
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) with max(c).
compute_min_speed() -> Float[Array, '']

Minimum sound speed on grid.

Returns:

Type Description
ndarray

Scalar (0-D array) with min(c).
generate_meshgrid() -> Tuple[list[Float[Array, '*N']], list[Int[Array, '*N']]]

Spatial and Fourier meshgrids.

Returns:

Type Description
(spatial_meshgrid, fourier_meshgrid)

Each a tuple of length d with arrays shaped N[i] per axis. - Spatial coordinates: 0..(N[i]-1) * dx[i]. - Fourier indices: -N[i]//2 .. N[i]//2 - 1.
compute_max_freq() -> Float[Array, '']

Max frequency allowed by grid / CFL proxy.

Returns:

Type Description
ndarray

Scalar ≈ max(c) / (2 * min(dx)).
generate_time_domain() -> Float[Array, ' Nt']

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

Returns:

Type Description
(ndarray, shape(Nt))

With dt = cfl * min(dx) / max(c) and tmax = ||grid_size|| / min(c), 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 whatever you provided 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