jaxpi.models

PINN base classes, model/optimizer factories, and the sharded training step.

TrainState [source]

Bases: TrainState

TrainState(step: int | jax.Array, apply_fn: collections.abc.Callable, params: flax.core.frozen_dict.FrozenDict[str, typing.Any], tx: optax._src.base.GradientTransformation, opt_state: Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, bool, int, float, complex, Iterable[ForwardRef('ArrayTree')], Mapping[Any, ForwardRef('ArrayTree')]], loss_weights: Dict, pts_weights: Dict, momentum: float, pts_momentum: float, prev_params: Any = None)

TrainState.apply_loss_weights() [source]

TrainState.apply_loss_weights(self, loss_weights, **kwargs)

TrainState.apply_pts_weights() [source]

TrainState.apply_pts_weights(self, pts_weights, **kwargs)

TrainState.replace()

TrainState.replace(self, **updates)

Returns a new object replacing the specified fields with new values.

create_arch() [source]

create_arch(config)

create_lr_schedule() [source]

create_lr_schedule(config)

create_optimizer() [source]

create_optimizer(config, lr)

create_model() [source]

create_model(config, model_cls, *model_args, params=None, **model_kwargs)

Build a model from its config: lr schedule, optimizer, arch, state.

model_args / model_kwargs are forwarded to the model constructor (problem-specific arguments such as initial conditions or physical parameters). params warm-starts the train state (transfer learning).

create_train_state() [source]

create_train_state(
    config,
    tx,
    arch,
    params=None,
    train_state_cls=<class 'jaxpi.models.TrainState'>,
)

PINN [source]

PINN.neural_net() [source]

PINN.neural_net(self, params, *args)

PINN.r_net() [source]

PINN.r_net(self, params, *args)

PINN.sol_net() [source]

PINN.sol_net(self, params, *args)

neural_net's outputs as a dict keyed by variables.

With residuals keyed by variable name, the pseudo-time damping pairs each residual with its own solution component by key — immune to any ordering (JAX sorts dict keys when flattening pytrees). No override needed: declaring variables is sufficient.

PINN.losses() [source]

PINN.losses(self, params, state, batch)

PINN.compute_pts_weights() [source]

PINN.compute_pts_weights(self, state, init_state, batch)

PINN.compute_loss_weights() [source]

PINN.compute_loss_weights(self, state, batch)

Balance losses based on the gradient norms of each loss.

PINN.loss()

PINN.loss(self, params, state, batch)

PINN.create_step_fn() [source]

PINN.create_step_fn(self)

PINN.create_update_loss_weights_fn() [source]

PINN.create_update_loss_weights_fn(self)

PINN.create_compute_raw_residual_losses_fn() [source]

PINN.create_compute_raw_residual_losses_fn(self)

Sharded unweighted residual losses (no pseudo-time, no causal) — used by evaluators, so that full multi-GPU batches fit in memory.

PINN.create_compute_grad_norms_fn() [source]

PINN.create_compute_grad_norms_fn(self)

Sharded per-term gradient norms — used by evaluators.

PINN.create_update_pts_weights_fn() [source]

PINN.create_update_pts_weights_fn(self)

ForwardIVP [source]

Bases: PINN

ForwardIVP.create_compute_causal_weights_fn() [source]

ForwardIVP.create_compute_causal_weights_fn(self)

ForwardIVP.compute_residual_losses() [source]

ForwardIVP.compute_residual_losses(
    self,
    params,
    state,
    batch,
    pseudo_time=False,
    causal=False,
)

ForwardBVP [source]

Bases: PINN

ForwardBVP.compute_residual_losses() [source]

ForwardBVP.compute_residual_losses(self, params, state, batch, pseudo_time=False)