Merge pull request #302 from asogaard/restructure-models

Restructure models

Author: asogaard

examples/train_model.py

@@ -93,6 +93,7 @@ def main():
     )
     gnn = DynEdge(
         nb_inputs=detector.nb_outputs,
+        global_pooling_schemes=["min", "max", "mean", "sum"],
     )
     task = EnergyReconstruction(
         hidden_size=gnn.nb_outputs,

src/graphnet/models/components/layers.py

@@ -1,4 +1,4 @@
-from typing import Callable, Optional, Sequence
+from typing import Callable, List, Optional, Sequence, Union
 
 from torch.functional import Tensor
 
@@ -13,7 +13,7 @@ def __init__(
         nn: Callable,
         aggr: str = "max",
         nb_neighbors: int = 8,
-        features_subset: Optional[Sequence] = None,
+        features_subset: Optional[Union[Sequence[int], List[int]]] = None,
         **kwargs,
     ):
         # Check(s)

src/graphnet/models/gnn/__init__.py

@@ -1,2 +1,3 @@
-from .dynedge import DynEdge, DynEdge_V2, DynEdge_V3
 from .convnet import ConvNet
+from .dynedge import DynEdge
+from .dynedge_jinst import DynEdgeJINST