Add a deepmass implementation using the same architecture as our denoiser

[EiffL]

This issue is to try to reimplement DeepMass in this framework, to allow for comparison with the previously published method.

This would also allow us to compare mean posteriors with DeepMass predictions, and possibly also check @NiallJeffrey's moments networks.

TODO:

• [x] : Decide if we want to keep a Wiener filter step (I'd say meh)
• [x] : Add a script to train a DeepMass model (can use same UNet model as DSM)

[EiffL]

@NiallJeffrey would you like to take this on? Otherwise we are happy to do it :-)

[EiffL]

@NiallJeffrey , @b-remy is making very good progress on using @liuxx479 's new KappaTNG simulations, so I think we could be able to get our main results for the paper within a couple of weeks at this point.

So, we wanted to run a comparison to DeepMass, and potentially your moments networks. Would you be interested in looking into this? Or should we go ahead and make a deepmass-like model on our own?

It might be slightly still subject to change, but we have a U-Net model that we use as our score network, I think it can be used as-is for deepmass. Only question is whether we want to include the Wiener-filter step or not....

[NiallJeffrey]

Did you make a DeepMass model of you're own? If not, I can adapt the existing one. I wanted to chat with Ben regarding Moment Networks for mass mapping, but he's taking holiday at the moment, so I haven't been able to

[EiffL]

I've added a simple DeepMass-like implementation in branch deepmass. What I do is reuse the exact same UNet as for our score learning approach, and just change the loss function, to train it so that given a masked, noisy, KS map, it tries to return the true underlying kappa map, under an l2 loss. Code is here: https://github.com/CosmoStat/jax-lensing/blob/deepmass/scripts/train_deepmass.py But not super tested yet.

So, this differs from the original deepmass on that I don't actually try to a Wiener filter step.... @NiallJeffrey what do you think? also, is the wiener filtering doable on the fly? Or did you do it offline in the deepmass paper?

[NiallJeffrey]

In the DeepMass paper I did it in advance, but one could do it on the fly. I'll look at the new implementation

[EiffL]

@b-remy I can't find a mask file for our new dataset, do you have an equivalent of https://github.com/CosmoStat/jax-lensing/blob/master/data/cosmos_mask_0.4arcmin.fits ?

[b-remy]

Indeed, there was no mask up to date. There is one here now: https://github.com/CosmoStat/jax-lensing/tree/comos-data/data/COSMOS

[EiffL]

I think this is done?

[EiffL]

But @b-remy could you push the latest version of your code to the repo?

[b-remy]

Hi, I just pushed the updated code for DeepMass. The U-net architecture is the same we used to train our denoiser, on which I removed the division by the noise level (which is zero here).

Training script can be found here.

and notebook result here.