ADEPTModule

class adept.ADEPTModule(cfg)Source

Bases: object

This class is the base class for all the ADEPT modules. It defines the interface that all the ADEPT modules must implement so that the ergoExo class can call them in the right order.

Parameters:

cfg – The configuration dictionary

post_process(run_output: dict, td: str) → dictSource

This function is responsible for post-processing the results of the simulation. It is called after the simulation is run and the results are available.

Parameters:

• run_output (Dict) – The output of the simulation

• td (str) – The temporary directory where the results are stored

Returns:

A dictionary of the post-processed results. This can include the metrics, the xarray datasets, and any other information that is relevant to the simulation

write_units() → dictSource

This function is responsible for writing the units, normalizing constants, and other important physical quantities to a dictionary. This dictionary is then dumped to a yaml file and logged to mlflow by the ergoExo class.

Returns:

A dictionary of the units

init_diffeqsolve() → dictSource

This function is responsible for initializing the differential equation solver diffrax.diffeqsolve. It sets up the time quantities, the solver quantities, and the save function.

Returns:

A dictionary of the differential equation solver quantities

get_derived_quantities() → dictSource

This function is responsible for getting the derived quantities from the configuration dictionary. This is needed for running the simulation. These quantities do get logged to mlflow by the ergoExo class.

Returns:

An updated configuration dictionary

get_solver_quantities()Source

This function is responsible for getting the solver quantities from the configuration dictionary. This is needed for running the simulation. These quantities do NOT get logged to mlflow because they are often arrays

Returns:

An updated configuration dictionary

get_save_func()Source

This function is responsible for getting the save function for the differential equation solver. This is needed for running the simulation. This function lets you subsample your simulation state so as to not save the entire thing at every timestep.

This dictionary is set as a class attribute for the ADEPTModule and are used in the __call__ function

init_state_and_args()Source

This function initializes the state and the arguments that are required to run the simulation. The state is the initial conditions of the simulation and the arguments are often the drivers

These are set as class attributes for the ADEPTModule and are used in the __call__ function

init_modules() → dict[str, Module]Source

This function initializes the necessary (trainable) physics 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 usually `eqx.Module`s so that you can take derivatives against the (parameters of the) modules.

Returns:

Dict – A dictionary of the (trainable) modules that are required to run the simulation

__call__(trainable_modules: dict, args: dict)Source

Call self as a function.