Merge pull request #639 from graphnet-team/Task-Refactor

Task refactor

src/graphnet/models/standard_model.py

@@ -18,7 +18,7 @@
 from graphnet.models.graphs import GraphDefinition
 from graphnet.models.gnn.gnn import GNN
 from graphnet.models.model import Model
-from graphnet.models.task import Task
+from graphnet.models.task import StandardLearnedTask
 
 
 class StandardModel(Model):
@@ -33,7 +33,7 @@ def __init__(
         *,
         graph_definition: GraphDefinition,
         gnn: GNN,
-        tasks: Union[Task, List[Task]],
+        tasks: Union[StandardLearnedTask, List[StandardLearnedTask]],
         optimizer_class: Type[torch.optim.Optimizer] = Adam,
         optimizer_kwargs: Optional[Dict] = None,
         scheduler_class: Optional[type] = None,
@@ -45,10 +45,10 @@ def __init__(
         super().__init__(name=__name__, class_name=self.__class__.__name__)
 
         # Check(s)
-        if isinstance(tasks, Task):
+        if isinstance(tasks, StandardLearnedTask):
             tasks = [tasks]
         assert isinstance(tasks, (list, tuple))
-        assert all(isinstance(task, Task) for task in tasks)
+        assert all(isinstance(task, StandardLearnedTask) for task in tasks)
         assert isinstance(graph_definition, GraphDefinition)
         assert isinstance(gnn, GNN)
 

src/graphnet/models/task/__init__.py

@@ -1,3 +1,8 @@
 """Physics task-specific modules to be used as model "read-outs"."""
 
-from .task import Task, IdentityTask
+from .task import (
+    Task,
+    IdentityTask,
+    StandardLearnedTask,
+    StandardFlowTask,
+)

src/graphnet/models/task/classification.py

@@ -5,7 +5,7 @@
 import torch
 from torch import Tensor
 
-from graphnet.models.task import Task, IdentityTask
+from graphnet.models.task import IdentityTask, StandardLearnedTask
 
 
 class MulticlassClassificationTask(IdentityTask):
@@ -17,7 +17,7 @@ class MulticlassClassificationTask(IdentityTask):
     """
 
 
-class BinaryClassificationTask(Task):
+class BinaryClassificationTask(StandardLearnedTask):
     """Performs binary classification."""
 
     # Requires one feature, logit for being signal class.
@@ -30,7 +30,7 @@ def _forward(self, x: Tensor) -> Tensor:
         return torch.sigmoid(x)
 
 
-class BinaryClassificationTaskLogits(Task):
+class BinaryClassificationTaskLogits(StandardLearnedTask):
     """Performs binary classification form logits."""
 
     # Requires one feature, logit for being signal class.

src/graphnet/models/task/reconstruction.py

@@ -4,11 +4,11 @@
 import torch
 from torch import Tensor
 
-from graphnet.models.task import Task
+from graphnet.models.task import StandardLearnedTask
 from graphnet.utilities.maths import eps_like
 
 
-class AzimuthReconstructionWithKappa(Task):
+class AzimuthReconstructionWithKappa(StandardLearnedTask):
     """Reconstructs azimuthal angle and associated kappa (1/var)."""
 
     # Requires two features: untransformed points in (x,y)-space.
@@ -46,7 +46,7 @@ def _forward(self, x: Tensor) -> Tensor:
         return angle
 
 
-class DirectionReconstructionWithKappa(Task):
+class DirectionReconstructionWithKappa(StandardLearnedTask):
     """Reconstructs direction with kappa from the 3D-vMF distribution."""
 
     # Requires three features: untransformed points in (x,y,z)-space.
@@ -70,7 +70,7 @@ def _forward(self, x: Tensor) -> Tensor:
         return torch.stack((vec_x, vec_y, vec_z, kappa), dim=1)
 
 
-class ZenithReconstruction(Task):
+class ZenithReconstruction(StandardLearnedTask):
     """Reconstructs zenith angle."""
 
     # Requires two features: zenith angle itself.
@@ -98,7 +98,7 @@ def _forward(self, x: Tensor) -> Tensor:
         return torch.stack((angle, kappa), dim=1)
 
 
-class EnergyReconstruction(Task):
+class EnergyReconstruction(StandardLearnedTask):
     """Reconstructs energy using stable method."""
 
     # Requires one feature: untransformed energy
@@ -112,7 +112,7 @@ def _forward(self, x: Tensor) -> Tensor:
         return torch.nn.functional.softplus(x, beta=0.05) + eps_like(x)
 
 
-class EnergyReconstructionWithPower(Task):
+class EnergyReconstructionWithPower(StandardLearnedTask):
     """Reconstructs energy."""
 
     # Requires one feature: untransformed energy
@@ -125,7 +125,7 @@ def _forward(self, x: Tensor) -> Tensor:
         return torch.pow(10, x[:, 0] + 1.0).unsqueeze(1)
 
 
-class EnergyTCReconstruction(Task):
+class EnergyTCReconstruction(StandardLearnedTask):
     """Reconstructs track and cascade energies using stable method."""
 
     # Requires two features: untransformed energy for track and cascade
@@ -161,7 +161,7 @@ def _forward(self, x: Tensor) -> Tensor:
         return pred
 
 
-class VertexReconstruction(Task):
+class VertexReconstruction(StandardLearnedTask):
     """Reconstructs vertex position and time."""
 
     # Requires four features, x, y, z, and t.
@@ -183,7 +183,7 @@ def _forward(self, x: Tensor) -> Tensor:
         return x
 
 
-class PositionReconstruction(Task):
+class PositionReconstruction(StandardLearnedTask):
     """Reconstructs vertex position."""
 
     # Requires three features, x, y, and z.
@@ -204,7 +204,7 @@ def _forward(self, x: Tensor) -> Tensor:
         return x
 
 
-class TimeReconstruction(Task):
+class TimeReconstruction(StandardLearnedTask):
     """Reconstructs time."""
 
     # Requires one feature, time.
@@ -217,7 +217,7 @@ def _forward(self, x: Tensor) -> Tensor:
         return x
 
 
-class InelasticityReconstruction(Task):
+class InelasticityReconstruction(StandardLearnedTask):
     """Reconstructs interaction inelasticity.
 
     That is, 1-(track energy / hadronic energy).