Commit graph

121 commits

Author SHA1 Message Date
Wassim Kabalan
9b8e42aa00 updates for FullField 2025-02-05 14:26:28 +01:00
Wassim Kabalan
34bf577b1b small fix 2025-01-21 12:08:47 +01:00
Wassim Kabalan
d4f72c8244 format 2025-01-21 11:21:35 +01:00
Wassim Kabalan
6aa2ccf7b5 don't send static cosmo into ODETerm 2025-01-21 11:18:00 +01:00
Wassim Kabalan
6ab26ea1ec remove extra text 2025-01-21 10:15:11 +01:00
Wassim Kabalan
03610602a8 bump minimum jd version 2025-01-20 23:08:12 +01:00
Wassim Kabalan
a8f11a75ea format 2025-01-20 22:47:07 +01:00
Wassim Kabalan
1f5c619531 format 2025-01-20 22:41:19 +01:00
Wassim Kabalan
20fe25c324 update tests and add test for FWD REV gradient 2025-01-20 22:40:28 +01:00
Wassim Kabalan
151fa09247 power spec should accept pytrees 2025-01-20 22:40:01 +01:00
Wassim Kabalan
38f6599974 painting now accepts pytrees 2025-01-20 22:39:40 +01:00
Wassim Kabalan
9e203b5680 pm now accept pytrees 2025-01-18 01:13:51 +01:00
Wassim Kabalan
204a9526ec kernels now accept pytrees 2025-01-18 01:13:39 +01:00
Wassim Kabalan
f5755b4b5d update painting functions to accept pytrees 2025-01-18 01:13:24 +01:00
Wassim Kabalan
7c3577ea71 update distributed ops for new jd 2025-01-18 01:12:53 +01:00
Wassim KABALAN
cb2a7ab17f
Fix pfft gradients (#34)
* update jaxdecomp version and test gradients

* Prepare for DTO tests

* format
2024-12-22 12:47:42 -05:00
Francois Lanusse
d81a2529e7 minor typo fix 2024-12-21 15:28:20 -05:00
Francois Lanusse
15f2fb1ee6 adding notice 2024-12-21 15:26:53 -05:00
Francois Lanusse
ae0f439ae4 fixing formatting of notebook 2024-12-21 13:14:42 -05:00
Francois Lanusse
ea9fbf6aa8
Update README.md 2024-12-21 13:13:37 -05:00
Francois Lanusse
ad16a0659a Created using Colab 2024-12-21 13:10:15 -05:00
Francois Lanusse
f245a1f685
Pypi upload compatible version (#33)
* moving test dependencies separately

* adding manifest to remove unecessary files

* updating name of project

* Fixing formatting

* Adding badge for pypi version

* Adding very simple install instructions
2024-12-21 11:47:13 -05:00
Francois Lanusse
160b86eb71
Create python-publish.yml 2024-12-21 11:38:24 -05:00
Francois Lanusse
f14f0fe68e
Changing description (#32) 2024-12-21 10:45:12 -05:00
Francois Lanusse
70ab9f1931 remove deprecated stuff
Former-commit-id: da6fdea0556407c968dd31d44af832c997f2645c
2024-12-21 09:58:42 -05:00
Francois Lanusse
bc6e57532d remove massive file
Former-commit-id: 69f49ba2ed23c16bcea70024ea3bcff4d71b8a5b
2024-12-21 09:58:18 -05:00
Wassim KABALAN
4b4450d7d3 Remove huge notebook from history (#30)
* 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>

* Update README.md

* Deleting Animating notebook

* Update pyproject.toml

---------

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>
Former-commit-id: 960cbf28ebcaad2ef0624c92e8f7f0729b75dceb
2024-12-20 08:43:43 -05:00
Francois Lanusse
dd161f7364 Adding version (#29)
* adding version

* fixing version file

Former-commit-id: ba1206bce8a310f6ed06040c2e23b69d0ea502aa
2024-12-20 07:45:29 -05:00
Francois Lanusse
832f8cb71c Update README.md
Former-commit-id: fd6801c3e91d99f529f1697cb947a0d6e5d39869
2024-12-20 05:56:16 -05:00
Wassim KABALAN
df8602b318 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>
Former-commit-id: 8c2e823d4669eac712089bf7f85ffb7912e8232d
2024-12-20 05:44:02 -05:00
Wassim KABALAN
a0a79277e5 Merge pull request #28 from DifferentiableUniverseInitiative/all-contributors/add-aboucaud
docs: add aboucaud as a contributor for review
2024-12-07 20:55:09 +01:00
allcontributors[bot]
d2c02fa991 docs: update .all-contributorsrc [skip ci] 2024-12-07 19:54:33 +00:00
allcontributors[bot]
5f0e75e95c docs: update README.md [skip ci] 2024-12-07 19:54:32 +00:00
Francois Lanusse
b21d0cd576 Merge pull request #27 from DifferentiableUniverseInitiative/all-contributors/add-hsimonfroy
docs: add hsimonfroy as a contributor for code
2024-12-07 14:50:44 -05:00
allcontributors[bot]
03b8a7f49b docs: update .all-contributorsrc [skip ci] 2024-12-07 19:50:10 +00:00
allcontributors[bot]
1e1cf6aaf2 docs: update README.md [skip ci] 2024-12-07 19:50:09 +00:00
Wassim KABALAN
3e8f9d03e8 Merge pull request #26 from DifferentiableUniverseInitiative/all-contributors/add-ASKabalan
docs: add ASKabalan as a contributor for code, infra, and review
2024-12-06 22:42:40 +01:00
allcontributors[bot]
22731ac628 docs: update .all-contributorsrc [skip ci] 2024-12-06 18:57:02 +00:00
allcontributors[bot]
a1ccd0de3a docs: update README.md [skip ci] 2024-12-06 18:57:01 +00:00
Francois Lanusse
f57e32af7f remove deprecated stuff 2024-10-24 16:36:41 -04:00
Wassim KABALAN
ef7a7ef5c9 Merge pull request #23 from hsimonfroy/main
2lpt
2024-08-05 19:38:35 +02:00
Wassim KABALAN
30060e82ea roll back painting 2024-08-05 19:37:33 +02:00
Hugo Simonfroy
9b21eed3b5 2lpt, get_ode, invlaplace, docstrings 2024-07-31 00:46:53 +02:00
Francois Lanusse
b0b793e766 Merge pull request #22 from DifferentiableUniverseInitiative/all-contributors/add-dlanzieri
docs: add dlanzieri as a contributor for code
2024-07-19 10:52:37 -04:00
allcontributors[bot]
0d3bdda3ab docs: update .all-contributorsrc [skip ci] 2024-07-19 14:51:26 +00:00
allcontributors[bot]
ad2ce8d825 docs: update README.md [skip ci] 2024-07-19 14:51:25 +00:00
Francois Lanusse
0f81a89d57 Merge pull request #14 from DifferentiableUniverseInitiative/neural_ode
Fourier-space Neural Network  scheme
2024-07-19 10:50:36 -04:00
Francois Lanusse
b949827e92 Update jaxpm/pm.py 2024-07-19 10:49:51 -04:00
Francois Lanusse
9a279d2d6c Merge branch 'main' into neural_ode 2024-07-19 10:48:09 -04:00
EiffL
f28442bb48 Applying formatting 2024-07-09 14:54:34 -04:00