Format code with black and isort

compSPI · Jul 9, 2021 · 7e9b58a · 7e9b58a
1 parent e07dc29
commit 7e9b58a
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 15 deletions.
diff --git a/sim_volumes.py b/sim_volumes.py
@@ -8,7 +8,7 @@
 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
 
 
@@ -84,10 +84,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
 
 
@@ -117,7 +125,7 @@ def simulate_volumes(particules, n_volumes, vol_size, center=2):
     rots, qs = 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)
@@ -157,9 +165,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,6 +1,8 @@
 """Test sim_volumes."""
-import numpy as np
 import os
+
+import numpy as np
+
 import sim_volumes
 
 os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
@@ -34,11 +36,10 @@ def test_simulate_volumes():
         particules = np.asarray([[1, 0], [0, 8], [0, 0]])
         n_volumes = 1
         vol_size = 64
-        volumes, qs = sim_volumes.simulate_volumes(
-            particules, n_volumes, vol_size)
+        volumes, qs = sim_volumes.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():
@@ -48,5 +49,6 @@ def test_save_volume():
         main_dir = ""
         name = "2particules"
         volumes, qs = sim_volumes.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)