Add DirectionReconstructionWithLogKappa

src/graphnet/models/task/reconstruction.py

@@ -1,5 +1,7 @@
 """Reconstruction-specific `Model` class(es)."""
 
+from typing import Any
+
 import numpy as np
 import torch
 from torch import Tensor
@@ -47,7 +49,14 @@ def _forward(self, x: Tensor) -> Tensor:
 
 
 class DirectionReconstructionWithKappa(StandardLearnedTask):
-    """Reconstructs direction with kappa from the 3D-vMF distribution."""
+    """Reconstructs direction with kappa from the 3D-vMF distribution.
+
+    Kappa is the length of the three outputs, so it can only grow as far as
+    the head's output scale does in training. Trained models hit a kappa
+    ceiling (~1e3 in practice) that bright events exceed, and their
+    predicted distributions come out underconfident. Prefer
+    `DirectionReconstructionWithLogKappa`, which predicts kappa separately.
+    """
 
     # Requires three features: untransformed points in (x,y,z)-space.
     default_target_labels = ["direction"]  # contains dir_x, dir_y, dir_z
@@ -69,6 +78,49 @@ def _forward(self, x: Tensor) -> Tensor:
         return torch.stack((vec_x, vec_y, vec_z, kappa), dim=1)
 
 
+class DirectionReconstructionWithLogKappa(StandardLearnedTask):
+    """Reconstructs direction with log-kappa from the 3D-vMF distribution.
+
+    The direction is the normalised first three outputs and
+    `kappa = kappa_min + exp(x[:, 3])`. The floor `kappa_min` keeps training
+    from settling on the uniform distribution at initialisation, where the
+    direction's gradient would otherwise vanish with kappa.
+    """
+
+    default_target_labels = ["direction"]
+    default_prediction_labels = [
+        "dir_x_pred",
+        "dir_y_pred",
+        "dir_z_pred",
+        "direction_kappa",
+    ]
+    nb_inputs = 4
+    _log_kappa_max = 20.0
+
+    def __init__(self, *args: Any, kappa_min: float = 1.0, **kwargs: Any):
+        """Construct `DirectionReconstructionWithLogKappa`.
+
+        Args:
+            *args: Passed to `StandardLearnedTask`.
+            kappa_min: Lower bound on the predicted kappa; must be positive.
+            **kwargs: Passed to `StandardLearnedTask`.
+        """
+        if kappa_min <= 0:
+            raise ValueError(f"kappa_min must be positive, got {kappa_min}")
+        super().__init__(*args, **kwargs)
+        self._kappa_min = kappa_min
+
+    def _forward(self, x: Tensor) -> Tensor:
+        direction = x[:, :3] / (
+            torch.linalg.vector_norm(x[:, :3], dim=1, keepdim=True)
+            + eps_like(x)
+        )
+        kappa = self._kappa_min + torch.exp(
+            x[:, 3:4].clamp(max=self._log_kappa_max)
+        )
+        return torch.cat((direction, kappa), dim=1)
+
+
 class ZenithReconstruction(StandardLearnedTask):
     """Reconstructs zenith angle."""
 

tests/models/test_task.py

@@ -1,13 +1,18 @@
 """Unit tests for Task classes."""
 
+from typing import Any
+
 import pytest
 import torch
 
 from graphnet.data.constants import FEATURES
 from graphnet.models.detector.icecube import IceCube86
 from graphnet.models.gnn import DynEdge
-from graphnet.models.task.reconstruction import EnergyReconstruction
-from graphnet.training.loss_functions import LogCoshLoss
+from graphnet.models.task.reconstruction import (
+    DirectionReconstructionWithLogKappa,
+    EnergyReconstruction,
+)
+from graphnet.training.loss_functions import LogCoshLoss, VonMisesFisher3DLoss
 from graphnet.models.graphs import KNNGraph
 from graphnet.models.graphs.nodes import NodesAsPulses
 
@@ -40,3 +45,52 @@ def test_transform_prediction_and_target() -> None:
             transform_target=torch.log10,
             transform_inference=torch.log10,
         )
+
+
+def _log_kappa_task(**kwargs: Any) -> DirectionReconstructionWithLogKappa:
+    return DirectionReconstructionWithLogKappa(
+        hidden_size=8,
+        target_labels="direction",
+        loss_function=VonMisesFisher3DLoss(),
+        **kwargs,
+    )
+
+
+def test_log_kappa_direction_output() -> None:
+    """Unit directions, kappa above its floor, the stock output layout."""
+    task = _log_kappa_task(kappa_min=1.0)
+    assert task.nb_inputs == 4
+    x = torch.randn(64, 4)
+    pred = task._forward(x)
+    assert pred.shape == (64, 4)
+    torch.testing.assert_close(
+        torch.linalg.vector_norm(pred[:, :3], dim=1), torch.ones(64)
+    )
+    torch.testing.assert_close(pred[:, 3], 1.0 + torch.exp(x[:, 3]))
+
+
+def test_log_kappa_no_ceiling_and_no_overflow() -> None:
+    """Kappa grows past the stock head's typical ceiling and stays finite."""
+    task = _log_kappa_task()
+    x = torch.tensor([[1.0, 0.0, 0.0, 10.0], [0.0, 1.0, 0.0, 1e4]])
+    kappa = task._forward(x)[:, 3]
+    assert kappa[0] > 2e4
+    assert torch.isfinite(kappa).all()
+
+
+def test_log_kappa_floor_must_be_positive() -> None:
+    """A zero floor lets training collapse, so it is rejected."""
+    with pytest.raises(ValueError):
+        _log_kappa_task(kappa_min=0.0)
+
+
+def test_log_kappa_loss_and_gradients() -> None:
+    """The vMF loss is finite and its gradient reaches all four outputs."""
+    task = _log_kappa_task()
+    x = torch.randn(32, 4, requires_grad=True)
+    target = torch.nn.functional.normalize(torch.randn(32, 3), dim=1)
+    loss = task._loss_function(task._forward(x), target)
+    assert torch.isfinite(loss)
+    loss.backward()
+    assert x.grad is not None
+    assert (x.grad.abs().sum(dim=0) > 0).all()