initial commit v0

README.md

@@ -1,3 +1,12 @@
 ---
 license: mit
 ---
+
+HyperTrack: Neural Combinatorics for High Energy Physics
+
+https://arxiv.org/abs/2309.14113
+
+Download the main repository from:
+github.com/mieskolainen/hypertrack
+
+Download this repository to the same path.

hypertrack/models/global_tune-5.py

@@ -0,0 +1,308 @@
+# HyperTrack model and training loss parameters
+# 
+# match with the corresponding 'models_<ID>.py' under 'hypertrack/models/'
+#
+# [email protected], 2023
+
+import numpy as np
+
+# -------------------------------------------------------------------------
+# Input normalization
+# (e.g. can accelerate training, and mitigate float scale problems, but not necessarily needed)
+normalize_input = False
+
+"""
+- coord[0] (min,max,mean,std): -1025.3399658203125 | 1025.3399658203125 | 1.0586246252059937 | 266.20428466796875
+- coord[1] (min,max,mean,std): -1025.3399658203125 | 1025.3399658203125 | -0.022702794522047043 | 267.56085205078125
+- coord[2] (min,max,mean,std): -2955.5 | 2955.5 | 1.6228374242782593 | 1064.4954833984375
+"""
+
+def feature_scaler(X):
+    mu    = [1.06, -0.023, 1.62]
+    sigma = [266.2, 267.6, 1064.5]
+    
+    for i in range(len(mu)):
+        X[:,i] = (X[:,i] - mu[i]) / sigma[i]
+
+# -------------------------------------------------------------------------
+
+# ** Training only parameters **
+train_param = {
+    
+    # Total loss weights per each individual loss
+    'beta': {
+        
+        'net': {
+            'edge_BCE'           : 0.2,    # 0.2
+            'edge_contrastive'   : 1.0,    # 1.0
+            'cluster_BCE'        : 0.2,    # 0.2
+            'cluster_contrastive': 0.2,    # 1.0
+            'cluster_neglogpdf':   0.0,    # [EXPERIMENTAL] (keep it zero)
+        },
+        
+        'pdf': {
+            'track_neglogpdf':     1.0,    # [EXPERIMENTAL]
+        }
+    },
+    
+    # Edge loss
+    'edge_BCE': {
+        'type':            'Focal',    # 'Focal', 'BCE', 'BCE+Hinge'
+        'gamma':               1.0,    # For 'Focal' (entropy exponent)
+        'delta':              0.05,    # For 'BCE+Hinge' (proportion)
+        'remove_self_edges':  False,   # Remove self-edges
+        'edge_balance':        True    # true/false edge balance unity re-weight
+    },
+    
+    # Contrastive loss per particle
+    'edge_contrastive': {
+        'weights':  True,            # TrackML hit weights ok with this
+        'type':   'softmax',
+        'tau':      0.3,             # temperature (see: https://arxiv.org/abs/2012.09740, https://openreview.net/pdf?id=vnOHGQY4FP1)
+        'sub_sample': 300,           # memory constraint (maximum number of target objects to compute the loss per event)
+        
+        'min_prob':  1e-3,           # minimum edge prob. score to be included in the loss [EXPERIMENTAL]
+    },                               # (higher values push towards purity, but can weaken efficiency for e.g. high multiplicity clusters)
+    
+    # Cluster hit binary cross entropy loss
+    'cluster_BCE': {
+        'weights': False,            # TrackML hit weights (0 for noise) not exactly compatible
+        'type':          'Focal',    # 'BCE', 'BCE+Hinge', 'Focal'
+        'gamma':             1.0,    # For 'Focal' (entropy exponent)
+        'delta':            0.05,    # For 'BCE+Hinge' (proportion)
+    },
+    
+    # Cluster set hit loss
+    'cluster_contrastive': {
+        'weights':  False,        # TrackML hit weights (0 for noise) not exactly compatible
+        'type':   'intersect',    # 'intersect', 'dice', 'jaccard'
+        'smooth':    1.0          # regularization for 'dice' and 'jaccard'
+    },
+    
+    # Cluster meta-supervision target
+    'meta_target':  'pivotmajor'   # 'major' (vote from all nodes ground truth) or 'pivotmajor' (vote from pivots ground truth)
+}
+
+# -------------------------------------------------------------------------
+
+# These algorithm parameters can be changed after training, but
+# note that the transformer network may adapt (learn) its weights according
+# to the values set here during the training
+cluster_param = {
+
+    # These are set from the command line interface
+    'algorithm':            None,
+    'edge_threshold':       None,
+
+    
+    ## Cut clustering & Transformer clustering input
+    'min_graph':               4,     # Minimum subgraph size after the threshold and WCC search, the rest are treated as noise
+    
+    ## DBSCAN clustering
+    'dbscan': {
+        'eps':               0.2,     
+        'min_samples':         3,
+    },
+    
+    ## HDBSCAN clustering
+    # https://hdbscan.readthedocs.io/en/latest/api.html
+    'hdbscan': {
+        'algorithm':         'generic',
+        'cluster_selection_epsilon': 0.0,
+        'cluster_selection_method': 'eom',  # 'eom' or 'leaf'
+        'alpha':             1.0,
+        'min_samples':         2,     # Keep it 2
+        'min_cluster_size':    4,    
+        'max_dist':           1.0     # Keep it 1.0
+    },
+    
+    ## Transformer clustering
+    'worker_split':              4,   # GPU Memory <-> GPU latency tradeoff (no accuracy impact)
+    
+    'transformer': {
+        'seed_strategy':  'random',   # 'random', 'max' (max norm), 'max_T (transverse max), 'min' (min norm), 'min_T' (transverse min)
+        'seed_ktop_max':         2,   # Number of pivot walk (seed) candidates (higher -> better accuracy but slower)
+        
+        'N_pivots':              3,   # Number of pivotal hits to search per cluster (>> 1)
+        'random_paths':          1,   # (Put >> 1 for MC sampled random walk, and 1 for greedy max-prob walk)
+        
+        'max_failures':          2,   # Maximum number of failures per pivot list nodes (put 1+ for greedy, >> 1 for MC walk)
+        'diffuse_threshold':   0.4,   # Diffusion connectivity ~ Pivot quality threshold
+        
+        # Micrograph extension type: 'pivot-spanned' (ok with 'hyper' adjacency), 'full' (for other than 'hyper' needed, more inclusive but possibly unstable)
+        'micrograph_mode':   'pivot-spanned',
+        
+        'threshold_algo':    'fixed', # 'soft' (learnable), 'fisher' (batch-by-batch 1D-Fisher rule adaptive) or 'fixed'
+        
+        'tau':                 0.001, # 'soft':: Sigmoid 'temperature' (tau -> 0 ~ heaviside step)
+        
+        'ktop_max':           30,     # 'fisher':: Maximum cluster size (how many are considered from Transformer output), ranked by mask score
+        'fisher_threshold': np.linspace(0.4,0.6, 0), # 'fisher':: Threshold values tested
+        
+        'fixed_threshold':     0.5,   # 'fixed':: Note, if this is too high -> training may be unstable (first transformer iterations are bad)
+        
+        'min_cluster_size':      4,   # Require at least this many constituents per cluster
+    }
+}
+
+
+# -------------------------------------------------------------------------
+
+### Geometric adjacency estimator
+geom_param  = {
+    
+    # Use pre-trained 'voxdyn' or 'neurodyn' (experimental)
+    'algorithm':  'voxdyn',
+    
+    # Print adjacency metrics (this will slow down significantly)
+    'verbose': False,
+    
+    #'device':  'cuda', # CUDA not working with Faiss from conda atm (CUDA 11.4)
+    'device': 'cpu',
+
+    # 'neurodyn' parameters (PLACEHOLDER; not implemented)
+    'neural_param': {
+        'layers':  [6, 128, 64, 1],
+        'act':     'silu',
+        'bn':        True,
+        'dropout':    0.0,
+        'last_act':  False
+    },
+    
+    'neural_path': 'models/neurodyn'
+}
+
+
+# -------------------------------------------------------------------------
+
+### GNN + Transformer model parameters
+net_model_param = {
+    
+    # GNN predictor block
+    'graph_block_param': {
+        
+        'GNN_model' : 'SuperEdgeConv',   # 'SuperEdgeConv', 'GaugeEdgeConv'
+        'nstack':           5,           # Number of GNN message passing layers
+        
+        'coord_dim':        3,           # Input dimension
+        'h_dim':           64,           # Intermediate latent embedding dimension
+        'z_dim':           61,           # Final latent embedding dimension
+        
+        # https://pytorch-geometric.readthedocs.io/en/latest/modules/nn.html#aggregation-operators
+        'SuperEdgeConv': {
+            'm_dim':                   64,
+            'aggr':             ['mean']*5, # 'mean' (seems best memory/accuracy wise), 'sum', 'max', 'softmax', 'multi-aggregation', 'set-transformer'
+            'use_residual':          True,
+        },
+        
+        'GaugeEdgeConv': {
+            'm_dim':                   64,
+            'aggr':                  ['mean']*5, # As many as 'nstack'
+            'norm_coord':             False,
+            'norm_coord_scale_init':   1e-2,
+        },
+        
+        # Edge prediction (correlation MLP) type: 'symmetric-dot', 'symmetrized', 'asymmetric'
+        # (clustering Transformer should prefer 'symmetric-dot')
+        'edge_type':   'symmetric-dot',
+        
+        ## Convolution (message passing) MLPs
+        
+        'MLP_GNN_edge': {
+            'act':         'silu',       # 'relu', 'tanh', 'silu', 'elu'
+            'bn':            True,
+            'dropout':        0.0,
+            'last_act':      True,
+        },
+        
+        #'MLP_GNN_coord': {               # Only for 'GaugeEdgeConv'
+        #    'act':         'silu',
+        #    'bn':            True,
+        #    'dropout':        0.0,
+        #    'last_act':      True,
+        #},
+
+        'MLP_GNN_latent': {
+            'act':         'silu',
+            'bn':            True,
+            'dropout':        0.0,
+            'last_act':      True,
+        },
+        
+        ## Latent Fusion MLP
+        'MLP_fusion': {
+            'act':         'silu',
+            'bn':            True,
+            'dropout':        0.0,
+            'last_act':      True,
+        },
+
+        ## 2-pt edge correlation MLP
+        'MLP_correlate': {
+            'act':         'silu',
+            'bn':            True,
+            'dropout':        0.0,
+            'last_act':     False,
+        },
+    },
+    
+    # Transformer clusterization block
+    'cluster_block_param': {
+        'in_dim':         64,       # Same as GNN 'zdim' + 3 (for 3D coordinates)
+        'h_dim':          64,       # Latent dim, needs to be divisible by num_heads
+        'output_dim':      1,       # Always 1
+        'nstack_dec':      4,       # Number of self-attention layers
+        
+        'MLP_enc': {                # First encoder MLP
+            'act':         'silu',  
+            'bn':           False,
+            'dropout':        0.0,
+            'last_act':     False,
+        },
+        
+        'MAB_dec': {                # Transformer decoder MAB
+            'num_heads':        4,
+            'ln':            True,
+            'dropout':        0.0,
+            'MLP_param':{
+                'act':         'silu',  
+                'bn':           False,
+                'dropout':        0.0,
+                'last_act':      True,
+            }
+        },
+        
+        'SAB_dec': {                # Transformer decoder SAB
+            'num_heads':        4,
+            'ln':            True,
+            'dropout':        0.0,
+            'MLP_param':{
+                'act':         'silu',  
+                'bn':           False,
+                'dropout':        0.0,
+                'last_act':      True,
+            }
+        },
+        
+        'MLP_mask': {               # Mask decoder MLP
+            'act':         'silu',  
+            'bn':           False,
+            'dropout':        0.0,
+            'last_act':     False,
+        }
+    }
+}
+
+# -------------------------------------------------------------------------
+# [EXPERIMENTAL] -- normalizing flow
+
+# Conditional data array indices (see /hypertrack/trackml.py)
+cond_ind = [0,1,2,3,4,5,6]
+
+pdf_model_param = {
+    'in_dim':          61,
+    'num_cond_inputs':  len(cond_ind),
+    'h_dim':          196,
+    'nblocks':          4,
+    'act':          'tanh'
+}

hypertrack/models/models_tune-5.py

@@ -0,0 +1,635 @@
+# HyperTrack neural model torch classes
+#
+# [email protected], 2023
+
+from typing import Callable, Union, Optional
+
+import math
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import torch_geometric
+import torch_geometric.transforms as T
+
+from torch_geometric.nn import MessagePassing
+from torch_geometric.typing import Size, Tensor, OptTensor, PairTensor, PairOptTensor, OptPairTensor, Adj
+
+from hypertrack.dmlp import MLP
+import hypertrack.flows as fnn
+
+
+class SuperEdgeConv(MessagePassing):
+    r"""
+    Custom GNN convolution operator aka 'generalized EdgeConv' (original EdgeConv: arxiv.org/abs/1801.07829)
+    """
+    
+    def __init__(self, mlp_edge: Callable, mlp_latent: Callable, aggr: str='mean',
+                 mp_attn_dim: int=0, use_residual=True, **kwargs):
+        
+        if aggr == 'multi-aggregation':
+            aggr = torch_geometric.nn.aggr.MultiAggregation(aggrs=['sum', 'mean', 'std', 'max', 'min'], mode='attn',
+                    mode_kwargs={'in_channels': mp_attn_dim, 'out_channels': mp_attn_dim, 'num_heads': 1})
+        
+        if aggr == 'set-transformer':
+            aggr = torch_geometric.nn.aggr.SetTransformerAggregation(channels=mp_attn_dim, num_seed_points=1, 
+                    num_encoder_blocks=1, num_decoder_blocks=1, heads=1, concat=False,
+                    layer_norm=False, dropout=0.0)
+
+        super().__init__(aggr=aggr, **kwargs)
+        self.nn           = mlp_edge
+        self.nn_final     = mlp_latent
+        self.use_residual = use_residual
+        
+        self.reset_parameters()
+
+        self.apply(self.init_)
+
+    def init_(self, module):
+        if type(module) in {nn.Linear}:
+            #print(__name__ + f'.SuperEdgeConv: Initializing module: {module}')
+            nn.init.xavier_normal_(module.weight)
+            nn.init.zeros_(module.bias)
+    
+    def reset_parameters(self):
+        torch_geometric.nn.inits.reset(self.nn)
+        torch_geometric.nn.inits.reset(self.nn_final)
+    
+    def forward(self, x: Union[Tensor, PairTensor], edge_index: Adj,
+            edge_attr: OptTensor = None, edge_weight: OptTensor = None, size: Size = None) -> Tensor:
+
+        if edge_attr is not None and len(edge_attr.shape) == 1: # if 1-dim edge_attributes
+            edge_attr = edge_attr[:,None]
+        
+        # Message passing
+        m = self.propagate(edge_index, x=x, edge_attr=edge_attr, edge_weight=edge_weight, size=None)
+        
+        # Final MLP
+        y = self.nn_final(torch.concat([x, m], dim=-1))
+        
+        # Residual connections
+        if self.use_residual and (y.shape[-1] == x.shape[-1]):
+            y = y + x
+        
+        return y
+
+    def message(self, x_i: Tensor, x_j: Tensor, edge_attr: OptTensor, edge_weight: OptTensor) -> Tensor:
+        
+        # Edge features
+        e1 = torch.norm(x_j - x_i, dim=-1) # Norm of the difference (invariant under rotations and translations)
+        e2 = torch.sum(x_j * x_i,  dim=-1) # Dot-product (invariant under rotations but not translations)
+        
+        if len(e1.shape) == 1:
+            e1 = e1[:,None]
+            e2 = e2[:,None]
+        
+        if edge_attr is not None:
+            m = self.nn(torch.cat([x_i, x_j - x_i, x_j * x_i, e1, e2, edge_attr], dim=-1))
+        else:
+            m = self.nn(torch.cat([x_i, x_j - x_i, x_j * x_i, e1, e2], dim=-1))
+        
+        return m if edge_weight is None else m * edge_weight.view(-1, 1)
+    
+    def __repr__(self):
+        return f'{self.__class__.__name__} (nn={self.nn}, nn_final={self.nn_final})'
+
+
+class CoordNorm(nn.Module):
+    """
+    Coordinate normalization for stability with GaugeEdgeConv
+    """
+    def __init__(self, eps = 1e-8, scale_init = 1.0):
+        super().__init__()
+        self.eps   = eps
+        scale      = torch.zeros(1).fill_(scale_init)
+        self.scale = nn.Parameter(scale)
+
+    def forward(self, coord):
+        norm = coord.norm(dim = -1, keepdim = True)
+        normed_coord = coord / norm.clamp(min = self.eps)
+        return normed_coord * self.scale
+
+
+class GaugeEdgeConv(MessagePassing):
+    r"""
+    Custom GNN convolution operator aka 'E(N) equivariant GNN' (arxiv.org/abs/2102.09844)
+    """
+    
+    def __init__(self, mlp_edge: Callable, mlp_coord: Callable, mlp_latent: Callable, coord_dim: int=0,
+        update_coord: bool=True, update_latent: bool=True, aggr: str='mean', mp_attn_dim: int=0, 
+        norm_coord=False, norm_coord_scale_init = 1e-2, **kwargs):
+        
+        kwargs.setdefault('aggr', aggr)
+        super(GaugeEdgeConv, self).__init__(**kwargs)
+        
+        self.mlp_edge      = mlp_edge
+        self.mlp_coord     = mlp_coord
+        self.mlp_latent    = mlp_latent
+        
+        self.update_coord  = update_coord
+        self.update_latent = update_latent
+        
+        self.coord_dim     = coord_dim
+        
+        # Coordinate normalization
+        self.coors_norm = CoordNorm(scale_init = norm_coord_scale_init) if norm_coord else nn.Identity()
+        
+        self.reset_parameters()
+
+        self.apply(self.init_)
+
+    def init_(self, module):
+        if type(module) in {nn.Linear}:
+            #print(__name__ + f'.GaugeEdgeConv: Initializing module: {module}')
+            nn.init.xavier_normal_(module.weight)
+            nn.init.zeros_(module.bias)
+    
+    def reset_parameters(self):
+        torch_geometric.nn.inits.reset(self.mlp_edge)
+        torch_geometric.nn.inits.reset(self.mlp_coord)
+        torch_geometric.nn.inits.reset(self.mlp_latent)
+        
+
+    def forward(self, x: Union[Tensor, PairTensor], edge_index: Adj,
+            edge_attr: OptTensor = None, edge_weight: OptTensor = None, size: Size = None) -> Tensor:
+        """
+        Forward function
+        """
+        
+        # Separate spatial (e.g. 3D-coordinates) and features
+        coord, feats = x[..., 0:self.coord_dim], x[..., self.coord_dim:]
+        
+        # Coordinate difference: x_i - x_j
+        diff_coord = coord[edge_index[0]] - coord[edge_index[1]]
+        diff_norm2 = (diff_coord ** 2).sum(dim=-1, keepdim=True)
+        
+        if edge_attr is not None:
+            if len(edge_attr.shape) == 1: # if 1-dim edge_attributes
+                edge_attr = edge_attr[:,None]
+
+            edge_attr_feats = torch.cat([edge_attr, diff_norm2], dim=-1)
+        else:
+            edge_attr_feats = diff_norm2
+        
+        # Propagation
+        latent_out, coord_out = self.propagate(edge_index, x=feats, edge_attr=edge_attr_feats,
+                                    edge_weight=edge_weight, coord=coord, diff_coord=diff_coord, size=None)
+
+        return torch.cat([coord_out, latent_out], dim=-1)
+    
+    def message(self, x_i: Tensor, x_j: Tensor, edge_attr: OptTensor, edge_weight: OptTensor) -> Tensor:
+        """
+        Message passing core operation between nodes (i,j)
+        """
+        
+        m_ij = self.mlp_edge(torch.cat([x_i, x_j, edge_attr], dim=-1))
+
+        return m_ij if edge_weight is None else m_ij * edge_weight.view(-1, 1)
+
+    def propagate(self, edge_index: Adj, size: Size = None, **kwargs):
+        """
+        The initial call to start propagating messages.
+        
+        Args:
+            edge_index: holds the indices of a general (sparse)
+                        assignment matrix of shape :obj:`[N, M]`.
+            size:       (tuple, optional) if none, the size will be inferred
+                        and assumed to be quadratic.
+            **kwargs:   Any additional data which is needed to construct and
+                        aggregate messages, and to update node embeddings.
+        """
+        
+        # Check:
+        # https://github.com/pyg-team/pytorch_geometric/blob/master/torch_geometric/nn/conv/message_passing.py
+        
+        size          = self._check_input(edge_index, size)
+        coll_dict     = self._collect(self._user_args, edge_index, size, kwargs)
+        msg_kwargs    = self.inspector.distribute('message',   coll_dict)
+        aggr_kwargs   = self.inspector.distribute('aggregate', coll_dict)
+        update_kwargs = self.inspector.distribute('update',    coll_dict)
+        
+        # Message passing of node latent embeddings
+        m_ij = self.message(**msg_kwargs)
+        
+        if self.update_coord:
+            
+            # Normalize
+            kwargs["diff_coord"] = self.coors_norm(kwargs["diff_coord"])
+            
+            # Aggregate weighted coordinates
+            mhat_i    = self.aggregate(kwargs["diff_coord"] * self.mlp_coord(m_ij), **aggr_kwargs)
+            coord_out = kwargs["coord"] + mhat_i # Residual connection
+            
+        else:
+            coord_out = kwargs["coord"]
+        
+        
+        if self.update_latent:
+            
+            # Aggregate message passing results
+            m_i        = self.aggregate(m_ij, **aggr_kwargs)
+            
+            # Update latent representation
+            latent_out = self.mlp_latent(torch.cat([kwargs["x"], m_i], dim = -1))
+            latent_out = kwargs["x"] + latent_out # Residual connection
+        else:
+            latent_out = kwargs["x"]
+        
+        # Return tuple
+        return self.update((latent_out, coord_out), **update_kwargs)
+
+    def __repr__(self):
+        return f'{self.__class__.__name__}(GaugeEdgeConv = {self.mlp_edge} | {self.mlp_coord} | {self.mlp_latent})'
+
+
+class InverterNet(torch.nn.Module):
+    """
+    HyperTrack neural model "umbrella" class, encapsulating GNN and Transformer etc.
+    """
+    def __init__(self, graph_block_param={}, cluster_block_param={}):
+        
+        """
+        conv_aggr: 'mean' seems very crucial.
+        """
+        super().__init__()
+        
+        self.training_on   = True
+        
+        self.coord_dim     = graph_block_param['coord_dim']   # Input dimension
+        self.h_dim         = graph_block_param['h_dim']       # Intermediate latent dimension
+        self.z_dim         = graph_block_param['z_dim']       # Final latent dimension
+        
+        self.GNN_model     = graph_block_param['GNN_model']
+        self.nstack        = graph_block_param['nstack']
+        
+        self.edge_type     = graph_block_param['edge_type']
+        MLP_fusion         = graph_block_param['MLP_fusion']
+        
+        self.num_edge_attr = 1 # cf. custom edge features constructed in self.encode()
+        self.conv_gnn_edx  = nn.ModuleList()
+        
+        # Transformer node mask learnable "soft" threshold
+        self.thr = nn.Parameter(torch.Tensor(1))
+        nn.init.constant_(self.thr, 0.5)
+        
+        # 1. GNN encoder
+        
+        ## Model type
+        if self.GNN_model == 'GaugeEdgeConv':
+            
+            self.m_dim     = graph_block_param['GaugeEdgeConv']['m_dim']
+
+            MLP_GNN_edge   = graph_block_param['MLP_GNN_edge']
+            MLP_GNN_coord  = graph_block_param['MLP_GNN_coord']
+            MLP_GNN_latent = graph_block_param['MLP_GNN_latent']
+            
+            num_intrinsic_attr = 1 # distance
+            
+            for i in range(self.nstack):
+                self.conv_gnn_edx.append(GaugeEdgeConv(
+                    mlp_edge    = MLP([2*self.h_dim + self.num_edge_attr + num_intrinsic_attr, 2*self.m_dim, self.m_dim], **MLP_GNN_edge),
+                    mlp_coord   = MLP([self.m_dim, 2*self.m_dim, 1], **MLP_GNN_coord),
+                    mlp_latent  = MLP([self.h_dim + self.m_dim, 2*self.h_dim, self.h_dim], **MLP_GNN_latent),
+                    aggr=graph_block_param['GaugeEdgeConv']['aggr'][i],
+                    norm_coord=graph_block_param['GaugeEdgeConv']['norm_coord'],
+                    norm_coord_scale_init=graph_block_param['GaugeEdgeConv']['norm_coord_scale_init'],
+                    coord_dim=self.coord_dim))
+
+            self.mlp_fusion_edx = MLP([self.nstack * (self.coord_dim + self.h_dim), self.h_dim, self.z_dim], **MLP_fusion)
+        
+        ## Model type
+        elif self.GNN_model == 'SuperEdgeConv':
+            
+            self.m_dim     = graph_block_param['SuperEdgeConv']['m_dim']
+            
+            MLP_GNN_edge   = graph_block_param['MLP_GNN_edge']
+            MLP_GNN_latent = graph_block_param['MLP_GNN_latent']
+            
+            num_intrinsic_attr = 2 # distance and dot-product
+            
+            self.conv_gnn_edx.append(SuperEdgeConv(
+                    mlp_edge    = MLP([3*self.coord_dim + self.num_edge_attr + num_intrinsic_attr, self.m_dim, self.m_dim], **MLP_GNN_edge),
+                    mlp_latent  = MLP([self.coord_dim + self.m_dim, self.h_dim, self.h_dim], **MLP_GNN_latent),
+                    mp_attn_dim=self.h_dim, aggr=graph_block_param['SuperEdgeConv']['aggr'][0], use_residual=False))
+            
+            for i in range(1,self.nstack):
+                self.conv_gnn_edx.append(SuperEdgeConv(
+                    mlp_edge    = MLP([3*self.h_dim + self.num_edge_attr + num_intrinsic_attr, self.m_dim, self.m_dim], **MLP_GNN_edge),
+                    mlp_latent  = MLP([self.h_dim + self.m_dim, self.h_dim, self.h_dim], **MLP_GNN_latent),
+                    mp_attn_dim=self.h_dim, aggr=graph_block_param['SuperEdgeConv']['aggr'][i], use_residual=graph_block_param['SuperEdgeConv']['use_residual']))
+            
+            self.mlp_fusion_edx = MLP([self.nstack * self.h_dim, self.h_dim, self.z_dim], **MLP_fusion)
+
+        else:
+            raise Exception(__name__ + '.__init__: Unknown GNN_model chosen')
+        
+        # 2. Graph edge predictor
+        MLP_correlate = graph_block_param['MLP_correlate']
+        
+        if   self.edge_type == 'symmetric-dot':
+            self.mlp_2pt_edx = MLP([self.z_dim,   self.z_dim//2, self.z_dim//2, 1], **MLP_correlate)
+        elif self.edge_type == 'symmetrized' or self.edge_type == 'asymmetric':
+            self.mlp_2pt_edx = MLP([2*self.z_dim, self.z_dim//2, self.z_dim//2, 1], **MLP_correlate)
+        
+        # 3. Clustering predictor
+        self.transformer_ccx = STransformer(**cluster_block_param)
+    
+    def encode(self, x, edge_index):
+        """
+        Encoder GNN
+        """
+        # Compute node degree 'custom feature' between edges
+        d         = torch_geometric.utils.degree(edge_index[0,:])
+        edge_attr = (d[edge_index[0,:]] - d[edge_index[1,:]]) / torch.mean(d)
+        edge_attr = edge_attr.to(x.dtype)
+        
+        # We take each output for the parallel fusion
+        x_out = [None] * self.nstack
+        
+        # First input [x; one-vector for the first embeddings (latents)]
+        if self.GNN_model == 'GaugeEdgeConv':
+            x_ = torch.cat([x, torch.ones((x.shape[0], self.h_dim), device=x.device)], dim=-1)
+        else:
+            x_ = x
+        
+        # Apply GNN layers
+        x_out[0] = self.conv_gnn_edx[0](x_, edge_index, edge_attr)
+        for i in range(1,self.nstack):
+            x_out[i] = self.conv_gnn_edx[i](x_out[i-1], edge_index, edge_attr)
+        
+        return self.mlp_fusion_edx(torch.cat(x_out, dim=-1))
+    
+    def decode(self, z, edge_index):
+        """
+        Decoder of two-point correlations (edges)
+        """
+        if   self.edge_type == 'symmetric-dot':
+            return self.mlp_2pt_edx(z[edge_index[0],:] * z[edge_index[1],:])
+        
+        elif self.edge_type == 'symmetrized':
+            a = self.mlp_2pt_edx(torch.cat([z[edge_index[0],:], z[edge_index[1],:]], dim=-1))
+            b = self.mlp_2pt_edx(torch.cat([z[edge_index[1],:], z[edge_index[0],:]], dim=-1))
+            return (a + b) / 2.0
+
+        elif self.edge_type == 'asymmetric':
+            a = self.mlp_2pt_edx(torch.cat([z[edge_index[0],:], z[edge_index[1],:]], dim=-1))
+            return a
+
+    def decode_cc_ind(self, X, X_pivot, X_mask=None, X_pivot_mask=None):
+        """
+        Decoder of N-point node mask
+        """
+        
+        return self.transformer_ccx(X=X, X_pivot=X_pivot, X_mask=X_mask, X_pivot_mask=X_pivot_mask)
+    
+    def set_model_param_grad(self, string_id='edx', requires_grad=True):
+        """
+        Freeze or unfreeze model parameters (for the gradient descent)
+        """
+        
+        for name, W in self.named_parameters():   
+            if string_id in name:
+                W.requires_grad = requires_grad
+                #print(f'Setting requires_grad={W.requires_grad} of the parameter <{name}>')
+        return
+
+    def get_model_param_grad(self, string_id='edx'):
+        """
+        Get model parameter state (for the gradient descent)
+        """
+        
+        for name, W in self.named_parameters():   
+            if string_id in name: # Return the state of the first
+                return W.requires_grad
+
+
+class MAB(nn.Module):
+    """
+    Attention based set Transformer block (arxiv.org/abs/1810.00825)
+    """
+    def __init__(self, dim_Q, dim_K, dim_V, num_heads=4, ln=True, dropout=0.0,
+                 MLP_param={'act': 'relu', 'bn': False, 'dropout': 0.0, 'last_act': True}):
+        super(MAB, self).__init__()
+        assert dim_V % num_heads == 0, "MAB: dim_V must be divisible by num_heads"
+        self.dim_V = dim_V
+        self.num_heads = num_heads
+        self.W_q = nn.Linear(dim_Q, dim_V)
+        self.W_k = nn.Linear(dim_K, dim_V)
+        self.W_v = nn.Linear(dim_K, dim_V)
+        self.W_o = nn.Linear(dim_V, dim_V) # Projection layer
+        
+        if ln:
+            self.ln0 = nn.LayerNorm(dim_V)
+            self.ln1 = nn.LayerNorm(dim_V)
+        
+        if dropout > 0:
+            self.Dropout = nn.Dropout(dropout)
+        
+        self.MLP = MLP([dim_V, dim_V, dim_V], **MLP_param)
+    
+        # We use torch default initialization here
+    
+    def reshape_attention_mask(self, Q, K, mask):
+        """
+        Reshape attention masks
+        """
+        total_mask = None
+
+        if mask[0] is not None:
+            qmask = mask[0].repeat(self.num_heads,1)[:,:,None]  # New shape = [# heads x # batches, # queries, 1]
+
+        if mask[1] is not None:
+            kmask = mask[1].repeat(self.num_heads,1)[:,None,:]  # New shape = [# heads x # batches, 1, # keys]
+        
+        if   mask[0] is None and mask[1] is not None:
+            total_mask = kmask.repeat(1,Q.shape[1],1)
+        elif mask[0] is not None and mask[1] is None:
+            total_mask = qmask.repeat(1,1,K.shape[1])
+        elif mask[0] is not None and mask[1] is not None:
+            total_mask = qmask & kmask # will auto broadcast dimensions to [# heads x # batches, # queries, # keys]
+        
+        return total_mask
+
+    def forward(self, Q, K, mask = (None, None)):
+        """
+        Q:    queries
+        K:    keys
+        mask: query mask [#batches x #queries], keys mask [#batches x #keys]
+        """
+        dim_split = self.dim_V // self.num_heads
+        
+        # Apply Matrix-vector multiplications
+        # and do multihead splittings (reshaping)
+        Q_ = torch.cat(self.W_q(Q).split(dim_split, -1), 0)
+        K_ = torch.cat(self.W_k(K).split(dim_split, -1), 0)
+        V_ = torch.cat(self.W_v(K).split(dim_split, -1), 0)
+        
+        # Dot-product attention: softmax(QK^T / sqrt(dim_V))
+        QK = Q_.bmm(K_.transpose(-1,-2)) / math.sqrt(self.dim_V)  # bmm does batched matrix multiplication
+        
+        # Attention mask
+        total_mask = self.reshape_attention_mask(Q=Q, K=K, mask=mask)
+        if total_mask is not None:                 
+            QK.masked_fill_(~total_mask, float('-1E6'))
+        
+        # Compute attention probabilities
+        A = torch.softmax(QK,-1)
+        
+        # Residual connection of Q + multi-head attention A weighted V result
+        H = Q + self.W_o(torch.cat((A.bmm(V_)).split(Q.size(0), 0),-1))
+        
+        # First layer normalization + Dropout
+        H = H if getattr(self, 'ln0', None) is None else self.ln0(H)
+        H = H if getattr(self, 'Dropout', None) is None else self.Dropout(H)
+        
+        # Residual connection of H + feed-forward net
+        H = H + self.MLP(H)
+        
+        # Second layer normalization + Dropout
+        H = H if getattr(self, 'ln1', None) is None else self.ln1(H)
+        H = H if getattr(self, 'Dropout', None) is None else self.Dropout(H)
+        
+        return H
+
+
+class SAB(nn.Module):
+    """
+    Full self-attention MAB(X,X)
+    ~O(N^2)
+    """
+    def __init__(self, dim_in, dim_out, num_heads=4, ln=True, dropout=0.0,
+                 MLP_param={'act': 'relu', 'bn': False, 'dropout': 0.0, 'last_act': True}):
+        super(SAB, self).__init__()
+        self.mab = MAB(dim_Q=dim_in, dim_K=dim_in, dim_V=dim_out, 
+                       num_heads=num_heads, ln=ln, dropout=dropout, MLP_param=MLP_param)
+
+    def forward(self, X, mask=None):
+        return self.mab(Q=X, K=X, mask=(mask, mask))
+
+class ISAB(nn.Module):
+    """
+    Faster version of SAB with inducing points
+    """
+    def __init__(self, dim_in, dim_out, num_inds, num_heads=4, ln=True, dropout=0.0,
+                 MLP_param={'act': 'relu', 'bn': False, 'dropout': 0.0, 'last_act': True}):
+        super(ISAB, self).__init__()
+        self.I = nn.Parameter(torch.Tensor(1, num_inds, dim_out))
+        nn.init.xavier_uniform_(self.I)
+        self.mab0 = MAB(dim_Q=dim_out, dim_K=dim_in, dim_V=dim_out,
+                        num_heads=num_heads, ln=ln, dropout=dropout, MLP_param=MLP_param)
+        self.mab1 = MAB(dim_Q=dim_in, dim_K=dim_out, dim_V=dim_out,
+                        num_heads=num_heads, ln=ln, dropout=dropout, MLP_param=MLP_param)
+
+    def forward(self, X, mask=None):
+        H = self.mab0(Q=self.I.expand(X.size(0),-1,-1), K=X, mask=(None, mask))
+        H = self.mab1(Q=X, K=H, mask=(mask, None))
+        return H
+
+class PMA(nn.Module):
+    """
+    Adaptive pooling with "k" > 1 option (several learnable reference vectors)
+    """
+    def __init__(self, dim, k=1, num_heads=4, ln=True, dropout=0.0,
+                 MLP_param={'act': 'relu', 'bn': False, 'dropout': 0.0, 'last_act': True}):
+        super(PMA, self).__init__()
+        self.S   = nn.Parameter(torch.Tensor(1, k, dim))
+        nn.init.xavier_uniform_(self.S)
+        self.mab = MAB(dim_Q=dim, dim_K=dim, dim_V=dim,
+                       num_heads=num_heads, ln=ln, dropout=dropout, MLP_param=MLP_param)
+    
+    def forward(self, X, mask):
+        return self.mab(Q=self.S.expand(X.size(0),-1,-1), K=X, mask=(None, mask))
+
+class STransformer(nn.Module):
+    """
+    Set Transformer based clustering network
+    """
+
+    class mySequential(nn.Sequential):
+        """
+        Multiple inputs version of nn.Sequential customized for
+        multiple self-attention layers with a (same) mask
+        """
+        def forward(self, *inputs):
+            
+            X, mask = inputs[0], inputs[1]
+            for module in self._modules.values():
+                X = module(X,mask)
+            return X
+    
+    def __init__(self, in_dim, h_dim, output_dim, nstack_dec=4, 
+                 MLP_enc={}, MAB_dec={}, SAB_dec={}, MLP_mask={}):
+        
+        super().__init__()
+        
+        # Encoder MLP
+        self.MLP_E = MLP([in_dim, in_dim, h_dim], **MLP_enc)
+        
+        # Decoder
+        self.mab_D = MAB(dim_Q=h_dim, dim_K=h_dim, dim_V=h_dim, **MAB_dec)
+
+        # Decoder self-attention layers
+        self.sab_stack_D = self.mySequential(*[
+            self.mySequential(
+                SAB(dim_in=h_dim, dim_out=h_dim, **SAB_dec)
+            )
+            for i in range(nstack_dec)
+        ])
+        
+        # Final mask MLP
+        self.MLP_m = MLP([h_dim, h_dim//2, h_dim//4, output_dim], **MLP_mask)
+    
+    def forward(self, X, X_pivot, X_mask = None, X_pivot_mask = None):
+        """
+        X:            input data vectors per row
+        X_pivot:      pivotal data (at least one per batch)
+        X_mask:       boolean (batch) mask for X (set 0 for zero-padded null elements)
+        X_pivot_mask: boolean (batch) mask for X_pivot
+        """
+        
+        # Simple encoder
+        G       = self.MLP_E(X)
+        G_pivot = self.MLP_E(X_pivot)
+        
+        # Compute cross-attention and self-attention
+        H_m = self.sab_stack_D(self.mab_D(Q=G, K=G_pivot, mask=(X_mask, X_pivot_mask)), X_mask)
+        
+        # Decode logits
+        return self.MLP_m(H_m)
+
+class TrackFlowNet(torch.nn.Module):
+    """
+    Normalizing Flow Network [experimental]
+    """
+    def __init__(self, in_dim, num_cond_inputs=None, h_dim=64, nblocks=4, act='tanh'):
+        """
+        conv_aggr: 'mean' in GNN seems to work ok with Flow!
+        """
+        super().__init__()
+
+        self.training_on = True
+        
+        # -----------------------------------
+        # MAF density estimator
+
+        modules = []
+        for _ in range(nblocks):
+            modules += [
+                fnn.MADE(num_inputs=in_dim, num_hidden=h_dim, num_cond_inputs=num_cond_inputs, act=act),
+                #fnn.BatchNormFlow(in_dim), # May cause problems in recursive use
+                fnn.Reverse(in_dim)
+            ]
+        
+        self.track_pdf = fnn.FlowSequential(*modules)
+        # -----------------------------------
+
+    def set_model_param_grad(self, string_id='pdf', requires_grad=True):
+        """
+        Freeze or unfreeze model parameters (for the gradient descent)
+        """
+        
+        for name, W in self.named_parameters():   
+            if string_id in name:
+                W.requires_grad = requires_grad
+                #print(f'Setting requires_grad={W.requires_grad} of the parameter <{name}>')
+        return

models/tag_f-0p01-hyper-5/model_net_epoch_834169.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:57da49d764e4fee13e9556a63b856cc254b28cd99397a99787a841baab1ee3c7
+size 143403488

models/tag_f-0p01-hyper-5/model_pdf_epoch_834169.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:365a51ad4bafce3aae0fb1f5442b1aa632d473ed866aca3d720bb0b4ba6401d1
+size 144523939

models/tag_f-0p01-hyper-5/optimizer_net_epoch_834169.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:91c9573a6a4a193be373cb523dd814c6340f2bfcc3ace913154009928a352944
+size 2548315

models/tag_f-0p01-hyper-5/optimizer_pdf_epoch_834169.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f9af437f10a548c60026915058d0202e947243d94638aa55569469e1e0ebf711
+size 2461899

models/tag_f-0p01-hyper-5/scheduler_net_epoch_834169.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:66896ef7df452a76f6482bd1b62bca415ae69d4fe88131a1c742ad638816d33c
+size 789

models/tag_f-0p01-hyper-5/scheduler_pdf_epoch_834169.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:80c31de2a96c990896e7a996c4226ad29b9b73d8230d0d7e26d51fee634479b7
+size 789

models/tag_f-0p1-hyper-5/model_net_epoch_752039.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b504e40740146fab5b809ad76b1c7814751a6b3594889e6268e5762d15041a5c
+size 129397024

models/tag_f-0p1-hyper-5/model_pdf_epoch_752039.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:57c7d332c93ed228129c93a115d727baae508835dc3179b2f37535f9f8f90bab
+size 130517539

models/tag_f-0p1-hyper-5/optimizer_net_epoch_752039.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:71def7ec1d75381d14d37f6149c98a16ba6d1fee5581f2b234b24a05ba60e5be
+size 2543579

models/tag_f-0p1-hyper-5/optimizer_pdf_epoch_752039.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c21ef4caf9621c3f385d00f08fc0a578e8a12ca02d015a67a1b0072c93a9f4a9
+size 981

models/tag_f-0p1-hyper-5/scheduler_net_epoch_752039.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8089ea387f0a81570664dbeec2f0a7201b1adea1195a7459f943792b783f6bc1
+size 789

models/tag_f-0p1-hyper-5/scheduler_pdf_epoch_752039.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4ec10834480a2d82e28eee0a60f05f078262b364b2587dd2d02caa51de97fda3
+size 789

models/tag_f-0p3-hyper-5/model_net_epoch_484878.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:80c29152b67cd037a07bc0de75e0f15a024f270b1f6b90205196fc8a186f5a75
+size 83839328

models/tag_f-0p3-hyper-5/model_pdf_epoch_484878.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8557c0ebffa2f335cdf7d94bd3025b5df220558b948d23484d706fb333d04913
+size 84959843

models/tag_f-0p3-hyper-5/optimizer_net_epoch_484878.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:62204519b3e2b84bd91396b84919ba7c4c680171988c255599969d2f7f45a734
+size 2543579

models/tag_f-0p3-hyper-5/optimizer_pdf_epoch_484878.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bb83aca384ca99600390f01b559ecca04b78fc9b8f2222c5f3d5deb99c1527bf
+size 981

models/tag_f-0p3-hyper-5/scheduler_net_epoch_484878.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:688e4074963740abec8f59b5f54c219377299feeabcea21c0f79015f47040e2a
+size 789

models/tag_f-0p3-hyper-5/scheduler_pdf_epoch_484878.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7d2e1e0f5cc38306cbdb00ae64b1dec0532852790dd48da7503310dfa97e119c
+size 789

models/voxdyn_node2node_hyper_ncell_131072.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7744317ce56a104d8d7357272e2b117a50d30db962fc7804fbd1e9dd658d7b15
+size 696889444

models/voxdyn_node2node_hyper_ncell_262144.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b9689592db5542078f033c1efcbbf3bd61157c0ac6b717542717189ca4d77718
+size 1248245246

models/voxdyn_node2node_hyper_ncell_32768.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c13e5c1574714a6e71c7c95c93574936bfa7ddef65938361facb6bc0639340ba
+size 215949890

models/voxdyn_node2node_hyper_ncell_524288.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ccc5926bfef73cde831fffb8b5c86b562b3176353cb61153ca6aab58eb0fe61a
+size 2135411740

models/voxdyn_node2node_hyper_ncell_65536.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c1d234b35eae47cbc7f99506986b4ed2b9c6f581ab9094d05f1bf4246fb86f04
+size 388562698