Constant noise assumption in sampling script

If I read this correctly, it looks like we are assuming constant noise on gamma throughout the field in the likelihood term:

jax-lensing/scripts/sample_maps.py

Line 213 in 13a4b8d

    
           sl = likelihood_score(x, sigma, masked_true_shear, sigma_mask).reshape(-1, 360*360,2)

But we know we have some variations, and that not accounting for it can throw off a model like DeepMass. If I read this part of the code correctly:

jax-lensing/scripts/train_deepmass.py

Line 101 in 13a4b8d

g1 = mask * (g1 + std1*jax.random.normal(key1, g1.shape))

DeepMass uses the non-constant noise variance, but not our DSM sampling? That would be inconsistent and we should fix that.

@b-remy can you let me know if I missed something?

Right, with the last PR, the likelihood will be always using the diagonal covariance matrix. Thanks for spotting this @EiffL

With PR #50, we can train DeepMass assuming a noise realisation from the COMOS catalog, by randomly rotating the ellipticities. This is a better assumption than a constant sigma over all pixels.

This changes a bit the results for the DeepMass estimation as you can see in the plot bellow. Visually, I can identify structures in the target that are present in the COSMOS noise assumption map and not in the constant assumption map, but also the opposite....

and slightly improves the metrics (RMSE and r are computed in the non-masked region):

Constant noise assumption

RMSE: 2.20e-02
r: 0.67364645

COSMOS noise assumption

RMSE: 2.17e-02
r: 0.68455

Finally, I think we can keep the COSMOS noise assumption for our DeepMass model.