Calibrator Plugins

The following tutorial demonstrates how one may distribute calibrators as plugins via calisim's extensible plugin system.

from calisim.data_model import (
	DistributionModel,
	ParameterDataType,
	ParameterSpecification,
)
from calisim.example_models import LotkaVolterraModel
from calisim.optimisation import OptimisationMethod, OptimisationMethodModel
from calisim.optimisation.optuna_wrapper import OptunaOptimisation as OptunaOptimisationBase
from calisim.statistics import MeanSquaredError
import optuna.samplers as opt_samplers
from importlib.metadata import entry_points
import numpy as np
import pandas as pd

import warnings
warnings.filterwarnings("ignore")

Registering Plugins

First, let us suppose that you have extended the Optuna wrapper within calisim and created the ExtendedOptunaOptimisation class. For instance, let’s add support for the NSGAIIISampler sampling algorithm within the specify() method.

class ExtendedOptunaOptimisation(OptunaOptimisationBase):

    def specify(self) -> None:
        """Specify the parameters of the model calibration procedure."""
        sampler_name = self.specification.method
        supported_samplers = dict(
            tpes=opt_samplers.TPESampler,
            cmaes=opt_samplers.CmaEsSampler,
            nsgaii=opt_samplers.NSGAIISampler,
            nsgaiii=opt_samplers.NSGAIIISampler, # Adding support for NSGAIIISampler.
            qmc=opt_samplers.QMCSampler,
            gp=opt_samplers.GPSampler,
        )
        sampler_class = supported_samplers.get(sampler_name, None)
        if sampler_class is None:
            raise ValueError(
                f"Unsupported Optuna sampler: {sampler_name}.",
                f"Supported Optuna samplers are {', '.join(supported_samplers)}",
            )
        sampler_kwargs = self.specification.method_kwargs
        if sampler_kwargs is None:
            sampler_kwargs = {}
        self.sampler = sampler_class(**sampler_kwargs)
        
        self.study = optuna.create_study(
            sampler=self.sampler,
            study_name=self.specification.experiment_name,
            directions=self.specification.directions,
            storage=self.specification.storage,
            load_if_exists=True,
        )

We can register ExtendedOptunaOptimisation as a new engine called optuna_extended under the OptimisationMethod calibration class. To do this, we must include ExtendedOptunaOptimisation as a plugin by modifying your pyproject.toml file.

[project.entry-points."calisim.external.optimisation"]
optuna_extended = "calisim.optimisation.optuna_wrapper:ExtendedOptunaOptimisation"

If you are using Poetry, you may need to add the following within your pyproject.toml instead:

[tool.poetry.plugins."calisim.external.optimisation"]
optuna_extended = "calisim.optimisation.optuna_wrapper:ExtendedOptunaOptimisation"

This assumes that the module path for ExtendedOptunaOptimisation is calisim.optimisation.optuna_wrapper. Naturally, the module path will vary depending on where your calibrator plugin is located. Let’s check that ExtendedOptunaOptimisation was successfully added as a plugin.

pd.DataFrame([ 
    { 
        "name": entrypoint.name, 
        "module": entrypoint.value.split(":")[0],
        "target": entrypoint.value.split(":")[1],  
        "group": entrypoint.group 
    }
    for entrypoint in entry_points().select(group="calisim.external.optimisation") 
])

	name	module	target	group
0	optuna_extended	calisim.optimisation.optuna_wrapper	ExtendedOptunaOptimisation	calisim.external.optimisation

It looks like the optuna_extended plugin associated with the ExtendedOptunaOptimisation class was successfully added under the calisim.external.optimisation group.

One can similarly add other plugins under different calisim modules. For instance, a plugin for the sensitivity module would be included under the calisim.external.sensitivity group. And so on.

Performing Calibration

Let’s run a calibration workflow using the optuna_extended engine. Let’s try calibrating an example model: LotkaVolterraModel.

model = LotkaVolterraModel()
observed_data = model.get_observed_data()
observed_data.head(5)

	year	lynx	hare
0	1900.0	4.0	30.0
1	1901.0	6.1	47.2
2	1902.0	9.8	70.2
3	1903.0	35.2	77.4
4	1904.0	59.4	36.3

We’ll define our parameter specification containing parameter names, data types, distribution types, and bounds.

parameter_spec = ParameterSpecification(
	parameters=[
		DistributionModel(
			name="alpha",
			distribution_name="uniform",
			distribution_args=[0.45, 0.55],
			data_type=ParameterDataType.CONTINUOUS,
		),
		DistributionModel(
			name="beta",
			distribution_name="uniform",
			distribution_args=[0.02, 0.03],
			data_type=ParameterDataType.CONTINUOUS,
		),
	]
)

We’ll define our objective function. We aim to minimise the discrepancy between simulated and observed data using the MeanSquaredError metric as our loss.

def objective(
	parameters: dict, simulation_id: str, observed_data: np.ndarray | None, t: pd.Series
) -> float | list[float]:
	simulation_parameters = dict(h0=34.0, l0=5.9, t=t, gamma=0.84, delta=0.026)

	for k in ["alpha", "beta"]:
		simulation_parameters[k] = parameters[k]

	simulated_data = model.simulate(simulation_parameters).lynx.values
	metric = MeanSquaredError()
	discrepancy = metric.calculate(observed_data, simulated_data)
	return discrepancy

We’ll define the specification for the Optuna optimisation procedure, making use of the NSGAIIISampler sampling algorithm. To do this, we’ll set method to nsgaiii.

specification = OptimisationMethodModel(
	experiment_name="optuna_extended_optimisation",
	parameter_spec=parameter_spec,
	observed_data=observed_data.lynx.values,
	method="nsgaiii",
	directions=["minimize"],
	n_iterations=100,
	calibration_func_kwargs=dict(t=observed_data.year),
)

Finally, let’s run the calibration workflow by instantiating an OptimisationMethod calibrator. Note that the engine will be set to optuna_extended.

calibrator = OptimisationMethod(
	calibration_func=objective, specification=specification, engine="optuna_extended"
)
type(calibrator.implementation)

calisim.optimisation.optuna_wrapper.ExtendedOptunaOptimisation

Finally, let’s check that our calibrator is using the NSGAIIISampler sampling algorithm.

calibrator.specify()
type(calibrator.implementation.sampler)

optuna.samplers._nsgaiii._sampler.NSGAIIISampler

Excellent. We can see that the OptimisationMethod calibration workflow is using the optuna_extended engine and ExtendedOptunaOptimisation class as defined by our newly installed plugin within the calisim.external.optimisation group.