vs_filtered_inversion.py

import os
import glob
import re
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from scipy.interpolate import interp1d
from scipy.ndimage import gaussian_filter

# ============================================================
# USER PARAMETERS
# ============================================================

input_dir = "/home/lea/Desktop/code/ag835/seismic/output_vs_inversion_from_csv/"
summary_csv = os.path.join(input_dir, "vs_inversion_summary.csv")
output_dir = os.path.join(input_dir, "filtered_section")
os.makedirs(output_dir, exist_ok=True)

# ------------------------------------------------------------
# FILTERS
# ------------------------------------------------------------
max_rms_misfit = 60.0
min_n_picks = 90
min_freq_span_hz = 15.0

# Optional bounds on median picked velocity
use_median_velocity_bounds = False
median_vmin = 1500.0
median_vmax = 1750.0

# ------------------------------------------------------------
# DEPTH GRID
# ------------------------------------------------------------
zmax_m = 220.0
dz_m = 2.0
z_grid = np.arange(0.0, zmax_m + dz_m, dz_m)

# ------------------------------------------------------------
# SMOOTHING
# ------------------------------------------------------------
apply_smoothing = True
smooth_sigma_z = 1.0
smooth_sigma_x = 1.0

# ------------------------------------------------------------
# PLOT
# ------------------------------------------------------------
cmap = "viridis"
vmin_plot = None
vmax_plot = None


# ============================================================
# HELPERS
# ============================================================

def interp_profile(z_nodes, vs_model, z_grid):
    f = interp1d(z_nodes, vs_model, kind="linear", bounds_error=False, fill_value=np.nan)
    return f(z_grid)

def load_inversion_npz(path):
    d = np.load(path, allow_pickle=True)
    return {
        "shot_id": int(np.ravel(d["shot_id"])[0]),
        "source_x_m": float(np.ravel(d["source_x_m"])[0]),
        "z_nodes_m": np.asarray(d["z_nodes_m"], dtype=float),
        "vs_model_mps": np.asarray(d["vs_model_mps"], dtype=float),
        "rms_misfit_mps": float(np.ravel(d["rms_misfit_mps"])[0]),
        "quality_flag": str(np.ravel(d["quality_flag"])[0]),
    }

def plot_section(section_x, z_grid, section_vs, outfile, title):
    fig, ax = plt.subplots(figsize=(14, 7))

    masked = np.ma.masked_invalid(section_vs)
    im = ax.imshow(
        masked,
        aspect="auto",
        origin="upper",
        extent=[section_x.min(), section_x.max(), z_grid.max(), z_grid.min()],
        cmap=cmap,
        vmin=vmin_plot,
        vmax=vmax_plot
    )

    ax.set_xlabel("Distance along line (m)")
    ax.set_ylabel("Depth (m)")
    ax.set_title(title)

    cbar = plt.colorbar(im, ax=ax)
    cbar.set_label("Vs (m/s)")

    plt.tight_layout()
    plt.savefig(outfile, dpi=220)
    plt.close(fig)


# ============================================================
# LOAD SUMMARY
# ============================================================

summary = pd.read_csv(summary_csv)

keep = (
    (summary["quality_flag"] == "good") &
    (summary["rms_misfit_mps"] <= max_rms_misfit) &
    (summary["n_picks"] >= min_n_picks) &
    (summary["freq_span_hz"] >= min_freq_span_hz)
)

if use_median_velocity_bounds:
    keep &= (
        (summary["picked_velocity_median_mps"] >= median_vmin) &
        (summary["picked_velocity_median_mps"] <= median_vmax)
    )

summary_keep = summary.loc[keep].copy().sort_values("source_x_m").reset_index(drop=True)

print("Total shots in summary :", len(summary))
print("Shots kept after filter:", len(summary_keep))

if len(summary_keep) == 0:
    raise RuntimeError("No shots passed the filter.")

# Save filtered summary
filtered_summary_csv = os.path.join(output_dir, "vs_inversion_summary_filtered.csv")
summary_keep.to_csv(filtered_summary_csv, index=False)
print("Saved:", filtered_summary_csv)

# ============================================================
# LOAD KEPT INVERSION FILES
# ============================================================

section_x = []
section_profiles = []

for _, row in summary_keep.iterrows():
    sid = int(row["shot_id"])
    npz_file = os.path.join(input_dir, f"shot_{sid}_vs_inversion.npz")

    if not os.path.exists(npz_file):
        print(f"Missing inversion file for shot {sid}, skipped.")
        continue

    inv = load_inversion_npz(npz_file)

    vs_on_grid = interp_profile(inv["z_nodes_m"], inv["vs_model_mps"], z_grid)

    section_x.append(inv["source_x_m"])
    section_profiles.append(vs_on_grid)

if len(section_x) == 0:
    raise RuntimeError("No inversion files could be loaded for kept shots.")

section_x = np.asarray(section_x, dtype=float)
section_profiles = np.asarray(section_profiles, dtype=float).T  # [nz, nx]

# Sort by x
order = np.argsort(section_x)
section_x = section_x[order]
section_profiles = section_profiles[:, order]

# ============================================================
# OPTIONAL SMOOTHING
# ============================================================

section_vs_raw = section_profiles.copy()

if apply_smoothing:
    tmp = section_profiles.copy()
    nanmask = ~np.isfinite(tmp)

    if np.any(np.isfinite(tmp)):
        fill_value = np.nanmedian(tmp[np.isfinite(tmp)])
        tmp[nanmask] = fill_value
        tmp = gaussian_filter(tmp, sigma=(smooth_sigma_z, smooth_sigma_x))
        tmp[nanmask] = np.nan
        section_profiles = tmp

section_vs_smooth = section_profiles

# ============================================================
# SAVE ARRAYS
# ============================================================

np.savez_compressed(
    os.path.join(output_dir, "vs_section_filtered.npz"),
    section_x_m=section_x,
    depth_grid_m=z_grid,
    vs_section_raw_mps=section_vs_raw,
    vs_section_smooth_mps=section_vs_smooth
)

# ============================================================
# PLOTS
# ============================================================

plot_section(
    section_x,
    z_grid,
    section_vs_raw,
    outfile=os.path.join(output_dir, "vs_section_filtered_raw.png"),
    title="Filtered Vs section (raw profiles)"
)

plot_section(
    section_x,
    z_grid,
    section_vs_smooth,
    outfile=os.path.join(output_dir, "vs_section_filtered_smooth.png"),
    title="Filtered Vs section (smoothed)"
)

# Also plot kept shot positions
fig, ax = plt.subplots(figsize=(12, 3))
ax.plot(summary["source_x_m"], np.zeros(len(summary)), "|", ms=18, color="lightgray", label="all shots")
ax.plot(summary_keep["source_x_m"], np.ones(len(summary_keep)), "|", ms=18, color="green", label="kept shots")
ax.set_yticks([0, 1])
ax.set_yticklabels(["all", "kept"])
ax.set_xlabel("Distance along line (m)")
ax.set_title("Shot filtering overview")
ax.grid(True, axis="x", alpha=0.3)
ax.legend()
plt.tight_layout()
plt.savefig(os.path.join(output_dir, "shot_filtering_overview.png"), dpi=200)
plt.close()

print("Saved:")
print(os.path.join(output_dir, "vs_section_filtered_raw.png"))
print(os.path.join(output_dir, "vs_section_filtered_smooth.png"))
print(os.path.join(output_dir, "shot_filtering_overview.png"))
print("Done.")