Equiv-DPnets working on high dimensional spaces.

Author: Danfoa

cfg/model/dae.yaml

@@ -3,13 +3,13 @@ defaults:
 
 name: DAE
 # Model hyperparameters
-obs_pred_w: 1.0                                  # Cost function weight for prediction in observation space Z
+obs_pred_w: 5.0                                  # Cost function weight for prediction in observation space Z
 orth_w: 1.0                                     # Weight of the orthonormal regularization term in the loss function
 corr_w: 0.0
 
 # Optimization hyperparameters parameters
 lr: 1e-3
-batch_size: 1024
+actiavtion: ELU
 equivariant: False
 
 summary: ${model.name}-Obs_w:${model.obs_pred_w}-Orth_w:${model.orth_w}-Act:${model.activation}-B:${model.bias}-BN:${model.batch_norm}-LR:${model.lr}-L:${model.num_layers}-${model.num_hidden_units}

cfg/model/dpnet.yaml

@@ -7,14 +7,13 @@ name: DPNet
 equivariant: False
 
 # Model hyperparameters
-activation: ReLU
+activation: ELU
 num_layers: 5                                   # Number MLPs' layers (including input and output layers)
 num_hidden_units: 128                           # Number of hidden units in each layer
 batch_norm: True
 bias: False
 # Optimization hyperparameters parameters
 lr: 1e-3
-batch_size: 1024
 
 max_ck_window_length: ${system.pred_horizon}    # Maximum length of the Chapman-Kolmogorov window
 ck_w: 0.0                                       # Weight of the Chapman-Kolmogorov regularization term in the loss function

cfg/model/edae.yaml

@@ -9,5 +9,7 @@ corr_w: 0.0
 state_dependent_obs_dyn: False                  # Whether to use state-dependent observation dynamics
 group_avg_trick: True
 
+
+actiavtion: ELU
 equivariant: True
 

nn/DeepProjections.py

@@ -129,6 +129,7 @@ def pre_process_obs_state(self,
         obs_state_traj_aux = super().pre_process_obs_state(obs_state_traj_aux)['obs_state_traj']
         return dict(obs_state_traj=obs_state_traj, obs_state_traj_aux=obs_state_traj_aux)
 
+
     def compute_loss_and_metrics(self,
                                  obs_state_traj: Tensor,
                                  obs_state_traj_aux: Tensor,

nn/DynamicsAutoEncoder.py

@@ -87,11 +87,11 @@ def forecast(self, state: Tensor, n_steps: int = 1, **kwargs) -> [dict[str, Tens
             f"{pred_obs_state_traj.shape}!=({self._batch_size}, {time_horizon}, {self.obs_state_dim})"
         return pred_state_traj, pred_obs_state_traj
 
-    def post_process_obs_state(self, pred_state_traj: Tensor, pred_state_one_step: Tensor) -> dict[str, Tensor]:
+    def post_process_obs_state(self, obs_state_traj: Tensor, pred_state_one_step: Tensor) -> dict[str, Tensor]:
         """ Post-process the predicted observable state trajectory given by the observable state dynamics.
 
         Args:
-            pred_state_traj: (batch, time, obs_state_dim) Trajectory of the predicted (time -1) observable states
+            obs_state_traj: (batch, time, obs_state_dim) Trajectory of the predicted (time -1) observable states
              predicted by the transfer operator.
             pred_state_one_step: (batch, time, obs_state_dim) Trajectory of the predicted one-step ahead (time)
              observable states predicted by the transfer operator.
@@ -101,7 +101,7 @@ def post_process_obs_state(self, pred_state_traj: Tensor, pred_state_one_step: T
                 - pred_obs_state_traj: (batch * time, obs_state_dim) Trajectory
                 - pred_obs_state_one_step: (batch, time, obs_state_dim) Trajectory
         """
-        batched_pred_obs_state_traj = batched_to_flat_trajectory(pred_state_traj)
+        batched_pred_obs_state_traj = batched_to_flat_trajectory(obs_state_traj)
         return dict(pred_obs_state_traj=batched_pred_obs_state_traj,
                     pred_obs_state_one_step=pred_state_one_step)
 

nn/EquivDeepPojections.py

@@ -21,7 +21,7 @@
 from data.DynamicsDataModule import DynamicsDataModule
 from nn.DeepProjections import DPNet
 from nn.EquivLinearDynamics import EquivLinearDynamics
-from nn.TwinMLP import TwinMLP
+from nn.ObservableNet import ObservableNet
 from nn.emlp import EMLP
 from nn.markov_dynamics import MarkovDynamics
 from utils.losses_and_metrics import forecasting_loss_and_metrics, obs_state_space_metrics
@@ -38,7 +38,7 @@ class EquivDPNet(DPNet):
         activation="p_elu",
         batch_norm=True,
         bias=False,
-        backbone_layers=-2  # num_layers - 2
+        # backbone_layers=-2  # num_layers - 2
         )
 
     def __init__(self,
@@ -257,40 +257,27 @@ def empirical_lin_inverse_projector(self, state: Tensor, obs_state: Tensor):
         return A, metrics
 
     def build_obs_fn(self, num_layers, **kwargs):
-        num_backbone_layers = kwargs.pop('backbone_layers', num_layers - 2 if self.aux_obs_space else 0)
-        if num_backbone_layers < 0:
-            num_backbone_layers = num_layers - num_backbone_layers
-        backbone_params = None
-        if num_backbone_layers > 0 and self.aux_obs_space:
-            num_hidden_units = kwargs.get('num_hidden_units')
-            activation_type = kwargs.pop('activation')
-            num_hidden_regular_fields = int(np.ceil(num_hidden_units // self.state_type_iso.size))
-            act = EMLP.get_activation(activation=activation_type,
-                                      in_type=self.state_type_iso,
-                                      channels=num_hidden_regular_fields)
-            backbone_params = dict(in_type=self.state_type_iso,
-                                   out_type=act.out_type,
-                                   activation=act,
-                                   num_layers=num_backbone_layers,
-                                   head_with_activation=True,
-                                   **copy.copy(kwargs))
-            kwargs['bias'] = False
-            kwargs['batch_norm'] = False
-            obs_fn_params = dict(in_type=act.out_type, out_type=self.obs_state_type,
-                                 num_layers=num_layers - num_backbone_layers,
-                                 activation=act,
-                                 head_with_activation=False, **kwargs)
-        else:
-            obs_fn_params = dict(in_type=self.state_type_iso,
-                                 out_type=self.obs_state_type,
-                                 num_layers=num_layers,
-                                 head_with_activation=False,
-                                 **kwargs)
 
-        return TwinMLP(net_kwargs=obs_fn_params,
-                       backbone_kwargs=backbone_params,
-                       fake_aux_fn=not self.aux_obs_space,
-                       equivariant=True)
+        num_hidden_units = kwargs.get('num_hidden_units')
+        activation_type = kwargs.pop('activation')
+        act = EMLP.get_activation(activation=activation_type,
+                                  in_type=self.state_type_iso,
+                                  desired_hidden_units=num_hidden_units)
+
+        obs_fn = EMLP(in_type=self.state_type_iso,
+                      out_type=self.obs_state_type,
+                      num_layers=num_layers,
+                      activation=act,
+                      **kwargs)
+        obs_fn_aux = None
+        if self.aux_obs_space:
+            obs_fn_aux = EMLP(in_type=self.state_type_iso,
+                              out_type=self.obs_state_type,
+                              num_layers=num_layers,
+                              activation=act,
+                              **kwargs)
+
+        return ObservableNet(obs_fn=obs_fn, obs_fn_aux=obs_fn_aux)
 
     def build_inv_obs_fn(self, num_layers, linear_decoder: bool, **kwargs):
         if linear_decoder:
@@ -307,7 +294,7 @@ def decoder(dpnet: DPNet, obs_state: Tensor):
                         **kwargs)
 
     def build_obs_dyn_module(self) -> MarkovDynamics:
-        return EquivLinearDynamics(state_rep=self.obs_state_type.representation,
+        return EquivLinearDynamics(state_type=self.obs_state_type,
                                    dt=self.dt,
                                    trainable=False,
                                    group_avg_trick=self.group_avg_trick)

nn/EquivDynamicsAutoencoder.py

@@ -107,19 +107,19 @@ def pre_process_obs_state(self,
                               ) -> dict[str, Tensor]:
         return super().pre_process_obs_state(obs_state_traj.tensor)
 
-    def post_process_obs_state(self, pred_state_traj: Tensor, **kwargs) -> dict[str, GeometricTensor]:
+    def post_process_obs_state(self, obs_state_traj: Tensor, **kwargs) -> dict[str, GeometricTensor]:
         """ Post-process the predicted observable state trajectory given by the observable state dynamics.
 
         Args:
-            pred_state_traj: (batch, time, obs_state_dim) Trajectory of the predicted (time -1) observable states
+            obs_state_traj: (batch, time, obs_state_dim) Trajectory of the predicted (time -1) observable states
              predicted by the transfer operator.
             **kwargs:
         Returns:
             Dictionary contraining
                 - pred_obs_state_traj: (batch * time, obs_state_dim) Geometric Tensor Trajectory
         """
-        flat_pred_obs_state_traj = batched_to_flat_trajectory(pred_state_traj)
-        return dict(pred_obs_state_traj=self.obs_state_type(flat_pred_obs_state_traj))
+        flat_obs_state_traj = batched_to_flat_trajectory(obs_state_traj)
+        return dict(obs_state_traj=self.obs_state_type(flat_obs_state_traj))
 
     def post_process_state(self, state_traj: GeometricTensor) -> Tensor:
         state_traj_input_basis = super().post_process_state(state_traj=state_traj.tensor)
@@ -138,7 +138,7 @@ def build_inv_obs_fn(self, num_layers: int, **kwargs):
                     **kwargs)
 
     def build_obs_dyn_module(self) -> MarkovDynamics:
-        return EquivLinearDynamics(state_type=self.state_type_iso,
+        return EquivLinearDynamics(state_type=self.obs_state_type,
                                    dt=self.dt,
                                    trainable=True,
                                    group_avg_trick=self.group_avg_trick)

nn/EquivLinearDynamics.py

@@ -46,7 +46,8 @@ def __init__(self,
         Q_iso2state = Tensor(Q_iso2state)
         Q_state2iso = Tensor(np.linalg.inv(Q_iso2state))
 
-        super(EquivLinearDynamics, self).__init__(state_rep=state_type.representation,
+        super(EquivLinearDynamics, self).__init__(state_dim=state_type.size,
+                                                  state_rep=state_type.representation,
                                                   dt=dt,
                                                   trainable=trainable,
                                                   dmd_algorithm=dmd_algorithm,
@@ -77,13 +78,10 @@ def forcast(self, state: GeometricTensor, n_steps: int = 1, **kwargs) -> Tensor:
                 next_obs_state = self.transfer_op(current_state)
             else:
                 transfer_op = self.get_transfer_op()
-                next_obs_state = torch.nn.functional.linear(current_state, transfer_op)
+                next_obs_state = self.state_type((transfer_op @ current_state.tensor.T).T)
             pred_state_traj.append(next_obs_state)
 
-        if self.is_trainable:
-            pred_state_traj = torch.stack([gt.tensor for gt in pred_state_traj], dim=1)
-        else:
-            pred_state_traj = torch.stack(pred_state_traj, dim=1)
+        pred_state_traj = torch.stack([gt.tensor for gt in pred_state_traj], dim=1)
         assert pred_state_traj.shape == (batch, n_steps + 1, state_dim)
         return pred_state_traj
 

nn/LightningLatentMarkovDynamics.py

@@ -116,10 +116,6 @@ def on_fit_start(self) -> None:
         self._loss_metrics_fn = loss_metrics_fn
 
     def on_train_start(self):
-        # TODO: Add number of layers and hidden channels dimensions.
-        hparams = flatten_dict(self._run_hps)
-        if hasattr(self.model, "get_hparams"):
-            hparams.update(flatten_dict(self.model.get_hparams()))
 
         if hasattr(self.model, "approximate_transfer_operator"):
             metrics = self.model.approximate_transfer_operator(self.trainer.datamodule.predict_dataloader())

nn/LinearDynamics.py

@@ -91,7 +91,6 @@ def forcast(self, state: Tensor, n_steps: int = 1, **kwargs) -> Tensor:
             pred_state_traj.append(next_obs_state)
 
         pred_state_traj = torch.stack(pred_state_traj, dim=1)
-        # a = pred_state_traj.detach().cpu().numpy()
         assert pred_state_traj.shape == (batch, n_steps + 1, state_dim)
         return pred_state_traj
 

nn/ObservableNet.py

@@ -1,46 +1,46 @@
-from typing import Optional
+import copy
+from typing import Optional, Union
 
+import escnn.nn
 import torch.nn
+from escnn.nn import EquivariantModule
 
+from nn.EquivLinearDynamics import EquivLinearDynamics
+from nn.LinearDynamics import LinearDynamics
 from nn.mlp import MLP
 from nn.emlp import EMLP
 
 
-class TwinMLP(torch.nn.Module):
-    """Auxiliary class to construct Twin MLPs with a potentially shared backbone."""
+class ObservableNet(torch.nn.Module):
 
-    def __init__(self, net_kwargs: dict, backbone_kwargs: Optional[dict] = None, equivariant=False, fake_aux_fn=False):
+    def __init__(self, 
+                 obs_fn: Union[torch.nn.Module, EquivariantModule],
+                 obs_fn_aux: Optional[Union[torch.nn.Module, EquivariantModule]] = None):
         super().__init__()
-        self.fake_aux_fn = fake_aux_fn
-        self.shared_backbone = backbone_kwargs is not None
-        mlp_class = MLP if not equivariant else EMLP # SO2MLP
+        self.equivariant = isinstance(obs_fn, EquivariantModule)
+        self.use_aux_obs_fn = obs_fn_aux is not None
 
-        if self.shared_backbone:
-            self.backbone = mlp_class(**backbone_kwargs)
-
-        self.fn1 = mlp_class(**net_kwargs)
-        if not fake_aux_fn:
-            self.fn2 = mlp_class(**net_kwargs)
+        self.obs = obs_fn
+        self.obs_aux = None
+        if self.use_aux_obs_fn:  # Use two twin networks to compute the main and auxiliary observable space.
+            self.obs_aux = obs_fn_aux
         else:
-            pass
+            if self.equivariant:
+                self.transfer_op_H_H_prime = escnn.nn.Linear(
+                    in_type=self.obs.out_type, out_type=self.obs.out_type, bias=False)
+            else:
+                self.transfer_op_H_H_prime = torch.nn.Linear(
+                    in_features=self.obs.out_dim, out_features=self.obs.out_dim, bias=False)
 
     def forward(self, input):
 
-        if self.shared_backbone:
-            backbone_output = self.backbone(input)
-            output1 = self.fn1(backbone_output)
-            output2 = self.fn2(backbone_output)
-        else:
-            if self.fake_aux_fn:
-                output1 = self.fn1(input)
-                output2 = output1
-            else:
-                output1 = self.fn1(input)
-                output2 = self.fn2(input)
+        obs_state = self.obs(input)
 
-        return output1, output2
+        if self.use_aux_obs_fn:
+            obs_aux_state = self.obs_aux(input)
+        else:
+            obs_aux_state = self.transfer_op_H_H_prime(obs_state)
 
-    def get_hparams(self):
-        return {}
+        return obs_state, obs_aux_state
 
 

nn/emlp.py

@@ -58,14 +58,11 @@ def __init__(self,
 
         if isinstance(activation, str):
             # Approximate the num of neurons as the num of signals in the space spawned by the irreps of the input type
-            self.num_hidden_regular_fields = int(np.ceil(num_hidden_units // self.in_type.size))
             # To compute the signal over the group we use all elements for finite groups
-            activation = self.get_activation(activation, in_type=in_type, channels=self.num_hidden_regular_fields)
+            activation = self.get_activation(activation, in_type=in_type, desired_hidden_units=num_hidden_units)
             hidden_type = activation.in_type
-            self.activation = activation
         elif isinstance(activation, EquivariantModule):
             hidden_type = activation.in_type
-            self.activation = activation
         else:
             raise ValueError(f"Activation type {type(activation)} not supported.")
 
@@ -104,17 +101,22 @@ def __init__(self,
         # self.net.check_equivariance()
 
     @staticmethod
-    def get_activation(activation, in_type: FieldType, channels: int):
+    def get_activation(activation, in_type: FieldType, desired_hidden_units: int):
         gspace = in_type.gspace
         group = gspace.fibergroup
         grid_length = group.order() if not group.continuous else 20
+
+        unique_irreps = set(in_type.irreps)
+        unique_irreps_dim = sum([group.irrep(*id).size for id in set(in_type.irreps)])
+        scale = in_type.size // unique_irreps_dim
+        channels = int(np.ceil(desired_hidden_units // unique_irreps_dim // scale))
         if "identity" in activation.lower():
             raise NotImplementedError("Identity activation not implemented yet")
             # return escnn.nn.IdentityModule()
         else:
             return escnn.nn.FourierPointwise(gspace,
                                              channels=channels,
-                                             irreps=in_type.irreps,
+                                             irreps=list(unique_irreps),
                                              function=f"p_{activation.lower()}",
                                              inplace=True,
                                              type='regular' if not group.continuous else 'rand',
@@ -124,12 +126,6 @@ def forward(self, x):
         """Forward pass of the EMLP model."""
         return self.net(x)
 
-    def get_hparams(self):
-        return {'num_layers': self.num_layers,
-                'hidden_ch':  self.num_hidden_regular_fields,
-                'activation': str(self.activation.__class__.__name__),
-                }
-
     def reset_parameters(self, init_mode=None):
         """Initialize weights and biases of E-MLP model."""
         raise NotImplementedError()