Fix missing detach in compute_metrics

Author: hummerichsander

README.md

@@ -17,7 +17,7 @@
 Split-flows provide a probabilistic bridge between molecular resolutions, enabling conditional backmapping and direct measurement of the configuration-dependent (local) information loss.
 
 <div align="center">
-    <img src="figures/flow_trajectory.gif" alt="Flow Trajectory" style="max-width: 300px; width: 50%;">
+    <img src="figures/flow_trajectory.gif" alt="Flow Trajectory" style="max-width: 300px; width: 50%; animation-iteration-count: 1;">
 </div>
 
 ## Installation

split_flows/models/split_flow.py

@@ -75,9 +75,9 @@ def augment_gmm(self, R: Tensor, gmm: GaussianMixture) -> Tensor:
         :return: Full set of coordinates with noise."""
 
         z = torch.empty((R.shape[0], self.num_particles, R.shape[2]), device=R.device)
-        z_gmm = torch.tensor(gmm.sample(R.shape[0])[0], dtype=R.dtype, device=R.device).view(
-            R.shape[0], -1, 3
-        )
+        z_gmm = torch.tensor(
+            gmm.sample(R.shape[0])[0], dtype=R.dtype, device=R.device
+        ).view(R.shape[0], -1, 3)
 
         start_idx = 0
         for i, (cg_idx, noise_idx) in enumerate(self.latent_groupings):
@@ -118,7 +118,9 @@ def log_prob(self, value: Tensor) -> Tensor:
             R_cg = value[:, cg_idx, :]
             R_noise = value[:, noise_idx, :]
             exponential_term = (
-                -0.5 * sum_except_batch((R_noise - R_cg[:, None, :]) ** 2) / self.scale**2
+                -0.5
+                * sum_except_batch((R_noise - R_cg[:, None, :]) ** 2)
+                / self.scale**2
             )
             normalization_term = -torch.log(Z) * R_noise.shape[1]
             log_prob += exponential_term + normalization_term
@@ -188,7 +190,9 @@ def _init_weights(self):
 
                 # Coordinate MLP - Xavier uniform for all but last layer
                 coord_mlp_layers = list(layer.coors_mlp.modules())
-                linear_layers = [m for m in coord_mlp_layers if isinstance(m, nn.Linear)]
+                linear_layers = [
+                    m for m in coord_mlp_layers if isinstance(m, nn.Linear)
+                ]
 
                 for i, m in enumerate(linear_layers):
                     if i == len(linear_layers) - 1:
@@ -312,7 +316,9 @@ def velocity(self, xt: Tensor, t: Tensor) -> Tensor:
 
         return self.velo_net(xt, t)
 
-    def compute_metrics(self, batch: tuple[Tensor, ...], batch_idx: int) -> dict[str, Tensor]:
+    def compute_metrics(
+        self, batch: tuple[Tensor, ...], batch_idx: int
+    ) -> dict[str, Tensor]:
         """Compute training/validation metrics.
 
         :param batch: Batch data tuple, expecting (r,) where r is a Tensor.
@@ -337,8 +343,10 @@ def compute_metrics(self, batch: tuple[Tensor, ...], batch_idx: int) -> dict[str
 
         if not self.training:
             x1 = self.compute_flow(x0, return_intermediate=False, verbose=False)
-            traj = md.Trajectory(x1.cpu().numpy(), self.top_aa)
-            metrics["ged"] = torch.mean(torch.tensor(graph_edit_distance(traj=traj, verbose=False)))
+            traj = md.Trajectory(x1.detach().cpu().numpy(), self.top_aa)
+            metrics["ged"] = torch.mean(
+                torch.tensor(graph_edit_distance(traj=traj, verbose=False))
+            )
 
         return metrics
 
@@ -395,10 +403,12 @@ def fit_latent_gmm(
 
         with torch.no_grad():
             x1 = r.to(self.device)
-            x0 = self.compute_flow(x1, reverse=True, chunk_size=chunk_size, verbose=verbose).cpu()
-            eps_sn = self.noise.to_standard_normal(x0)[:, self.indices_split[1].cpu(), :].view(
-                x0.shape[0], -1
-            )
+            x0 = self.compute_flow(
+                x1, reverse=True, chunk_size=chunk_size, verbose=verbose
+            ).cpu()
+            eps_sn = self.noise.to_standard_normal(x0)[
+                :, self.indices_split[1].cpu(), :
+            ].view(x0.shape[0], -1)
 
         gmm = GaussianMixture(n_components=n_components, *args, **kwargs)
         gmm.fit(eps_sn.numpy())

split_flows/utils/metrics.py

@@ -147,7 +147,9 @@ def compute_bond_cutoff_mdtraj(topology, scale=1.3):
     """Compute bond cutoffs for MDTraj topology"""
     atomic_nums = [atom.element.atomic_number for atom in topology.atoms]
     # COVCUTOFFTABLE values are in Angstroms, convert to nanometers for MDTraj
-    vdw_array = torch.Tensor([COVCUTOFFTABLE[int(el)] / 10.0 for el in atomic_nums])  # Å to nm
+    vdw_array = torch.Tensor(
+        [COVCUTOFFTABLE[int(el)] / 10.0 for el in atomic_nums]
+    )  # Å to nm
 
     cutoff_array = (vdw_array[None, :] + vdw_array[:, None]) * scale
 

split_flows/utils/utils.py

@@ -80,5 +80,7 @@ def gradient(
 
     if grad_outputs is None:
         grad_outputs = torch.ones_like(output).detach()
-    grad = torch.autograd.grad(output, x, grad_outputs=grad_outputs, create_graph=create_graph)[0]
+    grad = torch.autograd.grad(
+        output, x, grad_outputs=grad_outputs, create_graph=create_graph
+    )[0]
     return grad