Added rotation angle parametrization

[mjyb16]

Adding ability to create an external shear from angle and magnitude instead of gamma 1 and gamma 2.

[ConnorStoneAstro]

Looks great! Just needs some tests. You can basically copy paste any shear related tests.

[AlexandreAdam]

We don't want an extra class, we need to choose the parametrization at the initialization of the class.

Furthermore, we need to remove the x0 and y0 in the parameterization. They are not standard parameters. We should make them default to 0 at the very least.

The way to do it is to use a flag in the initialization to specify the parametrization, either polar or cartesian. Cartesian is useful for model fitting, while polar is more intuitive. The method for the extra class should be folded into the original class as private method. Then, use the flag to decide which method you want to use at inference time.

If you want to avoid control flows at inference time, then you can always attach the method you know will be used to some variable at initialization. Here is an example.

class Template:
     def init(self, some_flag=True):
           if some_flag:
                  self.method_to_use = self._method_a
           else:
                  self.method_to_use = self._method_b

Make sure the angle is defined east of North, East pointing to the left.

[ConnorStoneAstro]

I don't agree about merging the classes (though I can be convinced). For one, the docs can be a lot clearer this way, each class describes its parameterization, we don't need to clarify for every function "when toggle true it means this, when toggle false it means that". For two, having a state like this can create silent errors, when someone thinks they are in one parameterization but they aren't, and you don't know it because either way there are two numbers in the "x" vector. For three, we need to think about other classes that may face this situation. There are other profiles with different parameterizations, and they might not all be so easy to swap between, and there may be more than two parameterizations. Making a new class for a new parameterization scales to as many cases as we want, while adding multiple toggles can get clunky fast.

What do you think?

[AlexandreAdam]

Having multiple classes to deal with different parametrization is not a good solution for 2 reasons

1. Adds complexity for maintenance. The classes are separate and any changes to the class would need to be tracked across the different classes which are independent
2. Two different parametrizations are not different objects. Semantically, two parameterizations point to the same underlying representation of the object. So the different parameterizations belong in the same class.

This second point is quite important to me. Instead, we can modify how the object behaves internally to solve the 2 issues with the previous approach I just mentioned. I propose the following implementation, which is fully implemented and tested for the sake of argument. I removed the x0 and y0 in the class, which are not parameters of the external shear.

from typing import Optional, Union, Annotated, Callable
from functools import wraps

from torch import Tensor
import torch

from .base import ThinLens, CosmologyType, NameType, ZLType
from ..parametrized import unpack
from ..packed import Packed

__all__ = ("ExternalShear",)


def convert_params(method: Callable):
    @wraps(method)
    def wrapper(self, *args, **kwargs):
        if hasattr(self, '_convert_params'):
            kwargs = self._convert_params(*args, **kwargs)
        return method(self, *args, **kwargs)
    return wrapper


class ExternalShear(ThinLens):
    def __init__(
        self,
        cosmology: CosmologyType,
        z_l: ZLType = None,
        gamma_1: Annotated[
            Optional[Union[Tensor, float]], "Shear component in the x-direction", True
        ] = None,
        gamma_2: Annotated[
            Optional[Union[Tensor, float]], "Shear component in the y-direction", True
        ] = None,
        gamma: Annotated[
            Optional[Union[Tensor, float]], "Shear magnitude", True
        ] = None,
        phi: Annotated[
            Optional[Union[Tensor, float]], "Shear angle", True
        ] = None,
        parametrization: Annotated[
            str, "Parametrization of the shear field", {"cartesian", "polar"}
        ] = "cartesian",
        s: Annotated[
            float, "Softening length for the elliptical power-law profile"
        ] = 0.0,
        name: NameType = None,
    ):
        super().__init__(cosmology, z_l, name=name)
        if parametrization.lower() == "cartesian":
            self.add_param("gamma_1", gamma_1)
            self.add_param("gamma_2", gamma_2)
        elif parametrization.lower() == "polar":
            self.add_param("gamma", gamma)
            self.add_param("phi", phi)
        else:
            raise ValueError(
                f"parametrization must be either 'cartesian' or 'polar', got {parametrization}"
            )
        self.s = s
        if parametrization.lower() == "cartesian":
            self._convert_params = lambda self, *args, **kwargs: kwargs # do nothing
        elif parametrization.lower() == "polar":
            self._convert_params = self._convert_polar_to_cartesian # convert polar parameters to cartesian
    
    def _convert_polar_to_cartesian(self, *args, **kwargs):
        gamma = kwargs.get('gamma')
        phi = kwargs.get('phi')
        # This breaks if gamma or phi are not provided (or are None)
        gamma_1, gamma_2 = self._polar_to_cartesian(gamma, phi)
        kwargs['gamma_1'] = gamma_1
        kwargs['gamma_2'] = gamma_2
        return kwargs
    
    @staticmethod
    def _polar_to_cartesian(gamma: Tensor, phi: Tensor) -> tuple[Tensor, Tensor]:
        gamma_1 = gamma * torch.cos(2*phi)
        gamma_2 = gamma * torch.sin(2*phi)
        return gamma_1, gamma_2
    
    @staticmethod
    def _cartesian_to_polar(gamma_1: Tensor, gamma_2: Tensor) -> tuple[Tensor, Tensor]:
        gamma = torch.sqrt(gamma_1**2 + gamma_2**2)
        phi = 0.5 * torch.atan2(gamma_2, gamma_1)
        return gamma, phi

    @unpack
    @convert_params
    def potential(
        self,
        x: Tensor,
        y: Tensor,
        z_s: Tensor,
        *args,
        params: Optional["Packed"] = None,
        z_l: Optional[Tensor] = None,
        gamma_1: Optional[Tensor] = None,
        gamma_2: Optional[Tensor] = None,
        gamma: Optional[Tensor] = None,
        phi: Optional[Tensor] = None,
        **kwargs,
    ) -> Tensor:
        # Equation 5.127 of Meneghetti et al. 2019
        return 0.5 * gamma_1 * (x**2 - y**2) + gamma_2 * x * y
        
    
    @unpack
    @convert_params
    def reduced_deflection_angle(
        self,
        x: Tensor,
        y: Tensor,
        z_s: Tensor,
        *args,
        params: Optional["Packed"] = None,
        z_l: Optional[Tensor] = None,
        gamma_1: Optional[Tensor] = None,
        gamma_2: Optional[Tensor] = None,
        gamma: Optional[Tensor] = None,
        phi: Optional[Tensor] = None,
        **kwargs,
    ) -> tuple[Tensor, Tensor]:
        # Derivative of the potential
        a1 = x * gamma_1 + y * gamma_2
        a2 = x * gamma_2 - y * gamma_1
        return a1, a2

    @unpack
    @convert_params
    def convergence(
        self,
        x: Tensor,
        y: Tensor,
        z_s: Tensor,
        *args,
        params: Optional["Packed"] = None,
        z_l: Optional[Tensor] = None,
        gamma_1: Optional[Tensor] = None,
        gamma_2: Optional[Tensor] = None,
        gamma: Optional[Tensor] = None,
        phi: Optional[Tensor] = None,
        **kwargs,
    ) -> Tensor:
        # By definition, convergence is zero for external shear
        return torch.zeros_like(x)

As a summary, my solution can be described as follows:

1. Define a bijective transformation between the two parametrizations. These will always exist by construction.
2. Define a fiducial parametrization (in that case cartesian), with which we implement all the methods. This is to avoid the problem of maintaining multiple methods/classes.
3. Use the transformation to map to the fiducial parametrization. This is done statically at the level of init, so there is no control flow added at inference time.

My approach is complete and numerically efficient. I think we should adopt it for all other similar cases. I am thinking about the NFW and issue #174 in particular.

I will submit this in a commit to this branch. I also fixed the test of the external convergence to remove x0 and y0 in the parametrization and added a test to check that my implementation works and both parametrizations yield the same answer

def test_parametrization():
    cosmo = FlatLambdaCDM(name="cosmo")
    lens = ExternalShear(cosmology=cosmo, z_l=0.5, parametrization="cartesian")
    lens_polar = ExternalShear(cosmology=cosmo, z_l=0.5, parametrization="polar")
    
    gamma_1 = torch.tensor(0.1)
    gamma_2 = torch.tensor(0.2)
    gamma = torch.sqrt(gamma_1**2 + gamma_2**2)
    phi = 0.5 * torch.atan2(gamma_2, gamma_1)
    
    # Check that the conversion yields the same results in the deflection angle
    x = torch.tensor([0.1, 0.2])
    y = torch.tensor([0.2, 0.1])
    z_s = torch.tensor(2.0)
    
    a1, a2 = lens.reduced_deflection_angle(x, y, z_s, gamma_1=gamma_1, gamma_2=gamma_2)
    a1_p, a2_p = lens_polar.reduced_deflection_angle(x, y, z_s, gamma=gamma, phi=phi)
    assert torch.allclose(a1, a1_p)
    assert torch.allclose(a2, a2_p)

[AlexandreAdam]

For me, what remain to be discussed is simply where to place the decorator I introduced. This will be a useful tool when we do this to other classes. We could place it in utils or somewhere like that.

Once this PR is passed, I'll open one that is more general to start implementing different parametrization of other profiles like NFW to answer issue #174

[ConnorStoneAstro]

I strongly disagree with this variant. Now every function has to accept kwargs for both parameterizations, this will make the docs very messy. It's also harder to follow what's happening.

I would rather the variant where we define two deflection angle, potential, and convergence functions then just set the appropriate one at init time.

I appreciate the concern about having diverging implementations, making maintenance more changing. But this is actually fixed by the func update I'm working on, in that case both variants can point to the same underlying func functions to do the core calculation.

I think this needs more discussion.

[AlexandreAdam]

This approach also points to a single implementation which can point to a functional version of the code.

As for the kwargs, well I’d argue this problem is not as bad as it looks. All the kwargs are optional. Only some are expected depending on the choice of parantrization. This implemention works well with the typical workflow for inference and simulations.

I’m not sure what the issue is about the kwargs expanding exactly. The docs have not changed all that much and they make sense to me.

[andreasfilipp]

I do agree with Alex that we should not go with two different implementations, but rather use a single implementation that could handle both. I am less concerned about the maintenance, since I guess most would be copy and paste.
But as Alex did point out and say: Two different parametrizations are not two different objects.
To me it does not make sense to have the same physical object as two different implementations. Besides that being potentially very misleading for Users, it would also be annoying to call shear_polar or shear_cartesian (or later NFW_scaleradius/NFW_concentration) instead of just shear and passing polar/cartesian as a keyword (in case one does not want to use the default option).
I cannot comment on the implementation in the code tho, because I would naively just have done it with if statements, what we do not want I assume

[ConnorStoneAstro]

@mjyb16 please convert this to a Draft PR while we sort out the API/implementation