more tutorials

Author: pmocz

docs/index.rst

@@ -106,6 +106,7 @@ jaxion
     pages/tutorial3
     pages/tutorial4
     pages/tutorial5
+    pages/tutorial6
 
 .. toctree::
     :maxdepth: 1

docs/pages/tutorial1.rst

@@ -25,6 +25,13 @@ It will also often be useful to import `jax.numpy <https://docs.jax.dev/en/lates
 
 to set up initial conditions, which as JAX arrays.
 
+By default, JAX runs in single-precision mode, which is sufficient for many Jaxion simulations.
+But if you'd like to run in double-precision mode, you can set:
+
+.. code-block:: python
+
+    jax.config.update("jax_enable_x64", True)
+
 The basic steps are:
 
 (1) set simulation parameters,
@@ -59,6 +66,7 @@ Let's look at the parameters for this example:
         "time": {
             "start": 0.0,
             "end": 1.0,
+            "safety_factor": 1.0
         },
         "output": {
             "path": f"./checkpoints",
@@ -85,6 +93,8 @@ with a base resolution of 32 grid cells per side,
 and a resolution multiplier of 2 (meaning the effective resolution is 64 grid cells per side).
 
 In the ``"time"`` section, we set the simulation to start at time 0.0 and end at time 1.0 (in code units (kpc / (km/s))).
+The time step in a Jaxion simulation is, by default, the quantum kinetically-limited time step,
+dt = (m_per_hbar / 6.0) * (dx * dx), which can be scaled by the ``safety_factor`` parameter (here set to 1.0).
 
 In the ``"output"`` section, we specify the output path for saving simulation checkpoints,
 the number of checkpoints to save (100), and enable saving.

docs/pages/tutorial3.rst

@@ -1,4 +1,29 @@
-Tutorial 3: Custom Callbacks
-============================
+Tutorial 3: Sponge Boundary Conditions
+======================================
 
-TODO XXX
+This tutorial describes how to add sponge boundary conditions.
+By default, Jaxion uses periodic boundary conditions.
+Sponge boundary conditions can be used to absorb outgoing waves at the domain boundaries.
+
+An example of sponge boundary conditions is provided in `examples/tidal_stripping <https://github.com/JaxionProject/jaxion/tree/main/examples/tidal_stripping>`_
+
+It is created by adding an external potential to the simulation.
+We need to set ``params["physics"]["external_potential"] = True``.
+
+.. code-block:: python
+
+    # V_0 is the depth of the sponge potential ~ G * M_tot / box_size
+    # r_N, r_p, and r_s define the sponge region
+    r_N = 0.5 * box_size
+    r_p = (7 / 8) * r_N
+    r_s = 0.5 * (r_N + r_p)
+    delta = r_N - r_p
+    V_sponge = (
+        -0.5j
+        * V_0
+        * (2 + jnp.tanh((R - r_s) / delta) - jnp.tanh(r_s / delta))
+        * jnp.heaviside(R - r_p, 0.0)
+    )
+    sim.state["V_ext"] = V_sponge
+
+Since the sponge potential is imaginary, it will absorb outgoing waves in the sponge region.

docs/pages/tutorial4.rst

@@ -1,4 +1,39 @@
-Tutorial 4: Custom Fields
-=========================
+Tutorial 4: Custom Callbacks
+============================
 
-TODO XXX
+This tutorial describes how to create custom callbacks in Jaxion simulations.
+The callback function is called at each time step during the simulation.
+It can be used to perform custom calculations or save info on the simulation state at each step.
+
+An example of a custom callback is provided in `examples/black_hole_accretion <https://github.com/JaxionProject/jaxion/tree/main/examples/black_hole_accretion>`_
+where the callback is used to save the mass of a black hole at each time step.
+
+It is simply done by creating new simulation state variables and defining a callback function:
+
+.. code-block:: python
+
+    # add callback to record info about state
+    n_buffer = sim.nt + 1
+    sim.state["tt"] = jnp.full((n_buffer,), jnp.nan)
+    sim.state["m_bh"] = jnp.full((n_buffer,), jnp.nan)
+    sim.state["tt"] = sim.state["tt"].at[0].set(0.0)
+    sim.state["m_bh"] = sim.state["m_bh"].at[0].set(M_bh)
+    sim.callback = callback
+
+    def callback(i, state):
+        # record the black hole mass at end of timestep i
+        state["tt"] = state["tt"].at[i + 1].set(state["t"])
+        state["m_bh"] = state["m_bh"].at[i + 1].set(state["mass"][0])
+        return state
+
+The callback function takes two arguments:
+- `i`: the current time step index
+- `state`: the current simulation state
+The function should return the updated state.
+
+In this example, the black hole mass looks something like this:
+
+.. figure:: ../../examples/black_hole_accretion/callback.png
+  :width: 300px
+  :align: center
+  :alt: black_hole_accretion callback

docs/pages/tutorial5.rst

@@ -1,4 +1,58 @@
-Tutorial 5: Multiple GPUs
+Tutorial 5: Custom Fields
 =========================
 
-TODO XXX
+This tutorial describes how to create custom quantum dark matter fields in Jaxion simulations.
+An example is provided in `examples/two_field <https://github.com/JaxionProject/jaxion/tree/main/examples/two_field>`_
+
+First, ``params["physics"]["quantum"]`` should be set to ``False`` since we will be adding out own fields.
+
+For example, to create a two-field simulation, we can add the simulation states:
+``sim.state["psi1"]`` and ``sim.state["psi2"]``.
+
+Then, we need to define custom kick, drift, and total density functions, and attach them to the simulation object:
+
+.. code-block:: python
+
+    def custom_density(state):
+        return jnp.abs(state["psi1"]) ** 2 + jnp.abs(state["psi2"]) ** 2
+
+    def custom_kick(state, V, dt):
+        state["psi1"] = jnp.exp(-1j * m1_per_hbar * dt * V) * state["psi1"]
+        state["psi2"] = jnp.exp(-1j * m2_per_hbar * dt * V) * state["psi2"]
+
+        return state
+
+    def custom_drift(state, k_sq, dt):
+        psi1_hat = jd.fft.pfft3d(state["psi1"])
+        psi1_hat = jnp.exp(dt * (-1.0j * k_sq / m1_per_hbar / 2.0)) * psi1_hat
+        state["psi1"] = jd.fft.pifft3d(psi1_hat)
+
+        psi2_hat = jd.fft.pfft3d(state["psi2"])
+        psi2_hat = jnp.exp(dt * (-1.0j * k_sq / m2_per_hbar / 2.0)) * psi2_hat
+        state["psi2"] = jd.fft.pifft3d(psi2_hat)
+
+        return state
+
+    sim.custom_density = custom_density
+    sim.custom_kick = custom_kick
+    sim.custom_drift = custom_drift
+
+In this example, the custom density function is used to compute the total density from both fields.
+The custom kick and drift functions are used to update the fields during the simulation.
+
+The simulation can then be run as usual with ``sim.run()``.
+
+It is also possible to add custom plotting for these fields by defining a custom plotting function and attaching it to the simulation object:
+``custom_plot(state, checkpoint_dir, i, params)``.
+
+The Two-Field example evolves two quantum fields with different masses that interact gravitationally, and looks something like this:
+
+.. figure:: ../../examples/two_field/rho1_070.png
+  :width: 300px
+  :align: center
+  :alt: two_field field 1
+
+.. figure:: ../../examples/two_field/rho2_070.png
+  :width: 300px
+  :align: center
+  :alt: two_field field 2

docs/pages/tutorial6.rst

@@ -0,0 +1,68 @@
+Tutorial 6: Multiple GPUs
+=========================
+
+This tutorial describes how to run Jaxion simulations on distributed GPUs.
+An example is provided in `examples/soliton_binary_merger <https://github.com/JaxionProject/jaxion/tree/main/examples/soliton_binary_merger>`_
+
+First, ensure that JAX is initialized for distributed GPU use, by calling:
+
+.. code-block:: python
+
+    import jax
+
+    jax.distributed.initialize()
+
+On some machines, you may need to specify additional parameters to the ``initialize()`` function.
+For most SLURM clusters, you will not need to specify anything.
+
+If you'd like your Python script to print info, it should by guarded by:
+
+.. code-block:: python
+
+    if jax.process_index() == 0:
+        print("Using distributed GPU mode")
+
+to prevent multiple processes from printing simultaneously.
+
+
+Sharding
+--------
+
+We need to set up sharding for the simulation state arrays.
+Sharding splits the arrays across multiple devices for distributed computation.
+This can be done as follows:
+
+.. code-block:: python
+
+    from jax.experimental import mesh_utils
+    from jax.sharding import Mesh, PartitionSpec, NamedSharding
+
+    # Create mesh and sharding for distributed computation
+    n_devices = jax.device_count()
+    devices = mesh_utils.create_device_mesh((1, n_devices))
+    mesh = Mesh(devices, axis_names=("x", "y"))
+    sharding = NamedSharding(mesh, PartitionSpec("x", "y"))
+
+In the example above, we create a virtual device mesh with one row and ``n_devices`` columns.
+Arrays (2D and 3D) are split along the "y" axis across the devices.
+
+When creating the simulation object, we need to pass the sharding to it:
+
+.. code-block:: python
+
+    sim = jaxion.Simulation(params, sharding=sharding)
+
+And that's it! Jaxion will now run the simulation on multiple GPUs.
+
+If you grab arrays from the simulation, such as the grid (``sim.grid``) or spectral grid (``sim.kgrid``),
+these will be sharded arrays too.
+
+
+Slurm Example
+-------------
+
+An example SLURM submission script for running our example on the Flatiron Rusty cluster
+is provided below:
+
+.. literalinclude:: ../../examples/soliton_binary_merger/sbatch_rusty_distributed.sh
+  :language: bash