Rescale TGLFNN predictive variances correctly

fusion_surrogates/tglfnn_ukaea/tglfnn_ukaea_model.py

@@ -162,18 +162,25 @@ def predict(self, inputs: jax.Array) -> Mapping[str, jax.Array]:
     # Combine the ensemble members
     mean = jnp.mean(predictions[..., 0], axis=1)  # pyrefly: ignore[bad-index]
     # Aleatoric uncertainty = mean of the predicted variances
-    aleatoric = jnp.mean(predictions[..., 1], axis=1)  # pyrefly: ignore[bad-index]
+    aleatoric = jnp.mean(
+        predictions[..., 1], axis=1
+    )  # pyrefly: ignore[bad-index]
     # Epistemic uncertainty = variance of the predicted means
-    epistemic = jnp.var(predictions[..., 0], axis=1)  # pyrefly: ignore[bad-index]
+    epistemic = jnp.var(
+        predictions[..., 0], axis=1
+    )  # pyrefly: ignore[bad-index]
     normalized_predictions = jnp.stack([mean, aleatoric + epistemic], axis=-1)
 
+    # Variances have no mean offset and scale with the squared stddev.
     broadcast_means = jnp.expand_dims(
-        self._output_means,
-        axis=tuple(range(1, normalized_predictions.ndim)),
+        jnp.stack(
+            [self._output_means, jnp.zeros_like(self._output_means)], axis=-1
+        ),
+        axis=tuple(range(1, normalized_predictions.ndim - 1)),
     )
     broadcast_stds = jnp.expand_dims(
-        self._output_stds,
-        axis=tuple(range(1, normalized_predictions.ndim)),
+        jnp.stack([self._output_stds, self._output_stds**2], axis=-1),
+        axis=tuple(range(1, normalized_predictions.ndim - 1)),
     )
 
     # Unnormalize and explicitly cast back to original dtype.

fusion_surrogates/tglfnn_ukaea/tglfnn_ukaea_model_test.py

@@ -14,10 +14,14 @@
 
 """Tests for UKAEA's TGLFNN surrogate."""
 
+from unittest import mock
+
 from absl.testing import absltest
 from absl.testing import parameterized
 from fusion_surrogates.tglfnn_ukaea import tglfnn_ukaea_model
+import jax
 import jax.numpy as jnp
+import numpy as np
 
 
 class TGLFNNukaeaModelTest(parameterized.TestCase):
@@ -53,6 +57,57 @@ def test_predict_shape(self, input_shape, expected_output_shape):
     for label in model.output_labels:
       self.assertEqual(predictions[label].shape, expected_output_shape)
 
+  @parameterized.product(
+      input_shape=((2,), (4, 2), (2, 3, 2)),
+      input_dtype=(jnp.float32, jnp.float64),
+      use_jit=(False, True),
+  )
+  def test_predict_rescales_mean_and_variance(
+      self, input_shape, input_dtype, use_jit
+  ):
+    # Two equally weighted members have means 1 and 3, variance log(2).
+    model_dict = {
+        "input_labels": ["input_a", "input_b"],
+        "config": {"num_estimators": 2, "model_size": 1},
+        "params": {
+            label: {
+                f"MLP_{i}": {
+                    "FullyConnectedLayer_0": {
+                        "weight": np.zeros((2, 2)),
+                        "bias": np.array([mean, 0.0]),
+                    }
+                }
+                for i, mean in enumerate((1.0, 3.0))
+            }
+            for label in ("flux_a", "flux_b")
+        },
+        "stats": {
+            "input_a": {"mean": 0.0, "std": 1.0},
+            "input_b": {"mean": 0.0, "std": 1.0},
+            "flux_a": {"mean": -10.0, "std": 3.0},
+            "flux_b": {"mean": 5.0, "std": 0.0},
+        },
+    }
+    self.addCleanup(jax.config.update, "jax_enable_x64", jax.config.x64_enabled)
+    jax.config.update("jax_enable_x64", True)
+    with mock.patch.object(
+        tglfnn_ukaea_model.tglfnn_ukaea_lib,
+        "load",
+        return_value=model_dict,
+    ):
+      model = tglfnn_ukaea_model.TGLFNNukaeaModel()
+    predict = jax.jit(model.predict) if use_jit else model.predict
+    predictions = predict(jnp.ones(input_shape, dtype=input_dtype))
+    for label, offset, scale in (("flux_a", -10.0, 3.0), ("flux_b", 5.0, 1.0)):
+      result = predictions[label]
+      self.assertEqual(result.shape, input_shape[:-1] + (2,))
+      self.assertEqual(result.dtype, input_dtype)
+      np.testing.assert_allclose(result[..., 0], 2.0 * scale + offset)
+      np.testing.assert_allclose(
+          result[..., 1], (1.0 + np.log(2.0)) * scale**2, rtol=1e-6
+      )
+      self.assertTrue(np.all(result[..., 1] >= 0.0))
+
 
 if __name__ == "__main__":
   absltest.main()