add alternative build anim function

Author: JoelLucaAdams

docs/epoch_workshop_2026/animating/animate_2D_dist_fn_multispecies_alternative.py

@@ -0,0 +1,56 @@
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+from matplotlib.animation import FuncAnimation
+
+import sdf_xarray as sdfxr
+
+input_dir = Path("datasets/2_1_two_stream_instability")
+ds = sdfxr.open_mfdataset(
+    input_dir, data_vars=["dist_fn_x_px_Left", "dist_fn_x_px_Right"]
+)
+
+# Rescale coords to account for kilometers
+ds = ds.epoch.rescale_coords(1e-3, "km", ["X_x_px_Left"])
+
+# Sum phase-space of species "Left" and "Right" in "x_px" distribution function
+# NOTE: We only use the values from the right distribution function as if we inherit
+# the coords we from the right we end up with 4 coords instead of 2
+total_phase_space = ds["dist_fn_x_px_Left"] + ds["dist_fn_x_px_Right"].values
+total_phase_space.attrs["long_name"] = "Phase Space Distribution"
+total_phase_space.attrs["units"] = "kg.m/s"
+
+fig, ax = plt.subplots()
+
+# construct first frame
+# Data needs to be transposed so that axes match
+total_phase_space = total_phase_space.T
+col_min = total_phase_space.values.min()
+col_max = total_phase_space.values.max()
+gif_frame = ax.pcolormesh(
+    total_phase_space["X_x_px_Left"],
+    total_phase_space["Px_x_px_Left"],
+    total_phase_space.isel(time=0),
+    vmin=col_min,
+    vmax=col_max,
+)
+ax.set_xlabel(
+    f"{total_phase_space['X_x_px_Left'].attrs['long_name']} [{total_phase_space['X_x_px_Left'].attrs['units']}]"
+)
+ax.set_ylabel(
+    f"{total_phase_space.attrs['long_name']} [{total_phase_space.attrs['units']}]"
+)
+ax.set_title(f"t = {ds['time'].values[0]:.3f} [{ds['time'].attrs['units']}]")
+cbar = fig.colorbar(gif_frame, ax=ax)
+cbar.set_label("Summed particle weight per bin")
+
+
+# Make GIF
+def update(i):
+    gif_frame.set_array(total_phase_space.isel(time=i))
+    ax.set_title(f"t = {ds['time'].values[i]:.3f} [{ds['time'].attrs['units']}]")
+    return (gif_frame,)
+
+
+anim = FuncAnimation(fig, update, frames=ds.sizes["time"])
+anim.save(input_dir / "phase_space.gif")