Adding a custom model

In this section, we will walk through the key steps required to add a new model. As an example, we’ll demonstrate by adding a simple dummy model which, after training, returns a trivial potential, internal, and interaction energy.

The model interface

First, we create a DummyModel class. This class will be the one used to create the object model in the training script.

class DummyModel:
    def __init__(self, config: dict, data_dim: int, tau: float) -> None:
        self.data_dim = data_dim
        self.layers = config['energy']['model']['layers']
        self.config_optimizer = config['energy']['optim']

        # create energy models
        self.model_potential = MLP(self.layers)

Next, we need a method to load the dataset. There are several dataset classes already provided for different input data format. Refer to the Dataset module for more information on the different dataset classes.

For this tutorial, we will use the CouplingsDataset class, which returns information about the couplings and the density, and is the dataset used by the JKOnet* model.

def load_dataset(self, dataset_name: str) -> CouplingsDataset:
    return CouplingsDataset(dataset_name)

Here, dataset_name is the name of the dataset to be loaded. The dataset must be in the data/ directory. See also the Dataset module.

JAX is a functional framework and we use classes only to provide a common interface. To mantain the state, we use a separate state variable. Thus, we neeed to define a method to create the state for our energy models.

def create_state(self, rng: jax.random.PRNGKey) -> Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]:
    potential = create_train_state(
        rng, self.model_potential, get_optimizer(self.config_optimizer), self.data_dim)
    return potential, _, _

Note that we have three return values, one for each energy term. In this case, we are only using the potential term.

Then, we need to provide a loss function for the model. In this case, since it is a dummy model, we generate a loss function with no physical meaning. When we minimize the loss and update the parameters, the function will converge to a potential that always returns 0.

def loss(
    self,
    potential_params: dict,
    xs: jnp.ndarray,
    ys: jnp.ndarray,
) -> jnp.ndarray:
    return jnp.abs(jnp.sum(self.model_potential.apply({'params': potential_params}, ys)))

We must also define a train_step method. This function must contain the calculation of the loss and gradients, as well as the subsequent state update.

def train_step(
    self,
    state: Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState],
    sample: Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]
) -> Tuple[jnp.ndarray, Tuple[train_state.TrainState, train_state.TrainState, train_state.TrainState]]:

    xs, ys, t, ws, rho, rho_grad = sample
    potential, internal, interaction = state
    loss, grads = jax.value_and_grad(
            self.loss, argnums=(0, 1, 2))(
                potential.params, xs, ys)
    potential = potential.apply_gradients(grads=grads[0])
    return loss, (potential, _, _)

Finally, the model must provide methods to get the potential, beta, and interaction energies. In this case, we will return a trivial potential and interaction energy, and a beta of 0 using the methods: get_potential, get_beta, and get_interaction.

def get_potential(self, state):
    potential, _, _ = state
    return lambda x: potential.apply_fn({'params': potential.params}, x)

def get_beta(self, state):
    return 0.

def get_interaction(self, state):
    return lambda x: 0.

Adding the DummyModel to the solvers

Now that the model is ready, we need to add it to the list of available solvers.

For this, we edit the __init__.py file in the models directory:

class EnumMethod(Enum):
    ... # Other models here
    DUMMY = 'dummy-model'  # The dummy model we're adding.

def get_model(
    solver: EnumMethod,
    config: dict,
    data_dim: int,
    dt: float):

    if solver == EnumMethod.DUMMY:
        from models.jkonet_star import DummyModel
        cls = DummyModel

    # Other models retrieval logic here
    # ...

    return cls(config, data_dim, dt)

Adding a Colormap for the DummyModel

Finally, add a colormap specific to the model in the style.yaml file. This will be used when plotting the predictions.

# training
groundtruth:
  light: '#F1F1F1'
  dark: '#C7B7A3'
  marker: 'o'

jkonet-star:
  light: '#CDF5FD'
  dark: '#A0E9FF'
  marker: '+'

dummy-model:
  light: '#FFC1C1'
  dark: '#FF6666'
  marker: '+'

Great! You are now ready to train your first DummyModel!