diff --git a/Examples/2D_magnetic_reconnection.py b/Examples/2D_magnetic_reconnection.py new file mode 100644 index 0000000..cae6f77 --- /dev/null +++ b/Examples/2D_magnetic_reconnection.py @@ -0,0 +1,258 @@ +"""Periodic double-Harris-sheet magnetic reconnection. + +The short defaults are a smoke test. For a clear nonlinear island movie, run + + python Examples/2D_magnetic_reconnection.py --nx 65 --ny 17 \ + --t-max 500 --chunk-time 50 --snapshots 51 --perturbation 0.1 + +Use ``--perturbation 1e-4`` instead for a linear tearing-growth diagnostic. + +Growth observable: the amplitude of the first y Fourier mode (ky = 2 pi / Ly) of +(Bx, By), taken over the left half of the box, which contains the sheet at Lx/4. +gamma is a least-squares fit of its logarithm over ``--fit-window`` (default: the +second half of the run). The script also prints the fit's rms log residual and the +growth rate of each half of the window; a linear eigenmode needs them to agree. + +Earlier measurements (ny = 5, nu = 1, t = 1000) against Camporeale et al.'s +0.256 Omega_ci = 5.33e-3 omega_pe: + + nx 65, H 10: 2.39e-3 (-55 %) nx 129, H 6: 4.76e-3 (-11 %, 53 min) + nx 129, H 8: 4.84e-3 ( -9 %) nx 257, H 6: 5.07e-3 (-4.9 %, 3.9 h) + +They used an earlier fit (log-linear through local maxima of the amplitude with +t >= t_max/4) and were not re-measured with this window fit. The rate rises with nx, +but no spatial convergence order has been measured, so no extrapolated value is +quoted. Times are on 36 shared CPU cores. +The parameters follow Camporeale et al. (2006), section IV.C; the periodic +double-sheet geometry and half-sheet magnetic diagnostic follow Koshkarov et +al. (2021). +""" + +import argparse +import json +import time +from pathlib import Path + +import jax +import jax.numpy as jnp +import matplotlib.pyplot as plt +import numpy as np +from diffrax import Dopri8 +from matplotlib.animation import FuncAnimation, PillowWriter + +from spectrax import hermite_tail_fraction, plasma_current, simulation + + +def initial_conditions(nx, ny, hermite, perturbation, nu): + a = jnp.sqrt(2.0) + lx, ly = 32.0 * a, 4.0 * jnp.pi * a + x = jnp.linspace(0.0, lx, nx, endpoint=False) + y = jnp.linspace(0.0, ly, ny, endpoint=False) + X, Y, Z = jnp.meshgrid(x, y, jnp.zeros(1), indexing="xy") + sheets = ( + jnp.cosh((X - lx / 4.0) / a) ** -2, + jnp.cosh((X - 3.0 * lx / 4.0) / a) ** -2, + ) + density = jnp.stack((sheets[0], sheets[1], sheets[0], sheets[1])) + alpha = a * jnp.array(([1.0 / 6.0] * 6) + ([1.0 / 24.0] * 6)) + coefficients = jnp.zeros( + (4, hermite, hermite, hermite, ny, nx, 1), dtype=jnp.float64 + ) + coefficients = coefficients.at[:, 0, 0, 0].set( + density / jnp.prod(alpha.reshape(4, 3), axis=1)[:, None, None, None] + ) + + kx, ky = 2.0 * jnp.pi / lx, 2.0 * jnp.pi / ly + magnetic = jnp.array( + [ + -perturbation * ky * jnp.sin(ky * Y) * jnp.sin(kx * X), + jnp.tanh((X - lx / 4.0) / a) + - jnp.tanh((X - 3.0 * lx / 4.0) / a) + - 1.0 + - perturbation * kx * jnp.cos(ky * Y) * jnp.cos(kx * X), + 5.0 * jnp.ones_like(X), + ] + ) + transform = lambda value: jnp.fft.rfftn( + value, axes=(-1, -3, -2), norm="forward" + ) + return { + "Lx": lx, + "Ly": ly, + "Lz": 1.0, + "mi_me": 16.0, + "qs": jnp.array((-1.0, -1.0, 1.0, 1.0)), + "alpha_s": alpha, + "u_s": jnp.array( + (0.0, 0.0, -1.0 / 6.0, 0.0, 0.0, 1.0 / 6.0, + 0.0, 0.0, 1.0 / 6.0, 0.0, 0.0, -1.0 / 6.0) + ) / a, + "Omega_cs": jnp.array((1.0 / 3.0, 1.0 / 3.0, 1.0 / 48.0, 1.0 / 48.0)), + "nu": nu, + "D": 0.0, + "ode_tolerance": 1e-7, + "Ck_0": transform(coefficients), + "Fk_0": transform(jnp.concatenate((jnp.zeros_like(magnetic), magnetic))), + } + + +def fields_and_flux(fk, nx, lx, ly): + ny = fk.shape[-3] + fields = np.fft.irfftn( + fk, s=(1, ny, nx), axes=(-1, -3, -2), norm="forward" + ).real + kx = 2.0 * np.pi * np.fft.rfftfreq(nx, d=lx / nx) + ky = 2.0 * np.pi * np.fft.fftfreq(ny, d=ly / ny) + KX, KY = np.meshgrid(kx, ky) + bx, by = fk[:, 3, :, :, 0], fk[:, 4, :, :, 0] + k2 = KX**2 + KY**2 + az_k = np.zeros_like(bx) + np.divide(-1j * KY * bx + 1j * KX * by, k2, out=az_k, where=k2 != 0.0) + az = np.fft.irfftn(az_k, s=(ny, nx), axes=(-2, -1), norm="forward").real + return fields, az + + +def kinetic_current_z(output, hermite, omega_ce): + """Ampere source J_z / Omega_ce from the Hermite moments. It equals (curl B)_z only + when the displacement current dE_z/dt vanishes.""" + current = jax.vmap(lambda ck: plasma_current( + output["qs"], output["alpha_s"], output["u_s"], ck, + hermite, hermite, hermite, 4))(output["Ck"]) + return np.asarray(current[:, 2, :, :, 0]) / omega_ce + + +def fit_growth(t, amplitude, window): + inside = (t >= window[0]) & (t <= window[1]) & (amplitude > 0) + if inside.sum() < 4: + return np.nan, np.nan, (np.nan, np.nan) + ts, logs = t[inside], np.log(amplitude[inside]) + gamma, intercept = np.polyfit(ts, logs, 1) + rms = float(np.sqrt(np.mean((logs - intercept - gamma * ts) ** 2))) + half = len(ts) // 2 + halves = tuple(float(np.polyfit(ts[s], logs[s], 1)[0]) for s in (slice(None, half + 1), slice(half, None))) + return float(gamma), rms, halves + + +def save_diagnostics(output_dir, t, fk, jz_k, nx, lx, ly, window): + fields, az = fields_and_flux(fk, nx, lx, ly) + jz = np.fft.irfftn(jz_k, s=(fk.shape[-3], nx), axes=(-2, -1), norm="forward").real + half = (nx + 1) // 2 + magnetic = fields[:, 3:, :, :half, 0] + phase = np.exp(2j * np.pi * np.arange(magnetic.shape[2]) / magnetic.shape[2]) + amplitude = np.linalg.norm( + np.sum(magnetic * phase[None, None, :, None], axis=(-2, -1)), axis=1 + ) / np.sqrt(magnetic.shape[2] * magnetic.shape[3]) + gamma, rms, halves = fit_growth(t, amplitude, window) + + sheet = np.argmin(abs(np.linspace(0.0, lx, nx, endpoint=False) - lx / 4.0)) + flux = np.ptp(az[:, :, sheet], axis=1) + fig, axes = plt.subplots(2, 1, figsize=(6.4, 6.0), sharex=True, constrained_layout=True) + axes[0].semilogy(t, amplitude, color="0.65", label="half-sheet magnetic mode") + if np.isfinite(gamma): + axes[0].axvspan(*window, color="0.9", zorder=0) + axes[0].set_title(rf"window fit $\gamma={gamma:.5f}$", fontsize=9) + axes[0].set_ylabel("magnetic-mode amplitude") + axes[0].legend(fontsize=8) + axes[1].plot(t, flux) + axes[1].set(xlabel=r"$\omega_{pe}t$", ylabel=r"reconnected flux $\Delta A_z$") + fig.savefig(output_dir / "reconnection-growth.png", dpi=180) + plt.close(fig) + + x = np.linspace(0.0, lx, nx, endpoint=False)[:half] + y = np.linspace(-ly / 2.0, ly / 2.0, fk.shape[-3], endpoint=False) + az = np.roll(az, fk.shape[-3] // 2, axis=1)[:, :, :half] + jz = np.roll(jz, fk.shape[-3] // 2, axis=1)[:, :, :half] + window = (x >= lx / 8.0) & (x <= 3.0 * lx / 8.0) + x, az, jz = x[window], az[:, :, window], jz[:, :, window] + jz -= jz.mean(axis=1, keepdims=True) + limit = max(np.quantile(abs(jz[-1]), 0.99), 1e-12) + sheet = np.argmin(abs(x - lx / 4.0)) + fig, ax = plt.subplots(figsize=(6.4, 3.8), constrained_layout=True) + image = ax.pcolormesh(x, y, jz[0], shading="auto", cmap="RdBu_r", vmin=-limit, vmax=limit) + contours = [None] + title = ax.set_title("") + ax.set(xlabel="x", ylabel="y") + fig.colorbar(image, ax=ax, label=r"kinetic $(J_z-\langle J_z\rangle_y)/\Omega_{ce}$") + + def update(frame): + image.set_array(jz[frame].ravel()) + if contours[0] is not None: + contours[0].remove() + center = az[frame, :, sheet].mean() + span = max(1.25 * np.ptp(az[frame, :, sheet]), 1e-12) + contours[0] = ax.contour( + x, y, az[frame], center + np.linspace(-span, span, 19), + colors="k", linewidths=0.55, + ) + title.set_text(rf"$\omega_{{pe}}t={t[frame]:.1f}$") + return image, title, contours[0] + + frames = range(0, len(t), max(1, len(t) // 100)) + animation = FuncAnimation(fig, update, frames=frames, interval=80, blit=False) + animation.save(output_dir / "reconnection-topology.gif", PillowWriter(fps=12), dpi=110) + plt.close(fig) + return amplitude, gamma, rms, halves + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--nx", type=int, default=33) + parser.add_argument("--ny", type=int, default=9) + parser.add_argument("--hermite", type=int, default=4) + parser.add_argument("--nu", type=float, default=1.0) + parser.add_argument("--perturbation", type=float, default=1e-4) + parser.add_argument("--t-max", type=float, default=10.0) + parser.add_argument("--chunk-time", type=float, default=10.0) + parser.add_argument("--snapshots", type=int, default=21) + parser.add_argument("--fit-window", type=float, nargs=2, default=None, + help="t_start t_end of the growth fit (default: second half of the run)") + parser.add_argument("--output-dir", type=Path, default=Path("reconnection_output")) + args = parser.parse_args() + args.output_dir.mkdir(parents=True, exist_ok=True) + + values = initial_conditions(args.nx, args.ny, args.hermite, args.perturbation, args.nu) + window = args.fit_window or (args.t_max / 2.0, args.t_max) + omega_ce = float(values["Omega_cs"][0]) + times, fields, currents, offset = [], [], [], 0.0 + start = time.perf_counter() + while offset < args.t_max: + duration = min(args.chunk_time, args.t_max - offset) + values["t_max"] = duration + output = simulation( + values, Nx=args.nx, Ny=args.ny, Nz=1, + Nn=args.hermite, Nm=args.hermite, Np=args.hermite, Ns=4, + timesteps=args.snapshots, dt=1e-3, solver=Dopri8(), + ) + jax.block_until_ready(output["Ck"]) + first = 0 if not times else 1 + times.append(np.asarray(output["time"])[first:] + offset) + fields.append(np.asarray(output["Fk"])[first:]) + currents.append(kinetic_current_z(output, args.hermite, omega_ce)[first:]) + tails = np.asarray(hermite_tail_fraction(output)[-1]) + values["Ck_0"] = jnp.asarray(output["Ck"][-1]) + values["Fk_0"] = jnp.asarray(output["Fk"][-1]) + del output + offset += duration + print(f"t={offset:g}") + + t, fk = np.concatenate(times), np.concatenate(fields) + amplitude, gamma, rms, halves = save_diagnostics( + args.output_dir, t, fk, np.concatenate(currents), args.nx, + float(values["Lx"]), float(values["Ly"]), window, + ) + np.savez(args.output_dir / "reconnection.npz", time=t, amplitude=amplitude, Fk=fk) + linear = args.perturbation <= 1e-3 + print(json.dumps({ + "runtime_s": time.perf_counter() - start, + "regime": "linear diagnostic" if linear else "finite-seed nonlinear", + "gamma": None if not np.isfinite(gamma) else float(gamma), + "fit_window": list(window), + "fit_rms_log_residual": None if not np.isfinite(rms) else rms, + "gamma_first_and_second_half": [h if np.isfinite(h) else None for h in halves], + "expected_gamma": 0.256 / 48.0 if linear else None, + "hermite_tail_fraction_per_population": tails.tolist(), + }, sort_keys=True)) + + +if __name__ == "__main__": + main()