diff --git a/src/caustics/lenses/func/__init__.py b/src/caustics/lenses/func/__init__.py index bb8b71c3..0373f357 100644 --- a/src/caustics/lenses/func/__init__.py +++ b/src/caustics/lenses/func/__init__.py @@ -53,7 +53,11 @@ concentration_tnfw, ) -from .multipole import reduced_deflection_angle_multipole, potential_multipole, convergence_multipole +from .multipole import ( + reduced_deflection_angle_multipole, + potential_multipole, + convergence_multipole, +) __all__ = ( "forward_raytrace", diff --git a/src/caustics/lenses/func/multipole.py b/src/caustics/lenses/func/multipole.py index 0028d830..b7e856aa 100644 --- a/src/caustics/lenses/func/multipole.py +++ b/src/caustics/lenses/func/multipole.py @@ -34,17 +34,21 @@ def reduced_deflection_angle_multipole(x0, y0, m, a_m, phi_m, x, y): Equation (B11) and (B12) https://arxiv.org/pdf/1307.4220, Xu et al. 2014 """ x, y = translate_rotate(x, y, x0, y0) - + phi = torch.arctan2(y, x) - ax = torch.cos(phi) * a_m / (1 - m**2) * torch.cos(m * (phi - phi_m)) + torch.sin(phi) * m * a_m / (1 - m**2) * torch.sin(m * (phi - phi_m)) - ay = torch.sin(phi) * a_m / (1 - m**2) * torch.cos(m * (phi - phi_m)) - torch.cos(phi) * m * a_m / (1 - m**2) * torch.sin(m * (phi - phi_m)) + ax = torch.cos(phi) * a_m / (1 - m**2) * torch.cos(m * (phi - phi_m)) + torch.sin( + phi + ) * m * a_m / (1 - m**2) * torch.sin(m * (phi - phi_m)) + ay = torch.sin(phi) * a_m / (1 - m**2) * torch.cos(m * (phi - phi_m)) - torch.cos( + phi + ) * m * a_m / (1 - m**2) * torch.sin(m * (phi - phi_m)) return derotate(ax, ay, phi) def potential_multipole(x0, y0, m, a_m, phi_m, x, y): """ - Compute the lensing potential. + Compute the lensing potential. Parameters ---------- @@ -78,12 +82,12 @@ def potential_multipole(x0, y0, m, a_m, phi_m, x, y): x, y = translate_rotate(x, y, x0, y0) r = torch.sqrt(x**2 + y**2) phi = torch.arctan2(y, x) - return r * a_m / (1 - m**2) * torch.cos(m * (phi - phi_m)) + return r * a_m / (1 - m**2) * torch.cos(m * (phi - phi_m)) def convergence_multipole(x0, y0, m, a_m, phi_m, x, y): """ - Compute the lensing convergence. + Compute the lensing convergence. Parameters ---------- @@ -117,4 +121,4 @@ def convergence_multipole(x0, y0, m, a_m, phi_m, x, y): x, y = translate_rotate(x, y, x0, y0) r = torch.sqrt(x**2 + y**2) phi = torch.arctan2(y, x) - return 1/r * a_m * torch.cos(m * (phi - phi_m)) + return 1 / r * a_m * torch.cos(m * (phi - phi_m)) diff --git a/src/caustics/lenses/multipole.py b/src/caustics/lenses/multipole.py index 52311129..7381a1c4 100644 --- a/src/caustics/lenses/multipole.py +++ b/src/caustics/lenses/multipole.py @@ -31,22 +31,15 @@ class Multipole(ThinLens): Orientation of multiple. m: Optional[Union[Tensor, int]] Order of multipole. - + """ - _null_params = { - "x0": 0.0, - "y0": 0.0, - "a_m": 0.1, - "phi_m": 0., - "m": 3 - } + _null_params = {"x0": 0.0, "y0": 0.0, "a_m": 0.1, "phi_m": 0.0, "m": 3} def __init__( self, cosmology: CosmologyType, z_l: ZLType = None, - x0: Annotated[ Optional[Union[Tensor, float]], "The x-coordinate of the lens center", True ] = None, @@ -228,7 +221,7 @@ def convergence( The projected mass density. *Unit: unitless* - + Equation (B10) and (B3) https://arxiv.org/pdf/1307.4220, Xu et al. 2014 """ diff --git a/src/caustics/models/utils.py b/src/caustics/models/utils.py index 40d1c026..eecb5bd2 100644 --- a/src/caustics/models/utils.py +++ b/src/caustics/models/utils.py @@ -48,7 +48,7 @@ def _get_kwargs_field_definitions( if k != "name": anno = v.annotation dtype = anno.__origin__ - cls_param = ClassParam(*anno.__metadata__) # gooooo + cls_param = ClassParam(*anno.__metadata__) # gooooo if cls_param.isParam: kwargs_field_definitions[PARAMS][k] = ( dtype, diff --git a/tests/test_multipole.py b/tests/test_multipole.py index 6894df9a..3e37637c 100644 --- a/tests/test_multipole.py +++ b/tests/test_multipole.py @@ -1,6 +1,11 @@ import torch from lenstronomy.LensModel.lens_model import LensModel -from utils import lens_test_helper, alpha_test_helper, kappa_test_helper, Psi_test_helper +from utils import ( + lens_test_helper, + alpha_test_helper, + kappa_test_helper, + Psi_test_helper, +) import yaml from caustics.cosmology import FlatLambdaCDM @@ -37,9 +42,15 @@ def test(sim_source, device, lens_models): # Parameters m, a_m ,phi_m z_s = torch.tensor(1.2) - x = torch.tensor([-0.342, 0.51, 0.1, 3.14/4, 3]) + x = torch.tensor([-0.342, 0.51, 0.1, 3.14 / 4, 3]) kwargs_ls = [ - {"center_x": x[0].item(), "center_y": x[1].item(), "a_m": x[2].item(), "phi_m": x[3].item(), "m": x[4].item()} + { + "center_x": x[0].item(), + "center_y": x[1].item(), + "a_m": x[2].item(), + "phi_m": x[3].item(), + "m": x[4].item(), + } ] lens_test_helper(lens, lens_ls, z_s, x, kwargs_ls, rtol, atol, device=device)