Multi-plane lensing

[CKrawczyk]

A draft PR with the start of the MP-lensing code. I still need to add some doc-strings and check that each class works as expected.

[coveralls]

Pull Request Test Coverage Report for Build 10490896144

Warning: This coverage report may be inaccurate.

This pull request's base commit is no longer the HEAD commit of its target branch. This means it includes changes from outside the original pull request, including, potentially, unrelated coverage changes.

• For more information on this, see Tracking coverage changes with pull request builds.
• To avoid this issue with future PRs, see these Recommended CI Configurations.
• For a quick fix, rebase this PR at GitHub. Your next report should be accurate.

Details

• 0 of 249 (0.0%) changed or added relevant lines in 3 files are covered.
• 104 unchanged lines in 2 files lost coverage.
• Overall coverage decreased (-0.8%) to 23.212%

---

Changes Missing Coverage	Covered Lines	Changed/Added Lines	%
herculens/LightModel/light_model_multiplane.py	0	56	0.0%
herculens/MassModel/mass_model_multiplane.py	0	64	0.0%
herculens/LensImage/lens_image_multiplane.py	0	129	0.0%

Files with Coverage Reduction	New Missed Lines	%
herculens/Util/param_util.py	26	27.27%
herculens/Analysis/plot.py	78	0.0%

Totals
Change from base Build 10171361525:	-0.8%
Covered Lines:	1603
Relevant Lines:	6906

---

💛 - Coveralls

[aymgal]

@CKrawczyk I think the multiplication by the arc masks should be performed after convolution, to avoid introducing artifacts in the model at the edge of the masks.

[CKrawczyk]

We did it this way because the convolution happens once on the final image rather than each slice independently. To get around the artifact a "soft mask" can be used (gaussian blur the mask so it is a smooth gradient down to zero at the edge).

[CKrawczyk]

Just added some methods to calculate the area distortion matrix for the MP mass model using the auto grad of the ray_shooting function directly. This removes the need for the bookkeeping that would be needed to calculate this value directly from the Hessians of each mass plane. It also vectorizes quite nicely so that each mass plane's matrix is calculated in the same pass.

Example output using a simple model with two SIS masses at different redshift:

[CKrawczyk]

For completion, I added a "direct" calculation method for the distortion matrix that uses the hessian method for each mass plane directly rather than the auto grad. It might be useful if the alpha's can't be auto-graded for any reason but the individual hessian functions work as expected.

[aymgal]

Thanks for the push @CKrawczyk ! These features will be very handy.
Let me know if you require something from my side for the merge at some point.

[CKrawczyk]

I think documentation is the only thing missing at the moment. I will try to take some time today to get that mostly sorted. I noticed that you use several different doc string formats throughout the code base, what one should I aim to use?

[aymgal]

@CKrawczyk The one that you just pushed is the right one, it's longer but much easier to read than other formats 👍

[CKrawczyk]

This is ready for review and testing. Every public-facing class method should have a doc string.

[CKrawczyk]

For testing above I was using this model setup:

mass_model = MPMassModel([
    ['SIS'],
    ['SIS']
])

light_model=MPLightModel([
    [],
    [],
    ['GAUSSIAN']
], [{}, {}, {}])

kwargs_numerics = {'supersampling_factor': 5}
lens_image = MPLensImage(
    pixel_grid, psf, noise,
    mass_model, light_model,
    kwargs_numerics=kwargs_numerics
)

kwargs_mass = [
    [{'theta_E': 1.0, 'center_x': 0.0, 'center_y': 0.0}],
    [{'theta_E': 0.9, 'center_x': 0.9, 'center_y': 0.0}]
]

kwargs_light = [
    [],
    [],
    [{'amp': 2, 'sigma': 0.05, 'center_x': 0.9, 'center_y': 0.05}]
]

eta_flat=jnp.array([1.05])

image = lens_image.model(
    eta_flat=eta_flat,
    kwargs_mass=kwargs_mass,
    kwargs_light=kwargs_light,
)

[CKrawczyk]

I have done a local test with an example notebook (the one used to make the plot above). Just found one small bug that I fixed with how the A matrix reshapes after calculation, I assumed 2D inputs, but now it will work as expected regardless of the input shape. The output will now always be (number light planes, *input shape, 2, 2) regardless of how many terms are in input shape.

[CKrawczyk]

For clarification, the $$\eta$$ matrix that the MP class uses is:

$$ \begin{equation} \begin{pmatrix} 0 & 1 & \eta_{02} & \eta_{03} & \eta_{04} & \dots \\ 0 & 0 & 1 & \eta_{13} & \eta_{14} & \dots \\ 0 & 0 & 0 & 1 & \eta_{24} & \dots \\ 0 & 0 & 0 & 0 & 1 & \dots \\ \vdots & \vdots & \vdots & \vdots & \vdots & \ddots \\ 0 & 0 & 0 & 0 & 0 & \dots \\ \end{pmatrix} \end{equation} $$

where

$$ \begin{equation} \eta_{ij} = \frac{D_{i,j} * D_{i+1}}{D_{j} * D_{i,i+1}} \end{equation} $$

Where D are the angular diameter distances, and plane 0 is the forward-most mass plane. Also note that by definition $\eta_{i, i+1} = 1$.

As the eta matrix is upper triangular we only need to pass in the values that are not known to be zero or one. These are passed in as a flat list that will fill the values row-wise into the matrix. So for a 4-plane lens eta_flat is:

$$ \begin{equation} \eta_{flat} = [\eta_{02}, \eta_{03}, \eta_{04}, \eta_{13}, \eta_{14}, \eta_{24}] \end{equation} $$

The MP lensing equation is given by:

$$ \begin{equation} \vec{x}_j = \vec{x}_0 - \sum _{i=0}^{j-1} \vec{\alpha} _{i}(\vec{x} _i) \eta _{ij} \end{equation} $$

By expanding this out and padding with some zeros we can see these values can be calculated by iterating over the rows of the eta matrix:

$$ \begin{eqnarray} \vec{x} _0 &=& \vec{x} _0 &-& \vec{\alpha} _{0}(\vec{x} _0) * 0 &-& \vec{\alpha} _{1}(\vec{x} _1) * 0 &-& \vec{\alpha} _{2}(\vec{x} _2) * 0 &-& \vec{\alpha} _{3}(\vec{x} _3) * 0 \\ \vec{x} _1 &=& \vec{x} _0 &-& \vec{\alpha} _{0}(\vec{x} _0) * 1 &-& \vec{\alpha} _{1}(\vec{x} _1) * 0 &-& \vec{\alpha} _{2}(\vec{x} _2) * 0 &-& \vec{\alpha} _{3}(\vec{x} _3) * 0 \\ \vec{x} _2 &=& \vec{x} _0 &-& \vec{\alpha} _{0}(\vec{x} _0) * \eta _{02} &-& \vec{\alpha} _{1}(\vec{x} _1) * 1 &-& \vec{\alpha} _{2}(\vec{x} _2) * 0 &-& \vec{\alpha} _{3}(\vec{x} _3) * 0 \\ \vec{x} _3 &=& \vec{x} _0 &-& \vec{\alpha} _{0}(\vec{x} _0) * \eta _{03} &-& \vec{\alpha} _{1}(\vec{x} _1) * \eta _{13} &-& \vec{\alpha} _{2}(\vec{x} _2) * 1 &-& \vec{\alpha} _{3}(\vec{x} _3) * 0 \\ \vec{x} _4 &=& \vec{x} _0 &-& \vec{\alpha} _{0}(\vec{x} _0) * \eta _{04} &-& \vec{\alpha} _{1}(\vec{x} _1) * \eta _{14} &-& \vec{\alpha} _{2}(\vec{x} _2) * \eta _{24} &-& \vec{\alpha} _{3}(\vec{x} _3) * 1 \\ \end{eqnarray} $$

The rows of eta are used to calculate each "column" of the above equation array from left to right so the $\alpha$ values only need to be calculated once on each plane (at the previous iteration's result).

The row of zeros at the bottom of the matrix is skipped in the calculation (as it doesn't change anything), it is only kept around to make the matrix square. The first column of zeros is a statement that the observed image stays unchanged by lensing, it is kept around to make the iteration loop easier to write.

[aymgal]

Thank you very much @CKrawczyk for the clarification and equation 👌 As soon as I have some time I'll make sure to merge this PR and then work on the integration with the other existing single-plane classes.