apply formatting

Author: Mayukhdeb

tests/test_laplacian_pyramid.py

@@ -6,36 +6,54 @@
 from topoloss.losses.laplacian_pyramid import laplacian_pyramid_loss
 
 supported_dtypes = [
-    torch.float32, 
-    # torch.float16, 
+    torch.float32,
+    # torch.float16,
     # torch.bfloat16
 ]
 
+
 @pytest.mark.parametrize("dtype1", supported_dtypes)
 @pytest.mark.parametrize("dtype2", supported_dtypes)
 @pytest.mark.parametrize("height, width", [(16, 16), (32, 32)])
 @pytest.mark.parametrize("factor_w, factor_h", [(2.0, 2.0), (3.0, 3.0), (4.0, 4.0)])
 @pytest.mark.parametrize("interpolation", ["bilinear", "nearest"])
-def test_laplacian_pyramid_loss_precision(dtype1, dtype2, height, width, factor_w, factor_h, interpolation):
+def test_laplacian_pyramid_loss_precision(
+    dtype1, dtype2, height, width, factor_w, factor_h, interpolation
+):
     # Create a sample cortical_sheet tensor
     e = 16  # Example depth
     torch.manual_seed(42)  # Set seed for reproducibility
     cortical_sheet = torch.rand(height, width, e).to(dtype1)
 
     # Call the function with the given precision
-    loss = laplacian_pyramid_loss(cortical_sheet, factor_w=factor_w, factor_h=factor_h, interpolation=interpolation)
+    loss = laplacian_pyramid_loss(
+        cortical_sheet,
+        factor_w=factor_w,
+        factor_h=factor_h,
+        interpolation=interpolation,
+    )
 
     # Ensure the loss is finite
     assert torch.isfinite(loss), f"Loss is not finite for dtype {dtype1}"
-    
+
     # Check type consistency
-    assert loss.dtype == dtype1, f"Loss dtype {loss.dtype} does not match input dtype {dtype1}."
+    assert (
+        loss.dtype == dtype1
+    ), f"Loss dtype {loss.dtype} does not match input dtype {dtype1}."
 
     # Compare results to float32 (considered ground truth for higher precision)
-    if dtype1 !=dtype2:
+    if dtype1 != dtype2:
         float32_sheet = cortical_sheet.to(dtype2)
-        expected_loss = laplacian_pyramid_loss(cortical_sheet=float32_sheet, factor_w=factor_w, factor_h=factor_h, interpolation=interpolation)
+        expected_loss = laplacian_pyramid_loss(
+            cortical_sheet=float32_sheet,
+            factor_w=factor_w,
+            factor_h=factor_h,
+            interpolation=interpolation,
+        )
         assert_close(
-            loss.to(dtype2), expected_loss, rtol=1e-3, atol=1e-4,
-            msg=f"Loss mismatch for dtype1 {dtype1} loss: {loss.to(dtype2)} and dtype2 {dtype2} loss2: {expected_loss}"
-        )
+            loss.to(dtype2),
+            expected_loss,
+            rtol=1e-3,
+            atol=1e-4,
+            msg=f"Loss mismatch for dtype1 {dtype1} loss: {loss.to(dtype2)} and dtype2 {dtype2} loss2: {expected_loss}",
+        )

tests/test_loss_conv.py

@@ -7,17 +7,22 @@
 supported_dtypes = [
     torch.float32,
     # torch.float16,
-    #torch.bfloat16,
+    # torch.bfloat16,
 ]
 
+
 # Define the fixture that provides the num_steps argument
 @pytest.mark.parametrize("num_steps", [2, 9])
 @pytest.mark.parametrize("hidden_channels", [16, 32])
 @pytest.mark.parametrize("init_from_layer", [True, False])
 @pytest.mark.parametrize("dtype", supported_dtypes)
 @pytest.mark.parametrize("interpolation", ["bilinear", "nearest"])
 def test_loss_conv(
-    num_steps: int, hidden_channels: int, init_from_layer: bool, dtype, interpolation: str
+    num_steps: int,
+    hidden_channels: int,
+    init_from_layer: bool,
+    dtype,
+    interpolation: str,
 ):  # num_steps is now passed by the fixture
 
     # Define the model
@@ -31,16 +36,38 @@ def test_loss_conv(
     if init_from_layer:
         losses = [
             LaplacianPyramid.from_layer(
-                model=model, layer=model[0], scale=1.0, factor_h=3.0, factor_w=3.0, interpolation=interpolation
+                model=model,
+                layer=model[0],
+                scale=1.0,
+                factor_h=3.0,
+                factor_w=3.0,
+                interpolation=interpolation,
             ),
             LaplacianPyramid.from_layer(
-                model=model, layer=model[2], scale=1.0, factor_h=3.0, factor_w=3.0, interpolation=interpolation
+                model=model,
+                layer=model[2],
+                scale=1.0,
+                factor_h=3.0,
+                factor_w=3.0,
+                interpolation=interpolation,
             ),
         ]
     else:
         losses = [
-            LaplacianPyramid(layer_name="0", scale=1.0, factor_h=3.0, factor_w=3.0, interpolation=interpolation),
-            LaplacianPyramid(layer_name="2", scale=1.0, factor_h=3.0, factor_w=3.0, interpolation=interpolation),
+            LaplacianPyramid(
+                layer_name="0",
+                scale=1.0,
+                factor_h=3.0,
+                factor_w=3.0,
+                interpolation=interpolation,
+            ),
+            LaplacianPyramid(
+                layer_name="2",
+                scale=1.0,
+                factor_h=3.0,
+                factor_w=3.0,
+                interpolation=interpolation,
+            ),
         ]
 
     # Define the TopoLoss

tests/test_loss_linear.py

@@ -10,6 +10,7 @@
     # torch.bfloat16, ### tests fail on bfloat16
 ]
 
+
 # Define the fixture that provides the num_steps argument
 @pytest.mark.parametrize("num_steps", [2, 9])
 @pytest.mark.parametrize("hidden_size", [30, 25])

topoloss/__init__.py

@@ -1 +1,6 @@
-from .core import TopoLoss, LaplacianPyramid, LaplacianPyramidOnBias, LaplacianPyramidOnInput
+from .core import (
+    TopoLoss,
+    LaplacianPyramid,
+    LaplacianPyramidOnBias,
+    LaplacianPyramidOnInput,
+)

topoloss/cortical_sheet/input.py

@@ -13,7 +13,7 @@ def get_cortical_sheet_linear_input(layer: nn.Linear):
     return rearrange(
         weight,
         "o (h w) -> h w o",
-        h = cortical_sheet_size.height,
-        w = cortical_sheet_size.width,
-        o = weight.shape[0]
+        h=cortical_sheet_size.height,
+        w=cortical_sheet_size.width,
+        o=weight.shape[0],
     )

topoloss/losses/laplacian_pyramid.py

@@ -9,8 +9,12 @@
 from dataclasses import dataclass, field
 import torch
 
+
 def laplacian_pyramid_loss(
-    cortical_sheet: TensorType["height", "width", "e"], factor_w: float, factor_h: float, interpolation: str = "bilinear"
+    cortical_sheet: TensorType["height", "width", "e"],
+    factor_w: float,
+    factor_h: float,
+    interpolation: str = "bilinear",
 ):
     grid = cortical_sheet
     assert grid.ndim == 3, "Expected grid to be a 3d tensor of shape (h, w, e)"
@@ -27,7 +31,9 @@ def laplacian_pyramid_loss(
         grid, scale_factor=(1 / factor_h, 1 / factor_w), mode=interpolation
     )
     # Upscale the downscaled grid tensor
-    upscaled_grid = F.interpolate(downscaled_grid, size=grid.shape[2:], mode=interpolation)
+    upscaled_grid = F.interpolate(
+        downscaled_grid, size=grid.shape[2:], mode=interpolation
+    )
 
     # Calculate the MSE loss between the original grid and upscaled grid
     # loss = F.mse_loss(upscaled_grid, grid)
@@ -54,19 +60,30 @@ class LaplacianPyramid:
     scale: Optional[Union[None, float]] = field(default=1.0)
 
     @classmethod
-    def from_layer(cls, model, layer, factor_h, factor_w, scale=1.0, interpolation: str ="bilinear"):
+    def from_layer(
+        cls,
+        model,
+        layer,
+        factor_h,
+        factor_w,
+        scale=1.0,
+        interpolation: str = "bilinear",
+    ):
         layer_name = get_name_by_layer(model=model, layer=layer)
         return cls(
             layer_name=layer_name,
             scale=scale,
             factor_h=factor_h,
             factor_w=factor_w,
-            interpolation=interpolation
+            interpolation=interpolation,
         )
 
 
 def laplacian_pyramid_loss_on_bias(
-    cortical_sheet: TensorType["h", "w"], factor_w: float, factor_h: float, interpolation: str = "bilinear"
+    cortical_sheet: TensorType["h", "w"],
+    factor_w: float,
+    factor_h: float,
+    interpolation: str = "bilinear",
 ):
 
     grid = cortical_sheet
@@ -86,7 +103,9 @@ def laplacian_pyramid_loss_on_bias(
         grid, scale_factor=(1 / factor_h, 1 / factor_w), mode=interpolation
     )
     # Upscale the downscaled grid tensor
-    upscaled_grid = F.interpolate(downscaled_grid, size=grid.shape[2:], mode=interpolation)
+    upscaled_grid = F.interpolate(
+        downscaled_grid, size=grid.shape[2:], mode=interpolation
+    )
 
     grid = rearrange(grid.squeeze(0).squeeze(0), "h w -> (h w)").unsqueeze(0)
     upscaled_grid = rearrange(
@@ -113,7 +132,15 @@ class LaplacianPyramidOnBias:
     scale: Optional[Union[None, float]] = field(default=1.0)
 
     @classmethod
-    def from_layer(cls, model, layer, factor_h, factor_w, scale=1.0, interpolation: str ="bilinear"):
+    def from_layer(
+        cls,
+        model,
+        layer,
+        factor_h,
+        factor_w,
+        scale=1.0,
+        interpolation: str = "bilinear",
+    ):
         assert (
             layer.bias is not None
         ), "Expected layer to have a bias, but got None. *sad sad sad*"
@@ -123,9 +150,10 @@ def from_layer(cls, model, layer, factor_h, factor_w, scale=1.0, interpolation:
             scale=scale,
             factor_h=factor_h,
             factor_w=factor_w,
-            interpolation=interpolation
+            interpolation=interpolation,
         )
 
+
 @dataclass
 class LaplacianPyramidOnInput:
     """
@@ -144,12 +172,20 @@ class LaplacianPyramidOnInput:
     scale: Optional[Union[None, float]] = field(default=1.0)
 
     @classmethod
-    def from_layer(cls, model, layer, factor_h, factor_w, scale=1.0, interpolation: str ="bilinear"):
+    def from_layer(
+        cls,
+        model,
+        layer,
+        factor_h,
+        factor_w,
+        scale=1.0,
+        interpolation: str = "bilinear",
+    ):
         layer_name = get_name_by_layer(model=model, layer=layer)
         return cls(
             layer_name=layer_name,
             scale=scale,
             factor_h=factor_h,
             factor_w=factor_w,
-            interpolation=interpolation
-        )
+            interpolation=interpolation,
+        )