Commit graph

41 commits

Author SHA1 Message Date
Wassim KABALAN
bf44dfdea9
jaxdecomp proto (#21)
* adding example of distributed solution

* put back old functgion

* update formatting

* add halo exchange and slice pad

* apply formatting

* implement distributed optimized cic_paint

* Use new cic_paint with halo

* Fix seed for distributed normal

* Wrap interpolation function to avoid all gather

* Return normal order frequencies for single GPU

* add example

* format

* add optimised bench script

* times in ms

* add lpt2

* update benchmark and add slurm

* Visualize only final field

* Update scripts/distributed_pm.py

Co-authored-by: Francois Lanusse <EiffL@users.noreply.github.com>

* Adjust pencil type for frequencies

* fix painting issue with slabs

* Shared operation in fourrier space now take inverted sharding axis for
slabs

* add assert to make pyright happy

* adjust test for hpc-plotter

* add PMWD test

* bench

* format

* added github workflow

* fix formatting from main

* Update for jaxDecomp pure JAX

* revert single halo extent change

* update for latest jaxDecomp

* remove fourrier_space in autoshmap

* make normal_field work with single controller

* format

* make distributed pm work in single controller

* merge bench_pm

* update to leapfrog

* add a strict dependency on jaxdecomp

* global mesh no longer needed

* kernels.py no longer uses global mesh

* quick fix in distributed

* pm.py no longer uses global mesh

* painting.py no longer uses global mesh

* update demo script

* quick fix in kernels

* quick fix in distributed

* update demo

* merge hugos LPT2 code

* format

* Small fix

* format

* remove duplicate get_ode_fn

* update visualizer

* update compensate CIC

* By default check_rep is false for shard_map

* remove experimental distributed code

* update PGDCorrection and neural ode to use new fft3d

* jaxDecomp pfft3d promotes to complex automatically

* remove deprecated stuff

* fix painting issue with read_cic

* use jnp interp instead of jc interp

* delete old slurms

* add notebook examples

* apply formatting

* add distributed zeros

* fix code in LPT2

* jit cic_paint

* update notebooks

* apply formating

* get local shape and zeros can be used by users

* add a user facing function to create uniform particle grid

* use jax interp instead of jax_cosmo

* use float64 for enmeshing

* Allow applying weights with relative cic paint

* Weights can be traced

* remove script folder

* update example notebooks

* delete outdated design file

* add readme for tutorials

* update readme

* fix small error

* forgot particles in multi host

* clarifying why cic_paint_dx is slower

* clarifying the halo size dependence on the box size

* ability to choose snapshots number with MultiHost script

* Adding animation notebook

* Put plotting in package

* Add finite difference laplace kernel + powerspec functions from Hugo

Co-authored-by: Hugo Simonfroy <hugo.simonfroy@gmail.com>

* Put plotting utils in package

* By default use absoulute painting with

* update code

* update notebooks

* add tests

* Upgrade setup.py to pyproject

* Format

* format tests

* update test dependencies

* add test workflow

* fix deprecated FftType in jaxpm.kernels

* Add aboucaud comments

* JAX version is 0.4.35 until Diffrax new release

* add numpy explicitly as dependency for tests

* fix install order for tests

* add numpy to be installed

* enforce no build isolation for fastpm

* pip install jaxpm test without build isolation

* bump jaxdecomp version

* revert test workflow

* remove outdated tests

---------

Co-authored-by: EiffL <fr.eiffel@gmail.com>
Co-authored-by: Francois Lanusse <EiffL@users.noreply.github.com>
Co-authored-by: Wassim KABALAN <wassim@apc.in2p3.fr>
Co-authored-by: Hugo Simonfroy <hugo.simonfroy@gmail.com>
2024-12-20 05:44:02 -05:00
Francois Lanusse
a848d43b56
remove deprecated stuff 2024-10-24 16:36:41 -04:00
Wassim KABALAN
735699de77 roll back painting 2024-08-05 19:37:33 +02:00
Hugo Simonfroy
d85867e88d 2lpt, get_ode, invlaplace, docstrings 2024-07-31 00:46:53 +02:00
Francois Lanusse
9cbdf18932
Update jaxpm/pm.py 2024-07-19 10:49:51 -04:00
Francois Lanusse
945dae4f80
Merge branch 'main' into neural_ode 2024-07-19 10:48:09 -04:00
EiffL
a2811c0606 Applying formatting 2024-07-09 14:54:34 -04:00
Francois Lanusse
5f463450d1
Adding option to have weights in the 3d cic paint 2023-05-05 19:00:08 +02:00
EiffL
d9578b1220 removed old files 2022-10-22 07:17:44 -04:00
EiffL
a6096ae04b moved xmap stuff to experimental 2022-10-22 07:17:29 -04:00
EiffL
a2e8017bc8 added things 2022-10-20 21:22:39 -07:00
EiffL
953fc48cc0 update data 2022-10-20 21:03:15 -07:00
EiffL
7f3de80f7a adding distributed ops 2022-10-20 23:05:39 -04:00
denise lanzieri
b26cfcd01a creoss correlation function 2022-06-18 18:23:46 +02:00
denise lanzieri
f185b8b913 few adjustments to PGD correction 2022-06-13 17:17:19 +02:00
denise lanzieri
5b72a73448 neural ode added 2022-06-11 14:28:30 +02:00
Francois Lanusse
101ef80d9c
Merge pull request #13 from DifferentiableUniverseInitiative/u/EiffL/lensing
U/eiff l/lensing
2022-05-18 15:55:40 +02:00
EiffL
1b7c797286 changes definition of lensplanes 2022-05-18 15:40:11 +02:00
Francois Lanusse
e7fa6c9e57
Merge pull request #11 from DifferentiableUniverseInitiative/u/EiffL/lensing
Adds basic utilities for Born lensing
2022-05-18 14:35:06 +02:00
Francois Lanusse
4a584b1d9c
Update jaxpm/painting.py 2022-05-18 10:22:21 +02:00
Denise Lanzieri
5394abbed1
Update jaxpm/pm.py
Co-authored-by: Francois Lanusse <EiffL@users.noreply.github.com>
2022-05-18 09:59:59 +02:00
Denise Lanzieri
782c9dbc6c
Update jaxpm/pm.py
Co-authored-by: Francois Lanusse <EiffL@users.noreply.github.com>
2022-05-18 09:58:46 +02:00
Denise Lanzieri
8266f11369
Update jaxpm/pm.py
Co-authored-by: Francois Lanusse <EiffL@users.noreply.github.com>
2022-05-18 09:58:34 +02:00
Denise Lanzieri
b737631346
Update jaxpm/pm.py
Co-authored-by: Francois Lanusse <EiffL@users.noreply.github.com>
2022-05-18 09:55:31 +02:00
EiffL
0f6fb39966 small fix 2022-05-17 23:49:12 +02:00
EiffL
a52287cd97 change impor 2022-05-17 23:42:57 +02:00
EiffL
88cff99736 adds fix to make code jittablel 2022-05-17 23:37:55 +02:00
EiffL
80eb8cb2a3 minor correction to gaussian smoothing 2022-05-17 23:02:01 +02:00
EiffL
0008f8549b adds utilities for simple lensing 2022-05-17 17:55:06 +02:00
denise lanzieri
d3026f7849 PGD 2022-05-17 15:28:30 +02:00
EiffL
e45bab2ecf adding density plane cutting code 2022-05-17 11:26:38 +02:00
EiffL
e188d5efb6 adding function for doing 2d paintinng 2022-05-17 11:19:56 +02:00
EiffL
d0a15e8c78 adding hamiltonian gnn demo 2022-04-28 00:21:46 +02:00
EiffL
5195a28582 fix minor issue 2022-03-26 03:03:59 +01:00
EiffL
907dc42bdf adding neural network 2022-03-26 02:59:39 +01:00
EiffL
d8a5b702a9 adding cic compensation tools 2022-03-26 00:06:34 +01:00
EiffL
5549568295 fix normalization of init cond 2022-03-25 22:34:13 +01:00
EiffL
b08641d51d Add utility to compute the power spectrum 2022-03-25 21:29:32 +01:00
EiffL
25bde5581e Adds demo and notebooks 2022-02-14 01:59:12 +01:00
EiffL
c9a42686eb adds growth functions from Chirag 2022-02-14 00:37:35 +01:00
EiffL
afdfcfa644 Adds a trivial jaxpm implementation 2022-02-13 21:36:03 +01:00