Feature Request: Add Functionality to Apply Constraints to Predictions

[sadamov]

I am proposing the addition of a new method to our model class, designed to apply constraints to predictions to ensure that the values fall within specified bounds. This functionality would be useful for maintaining the integrity of our model's predictions in scenarios where certain variables have predefined limits.

Proposed Method:

def apply_constraints(self, prediction):
    """
    Apply constraints to prediction to ensure values are within the
    specified bounds
    """
    for param, (min_val, max_val) in constants.PARAM_CONSTRAINTS.items():
        indices = self.variable_indices[param]
        for index in indices:
            # Apply clamping to ensure values are within the specified
            # bounds
            prediction[:, :, :, index] = torch.clamp(
                prediction[:, :, :, index],
                min=min_val,
                max=max_val if max_val is not None else float("inf"),
            )
    return prediction

Rationale:

Data Integrity: Ensuring that predictions adhere to real-world constraints is essential for the reliability of our model's outputs. This method would allow us to enforce these constraints directly on the prediction tensor.
Flexibility: By dynamically applying constraints based on the variable indices, we can maintain a high degree of flexibility in how we handle different variables with varying constraints.

The method could be added to the ARModel class, which is our primary model class.
The constants.PARAM_CONSTRAINTS dictionary, which maps variable names to their minimum and maximum values, should be used to determine the constraints for each variable.

PARAM_CONSTRAINTS = {
    "RELHUM": (0, 100),
    "CLCT": (0, 100),
    "TOT_PREC": (0, None),
}

This feature is closely related to #18

[joeloskarsson]

I would want to do this by constraining the model output itself, also for what is used during training (a clamp does not give gradient). Specifically, variables with constraints should be handled by rescaling a sigmoid output or a softplus (for > 0). As this applies to all models, it would be nice with it included on a high level, so it does not have to be handled separately in each model.

One option would be to have a function constrain_outputs in the ARModel class that applies these activation functions. Then everything that inherits from this can just apply this to the state after making a prediction.

[SimonKamuk]

@joeloskarsson @sadamov
I agree with Joel that simply clipping values would impact model training/performance, and that using activation functions is the solution. In the development meeting it sounded like no one was really working on this, so if you want I can make a PR to implement this?

[sadamov]

I agree with both of you, and that would be highly appreciated!

[joeloskarsson]

That would be wonderful @SimonKamuk! I had a though about this: As the model currently predicts residuals, it's not obvious how to do this in that case. Since the predicted variables are $X^{t} + \text{model}(\dots)$ we can't just throw e.g. a sigmoid on top of that One solution to this is to just stop predicting residuals and predicting full states. But that's quite a big change and maybe something we would want to experiment with before. Maybe you can come up with some other solution.

Another consideration here is that the variable constraints apply in the original scale of the data, not for the standardized data. So e.g. a constraint that a variable should be in $[0,1]$ does not mean that the corresponding model output should be in that range, but some other range $[a,b]$.

Happy to discuss this more!