Format code with black and isort

compSPI · Jul 9, 2021 · cece6a2 · cece6a2
1 parent 5b397d2
commit cece6a2
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 14 deletions.
diff --git a/sim_volumes.py b/sim_volumes.py
@@ -1,13 +1,13 @@
 """Generate 3D map of molecules."""
 
 import numpy as np
-import coords
 
+import coords
 
 PARTICULES = np.asarray([[1, 0, 0, 3], [0, 8, 0, 5], [0, 0, 7, 9]])
 N_VOLUMES = 2000
 VOL_SIZE = 64
-NAME = '4_points_3d'
+NAME = "4_points_3d"
 CENTER = 2
 
 
@@ -36,10 +36,18 @@ def modify_weight(points, volume, vol_size, center):
         for i in range(vol_size):
             for j in range(vol_size):
                 for k in range(vol_size):
-                    volume[i][j][k] += np.exp(-np.linalg.norm(
-                        [i/center-vol_size/center/2, j/center -
-                         vol_size/center/2, k/center-vol_size/center/2] -
-                        point)**2/2)
+                    volume[i][j][k] += np.exp(
+                        -np.linalg.norm(
+                            [
+                                i / center - vol_size / center / 2,
+                                j / center - vol_size / center / 2,
+                                k / center - vol_size / center / 2,
+                            ]
+                            - point
+                        )
+                        ** 2
+                        / 2
+                    )
     return volume
 
 
@@ -69,7 +77,7 @@ def simulate_volumes(particules, n_volumes, vol_size, center=2):
     rots, qs = coords.uniform_rotations(n_volumes)
     volumes = np.zeros((n_volumes,) + (vol_size,) * 3)
     for idx in range(n_volumes):
-        if idx % (n_volumes/10) == 0:
+        if idx % (n_volumes / 10) == 0:
             print(idx)
         points = rots[idx].dot(particules)
         volumes[idx] = modify_weight(points, volumes[idx], vol_size, center)
@@ -109,9 +117,8 @@ def save_volume(particules, n_volumes, vol_size, main_dir, name, center=2):
     >>>import coords
     _ = save_volume(PARTICULES,N_VOLUMES,VOL_SIZE, dir, NAME, CENTER)
     """
-    volumes, labels = simulate_volumes(
-        particules, n_volumes, vol_size, center)
-    np.save(main_dir + name + '_molecules.npy', volumes)
-    np.save(main_dir + name + '_labels.npy', volumes)
+    volumes, labels = simulate_volumes(particules, n_volumes, vol_size, center)
+    np.save(main_dir + name + "_molecules.npy", volumes)
+    np.save(main_dir + name + "_labels.npy", volumes)
 
     return volumes, labels
diff --git a/test_sim_volumes.py b/test_sim_volumes.py
@@ -1,7 +1,9 @@
 """Test sim_volumes."""
-import numpy as np
 import os
+
+import numpy as np
 import torch
+
 import sim_volumes as sv
 
 os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
@@ -26,7 +28,7 @@ def test_simulate_volumes(self):
         volumes, qs = sv.simulate_volumes(particules, n_volumes, vol_size)
         assert volumes.shape == (1, 64, 64, 64)
         assert len(qs) == 1
-        assert qs[0][0] == '['
+        assert qs[0][0] == "["
 
     @staticmethod
     def test_save_volume(self):
@@ -36,5 +38,6 @@ def test_save_volume(self):
         main_dir = ""
         name = "2particules"
         volumes, qs = sv.save_volume(
-            particules, n_volumes, vol_size, main_dir, name, center=2)
+            particules, n_volumes, vol_size, main_dir, name, center=2
+        )
         assert volumes.shape == (1, 64, 64, 64)