formatted code with black

src/cryo_sbi/inference/priors.py

@@ -15,7 +15,7 @@ def gen_quat() -> torch.Tensor:
     count = 0
     while count < 1:
         quat = 2 * torch.rand(size=(4,)) - 1
-        norm = torch.sqrt(torch.sum(quat**2))
+        norm = torch.sqrt(torch.sum(quat ** 2))
         if 0.2 <= norm <= 1.0:
             quat /= norm
             count += 1
@@ -207,7 +207,7 @@ class PriorLoader(DataLoader):
     def __init__(
         self,
         prior: Distribution,
-        batch_size: int = 2**8,  # 256
+        batch_size: int = 2 ** 8,  # 256
         **kwargs,
     ):
         super().__init__(

src/cryo_sbi/utils/image_utils.py

@@ -25,9 +25,9 @@ def circular_mask(n_pixels: int, radius: int, inside: bool = True) -> torch.Tens
     r_2d = grid[None, :] ** 2 + grid[:, None] ** 2
 
     if inside is True:
-        mask = r_2d < radius**2
+        mask = r_2d < radius ** 2
     else:
-        mask = r_2d > radius**2
+        mask = r_2d > radius ** 2
 
     return mask
 
@@ -183,7 +183,7 @@ def __init__(self, image_size: int, sigma: int):
             -0.5 * (image_size - 1), 0.5 * (image_size - 1), image_size
         )
         self._r_2d = self._grid[None, :] ** 2 + self._grid[:, None] ** 2
-        self._mask = torch.exp(-self._r_2d / (2 * sigma**2))
+        self._mask = torch.exp(-self._r_2d / (2 * sigma ** 2))
 
     def __call__(self, image: torch.Tensor) -> torch.Tensor:
         """

src/cryo_sbi/wpa_simulator/ctf.py

@@ -21,7 +21,7 @@ def apply_ctf(image: torch.Tensor, defocus, b_factor, amp, pixel_size) -> torch.
     freq_pix_1d = torch.fft.fftfreq(num_pixels, d=pixel_size, device=image.device)
     x, y = torch.meshgrid(freq_pix_1d, freq_pix_1d, indexing="ij")
 
-    freq2_2d = x**2 + y**2
+    freq2_2d = x ** 2 + y ** 2
     freq2_2d = freq2_2d.expand(num_batch, -1, -1)
     imag = torch.zeros_like(freq2_2d, device=image.device) * 1j
 
@@ -30,7 +30,7 @@ def apply_ctf(image: torch.Tensor, defocus, b_factor, amp, pixel_size) -> torch.
 
     ctf = (
         -amp * torch.cos(phase * freq2_2d * 0.5)
-        - torch.sqrt(1 - amp**2) * torch.sin(phase * freq2_2d * 0.5)
+        - torch.sqrt(1 - amp ** 2) * torch.sin(phase * freq2_2d * 0.5)
         + imag
     )
     ctf = ctf * env / amp

src/cryo_sbi/wpa_simulator/image_generation.py

@@ -13,7 +13,7 @@ def gen_quat() -> torch.Tensor:
     count = 0
     while count < 1:
         quat = 2 * torch.rand(size=(4,)) - 1
-        norm = torch.sqrt(torch.sum(quat**2))
+        norm = torch.sqrt(torch.sum(quat ** 2))
         if 0.2 <= norm <= 1.0:
             quat /= norm
             count += 1
@@ -72,7 +72,7 @@ def project_density(
     """
 
     num_batch, _, num_atoms = coords.shape
-    norm = 1 / (2 * torch.pi * sigma**2 * num_atoms)
+    norm = 1 / (2 * torch.pi * sigma ** 2 * num_atoms)
 
     grid_min = -pixel_size * num_pixels * 0.5
     grid_max = pixel_size * num_pixels * 0.5

src/cryo_sbi/wpa_simulator/noise.py

@@ -19,7 +19,7 @@ def circular_mask(n_pixels: int, radius: int, device: str = "cpu") -> torch.Tens
         -0.5 * (n_pixels - 1), 0.5 * (n_pixels - 1), n_pixels, device=device
     )
     r_2d = grid[None, :] ** 2 + grid[:, None] ** 2
-    mask = r_2d < radius**2
+    mask = r_2d < radius ** 2
 
     return mask
 

tests/test_image_utils.py

@@ -11,9 +11,9 @@ def test_circular_mask():
 
     assert inside_mask.shape == (n_pixels, n_pixels)
     assert outside_mask.shape == (n_pixels, n_pixels)
-    assert inside_mask.sum().item() == pytest.approx(radius**2 * 3.14159, abs=10)
+    assert inside_mask.sum().item() == pytest.approx(radius ** 2 * 3.14159, abs=10)
     assert outside_mask.sum().item() == pytest.approx(
-        n_pixels**2 - radius**2 * 3.14159, abs=10
+        n_pixels ** 2 - radius ** 2 * 3.14159, abs=10
     )
 
 
@@ -27,7 +27,7 @@ def test_mask_class():
     masked_image = mask(image)
     assert masked_image.shape == (image_size, image_size)
     assert masked_image[inside].sum().item() == pytest.approx(
-        image_size**2 - radius**2 * 3.14159, abs=10
+        image_size ** 2 - radius ** 2 * 3.14159, abs=10
     )
 
 

tests/test_simulator.py

@@ -80,6 +80,7 @@ def test_get_snr(noise_std, num_images):
         snr.flatten(), noise_std * torch.ones(images.shape[0]), atol=1e-01
     ), "SNR is not correct"
 
+
 @pytest.mark.parametrize(("num_images"), [2, 3, 10])
 def test_simulator_default_settings(num_images):
     sim = CryoEmSimulator("tests/config_files/image_params_testing.json")
@@ -94,5 +95,3 @@ def test_simulator_custom_indices():
 
     assert (parameters[0] == test_indices).all().item()
     assert images.shape == torch.Size([6, 64, 64])
-
-