ergoExo

class adept.ergoExo(mlflow_run_id: str | None = None, mlflow_nested: bool = False, parent_run_id: str | None = None)Source

Bases: object

This class is the main interface for running a simulation. It is responsible for calling all the ADEPT modules in the right order and logging parameters and results to mlflow.

This approach helps decouple the numerical solvers from the experiment management and facilitates the addition of new solvers

Typical usage is as follows

exoskeleton = ergoExo()
modules = exoskeleton.setup(cfg)
run_output, post_processing_output, mlflow_run_id = exoskeleton(modules, args=None)

If you are resuming an existing mlflow run, you can do the following

exoskeleton = ergoExo(mlflow_run_id=mlflow_run_id)
modules = exoskeleton.setup(cfg)
run_output, post_processing_output, mlflow_run_id = exoskeleton(modules, args=None)

If you are introducing a custom ADEPTModule, you can do the following

exoskeleton = ergoExo()
modules = exoskeleton.setup(cfg, exoskeleton_module=custom_module)
run_output, post_processing_output, mlflow_run_id = adept(modules, args=None)
setup(cfg: dict, adept_module: ADEPTModule | None = None) → dict[str, Module]Source

This function sets up the differentiable simulation by getting the chosen solver and setting it up At this point in time, the setup includes

  1. initializing the mlflow run and setting the runid or resuming an existing run

  2. getting the right ADEPTModule or using the one passed in. This gets assigned to self.adept_module.

  3. updating the config, units, derived quantities, and array config as defined by the ADEPTModule. It also dumps this information to the temporary directory, which will be logged later, and logging the parameters to mlflow

  4. initializing the state and args as defined by the ADEPTModule

  5. initializing the diffeqsolve as defined by the ADEPTModule

  6. initializing the necessary (trainable) physics modules as defined by the ADEPTModule

Parameters:

cfg – The configuration dictionary

Returns:

A dictionary of trainable modules (Dict[str, eqx.Module])

This is a dictionary of the (trainable) modules that are required to run the simulation. These can be modules that change the initial conditions, or the driver (boundary conditions), or the metric calculation. These modules are equinox modules in order to play nice with diffrax

__call__(modules: dict | None = None, args: dict | None = None, export=True) → tuple[Solution, dict, str]Source

This function is the main entry point for running a simulation. It takes a configuration dictionary and returns a diffrax.Solution object and a dictionary of datasets. It calls the self.adept_module’s __call__ function.

It is also responsible for logging the artifacts and metrics to mlflow.

Parameters:

  • modules (Dict(str, eqx.Module)) – The trainable modules that are required to run the simulation.

  • call (All the other parameters are static and initialized during the setup)

Returns:

a tuple of the run_output (diffrax.Solution), post_processing_output (Dict[str, xarray.dataset]), and the mlflow_run_id (str).

The run_output comes from the __call__ function of the self.adept_module. The post_processing_output comes from the post_process method of the self.adept_module. The mlflow_run_id is the id of the mlflow run that was created during the setup call or passed in during the initialization of the class

val_and_grad(modules: dict | None = None, args: dict | None = None, export=True) → tuple[float, dict, tuple[Solution, dict, str]]Source

This function is the value and gradient of the simulation. This is a very similar looking function to the __call__ function but calls the self.adept_module.vg rather than the self.adept_module.__call__.

It is also responsible for logging the artifacts and metrics to mlflow.

Parameters:

  • modules – The (trainable) modules that are required to run the simulation and take the gradient against.

  • call (All the other parameters are static and initialized during the setup)

Returns:

a tuple of the value (float), gradient (Dict), and a tuple of the run_output (diffrax.Solution), post_processing_output (Dict[str, xarray.dataset]), and the mlflow_run_id (str).

The value and gradient, and run_output come from the adept_module.vg function. The run_output is the same as that from __call__ function of the self.adept_module. The post_processing_output comes from the post_process method of the self.adept_module. The mlflow_run_id is the id of the mlflow run that was created during the setup call or passed in during the initialization