add mean cell back to the return value

curtischong · Apr 30, 2024 · 54bcd51 · 54bcd51
1 parent 08a8e83
commit 54bcd51
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/diffusion/diffusion_helpers.py b/diffusion/diffusion_helpers.py
@@ -178,7 +178,7 @@ def forward(self, l0_vector, t, num_atoms):
         #     * eps
         # )
         ht = mean + variance
-        return ht, eps
+        return ht, eps, mean_cell
 
     def reverse(self, lt, predicted_symmetric_vector_noise, t):
         alpha = 1 - self.betas[t]

diff --git a/diffusion/diffusion_loss.py b/diffusion/diffusion_loss.py
@@ -199,7 +199,7 @@ def diffuse_lattice_params(
         rotation_matrix, symmetric_matrix = polar_decomposition(lattice)
         symmetric_matrix_vector = symmetric_matrix_to_vector(symmetric_matrix)
 
-        noisy_symmetric_vector, noise_vector = self.lattice_diffusion(
+        noisy_symmetric_vector, noise_vector, _mean_cell = self.lattice_diffusion(
             symmetric_matrix_vector, t_int, num_atoms
         )
         noisy_symmetric_matrix = vector_to_symmetric_matrix(noisy_symmetric_vector)

diff --git a/exploration/verify_vp_limited_mean_and_var.py b/exploration/verify_vp_limited_mean_and_var.py
@@ -35,7 +35,7 @@ def main():
         symmetric_matrix_vector = symmetric_matrix_to_vector(symmetric_matrix)
         num_atoms = torch.tensor([8])
         noisy_symmetric_vector, _symmetric_vector_noise, mean_cell = vp(
-            symmetric_matrix_vector, t, num_atoms, real_lattice
+            symmetric_matrix_vector, t, num_atoms
         )
 
         noisy_symmetric_matrix = vector_to_symmetric_matrix(noisy_symmetric_vector)