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() [source]

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.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_update_pts_weights_fn() [source]

PINN.create_update_pts_weights_fn(self)

ForwardIVP [source]

Bases: PINN

ForwardIVP.compute_causal_weights()

ForwardIVP.compute_causal_weights(self, state, batch)

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)