File size: 2,404 Bytes
68d9039 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | 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)
|