Main Functionalities

NNFE

The top-level solver object. Construct it from a YAML config file with NNFE.from_yaml or from an NNFEConfig dataclass with NNFE.from_config.

nnfe.nnfe_object.NNFE

Top-level NNFE solver object.

Couples a finite-element problem with a neural network that is trained to minimise the FE residual over a prescribed parameter space. The class orchestrates:

• Problem definition — mesh, boundary conditions, and residual evaluation.
• Machine learning model — network, optimizer, and training state.
• Sampler — control-variable point sets for training and testing.
• Training and testing loops.
• Saving / loading model weights and configuration.

Prefer :meth:from_yaml or :meth:from_config for typical use. The direct constructor is available for cases where sub-components are built manually.

Parameters:

Name	Type	Description	Default
problem_manager	ProblemManager	Wraps the FE problem definition (mesh, BCs, and residual computation).	required
ml_manager	MLManager	Owns the neural network, optimizer, and training state.	required
sampler	Sampler	Provides training and testing point sets.	required
plotter	Plotter	Handles diagnostic plot generation.	required
project_manager	ProjectManager	Manages the run directory, keys, and saving.	required
nnfe_params	NNFEParamsConfig	Specifies which FE variables the network controls and their ordering in the output vector.	required
config	NNFEConfig	The full :class:~nnfe.nnfe_config.NNFEConfig used to build this object; stored for later serialisation.	None

from_yaml(path) classmethod

from_config(config: NNFEConfig) classmethod

Train

nnfe.nnfe_object.NNFE.train()

Main training loop for NNFE

Test

nnfe.nnfe_object.NNFE.test(x)

Test the accuracy of the model after training

Evaluate

nnfe.nnfe_object.NNFE.evaluate(x: typing.Union[jax.Array, numpy.ndarray, numpy.bool, numpy.number, bool, int, float, complex])

Evaluate the model at a given point x. Args: x: The input point to evaluate the model at. Returns: The output of the model at the input point x.

Save

nnfe.nnfe_object.NNFE.save()

Responsible for saving the model and other things May want to create a temp save and final save to split up training process just in case

Dump Config

nnfe.nnfe_object.NNFE.dump_config(save_dir: Path, filename: Path)

Dumps this specific manager's configuration to a YAML file.