diff --git a/notebooks/README.md b/notebooks/README.md index 43d9d0b..872fdd4 100644 --- a/notebooks/README.md +++ b/notebooks/README.md @@ -37,3 +37,50 @@ Each notebook includes installation instructions and guidelines for configuring - **SLURM** for job scheduling on clusters (if running multi-host setups) > **Note**: These notebooks are tested on the **Jean Zay** supercomputer and may require configuration changes for different HPC clusters. + +## Caveats + +### Cloud-in-Cell (CIC) Painting (Single Device) + +There is two ways to perform the CIC painting in JAXPM. The first one is to use the `cic_paint` which paints absolute particle positions to the mesh. The second one is to use the `cic_paint_dx` which paints relative particle positions to the mesh (using uniform particles). The absolute version is faster at the cost of more memory usage. + +inorder to use relative painting you need to : + + - Set the `particles` argument in `lpt` function from `jaxpm.pm` to `None` + - Set `paint_absolute_pos` to `False` in `make_ode_fn` or `make_diffrax_ode` function from `jaxpm.pm` (it is True by default) + +Otherwise you set `particles` to the starting particles of your choice and leave `paint_absolute_pos` to `True` (default value). + +### Cloud-in-Cell (CIC) Painting (Multi Device) + +Both `cic_paint` and `cic_paint_dx` functions are available in multi-device mode. + +You need to set the arguments `sharding` and `halo_size` which is explained in the notebook [03-MultiGPU_PM_Halo.ipynb](03-MultiGPU_PM_Halo.ipynb). + +One thing to note that `cic_paint` is not as accurate as `cic_paint_dx` in multi-device mode and therefor is not recommended. + +Using relative painting in multi-device mode is just like in single device mode.\ +You need to set the `particles` argument in `lpt` function from `jaxpm.pm` to `None` and set `paint_absolute_pos` to `False` + +### Distributed PM + +To run a distributed PM follow the examples in notebooks [03](03-MultiGPU_PM_Halo.ipynb) and [05](05-MultiHost_PM.ipynb) for multi-host. + +In short you need to set the arguments `sharding` and `halo_size` in `lpt` , `linear_field` the `make_ode` functions and `pm_forces` if you use it. + +Missmatching the shardings will give you errors and unexpected results. + +You can also use `normal_field` and `uniform_particles` from `jaxpm.pm.distributed` to create the fields and particles with a sharding. + +### Choosing the right pdims + +pdims are processor dimensions.\ +Explained more in the jaxdecomp paper [here](https://github.com/DifferentiableUniverseInitiative/jaxDecomp). + +For 8 devices there are three decompositions that are possible: +- (1 , 8) +- (2 , 4) , (4 , 2) +- (8 , 1) + +(1 , X) should be the fastest (2 , X) or (X , 2) is more accurate but slightly slower.\ +and (X , 1) is giving the least accurate results for some reason so it is not recommended.