Equiv-DAE operational

This commit restructures the DPNET and DAE algorithms under the same base class LatendMarkovDynamics module. Both DPNet and DAE have their own equivariant implementations operational in this commit.

Author: Danfoa

cfg/config.yaml

@@ -67,6 +67,8 @@ hydra:
           - system.pred_horizon
           - system.eval_pred_horizon
           - system.data_path
+          - system.group
+          - system.noise_level
           - debug
           - debug_loops
           - seed

cfg/model/base_model.yaml

@@ -3,12 +3,12 @@
 name: ??
 # Dataset related parameters
 augment: False          # Weather to use data-augmentation of the dataset if the system has symmetries.
-# Architectural and Task parameters
-equivariant: True      # Impose equivariance constraints on NNs
+# Model hyperparameters
 activation: ReLU
-num_layers: 4             # Number MLPs' layers (including input and output layers)
-bias: True              # Use bias in the MLPs
-batch_norm: True        # Use batch normalization layers before activation functions in MLPs
+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

cfg/model/dae.yaml

@@ -3,14 +3,13 @@ defaults:
 
 name: DAE
 # Model hyperparameters
-loss_pred_w: 0.01    # Cost function weight for prediction in observation space Z
-equivariant: True
-activation: ReLU
-n_layers: 4             # Number MLPs' layers (including input and output layers)
+obs_pred_w: 1.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
+equivariant: False
 
-#eigval_init: "stable"   # Initialization of eigenvalues [stable, ]
-#eigval_constraint: "unconstrained"   # [unconstrained or unit_circle]
+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/edae.yaml

@@ -0,0 +1,13 @@
+defaults:
+  - dae
+
+name: E-DAE
+# Model hyperparameters
+obs_pred_w: 1.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
+state_dependent_obs_dyn: False                  # Whether to use state-dependent observation dynamics
+group_avg_trick: True
+
+equivariant: True
+

cfg/system/linear_system.yaml

@@ -8,12 +8,12 @@ obs_state_dim: 2             # Dimension of the system's observable state
 
 frames_per_state: 1          # Number of time-frames to use as a Markov Process state time step
 pred_horizon: 10             # Number (or percentage) of Markov Process state time steps to predict into the future
-eval_pred_horizon: 100       # Number (or percentage) of Markov Process state time steps to predict into the future
+eval_pred_horizon: .5        # Number (or percentage) of Markov Process state time steps to predict into the future
 
 group: SO(2)
-noise_level: 0
+noise_level: 1
 n_constraints: 0
 
-data_path: '${system.name}/group=${system.group}-dim=${system.state_dim}/'
+data_path: '${system.name}/group=${system.group}-dim=${system.state_dim}/n_constraints=${system.n_constraints}'
 
 summary: S:${system.state_dim}-OS:${system.obs_state_dim}-H:${system.pred_horizon}-G:${system.group}-N:${system.noise_level}

dynamical_systems/stable_linear_system.py

@@ -6,6 +6,7 @@
 import numpy as np
 import scipy
 from escnn.group import Representation
+from tqdm import tqdm
 
 from data.DynamicsRecording import DynamicsRecording
 from utils.mysc import companion_matrix, matrix_average_trick, random_orthogonal_matrix
@@ -26,6 +27,7 @@ def sample_initial_condition(state_dim, P=None, z=None):
     distance_from_origin = max(distance_from_origin, MIN_DISTANCE_FROM_ORIGIN)  # Truncate unlikely low values
     x0 = distance_from_origin * direction
 
+    trials = 500
     if P is not None:
         violation = P @ x0 < z
         is_constraint_violated = np.any(violation)
@@ -42,6 +44,10 @@ def sample_initial_condition(state_dim, P=None, z=None):
 
             violation = P @ x0 < z
             is_constraint_violated = np.any(violation)
+
+            trials -= 1
+            if trials == 0:
+                raise RuntimeError("Too constrained.")
         if np.linalg.norm(x0) < MIN_DISTANCE_FROM_ORIGIN:  # If sample is too close to zero ignore it.
             x0 = sample_initial_condition(state_dim, P=P, z=z)
     return x0
@@ -152,7 +158,7 @@ def stable_equivariant_lin_dynamics(rep_X: Representation, time_constant=1, min_
         iso_state_dim = rep_iso.size
         A_iso = stable_lin_dynamics(rep_iso,
                                     time_constant=time_constant,
-                                    stable_eigval_prob=1 / (iso_state_dim) if state_dim > 1 else 0.0,
+                                    stable_eigval_prob=1 / (iso_state_dim + 1) if state_dim > 1 else 0.0,
                                     min_period=min_period,
                                     max_period=max_period)
         # Enforce G-equivariance
@@ -247,7 +253,7 @@ def evolve_linear_dynamics(A: np.ndarray, init_state: np.ndarray, dt: float, sim
 if __name__ == '__main__':
     np.set_printoptions(precision=3)
 
-    order = 3
+    order = 2
     subgroups_ids = dict(C2=('cone', 1),
                          Tetrahedral=('fulltetra',),
                          Octahedral=(True, 'octa',),
@@ -264,7 +270,8 @@ def evolve_linear_dynamics(A: np.ndarray, init_state: np.ndarray, dt: float, sim
     G, g_dynamics_2_Gsub_domain, g_domain_2_g_dynamics = G_domain.subgroup(G_id)
 
     # Define the state representation.
-    rep_X = G.standard_representation()  # + G.irrep(1)
+    # rep_X = G.regular_representation # + G.irrep(1)
+    rep_X = G.irrep(0) + G.standard_representation()  # + G.irrep(1)
     # rep_X = G.irrep(1) + G.irrep(2) #+ G.irrep(1) #+ G.irrep(0)
     #
     # Generate stable equivariant linear dynamics withing a range of fast and slow dynamics
@@ -279,11 +286,11 @@ def evolve_linear_dynamics(A: np.ndarray, init_state: np.ndarray, dt: float, sim
         T = fastest_period   # Simulate until the slowest stable mode has completed a full period.
     else:  # System has transient dynamics that vanish to 36.8% in fastest_time_constant seconds.
         T = 6 * fastest_time_constant        # Required time for this transient dynamics to vanish.
-    dt = T * 0.005  # Sample time to obtain 200 samples per trajectory
+    dt = T * 0.005  # Sample time to obtain 100 samples per trajectory
 
     # Generate trajectories of the system dynamics
     n_constraints = 0
-    n_trajs = 100
+    n_trajs = 120
     # Generate hyperplanes that constraint outer region of space
     P_symm, offset = None, None
     if n_constraints > 0:
@@ -293,12 +300,12 @@ def evolve_linear_dynamics(A: np.ndarray, init_state: np.ndarray, dt: float, sim
         for normal_plane in normal_planes:
             normal_orbit = np.vstack([np.linalg.det(rep_X(g)) * (rep_X(g) @ normal_plane) for g in G.elements])
             # Fix point of linear systems is the origin
-            offset_orbit = np.asarray([-np.random.uniform(-0.05, 0.6)] * normal_orbit.shape[0])
+            offset_orbit = np.asarray([-np.random.uniform(-0.1, 0.3)] * normal_orbit.shape[0])
             P_symm = np.vstack((P_symm, normal_orbit)) if P_symm is not None else normal_orbit
             offset = np.concatenate((offset, offset_orbit)) if offset is not None else offset_orbit
 
     trajs_per_noise_level = []
-    for noise_level in range(10):
+    for noise_level in tqdm(range(10), desc="noise level"):
         sigma = T * 0.005 * noise_level
         state_trajs = []
         for _ in range(n_trajs):
@@ -380,9 +387,12 @@ def evolve_linear_dynamics(A: np.ndarray, init_state: np.ndarray, dt: float, sim
                                      fig=fig, constraint_matrix=P_symm, constraint_offset=offset,
                                      traj_colorscale='Agsunset', init_state_color='yellow',
                                      legendgroup="val")
+            else:
+                pass
 
-            fig.write_html(path_2_system / 'test_trajectories.html')
-        if noise_level == 0 and fig is not None:
+            if fig is not None:
+                fig.write_html(path_2_system / 'test_trajectories.html')
+        if noise_level == 1 and fig is not None:
             fig.show()
     # fig.show()
     print(f"Recordings saved to {path_2_system}")

nn/DPNet.py nn/DeepProjections.py

@@ -89,7 +89,6 @@ def __init__(
                                     dt=dt,
                                     **markov_dyn_params)
 
-
     def forecast(self, state: Tensor, n_steps: int = 1, **kwargs) -> [dict[str, Tensor]]:
         """Forward pass of the dynamics model, producing a prediction of the next `n_steps` states.
 
@@ -107,15 +106,33 @@ def forecast(self, state: Tensor, n_steps: int = 1, **kwargs) -> [dict[str, Tens
         obs_state = self.obs_fn(state)
         pred_obs_state_traj = self.obs_state_dynamics.forcast(state=obs_state, n_steps=n_steps)
         pred_state_traj = self.inv_obs_fn(pred_obs_state_traj)
-        if self.transfer_op is None:
-            raise RuntimeError("The transfer operator not approximated yet. Call `approximate_transfer_operator`")
 
         assert pred_state_traj.shape == (self._batch_size, time_horizon, self.state_dim), \
             f"{pred_state_traj.shape}!=({self._batch_size}, {time_horizon}, {self.state_dim})"
         assert pred_obs_state_traj.shape == (self._batch_size, time_horizon, self.obs_state_dim), \
             f"{pred_obs_state_traj.shape}!=({self._batch_size}, {time_horizon}, {self.obs_state_dim})"
+        raise NotImplementedError("This function needs to handle pre/post state processing")
         return pred_state_traj, pred_obs_state_traj
 
+    def pre_process_obs_state(self,
+                              obs_state_traj: Tensor,
+                              obs_state_traj_aux: Optional[Tensor] = None) -> dict[str, Tensor]:
+        """ Apply transformations to the observable state trajectory.
+        Args:
+            obs_state_traj: (batch * time, obs_state_dim) or (batch, time, obs_state_dim)
+            obs_state_traj_aux: (batch * time, obs_state_dim) or (batch, time, obs_state_dim)
+        Returns:
+            Directory containing
+                - obs_state_traj: (batch, time, obs_state_dim) tensor.
+                - obs_state_traj_aux: (batch, time, obs_state_dim) tensor.
+        """
+        obs_state_traj = super().pre_process_obs_state(obs_state_traj)['obs_state_traj']
+        if obs_state_traj_aux is not None:
+            obs_state_traj_aux = super().pre_process_obs_state(obs_state_traj_aux)['obs_state_traj']
+        else:
+            obs_state_traj_aux = None
+        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,
@@ -277,7 +294,7 @@ def approximate_transfer_operator(self, train_data_loader: DataLoader):
         # Obtain the observable state
         obs_fn_output = self.obs_fn(state_traj)
         # Post process observable state
-        obs_state_trajs = self.post_process_obs_state(*obs_fn_output)
+        obs_state_trajs = self.pre_process_obs_state(*obs_fn_output)
         obs_state_traj = obs_state_trajs.pop('obs_state_traj')
 
         assert obs_state_traj.shape[1] == 2, f"Expected single step datapoints, got {obs_state_traj.shape[1]} steps."
@@ -320,7 +337,7 @@ def build_obs_fn(self, num_layers, identity=False, **kwargs):
 
         return TwinMLP(net_kwargs=obs_fn_params, backbone_kwargs=backbone_params, fake_aux_fn=not self.aux_obs_space)
 
-    def build_inv_obs_fn(self, num_layers, linear_decoder: bool,  identity=False, **kwargs):
+    def build_inv_obs_fn(self, num_layers, linear_decoder: bool, identity=False, **kwargs):
         if identity:
             return lambda x: x
 
@@ -385,7 +402,7 @@ def get_hparams(self):
     state_dim = 2
     obs_state_dim = state_dim
 
-    change_of_basis = None# Tensor(random_orthogonal_matrix(state_dim))
+    change_of_basis = None  # Tensor(random_orthogonal_matrix(state_dim))
 
     test_dpnet = DPNet(state_dim=state_dim,
                        obs_state_dim=obs_state_dim,
@@ -398,5 +415,5 @@ def get_hparams(self):
     pred_state_traj = out['pred_state_traj']
 
     assert pred_state_traj.shape == random_state_traj.shape, f"{pred_state_traj.shape} != {random_state_traj.shape}"
-    assert torch.allclose(pred_state_traj, random_state_traj, rtol=1e-5, atol=1e-5), f"{pred_state_traj - random_state_traj}"
-
+    assert torch.allclose(pred_state_traj, random_state_traj, rtol=1e-5,
+                          atol=1e-5), f"{pred_state_traj - random_state_traj}"