fol.deep_neural_networks

The fol.deep_neural_networks module provides a comprehensive collection of deep learning architectures and operator-learning frameworks designed for learning mappings between function spaces, with a strong focus on physics-informed modeling, parametric systems, and scientific machine learning.

This module implements state-of-the-art neural operators such as DeepONets and Fourier Neural Operators (FNOs), together with Conditional Neural Fields and meta-learning strategies. The provided models enable the approximation of complex solution operators and continuous fields arising in partial differential equations (PDEs), multiphysics simulations, and high-dimensional parametric problems, where classical surrogate models are often insufficient.

Key capabilities of this module include:

• Learning nonlinear operators mapping input functions, boundary conditions, or parameters to solution fields.

• Supporting data-driven, physics-informed, and hybrid training paradigms.

• Handling explicit and implicit parametric operator learning formulations.

• Neural operators based on DeepONets and Fourier Neural Operators for learning mappings between infinite-dimensional spaces.

• Neural field representations using coordinate-based MLPs, including SIREN-style sinusoidal networks for high-frequency signal and geometry representation.

• Conditional neural fields that model families of functions conditioned on parameters, latent variables, or context embeddings.

• Autoencoding architectures for learning compact latent representations of solution manifolds and parametric fields.

• Meta-learning and latent-variable methods for fast adaptation across tasks, parameters, and operating regimes.

• Efficient spectral representations and convolutions for scalable, high-dimensional problems.

The module is designed to be extensible and modular, enabling seamless integration into scientific computing pipelines while leveraging modern deep learning frameworks.

Deep network base class

Authors: Reza Najian Asl, RezaNajian Date: April, 2024 License: FOL/LICENSE

class fol.deep_neural_networks.deep_network.DeepNetwork(name, loss_function, flax_neural_network, optax_optimizer)

Bases: ABC

Base abstract class for deep learning models.

This class provides a common training/testing workflow for Flax/NNX neural networks optimized with Optax. It centralizes model initialization, JIT-compiled train/test steps, convergence checks, plotting of training curves, and Orbax checkpointing.

Subclasses must implement ComputeBatchLossValue() to define how a batch loss is computed (including any auxiliary loss terms and metrics) and Finalize() to run any end-of-training procedures.

The training loop performs:

• Optional restoration of a previously saved model state.

• Optional data/model sharding across devices.

• Batched training with randomized batch indices.

• Periodic evaluation on a test set.

• Convergence checks based on configured criteria.

• Optional checkpointing (least-loss, interval, and final state).

• Optional plotting of loss histories.

Parameters:

• name (str) – Name identifier for the model instance. Used for logging and to organize checkpoint folders.

• loss_function (Loss) – Loss function object that defines the optimization objective and provides any required initialization.

• flax_neural_network (nnx.Module) – Flax/NNX module defining the neural network architecture.

• optax_optimizer (GradientTransformation) – Optax optimizer transformation used to update the model parameters.

name

Model name used for identification and checkpointing.

Type:

str

loss_function

Loss function used to compute objective values.

Type:

Loss

flax_neural_network

Flax/NNX neural network module.

Type:

nnx.Module

optax_optimizer

Optax optimizer transformation.

Type:

GradientTransformation

initialized

Indicates whether Initialize() has been called.

Type:

bool

checkpointer

Orbax checkpointer used for saving/restoring model states.

Type:

ocp.StandardCheckpointer

nnx_optimizer

NNX optimizer wrapper holding parameters and optimizer state.

Type:

nnx.Optimizer

Notes

The subclass implementation of ComputeBatchLossValue() must return a metrics dictionary containing the key "total_loss". This value is used for logging, convergence checks, plotting, and checkpointing.

CheckConvergence(train_history_dict, convergence_settings)

Determine whether training has converged based on loss history.

Convergence is checked using: - Absolute threshold on the most recent value of the configured criterion, - Relative change between the last two criterion values, - Or reaching the configured maximum number of epochs.

Parameters:

• train_history_dict (dict) – Dictionary of training histories. Must contain the key specified by convergence_settings["convergence_criterion"] whose value is a list of scalar history values.

• convergence_settings (dict) – Dictionary containing: - "convergence_criterion" (str): key in train_history_dict - "absolute_error" (float): absolute threshold - "relative_error" (float): relative change threshold - "num_epochs" (int): maximum epochs

Returns:

True if the convergence conditions are met, otherwise False.

Return type:

bool

Raises:

KeyError – If the configured convergence criterion is missing from train_history_dict.

abstract ComputeBatchLossValue(batch_set, nn_model)

Compute the loss for a single batch and return loss metrics.

Subclasses must implement this method. It is used by both training and testing steps. Implementations are expected to perform a forward pass of nn_model on the batch input, evaluate the configured loss function and any auxiliary loss terms, and return a scalar batch loss together with a metrics dictionary.

The returned metrics dictionary must include the key "total_loss" because the training loop uses it for logging, convergence checks, plotting, and checkpoint decisions.

Parameters:

• batch_set (Tuple[jax.numpy.ndarray, jax.numpy.ndarray]) – Batch tuple (inputs, targets). Exact shapes depend on the problem setup.

• nn_model (nnx.Module) – The Flax/NNX model to evaluate.

Returns:

A tuple (loss_value, metrics_dict) where loss_value is the scalar loss used for differentiation and optimization, and metrics_dict contains batch-level metrics including the required key "total_loss".

Return type:

Tuple[jax.numpy.ndarray, dict]

Raises:

KeyError – If the returned metrics dictionary does not contain "total_loss".

abstract Finalize()

Finalize the model after training.

Subclasses implement this hook to perform any final steps that must occur once at the end of training (for example, releasing resources, computing final diagnostics, exporting artifacts, or post-processing).

Returns:

None

GetName()

Return the model name.

Returns:

The name identifier of this model instance.

Return type:

str

GetState()

Return the state tuple consumed by TrainStep() and TestStep().

Returns:

A tuple (self.flax_neural_network, self.nnx_optimizer).

Return type:

tuple[nnx.Module, nnx.Optimizer]

Initialize(reinitialize=False)

Initialize the loss function, checkpointer, and NNX optimizer wrapper.

This method prepares the model for training and testing. It initializes self.loss_function (if needed or if reinitialize=True), creates an Orbax StandardCheckpointer, and constructs an nnx.Optimizer that couples the Flax/NNX model with the Optax optimizer transformation.

Parameters:

reinitialize (bool, optional) – If True, force reinitialization of the loss function even if it was previously initialized. Default is False.

Returns:

None

PlotHistoryDict(plot_settings, train_history_dict, test_history_dict)

Plot and save training/test history curves.

This method creates a semilog-y plot for selected metrics from train_history_dict and test_history_dict and saves the figure as training_history.png in plot_settings["save_directory"].

Parameters:

• plot_settings (dict) – Dictionary containing plot configuration. Expected entries include plot_frequency, plot_list, save_directory, and test_frequency.

• train_history_dict (dict) – Training history mapping metric names to lists of recorded values.

• test_history_dict (dict) – Test history mapping metric names to lists of recorded values.

Returns:

None

Raises:

KeyError – If required entries are missing from plot_settings.

RestoreState(restore_state_directory)

Restore the model state from a checkpoint directory.

This loads a previously saved NNX state using Orbax and updates the in-memory Flax/NNX model parameters accordingly.

Parameters:

restore_state_directory (str) – Directory containing the saved Orbax checkpoint state.

Returns:

None

Raises:

• FileNotFoundError – If the checkpoint directory does not exist (may be raised by the checkpointer backend).

• ValueError – If the restored state is incompatible with the current model structure.

SaveCheckPoint(check_point_type, checkpoint_state_dir)

Save the current model state to a checkpoint directory.

This writes the current NNX model state (parameters and related state) using Orbax to the provided directory and forces the write.

Parameters:

• check_point_type (str) – Human-readable label for logging (e.g., "train", "test", "final", or "interval <epoch>").

• checkpoint_state_dir (str) – Directory where the checkpoint will be saved.

Returns:

None

Raises:

OSError – If the directory cannot be created or written to.

TestStep(state, data)

Execute one JIT-compiled evaluation (test) step.

This step calls ComputeBatchLossValue() without updating parameters and returns the scalar "total_loss" for logging.

Parameters:

• state (tuple[nnx.Module, nnx.Optimizer]) – State tuple (model, optimizer). The optimizer is unused but kept for interface consistency.

• data (Tuple[jax.numpy.ndarray, jax.numpy.ndarray]) – Batch tuple (inputs, targets).

Returns:

Scalar total loss for the batch (metrics["total_loss"]).

Return type:

jax.numpy.ndarray

Train(train_set, test_set=(Array([], shape=(0,), dtype=float32), Array([], shape=(0,), dtype=float32)), test_frequency=100, batch_size=100, convergence_settings={}, plot_settings={}, restore_nnx_state_settings={}, train_checkpoint_settings={}, test_checkpoint_settings={}, save_nnx_state_settings={}, data_model_sharding_settings={}, working_directory='.')

Train the model using mini-batch optimization and optional evaluation.

This method orchestrates the full training loop: - Applies default settings and merges user settings dictionaries, - Optionally restores a saved model state, - Optionally shards data/model across devices, - Runs multiple epochs of randomized mini-batch updates, - Periodically evaluates on test_set, - Checks convergence criteria, - Optionally plots and saves loss histories, - Optionally saves checkpoints (least-loss, interval, and/or final state).

Parameters:

• train_set (Tuple[jax.numpy.ndarray, jax.numpy.ndarray]) – Training dataset tuple (inputs, targets).

• test_set (Tuple[jax.numpy.ndarray, jax.numpy.ndarray], optional) – Test dataset tuple (inputs, targets). If empty arrays are provided, testing is skipped. Default is empty arrays.

• test_frequency (int, optional) – Evaluate test loss every test_frequency epochs. Default is 100.

• batch_size (int, optional) – Requested batch size. May be adjusted to evenly divide the training set size for batching/parallelization. Default is 100.

• convergence_settings (dict, optional) – Convergence configuration (epochs, criterion name, tolerances). Missing keys are filled from defaults.

• plot_settings (dict, optional) – Plot configuration (metrics list, plot frequency, save paths). Missing keys are filled from defaults.

• restore_nnx_state_settings (dict, optional) – Restore configuration. If restore=True, restores model state from state_directory.

• train_checkpoint_settings (dict, optional) – Checkpoint configuration for tracking least training loss.

• test_checkpoint_settings (dict, optional) – Checkpoint configuration for tracking least test loss.

• save_nnx_state_settings (dict, optional) – Final and interval checkpoint configuration.

• data_model_sharding_settings (dict, optional) – Sharding configuration including device partitioning.

• working_directory (str, optional) – Base directory used to store plots and checkpoints. Default is ".".

Returns:

None

Raises:

ValueError – If provided settings dictionaries are missing required keys after default merging (implementation-dependent).

TrainStep(state, data)

Execute one JIT-compiled training step.

This step calls ComputeBatchLossValue() with gradients enabled, applies the parameter update using the internal nnx.Optimizer, and returns the scalar "total_loss" for logging.

Parameters:

• state (tuple[nnx.Module, nnx.Optimizer]) – Training state as (model, optimizer).

• data (Tuple[jax.numpy.ndarray, jax.numpy.ndarray]) – Batch tuple (inputs, targets).

Returns:

Scalar total loss for the batch (metrics["total_loss"]).

Return type:

jax.numpy.ndarray

Explicit parametric operator learning

This submodule implements explicit parametric operator learning on discretized fields, where a fixed-dimensional parametric input space (for example control variables or Fourier coefficients) is mapped directly to a fixed-dimensional discretized field such as temperature or displacement.

The learning is unsupervised or physics-informed: no direct target fields are required. Instead, predicted fields are evaluated using physics-based loss functions (for example weighted residual- or energy-based formulations), with boundary conditions applied explicitly at the field level.

Authors: Reza Najian Asl, RezaNajian Date: April, 2024 License: FOL/LICENSE

class fol.deep_neural_networks.explicit_parametric_operator_learning.ExplicitParametricOperatorLearning(name, control, loss_function, flax_neural_network, optax_optimizer)

Bases: DeepNetwork

Explicit parametric operator learning on discretized fields.

This class implements explicit parametric operator learning where both the input parameter space and the output field space are discretized and have fixed shape. The input typically represents a low-dimensional parametric space, such as control variables of a Fourier parametrization, while the output corresponds to a high-dimensional discretized field, for example a temperature or displacement field.

The neural network explicitly represents the mapping from parametric inputs to field unknowns. Learning is unsupervised or physics-informed: no direct target fields are required. Instead, training is driven by physical loss functions evaluated on the predicted fields, such as residual-based or energy-based formulations.

The Control object maps parametric inputs to controlled variables, and the Loss object evaluates the predicted fields after Dirichlet boundary conditions are applied. Consistency between network dimensions, control variables, and loss unknowns is enforced during initialization.

Parameters:

• name (str) – Name identifier for the model instance, used for logging and checkpointing.

• control (Control) – Control object defining the parametric input space and how raw parameters are mapped to controlled variables.

• loss_function (Loss) – Physics-based loss function used to evaluate predicted fields. The loss operates on full discretized fields and does not require supervised target data.

• flax_neural_network (nnx.Module) – Flax/NNX neural network mapping parametric inputs to unknown field degrees of freedom. The network must expose in_features and out_features attributes.

• optax_optimizer (GradientTransformation) – Optax optimizer transformation used to construct the optimizer for training.

Raises:

RuntimeError – If the neural network input or output dimensions are inconsistent with the control variable size or the number of unknowns defined by the loss function.

ComputeBatchLossValue(batch, nn_model)

Compute the batch loss for explicit parametric operator learning.

This method evaluates the loss in an unsupervised or physics-informed setting. The batch is provided as a tuple for consistency with the DeepNetwork base class interface, but only the first entry is used. The second entry of the batch tuple is expected to be None and does not represent supervised targets.

The computation proceeds by mapping the batch of parametric inputs to controlled variables using the associated Control object, then predicting the unknown degrees of freedom with the neural network. The predicted unknowns are inserted into full discretized field vectors, and the loss function is evaluated directly on these fields using physical constraints or residual-based objectives.

The returned metrics dictionary always includes the key "total_loss", which is required by the training loop for logging, convergence checks, plotting, and checkpointing.

Parameters:

• batch (Tuple[jax.numpy.ndarray, None]) – Batch tuple (batch_X, None) where batch_X contains the parametric input samples. The second entry is unused and present only to maintain a consistent interface with the base class.

• nn_model (nnx.Module) – Neural network used to infer the unknown degrees of freedom from the parametric inputs.

Returns:

A tuple (batch_loss, metrics_dict) where batch_loss is the scalar loss aggregated over the batch, and metrics_dict contains loss statistics including the mandatory key "total_loss".

Return type:

Tuple[jax.numpy.ndarray, dict]

Raises:

None – This method assumes consistency between the control, loss function, and network dimensions, which is enforced during initialization.

ComputeBatchPredictions(batch_X, nn_model)

Compute network predictions for a batch of parametric inputs.

This helper applies the provided neural network to a batch of parametric inputs batch_X and returns the corresponding batch of predicted unknown DOFs. It does not insert Dirichlet values or build full field vectors; that is handled at the loss level.

Parameters:

• batch_X (jax.numpy.ndarray) – Batch of parametric inputs (for example control variables entering the operator). The leading dimension is the batch size, and the second dimension must match nn_model.in_features.

• nn_model (nnx.Module) – Neural network that maps parametric inputs to unknown DOFs.

Returns:

Batch of predicted unknown DOF vectors, one per row in batch_X.

Return type:

jax.numpy.ndarray

Raises:

None – This method is a thin wrapper around the network call and does not introduce additional failure modes beyond those of nn_model itself.

Finalize()

Finalize the model after training.

Subclasses implement this hook to perform any final steps that must occur once at the end of training (for example, releasing resources, computing final diagnostics, exporting artifacts, or post-processing).

Returns:

None

Initialize(reinitialize=False)

Initialize model, loss, control, optimizer, and consistency checks.

This method first runs the base-class initialization, which prepares the loss function, constructs the Orbax checkpointer, and builds the nnx.Optimizer around the provided network and Optax transformation. It then ensures that the associated control object is initialized and finally checks that the neural network input and output dimensions are consistent with the control and loss function sizes.

Concretely, the following consistency checks are performed:

• flax_neural_network.in_features is compared to control.GetNumberOfVariables() to ensure that the parametric input dimension matches the control space.

• flax_neural_network.out_features is compared to loss_function.GetNumberOfUnknowns() to ensure that the network output dimension matches the number of unknown field DOFs.

Parameters:

reinitialize (bool, optional) – If True, force reinitialization of all components, even if they have already been initialized. Default is False.

Returns:

None

Raises:

RuntimeError – If the network does not expose in_features or out_features, or if these dimensions are inconsistent with the control or loss function sizes. The underlying implementation uses fol_error to signal these conditions.

Predict(batch_X)

Perform inference for a batch of parametric inputs and apply Dirichlet boundary conditions.

This method runs inference by evaluating the trained neural network on a batch of parametric inputs to predict the unknown degrees of freedom. The predicted unknowns are then embedded into full discretized field vectors, and Dirichlet boundary conditions are applied by setting the prescribed boundary values at the corresponding DOF indices for every sample in the batch.

The result is a batch of full field predictions that satisfy the imposed boundary conditions and are suitable for post-processing or evaluation.

Parameters:

batch_X (jax.numpy.ndarray) – Batch of parametric inputs used for inference. The leading dimension corresponds to the batch size, and the feature dimension must match the number of control variables.

Returns:

Batch of full discretized field vectors obtained by inference, with Dirichlet boundary conditions applied consistently across the batch.

Return type:

jax.numpy.ndarray

Raises:

None –

Implicit parametric operator learning

This submodule implements implicit parametric operator learning using coordinate-based neural fields. A fixed-dimensional parametric input (for example control variables or parameterization features such as Fourier coefficients) conditions a neural field represented by a coordinate-based MLP (the synthesizer), with conditioning provided by a modulator network according to the coupling modes implemented by fol.deep_neural_networks.nns.HyperNetwork.

The learning is unsupervised or physics-informed: predicted discretized fields are evaluated using physics-based loss functionals (for example residual- or energy-based formulations), and Dirichlet boundary conditions are enforced by explicitly inserting prescribed values across the batch.

Although training is typically performed on a fixed FE mesh, the coordinate-based synthesizer enables multi-resolution inference (and, in principle, multi-resolution training) by evaluating the conditioned neural field on alternative coordinate sets.

Authors: Reza Najian Asl, RezaNajian Date: October, 2024 License: FOL/LICENSE

class fol.deep_neural_networks.implicit_parametric_operator_learning.ImplicitParametricOperatorLearning(name, control, loss_function, flax_neural_network, optax_optimizer)

Bases: DeepNetwork

Implicit parametric operator learning via hypernetwork-conditioned neural fields.

This class implements an implicit (coordinate-based) parametric operator-learning workflow in which a neural field is represented by a coordinate-based MLP (the synthesizer network) and is conditioned by a modulator network. The synthesizer performs the neural-field task by mapping spatial coordinates to discretized field values, while the modulator encodes parametric information from a fixed-dimensional input space and conditions the synthesizer according to the coupling modes supported by fol.deep_neural_networks.nns.HyperNetwork.

Conceptually, the model represents a conditional neural field. The synthesizer evaluates a coordinate-based representation of the solution field, and the modulator injects parametric context (for example control variables or feature points of a parameterization such as Fourier coefficients/frequencies) that selects or shifts the field within a family of solutions.

Training is typically unsupervised or physics-informed. Supervised target fields are not required; instead, discretized field predictions are evaluated using a physics-based loss functional (for example residual- or energy-based losses). During inference, Dirichlet boundary conditions are enforced by inserting prescribed values at Dirichlet indices across the full batch using fol.loss_functions.loss.Loss.GetFullDofVector().

Although training is commonly performed on a fixed FE mesh (with a fixed discretization used inside the loss evaluation), the coordinate-based nature of the synthesizer network makes multi-resolution inference (and, in principle, multi-resolution training) possible. In such cases, the synthesizer can be evaluated on different coordinate sets than those used during training, while maintaining the same conditioning signal from the modulator.

The base fol.deep_neural_networks.deep_network.DeepNetwork provides optimizer and checkpoint integration. This class adds coupling to a fol.controls.control.Control object and a prediction interface that produces discretized fields and then applies boundary conditions.

Parameters:

• name (str) – Name identifier for the model instance (used for logging and checkpointing).

• control (Control) – Control object defining the parametric input space and mapping raw parameters to controlled variables used by the loss and boundary condition enforcement. Typical inputs are parameterization features such as Fourier coefficients/frequencies or other low-dimensional control points.

• loss_function (Loss) – Physics-based loss function evaluated on predicted discretized fields. The loss defines the FE mesh, DOF structure, and boundary-condition handling required for training and inference.

• flax_neural_network (nnx.Module) – Flax/NNX model used for implicit evaluation. In the common setup this is a hypernetwork-based model in which a modulator conditions a coordinate-based synthesizer (neural field) using coupling modes implemented by fol.deep_neural_networks.nns.HyperNetwork. The callable interface must accept coordinates for neural-field evaluation, and the output dimension must be consistent with loss_function.dofs.

• optax_optimizer (GradientTransformation) – Optax optimizer transformation used to construct the optimizer for training.

Raises:

RuntimeError – If the neural network does not expose required attributes (for example in_features or out_features) or if its output dimension is inconsistent with the DOF definition required by the loss function.

ComputeBatchLossValue(batch, nn_model)

Compute the batch loss and return loss metrics.

This method is designed for unsupervised or physics-informed learning. The batch is provided as a tuple for interface consistency with the base class. In typical physics-based training, the second entry of the tuple is unused and may be None because there are no supervised targets.

The loss is computed by first mapping parametric inputs to controlled variables using Control, then predicting discretized fields with the neural network, and finally evaluating the physics-based loss on the predicted fields.

Parameters:

• batch (Tuple[jax.numpy.ndarray, jax.numpy.ndarray]) – Batch tuple (batch_X, batch_y) used for interface consistency. For physics-informed learning, batch_y is typically None and is not used by this method.

• nn_model (nnx.Module) – Neural network used to produce predictions for the batch.

Returns:

A tuple (batch_loss, metrics_dict) where batch_loss is a scalar aggregated over the batch and metrics_dict contains loss statistics including the mandatory key "total_loss".

Return type:

Tuple[jax.numpy.ndarray, dict]

Raises:

None –

ComputeBatchPredictions(batch_X, nn_model)

Compute implicit neural-field predictions for a batch of parametric inputs.

This method evaluates a coordinate-based neural field (the synthesizer network) conditioned by the parametric inputs. The parametric inputs batch_X provide the conditioning signal (typically through a modulator network), while the synthesizer is evaluated on the mesh node coordinates supplied by the loss function. The result is a batch of discretized field values defined on the training mesh.

Although the mesh coordinates used here correspond to the discretization associated with the loss function, the coordinate-based formulation allows the same conditioned neural field to be evaluated on alternative coordinate sets for multi-resolution inference.

Parameters:

• batch_X (jax.numpy.ndarray) – Batch of parametric inputs used to condition the neural field. The leading dimension corresponds to the batch size.

• nn_model (nnx.Module) – Neural network implementing the conditioned neural field and evaluated as nn_model(batch_X, coords).

Returns:

Batch of discretized field predictions produced by the conditioned neural field on the provided mesh coordinates.

Return type:

jax.numpy.ndarray

Raises:

None –

Finalize()

Finalize the model after training.

Subclasses implement this hook to perform any final steps that must occur once at the end of training (for example, releasing resources, computing final diagnostics, exporting artifacts, or post-processing).

Returns:

None

Initialize(reinitialize=False)

Initialize model components and validate dimensional consistency.

This method initializes the base DeepNetwork components and then initializes the associated Control object. It also validates that the provided neural network is compatible with the loss DOF definition.

In the current implementation, compatibility is checked by requiring that the neural network exposes out_features and that it matches the number of DOF components specified by the loss function.

Parameters:

reinitialize (bool, optional) – If True, force reinitialization even if the instance was previously initialized. Default is False.

Returns:

None

Raises:

RuntimeError – If the neural network does not expose in_features or out_features attributes, or if out_features does not match the number of DOF components required by the loss.

Predict(batch_X)

Perform inference for a batch of parametric inputs and apply Dirichlet boundary conditions.

This method runs inference by computing controlled variables from the batch inputs, evaluating the neural network on the mesh node coordinates, and applying Dirichlet boundary conditions by inserting prescribed values at the Dirichlet indices for every sample in the batch using Loss.GetFullDofVector().

Parameters:

batch_X (jax.numpy.ndarray) – Batch of parametric inputs used for inference. The leading dimension corresponds to the batch size.

Returns:

Batch of full discretized field vectors obtained by inference, with Dirichlet boundary conditions applied consistently across the batch.

Return type:

jax.numpy.ndarray

Raises:

None –

PredictDynamics(initial_Batch, num_steps)

Perform autoregressive inference over multiple prediction steps.

Starting from an initial batch of parametric inputs or states, this method repeatedly applies Predict() to generate a trajectory. The returned array stacks the initial batch as the first entry followed by num_steps predicted batches.

Parameters:

• initial_Batch (jax.numpy.ndarray) – Initial batch used to start the rollout. The exact interpretation depends on the model usage, but the leading dimension must correspond to the batch size expected by Predict().

• num_steps (int) – Number of autoregressive prediction steps to perform.

Returns:

Stacked trajectory array containing the initial batch followed by the predicted batches. The first axis indexes time steps.

Return type:

jax.numpy.ndarray

Raises:

ValueError – If num_steps is negative.

Meta-implicit parametric operator learning

This submodule implements meta-implicit parametric operator learning, which extends implicit parametric operator learning by introducing per-sample latent adaptation. Instead of directly conditioning the neural field with parametric inputs, latent variables are optimized in an inner loop to minimize the physics-based loss, enabling fast adaptation without updating the main network weights.

The neural field is represented by a coordinate-based synthesizer MLP and is conditioned through a modulator network using the coupling modes implemented by fol.deep_neural_networks.nns.HyperNetwork. Training remains unsupervised or physics-informed, with explicit boundary-condition enforcement and support for multi-resolution inference.

Authors: Reza Najian Asl, RezaNajian Date: December, 2024 License: FOL/LICENSE

class fol.deep_neural_networks.meta_implicit_parametric_operator_learning.MetaImplicitParametricOperatorLearning(name, control, loss_function, flax_neural_network, main_loop_optax_optimizer, latent_step_size=0.01, num_latent_iterations=3)

Bases: ImplicitParametricOperatorLearning

Meta-learning extension of implicit parametric operator learning with latent-code adaptation.

This class derives from ImplicitParametricOperatorLearning and adds a meta-learning style inner-loop optimization over latent codes. In the base implicit formulation, the parametric input directly conditions a coordinate-based neural field (typically via a modulator–synthesizer structure). In this meta formulation, the parametric input is first mapped to controlled variables, and then a latent code is optimized per sample to minimize the physics-based loss. The optimized latent code is subsequently used to evaluate the neural field on the mesh coordinates.

The practical difference is that the conditioning signal is not applied only by feeding the parametric input into the network. Instead, the model performs test-time (or per-batch) adaptation by iteratively updating latent variables using gradients of the physics loss. This enables fast adaptation to new parameter regimes or tasks while keeping the main network weights fixed during the inner loop.

The provided flax_neural_network is expected to be a HyperNetwork, where the synthesizer represents a coordinate-based MLP neural field evaluated on FE mesh node coordinates and the modulator provides conditioning according to supported coupling modes. Training and loss evaluation are typically performed on a fixed FE mesh, but the coordinate-based synthesizer formulation makes multi-resolution inference (and, in principle, multi-resolution training) possible by evaluating the conditioned neural field on alternative coordinate sets.

Parameters:

• name (str) – Name identifier for the model instance (used for logging and checkpointing).

• control (Control) – Control object defining the parametric input space and mapping raw parameters to controlled variables used by the loss. Typical inputs are low-dimensional parameterization features such as Fourier coefficients or other control points.

• loss_function (Loss) – Physics-based loss function evaluated on predicted discretized fields. The loss defines the FE mesh, DOF structure, and boundary-condition handling required for training and inference.

• flax_neural_network (HyperNetwork) – Hypernetwork-based model used for implicit evaluation. The network is evaluated as nn_model(latent_codes, coords), where latent_codes are adapted by the inner loop and coords are FE mesh node coordinates. The synthesizer component performs the coordinate-based neural-field mapping and the modulator provides conditioning according to the hypernetwork coupling settings.

• main_loop_optax_optimizer (GradientTransformation) – Optax optimizer transformation used for the outer (main) optimization loop over network parameters.

• latent_step_size (float, optional) – Step size used for gradient-based latent-code updates in the inner loop. Default is 1e-2.

• num_latent_iterations (int, optional) – Number of inner-loop latent optimization iterations performed per batch. Default is 3.

Raises:

RuntimeError – If the provided neural network does not expose required attributes (for example in_features) needed to size the latent codes, or if the network interface is incompatible with nn_model(latent, coords).

ComputeBatchPredictions(batch_X, nn_model)

Compute batch predictions using latent-code inner-loop optimization.

This method overrides the base implicit prediction behavior by introducing a per-sample latent adaptation step. For each parametric input in batch_X, the method initializes a latent code and iteratively updates it using the gradient of the physics-based loss with respect to the latent variables. The neural network is evaluated on the FE mesh node coordinates at each iteration to compute the loss, and the latent code is updated using a normalized gradient step. After num_latent_iterations updates, the final latent codes are used to produce the discretized field predictions on the mesh.

Compared to ComputeBatchPredictions(), where parametric inputs directly condition the neural field through the network interface, this meta formulation conditions the neural field indirectly through optimized latent variables that minimize the physics loss for the given control outputs.

Parameters:

• batch_X (jax.numpy.ndarray) – Batch of parametric inputs. The leading dimension corresponds to the batch size. These inputs are transformed into controlled variables by self.control for physics-based loss evaluation.

• nn_model (nnx.Module) – Neural model evaluated as nn_model(latent_codes, coords), where latent_codes are adapted in the inner loop and coords are FE mesh node coordinates obtained from the loss function mesh. The model must expose in_features to define the latent code dimension.

Returns:

Batch of discretized field predictions evaluated on the FE mesh node coordinates using the optimized latent codes.

Return type:

jax.numpy.ndarray

Raises:

• RuntimeError – If nn_model does not expose in_features required to size the latent-code array, or if the model cannot be called as nn_model(latent_codes, coords).

• ValueError – If self.num_latent_iterations is negative or if the latent-gradient normalization encounters invalid norms.

Meta-alpha-meta implicit parametric operator learning

This submodule implements meta-alpha-meta implicit parametric operator learning, which further extends meta-implicit parametric operator learning by introducing a learnable latent-step size in the inner-loop adaptation. As in the meta-implicit formulation, latent variables are optimized per sample to minimize the physics-based loss, but in this variant the magnitude of the latent update itself is learned jointly with the network parameters.

The neural field is represented by a coordinate-based synthesizer MLP and is conditioned through a modulator network using the coupling modes implemented by fol.deep_neural_networks.nns.HyperNetwork. The latent codes are adapted using gradient-based updates, while a dedicated trainable step model controls the latent update size, enabling improved robustness and adaptability across problem instances.

Training remains unsupervised or physics-informed, with explicit enforcement of boundary conditions and preservation of the coordinate-based formulation, which allows multi-resolution inference by evaluating the conditioned neural field on alternative coordinate sets.

Authors: Reza Najian Asl, RezaNajian Date: December, 2024 License: FOL/LICENSE

class fol.deep_neural_networks.meta_alpha_meta_implicit_parametric_operator_learning.LatentStepModel(*args, **kwargs)

Bases: Module

Trainable latent-step wrapper used for meta-learning the inner-loop step size.

This small NNX module stores a single scalar parameter latent_step as an flax.nnx.Param. It is used by meta-implicit learning methods to learn the step size applied during latent-code updates in the inner loop.

Parameters:

init_latent_step_value (float or jax.Array) – Initial value for the latent-step parameter.

Returns:

None

Raises:

TypeError – If init_latent_step_value cannot be stored as an NNX parameter.

class fol.deep_neural_networks.meta_alpha_meta_implicit_parametric_operator_learning.MetaAlphaMetaImplicitParametricOperatorLearning(name, control, loss_function, flax_neural_network, main_loop_optax_optimizer, latent_step_optax_optimizer, latent_step_size=0.01, num_latent_iterations=3)

Bases: ImplicitParametricOperatorLearning

Meta-implicit parametric operator learning with a learnable latent-step size.

This class derives from ImplicitParametricOperatorLearning and implements an implicit, coordinate-based operator learning workflow where a neural-field synthesizer is evaluated on FE mesh coordinates. The model is typically a HyperNetwork, meaning the synthesizer performs the coordinate-based neural-field mapping while a modulator conditions the synthesizer according to the coupling modes defined by HyperNetwork.

The main difference from the base implicit class is that predictions are not produced directly from parametric inputs. Instead, a latent code is optimized per sample in an inner loop by minimizing the physics-based loss. The main difference from MetaImplicitParametricOperatorLearning is that this class also learns the inner-loop update magnitude. A dedicated trainable scalar step model is optimized jointly with the network parameters, so the latent step size adapts during training.

Training is usually unsupervised or physics-informed. The loss is evaluated on discretized field predictions, and Dirichlet boundary conditions can be enforced at inference time by inserting prescribed values across the batch using GetFullDofVector().

Parameters:

• name (str) – Name identifier used for logging and checkpointing.

• control (Control) – Control object that maps raw parametric inputs to controlled variables used by the loss functional.

• loss_function (Loss) – Physics-based loss functional evaluated on predicted discretized fields. The loss provides access to the FE mesh coordinates and DOF handling.

• flax_neural_network (HyperNetwork) – HyperNetwork used for implicit neural-field prediction. The synthesizer is expected to be coordinate-based, and the modulator provides conditioning according to HyperNetwork coupling settings.

• main_loop_optax_optimizer (GradientTransformation) – Optimizer transformation for updating the neural network parameters.

• latent_step_optax_optimizer (GradientTransformation) – Optimizer transformation for updating the latent-step parameter.

• latent_step_size (float, optional) – Initial value for the learnable latent step size. Default is 1e-2.

• num_latent_iterations (int, optional) – Number of latent inner-loop update iterations used during prediction. Default is 3.

Returns:

None

Raises:

• RuntimeError – If the model is used before initialization in workflows that require initialized loss/control components.

• ValueError – If num_latent_iterations is negative.

ComputeBatchPredictions(batch_X, meta_model)

Compute batch predictions using latent-code adaptation with a learnable step size.

This method performs per-sample latent optimization in an inner loop. For each parametric input, controlled variables are computed using self.control. Latent codes are initialized and iteratively updated by descending the gradient of the physics-based loss with respect to the latent variables. Unlike the simpler meta-implicit variant where the latent update step size is fixed, this class uses a trainable step model latent_step that is optimized during the outer training loop.

The neural field is evaluated on FE mesh node coordinates obtained from the loss function mesh. The final latent codes are used to produce discretized field predictions on the same coordinate set.

Parameters:

• batch_X (jax.numpy.ndarray) – Batch of parametric inputs. The leading dimension corresponds to the batch size.

• meta_model (Tuple[nnx.Module, nnx.Module]) – Tuple (nn_model, latent_step_model) where nn_model is the HyperNetwork used for predictions and latent_step_model returns the current learnable step size via latent_step_model().

Returns:

Batch of discretized field predictions evaluated on the FE mesh node coordinates.

Return type:

jax.numpy.ndarray

Raises:

• RuntimeError – If nn_model does not expose in_features required to size the latent-code array, or if the model cannot be called as nn_model(latent_codes, coords).

• ValueError – If self.num_latent_iterations is negative.

GetState()

Return the full meta-learning state required for training and checkpointing.

The returned tuple includes the main network, its optimizer state, the latent-step model, and the latent-step optimizer state.

Parameters:

None –

Returns:

Tuple (flax_neural_network, nnx_optimizer, latent_step_model, latent_step_optimizer).

Return type:

Tuple[nnx.Module, nnx.Optimizer, nnx.Module, nnx.Optimizer]

Raises:

None –

Predict(batch_X)

Perform inference for a batch and apply Dirichlet boundary conditions.

This method computes controlled variables from batch_X, predicts the discretized field using latent-code adaptation with a learnable step size, and returns the full DOF vector with Dirichlet boundary values inserted consistently across the batch via Loss.GetFullDofVector().

Parameters:

batch_X (jax.numpy.ndarray) – Batch of parametric inputs for inference.

Returns:

Batch of full discretized field vectors with Dirichlet boundary conditions applied.

Return type:

jax.numpy.ndarray

Raises:

RuntimeError – If prediction fails due to incompatible model signatures or loss DOF handling.

RestoreState(restore_state_directory)

Restore both the main network state and the latent-step state from disk.

This method restores the HyperNetwork parameters from the nn subdirectory and restores the latent-step model parameter from the latent subdirectory. After restoration, the in-memory modules are updated in place.

Parameters:

restore_state_directory (str) – Directory containing the saved meta checkpoint. The directory is expected to include nn and latent subdirectories.

Returns:

None

Raises:

• FileNotFoundError – If the expected checkpoint directories do not exist.

• RuntimeError – If the checkpointer fails to restore or update the module states.

SaveCheckPoint(check_point_type, checkpoint_state_dir)

Save both the main network state and the latent-step state to disk.

This method writes two checkpoint subdirectories under checkpoint_state_dir. One subdirectory stores the main HyperNetwork state and the other stores the latent-step model state. This allows full restoration of meta-training and inference behavior, including the learned latent step size.

Parameters:

• check_point_type (str) – Label used for logging the checkpoint type, for example "best" or "latest".

• checkpoint_state_dir (str) – Directory path where the checkpoint subdirectories are created.

Returns:

None

Raises:

• OSError – If the checkpoint directories cannot be created or written.

• RuntimeError – If the underlying checkpointer fails to save the provided states.

TestStep(meta_state, data)

Compute the batch loss for evaluation using the current meta-state.

Parameters:

• meta_state (Tuple[nnx.Module, nnx.Optimizer, nnx.Module, nnx.Optimizer]) – Tuple containing (nn_model, main_optimizer, latent_step_model, latent_optimizer).

• data (Tuple[jax.numpy.ndarray, jax.numpy.ndarray]) – Batch tuple used for interface consistency. The second element may be None for physics-informed workflows.

Returns:

Scalar batch loss value.

Return type:

jax.numpy.ndarray

Raises:

RuntimeError – If loss evaluation fails.

TrainStep(meta_state, data)

Perform one meta-training step updating both network parameters and latent-step.

This step computes gradients of the batch physics loss with respect to both the main network parameters and the learnable latent-step parameter. The main optimizer updates the network, and the latent optimizer updates the step model.

Parameters:

• meta_state (Tuple[nnx.Module, nnx.Optimizer, nnx.Module, nnx.Optimizer]) – Tuple containing (nn_model, main_optimizer, latent_step_model, latent_optimizer).

• data (Tuple[jax.numpy.ndarray, jax.numpy.ndarray]) – Batch tuple passed through the base training interface. The second element is typically unused for physics-informed learning and may be None.

Returns:

Scalar batch loss value.

Return type:

jax.numpy.ndarray

Raises:

RuntimeError – If gradient computation fails due to incompatible model outputs or loss evaluation.

Fourier parametric operator learning

This submodule implements Fourier parametric operator learning using a Fourier Neural Operator (FNO) on discretized fields. A fixed-dimensional parametric input space (for example control variables or parameterization features) is mapped to grid-aligned input channels and processed by the FNO to produce discretized field outputs such as temperature or displacement.

The FNO is not bound to a specific mesh resolution. Once trained, the learned operator can be evaluated on different grid resolutions, as long as the mesh is structured and uniform (for example square grids in 2D or cubic grids in 3D). This enables resolution-invariant inference across compatible discretizations.

The learning can be data-driven or physics-informed, depending on the chosen loss function, with boundary conditions enforced explicitly through the loss.

Authors: Reza Najian Asl, RezaNajian Date: June, 2025 License: FOL/LICENSE

class fol.deep_neural_networks.fourier_parametric_operator_learning.DataDrivenFourierParametricOperatorLearning(name, control, loss_function, flax_neural_network, optax_optimizer)

Bases: FourierParametricOperatorLearning

Data-driven parametric operator learning using a Fourier Neural Operator.

This class specializes FourierParametricOperatorLearning for fully supervised, data-driven training. The Fourier Neural Operator is trained by directly comparing predicted discretized fields against provided target fields using a supervised loss function.

No physical constraints are enforced explicitly in the loss; instead, the operator is learned purely from input–output field data. This formulation is appropriate when high-fidelity simulation or experimental data are available.

Parameters:

• name (str) – Name identifier for the model instance, used for logging and checkpointing.

• control (Control) – Control object defining the parametric input space and mapping raw parameters to controlled variables arranged as grid-aligned inputs.

• loss_function (Loss) – Supervised loss function used to compare predicted discretized fields against target fields.

• flax_neural_network (nnx.Module) – Fourier Neural Operator mapping grid-aligned inputs to grid-aligned output fields.

• optax_optimizer (GradientTransformation) – Optax optimizer transformation used for training.

Raises:

None –

ComputeBatchLossValue(batch, nn_model)

Compute supervised batch loss for data-driven Fourier Neural Operator training.

The batch consists of parametric inputs and corresponding target discretized fields. Predictions are computed using the Fourier Neural Operator and are compared directly against the target fields using the provided loss function.

Parameters:

• batch (Tuple[jax.numpy.ndarray, jax.numpy.ndarray]) – Tuple (batch_X, batch_Y) where batch_X contains parametric inputs and batch_Y contains target discretized fields.

• nn_model (nnx.Module) – Fourier Neural Operator used to generate predictions.

Returns:

Batch loss value and a dictionary of loss statistics including the key "total_loss".

Return type:

Tuple[jax.numpy.ndarray, dict]

Raises:

None –

class fol.deep_neural_networks.fourier_parametric_operator_learning.FourierParametricOperatorLearning(name, control, loss_function, flax_neural_network, optax_optimizer)

Bases: DeepNetwork

Parametric operator learning using a Fourier Neural Operator (FNO).

This class implements parametric operator learning where the neural model is a Fourier Neural Operator operating on discretized fields defined on a structured grid. Parametric inputs are first mapped to controlled variables using a Control object and then reshaped into grid-aligned tensor representations suitable for FNO evaluation.

In this formulation, both training and inference assume a fixed grid resolution associated with the finite-element mesh used by the loss function. The FNO maps grid-aligned input channels to grid-aligned output fields in the spectral domain, after which the outputs are flattened and embedded into full DOF vectors with boundary conditions applied by the loss function.

This class provides the common infrastructure for Fourier Neural Operator– based parametric operator learning. Concrete training paradigms are defined in derived classes, such as data-driven and physics-informed variants.

Parameters:

• name (str) – Name identifier for the model instance, used for logging and checkpointing.

• control (Control) – Control object defining the parametric input space and mapping raw parameters to grid-aligned controlled variables.

• loss_function (Loss) – Loss function defining the discretized field representation, DOF structure, mesh topology, and boundary-condition handling.

• flax_neural_network (nnx.Module) – Fourier Neural Operator mapping grid-aligned inputs to grid-aligned output fields.

• optax_optimizer (GradientTransformation) – Optax optimizer transformation used for training.

Raises:

RuntimeError – If the control output or loss mesh is incompatible with the assumed structured grid required by the Fourier Neural Operator.

ComputeBatchPredictions(batch_X, nn_model)

Compute grid-based field predictions for a batch of parametric inputs.

This method reshapes the controlled parametric inputs into a structured grid format compatible with Fourier-based neural operators. The neural network is then evaluated on the grid-aligned inputs, and the resulting grid-aligned outputs are flattened into vectors of unknown DOFs.

Parameters:

• batch_X (jax.numpy.ndarray) – Batch of controlled parametric inputs. The array is expected to encode grid-aligned channels for each sample.

• nn_model (nnx.Module) – Fourier-based neural operator evaluated on reshaped grid inputs.

Returns:

Batch of flattened field predictions corresponding to the unknown degrees of freedom.

Return type:

jax.numpy.ndarray

Raises:

ValueError – If the input cannot be reshaped consistently with the mesh size inferred from the loss function.

Finalize()

Finalize the model after training.

Subclasses implement this hook to perform any final steps that must occur once at the end of training (for example, releasing resources, computing final diagnostics, exporting artifacts, or post-processing).

Returns:

None

Initialize(reinitialize=False)

Initialize the Fourier parametric operator learning model.

This method initializes the base DeepNetwork components and ensures that the associated Control object is also initialized. No additional dimensional consistency checks are performed here beyond those enforced by the base class and the loss function.

Parameters:

reinitialize (bool, optional) – If True, force reinitialization even if the instance was already initialized. Default is False.

Returns:

None

Raises:

None –

Predict(batch_control)

Perform inference and apply boundary conditions for a batch of inputs.

This method computes controlled variables from the raw parametric inputs, evaluates the Fourier-based neural operator to predict discretized fields, and applies Dirichlet boundary conditions by inserting prescribed values across the batch using the loss function.

Parameters:

batch_control (jax.numpy.ndarray) – Batch of parametric inputs used for inference.

Returns:

Batch of full discretized field vectors with boundary conditions applied.

Return type:

jax.numpy.ndarray

Raises:

None –

class fol.deep_neural_networks.fourier_parametric_operator_learning.PhysicsInformedFourierParametricOperatorLearning(name, control, loss_function, flax_neural_network, optax_optimizer)

Bases: FourierParametricOperatorLearning

Physics-informed parametric operator learning using a Fourier Neural Operator.

This class specializes FourierParametricOperatorLearning for unsupervised or physics-informed training. The Fourier Neural Operator predicts discretized fields that are evaluated using physics-based loss functionals, such as residual- or energy-based formulations, rather than supervised target data.

This formulation is suitable when governing equations are known and labeled training data are limited or unavailable. Boundary conditions are enforced explicitly through the loss function.

Parameters:

• name (str) – Name identifier for the model instance, used for logging and checkpointing.

• control (Control) – Control object defining the parametric input space and mapping raw parameters to controlled variables.

• loss_function (Loss) – Physics-based loss function evaluated on predicted discretized fields.

• flax_neural_network (nnx.Module) – Fourier Neural Operator mapping grid-aligned inputs to grid-aligned output fields.

• optax_optimizer (GradientTransformation) – Optax optimizer transformation used for training.

Raises:

None –

ComputeBatchLossValue(batch, nn_model)

Compute physics-informed batch loss for Fourier Neural Operator learning.

The batch is provided as a tuple for interface consistency, but only the parametric inputs are used. Predicted discretized fields are evaluated using a physics-based loss functional without supervised target fields.

Parameters:

• batch (Tuple[jax.numpy.ndarray, jax.numpy.ndarray]) – Tuple (batch_X, None) where batch_X contains parametric inputs. The second entry is unused.

• nn_model (nnx.Module) – Fourier Neural Operator used to generate predictions.

Returns:

Batch loss value and a dictionary of loss statistics including the key "total_loss".

Return type:

Tuple[jax.numpy.ndarray, dict]

Raises:

None –

DeepONet parametric operator learning

This submodule implements DeepONet-based parametric operator learning on discretized fields. A fixed-dimensional parametric input space conditions a DeepONet that is evaluated on FE mesh node coordinates to produce discretized field outputs such as temperature or displacement.

The learning can be data-driven or physics-informed, depending on the chosen loss function, with boundary conditions enforced explicitly through the loss and inference utilities.

Authors: Reza Najian Asl, RezaNajian Date: June, 2025 License: FOL/LICENSE

class fol.deep_neural_networks.deep_o_net_parametric_operator_learning.DataDrivenDeepONetParametricOperatorLearning(name, control, loss_function, flax_neural_network, optax_optimizer)

Bases: DeepONetParametricOperatorLearning

Data-driven parametric operator learning using a DeepONet.

This class specializes DeepONetParametricOperatorLearning for fully supervised, data-driven training. The DeepONet predicts discretized fields that are directly compared against provided target fields using a supervised loss function.

In this formulation, the operator is learned purely from input–output data pairs. No physical constraints are enforced explicitly beyond those embedded in the training data. Boundary-condition handling is delegated to the loss function and to the inference routine that assembles full DOF vectors.

Parameters:

• name (str) – Name identifier for the model instance, used for logging and checkpointing.

• control (Control) – Control object defining the parametric input space and mapping raw parameters to controlled variables supplied to the DeepONet.

• loss_function (Loss) – Supervised loss function used to compare predicted discretized fields against target fields provided in the training data.

• flax_neural_network (nnx.Module) – DeepONet model evaluated as nn_model(inputs, coords), where coords are FE mesh node coordinates.

• optax_optimizer (GradientTransformation) – Optax optimizer transformation used for training.

Raises:

RuntimeError – If the predicted field shape is incompatible with the target field shape expected by the loss function.

ComputeBatchLossValue(batch, nn_model)

Compute data-driven batch loss and return loss metrics.

The batch is interpreted as (batch_X, batch_y) where batch_y is a target discretized field. Predictions are computed from controlled inputs and compared against the targets using the configured loss function.

Parameters:

• batch (Tuple[jax.numpy.ndarray, jax.numpy.ndarray]) – Training batch tuple (batch_X, batch_y). batch_X contains raw parametric inputs and batch_y contains target fields.

• nn_model (nnx.Module) – DeepONet used to produce predictions for the batch.

Returns:

A tuple (batch_loss, metrics_dict) where batch_loss is a scalar aggregated over the batch and metrics_dict contains loss statistics including the mandatory key "total_loss".

Return type:

Tuple[jax.numpy.ndarray, dict]

Raises:

RuntimeError – If target tensors in batch_y are incompatible with the loss function expectations or the prediction shape.

class fol.deep_neural_networks.deep_o_net_parametric_operator_learning.DeepONetParametricOperatorLearning(name, control, loss_function, flax_neural_network, optax_optimizer)

Bases: DeepNetwork

Parametric operator learning using a DeepONet evaluated on FE mesh coordinates.

This class implements a DeepONet-based operator-learning workflow where a fixed-dimensional parametric input (for example control variables or parameterization features) is mapped to a discretized field by evaluating a DeepONet on the node coordinates of the FE mesh provided by the loss function. The model therefore learns an operator from a parametric space to a field space discretized on a mesh.

The base DeepNetwork functionality is used for optimizer and checkpoint integration. This class adds coupling to a Control object and defines a prediction interface that supplies the mesh coordinates to the DeepONet.

During inference, Dirichlet boundary conditions are enforced by inserting prescribed values at the Dirichlet indices across the full batch using Loss.GetFullDofVector().

Parameters:

• name (str) – Name identifier for the model instance, used for logging and checkpointing.

• control (Control) – Control object defining the parametric input space and mapping raw parameters to controlled variables used by the network and loss.

• loss_function (Loss) – Loss object defining the FE mesh discretization (node coordinates) and the DOF/boundary-condition handling used to assemble full field vectors.

• flax_neural_network (nnx.Module) – DeepONet model expected to be callable as nn_model(batch_inputs, coords) where coords are the FE mesh node coordinates. The network output is reshaped to (batch_size, -1) to represent discretized field values.

• optax_optimizer (GradientTransformation) – Optax optimizer transformation used to construct the optimizer for training.

Raises:

RuntimeError – If the DeepONet or loss configuration is incompatible with the FE mesh discretization or boundary-condition assembly required by Loss.GetFullDofVector().

ComputeBatchPredictions(batch_X, nn_model)

Compute DeepONet predictions for a batch of parametric inputs.

The DeepONet is evaluated on the FE mesh node coordinates retrieved from the loss function. The returned output is reshaped to (batch_size, -1) so it can be treated as a batch of discretized field vectors.

Parameters:

• batch_X (jax.numpy.ndarray) – Batch of parametric (or controlled) inputs. The first dimension is the batch size.

• nn_model (nnx.Module) – DeepONet model evaluated as nn_model(batch_X, coords).

Returns:

Batch of discretized field predictions with shape (batch_size, num_predicted_dofs).

Return type:

jax.numpy.ndarray

Raises:

RuntimeError – If the FE mesh coordinates cannot be obtained from the loss function, or if the network output cannot be reshaped to a 2D batch array.

Finalize()

Finalize the model after training.

Subclasses implement this hook to perform any final steps that must occur once at the end of training (for example, releasing resources, computing final diagnostics, exporting artifacts, or post-processing).

Returns:

None

Initialize(reinitialize=False)

Initialize model components and the associated control object.

This method initializes the base DeepNetwork components and then initializes the associated Control instance. Initialization is skipped if the model was already initialized, unless reinitialize is set to True.

Parameters:

reinitialize (bool, optional) – If True, force reinitialization even if the instance was previously initialized. Default is False.

Returns:

None

Raises:

None –

Predict(batch_X)

Perform inference for a batch of inputs and apply Dirichlet constraints.

This method computes controlled variables from batch_X, evaluates the DeepONet on the FE mesh coordinates, and assembles the full DOF vector by applying Dirichlet boundary conditions across the batch using Loss.GetFullDofVector().

Parameters:

batch_X (jax.numpy.ndarray) – Batch of raw parametric inputs. The first dimension is the batch size.

Returns:

Batch of full discretized field vectors with Dirichlet boundary conditions enforced.

Return type:

jax.numpy.ndarray

Raises:

RuntimeError – If boundary-condition assembly via Loss.GetFullDofVector() fails due to inconsistent DOF definitions or incompatible shapes.

class fol.deep_neural_networks.deep_o_net_parametric_operator_learning.PhysicsInformedDeepONetParametricOperatorLearning(name, control, loss_function, flax_neural_network, optax_optimizer)

Bases: DeepONetParametricOperatorLearning

Physics-informed parametric operator learning using a DeepONet.

This class specializes DeepONetParametricOperatorLearning for unsupervised or physics-informed training. No supervised target fields are required. Instead, discretized field predictions produced by the DeepONet are evaluated using physics-based loss functionals, such as residual- or energy-based formulations.

The DeepONet is evaluated on FE mesh node coordinates, and the resulting predictions are passed directly to the loss function together with the controlled variables. Dirichlet boundary conditions are enforced through the loss and during inference by assembling full DOF vectors.

Parameters:

• name (str) – Name identifier for the model instance, used for logging and checkpointing.

• control (Control) – Control object defining the parametric input space and mapping raw parameters to controlled variables used by the physics-based loss.

• loss_function (Loss) – Physics-based loss function used to evaluate discretized field predictions without supervised target data.

• flax_neural_network (nnx.Module) – DeepONet model evaluated as nn_model(inputs, coords), where coords are FE mesh node coordinates.

• optax_optimizer (GradientTransformation) – Optax optimizer transformation used for training.

Raises:

RuntimeError – If the physics-based loss evaluation fails due to incompatible control variables, mesh definitions, or predicted field shapes.

ComputeBatchLossValue(batch, nn_model)

Compute physics-informed batch loss and return loss metrics.

The physics-informed loss is computed by first mapping raw inputs to controlled variables, predicting discretized fields using the DeepONet, and then evaluating the physics-based loss against the controlled variables (rather than supervised targets).

Parameters:

• batch (Tuple[jax.numpy.ndarray, jax.numpy.ndarray]) – Training batch tuple (batch_X, batch_y) kept for interface consistency. batch_y is typically None and is not used.

• nn_model (nnx.Module) – DeepONet used to produce predictions for the batch.

Returns:

A tuple (batch_loss, metrics_dict) where batch_loss is a scalar aggregated over the batch and metrics_dict contains loss statistics including the mandatory key "total_loss".

Return type:

Tuple[jax.numpy.ndarray, dict]

Raises:

RuntimeError – If the physics loss evaluation fails due to inconsistent control outputs, mesh definitions, or prediction shapes.

Neural fields and hypernetworks

This module provides building blocks for neural field models (e.g., SIREN-style MLPs and Fourier-feature MLPs) and DeepONets, and hypernetworks used to modulate or generate parameters for coordinate-based models. These components are commonly used in implicit neural representations, conditional neural fields, and meta-learning workflows.

Authors: Reza Najian Asl, RezaNajian Date: November, 2024 License: FOL/LICENSE

class fol.deep_neural_networks.nns.HyperNetwork(*args, **kwargs)

Bases: Module

Hypernetwork that modulates a synthesizer MLP using a modulator MLP.

The hypernetwork couples a modulator network to a synthesizer network and applies modulation during the synthesizer forward pass. In this implementation, the supported coupled variable is a bias-like shift that is added to synthesizer layer activations.

Coupling modes supported by coupling_settings["modulator_to_synthesizer_coupling_mode"] are:

1. "all_to_all": layer-wise coupling between modulator and synthesizer.

2. "last_to_all": modulator final output is injected into every synthesizer layer.

3. "one_modulator_per_synthesizer_layer": a separate modulator is created per synthesizer layer.

Parameters:

• name (str) – Name identifier for the hypernetwork instance.

• modulator_nn (MLP) – Modulator network producing conditioning signals.

• synthesizer_nn (MLP) – Synthesizer network producing task outputs.

• coupling_settings (dict, optional) – Coupling configuration dictionary. Common keys are "coupled_variable" and "modulator_to_synthesizer_coupling_mode". Default is {}.

Raises:

• KeyError – If required entries are missing from coupling_settings.

• ValueError – If an unsupported coupling mode is requested or if network shapes are inconsistent with the chosen coupling mode.

CountTrainableParams()

Count trainable parameters registered as nnx.Param in this module.

Parameters:

None –

Returns:

Total number of trainable scalar parameters.

Return type:

int

Raises:

None –

GetName()

Return the name identifier of this hypernetwork instance.

Parameters:

None –

Returns:

Name of the hypernetwork instance.

Return type:

str

Raises:

None –

all_to_all_fw(latent_array, coord_matrix, modulator_nn, synthesizer_nn)

Forward pass for the "all_to_all" coupling mode.

In this mode, the modulator and synthesizer networks are assumed to have identical layer structures. At each layer, the modulator activation is added to the synthesizer activation (shift coupling), and both networks apply their respective activation functions.

The computation proceeds as follows.

Step 1:

Reshape latent_array to a row-vector input for the modulator network and initialize synthesizer activations using coord_matrix.

Step 2:

For each coupled layer, compute the modulator pre-activation (with optional skip connections), compute the synthesizer pre-activation (with optional skip connections), add the modulator output to the synthesizer activations, and apply the activation functions for both networks.

Step 3:

Apply the final synthesizer linear layer to the last synthesizer activations and return the resulting output.

Parameters:

• latent_array (jax.Array) – Conditioning input for the modulator network.

• coord_matrix (jax.Array) – Coordinate or feature input for the synthesizer network.

• modulator_nn (MLP) – Modulator network providing per-layer conditioning.

• synthesizer_nn (MLP) – Synthesizer network producing the final output.

Returns:

Output of the synthesizer network after applying all-to-all coupling.

Return type:

jax.Array

Raises:

ValueError – If the modulator and synthesizer architectures are incompatible for layer-wise coupling.

last_to_all_fw(latent_array, coord_matrix, modulator_nn, synthesizer_nn)

Forward pass for the "last_to_all" coupling mode.

In this mode, the modulator is first evaluated once to produce a single modulation vector. That vector is then split into chunks matching the bias sizes of the synthesizer hidden layers, and each chunk is added to the corresponding synthesizer layer activations (shift coupling).

The computation proceeds as follows.

Step 1:

Forward propagate latent_array through the modulator network to produce a global modulation vector x_modul.

Step 2:

Initialize synthesizer activations using coord_matrix.

Step 3:

For each synthesizer hidden layer, compute the layer pre-activation (with optional skip connections), take the slice of x_modul matching that layer width, add this slice to the synthesizer activations, and apply the synthesizer activation function.

Step 4:

Apply the final synthesizer linear layer and return the output.

Parameters:

• latent_array (jax.Array) – Conditioning input for the modulator network.

• coord_matrix (jax.Array) – Coordinate or feature input for the synthesizer network.

• modulator_nn (MLP) – Modulator network producing a global modulation vector.

• synthesizer_nn (MLP) – Synthesizer network producing the final output.

Returns:

Output of the synthesizer network after applying last-to-all coupling.

Return type:

jax.Array

Raises:

ValueError – If the modulator output size does not match the total number of modulated synthesizer hidden biases.

one_modulator_per_synthesizer_layer_fw(latent_array, coord_matrix, modulator_nns, synthesizer_nn)

Forward pass for "one_modulator_per_synthesizer_layer" coupling.

In this mode, each synthesizer hidden layer has a dedicated modulator network that is evaluated on the same latent_array. The resulting modulation vector is added to the corresponding synthesizer layer activations (shift coupling).

The computation proceeds as follows.

Step 1:

Initialize synthesizer activations using coord_matrix.

Step 2:

For each synthesizer hidden layer, compute the synthesizer layer pre-activation (with optional skip connections), evaluate the corresponding modulator to obtain a modulation vector of matching size, add this modulation to the synthesizer activations, and apply the synthesizer activation function.

Step 3:

Apply the final synthesizer linear layer and return the output.

Parameters:

• latent_array (jax.Array) – Conditioning input shared across all modulators.

• coord_matrix (jax.Array) – Coordinate or feature input for the synthesizer network.

• modulator_nns (list[MLP]) – List of modulator networks, one per synthesizer hidden layer.

• synthesizer_nn (MLP) – Synthesizer network producing the final output.

Returns:

Output of the synthesizer network after applying per-layer modulator coupling.

Return type:

jax.Array

Raises:

ValueError – If the number of modulators does not match the number of synthesizer hidden layers.

class fol.deep_neural_networks.nns.MLP(*args, **kwargs)

Bases: Module

General-purpose multi-layer perceptron (MLP) building block.

This module can be used to construct a wide range of MLP-based architectures, including neural fields, DeepONets, and standard feed-forward networks. It supports a variety of activation functions such as "relu", "leaky_relu", "elu", and "tanh", with a sinusoidal activation (SIREN-style "sin") as the default option. The sinusoidal activation is particularly suited for neural fields and implicit representations that must capture high-frequency behavior.

In addition to standard dense layers, the MLP can:

1. Use optional skip connections with a configurable frequency to improve gradient flow and expressivity.

2. Apply Fourier feature mappings at the input layer, optionally with

   learnable Fourier frequencies, to better represent high-frequency signals in coordinate-based models.

These capabilities make the class suitable as a reusable backbone for conditional neural fields, operator-learning networks such as DeepONets, and other coordinate-based models.

Parameters:

• name (str) – Name identifier for the network instance.

• input_size (int, optional) – Number of input features. Default is 0.

• output_size (int, optional) – Number of output features. If 0, the network can be configured as a feature extractor without a final output layer. Default is 0.

• hidden_layers (list, optional) – Hidden layer widths. Default is [].

• activation_settings (dict, optional) – Activation configuration dictionary. Common keys include "type", "prediction_gain", and "initialization_gain". Missing entries are filled using internal defaults. Default is {}.

• use_bias (bool, optional) – If True, create trainable biases. If False, biases are fixed zeros. Default is True.

• skip_connections_settings (dict, optional) – Skip connection configuration dictionary. Expected keys are "active" and "frequency". Missing entries are filled using internal defaults. Default is {}.

• fourier_feature_settings (dict, optional) – Fourier feature mapping configuration dictionary. Expected keys are "active", "type", "size", "frequency_scale", and "learn_frequency". Missing entries are filled using internal defaults. Default is {}.

Raises:

• KeyError – If required activation configuration keys are missing after defaults are applied.

• ValueError – If any provided dimension (input, output, hidden) is negative or if Fourier feature settings are inconsistent.

ComputeX(w, prev_x, b)

Compute an affine layer output without skip connections.

Parameters:

• w (nnx.Param) – Weight matrix for the current layer.

• prev_x (jax.Array) – Input activations to the current layer.

• b (nnx.Param) – Bias vector for the current layer.

Returns:

Layer pre-activation output computed as prev_x @ w + b.

Return type:

jax.Array

Raises:

ValueError – If the input dimensions are incompatible for matrix multiplication.

ComputeXSkip(w, prev_x, in_x, b)

Compute an affine layer output with a skip connection.

This method concatenates the current activation prev_x with the original network input in_x along the feature axis and applies an affine transform.

Parameters:

• w (nnx.Param) – Weight matrix for the current layer.

• prev_x (jax.Array) – Current layer activations.

• in_x (jax.Array) – Original input to the network that is injected via a skip path.

• b (nnx.Param) – Bias vector for the current layer.

Returns:

Layer pre-activation output computed from the concatenated input.

Return type:

jax.Array

Raises:

ValueError – If concatenation or matrix multiplication dimensions are incompatible.

CountTrainableParams()

Count trainable parameters registered as nnx.Param in this module.

This method extracts the module parameter state and sums the number of scalar entries across all leaves to compute the total number of trainable parameters.

Parameters:

None –

Returns:

Total number of trainable scalar parameters.

Return type:

int

Raises:

None –

Forward(x, nn_params)

Forward pass through the MLP without skip connections.

For each hidden layer, this method applies an affine transform followed by the configured activation function (and gain if applicable). The final layer is applied as a linear transform without an activation.

Parameters:

• x (jax.Array) – Input array to the network.

• nn_params (list[tuple[nnx.Param, nnx.Param]]) – Sequence of (weights, biases) for each layer.

Returns:

Network output after applying all layers.

Return type:

jax.Array

Raises:

ValueError – If parameter shapes do not match the input activation shapes.

ForwardSkip(x, nn_params)

Forward pass through the MLP with periodic skip connections.

Skip connections are applied by concatenating the original input to the network with the current activation every skip_connections_settings["frequency"] layers (after the first layer). Hidden layers apply the configured activation, while the last layer remains linear.

Parameters:

• x (jax.Array) – Input array to the network (before Fourier mapping if enabled).

• nn_params (list[tuple[nnx.Param, nnx.Param]]) – Sequence of (weights, biases) for each layer.

Returns:

Network output after applying all layers with skip connections.

Return type:

jax.Array

Raises:

• KeyError – If "frequency" is missing from skip_connections_settings.

• ValueError – If parameter shapes do not match the concatenated activation shapes.

GetName()

Return the name identifier of this MLP instance.

Parameters:

None –

Returns:

Name of the network instance.

Return type:

str

Raises:

None –

InitialNetworkParameters()

Initialize network parameters and optional Fourier feature mapping.

This method constructs the layer size list, initializes each layer weight matrix and bias vector using layer_init_factopry(), and stores the parameters in self.nn_params. When Fourier features are enabled, it also creates the frequency matrix B and defines self.input_mapping to map the input coordinates to sinusoidal features.

The initialization accounts for skip connections by increasing the input dimension of layers that receive concatenated inputs.

Parameters:

None –

Returns:

None

Raises:

• KeyError – If required entries are missing from activation_settings or fourier_feature_settings.

• ValueError – If hidden layer sizes are invalid or if Fourier feature dimensions are inconsistent with input_size.

fol.deep_neural_networks.nns.layer_init_factopry(key, in_dim, out_dim, activation_settings)

Initialize a layer weight matrix and bias vector.

This helper selects an initialization scheme based on the requested activation type. For sinusoidal activations it delegates to siren_init(). For other activations it chooses a standard initializer commonly used for stable training.

Parameters:

• key (Array) – PRNG key used to sample random initial parameters.

• in_dim (int) – Input feature dimension of the layer.

• out_dim (int) – Output feature dimension of the layer.

• activation_settings (dict) – Activation configuration dictionary. Must include the key "type". For "type" == "sin", the settings must also include the SIREN parameters required by siren_init().

Returns:

A tuple (weights, biases) where weights has shape (in_dim, out_dim) and biases has shape (out_dim,).

Return type:

Tuple[Array, Array]

Raises:

• KeyError – If "type" is missing from activation_settings or if required SIREN keys are missing when type == "sin".

• ValueError – If in_dim or out_dim is not positive.

fol.deep_neural_networks.nns.siren_init(key, in_dim, out_dim, activation_settings)

Initialize weights and biases for a SIREN (sinusoidal) layer.

This initializer is designed for sinusoidal representation networks and chooses a uniform distribution bound that depends on the layer index and the SIREN frequency parameter (omega). The behavior matches common SIREN initialization practice and supports gain scaling strategies.

References

Sitzmann et al. (2020), “Implicit neural representations with periodic activation functions”, NeurIPS. Yeom et al. (2024), “Fast Training of Sinusoidal Neural Fields via Scaling Initialization”, arXiv:2410.04779.

Parameters:

• key (Array) – PRNG key used to sample random initial parameters.

• in_dim (int) – Input feature dimension of the layer.

• out_dim (int) – Output feature dimension of the layer.

• activation_settings (dict) – Dictionary containing SIREN initialization parameters. Required keys are "current_layer_idx", "total_num_layers", "initialization_gain", and "prediction_gain".

Returns:

A tuple (weights, biases) where weights has shape (in_dim, out_dim) and biases has shape (out_dim,).

Return type:

Tuple[Array, Array]

Raises:

• KeyError – If any required key is missing from activation_settings.

• ValueError – If in_dim or out_dim is not positive, or if total_num_layers is inconsistent with current_layer_idx.