NEOPAX._geometry_models

Modular geometry models for NEOPAX. - Registry-based selection (like transport_flux_models.py) - JIT-compatible, differentiable (JAX-friendly)

Attributes

GEOMETRY_MODEL_REGISTRY

Classes

GeometryModelBase

Abstract base class for geometry models.

VmecBoozer

Abstract base class for geometry models.

Functions

register_geometry_model(→ None)

get_geometry_model(→ GeometryModelBase)

extend_grid_with_ghosts(grid)

Extend a 1D grid array with one ghost cell on each side (extrapolated).

extend_faces_with_ghosts(grid_half)

Extend a 1D face grid array for N+1 faces to N+3 faces (ghost faces).

Module Contents

NEOPAX._geometry_models.GEOMETRY_MODEL_REGISTRY: dict[str, Callable[Ellipsis, GeometryModelBase]]
NEOPAX._geometry_models.register_geometry_model(name: str, builder: Callable[Ellipsis, GeometryModelBase]) → None
NEOPAX._geometry_models.get_geometry_model(name: str, **kwargs) → GeometryModelBase
class NEOPAX._geometry_models.GeometryModelBase

Abstract base class for geometry models.

abstractmethod as_field()
class NEOPAX._geometry_models.VmecBoozer(n_r, vmec=None, booz=None, **kwargs)

Bases: GeometryModelBase

Abstract base class for geometry models.

n_r: int
a_b: float
Psia_value: float
rho_grid: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
rho_grid_half: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
r_grid: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
r_grid_half: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
full_grid_indices: jaxtyping.Int[jaxtyping.Array, VmecBoozer.n_r]
dr: float
Vprime: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
Vprime_half: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
overVprime: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
epsilon_t: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
B0: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
B_10: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
enlogation: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
iota: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
R0: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
B0prime: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
curvature: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
G_PS: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
sqrtg00_value: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
Bsqav: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
I_value: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
G_value: jaxtyping.Float[jaxtyping.Array, VmecBoozer.n_r]
NEOPAX._geometry_models.extend_grid_with_ghosts(grid)

Extend a 1D grid array with one ghost cell on each side (extrapolated).

NEOPAX._geometry_models.extend_faces_with_ghosts(grid_half)

Extend a 1D face grid array for N+1 faces to N+3 faces (ghost faces).