style: Mark equations where black formatting should not apply (#122)

* Mark lenses equations to skip formatting

* Mark fmt skip for rest of caustics

* Fix super long equations

* bug test_kappa swapped z_l and z_s

* Pseudojaffe potential in units of arcsec maybe arcsec squared

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

src/caustics/constants.py

@@ -13,10 +13,12 @@
     "km_to_Mpc",
 )
 
+# fmt: off
 rad_to_arcsec = 180 / pi * 60**2
 arcsec_to_rad = 1 / rad_to_arcsec
 c_km_s = float(_c_astropy.to("km/s").value)
 G = float(_G_astropy.to("pc * km^2 / (s^2 * solMass)").value)
 G_over_c2 = float((_G_astropy / _c_astropy**2).to("Mpc/solMass").value)  # type: ignore
 c_Mpc_s = float(_c_astropy.to("Mpc/s").value)
 km_to_Mpc = 3.2407792896664e-20  # TODO: use astropy
+# fmt: on

src/caustics/cosmology.py

@@ -272,7 +272,7 @@ def critical_surface_density(
         d_l = self.angular_diameter_distance(z_l, params)
         d_s = self.angular_diameter_distance(z_s, params)
         d_ls = self.angular_diameter_distance_z1z2(z_l, z_s, params)
-        return d_s / (4 * pi * G_over_c2 * d_l * d_ls)
+        return d_s / (4 * pi * G_over_c2 * d_l * d_ls)  # fmt: skip
 
 
 class FlatLambdaCDM(Cosmology):
@@ -368,7 +368,7 @@ def critical_density(
             Critical density at redshift z.
         """
         Ode0 = 1 - Om0
-        return central_critical_density * (Om0 * (1 + z) ** 3 + Ode0)
+        return central_critical_density * (Om0 * (1 + z) ** 3 + Ode0)  # fmt: skip
 
     @unpack(1)
     def _comoving_distance_helper(
@@ -420,14 +420,10 @@ def comoving_distance(
         """
         Ode0 = 1 - Om0
         ratio = (Om0 / Ode0) ** (1 / 3)
-        return (
-            self.hubble_distance(h0)
-            * (
-                self._comoving_distance_helper((1 + z) * ratio, params)
-                - self._comoving_distance_helper(ratio, params)
-            )
-            / (Om0 ** (1 / 3) * Ode0 ** (1 / 6))
-        )
+        DH = self.hubble_distance(h0)
+        DC1z = self._comoving_distance_helper((1 + z) * ratio, params)
+        DC = self._comoving_distance_helper(ratio, params)
+        return DH * (DC1z - DC) / (Om0 ** (1 / 3) * Ode0 ** (1 / 6))  # fmt: skip
 
     @unpack(1)
     def transverse_comoving_distance(

src/caustics/lenses/base.py

@@ -647,7 +647,10 @@ def reduced_deflection_angle(
         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
+        return (
+            (d_ls / d_s) * deflection_angle_x,
+            (d_ls / d_s) * deflection_angle_y,
+        )
 
     @unpack(3)
     def physical_deflection_angle(
@@ -684,7 +687,10 @@ def physical_deflection_angle(
         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
+        return (
+            (d_s / d_ls) * deflection_angle_x,
+            (d_s / d_ls) * deflection_angle_y,
+        )
 
     @abstractmethod
     @unpack(3)
@@ -783,7 +789,7 @@ def surface_density(
         critical_surface_density = self.cosmology.critical_surface_density(
             z_l, z_s, params
         )
-        return self.convergence(x, y, z_s, params) * critical_surface_density
+        return self.convergence(x, y, z_s, params) * critical_surface_density  # fmt: skip
 
     @unpack(3)
     def raytrace(
@@ -854,9 +860,9 @@ def time_delay(
         d_ls = self.cosmology.angular_diameter_distance_z1z2(z_l, z_s, params)
         ax, ay = self.reduced_deflection_angle(x, y, z_s, params)
         potential = self.potential(x, y, z_s, params)
-        factor = (1 + z_l) / c_Mpc_s * d_s * d_l / d_ls
-        fp = 0.5 * d_ls**2 / d_s**2 * (ax**2 + ay**2) - potential
-        return factor * fp * arcsec_to_rad**2
+        factor = (1 + z_l) / c_Mpc_s * d_s * d_l / d_ls  # fmt: skip
+        fp = 0.5 * d_ls**2 / d_s**2 * (ax**2 + ay**2) - potential  # fmt: skip
+        return factor * fp * arcsec_to_rad**2  # fmt: skip
 
     @unpack(4)
     def _jacobian_deflection_angle_finitediff(

src/caustics/lenses/epl.py

@@ -164,7 +164,7 @@ def reduced_deflection_angle(
 
         # Tessore et al 2015 (eq. 23)
         r_omega = self._r_omega(z, t, q)
-        alpha_c = 2.0 / (1.0 + q) * (b / r) ** t * r_omega
+        alpha_c = 2.0 / (1.0 + q) * (b / r) ** t * r_omega  # fmt: skip
 
         alpha_real = torch.nan_to_num(alpha_c.real, posinf=10**10, neginf=-(10**10))
         alpha_imag = torch.nan_to_num(alpha_c.imag, posinf=10**10, neginf=-(10**10))
@@ -197,9 +197,9 @@ def _r_omega(self, z, t, q):
         part_sum = omega_i
 
         for i in range(1, self.n_iter):
-            factor = (2.0 * i - (2.0 - t)) / (2.0 * i + (2.0 - t))
-            omega_i = -f * factor * phi * omega_i
-            part_sum = part_sum + omega_i
+            factor = (2.0 * i - (2.0 - t)) / (2.0 * i + (2.0 - t))  # fmt: skip
+            omega_i = -f * factor * phi * omega_i  # fmt: skip
+            part_sum = part_sum + omega_i  # fmt: skip
 
         return part_sum
 
@@ -281,5 +281,5 @@ def convergence(
             The convergence of the lens.
         """
         x, y = translate_rotate(x, y, x0, y0, phi)
-        psi = (q**2 * (x**2 + self.s**2) + y**2).sqrt()
-        return (2 - t) / 2 * (b / psi) ** t
+        psi = (q**2 * (x**2 + self.s**2) + y**2).sqrt()  # fmt: skip
+        return (2 - t) / 2 * (b / psi) ** t  # fmt: skip

src/caustics/lenses/mass_sheet.py

@@ -104,7 +104,7 @@ def potential(
         **kwargs,
     ) -> tuple[Tensor, Tensor]:
         # Meneghetti eq 3.81
-        return surface_density * 0.5 * (x**2 + y**2)
+        return surface_density * 0.5 * (x**2 + y**2)  # fmt: skip
 
     @unpack(3)
     def convergence(

src/caustics/lenses/multiplane.py

@@ -136,7 +136,7 @@ def raytrace_z1z2(
         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
+        X, Y = x * arcsec_to_rad * D, y * arcsec_to_rad * D  # fmt: skip
 
         # Initial angles are observation angles
         # (negative needed because of negative in propagation term)
@@ -168,7 +168,10 @@ def raytrace_z1z2(
 
         # Convert from physical position to angular position on the source plane
         D_end = self.cosmology.transverse_comoving_distance_z1z2(z_start, z_end, params)
-        return X * rad_to_arcsec / D_end, Y * rad_to_arcsec / D_end
+        return (
+            X * rad_to_arcsec / D_end,
+            Y * rad_to_arcsec / D_end,
+        )
 
     @unpack(3)
     def effective_reduced_deflection_angle(

src/caustics/lenses/nfw.py

@@ -145,7 +145,7 @@ def get_scale_radius(
             The scale radius of the lens in Mpc.
         """
         critical_density = self.cosmology.critical_density(z_l, params)
-        r_delta = (3 * m / (4 * pi * DELTA * critical_density)) ** (1 / 3)
+        r_delta = (3 * m / (4 * pi * DELTA * critical_density)) ** (1 / 3)  # fmt: skip
         return 1 / c * r_delta
 
     @unpack(0)
@@ -169,13 +169,8 @@ def get_scale_density(
         Tensor
             The scale density of the lens in solar masses per Mpc cubed.
         """
-        return (
-            DELTA
-            / 3
-            * self.cosmology.critical_density(z_l, params)
-            * c**3
-            / ((1 + c).log() - c / (1 + c))
-        )
+        sigma_crit = self.cosmology.critical_density(z_l, params)
+        return DELTA / 3 * sigma_crit * c**3 / ((1 + c).log() - c / (1 + c))  # fmt: skip
 
     @unpack(1)
     def get_convergence_s(
@@ -205,11 +200,7 @@ def get_convergence_s(
         critical_surface_density = self.cosmology.critical_surface_density(
             z_l, z_s, params
         )
-        return (
-            self.get_scale_density(params)
-            * self.get_scale_radius(params)
-            / critical_surface_density
-        )
+        return  self.get_scale_density(params) * self.get_scale_radius(params) / critical_surface_density  # fmt: skip
 
     @staticmethod
     def _f_differentiable(x: Tensor) -> Tensor:
@@ -228,18 +219,8 @@ def _f_differentiable(x: Tensor) -> Tensor:
         """
         # TODO: generalize beyond torch, or patch Tensor
         f = torch.zeros_like(x)
-        f[x > 1] = (
-            1
-            - 2
-            / (x[x > 1] ** 2 - 1).sqrt()
-            * ((x[x > 1] - 1) / (x[x > 1] + 1)).sqrt().arctan()
-        )
-        f[x < 1] = (
-            1
-            - 2
-            / (1 - x[x < 1] ** 2).sqrt()
-            * ((1 - x[x < 1]) / (1 + x[x < 1])).sqrt().arctanh()
-        )
+        f[x > 1] = 1 - 2 / (x[x > 1] ** 2 - 1).sqrt() * ((x[x > 1] - 1) / (x[x > 1] + 1)).sqrt().arctan()  # fmt: skip
+        f[x < 1] = 1 - 2 / (1 - x[x < 1] ** 2).sqrt() * ((1 - x[x < 1]) / (1 + x[x < 1])).sqrt().arctanh()  # fmt: skip
         return f
 
     @staticmethod
@@ -258,16 +239,19 @@ def _f_batchable(x: Tensor) -> Tensor:
             Result of the deflection angle computation.
         """
         # TODO: generalize beyond torch, or patch Tensor
+        # fmt: off
         return torch.where(
             x > 1,
             1 - 2 / (x**2 - 1).sqrt() * ((x - 1) / (x + 1)).sqrt().arctan(),
             torch.where(
                 x < 1,
                 1 - 2 / (1 - x**2).sqrt() * ((1 - x) / (1 + x)).sqrt().arctanh(),
-                torch.zeros_like(x),  # x == 1
+                torch.zeros_like(x),  # where: x == 1
             ),
         )
 
+    # fmt: on
+
     @staticmethod
     def _g_differentiable(x: Tensor) -> Tensor:
         """
@@ -286,8 +270,8 @@ def _g_differentiable(x: Tensor) -> Tensor:
         # TODO: generalize beyond torch, or patch Tensor
         term_1 = (x / 2).log() ** 2
         term_2 = torch.zeros_like(x)
-        term_2[x > 1] = (1 / x[x > 1]).arccos() ** 2
-        term_2[x < 1] = -(1 / x[x < 1]).arccosh() ** 2
+        term_2[x > 1] = (1 / x[x > 1]).arccos() ** 2  # fmt: skip
+        term_2[x < 1] = -(1 / x[x < 1]).arccosh() ** 2  # fmt: skip
         return term_1 + term_2
 
     @staticmethod
@@ -313,7 +297,7 @@ def _g_batchable(x: Tensor) -> Tensor:
             torch.where(
                 x < 1,
                 -(1 / x).arccosh() ** 2,
-                torch.zeros_like(x),  # x == 1
+                torch.zeros_like(x),  # where: x == 1
             ),
         )
         return term_1 + term_2
@@ -335,8 +319,8 @@ def _h_differentiable(x: Tensor) -> Tensor:
         """
         term_1 = (x / 2).log()
         term_2 = torch.ones_like(x)
-        term_2[x > 1] = (1 / x[x > 1]).arccos() * 1 / (x[x > 1] ** 2 - 1).sqrt()
-        term_2[x < 1] = (1 / x[x < 1]).arccosh() * 1 / (1 - x[x < 1] ** 2).sqrt()
+        term_2[x > 1] = (1 / x[x > 1]).arccos() * 1 / (x[x > 1] ** 2 - 1).sqrt()  # fmt: skip
+        term_2[x < 1] = (1 / x[x < 1]).arccosh() * 1 / (1 - x[x < 1] ** 2).sqrt()  # fmt: skip
         return term_1 + term_2
 
     @staticmethod
@@ -357,9 +341,9 @@ def _h_batchable(x: Tensor) -> Tensor:
         term_1 = (x / 2).log()
         term_2 = torch.where(
             x > 1,
-            (1 / x).arccos() * 1 / (x**2 - 1).sqrt(),
+            (1 / x).arccos() * 1 / (x**2 - 1).sqrt(),  # fmt: skip
             torch.where(
-                x < 1, (1 / x).arccosh() * 1 / (1 - x**2).sqrt(), torch.ones_like(x)
+                x < 1, (1 / x).arccosh() * 1 / (1 - x**2).sqrt(), torch.ones_like(x)  # fmt: skip
             ),
         )
         return term_1 + term_2
@@ -405,22 +389,13 @@ def reduced_deflection_angle(
         xi = d_l * th * arcsec_to_rad
         r = xi / scale_radius
 
-        deflection_angle = (
-            16
-            * pi
-            * G_over_c2
-            * self.get_scale_density(params)
-            * scale_radius**3
-            * self._h(r)
-            * rad_to_arcsec
-            / xi
-        )
+        deflection_angle = 16 * pi * G_over_c2 * self.get_scale_density(params) * scale_radius**3 * self._h(r) * rad_to_arcsec / xi  # fmt: skip
 
         ax = deflection_angle * x / th
         ay = deflection_angle * y / th
         d_s = self.cosmology.angular_diameter_distance(z_s, params)
         d_ls = self.cosmology.angular_diameter_distance_z1z2(z_l, z_s, params)
-        return ax * d_ls / d_s, ay * d_ls / d_s
+        return ax * d_ls / d_s, ay * d_ls / d_s  # fmt: skip
 
     @unpack(3)
     def convergence(
@@ -463,7 +438,7 @@ def convergence(
         xi = d_l * th * arcsec_to_rad
         r = xi / scale_radius  # xi / xi_0
         convergence_s = self.get_convergence_s(z_s, params)
-        return 2 * convergence_s * self._f(r) / (r**2 - 1)
+        return 2 * convergence_s * self._f(r) / (r**2 - 1)  # fmt: skip
 
     @unpack(3)
     def potential(
@@ -506,10 +481,4 @@ def potential(
         xi = d_l * th * arcsec_to_rad
         r = xi / scale_radius  # xi / xi_0
         convergence_s = self.get_convergence_s(z_s, params)
-        return (
-            2
-            * convergence_s
-            * self._g(r)
-            * scale_radius**2
-            / (d_l**2 * arcsec_to_rad**2)
-        )
+        return 2 * convergence_s * self._g(r) * scale_radius**2 / (d_l**2 * arcsec_to_rad**2)  # fmt: skip

src/caustics/lenses/pixelated_convergence.py

@@ -190,9 +190,7 @@ def _unpad_fft(self, x: Tensor) -> Tensor:
         Tensor
             The input tensor without padding.
         """
-        return torch.roll(x, (-self._s[0] // 2, -self._s[1] // 2), dims=(-2, -1))[
-            ..., : self.n_pix, : self.n_pix
-        ]
+        return torch.roll(x, (-self._s[0] // 2, -self._s[1] // 2), dims=(-2, -1))[..., : self.n_pix, : self.n_pix]  # fmt: skip
 
     def _unpad_conv2d(self, x: Tensor) -> Tensor:
         """
@@ -342,9 +340,8 @@ def _deflection_angle_conv2d(
         # we actually want the cross-correlation.
 
         2 * self.n_pix
-        convergence_map_flipped = convergence_map.flip((-1, -2))[
-            None, None
-        ]  # noqa: E501 F.pad(, ((pad - self.n_pix)//2, (pad - self.n_pix)//2, (pad - self.n_pix)//2, (pad - self.n_pix)//2), mode = self.padding_mode)
+        convergence_map_flipped = convergence_map.flip((-1, -2))[None, None]
+        # noqa: E501 F.pad(, ((pad - self.n_pix)//2, (pad - self.n_pix)//2, (pad - self.n_pix)//2, (pad - self.n_pix)//2), mode = self.padding_mode)
         deflection_angle_x = F.conv2d(
             self.ax_kernel[None, None], convergence_map_flipped, padding="same"
         ).squeeze() * (self.pixelscale**2 / pi)