add rayleigh taylor example

Author: pmocz

adirondax/gravity.py

@@ -9,8 +9,30 @@ def calculate_gravitational_potential(rho, k_sq, G, rho_bar):
     return V
 
 
-def get_acceleration(V, kx, ky):
-    V_hat = jnp.fft.fftn(V)
-    ax = -jnp.real(jnp.fft.ifftn(1.0j * kx * V_hat))
-    ay = -jnp.real(jnp.fft.ifftn(1.0j * ky * V_hat))
+def get_acceleration(V, kx, ky, dx, dy, bc_x_is_reflective, bc_y_is_reflective):
+    if not bc_x_is_reflective or not bc_y_is_reflective:
+        V_hat = jnp.fft.fftn(V)
+    if bc_x_is_reflective:
+        # 2nd order finite difference
+        ax = -(jnp.roll(V, -1, axis=0) - jnp.roll(V, 1, axis=0)) / (2.0 * dx)
+        # one-sided 2nd order difference at boundary
+        ax = ax.at[0, :].set(-(-3.0 * V[0, :] + 4.0 * V[1, :] - V[2, :]) / (2.0 * dx))
+        ax = ax.at[-1, :].set(
+            -(3.0 * V[-1, :] - 4.0 * V[-2, :] + V[-3, :]) / (2.0 * dx)
+        )
+    else:
+        # periodic
+        ax = -jnp.real(jnp.fft.ifftn(1.0j * kx * V_hat))
+    if bc_y_is_reflective:
+        # 2nd order finite difference
+        ay = -(jnp.roll(V, -1, axis=1) - jnp.roll(V, 1, axis=1)) / (2.0 * dy)
+        # one-sided 2nd order difference at boundary
+        ay = ay.at[:, 0].set(-(-3.0 * V[:, 0] + 4.0 * V[:, 1] - V[:, 2]) / (2.0 * dy))
+        ay = ay.at[:, -1].set(
+            -(3.0 * V[:, -1] - 4.0 * V[:, -2] + V[:, -3]) / (2.0 * dy)
+        )
+    else:
+        # periodic
+        ay = -jnp.real(jnp.fft.ifftn(1.0j * ky * V_hat))
+
     return ax, ay

adirondax/hydro/euler2d.py

@@ -100,11 +100,79 @@ def hydro_euler2d_timestep(rho, vx, vy, P, gamma, dx, dy):
     return dt
 
 
+def add_ghost_cells(rho, vx, vy, P, axis):
+    """Add ghost cells for reflective boundary conditions along given axis"""
+
+    if axis == 0:
+        # x-axis
+        rho_new = jnp.concatenate((rho[0:1, :], rho, rho[-1:, :]), axis=0)
+        vx_new = jnp.concatenate((-vx[0:1, :], vx, -vx[-1:, :]), axis=0)
+        vy_new = jnp.concatenate((vy[0:1, :], vy, vy[-1:, :]), axis=0)
+        P_new = jnp.concatenate((P[0:1, :], P, P[-1:, :]), axis=0)
+    elif axis == 1:
+        # y-axis
+        rho_new = jnp.concatenate((rho[:, 0:1], rho, rho[:, -1:]), axis=1)
+        vx_new = jnp.concatenate((vx[:, 0:1], vx, vx[:, -1:]), axis=1)
+        vy_new = jnp.concatenate((-vy[:, 0:1], vy, -vy[:, -1:]), axis=1)
+        P_new = jnp.concatenate((P[:, 0:1], P, P[:, -1:]), axis=1)
+
+    return rho_new, vx_new, vy_new, P_new
+
+
+def remove_ghost_cells(Mass, Momx, Momy, Energy, axis):
+    """Remove ghost cells for reflective boundary conditions along given axis"""
+
+    if axis == 0:
+        # x-axis
+        Mass_new = Mass[1:-1, :]
+        Momx_new = Momx[1:-1, :]
+        Momy_new = Momy[1:-1, :]
+        Energy_new = Energy[1:-1, :]
+    elif axis == 1:
+        # y-axis
+        Mass_new = Mass[:, 1:-1]
+        Momx_new = Momx[:, 1:-1]
+        Momy_new = Momy[:, 1:-1]
+        Energy_new = Energy[:, 1:-1]
+
+    return Mass_new, Momx_new, Momy_new, Energy_new
+
+
+def set_ghost_gradients(f_dx, axis):
+    """Set gradients in ghost cells to (-1) x  value of the first interior cell (f_dx already has ghost cells)"""
+
+    if axis == 0:
+        f_dx = f_dx.at[0, :].set(-f_dx[1, :])
+        f_dx = f_dx.at[-1, :].set(-f_dx[-2, :])
+    elif axis == 1:
+        f_dx = f_dx.at[:, 0].set(-f_dx[:, 1])
+        f_dx = f_dx.at[:, -1].set(-f_dx[:, -2])
+
+    return f_dx
+
+
 def hydro_euler2d_fluxes(
-    rho, vx, vy, P, gamma, dx, dy, dt, riemann_solver_type, use_slope_limiting
+    rho,
+    vx,
+    vy,
+    P,
+    gamma,
+    dx,
+    dy,
+    dt,
+    riemann_solver_type,
+    use_slope_limiting,
+    bc_x_is_reflective,
+    bc_y_is_reflective,
 ):
     """Take a simulation timestep"""
 
+    # Add Ghost Cells (if needed)
+    if bc_x_is_reflective:
+        rho, vx, vy, P = add_ghost_cells(rho, vx, vy, P, axis=0)
+    if bc_y_is_reflective:
+        rho, vx, vy, P = add_ghost_cells(rho, vx, vy, P, axis=1)
+
     # get Conserved variables
     Mass, Momx, Momy, Energy = get_conserved(rho, vx, vy, P, gamma, dx * dy)
 
@@ -121,6 +189,18 @@ def hydro_euler2d_fluxes(
         vy_dx, vy_dy = slope_limit(vy, vy_dx, vy_dy, dx, dy)
         P_dx, P_dy = slope_limit(P, P_dx, P_dy, dx, dy)
 
+    # set ghost cell gradients
+    if bc_x_is_reflective:
+        rho_dx = set_ghost_gradients(rho_dx, axis=0)
+        vx_dx = set_ghost_gradients(vx_dx, axis=0)
+        vy_dx = set_ghost_gradients(vy_dx, axis=0)
+        P_dx = set_ghost_gradients(P_dx, axis=0)
+    if bc_y_is_reflective:
+        rho_dy = set_ghost_gradients(rho_dy, axis=1)
+        vx_dy = set_ghost_gradients(vx_dy, axis=1)
+        vy_dy = set_ghost_gradients(vy_dy, axis=1)
+        P_dy = set_ghost_gradients(P_dy, axis=1)
+
     # extrapolate half-step in time
     rho_prime = rho - 0.5 * dt * (vx * rho_dx + rho * vx_dx + vy * rho_dy + rho * vy_dy)
     vx_prime = vx - 0.5 * dt * (vx * vx_dx + vy * vx_dy + (1.0 / rho) * P_dx)
@@ -167,15 +247,23 @@ def hydro_euler2d_fluxes(
     Momy = apply_fluxes(Momy, flux_Momy_X, flux_Momy_Y, dx, dy, dt)
     Energy = apply_fluxes(Energy, flux_Energy_X, flux_Energy_Y, dx, dy, dt)
 
+    # remove ghost cells
+    if bc_x_is_reflective:
+        Mass, Momx, Momy, Energy = remove_ghost_cells(Mass, Momx, Momy, Energy, axis=0)
+    if bc_y_is_reflective:
+        Mass, Momx, Momy, Energy = remove_ghost_cells(Mass, Momx, Momy, Energy, axis=1)
+
     rho, vx, vy, P = get_primitive(Mass, Momx, Momy, Energy, gamma, dx * dy)
 
     return rho, vx, vy, P
 
 
-def hydro_euler2d_accelerate(rho, vx, vy, P, ax, ay, gamma, dt):
-    e = P / ((gamma - 1.0) * rho)
-    e_new = e + dt * (ax * vx + ay * vy)
-    P_new = (gamma - 1.0) * rho * e_new
-    vx_new = vx + dt * ax
-    vy_new = vy + dt * ay
+def hydro_euler2d_accelerate(rho, vx, vy, P, ax, ay, gamma, dx, dy, dt):
+    Mass, Momx, Momy, Energy = get_conserved(rho, vx, vy, P, gamma, dx * dy)
+
+    Energy += dt * (Momx * ax + Momy * ay)
+    Momx += dt * Mass * ax
+    Momy += dt * Mass * ay
+
+    _, vx_new, vy_new, P_new = get_primitive(Mass, Momx, Momy, Energy, gamma, dx * dy)
     return vx_new, vy_new, P_new

adirondax/simulation.py

@@ -47,6 +47,19 @@ def __init__(self, params):
         ):
             raise ValueError("'hlld' riemann solver only exists for magnetic=True")
 
+        if (
+            self.params["mesh"]["boundary_condition"][0] != "periodic"
+            or self.params["mesh"]["boundary_condition"][1] != "periodic"
+        ):
+            if self.params["physics"]["quantum"]:
+                raise NotImplementedError(
+                    "Quantum only implemented for periodic boundary conditions."
+                )
+            if self.params["physics"]["gravity"]:
+                raise NotImplementedError(
+                    "Gravity only implemented for periodic boundary conditions."
+                )
+
         if self.params["output"]["save"] and self.params["time"]["num_timesteps"] > 0:
             if (
                 self.params["time"]["num_timesteps"]
@@ -198,10 +211,16 @@ def _evolve(self, state):
         dy = Ly / ny
         nt = self.params["time"]["num_timesteps"]
         t_span = self.params["time"]["span"]
+        bc_x = self.params["mesh"]["boundary_condition"][0]
+        bc_y = self.params["mesh"]["boundary_condition"][1]
 
         use_adaptive_timesteps = True if nt < 1 else False
         dt_ref = jnp.nan if use_adaptive_timesteps else t_span / nt
 
+        # boundary conditions
+        bc_x_is_reflective = True if bc_x == "reflective" else False
+        bc_y_is_reflective = True if bc_y == "reflective" else False
+
         # Physics flags
         use_hydro = self.params["physics"]["hydro"]
         use_magnetic = self.params["physics"]["magnetic"]
@@ -290,14 +309,25 @@ def _kick(state, k_sq, dt):
             # apply
             if use_gravity or use_external_potential:
                 if use_quantum:
-                    state["psi"] = quantum_kick(state["psi"], V, m_per_hbar, dt / 2.0)
+                    state["psi"] = quantum_kick(state["psi"], V, m_per_hbar, dt)
                 if use_hydro:
                     if use_magnetic:
                         raise NotImplementedError("implement me.")
                     kx, ky = self.kgrid
-                    ax, ay = get_acceleration(V, kx, ky)
+                    ax, ay = get_acceleration(
+                        V, kx, ky, dx, dy, bc_x_is_reflective, bc_y_is_reflective
+                    )
                     state["vx"], state["vy"], state["P"] = hydro_euler2d_accelerate(
-                        state["rho"], state["vx"], state["vy"], state["P"], ax, ay, dt
+                        state["rho"],
+                        state["vx"],
+                        state["vy"],
+                        state["P"],
+                        ax,
+                        ay,
+                        gamma,
+                        dx,
+                        dy,
+                        dt,
                     )
 
         def _drift(state, k_sq, dt):
@@ -342,6 +372,8 @@ def _drift(state, k_sq, dt):
                             dt,
                             riemann_solver_type,
                             use_slope_limiting,
+                            bc_x_is_reflective,
+                            bc_y_is_reflective,
                         )
                     )
 

examples/rayleigh_taylor/README.md

@@ -0,0 +1,20 @@
+# Rayleigh-Taylor
+
+Simulate the Rayleigh-Taylor Instability (Euler equations)
+
+Philip Mocz (2025)
+
+Usage:
+
+```console
+python rayleigh_taylor.py
+```
+
+Takes around 7 seconds to run on my macbook (cpu).
+
+
+## Simulation snapshots
+
+<div style="display:flex;flex-wrap:wrap;gap:8px">
+  <img src="output.png" alt="output" width="45%"/>
+</div>

examples/rayleigh_taylor/movie.gif

@@ -0,0 +1,103 @@
+import jax.numpy as jnp
+
+# TODO: REMOVE THE FOLLOWING LINES
+import sys
+
+sys.path.append("../../")
+
+import adirondax as adx
+import time
+import matplotlib.pyplot as plt
+
+"""
+Simulate the Rayleigh-Taylor Instability
+
+Philip Mocz (2025)
+"""
+
+
+def set_up_simulation():
+    # Define the parameters for the simulation
+    nx = 64
+    ny = 192
+    nt = 10000  #  -1
+    t_stop = 15.0
+
+    params = {
+        "physics": {
+            "hydro": True,
+            "external_potential": True,
+        },
+        "mesh": {
+            "type": "cartesian",
+            "resolution": [nx, ny],
+            "box_size": [0.5, 1.5],
+            "boundary_condition": ["periodic", "reflective"],
+        },
+        "time": {
+            "span": t_stop,
+            "num_timesteps": nt,
+        },
+        "output": {
+            "num_checkpoints": 100,
+            "save": True,
+            "plot_dynamic_range": 2.0,
+        },
+        "hydro": {
+            "eos": {"type": "ideal", "gamma": 1.4},
+            "slope_limiting": False,
+        },
+    }
+
+    # Initialize the simulation
+    sim = adx.Simulation(params)
+
+    # Set initial conditions
+    # (heavy fluid on top of light)
+    sim.state["t"] = 0.0
+    X, Y = sim.mesh
+    w0 = 0.0025
+    P0 = 2.5
+    g = 0.1
+    sim.state["rho"] = 1.0 + (Y > 0.75)
+    sim.state["vx"] = jnp.zeros(X.shape)
+    sim.state["vy"] = (
+        w0 * (1.0 - jnp.cos(4.0 * jnp.pi * X)) * (1.0 - jnp.cos(4.0 * jnp.pi * Y / 3.0))
+    )
+    sim.state["P"] = P0 - g * (Y - 0.75) * sim.state["rho"]
+
+    # external potential
+    def external_potential(x, y):
+        V = g * y
+        return V
+
+    sim.external_potential = external_potential
+
+    return sim
+
+
+def make_plot(sim):
+    # Plot the solution
+    plt.figure(figsize=(4, 6), dpi=80)
+    plt.imshow(sim.state["rho"].T, cmap="jet", vmin=0.8, vmax=2.2)
+    plt.gca().invert_yaxis()
+    plt.colorbar(label="density")
+    plt.tight_layout()
+    plt.savefig("output.png", dpi=240)
+    plt.show()
+
+
+def main():
+    sim = set_up_simulation()
+
+    # Evolve the system
+    t0 = time.time()
+    sim.run()
+    print("Run time (s): ", time.time() - t0)
+    print("Steps taken:", sim.steps_taken)
+
+    make_plot(sim)
+
+
+if __name__ == "__main__":
+    main()

examples/rayleigh_taylor/output.png

@@ -3,7 +3,7 @@
 # Usage: ./make_gif.sh /path/to/folder output.gif
 
 folder="$1"
-output="${2:-output.gif}"
+output="${2:-movie.gif}"
 
 if [ -z "$folder" ]; then
   echo "Usage: $0 /path/to/folder [output.gif]"