Writing Model Classes

Writing Model Classes#

In Getting Started and Fitting a line to data, we saw two examples of how simpple can be used to quickly build a model from a likelihood function and, optionally, a forward model. For more complex models, it may be preferable to write a custom class that can store additional information.

In this tutorial, we will implement a flexible polynomial model with a custom class, to show a simple example of how this API can be used.

Notice the following:

  • Only two methods are absolutely required: _forward() and _log_likelihood(). They must accept a dictionary of parameters as their first argument.

  • __init__() is where the model is built. It should accept a parameter dictionary or initialize it internally, and then call super().__init__() to run the usual boilerplate for simpple models. This is required so that the fixed and variable parameters are properly handled!.

  • __init__() is also where you should run additional checks and create extra attributes, such as order in the example below.

  • Model attributes can be re-used in the forward model or the log-likelihood. This is powerful if you want to store data or have configurable parameters (such as order here).

from simpple.model import ForwardModel
import simpple.distributions as sdist
import numpy as np


class PolyModel(ForwardModel):
    def __init__(self, parameters: dict[str, sdist.Distribution], order: int):
        super().__init__(parameters)
        self.order = order
        for i in range(self.order + 1):
            k = "a" + str(i)
            if k not in self.parameters:
                raise KeyError(
                    f"Parameters should have keys from a0 to a{self.order} for polynomial of order {self.order}. Key {k} not found."
                )

    def _forward(self, p, x):
        parr = np.array([p[f"a{i}"] for i in range(self.order + 1)])
        return np.vander(x, self.order + 1, increasing=True) @ parr

    def _log_likelihood(self, p, x, y, yerr):
        ymod = self.forward(p, x)
        var = yerr**2 + p["sigma"] ** 2
        return -0.5 * np.sum(np.log(2 * np.pi * var) + (y - ymod) ** 2 / var)

Let us now test our class by creating a model and calling it.

parameters = {
    "a0": sdist.Uniform(-10, 10),
    "a1": sdist.Uniform(-10, 10),
    "a2": sdist.Uniform(-10, 10),
    "sigma": sdist.LogUniform(1e-5, 100),
}
model = PolyModel(parameters, 2)
test_p = [0.0, 1.0, 2.0, 0.0]
test_x = np.array([1, 2, 3, 4])
test_y = model.forward(test_p, test_x)

And let us double check the implementation with Numpy.

np_y = np.polynomial.Polynomial(test_p)(test_x)
assert np.all(test_y == np_y)

That’s it! We could then re-use this model (with order 1) as a drop-in replacement in the Fitting a line to data notebook.