Upload code/plots.py

code/plots.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import matplotlib
+
+matplotlib.use("Agg")
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def plot_reconstruction_batch(
+    path: str | Path,
+    target_loglam,
+    target_flux,
+    pred_flux,
+    loss_mask,
+    valid_patch,
+    z_true,
+    z_pred,
+    max_items: int = 4,
+) -> None:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    bsz = min(max_items, len(target_flux))
+    fig, axes = plt.subplots(bsz, 1, figsize=(13, 3.2 * bsz), squeeze=False)
+    for i in range(bsz):
+        ax = axes[i, 0]
+        wave = np.exp(target_loglam[i])
+        valid = valid_patch[i].astype(bool)
+        masked = loss_mask[i].astype(bool)
+        ax.plot(wave[valid], target_flux[i][valid], color="black", linewidth=1.0, label="target")
+        ax.plot(wave[valid], pred_flux[i][valid], color="#d62728", linewidth=1.0, alpha=0.9, label="recon")
+        if masked.any():
+            ax.scatter(wave[masked], target_flux[i][masked], s=8, color="#1f77b4", alpha=0.55, label="masked target")
+        ax.set_ylabel("norm flux")
+        ax.set_title(f"z true={z_true[i]:.5f}  z pred={z_pred[i]:.5f}")
+        ax.grid(alpha=0.2)
+        if i == 0:
+            ax.legend(loc="best", fontsize=8)
+    axes[-1, 0].set_xlabel("wavelength Angstrom")
+    fig.tight_layout()
+    fig.savefig(path, dpi=150)
+    plt.close(fig)
+
+
+def plot_redshift_scatter(path: str | Path, y_true: np.ndarray, y_pred: np.ndarray) -> None:
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    z_true = np.expm1(y_true)
+    z_pred = np.expm1(y_pred)
+    fig, axes = plt.subplots(1, 2, figsize=(11, 4.2))
+    axes[0].scatter(z_true, z_pred, s=6, alpha=0.35)
+    lim = [float(np.nanmin(z_true)), float(np.nanmax(z_true))]
+    axes[0].plot(lim, lim, color="black", linewidth=1)
+    axes[0].set_xlabel("z true")
+    axes[0].set_ylabel("z pred")
+    axes[0].grid(alpha=0.2)
+    dz = (z_pred - z_true) / (1 + z_true)
+    axes[1].scatter(z_true, dz, s=6, alpha=0.35)
+    axes[1].axhline(0, color="black", linewidth=1)
+    axes[1].axhline(0.01, color="#d62728", linewidth=1, linestyle="--")
+    axes[1].axhline(-0.01, color="#d62728", linewidth=1, linestyle="--")
+    axes[1].set_xlabel("z true")
+    axes[1].set_ylabel("delta z / (1 + z)")
+    axes[1].grid(alpha=0.2)
+    fig.tight_layout()
+    fig.savefig(path, dpi=150)
+    plt.close(fig)