guilhem_lavaux/JaxPM
0
0
Fork 1
mirror of https://github.com/DifferentiableUniverseInitiative/JaxPM.git synced 2025-04-11 21:50:55 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

notebook update

Browse source
This commit is contained in:
Wassim Kabalan 2025-02-28 09:59:04 +01:00
parent b43cb373a0
commit f3b8f4160e
1 changed files with 2 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
notebooks/03-MultiGPU_PM_Halo.ipynb
View file

@ -177,7 +177,7 @@
"\n", "\n",
" # Evolve the simulation forward\n", " # Evolve the simulation forward\n",
" ode_fn = ODETerm(\n", " ode_fn = ODETerm(\n",
" make_diffrax_ode(cosmo, mesh_shape, paint_absolute_pos=False , sharding=sharding))\n", " make_diffrax_ode(mesh_shape, paint_absolute_pos=False , sharding=sharding))\n",
" solver = Tsit5()\n", " solver = Tsit5()\n",
"\n", "\n",
" stepsize_controller = PIDController(rtol=1e-6 , atol=1e-6)\n", " stepsize_controller = PIDController(rtol=1e-6 , atol=1e-6)\n",
@ -407,7 +407,7 @@
"\n", "\n",
" # Evolve the simulation forward\n", " # Evolve the simulation forward\n",
" ode_fn = ODETerm(\n", " ode_fn = ODETerm(\n",
" make_diffrax_ode(cosmo, mesh_shape, paint_absolute_pos=False , sharding=sharding , halo_size=halo_size))\n", " make_diffrax_ode(mesh_shape, paint_absolute_pos=False , sharding=sharding , halo_size=halo_size))\n",
" solver = Tsit5()\n", " solver = Tsit5()\n",
"\n", "\n",
" stepsize_controller = PIDController(rtol=1e-6 , atol=1e-6)\n", " stepsize_controller = PIDController(rtol=1e-6 , atol=1e-6)\n",