Repository: mllam/neural-lam
Issue category: Feature Request

Issue:
Feature Request: Add Functionality to Apply Constraints to Predictions
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 


Hints:
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.
@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?
I agree with both of you, and that would be highly appreciated!
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!

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index b16ee732..45bb97c9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -15,6 +15,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Add support for multi-node training.
 [\#103](https://github.com/mllam/neural-lam/pull/103) @simonkamuk @sadamov
 
+- Add option to clamp output prediction using limits specified in config file [\#92](https://github.com/mllam/neural-lam/pull/92) @SimonKamuk
+
 ### Fixed
 - Only print on rank 0 to avoid duplicates of all print statements.
 [\#103](https://github.com/mllam/neural-lam/pull/103) @simonkamuk @sadamov
diff --git a/README.md b/README.md
index d3e7379c..dfe9a137 100644
--- a/README.md
+++ b/README.md
@@ -148,12 +148,23 @@ training:
     weights:
       u100m: 1.0
       v100m: 1.0
-```
-
-For now the neural-lam config only defines two things: 1) the kind of data
-store and the path to its config, and 2) the weighting of different features in
-the loss function. If you don't define the state feature weighting it will default
-to weighting all features equally.
+      t2m: 1.0
+      r2m: 1.0
+  output_clamping:
+    lower:
+      t2m: 0.0
+      r2m: 0
+    upper:
+      r2m: 1.0
+```
+
+For now the neural-lam config only defines few things:
+
+1. The kind of datastore and the path to its config
+2. The weighting of different features in
+the loss function. If you don't define the state feature weighting it will default to
+weighting all features equally.
+3. Valid numerical range for output of each feature.The numerical range of all features default to $]-\infty, \infty[$.
 
 (This example is taken from the `tests/datastore_examples/mdp` directory.)
 
diff --git a/neural_lam/config.py b/neural_lam/config.py
index d3e09697..9deb15d6 100644
--- a/neural_lam/config.py
+++ b/neural_lam/config.py
@@ -68,6 +68,23 @@ class UniformFeatureWeighting:
     pass
 
 
+@dataclasses.dataclass
+class OutputClamping:
+    """
+    Configuration for clamping the output of the model.
+
+    Attributes
+    ----------
+    lower : Dict[str, float]
+        The minimum value to clamp each output feature to.
+    upper : Dict[str, float]
+        The maximum value to clamp each output feature to.
+    """
+
+    lower: Dict[str, float] = dataclasses.field(default_factory=dict)
+    upper: Dict[str, float] = dataclasses.field(default_factory=dict)
+
+
 @dataclasses.dataclass
 class TrainingConfig:
     """
@@ -86,6 +103,10 @@ class TrainingConfig:
         ManualStateFeatureWeighting, UniformFeatureWeighting
     ] = dataclasses.field(default_factory=UniformFeatureWeighting)
 
+    output_clamping: OutputClamping = dataclasses.field(
+        default_factory=OutputClamping
+    )
+
 
 @dataclasses.dataclass
 class NeuralLAMConfig(dataclass_wizard.JSONWizard, dataclass_wizard.YAMLWizard):
diff --git a/neural_lam/models/base_graph_model.py b/neural_lam/models/base_graph_model.py
index 728a9489..35b1ab12 100644
--- a/neural_lam/models/base_graph_model.py
+++ b/neural_lam/models/base_graph_model.py
@@ -79,6 +79,192 @@ def __init__(self, args, config: NeuralLAMConfig, datastore: BaseDatastore):
             layer_norm=False,
         )  # No layer norm on this one
 
+        # Compute indices and define clamping functions
+        self.prepare_clamping_params(config, datastore)
+
+    def prepare_clamping_params(
+        self, config: NeuralLAMConfig, datastore: BaseDatastore
+    ):
+        """
+        Prepare parameters for clamping predicted values to valid range
+        """
+
+        # Read configs
+        state_feature_names = datastore.get_vars_names(category="state")
+        lower_lims = config.training.output_clamping.lower
+        upper_lims = config.training.output_clamping.upper
+
+        # Check that limits in config are for valid features
+        unknown_features_lower = set(lower_lims.keys()) - set(
+            state_feature_names
+        )
+        unknown_features_upper = set(upper_lims.keys()) - set(
+            state_feature_names
+        )
+        if unknown_features_lower or unknown_features_upper:
+            raise ValueError(
+                "State feature limits were provided for unknown features: "
+                f"{unknown_features_lower.union(unknown_features_upper)}"
+            )
+
+        # Constant parameters for clamping
+        sigmoid_sharpness = 1
+        softplus_sharpness = 1
+        sigmoid_center = 0
+        softplus_center = 0
+
+        normalize_clamping_lim = (
+            lambda x, feature_idx: (x - self.state_mean[feature_idx])
+            / self.state_std[feature_idx]
+        )
+
+        # Check which clamping functions to use for each feature
+        sigmoid_lower_upper_idx = []
+        sigmoid_lower_lims = []
+        sigmoid_upper_lims = []
+
+        softplus_lower_idx = []
+        softplus_lower_lims = []
+
+        softplus_upper_idx = []
+        softplus_upper_lims = []
+
+        for feature_idx, feature in enumerate(state_feature_names):
+            if feature in lower_lims and feature in upper_lims:
+                assert (
+                    lower_lims[feature] < upper_lims[feature]
+                ), f'Invalid clamping limits for feature "{feature}",\
+                     lower: {lower_lims[feature]}, larger than\
+                     upper: {upper_lims[feature]}'
+                sigmoid_lower_upper_idx.append(feature_idx)
+                sigmoid_lower_lims.append(
+                    normalize_clamping_lim(lower_lims[feature], feature_idx)
+                )
+                sigmoid_upper_lims.append(
+                    normalize_clamping_lim(upper_lims[feature], feature_idx)
+                )
+            elif feature in lower_lims and feature not in upper_lims:
+                softplus_lower_idx.append(feature_idx)
+                softplus_lower_lims.append(
+                    normalize_clamping_lim(lower_lims[feature], feature_idx)
+                )
+            elif feature not in lower_lims and feature in upper_lims:
+                softplus_upper_idx.append(feature_idx)
+                softplus_upper_lims.append(
+                    normalize_clamping_lim(upper_lims[feature], feature_idx)
+                )
+
+        self.register_buffer(
+            "sigmoid_lower_lims", torch.tensor(sigmoid_lower_lims)
+        )
+        self.register_buffer(
+            "sigmoid_upper_lims", torch.tensor(sigmoid_upper_lims)
+        )
+        self.register_buffer(
+            "softplus_lower_lims", torch.tensor(softplus_lower_lims)
+        )
+        self.register_buffer(
+            "softplus_upper_lims", torch.tensor(softplus_upper_lims)
+        )
+
+        self.register_buffer(
+            "clamp_lower_upper_idx", torch.tensor(sigmoid_lower_upper_idx)
+        )
+        self.register_buffer(
+            "clamp_lower_idx", torch.tensor(softplus_lower_idx)
+        )
+        self.register_buffer(
+            "clamp_upper_idx", torch.tensor(softplus_upper_idx)
+        )
+
+        # Define clamping functions
+        self.clamp_lower_upper = lambda x: (
+            self.sigmoid_lower_lims
+            + (self.sigmoid_upper_lims - self.sigmoid_lower_lims)
+            * torch.sigmoid(sigmoid_sharpness * (x - sigmoid_center))
+        )
+        self.clamp_lower = lambda x: (
+            self.softplus_lower_lims
+            + torch.nn.functional.softplus(
+                x - softplus_center, beta=softplus_sharpness
+            )
+        )
+        self.clamp_upper = lambda x: (
+            self.softplus_upper_lims
+            - torch.nn.functional.softplus(
+                softplus_center - x, beta=softplus_sharpness
+            )
+        )
+
+        self.inverse_clamp_lower_upper = lambda x: (
+            sigmoid_center
+            + utils.inverse_sigmoid(
+                (x - self.sigmoid_lower_lims)
+                / (self.sigmoid_upper_lims - self.sigmoid_lower_lims)
+            )
+            / sigmoid_sharpness
+        )
+        self.inverse_clamp_lower = lambda x: (
+            utils.inverse_softplus(
+                x - self.softplus_lower_lims, beta=softplus_sharpness
+            )
+            + softplus_center
+        )
+        self.inverse_clamp_upper = lambda x: (
+            -utils.inverse_softplus(
+                self.softplus_upper_lims - x, beta=softplus_sharpness
+            )
+            + softplus_center
+        )
+
+    def get_clamped_new_state(self, state_delta, prev_state):
+        """
+        Clamp prediction to valid range supplied in config
+        Returns the clamped new state after adding delta to original state
+
+        Instead of the new state being computed as
+        $X_{t+1} = X_t + \\delta = X_t + model(\\{X_t,X_{t-1},...\\}, forcing)$
+        The clamped values will be
+        $f(f^{-1}(X_t) + model(\\{X_t, X_{t-1},... \\}, forcing))$
+        Which means the model will learn to output values in the range of the
+        inverse clamping function
+
+        state_delta: (B, num_grid_nodes, feature_dim)
+        prev_state: (B, num_grid_nodes, feature_dim)
+        """
+
+        # Assign new state, but overwrite clamped values of each type later
+        new_state = prev_state + state_delta
+
+        # Sigmoid/logistic clamps between ]a,b[
+        if self.clamp_lower_upper_idx.numel() > 0:
+            idx = self.clamp_lower_upper_idx
+
+            new_state[:, :, idx] = self.clamp_lower_upper(
+                self.inverse_clamp_lower_upper(prev_state[:, :, idx])
+                + state_delta[:, :, idx]
+            )
+
+        # Softplus clamps between ]a,infty[
+        if self.clamp_lower_idx.numel() > 0:
+            idx = self.clamp_lower_idx
+
+            new_state[:, :, idx] = self.clamp_lower(
+                self.inverse_clamp_lower(prev_state[:, :, idx])
+                + state_delta[:, :, idx]
+            )
+
+        # Softplus clamps between ]-infty,b[
+        if self.clamp_upper_idx.numel() > 0:
+            idx = self.clamp_upper_idx
+
+            new_state[:, :, idx] = self.clamp_upper(
+                self.inverse_clamp_upper(prev_state[:, :, idx])
+                + state_delta[:, :, idx]
+            )
+
+        return new_state
+
     def get_num_mesh(self):
         """
         Compute number of mesh nodes from loaded features,
@@ -173,5 +359,7 @@ def predict_step(self, prev_state, prev_prev_state, forcing):
         # Rescale with one-step difference statistics
         rescaled_delta_mean = pred_delta_mean * self.diff_std + self.diff_mean
 
-        # Residual connection for full state
-        return prev_state + rescaled_delta_mean, pred_std
+        # Clamp values to valid range (also add the delta to the previous state)
+        new_state = self.get_clamped_new_state(rescaled_delta_mean, prev_state)
+
+        return new_state, pred_std
diff --git a/neural_lam/utils.py b/neural_lam/utils.py
index ab213453..b1a277ea 100644
--- a/neural_lam/utils.py
+++ b/neural_lam/utils.py
@@ -307,3 +307,36 @@ def setup_training_logger(datastore, args, run_name):
         )
 
     return logger
+
+
+def inverse_softplus(x, beta=1, threshold=20):
+    """
+    Inverse of torch.nn.functional.softplus
+
+    For x*beta above threshold, returns linear function for numerical
+    stability.
+
+    Input is clamped to x > ln(1+1e-6)/beta which is approximately positive
+    values of x.
+    Note that this torch.clamp_min will make gradients 0, but this is not a
+    problem as values of x that are this close to 0 have gradients of 0 anyhow.
+    """
+    non_linear_part = (
+        torch.log(torch.clamp_min(torch.expm1(x * beta), 1e-6)) / beta
+    )
+    x = torch.where(x * beta <= threshold, non_linear_part, x)
+
+    return x
+
+
+def inverse_sigmoid(x):
+    """
+    Inverse of torch.sigmoid
+
+    Sigmoid output takes values in [0,1], this makes sure input is just within
+    this interval.
+    Note that this torch.clamp will make gradients 0, but this is not a problem
+    as values of x that are this close to 0 or 1 have gradients of 0 anyhow.
+    """
+    x_clamped = torch.clamp(x, min=1e-6, max=1 - 1e-6)
+    return torch.log(x_clamped / (1 - x_clamped))


Test patch:
diff --git a/tests/datastore_examples/mdp/danra_100m_winds/config.yaml b/tests/datastore_examples/mdp/danra_100m_winds/config.yaml
index 0bb5c5ec..8b3362e0 100644
--- a/tests/datastore_examples/mdp/danra_100m_winds/config.yaml
+++ b/tests/datastore_examples/mdp/danra_100m_winds/config.yaml
@@ -7,3 +7,12 @@ training:
     weights:
       u100m: 1.0
       v100m: 1.0
+      t2m: 1.0
+      r2m: 1.0
+  output_clamping:
+    lower:
+      t2m: 0.0
+      r2m: 0
+    upper:
+      r2m: 1.0
+      u100m: 100.0
diff --git a/tests/datastore_examples/mdp/danra_100m_winds/danra.datastore.yaml b/tests/datastore_examples/mdp/danra_100m_winds/danra.datastore.yaml
index 3f9bafac..876b1405 100644
--- a/tests/datastore_examples/mdp/danra_100m_winds/danra.datastore.yaml
+++ b/tests/datastore_examples/mdp/danra_100m_winds/danra.datastore.yaml
@@ -55,7 +55,7 @@ inputs:
         dims: [x, y]
     target_output_variable: state
 
-  danra_surface:
+  danra_surface_forcing:
     path: https://object-store.os-api.cci1.ecmwf.int/mllam-testdata/danra_cropped/v0.2.0/single_levels.zarr
     dims: [time, x, y]
     variables:
@@ -73,6 +73,24 @@ inputs:
         name_format: "{var_name}"
     target_output_variable: forcing
 
+  danra_surface:
+    path: https://object-store.os-api.cci1.ecmwf.int/mllam-testdata/danra_cropped/v0.2.0/single_levels.zarr
+    dims: [time, x, y]
+    variables:
+      - r2m
+      - t2m
+    dim_mapping:
+      time:
+        method: rename
+        dim: time
+      grid_index:
+        method: stack
+        dims: [x, y]
+      state_feature:
+        method: stack_variables_by_var_name
+        name_format: "{var_name}"
+    target_output_variable: state
+
   danra_static:
     path: https://object-store.os-api.cci1.ecmwf.int/mllam-testdata/danra_cropped/v0.2.0/single_levels.zarr
     dims: [x, y]
diff --git a/tests/test_clamping.py b/tests/test_clamping.py
new file mode 100644
index 00000000..f3f9365d
--- /dev/null
+++ b/tests/test_clamping.py
@@ -0,0 +1,283 @@
+# Standard library
+from pathlib import Path
+
+# Third-party
+import torch
+
+# First-party
+from neural_lam import config as nlconfig
+from neural_lam.create_graph import create_graph_from_datastore
+from neural_lam.datastore.mdp import MDPDatastore
+from neural_lam.models.graph_lam import GraphLAM
+from tests.conftest import init_datastore_example
+
+
+def test_clamping():
+    datastore = init_datastore_example(MDPDatastore.SHORT_NAME)
+
+    graph_name = "1level"
+
+    graph_dir_path = Path(datastore.root_path) / "graph" / graph_name
+
+    if not graph_dir_path.exists():
+        create_graph_from_datastore(
+            datastore=datastore,
+            output_root_path=str(graph_dir_path),
+            n_max_levels=1,
+        )
+
+    class ModelArgs:
+        output_std = False
+        loss = "mse"
+        restore_opt = False
+        n_example_pred = 1
+        graph = graph_name
+        hidden_dim = 4
+        hidden_layers = 1
+        processor_layers = 2
+        mesh_aggr = "sum"
+        lr = 1.0e-3
+        val_steps_to_log = [1, 3]
+        metrics_watch = []
+        num_past_forcing_steps = 1
+        num_future_forcing_steps = 1
+
+    model_args = ModelArgs()
+
+    config = nlconfig.NeuralLAMConfig(
+        datastore=nlconfig.DatastoreSelection(
+            kind=datastore.SHORT_NAME, config_path=datastore.root_path
+        ),
+        training=nlconfig.TrainingConfig(
+            output_clamping=nlconfig.OutputClamping(
+                lower={"t2m": 0.0, "r2m": 0.0},
+                upper={"r2m": 1.0, "u100m": 100.0},
+            )
+        ),
+    )
+
+    model = GraphLAM(
+        args=model_args,
+        datastore=datastore,
+        config=config,
+    )
+
+    features = datastore.get_vars_names(category="state")
+    original_state = torch.zeros(1, 1, len(features))
+    zero_delta = original_state.clone()
+
+    # Get a state well within the bounds
+    original_state[:, :, model.clamp_lower_upper_idx] = (
+        model.sigmoid_lower_lims + model.sigmoid_upper_lims
+    ) / 2
+    original_state[:, :, model.clamp_lower_idx] = model.softplus_lower_lims + 10
+    original_state[:, :, model.clamp_upper_idx] = model.softplus_upper_lims - 10
+
+    # Get a delta that tries to push the state out of bounds
+    delta = torch.ones_like(zero_delta)
+    delta[:, :, model.clamp_lower_upper_idx] = (
+        model.sigmoid_upper_lims - model.sigmoid_lower_lims
+    ) / 3
+    delta[:, :, model.clamp_lower_idx] = -5
+    delta[:, :, model.clamp_upper_idx] = 5
+
+    # Check that a delta of 0 gives unchanged state
+    zero_prediction = model.get_clamped_new_state(zero_delta, original_state)
+    assert (abs(original_state - zero_prediction) < 1e-6).all().item()
+
+    # Make predictions towards bounds for each feature
+    prediction = zero_prediction.clone()
+    n_loops = 100
+    for i in range(n_loops):
+        prediction = model.get_clamped_new_state(delta, prediction)
+
+    # check that unclamped states are as expected
+    # delta is 1, so they should be 1*n_loops
+    assert (
+        (
+            abs(
+                prediction[
+                    :,
+                    :,
+                    list(
+                        set(range(len(features)))
+                        - set(model.clamp_lower_upper_idx.tolist())
+                        - set(model.clamp_lower_idx.tolist())
+                        - set(model.clamp_upper_idx.tolist())
+                    ),
+                ]
+                - n_loops
+            )
+            < 1e-6
+        )
+        .all()
+        .item()
+    )
+
+    # Check that clamped states are within bounds
+    # they should not be at the bounds but allow it due to numerical precision
+    assert (
+        (
+            model.sigmoid_lower_lims
+            <= prediction[:, :, model.clamp_lower_upper_idx]
+            <= model.sigmoid_upper_lims
+        )
+        .all()
+        .item()
+    )
+    assert (
+        (model.softplus_lower_lims <= prediction[:, :, model.clamp_lower_idx])
+        .all()
+        .item()
+    )
+    assert (
+        (prediction[:, :, model.clamp_upper_idx] <= model.softplus_upper_lims)
+        .all()
+        .item()
+    )
+
+    # Check that prediction is within bounds in original non-normalized space
+    unscaled_prediction = prediction * model.state_std + model.state_mean
+    features_idx = {f: i for i, f in enumerate(features)}
+    lower_lims = {
+        features_idx[f]: lim
+        for f, lim in config.training.output_clamping.lower.items()
+    }
+    upper_lims = {
+        features_idx[f]: lim
+        for f, lim in config.training.output_clamping.upper.items()
+    }
+    assert (
+        (
+            torch.tensor(list(lower_lims.values()))
+            <= unscaled_prediction[:, :, list(lower_lims.keys())]
+        )
+        .all()
+        .item()
+    )
+    assert (
+        (
+            unscaled_prediction[:, :, list(upper_lims.keys())]
+            <= torch.tensor(list(upper_lims.values()))
+        )
+        .all()
+        .item()
+    )
+
+    # Check that a prediction from a state starting outside the bounds is also
+    # pushed within bounds. 3 delta should be enough to give an initial state
+    # out of bounds so 5 is well outside
+    invalid_state = original_state + 5 * delta
+    assert (
+        not (
+            model.sigmoid_lower_lims
+            <= invalid_state[:, :, model.clamp_lower_upper_idx]
+            <= model.sigmoid_upper_lims
+        )
+        .any()
+        .item()
+    )
+    assert (
+        not (
+            model.softplus_lower_lims
+            <= invalid_state[:, :, model.clamp_lower_idx]
+        )
+        .any()
+        .item()
+    )
+    assert (
+        not (
+            invalid_state[:, :, model.clamp_upper_idx]
+            <= model.softplus_upper_lims
+        )
+        .any()
+        .item()
+    )
+    invalid_prediction = model.get_clamped_new_state(zero_delta, invalid_state)
+    assert (
+        (
+            model.sigmoid_lower_lims
+            <= invalid_prediction[:, :, model.clamp_lower_upper_idx]
+            <= model.sigmoid_upper_lims
+        )
+        .all()
+        .item()
+    )
+    assert (
+        (
+            model.softplus_lower_lims
+            <= invalid_prediction[:, :, model.clamp_lower_idx]
+        )
+        .all()
+        .item()
+    )
+    assert (
+        (
+            invalid_prediction[:, :, model.clamp_upper_idx]
+            <= model.softplus_upper_lims
+        )
+        .all()
+        .item()
+    )
+
+    # Above tests only check the upper sigmoid limit.
+    # Repeat to check lower sigmoid limit
+
+    # Make predictions towards bounds for each feature
+    prediction = zero_prediction.clone()
+    n_loops = 100
+    for i in range(n_loops):
+        prediction = model.get_clamped_new_state(-delta, prediction)
+
+    # Check that clamped states are within bounds
+    assert (
+        (
+            model.sigmoid_lower_lims
+            <= prediction[:, :, model.clamp_lower_upper_idx]
+            <= model.sigmoid_upper_lims
+        )
+        .all()
+        .item()
+    )
+
+    # Check that prediction is within bounds in original non-normalized space
+    assert (
+        (
+            torch.tensor(list(lower_lims.values()))
+            <= unscaled_prediction[:, :, list(lower_lims.keys())]
+        )
+        .all()
+        .item()
+    )
+    assert (
+        (
+            unscaled_prediction[:, :, list(upper_lims.keys())]
+            <= torch.tensor(list(upper_lims.values()))
+        )
+        .all()
+        .item()
+    )
+
+    # Check that a prediction from a state starting outside the bounds is also
+    # pushed within bounds. 3 delta should be enough to give an initial state
+    # out of bounds so 5 is well outside
+    invalid_state = original_state - 5 * delta
+    assert (
+        not (
+            model.sigmoid_lower_lims
+            <= invalid_state[:, :, model.clamp_lower_upper_idx]
+            <= model.sigmoid_upper_lims
+        )
+        .any()
+        .item()
+    )
+    invalid_prediction = model.get_clamped_new_state(zero_delta, invalid_state)
+    assert (
+        (
+            model.sigmoid_lower_lims
+            <= invalid_prediction[:, :, model.clamp_lower_upper_idx]
+            <= model.sigmoid_upper_lims
+        )
+        .all()
+        .item()
+    )