Data

PhiJAX keeps application data construction on the host and accelerator placement in the trainer. A DataModule owns immutable pools and sampling policy; it does not need to know which device the trainer selected.

DataModule lifecycle

prepare_data()
setup("fit" | "predict")
train_batch_source(...) | predict_batch_source()
Trainer prepares source and places batches
teardown("fit" | "predict")

prepare_data may generate or obtain shared artifacts. Build process-local NumPy pools in setup. prepare_stage(stage) runs these methods at most once until teardown_stage(stage) releases the stage. The Trainer prepares fitting before requesting optional model normalization statistics, then retrieves sources and handles teardown. The base input_statistics() returns None; override it to opt into input normalization. Return None from predict_batch_source() to skip prediction. Prediction sources must be finite and repeatable; training sources may sample indefinitely from the global optimizer step.

PhiDataModule [source]

PhiDataModule()

Bases: ABC

Define a Lightning-inspired host data lifecycle for PhiJAX applications.

Application subclasses construct immutable pools in :meth:setup and expose JAX-compatible sources through :meth:train_batch_source and :meth:predict_batch_source. The module is mutable host-side orchestration state; model, optimizer, balancer, and PRNG state used by compiled training remain explicit and functional.

Attributes:

Name	Type	Description
`pools`	`Mapping[str, HostPool]`	Immutable application pools populated by :meth:`setup`.
`prepared_stage`	`DataStage \| None`	Currently prepared lifecycle stage, or `None` outside a Trainer task.

Initialize an empty application pool mapping.

prepare_data [source]

prepare_data() -> None

Prepare external artifacts before application pools are constructed.

The base implementation performs no work. Subclasses may generate or download files but should not assign process-local runtime state here. Runtime pools belong in :meth:setup.

setup [source] `abstractmethod`

setup(stage: DataStage) -> None

Populate immutable application pools for one execution stage.

Parameters:

Name	Type	Description	Default
`stage`	`DataStage`	Requested `fit` or `predict` lifecycle stage.	required

prepare_stage [source]

prepare_stage(stage: DataStage) -> None

Prepare data and set up one stage exactly once until teardown.

Parameters:

Name	Type	Description	Default
`stage`	`DataStage`	Requested `fit` or `predict` lifecycle stage.	required

Raises:

Type	Description
`BaseException`	Re-raises preparation or setup errors after best-effort application teardown.

teardown_stage [source]

teardown_stage(stage: DataStage) -> None

Tear down one prepared stage and clear its lifecycle marker.

Parameters:

Name	Type	Description	Default
`stage`	`DataStage`	Completed `fit` or `predict` lifecycle stage.	required

train_batch_source [source] `abstractmethod`

train_batch_source(
    batch_keys: tuple[str, ...], key: Array
) -> TrainingBatchSource

Build a deterministic step-indexed training batch source.

Parameters:

Name	Type	Description	Default
`batch_keys`	`tuple[str, ...]`	Objective batch keys in stable declaration order.	required
`key`	`Array`	Explicit root sampling key.	required

Returns:

Type	Description
`TrainingBatchSource`	Unprepared step-indexed source that the Trainer places before sampling.

predict_batch_source [source]

predict_batch_source() -> PredictionBatchSource | None

Build an optional finite ordered host-backed prediction batch source.

Returns:

Type	Description
`PredictionBatchSource \| None`	Re-iterable source whose batches the Trainer places before prediction, or `None` when this DataModule does
`PredictionBatchSource \| None`	not provide prediction data.

prediction_pool [source]

prediction_pool() -> HostPool | None

Return the optional ordered host pool represented by prediction batches.

Returns:

Type	Description
`HostPool \| None`	Prediction pool carrying reconstruction metadata and host targets, or `None` when prediction is absent.

normalization_pools [source]

normalization_pools() -> Mapping[str, HostPool]

Select host pools used to derive shared model input statistics.

The default excludes the prediction source's pool so a dense evaluation grid does not dominate training-coordinate statistics. Applications may override this policy when one pool best represents the physical domain.

Returns:

Type	Description
`Mapping[str, HostPool]`	Non-empty training-oriented host-pool mapping.

Raises:

Type	Description
`ValueError`	If excluding the prediction pool leaves no pools.

input_statistics [source]

input_statistics() -> (
    tuple[NDArray[Any], NDArray[Any]] | None
)

Return optional model-input normalization statistics.

The base implementation disables input normalization. Applications can opt in by overriding this method and returning empirical statistics from :func:phijax.data.input_statistics or exact statistics for a continuous sampling distribution.

Returns:

Type	Description
`tuple[NDArray[Any], NDArray[Any]] \| None`	Per-coordinate mean and safely positive standard deviation arrays, or `None` to skip normalization.

teardown [source]

teardown(stage: DataStage) -> None

Release application resources after one lifecycle stage.

Parameters:

Name	Type	Description	Default
`stage`	`DataStage`	Completed `fit` or `predict` lifecycle stage.	required

Host pools

inputs is a rank-two NumPy array. Omit targets for an automatically created zero-width target array. aux and metadata default to immutable empty mappings. When omitted, reference_shape becomes (sample_count,) and flat_index becomes arange(sample_count). Array-valued aux fields share the sample axis.

HostPool [source] `dataclass`

HostPool(
    inputs: NDArray[Any],
    targets: NDArray[Any] | None = None,
    aux: Mapping[str, Any] | None = None,
    metadata: Mapping[str, Any] | None = None,
    reference_shape: tuple[int, ...] | None = None,
    flat_index: NDArray[Any] | None = None,
)

Represent one immutable host-side subset with stable row ordering.

Copy arrays, fill safe defaults, validate dimensions, and freeze storage.

Parameters:

Name	Type	Description	Default
`inputs`	`NDArray[Any]`	Rank-two coordinate array with shape `[samples, input_features]`.	required
`targets`	`NDArray[Any] \| None`	Optional rank-two target array. Omission creates a zero-width array.	`None`
`aux`	`Mapping[str, Any] \| None`	Optional sample-wise auxiliary arrays.	`None`
`metadata`	`Mapping[str, Any] \| None`	Optional structural or reconstruction metadata.	`None`
`reference_shape`	`tuple[int, ...] \| None`	Optional dense reconstruction shape. Omission uses `(sample_count,)`.	`None`
`flat_index`	`NDArray[Any] \| None`	Optional flattened reconstruction indices. Omission uses sequential row indices.	`None`

Raises:

Type	Description
`ValueError`	If `inputs` or `targets` is not rank two, or if any sample-wise array has a mismatched leading row count.

fields [source]

fields() -> Mapping[str, NDArray[Any]]

Collect dense sample-wise fields for one device transfer.

Returns:

Type	Description
`Mapping[str, NDArray[Any]]`	Read-only mapping containing `inputs`, `targets`, and NumPy-array entries from `aux`.

input_statistics [source]

input_statistics(
    pools: Mapping[str, HostPool],
) -> tuple[NDArray[Any], NDArray[Any]]

Compute shared coordinate normalization statistics across named pools.

Parameters:

Name	Type	Description	Default
`pools`	`Mapping[str, HostPool]`	Non-empty named host pools with a common input width.	required

Returns:

Type	Description
`tuple[NDArray[Any], NDArray[Any]]`	Per-coordinate mean and safely floored standard deviation as `float32` arrays.

Raises:

Type	Description
`ValueError`	If no pools are supplied or their input widths differ.

reconstruct_predictions [source]

reconstruct_predictions(
    predictions: NDArray[Any], pool: HostPool
) -> NDArray[Any]

Scatter flat predictions back to a pool's dense reference grid.

Parameters:

Name	Type	Description	Default
`predictions`	`NDArray[Any]`	Flat prediction array with shape `[samples, output_features]`.	required
`pool`	`HostPool`	Host pool carrying reference shape and flattened row indices.	required

Returns:

Type	Description
`NDArray[Any]`	Dense prediction array with shape `[*reference_shape, output_features]`.

Raises:

Type	Description
`ValueError`	If prediction rows, indices, or reference shape are inconsistent.

Prediction artifacts

Prediction artifacts can be saved and loaded without callbacks. The canonical NPZ stores physical arrays, output scales, schema information, and safe application metadata. Optional MATLAB files store the same physical values and may use application-specific field names.

save_prediction_artifact [source]

save_prediction_artifact(
    output_path: str | Path,
    predictions: object,
    pool: HostPool,
    *,
    save_mat: bool = False,
    mat_field_names: Mapping[str, str] | None = None,
) -> Path

Save a canonical NumPy prediction artifact and an optional MATLAB sidecar.

Model predictions and compatible pool targets are multiplied by application-provided output_scales before serialization. Missing scales resolve to ones. Both formats store physical values, the applied scales, schema version 2, and value_space="physical".

Parameters:

Name	Type	Description	Default
`output_path`	`str \| Path`	Destination compressed NumPy artifact path ending in `.npz`.	required
`predictions`	`object`	Flat normalized model predictions ordered like `pool.inputs`.	required
`pool`	`HostPool`	Host prediction pool supplying normalized targets, reconstruction metadata, and optional `artifact_metadata`.	required
`save_mat`	`bool`	Whether to also save a `.mat` file beside the NumPy artifact.	`False`
`mat_field_names`	`Mapping[str, str] \| None`	Optional mapping from generic prediction-array names to MATLAB variable names.	`None`

Returns:

Type	Description
`Path`	Absolute path to the saved `.npz` artifact.

Raises:

Type	Description
`KeyError`	If `mat_field_names` refers to an array that is not exported.
`ValueError`	If the destination suffix, predictions, targets, output scales, or MATLAB names are invalid.

load_prediction_artifact [source]

load_prediction_artifact(
    path: str | Path,
) -> PredictionArtifact

Load and validate a physical PhiJAX prediction artifact without pickle support.

Missing output_scales default to one for every output channel, leaving stored physical values unchanged.

Parameters:

Name	Type	Description	Default
`path`	`str \| Path`	Schema-version-2 `.npz` prediction artifact.	required

Returns:

Type	Description
`PredictionArtifact`	Immutable artifact with combined arrays, named output views, and separated metadata.

Raises:

Type	Description
`FileNotFoundError`	If `path` does not identify a file.
`KeyError`	If a required schema array is absent.
`ValueError`	If schema metadata, shapes, names, indices, scales, or values are inconsistent.

PredictionArtifact [source] `dataclass`

PredictionArtifact(
    prediction: NDArray[Any],
    target: NDArray[Any],
    flat_prediction: NDArray[Any],
    flat_target: NDArray[Any],
    inputs: NDArray[Any],
    dense_inputs: NDArray[Any],
    mask: NDArray[bool_],
    flat_index: NDArray[int64],
    reference_shape: tuple[int, ...],
    coordinate_names: tuple[str, ...],
    output_names: tuple[str, ...],
    output_scales: NDArray[Any],
    outputs: Mapping[str, NDArray[Any]],
    targets: Mapping[str, NDArray[Any]],
    metadata: Mapping[str, NDArray[Any]],
)

Expose one validated physical prediction artifact and named channel views.

Attributes:

Name	Type	Description
`prediction`	`NDArray[Any]`	Dense physical predictions with shape `[*reference_shape, outputs]`.
`target`	`NDArray[Any]`	Dense physical targets with zero or all output channels.
`flat_prediction`	`NDArray[Any]`	Physical predictions at the rows represented by `flat_index`.
`flat_target`	`NDArray[Any]`	Physical targets at the rows represented by `flat_index`.
`inputs`	`NDArray[Any]`	Flat model coordinates in the application's configured coordinate space.
`dense_inputs`	`NDArray[Any]`	Coordinates reconstructed on the reference grid.
`mask`	`NDArray[bool_]`	Boolean reference-grid mask selecting represented rows.
`flat_index`	`NDArray[int64]`	Flat reference-grid index for every represented row.
`reference_shape`	`tuple[int, ...]`	Dense reference-grid shape.
`coordinate_names`	`tuple[str, ...]`	Semantic name of every input coordinate.
`output_names`	`tuple[str, ...]`	Semantic name of every predicted output.
`output_scales`	`NDArray[Any]`	Physical multiplier applied to each normalized model output during serialization.
`outputs`	`Mapping[str, NDArray[Any]]`	Read-only mapping from each output name to its dense physical prediction view.
`targets`	`Mapping[str, NDArray[Any]]`	Read-only mapping from each output name to its dense physical target view, or an empty mapping when the prediction pool has no targets.
`metadata`	`Mapping[str, NDArray[Any]]`	Read-only mapping containing schema, structural, and application-provided metadata arrays.

to_matlab_prediction_arrays [source]

to_matlab_prediction_arrays(
    arrays: Mapping[str, Any],
    *,
    field_names: Mapping[str, str] | None = None,
) -> dict[str, Any]

Rename reconstructed prediction arrays for MATLAB export.

Parameters:

Name	Type	Description	Default
`arrays`	`Mapping[str, Any]`	Reconstructed prediction arrays and optional nested metadata keyed by generic PhiJAX names.	required
`field_names`	`Mapping[str, str] \| None`	Optional mapping from generic top-level or metadata names to MATLAB variable names.	`None`

Returns:

Type	Description
`dict[str, Any]`	Independent mapping suitable for :func:`scipy.io.savemat`.

Raises:

Type	Description
`KeyError`	If `field_names` contains a source name absent from `arrays`.
`TypeError`	If a configured source or destination name is not a string.
`ValueError`	If a destination is not a MATLAB identifier or multiple arrays use the same destination.

Samplers

Name	Source	Typical PINN use
`random_rows`	Finite aligned pool	Initial, boundary, or supervised samples
`uniform_domain`	Continuous rectangular bounds	Fresh PDE collocation coordinates
`space_time`	Fixed spatial mesh + fresh times	Time-dependent PDE collocation

prepare(device) moves persistent arrays once. sample(key, batch_size) remains explicitly keyed and returns a fixed leading dimension.

BatchSampler [source]

Bases: ABC

Define one immutable explicit-key batch sampling policy.

prepare [source]

prepare(device: JaxDevice) -> BatchSampler

Prepare persistent sampler state for one trainer-selected device.

Stateless custom samplers may inherit this implementation. Samplers owning arrays should override it and return an immutable copy whose persistent storage resides on device.

Parameters:

Name	Type	Description	Default
`device`	`JaxDevice`	Trainer Strategy's process-local root device.	required

Returns:

Type	Description
`BatchSampler`	Device-ready sampler.

sample [source] `abstractmethod`

sample(key: Array, batch_size: int) -> DevicePool

Draw one fixed-size batch.

Parameters:

Name	Type	Description	Default
`key`	`Array`	Explicit JAX PRNG key controlling this draw.	required
`batch_size`	`int`	Positive number of requested rows.	required

Returns:

Type	Description
`DevicePool`	Sample-wise device arrays sharing the requested leading dimension.

sample_all [source]

sample_all() -> DevicePool

Return the complete finite source when supported.

Returns:

Type	Description
`DevicePool`	Complete sample-wise device arrays.

Raises:

Type	Description
`ValueError`	If the sampler represents a continuous distribution without a finite source.

RandomRowSampler [source] `dataclass`

RandomRowSampler(
    pool: Mapping[str, Any],
    replace: bool = True,
    sort_axis: int | None = None,
)

Bases: BatchSampler

Sample aligned rows from one finite pool.

Attributes:

Name	Type	Description
`pool`	`Mapping[str, Any]`	Sample-wise host or device arrays with a shared leading dimension.
`replace`	`bool`	Whether randomly drawn indices may repeat.
`sort_axis`	`int \| None`	Optional input-coordinate column used to sort every returned field consistently.

__post_init__ [source]

__post_init__() -> None

Validate the finite source and optional input sort axis.

Raises:

Type	Description
`KeyError`	If the pool has no `inputs` field.
`ValueError`	If pool fields have inconsistent row counts or `sort_axis` is invalid.

prepare [source]

prepare(device: JaxDevice) -> RandomRowSampler

Place the finite candidate pool once on the trainer-selected device.

Parameters:

Name	Type	Description	Default
`device`	`JaxDevice`	Trainer Strategy's process-local root device.	required

Returns:

Type	Description
`RandomRowSampler`	Immutable sampler sharing policy values with a device-resident pool.

sample [source]

sample(key: Array, batch_size: int) -> DevicePool

Sample aligned finite-pool rows.

Parameters:

Name	Type	Description	Default
`key`	`Array`	Explicit JAX PRNG key controlling row indices.	required
`batch_size`	`int`	Positive number of requested rows.	required

Returns:

Type	Description
`DevicePool`	Random aligned pool fields, optionally sorted by one input coordinate.

sample_all [source]

sample_all() -> DevicePool

Return every finite-pool row exactly once.

Returns:

Type	Description
`DevicePool`	Complete aligned pool fields, optionally sorted by one input coordinate.

UniformDomainSampler [source] `dataclass`

UniformDomainSampler(
    bounds: Any,
    templates: Mapping[str, Any],
    dtype: DTypeLike = jnp.float32,
    sort_axis: int | None = None,
)

Bases: BatchSampler

Generate fresh coordinates uniformly inside a rectangular domain.

Attributes:

Name	Type	Description
`bounds`	`Any`	Array with shape `[input_features, 2]` containing increasing lower and upper bounds.
`templates`	`Mapping[str, Any]`	Non-input field values broadcast to every generated coordinate.
`dtype`	`DTypeLike`	Floating dtype used for generated coordinates.
`sort_axis`	`int \| None`	Optional coordinate column used to sort each generated batch.

__post_init__ [source]

__post_init__() -> None

Validate domain bounds, templates, dtype, and sort policy.

Raises:

Type	Description
`ValueError`	If bounds are not finite increasing intervals, a template is named `inputs`, or sorting is invalid.

prepare [source]

prepare(device: JaxDevice) -> UniformDomainSampler

Place domain bounds and field templates once on the trainer-selected device.

Parameters:

Name	Type	Description	Default
`device`	`JaxDevice`	Trainer Strategy's process-local root device.	required

Returns:

Type	Description
`UniformDomainSampler`	Immutable sampler with device-resident generation metadata.

sample [source]

sample(key: Array, batch_size: int) -> DevicePool

Generate one fresh fixed-size domain batch.

Parameters:

Name	Type	Description	Default
`key`	`Array`	Explicit JAX PRNG key controlling coordinate generation.	required
`batch_size`	`int`	Positive number of requested coordinates.	required

Returns:

Type	Description
`DevicePool`	Generated `inputs` and broadcast template fields.

Raises:

Type	Description
`ValueError`	If `batch_size` is not positive.

SpaceTimeSampler [source] `dataclass`

SpaceTimeSampler(
    pool: Mapping[str, Any],
    time_bounds: tuple[float, float],
    time_axis: int = 0,
    replace: bool = True,
    sort_axis: int | None = 0,
)

Bases: BatchSampler

Pair fresh uniform times with rows sampled from a fixed spatial mesh.

Attributes:

Name	Type	Description
`pool`	`Mapping[str, Any]`	Finite host or device pool whose `inputs` contain a replaceable temporal coordinate column.
`time_bounds`	`tuple[float, float]`	Increasing lower and upper time bounds.
`time_axis`	`int`	Input-coordinate column replaced with fresh uniform times.
`replace`	`bool`	Whether sampled spatial rows may repeat.
`sort_axis`	`int \| None`	Optional input-coordinate column used to sort every returned field.

__post_init__ [source]

__post_init__() -> None

Validate the finite mesh and temporal settings.

Raises:

Type	Description
`ValueError`	If bounds, the temporal axis, target semantics, or sorting are invalid.

prepare [source]

prepare(device: JaxDevice) -> SpaceTimeSampler

Place the finite spatial mesh once on the trainer-selected device.

Parameters:

Name	Type	Description	Default
`device`	`JaxDevice`	Trainer Strategy's process-local root device.	required

Returns:

Type	Description
`SpaceTimeSampler`	Immutable sampler with device-resident mesh fields.

sample [source]

sample(key: Array, batch_size: int) -> DevicePool

Sample spatial rows and replace their temporal coordinate.

Parameters:

Name	Type	Description	Default
`key`	`Array`	Explicit JAX PRNG key split across mesh and temporal draws.	required
`batch_size`	`int`	Positive number of requested rows.	required

Returns:

Type	Description
`DevicePool`	Aligned pool fields with fresh uniformly generated time coordinates.

create_sampler [source]

create_sampler(
    name: SamplerName | str,
    pool: HostPool,
    *,
    options: Mapping[str, Any] | None = None,
) -> BatchSampler

Create a built-in sampler from an explicit stable name.

Sampler selection intentionally avoids Hydra _target_ values. The surrounding DataModule resolves the short names random_rows, uniform_domain, and space_time while retaining typed constructors internally.

Parameters:

Name	Type	Description	Default
`name`	`SamplerName \| str`	Built-in sampler name.	required
`pool`	`HostPool`	Immutable host pool supplying finite rows, field templates, or inferred coordinate bounds.	required
`options`	`Mapping[str, Any] \| None`	Sampler-specific options such as `replace`, `sort_axis`, `bounds`, `constant_fields`, `time_bounds`, or `time_axis`.	`None`

Returns:

Type	Description
`BatchSampler`	Configured immutable batch sampler.

Raises:

Type	Description
`TypeError`	If an option has an unsupported container or scalar type.
`ValueError`	If the name, bounds, templates, or sampler options are invalid.

sample_pool [source]

sample_pool(
    key: Array,
    pool: Mapping[str, Any],
    batch_size: int,
    *,
    replace: bool = True,
) -> DevicePool

Sample a fixed-shape device batch using an explicit PRNG key.

Parameters:

Name	Type	Description	Default
`key`	`Array`	JAX PRNG key controlling index selection.	required
`pool`	`Mapping[str, Any]`	Host or device field mapping with a shared leading sample dimension.	required
`batch_size`	`int`	Positive number of rows in the returned batch.	required
`replace`	`bool`	Whether indices may repeat.	`True`

Returns:

Type	Description
`DevicePool`	Field mapping with leading dimension `batch_size`.

Raises:

Type	Description
`ValueError`	If `batch_size` is not positive, the pool is empty, or sampling without replacement requests more rows than the pool contains.
`KeyError`	If `pool` does not contain `inputs`.

Batch sources

NamedBatchSource folds its root key with the global step, then splits keys in sampler declaration order. A restored checkpoint therefore continues the same sampling sequence. The "all" batch-size policy works only with finite row sources.

NamedBatchSource.from_pools() is the concise path for aligned finite datasets. It creates ordered RandomRowSampler instances and validates requested pool names and batch-size policies.

Define one batch-size policy per pool in the DataModule, then reuse the objective's inferred order:

self.batch_sizes = {"initial": "all", "boundary": 128, "pde": 4096}


def train_batch_source(self, batch_keys, key):
    return NamedBatchSource.from_pools(self.pools, self.batch_sizes, key, names=batch_keys)

An integer samples that many rows on every optimizer step. The "all" policy uses every row and is available only for finite pools.

ChunkedPredictionSource pads its final host batch and supplies a Boolean mask; prediction removes padded rows before output assembly.

TrainingBatchSource [source]

Bases: Protocol

Define deterministic step-indexed training and diagnostic batch delivery.

prepare [source]

prepare(device: JaxDevice) -> Self

Prepare persistent source state for one trainer-selected device.

Parameters:

Name	Type	Description	Default
`device`	`JaxDevice`	Trainer Strategy's process-local root device.	required

Returns:

Type	Description
`Self`	Device-ready source preserving sampling policy and deterministic key behavior.

call [source]

__call__(step: int) -> NamedBatches

Produce one named training batch for a global optimizer step.

Parameters:

Name	Type	Description	Default
`step`	`int`	Non-negative global optimizer step, including a restored checkpoint offset.	required

Returns:

Type	Description
`NamedBatches`	Named fixed-structure training batches.

sample [source]

sample(
    key: Array,
    batch_sizes: Mapping[str, BatchSize] | None = None,
) -> NamedBatches

Produce explicitly keyed batches with optional size overrides.

Parameters:

Name	Type	Description	Default
`key`	`Array`	Explicit root sampling key.	required
`batch_sizes`	`Mapping[str, BatchSize] \| None`	Optional per-stream batch-size overrides used by diagnostics such as exact NTK balancing.	`None`

Returns:

Type	Description
`NamedBatches`	Named fixed-structure batches.

PredictionBatchSource [source]

Bases: Protocol

Define finite ordered prediction batches and their reconstruction pool.

pool `property`

pool: HostPool

Return the immutable host pool used to reconstruct flat predictions.

iter [source]

__iter__() -> Iterator[Mapping[str, Any]]

Iterate over fixed-size host-backed padded prediction batches.

len [source]

__len__() -> int

Return the finite number of prediction batches.

NamedBatchSource [source] `dataclass`

NamedBatchSource(
    samplers: Mapping[str, BatchSampler],
    batch_sizes: Mapping[str, BatchSize],
    key: Array,
)

Compose named immutable samplers into deterministic step-indexed batches.

Attributes:

Name	Type	Description
`samplers`	`Mapping[str, BatchSampler]`	Named sampling policies aligned with objective batch keys.
`batch_sizes`	`Mapping[str, BatchSize]`	Per-sampler positive integer or `all` policies.
`key`	`Array`	Explicit root key folded with each global optimizer step.

from_pools [source] `classmethod`

from_pools(
    pools: Mapping[str, HostPool],
    batch_sizes: Mapping[str, BatchSize],
    key: Array,
    *,
    names: tuple[str, ...] | None = None,
    replace: bool = True,
    sort_axes: Mapping[str, int | None] | None = None,
) -> Self

Construct aligned finite-row samplers from named host pools.

Parameters:

Name	Type	Description	Default
`pools`	`Mapping[str, HostPool]`	Finite host pools keyed by application batch name.	required
`batch_sizes`	`Mapping[str, BatchSize]`	Sampling policies keyed by the same names.	required
`key`	`Array`	Explicit root key folded with each optimizer step.	required
`names`	`tuple[str, ...] \| None`	Optional stable subset and ordering. The batch-size declaration order is used when omitted.	`None`
`replace`	`bool`	Whether sampled rows may repeat.	`True`
`sort_axes`	`Mapping[str, int \| None] \| None`	Optional input-coordinate sort axis per selected pool.	`None`

Returns:

Type	Description
`Self`	Deterministic named source containing one :class:`RandomRowSampler` per selected pool.

Raises:

Type	Description
`KeyError`	If a selected name is absent from `pools` or `batch_sizes`.
`ValueError`	If selected names are empty or duplicated.

__post_init__ [source]

__post_init__() -> None

Validate sampler names and batch-size policies.

Raises:

Type	Description
`KeyError`	If a sampler has no batch-size policy.
`ValueError`	If no samplers are supplied or a policy is invalid for its sampler.

prepare [source]

prepare(device: JaxDevice) -> Self

Place persistent sampler arrays and the root key on a trainer-selected device.

Preparing once retains device-side row selection and generated-coordinate sampling without requiring an application DataModule to know which accelerator the Trainer selected.

Parameters:

Name	Type	Description	Default
`device`	`JaxDevice`	Trainer Strategy's process-local root device.	required

Returns:

Type	Description
`Self`	Immutable device-ready source.

call [source]

__call__(step: int) -> dict[str, DevicePool]

Sample every named policy for one global optimizer step.

Parameters:

Name	Type	Description	Default
`step`	`int`	Non-negative global optimizer step, including any restored checkpoint offset.	required

Returns:

Type	Description
`dict[str, DevicePool]`	Named fixed-structure training batches.

Raises:

Type	Description
`ValueError`	If `step` is negative.

sample [source]

sample(
    key: Array,
    batch_sizes: Mapping[str, BatchSize] | None = None,
) -> dict[str, DevicePool]

Sample named policies with an explicit key and optional size overrides.

Parameters:

Name	Type	Description	Default
`key`	`Array`	Explicit root key split in sampler declaration order.	required
`batch_sizes`	`Mapping[str, BatchSize] \| None`	Optional replacement policies, used for fixed balancer diagnostic batches.	`None`

Returns:

Type	Description
`dict[str, DevicePool]`	Named batches aligned with sampler declaration order.

Raises:

Type	Description
`KeyError`	If a sampler has no selected batch-size policy.
`ValueError`	If an override is invalid or requests `all` from a generated sampler.

ChunkedPredictionSource [source] `dataclass`

ChunkedPredictionSource(pool: HostPool, batch_size: int)

Lazily construct fixed-size prediction batches from an immutable host pool.

The source retains the complete dataset on the host. The Trainer places only the currently yielded padded batch on its selected device, reducing peak device memory while preserving deterministic row order.

Attributes:

Name	Type	Description
`pool`	`HostPool`	Immutable host prediction pool carrying dense reconstruction metadata.
`batch_size`	`int`	Positive fixed number of rows yielded per batch.

__post_init__ [source]

__post_init__() -> None

Validate prediction row and fixed batch sizes.

Raises:

Type	Description
`ValueError`	If pool rows or `batch_size` are invalid.

len [source]

__len__() -> int

Return the finite number of padded prediction batches.

Returns:

Type	Description
`int`	At least one batch, including for an empty prediction pool.

iter [source]

__iter__() -> Iterator[dict[str, np.ndarray]]

Yield host-sliced and padded prediction batches.

Yields:

Type	Description
`dict[str, ndarray]`	Field mappings with a fixed leading dimension and Boolean `mask` identifying valid rows.

Array IO and declarative pools

load_arrays supports NPZ, MATLAB, and HDF5 files. SciPy reads classic MATLAB files. mat73 restores MATLAB's logical array order from v7.3 HDF5 files, while generic HDF5 input keeps its stored order. get_array accepts slash-delimited nested keys. Use declarative pool construction for direct field mapping. Subclass PhiDataModule for custom generation or sampling.

load_arrays [source]

load_arrays(
    path: str | Path, *, file_format: ArrayFormat = "auto"
) -> dict[str, Any]

Load an NPZ, MATLAB, or HDF5 artifact into host memory.

File access and decoding remain entirely NumPy-based so this function can run without initializing a JAX backend. HDF5 groups are represented by nested dictionaries, while all HDF5 datasets are materialized before the file is closed. Automatic mode uses the file suffix and recognizes HDF5 content when the suffix is unknown.

Parameters:

Name	Type	Description	Default
`path`	`str \| Path`	Source `.npz`, `.mat`, `.h5`, `.hdf`, or `.hdf5` path.	required
`file_format`	`ArrayFormat`	Explicit file format or `"auto"` for detection.	`'auto'`

Returns:

Type	Description
`dict[str, Any]`	Mapping of source field names to NumPy arrays or nested field mappings.

Raises:

Type	Description
`FileNotFoundError`	If `path` is not an existing file.
`ValueError`	If the requested format is invalid or automatic detection fails.

get_array [source]

get_array(
    arrays: Mapping[str, Any], key: str
) -> NDArray[Any]

Read one array using a top-level key or slash-delimited nested path.

A direct mapping key takes precedence over path traversal so NPZ fields containing / remain addressable.

Parameters:

Name	Type	Description	Default
`arrays`	`Mapping[str, Any]`	Mapping returned by :func:`load_arrays` or an equivalent array tree.	required
`key`	`str`	Top-level field name or nested path such as `"results/velocity"`.	required

Returns:

Type	Description
`NDArray[Any]`	Requested value represented as a NumPy array.

Raises:

Type	Description
`KeyError`	If the key or one of its path components is absent.
`TypeError`	If the key resolves to a group rather than an array value.
`ValueError`	If `key` is empty.

build_array_pools [source]

build_array_pools(
    *,
    source: Mapping[str, Any],
    coordinates: Mapping[str, Any],
    pools: Mapping[str, Any],
    fields: Mapping[str, Any] | None = None,
    dtype: str = "float32",
) -> dict[str, HostPool]

Build named host pools declaratively from an NPZ, MATLAB, or HDF5 artifact.

Coordinate arrays define reusable grid axes and, unless explicitly overridden, their uniform-sampling bounds. File-backed pools form Cartesian grids, optionally sliced by a scalar index or half-open [start, stop] range. Uniform pools sample directly from the bounds of their configured input coordinates. Losses remain independent and share a sampled pool by referencing the same objective batch_key.

Parameters:

Name	Type	Description	Default
`source`	`Mapping[str, Any]`	Mapping with `path`, optional `file_format`, and optional callable `prepare` invoked before loading.	required
`coordinates`	`Mapping[str, Any]`	Named one-dimensional coordinate specifications. Each value is a source key or a mapping with `key` and optional explicit `bounds`.	required
`pools`	`Mapping[str, Any]`	Named grid or uniform-sampling pool specifications.	required
`fields`	`Mapping[str, Any] \| None`	Optional named source or constant field specifications. Source fields contain `key`, `axes`, and an optional `feature_axis`; constant fields contain exactly one of `key` and `value` and no sample axes.	`None`
`dtype`	`str`	Floating NumPy dtype used for coordinates, fields, and generated samples.	`'float32'`

Returns:

Type	Description
`dict[str, HostPool]`	Non-empty mapping of configured names to immutable :class:`HostPool` objects.

Raises:

Type	Description
`TypeError`	If a major configuration section has the wrong container type.
`ValueError`	If the dtype, coordinates, fields, pool references, selections, or sampling settings are invalid.

Host transforms

Reusable transforms operate on NumPy arrays while building host data. They derive statistics from finite values and preserve non-finite entries. They also accept fixed statistics, so prediction can reuse values derived from training data. Keep application-specific physical scaling in the application DataModule.

log_compress [source]

log_compress(
    values: ArrayLike,
    *,
    percentile: float = 99.0,
    dynamic_range: float = 40.0,
) -> NDArray[np.floating[Any]]

Compress power-like values into a normalized relative-decibel range.

Finite values above the percentile-derived reference are clipped, non-positive finite values are mapped to the floating-point epsilon before applying the logarithm, and the retained decibel interval is mapped to [0, 1]. Non-finite entries are preserved.

Parameters:

Name	Type	Description	Default
`values`	`ArrayLike`	Power-like values to compress.	required
`percentile`	`float`	Finite upper-reference percentile in `(0, 100]`.	`99.0`
`dynamic_range`	`float`	Positive relative-decibel range retained below the reference.	`40.0`

Returns:

Type	Description
`NDArray[floating[Any]]`	Floating array with finite values in `[0, 1]` and the same shape as `values`.

Raises:

Type	Description
`ValueError`	If `percentile` or `dynamic_range` is invalid.

minmax_scale [source]

minmax_scale(
    values: ArrayLike,
    *,
    minimum: ArrayLike | None = None,
    maximum: ArrayLike | None = None,
    feature_range: Sequence[float] = (0.0, 1.0),
) -> NDArray[np.floating[Any]]

Scale finite values from a supplied or data-derived interval.

Explicit minimum and maximum arrays may be reused across datasets and follow NumPy broadcasting. When either statistic is omitted, it is derived globally from finite entries in values. A constant source interval maps to the lower endpoint of feature_range, and non-finite entries are preserved.

Parameters:

Name	Type	Description	Default
`values`	`ArrayLike`	Values to scale.	required
`minimum`	`ArrayLike \| None`	Optional finite lower statistic. Defaults to the finite global minimum of `values`.	`None`
`maximum`	`ArrayLike \| None`	Optional finite upper statistic. Defaults to the finite global maximum of `values`.	`None`
`feature_range`	`Sequence[float]`	Two finite values defining the strictly increasing output interval.	`(0.0, 1.0)`

Returns:

Type	Description
`NDArray[floating[Any]]`	Floating array with the same shape as `values`.

Raises:

Type	Description
`ValueError`	If the output interval or explicit source statistics are invalid.

standardize [source]

standardize(
    values: ArrayLike,
    *,
    mean: ArrayLike | None = None,
    std: ArrayLike | None = None,
    eps: float = 1e-12,
) -> NDArray[np.floating[Any]]

Center and scale finite values with reusable statistics.

Explicit mean and std arrays follow NumPy broadcasting. Missing statistics are derived globally from finite entries in values. A standard deviation no greater than eps produces zeros for finite entries, and non-finite entries are preserved.

Parameters:

Name	Type	Description	Default
`values`	`ArrayLike`	Values to standardize.	required
`mean`	`ArrayLike \| None`	Optional finite centering statistic. Defaults to the finite global mean of `values`.	`None`
`std`	`ArrayLike \| None`	Optional finite non-negative scale statistic. Defaults to the finite global standard deviation.	`None`
`eps`	`float`	Non-negative threshold below which the scale is treated as zero.	`1e-12`

Returns:

Type	Description
`NDArray[floating[Any]]`	Floating array with the same shape as `values`.

Raises:

Type	Description
`ValueError`	If `eps`, `mean`, or `std` is invalid.

scale_by_max [source]

scale_by_max(
    values: ArrayLike,
    *,
    percentile: float | None = None,
    clip: bool | None = None,
) -> NDArray[np.floating[Any]]

Scale finite values by their maximum or an upper percentile.

Percentile scaling clips by default so outliers do not produce values above one. Exact-maximum scaling does not clip by default. Non-finite entries are excluded from statistic estimation and preserved in the returned array.

Parameters:

Name	Type	Description	Default
`values`	`ArrayLike`	Values to scale.	required
`percentile`	`float \| None`	Optional finite upper percentile in `(0, 100]` used instead of the maximum.	`None`
`clip`	`bool \| None`	Whether to clip finite values to the divisor before scaling. Defaults to whether `percentile` is set.	`None`

Returns:

Type	Description
`NDArray[floating[Any]]`	Floating array with the same shape as `values`. A zero divisor maps finite entries to zero.

Raises:

Type	Description
`ValueError`	If `percentile` is outside `(0, 100]`.

See Building an application DataModule for a complete implementation.

Data

DataModule lifecycle

PhiDataModule [source]

prepare_data [source]

setup [source] abstractmethod

prepare_stage [source]

teardown_stage [source]

train_batch_source [source] abstractmethod

predict_batch_source [source]

prediction_pool [source]

normalization_pools [source]

input_statistics [source]

teardown [source]

Host pools

HostPool [source] dataclass

fields [source]

input_statistics [source]

reconstruct_predictions [source]

Prediction artifacts

save_prediction_artifact [source]

load_prediction_artifact [source]

PredictionArtifact [source] dataclass

to_matlab_prediction_arrays [source]

Samplers

BatchSampler [source]

prepare [source]

sample [source] abstractmethod

sample_all [source]

RandomRowSampler [source] dataclass

__post_init__ [source]

prepare [source]

sample [source]

sample_all [source]

UniformDomainSampler [source] dataclass

__post_init__ [source]

prepare [source]

sample [source]

SpaceTimeSampler [source] dataclass

__post_init__ [source]

prepare [source]

sample [source]

create_sampler [source]

sample_pool [source]

Batch sources

TrainingBatchSource [source]

prepare [source]

__call__ [source]

sample [source]

PredictionBatchSource [source]

pool property

__iter__ [source]

__len__ [source]

NamedBatchSource [source] dataclass

from_pools [source] classmethod

__post_init__ [source]

prepare [source]

__call__ [source]

sample [source]

ChunkedPredictionSource [source] dataclass

__post_init__ [source]

__len__ [source]

__iter__ [source]

Array IO and declarative pools

load_arrays [source]

get_array [source]

build_array_pools [source]

Host transforms

log_compress [source]

minmax_scale [source]

standardize [source]

scale_by_max [source]

setup [source] `abstractmethod`

train_batch_source [source] `abstractmethod`

HostPool [source] `dataclass`

PredictionArtifact [source] `dataclass`

sample [source] `abstractmethod`

RandomRowSampler [source] `dataclass`

UniformDomainSampler [source] `dataclass`

SpaceTimeSampler [source] `dataclass`

call [source]

pool `property`

iter [source]

len [source]

NamedBatchSource [source] `dataclass`

from_pools [source] `classmethod`

call [source]

ChunkedPredictionSource [source] `dataclass`

len [source]

iter [source]