subsurfacegen/field-scale-dataset-preview

1#!/usr/bin/env python2"""Field-Scale Dataset Preview — NeurIPS-ready viewer for the 4-file sample.3 4Renders five publication-quality figures matching the manuscript's intro5figure style. Pure NumPy + Matplotlib + h5py + hdf5plugin.6 7Usage:8 9    pip install huggingface_hub h5py hdf5plugin numpy matplotlib scipy10    huggingface-cli download subsurfacegen/field-scale-dataset-preview \\11        --repo-type=dataset --local-dir=./preview_data12    python view_preview.py --plot all --input-dir ./preview_data --output-dir ./figs13 14Individual:15    python view_preview.py --plot cube       # 3D cutaway render of the volume16    python view_preview.py --plot orthogonal # 3 orthogonal slices w/ crosshairs17    python view_preview.py --plot slice      # 2D slice w/ source + streamer18    python view_preview.py --plot wavefield  # Wavefield panels w/ velocity bg19    python view_preview.py --plot gather     # Shot gather panels (Left/Middle/Right)20"""21from __future__ import annotations22 23import argparse24import sys25from pathlib import Path26 27# Optional: register HDF5 compression filters used by wavefield + gather files.28try:29    import hdf5plugin  # noqa: F40130except ImportError:31    print("WARNING: hdf5plugin not installed; wavefield + gather reads may fail. "32          "Install via: pip install hdf5plugin", flush=True)33 34import h5py35import matplotlib36 37matplotlib.use("Agg")38import matplotlib.colors as mcolors39import matplotlib.pyplot as plt40import numpy as np41from mpl_toolkits.mplot3d import Axes3D  # noqa: F401  (registers projection)42 43 44# ---------------- Paper-figure constants ----------------45 46# Match render_shot_gather_clean / render_wavefield_panel_clean conventions.47GATHER_SOURCE_INDICES = (8, 32, 56)  # paper: SG_SOURCE_INDICES48GATHER_POSITION_NAMES = ("Left", "Middle", "Right")49GATHER_SOURCE_COLORS = ("#B87100", "#004A77", "#A00000")  # bronze/navy/burgundy50 51# Wavefield time panels (more than the paper's 3 to show fuller progression)52WAVEFIELD_FRAMES = [53    ("0.20 s", 14),    # Start (paper)54    ("1.00 s", 71),55    ("1.96 s", 140),   # Middle (paper)56    ("3.00 s", 214),57    ("4.00 s", 285),58    ("5.00 s", 357),   # End (paper)59]60 61# 3D orthogonal-slices configuration: which inline + crossline + depth slice62INLINE_INDEX = 507     # the inline that the 2D slice was extracted at63CROSSLINE_INDEX = 500  # mid-y by convention64DEPTH_INDEX_FRAC = 0.5  # mid-depth65 66# Velocity colormap (paper convention: TwoSlopeNorm + seismic)67VEL_VMIN, VEL_VCENTER, VEL_VMAX = 1500.0, 2300.0, 4500.068 69# 3D-cube cutaway render config (matches cutaway_3d_d619_batch.py)70CUBE_VEL_VCENTER = 3100.071CUBE_DEPTH_RANGE = (80, 619)72CUBE_CROSSLINE_RANGE = (150, 1000)73CUBE_INLINE_RANGE = (150, 1000)74CUBE_VIEW_ELEV = 2075CUBE_VIEW_AZIM = -11576CUBE_MAX_FACE_DIM = 35077DX_KM = 0.01078 79# NeurIPS-ready typography (apply globally so every figure is consistent)80plt.rcParams.update({81    "font.size": 14,82    "axes.titlesize": 16,83    "axes.labelsize": 14,84    "xtick.labelsize": 12,85    "ytick.labelsize": 12,86    "legend.fontsize": 12,87    "figure.titlesize": 18,88    "axes.titleweight": "bold",89    "figure.titleweight": "bold",90    "axes.linewidth": 1.4,91    "xtick.major.width": 1.2,92    "ytick.major.width": 1.2,93    "xtick.major.size": 5,94    "ytick.major.size": 5,95    "lines.linewidth": 1.6,96    "savefig.bbox": "tight",97    "savefig.dpi": 150,98})99 100 101# ---------------- helpers ----------------102 103def _vel_norm() -> mcolors.TwoSlopeNorm:104    return mcolors.TwoSlopeNorm(vmin=VEL_VMIN, vcenter=VEL_VCENTER, vmax=VEL_VMAX)105 106 107def _wf_alpha_cmap() -> mcolors.ListedColormap:108    """Seismic with alpha proportional to |v|; tiny values become transparent."""109    base = matplotlib.colormaps["seismic"](np.linspace(0, 1, 256))110    t = np.linspace(-1.0, 1.0, 256)111    alphas = np.clip(np.abs(t) * 1.6, 0.0, 1.0)112    alphas[np.abs(t) < 0.04] = 0.0113    base[:, 3] = alphas114    return mcolors.ListedColormap(base, name="seismic_alpha")115 116 117def _load_h5_dataset(path: Path, key: str) -> tuple[np.ndarray, dict]:118    with h5py.File(path, "r") as f:119        ds = f[key]120        data = ds[...]121        attrs = {}122        for k, v in ds.attrs.items():123            attrs[k] = v.decode() if isinstance(v, bytes) else v124        for k, v in f.attrs.items():125            attrs[f"root.{k}"] = v.decode() if isinstance(v, bytes) else v126    return data, attrs127 128 129def _save(fig, output_dir: Path, name: str):130    output_dir.mkdir(parents=True, exist_ok=True)131    out = output_dir / name132    fig.savefig(out, dpi=150)133    plt.close(fig)134    print(f"  wrote {out}")135 136 137def _block_mean(arr: np.ndarray, max_dim: int) -> np.ndarray:138    h, w = arr.shape139    bh = max(1, int(np.ceil(h / max_dim)))140    bw = max(1, int(np.ceil(w / max_dim)))141    if bh == 1 and bw == 1:142        return arr.astype(np.float32, copy=False)143    h2 = (h // bh) * bh144    w2 = (w // bw) * bw145    return (arr[:h2, :w2].astype(np.float32, copy=False)146            .reshape(h2 // bh, bh, w2 // bw, bw).mean(axis=(1, 3)))147 148 149# ---------------- 3D cube cutaway (paper render_cube style) ----------------150 151def plot_3d_cube_cutaway(input_dir: Path, output_dir: Path):152    p = input_dir / "models/gom_d619/gom_151_sos.h5"153    print(f"Loading 3D volume for cutaway: {p}")154    with h5py.File(p, "r") as fh:155        vol = np.asarray(fh["velocity"][...], dtype=np.float32)156    nz, ny, nx = vol.shape157    print(f"  shape (nz, ny, nx) = ({nz}, {ny}, {nx})")158    norm = mcolors.TwoSlopeNorm(vmin=VEL_VMIN, vcenter=CUBE_VEL_VCENTER, vmax=VEL_VMAX)159    cmap = plt.get_cmap("seismic")160 161    # Take the cutaway sub-volume (top + 2 sides shaved per paper config)162    z0, z1 = CUBE_DEPTH_RANGE163    y0, y1 = CUBE_CROSSLINE_RANGE164    x0, x1 = CUBE_INLINE_RANGE165    raw_faces = {166        "top":    vol[z0,         y0:y1, x0:x1],167        "bottom": vol[z1 - 1,     y0:y1, x0:x1],168        "left":   vol[z0:z1,      y0:y1, x0],169        "right":  vol[z0:z1,      y0:y1, x1 - 1],170        "front":  vol[z0:z1,      y0,    x0:x1],171        "back":   vol[z0:z1,      y1 - 1, x0:x1],172    }173    faces = {k: _block_mean(v, CUBE_MAX_FACE_DIM) for k, v in raw_faces.items()}174    del vol175 176    x_km_lim = (x0 * DX_KM, (x1 - 1) * DX_KM)177    y_km_lim = (y0 * DX_KM, (y1 - 1) * DX_KM)178    z_km_lim = (z0 * DX_KM, (z1 - 1) * DX_KM)179 180    def coords(face_shape, a_lim, b_lim):181        nh, nw = face_shape182        aa = np.linspace(a_lim[0], a_lim[1], nh)183        bb = np.linspace(b_lim[0], b_lim[1], nw)184        A, B = np.meshgrid(aa, bb, indexing="ij")185        return A, B186 187    fig = plt.figure(figsize=(10, 9))188    ax = fig.add_subplot(111, projection="3d")189    ax.set_facecolor("white")190    surf_kw = dict(shade=False, rstride=1, cstride=1, linewidth=0,191                   antialiased=False, rasterized=True)192 193    Y, X = coords(faces["top"].shape, y_km_lim, x_km_lim)194    ax.plot_surface(X, Y, np.full_like(X, z_km_lim[0]),195                    facecolors=cmap(norm(faces["top"])), **surf_kw)196    Y, X = coords(faces["bottom"].shape, y_km_lim, x_km_lim)197    ax.plot_surface(X, Y, np.full_like(X, z_km_lim[1]),198                    facecolors=cmap(norm(faces["bottom"])), **surf_kw)199    Z2, Y2 = coords(faces["left"].shape, z_km_lim, y_km_lim)200    ax.plot_surface(np.full_like(Y2, x_km_lim[0]), Y2, Z2,201                    facecolors=cmap(norm(faces["left"])), **surf_kw)202    Z2, Y2 = coords(faces["right"].shape, z_km_lim, y_km_lim)203    ax.plot_surface(np.full_like(Y2, x_km_lim[1]), Y2, Z2,204                    facecolors=cmap(norm(faces["right"])), **surf_kw)205    Z3, X3 = coords(faces["front"].shape, z_km_lim, x_km_lim)206    ax.plot_surface(X3, np.full_like(X3, y_km_lim[0]), Z3,207                    facecolors=cmap(norm(faces["front"])), **surf_kw)208    Z3, X3 = coords(faces["back"].shape, z_km_lim, x_km_lim)209    ax.plot_surface(X3, np.full_like(X3, y_km_lim[1]), Z3,210                    facecolors=cmap(norm(faces["back"])), **surf_kw)211 212    ax.set_xlim(*x_km_lim); ax.set_ylim(*y_km_lim); ax.set_zlim(*z_km_lim)213    ax.invert_zaxis()214    ax.set_box_aspect((1, 1, 1))215    ax.view_init(elev=CUBE_VIEW_ELEV, azim=CUBE_VIEW_AZIM)216    ax.set_xticks([]); ax.set_yticks([]); ax.set_zticks([])217    for pane in (ax.xaxis.pane, ax.yaxis.pane, ax.zaxis.pane):218        pane.set_facecolor((1, 1, 1, 0)); pane.set_edgecolor((1, 1, 1, 0))219    ax.set_title("3D SOS-smoothed velocity volume (cutaway)\n"220                 "gom_151_sos.h5  —  Gulf of Mexico realization",221                 pad=18)222    # Manual colorbar223    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)224    sm.set_array([])225    cbar = fig.colorbar(sm, ax=ax, shrink=0.55, pad=0.05, aspect=20)226    cbar.set_label("velocity (m/s)", fontsize=14)227    cbar.ax.tick_params(labelsize=12)228    _save(fig, output_dir, "model_3d_cube_cutaway.png")229 230 231# ---------------- 3D orthogonal slices with crosshairs ----------------232 233def plot_3d_orthogonal_slices(input_dir: Path, output_dir: Path):234    p = input_dir / "models/gom_d619/gom_151_sos.h5"235    print(f"Loading 3D volume: {p}")236    vol, _ = _load_h5_dataset(p, "velocity")237    nz, ny, nx = vol.shape238    mid_z = int(round(nz * DEPTH_INDEX_FRAC))239    print(f"  shape (nz, ny, nx) = ({nz}, {ny}, {nx})")240    print(f"  inline_idx={INLINE_INDEX}, crossline_idx={CROSSLINE_INDEX}, depth_idx={mid_z}")241 242    norm = _vel_norm()243    crosshair_kw = dict(color="red", linestyle="--", linewidth=2.0,244                        alpha=0.95, dashes=(4, 3))245 246    fig, axes = plt.subplots(1, 3, figsize=(20, 6.5),247                             gridspec_kw={"wspace": 0.25})248 249    # Panel 1: inline slice (depth × x).  Crossline + depth slices appear as crosshair.250    inline = vol[:, INLINE_INDEX, :]251    im = axes[0].imshow(inline, aspect="auto", cmap="seismic", norm=norm,252                        extent=[0, nx, nz, 0], interpolation="lanczos")253    axes[0].set_title(f"Inline {INLINE_INDEX}\n(the 2D slice file)")254    axes[0].set_xlabel("crossline x (sample)")255    axes[0].set_ylabel("depth (sample)")256    axes[0].axvline(CROSSLINE_INDEX, **crosshair_kw)  # where crossline panel cuts257    axes[0].axhline(mid_z, **crosshair_kw)            # where depth panel cuts258 259    # Panel 2: crossline slice (depth × y).  Inline + depth slices as crosshair.260    cross = vol[:, :, CROSSLINE_INDEX]261    axes[1].imshow(cross, aspect="auto", cmap="seismic", norm=norm,262                   extent=[0, ny, nz, 0], interpolation="lanczos")263    axes[1].set_title(f"Crossline at x={CROSSLINE_INDEX}")264    axes[1].set_xlabel("inline y (sample)")265    axes[1].set_ylabel("depth (sample)")266    axes[1].axvline(INLINE_INDEX, **crosshair_kw)  # where inline panel cuts267    axes[1].axhline(mid_z, **crosshair_kw)         # where depth panel cuts268 269    # Panel 3: depth slice (y × x).  Inline + crossline cuts as crosshair.270    depth = vol[mid_z, :, :]271    axes[2].imshow(depth, aspect="auto", cmap="seismic", norm=norm,272                   extent=[0, nx, ny, 0], interpolation="lanczos")273    axes[2].set_title(f"Depth slice at z={mid_z}")274    axes[2].set_xlabel("crossline x (sample)")275    axes[2].set_ylabel("inline y (sample)")276    axes[2].axhline(INLINE_INDEX, **crosshair_kw)    # where inline panel cuts277    axes[2].axvline(CROSSLINE_INDEX, **crosshair_kw) # where crossline panel cuts278 279    # Colorbar with key above it explaining the dashed lines.280    cbar_ax = fig.add_axes([0.92, 0.13, 0.013, 0.66])281    cbar = fig.colorbar(im, cax=cbar_ax)282    cbar.set_label("velocity (m/s)", fontsize=14)283    cbar.ax.tick_params(labelsize=12)284    # Compact key just above the colorbar285    fig.text(0.926, 0.85,286             "‒‒ red dashed:\nlocations of\nother 2 panels",287             fontsize=11, color="red", weight="bold",288             ha="left", va="bottom")289 290    fig.suptitle("3D SOS volume — orthogonal slices through gom_151_sos.h5",291                 y=0.98)292    _save(fig, output_dir, "model_3d_orthogonal_slices.png")293 294 295# ---------------- 2D slice with source + streamer geometry ----------------296 297def plot_2d_slice(input_dir: Path, output_dir: Path):298    """2D velocity slice with seismic cmap, streamer + source geometry overlay."""299    p = input_dir / "slices/slice_gom_151_il_0507.h5"300    print(f"Loading 2D slice: {p}")301    sl, _ = _load_h5_dataset(p, "velocity")302    nz, nx = sl.shape303    print(f"  shape (nz, nx) = ({nz}, {nx})")304 305    # Physical extent (10 m grid spacing)306    extent = (0.0, nx * 0.01, nz * 0.01, 0.0)307    norm = _vel_norm()308 309    fig, ax = plt.subplots(1, 1, figsize=(13, 5.8))310    im = ax.imshow(sl, cmap="seismic", norm=norm,311                   extent=extent, aspect="auto", interpolation="lanczos")312    ax.set_xlabel("horizontal x (km)")313    ax.set_ylabel("depth z (km)")314    ax.set_title("2D velocity slice  —  slice_gom_151_il_0507.h5  (inline 507 of gom_151)")315 316    # Streamer (1000 receivers @ 10 m depth, 10 m horizontal spacing,317    # standard 0.5 km edge margin per the simulation config).318    rec_z_km = 0.01319    rec_x_km = np.linspace(0.6, 9.4, 1000)320    rec_x_show = rec_x_km[::25]321    ax.scatter(rec_x_show, np.full_like(rec_x_show, rec_z_km),322               marker="v", s=24, c="black", edgecolors="white",323               linewidths=0.4, zorder=10, label="receivers (every 25th)")324 325    # Source for the wavefield panel (paper figure)326    ax.scatter([4.958], [rec_z_km], marker="*", s=460,327               c="gold", edgecolors="black", linewidths=1.2,328               zorder=11, label="wavefield source @ x=4.958 km")329    ax.legend(loc="upper right", framealpha=0.92)330    cbar = fig.colorbar(im, ax=ax, pad=0.01, shrink=0.95)331    cbar.set_label("velocity (m/s)")332    _save(fig, output_dir, "slice_2d_velocity.png")333 334 335# ---------------- Wavefield panels (paper render style) ----------------336 337def plot_wavefield(input_dir: Path, output_dir: Path,338                   frames: list[tuple[str, int]] = WAVEFIELD_FRAMES):339    """Wavefield time progression with the velocity model drawn behind."""340    p = input_dir / "wavefields/5s/3-6Hz/wavefield_gom_151_il_0507_srchorizontal4.958km.h5"341    print(f"Loading wavefield: {p}")342    wf, _ = _load_h5_dataset(p, "wavefield")343    print(f"  wavefield shape (nt, _, _) = {wf.shape}")344    # Background velocity slice (loaded for the underlay)345    vp = input_dir / "slices/slice_gom_151_il_0507.h5"346    vel, _ = _load_h5_dataset(vp, "velocity")347    print(f"  velocity background shape (nz, nx) = {vel.shape}")348 349    extent = (0.0, 10.0, 6.19, 0.0)350    src_x_km = 4.958351    wf_cmap = _wf_alpha_cmap()352 353    n = len(frames)354    ncols = 3355    nrows = int(np.ceil(n / ncols))356    fig, axes = plt.subplots(nrows, ncols, figsize=(6.5 * ncols, 4.8 * nrows),357                             gridspec_kw={"wspace": 0.18, "hspace": 0.30})358    axes = np.atleast_2d(axes).flatten()359 360    for idx, (label, frame_idx) in enumerate(frames):361        ax = axes[idx]362        frame_idx = max(0, min(frame_idx, wf.shape[0] - 1))363        # Wavefield is stored (nt, nx, nz) where nx=1000, nz=619 (devito364        # convention).  Transpose to (nz, nx) so it aligns with the velocity365        # slice (nz=619, nx=1000) when overlaid on the same physical extent.366        wf_frame = wf[frame_idx].T367        # Underlay velocity model in gray, alpha 0.40368        ax.imshow(vel, cmap="gray", alpha=0.40, extent=extent,369                  aspect="auto", interpolation="nearest",370                  vmin=float(vel.min()), vmax=float(vel.max()))371        # Wavefield with seismic-alpha cmap (transparent near 0)372        clip = float(np.percentile(np.abs(wf_frame), 99.5))373        if clip <= 0:374            clip = float(np.abs(wf_frame).max() + 1e-12)375        ax.imshow(wf_frame, cmap=wf_cmap,376                  norm=mcolors.Normalize(-clip, clip),377                  extent=extent, aspect="auto", interpolation="nearest")378        # Source marker379        ax.scatter([src_x_km], [0.0], marker="*", s=320, c="gold",380                   edgecolors="black", linewidths=1.2, zorder=10)381        ax.set_xlim(0, 10); ax.set_ylim(6.19, 0)382        ax.set_xlabel("horizontal x (km)")383        ax.set_ylabel("depth z (km)")384        ax.set_title(f"t = {label}", fontsize=15)385 386    for j in range(n, len(axes)):387        axes[j].set_visible(False)388 389    fig.suptitle("Wavefield time progression  —  3-6 Hz, source @ x=4.958 km\n"390                 "(velocity model in grayscale; wavefield amplitude in red/blue)",391                 y=1.00, fontsize=17)392    _save(fig, output_dir, "wavefield_time_progression.png")393 394 395# ---------------- Shot gather (paper render style) ----------------396 397def plot_shot_gather(input_dir: Path, output_dir: Path,398                     source_indices=GATHER_SOURCE_INDICES,399                     position_names=GATHER_POSITION_NAMES,400                     source_colors=GATHER_SOURCE_COLORS):401    """1x4 layout: velocity model with color-coded source positions (left)402    + 3 shot-gather panels (Left/Middle/Right matching paper figure)."""403    p_gather = input_dir / "shot_gathers/8s/3-25Hz/shot_gather_cube_gom_151_il_0507.h5"404    p_slice = input_dir / "slices/slice_gom_151_il_0507.h5"405    print(f"Loading shot-gather cube: {p_gather}")406    cube, gather_attrs = _load_h5_dataset(p_gather, "shot_gather_cube")407    n_src, nt, n_rec = cube.shape408    print(f"  shape (n_src, nt, n_rec) = ({n_src}, {nt}, {n_rec})")409    print(f"Loading velocity slice for geometry panel: {p_slice}")410    vel, _ = _load_h5_dataset(p_slice, "velocity")411 412    # Source x-positions in km from the gather cube's flat attr (length 64).413    src_x_all = np.asarray(gather_attrs.get("source_horizontal_km", []), dtype=float)414    if src_x_all.size != n_src:415        # Fall back: linear sample if attr missing/wrong length416        src_x_all = np.linspace(0.6, 9.4, n_src)417    rec_z_km = float(gather_attrs.get("receiver_depth_km", 0.01))418 419    delta_t_stored = 14e-3420    total_s = nt * delta_t_stored421 422    # 1 x 4 grid:  velocity panel (wider so it stays roughly square)423    # then three gather panels.424    fig, axes = plt.subplots(425        1, 4,426        figsize=(26, 6.8),427        gridspec_kw={"width_ratios": [1.55, 1.0, 1.0, 1.0], "wspace": 0.22},428    )429 430    # ----- Panel 0: velocity model with source positions overlay -----431    ax0 = axes[0]432    extent_vel = (0.0, 10.0, 6.19, 0.0)433    ax0.imshow(vel, cmap="seismic", norm=_vel_norm(),434               extent=extent_vel, aspect="auto", interpolation="lanczos")435    # Streamer (1000 receivers @ 10 m depth, every 25th shown)436    rec_x_all = np.linspace(0.6, 9.4, 1000)437    rec_x_show = rec_x_all[::25]438    ax0.scatter(rec_x_show, np.full_like(rec_x_show, rec_z_km),439                marker="v", s=22, c="black", edgecolors="white",440                linewidths=0.4, zorder=10, label="receivers (every 25th)")441    # Highlight the 3 chosen sources at their actual horizontal_km442    for sidx, color, name in zip(source_indices, source_colors, position_names):443        sx = float(src_x_all[sidx])444        ax0.scatter([sx], [rec_z_km], marker="*", s=520,445                    c=color, edgecolors="black", linewidths=1.4, zorder=11,446                    label=f"{name}  (idx {sidx}, x={sx:.2f} km)")447    ax0.set_xlim(0, 10); ax0.set_ylim(6.19, 0)448    ax0.set_xlabel("horizontal x (km)")449    ax0.set_ylabel("depth z (km)")450    ax0.set_title("Source positions on the velocity model", fontsize=15)451    ax0.legend(loc="lower right", framealpha=0.92, fontsize=10)452 453    # ----- Panels 1-3: gathers -----454    clip = float(np.percentile(np.abs(cube), 97.0))455    if clip <= 0:456        clip = float(np.abs(cube).max() + 1e-12)457    for ax, sidx, label, color in zip(axes[1:], source_indices,458                                      position_names, source_colors):459        sidx = max(0, min(sidx, n_src - 1))460        gather = cube[sidx]461        ax.imshow(gather, cmap="seismic", vmin=-clip, vmax=clip,462                  aspect="auto", interpolation="lanczos",463                  extent=[0, n_rec, total_s, 0])464        ax.set_xlabel("receiver index")465        ax.set_ylabel("time (s)")466        for s in ax.spines.values():467            s.set_color(color); s.set_linewidth(2.4)468        ax.set_title(f"{label}  (source idx {sidx})",469                     color=color, fontweight="bold", fontsize=15)470 471    fig.suptitle("Shot-gather panels  —  3-25 Hz, 64-source cube  "472                 "(velocity model on left shows the 3 source positions; matches manuscript intro figure)",473                 y=1.02, fontsize=17)474    _save(fig, output_dir, "shot_gather_panels.png")475 476 477# ---------------- main ----------------478 479def main():480    ap = argparse.ArgumentParser(description=__doc__,481                                 formatter_class=argparse.RawDescriptionHelpFormatter)482    ap.add_argument("--input-dir", type=Path, default=Path("./preview_data"),483                    help="Dir holding the 4 h5 files. Default: ./preview_data")484    ap.add_argument("--output-dir", type=Path, default=Path("./figs"),485                    help="Where to save PNGs. Default: ./figs")486    ap.add_argument("--plot",487                    choices=("cube", "orthogonal", "slice", "wavefield", "gather", "all"),488                    default="all", help="Which figure to render. Default: all")489    args = ap.parse_args()490 491    if not args.input_dir.exists():492        sys.exit(f"--input-dir does not exist: {args.input_dir}")493 494    plots = ("cube", "orthogonal", "slice", "wavefield", "gather") \495        if args.plot == "all" else (args.plot,)496    for p in plots:497        if p == "cube": plot_3d_cube_cutaway(args.input_dir, args.output_dir)498        elif p == "orthogonal": plot_3d_orthogonal_slices(args.input_dir, args.output_dir)499        elif p == "slice": plot_2d_slice(args.input_dir, args.output_dir)500        elif p == "wavefield": plot_wavefield(args.input_dir, args.output_dir)501        elif p == "gather": plot_shot_gather(args.input_dir, args.output_dir)502    print(f"\nDone. PNGs in {args.output_dir}/")503 504 505if __name__ == "__main__":506    main()507