make_figures.py

#!/usr/bin/env python3
"""
Generate two publication-quality figures for NGC 6752:

  1) chromosome map with K-means populations  (high-DPI, clean title)
  2) the cluster's colour-magnitude (Hertzsprung-Russell) diagram,
     with the two populations highlighted on the RGB.

Reuses the analysis functions from multiple_populations.py so nothing is
recomputed differently.
"""
from __future__ import annotations

from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib import font_manager  # noqa: F401  (ensures mathtext available)

import multiple_populations as mp

CATALOG = Path("data/m6752_named.txt")
NAME = "NGC 6752"
OUT = Path("figures")
DPI = 300

C_1G = "#2c6fbb"   # first generation
C_2G = "#c0392b"   # second generation

plt.rcParams.update({
    "font.size": 12,
    "axes.titlesize": 15,
    "axes.titleweight": "bold",
    "axes.labelsize": 12,
    "figure.dpi": DPI,
    "savefig.dpi": DPI,
    "savefig.bbox": "tight",
})


def build() -> tuple[pd.DataFrame, pd.DataFrame, mp.Columns]:
    df = mp.load_hugs_catalog(CATALOG)
    cols = mp.resolve_columns(df)
    rgb = mp.select_red_giants(df, cols, membership_cut=0.9,
                               max_rms=0.03, min_qfit=0.9)
    cm = mp.chromosome_map(rgb, cols)
    cm = mp.cluster_populations(cm)
    return df, cm.join(rgb[[cols.f275w, cols.f336w, cols.f438w, cols.f814w]]), cols


def fig_chromosome(cm: pd.DataFrame) -> Path:
    OUT.mkdir(parents=True, exist_ok=True)
    fig, ax = plt.subplots(figsize=(7.5, 6.2))

    for label, color in (("1G (P1)", C_1G), ("2G (P2)", C_2G)):
        sub = cm[cm["population"] == label]
        ax.scatter(sub["dx"], sub["dy"], s=26, alpha=0.78,
                   edgecolors="white", linewidths=0.4, color=color,
                   label=f"{label}   n = {len(sub)}")

    frac = (cm["population"] == "2G (P2)").mean()
    ax.set_xlabel(r"$\Delta_{F275W,\,F814W}$  (verticalised pseudo-colour)")
    ax.set_ylabel(r"$\Delta_{C\,F275W,\,F336W,\,F438W}$")
    ax.set_title(f"Multiple stellar populations in {NAME}")
    ax.legend(frameon=False, loc="upper left", fontsize=11)
    ax.text(0.98, 0.04,
            f"Hubble / HUGS photometry\nunsupervised K-means  ·  "
            f"2G fraction \u2248 {frac:.0%}",
            transform=ax.transAxes, ha="right", va="bottom",
            fontsize=9, color="0.35")
    ax.grid(True, alpha=0.15)
    fig.tight_layout()
    p = OUT / "ngc6752_chromosome_map_kmeans.png"
    fig.savefig(p)
    plt.close(fig)
    return p


def fig_cmd(df: pd.DataFrame, cm: pd.DataFrame, c: mp.Columns) -> Path:
    """
    Colour-magnitude diagram (the observational H-R diagram).

    Background: all good-photometry member stars (grey). Foreground: the RGB
    stars, coloured by the population K-means assigned to them.
    """
    OUT.mkdir(parents=True, exist_ok=True)

    m606 = pd.to_numeric(df["F606W_mag"], errors="coerce")
    m814 = pd.to_numeric(df[c.f814w], errors="coerce")

    ok = (m606 > -50) & (m606 < 90) & (m814 > -50) & (m814 < 90)
    if c.membership is not None:
        memb = pd.to_numeric(df[c.membership], errors="coerce")
        scale = 100.0 if memb.max() > 1.5 else 1.0
        ok &= memb >= 0.9 * scale

    colour = (m606 - m814)[ok]
    mag = m814[ok]

    fig, ax = plt.subplots(figsize=(6.6, 7.6))
    ax.scatter(colour, mag, s=3, color="0.78", alpha=0.5,
               edgecolors="none", rasterized=True, label="cluster members")

    # Overlay the classified RGB stars in the same F606W-F814W plane.
    rgb_idx = cm.index
    rcol = (pd.to_numeric(df.loc[rgb_idx, "F606W_mag"], errors="coerce")
            - pd.to_numeric(df.loc[rgb_idx, c.f814w], errors="coerce"))
    rmag = pd.to_numeric(df.loc[rgb_idx, c.f814w], errors="coerce")
    for label, color in (("1G (P1)", C_1G), ("2G (P2)", C_2G)):
        m = (cm["population"] == label).values
        ax.scatter(rcol[m], rmag[m], s=18, alpha=0.85,
                   edgecolors="white", linewidths=0.3, color=color,
                   label=f"RGB · {label}")

    ax.invert_yaxis()  # brighter stars at top
    ax.set_xlabel(r"$F606W - F814W$  (colour)")
    ax.set_ylabel(r"$F814W$  (apparent magnitude)")
    ax.set_title(f"Colour-magnitude diagram of {NAME}")
    ax.legend(frameon=False, loc="upper right", fontsize=10)
    ax.text(0.02, 0.02, "Hubble / HUGS photometry",
            transform=ax.transAxes, ha="left", va="bottom",
            fontsize=9, color="0.35")
    ax.grid(True, alpha=0.15)
    fig.tight_layout()
    p = OUT / "ngc6752_cmd_hr_diagram.png"
    fig.savefig(p)
    plt.close(fig)
    return p


def fig_combined(df: pd.DataFrame, cm: pd.DataFrame, c: mp.Columns) -> Path:
    """Two-panel figure: CMD (left) + chromosome map (right), for one post."""
    OUT.mkdir(parents=True, exist_ok=True)
    fig, (axL, axR) = plt.subplots(1, 2, figsize=(13.4, 7.0),
                                   gridspec_kw={"width_ratios": [1.0, 1.18]})

    # ---- LEFT: colour-magnitude (H-R) diagram --------------------------------
    m606 = pd.to_numeric(df["F606W_mag"], errors="coerce")
    m814 = pd.to_numeric(df[c.f814w], errors="coerce")
    ok = (m606 > -50) & (m606 < 90) & (m814 > -50) & (m814 < 90)
    if c.membership is not None:
        memb = pd.to_numeric(df[c.membership], errors="coerce")
        scale = 100.0 if memb.max() > 1.5 else 1.0
        ok &= memb >= 0.9 * scale
    axL.scatter((m606 - m814)[ok], m814[ok], s=3, color="0.78", alpha=0.5,
                edgecolors="none", rasterized=True, label="cluster members")

    rgb_idx = cm.index
    rcol = (pd.to_numeric(df.loc[rgb_idx, "F606W_mag"], errors="coerce")
            - pd.to_numeric(df.loc[rgb_idx, c.f814w], errors="coerce"))
    rmag = pd.to_numeric(df.loc[rgb_idx, c.f814w], errors="coerce")
    for label, color in (("1G (P1)", C_1G), ("2G (P2)", C_2G)):
        m = (cm["population"] == label).values
        axL.scatter(rcol[m], rmag[m], s=16, alpha=0.85, edgecolors="white",
                    linewidths=0.3, color=color, label=f"RGB · {label}")
    axL.invert_yaxis()
    axL.set_xlabel(r"$F606W - F814W$  (colour)")
    axL.set_ylabel(r"$F814W$  (apparent magnitude)")
    axL.set_title("Ordinary view: one simple cluster")
    axL.legend(frameon=False, loc="upper right", fontsize=9.5)
    axL.grid(True, alpha=0.15)

    # ---- RIGHT: chromosome map ----------------------------------------------
    for label, color in (("1G (P1)", C_1G), ("2G (P2)", C_2G)):
        sub = cm[cm["population"] == label]
        axR.scatter(sub["dx"], sub["dy"], s=24, alpha=0.78, edgecolors="white",
                    linewidths=0.4, color=color, label=f"{label}   n = {len(sub)}")
    frac = (cm["population"] == "2G (P2)").mean()
    axR.set_xlabel(r"$\Delta_{F275W,\,F814W}$  (verticalised pseudo-colour)")
    axR.set_ylabel(r"$\Delta_{C\,F275W,\,F336W,\,F438W}$")
    axR.set_title("Hubble UV view: two hidden generations")
    axR.legend(frameon=False, loc="upper left", fontsize=10)
    axR.text(0.98, 0.04, f"unsupervised K-means  ·  2G fraction \u2248 {frac:.0%}",
             transform=axR.transAxes, ha="right", va="bottom",
             fontsize=9, color="0.35")
    axR.grid(True, alpha=0.15)

    fig.suptitle(f"Multiple stellar populations in {NAME}  —  Hubble / HUGS photometry",
                 fontsize=15, fontweight="bold", y=1.00)
    fig.tight_layout()
    p = OUT / "ngc6752_combined.png"
    fig.savefig(p)
    plt.close(fig)
    return p


def main() -> int:
    df, cm, cols = build()
    p1 = fig_chromosome(cm)
    p2 = fig_cmd(df, cm, cols)
    p3 = fig_combined(df, cm, cols)
    print("wrote:", p1)
    print("wrote:", p2)
    print("wrote:", p3)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())