Commit graph

95 commits

Author SHA1 Message Date
Wassim KABALAN
2ea05a1cd6 Update for jaxDecomp pure JAX 2024-10-18 20:45:08 +02:00
Wassim KABALAN
831291c1f9 bench 2024-08-02 23:39:09 +02:00
Wassim KABALAN
9af4659c81 add PMWD test 2024-08-02 21:22:56 +02:00
Wassim KABALAN
aebc3e72c0 adjust test for hpc-plotter 2024-08-02 21:22:45 +02:00
Wassim KABALAN
ccbfee3615 add assert to make pyright happy 2024-08-02 21:21:59 +02:00
Wassim KABALAN
75604d2396 Shared operation in fourrier space now take inverted sharding axis for
slabs
2024-08-02 21:21:06 +02:00
Wassim KABALAN
8c5bd76c33 fix painting issue with slabs 2024-08-02 21:20:15 +02:00
Wassim KABALAN
f25eb7d465 Adjust pencil type for frequencies 2024-07-28 13:44:03 +02:00
Wassim KABALAN
1f2035176f
Update scripts/distributed_pm.py
Co-authored-by: Francois Lanusse <EiffL@users.noreply.github.com>
2024-07-19 19:25:26 +02:00
Wassim KABALAN
5b7f595c32 Visualize only final field 2024-07-18 17:05:22 +02:00
Wassim KABALAN
0216837033 update benchmark and add slurm 2024-07-18 17:04:52 +02:00
Wassim KABALAN
ed8cf8e532 add lpt2 2024-07-18 17:04:05 +02:00
Wassim KABALAN
1f6b9c3217 times in ms 2024-07-18 13:23:53 +02:00
Wassim KABALAN
5f6d42eaeb add optimised bench script 2024-07-18 13:16:37 +02:00
Wassim KABALAN
4f508b7cb6 format 2024-07-18 13:16:14 +02:00
Wassim KABALAN
82be56836a add example 2024-07-18 12:45:56 +02:00
Wassim KABALAN
abde5439f6 Return normal order frequencies for single GPU 2024-07-18 12:44:18 +02:00
Wassim KABALAN
c81d4d2336 Wrap interpolation function to avoid all gather 2024-07-18 12:43:52 +02:00
Wassim KABALAN
7f48cfa8af Fix seed for distributed normal 2024-07-18 12:43:24 +02:00
Wassim KABALAN
7501b5bc6d Use new cic_paint with halo 2024-07-18 12:42:07 +02:00
Wassim KABALAN
5775a37550 implement distributed optimized cic_paint 2024-07-18 12:39:15 +02:00
Wassim KABALAN
e62cd84cbd apply formatting 2024-07-18 12:38:42 +02:00
Wassim KABALAN
ac86468c7c add halo exchange and slice pad 2024-07-10 00:25:52 +02:00
EiffL
319942a6bc update formatting 2024-07-09 18:02:57 -04:00
EiffL
6408aff1de put back old functgion 2024-07-09 18:00:35 -04:00
EiffL
a742065ffd adding example of distributed solution 2024-07-09 17:45:28 -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
Francois Lanusse
2db7c43ab8
Merge pull request #16 from DifferentiableUniverseInitiative/u/EiffL/distributed
U/eiff l/distributed
2022-10-22 07:23:51 -04: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
EiffL
2da21d8b5e Added prototype distributed LPT 2022-10-19 17:37:38 -07:00
EiffL
61578217fe Added example code to compute distributed fft 2022-10-17 22:07:12 -07: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
5c4f75a102 fixes the lensing demo 2022-05-18 15:43:45 +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
Francois Lanusse
9967dffd82
Merge pull request #12 from DifferentiableUniverseInitiative/PGD
PGD
2022-05-18 10:21:01 +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