style: pre-commit fixes

Ciela-Institute · Dec 16, 2023 · 687e651 · 687e651
1 parent de575e6
commit 687e651
Show file tree

Hide file tree

Showing 5 changed files with 98 additions and 192 deletions.
diff --git a/docs/source/tutorials/InvertLensEquation.ipynb b/docs/source/tutorials/InvertLensEquation.ipynb
diff --git a/docs/source/tutorials/MultiplaneDemo.ipynb b/docs/source/tutorials/MultiplaneDemo.ipynb
diff --git a/docs/source/tutorials/VisualizeCaustics.ipynb b/docs/source/tutorials/VisualizeCaustics.ipynb
diff --git a/src/caustics/lenses/base.py b/src/caustics/lenses/base.py
@@ -64,7 +64,9 @@ def jacobian_lens_equation(
                     "and known pixelscale. "
                     "Please include the pixelscale argument"
                 )
-            return self._jacobian_lens_equation_finitediff(x, y, z_s, pixelscale, params, **kwargs)
+            return self._jacobian_lens_equation_finitediff(
+                x, y, z_s, pixelscale, params, **kwargs
+            )
         else:
             raise ValueError("method should be one of: autograd, finitediff")
 
@@ -145,7 +147,9 @@ def forward_raytrace(
             raise ValueError("fov must be given to generate initial guesses")
 
         # Random starting points in image plane
-        guesses = torch.as_tensor(fov) * (torch.rand(n_init, 2) - 0.5)  # Has shape (n_init, Din:2)
+        guesses = torch.as_tensor(fov) * (
+            torch.rand(n_init, 2) - 0.5
+        )  # Has shape (n_init, Din:2)
 
         # Optimize guesses in image plane
         x, l, c = batch_lm(  # noqa: E741 Unused `l` variable
@@ -528,7 +532,9 @@ def _jacobian_lens_equation_autograd(
         This equates to a (2,2) matrix at each (x,y) point.
         """
         # Build Jacobian
-        J = self._jacobian_effective_deflection_angle_autograd(x, y, z_s, params, **kwargs)
+        J = self._jacobian_effective_deflection_angle_autograd(
+            x, y, z_s, params, **kwargs
+        )
         return torch.eye(2) - J.detach()
 
     @unpack(3)
@@ -638,7 +644,9 @@ def reduced_deflection_angle(
         """
         d_s = self.cosmology.angular_diameter_distance(z_s, params)
         d_ls = self.cosmology.angular_diameter_distance_z1z2(z_l, z_s, params)
-        deflection_angle_x, deflection_angle_y = self.physical_deflection_angle(x, y, z_s, params)
+        deflection_angle_x, deflection_angle_y = self.physical_deflection_angle(
+            x, y, z_s, params
+        )
         return (
             (d_ls / d_s) * deflection_angle_x,
             (d_ls / d_s) * deflection_angle_y,
@@ -676,7 +684,9 @@ def physical_deflection_angle(
         """
         d_s = self.cosmology.angular_diameter_distance(z_s, params)
         d_ls = self.cosmology.angular_diameter_distance_z1z2(z_l, z_s, params)
-        deflection_angle_x, deflection_angle_y = self.reduced_deflection_angle(x, y, z_s, params)
+        deflection_angle_x, deflection_angle_y = self.reduced_deflection_angle(
+            x, y, z_s, params
+        )
         return (
             (d_s / d_ls) * deflection_angle_x,
             (d_s / d_ls) * deflection_angle_y,
@@ -776,7 +786,9 @@ def surface_density(
         Tensor
             Surface mass density at the given coordinates in solar masses per Mpc^2.
         """
-        critical_surface_density = self.cosmology.critical_surface_density(z_l, z_s, params)
+        critical_surface_density = self.cosmology.critical_surface_density(
+            z_l, z_s, params
+        )
         return self.convergence(x, y, z_s, params) * critical_surface_density  # fmt: skip
 
     @unpack(3)
@@ -847,8 +859,12 @@ def time_delay_gravitational(
         d_ls = self.cosmology.angular_diameter_distance_z1z2(z_l, z_s, params)
         potential = self.potential(x, y, z_s, params)
         factor = (1 + z_l) / c_Mpc_s * d_s * d_l / d_ls
-        print(self.potential(0., 0., z_s, params))
-        return -factor * (potential - self.potential(0., 0., z_s, params)) * arcsec_to_rad**2
+        print(self.potential(0.0, 0.0, z_s, params))
+        return (
+            -factor
+            * (potential - self.potential(0.0, 0.0, z_s, params))
+            * arcsec_to_rad**2
+        )
 
     @unpack(5)
     def time_delay_geometric_source(
@@ -958,7 +974,6 @@ def time_delay_geometric(
         factor = (1 + z_l) / c_Mpc_s * d_s * d_l / d_ls
         fp = 0.5 * (ax**2 + ay**2)
         return factor * fp * arcsec_to_rad**2
-
 
     @unpack(3)
     def time_delay(
@@ -1091,7 +1106,9 @@ def jacobian_deflection_angle(
                     "Finite differences lensing jacobian requires regular grid "
                     "and known pixelscale. Please include the pixelscale argument"
                 )
-            return self._jacobian_deflection_angle_finitediff(x, y, z_s, pixelscale, params)
+            return self._jacobian_deflection_angle_finitediff(
+                x, y, z_s, pixelscale, params
+            )
         else:
             raise ValueError("method should be one of: autograd, finitediff")
 
@@ -1111,7 +1128,9 @@ def _jacobian_lens_equation_finitediff(
         This equates to a (2,2) matrix at each (x,y) point.
         """
         # Build Jacobian
-        J = self._jacobian_deflection_angle_finitediff(x, y, z_s, pixelscale, params, **kwargs)
+        J = self._jacobian_deflection_angle_finitediff(
+            x, y, z_s, pixelscale, params, **kwargs
+        )
         return torch.eye(2) - J
 
     @unpack(3)

diff --git a/src/caustics/lenses/multiplane.py b/src/caustics/lenses/multiplane.py
@@ -39,7 +39,9 @@ def __init__(self, cosmology: Cosmology, lenses: list[ThinLens], name: str = Non
             self.add_parametrized(lens)
 
     @unpack(0)
-    def get_z_ls(self, *args, params: Optional["Packed"] = None, **kwargs) -> list[Tensor]:
+    def get_z_ls(
+        self, *args, params: Optional["Packed"] = None, **kwargs
+    ) -> list[Tensor]:
         """
         Get the redshifts of each lens in the multiplane.
 
@@ -131,7 +133,9 @@ def raytrace_z1z2(
         lens_planes = [i for i, _ in sorted(enumerate(z_ls), key=itemgetter(1))]
 
         # Compute physical position on first lens plane
-        D = self.cosmology.transverse_comoving_distance_z1z2(z_start, z_ls[lens_planes[0]], params)
+        D = self.cosmology.transverse_comoving_distance_z1z2(
+            z_start, z_ls[lens_planes[0]], params
+        )
         X, Y = x * arcsec_to_rad * D, y * arcsec_to_rad * D  # fmt: skip
 
         # Initial angles are observation angles
@@ -140,7 +144,9 @@ def raytrace_z1z2(
 
         for i in lens_planes:
             # Compute deflection angle at current ray positions
-            D_l = self.cosmology.transverse_comoving_distance_z1z2(z_start, z_ls[i], params)
+            D_l = self.cosmology.transverse_comoving_distance_z1z2(
+                z_start, z_ls[i], params
+            )
             alpha_x, alpha_y = self.lenses[i].physical_deflection_angle(
                 X * rad_to_arcsec / D_l,
                 Y * rad_to_arcsec / D_l,
@@ -154,7 +160,9 @@ def raytrace_z1z2(
 
             # Propagate rays to next plane (basically eq 18)
             z_next = z_ls[i + 1] if i != lens_planes[-1] else z_end
-            D = self.cosmology.transverse_comoving_distance_z1z2(z_ls[i], z_next, params)
+            D = self.cosmology.transverse_comoving_distance_z1z2(
+                z_ls[i], z_next, params
+            )
             X = X + D * theta_x * arcsec_to_rad
             Y = Y + D * theta_y * arcsec_to_rad
 
@@ -214,7 +222,7 @@ def surface_density(
         """
         # TODO: rescale mass densities of each lens and sum
         raise NotImplementedError()
-    
+
     @unpack(3)
     def time_delay_gravitational(
         self,
@@ -268,8 +276,12 @@ def time_delay_gravitational(
             z_next = z_ls[i + 1] if i != lens_planes[-1] else z_s
             # Compute deflection angle at current ray positions
             D_l = self.cosmology.transverse_comoving_distance(z_ls[i], params)
-            D = self.cosmology.transverse_comoving_distance_z1z2(z_ls[i], z_next, params)
-            D_is = self.cosmology.transverse_comoving_distance_z1z2(z_ls[i], z_s, params)
+            D = self.cosmology.transverse_comoving_distance_z1z2(
+                z_ls[i], z_next, params
+            )
+            D_is = self.cosmology.transverse_comoving_distance_z1z2(
+                z_ls[i], z_s, params
+            )
             D_next = self.cosmology.transverse_comoving_distance(z_next, params)
             alpha_x, alpha_y = self.lenses[i].physical_deflection_angle(
                 X * rad_to_arcsec / D_l,
@@ -295,7 +307,7 @@ def time_delay_gravitational(
             Y = Y + D * theta_y * arcsec_to_rad
 
         return TD
-    
+
     @unpack(3)
     def time_delay_geometric(
         self,
@@ -349,8 +361,12 @@ def time_delay_geometric(
             z_next = z_ls[i + 1] if i != lens_planes[-1] else z_s
             # Compute deflection angle at current ray positions
             D_l = self.cosmology.transverse_comoving_distance(z_ls[i], params)
-            D = self.cosmology.transverse_comoving_distance_z1z2(z_ls[i], z_next, params)
-            D_is = self.cosmology.transverse_comoving_distance_z1z2(z_ls[i], z_s, params)
+            D = self.cosmology.transverse_comoving_distance_z1z2(
+                z_ls[i], z_next, params
+            )
+            D_is = self.cosmology.transverse_comoving_distance_z1z2(
+                z_ls[i], z_s, params
+            )
             D_next = self.cosmology.transverse_comoving_distance(z_next, params)
             alpha_x, alpha_y = self.lenses[i].physical_deflection_angle(
                 X * rad_to_arcsec / D_l,
@@ -427,8 +443,12 @@ def time_delay(
             z_next = z_ls[i + 1] if i != lens_planes[-1] else z_s
             # Compute deflection angle at current ray positions
             D_l = self.cosmology.transverse_comoving_distance(z_ls[i], params)
-            D = self.cosmology.transverse_comoving_distance_z1z2(z_ls[i], z_next, params)
-            D_is = self.cosmology.transverse_comoving_distance_z1z2(z_ls[i], z_s, params)
+            D = self.cosmology.transverse_comoving_distance_z1z2(
+                z_ls[i], z_next, params
+            )
+            D_is = self.cosmology.transverse_comoving_distance_z1z2(
+                z_ls[i], z_s, params
+            )
             D_next = self.cosmology.transverse_comoving_distance(z_next, params)
             alpha_x, alpha_y = self.lenses[i].physical_deflection_angle(
                 X * rad_to_arcsec / D_l,