Port Weak Lensing code to Mesh TensorFlow

[EiffL]

For now, the raytracing code we have developed with @dlanzieri is based on the non-distributed version of FlowPM. I'm opening this issue to document the porting of the necessary pieces to Mesh Tensorflow.

Here is the list of things I think we need:

• : 2D cic painting for pojecting snapshots

Things that are optional at this stage but would be nice to have

• : 2D interpolation on distributed tensors
• : Differentiable cosmology in the mesh nbody code
• : Differentiable initial conditions

[EiffL]

Ok, I've started working in branch mesh_lensing and I've got the distributed 2d cic painting working :-)

I am building a lensing script here: https://github.com/DifferentiableUniverseInitiative/flowpm/blob/mesh_lensing/scripts/pyramid_lensing.py

so, this all looks very feasible.... if all goes well, we'll have lensing capabilities by end of saturday

[EiffL]

almost there:

with the following settiings:
nc=512
box_size=100
field_size=5x5deg
field_npix=1024

[dlanzieri]

It's great! :O

[dlanzieri]

@EiffL By the way, I have some questions related to the differentiable initial conditions!