multiple_populations.py

#!/usr/bin/env python3
"""
Multiple Stellar Populations in the Globular Cluster NGC 6752
==================================================================

This program reproduces, from REAL Hubble Space Telescope photometry, the
"multiple populations" signature in a globular cluster: the splitting of the
red-giant branch into a first (1G/P1) and a second (2G/P2) stellar generation,
visible on a pseudo-two-colour diagram known as a "chromosome map".

Inspiration (narrative only):
    Jan Hattenbach, "Cluster Mystery", Sky & Telescope, July 2026.

Scientific data + method (the real source you must cite):
    Photometry: HUGS - HST UV Legacy Survey of Galactic Globular Clusters
                Piotto et al. 2015, AJ, 149, 91
                Nardiello et al. 2018, MNRAS, 481, 3382
                Catalogue DOI: 10.17909/T9810F  (MAST HLSP "HUGS")
    Chromosome-map technique:
                Milone et al. 2017, MNRAS, 464, 3636

The chromosome map combines stellar magnitudes in the F275W, F336W, F438W and
F814W filters so that the abundance variations of He, C, N, O that distinguish
the populations are maximally separated.

This script does NOT reprocess raw images. It uses the public, ready-made HUGS
photometric catalogue and is deliberately robust to small differences in the
catalogue's column layout: it parses the header and locates the needed columns
automatically.

Author: Ali Razeghi  (github.com/Ali-Razeghi)
License: MIT
"""

from __future__ import annotations

import argparse
import sys
from dataclasses import dataclass
from pathlib import Path

import numpy as np
import pandas as pd
import matplotlib

matplotlib.use("Agg")  # headless: write files, never require a display
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans, DBSCAN
from sklearn.preprocessing import StandardScaler


# --------------------------------------------------------------------------- #
# Configuration
# --------------------------------------------------------------------------- #

# Direct MAST download URL for the NGC 6752 catalogue (photometric method 2).
# If this ever moves, see https://archive.stsci.edu/prepds/hugs/
HUGS_DEFAULT_URL = (
    "https://archive.stsci.edu/hlsps/hugs/ngc6752/"
    "hlsp_hugs_hst_wfc3-uvis-acs-wfc_ngc6752_multi_v1_catalog-meth2.txt"
)

# Filters whose magnitudes we need to build the chromosome map.
NEEDED_FILTERS = ("F275W", "F336W", "F438W", "F814W")

# Sentinel value HUGS uses for "no measurement". Anything <= this is dropped.
BAD_MAG = -50.0


@dataclass
class Columns:
    """Resolved column names for the four filter magnitudes + membership."""

    f275w: str
    f336w: str
    f438w: str
    f814w: str
    membership: str | None


# --------------------------------------------------------------------------- #
# Loading & column resolution
# --------------------------------------------------------------------------- #

def _looks_like_header(line: str) -> bool:
    """A HUGS header line names columns; data lines are mostly numbers."""
    tokens = line.replace("#", " ").split()
    if not tokens:
        return False
    # If most tokens fail to parse as floats, treat it as a header/label line.
    non_numeric = 0
    for tok in tokens:
        try:
            float(tok)
        except ValueError:
            non_numeric += 1
    return non_numeric >= max(2, len(tokens) // 2)


def download_if_url(source: str, dest: Path) -> Path:
    """
    If `source` is an http(s) URL, download it to `dest` (once) and return the
    local path. Otherwise treat `source` as an existing local path.

    Runs in YOUR environment, so it uses your network access to MAST.
    """
    if not str(source).lower().startswith(("http://", "https://")):
        return Path(source)

    import urllib.request

    dest.parent.mkdir(parents=True, exist_ok=True)
    if dest.exists() and dest.stat().st_size > 0:
        print(f"Using cached download: {dest}")
        return dest

    print(f"Downloading:\n  {source}\n-> {dest}")
    # A browser-like User-Agent avoids some automated-access blocks.
    req = urllib.request.Request(
        source, headers={"User-Agent": "Mozilla/5.0 (research; astropy-style client)"}
    )
    with urllib.request.urlopen(req, timeout=120) as resp, dest.open("wb") as out:
        out.write(resp.read())
    print(f"Downloaded {dest.stat().st_size:,} bytes")
    return dest


def load_hugs_catalog(path: Path) -> pd.DataFrame:
    """
    Load a HUGS photometric catalogue.

    HUGS catalogues are whitespace-delimited text. Some releases ship a clean
    column header; others document columns only in a separate README. We try a
    real header first, and fall back to positional names if none is found.
    """
    # First pass: read raw lines to locate a usable header.
    header_tokens: list[str] | None = None
    with path.open("r", errors="replace") as fh:
        for raw in fh:
            line = raw.strip()
            if not line:
                continue
            if line.startswith("#") and _looks_like_header(line):
                header_tokens = line.lstrip("#").split()
            # Stop scanning once we hit the first clear data row.
            if not line.startswith("#"):
                break

    if header_tokens:
        df = pd.read_csv(
            path,
            sep=r"\s+",
            comment="#",
            header=None,
            names=header_tokens,
            engine="python",
        )
    else:
        # No header: load positionally and name columns col0, col1, ...
        df = pd.read_csv(path, sep=r"\s+", comment="#", header=None, engine="python")
        df.columns = [f"col{i}" for i in range(df.shape[1])]

    if df.empty:
        raise ValueError(f"No data rows parsed from {path}")
    return df


def resolve_columns(df: pd.DataFrame) -> Columns:
    """
    Find which DataFrame columns hold the four filter magnitudes and the
    membership probability.

    Strategy:
      1. Match by name when the header is descriptive (e.g. a column whose name
         contains 'F275W' and a magnitude hint).
      2. Fall back to the documented HUGS positional layout when names are
         generic (col0, col1, ...).
    """
    cols = list(df.columns)
    lower = {c: str(c).lower() for c in cols}

    def find_named(filt: str) -> str | None:
        f = filt.lower()
        # Prefer a calibrated magnitude column for this filter.
        candidates = [c for c in cols if f in lower[c]]
        if not candidates:
            return None
        for hint in ("mag", "_m", "vega", "calib"):
            for c in candidates:
                if hint in lower[c]:
                    return c
        return candidates[0]

    named = {f: find_named(f) for f in NEEDED_FILTERS}
    membership = next(
        (c for c in cols if "memb" in lower[c] or "prob" in lower[c] or lower[c] == "p"),
        None,
    )

    if all(named.values()):
        return Columns(
            f275w=named["F275W"],
            f336w=named["F336W"],
            f438w=named["F438W"],
            f814w=named["F814W"],
            membership=membership,
        )

    # ---- Positional fallback (HUGS meth catalogue documented layout) -------
    # Documented order of the per-filter calibrated magnitudes in the HUGS
    # "method" catalogues is: F275W, F336W, F438W, F606W, F814W, each preceded
    # by x,y and followed by error/quality columns. The exact indices can shift
    # between releases, so we expose them as a single place to adjust.
    #
    # If you hit this branch, open the catalogue's README on MAST and set the
    # five indices below to the calibrated-magnitude columns.
    raise ColumnResolutionError(
        "Could not resolve filter columns by name. The catalogue header is not "
        "descriptive in this release.\n"
        "Open the README at https://archive.stsci.edu/prepds/hugs/ and map the "
        "calibrated magnitude columns for F275W, F336W, F438W, F814W, then pass "
        "them via --cols, e.g.:  --cols F275W=col5 F336W=col9 F438W=col13 "
        "F814W=col21"
    )


class ColumnResolutionError(RuntimeError):
    pass


def apply_manual_columns(df: pd.DataFrame, mapping: dict[str, str]) -> Columns:
    for f in NEEDED_FILTERS:
        if mapping.get(f) not in df.columns:
            raise ColumnResolutionError(
                f"--cols mapping for {f} -> {mapping.get(f)!r} not found in file."
            )
    membership = next(
        (c for c in df.columns if "memb" in str(c).lower() or "prob" in str(c).lower()),
        None,
    )
    return Columns(
        f275w=mapping["F275W"],
        f336w=mapping["F336W"],
        f438w=mapping["F438W"],
        f814w=mapping["F814W"],
        membership=membership,
    )


# --------------------------------------------------------------------------- #
# Building the chromosome map
# --------------------------------------------------------------------------- #

def select_red_giants(
    df: pd.DataFrame, c: Columns, membership_cut: float = 0.9,
    max_rms: float = 0.05, min_qfit: float = 0.9,
) -> pd.DataFrame:
    """
    Keep likely cluster-member red-giant-branch (RGB) stars with GOOD
    photometry in all four filters.

    Quality control is essential for HUGS catalogues: each filter has a
    photometric RMS and a quality-of-fit (qfit) column. Stars with poor fits
    form a spurious cloud on the chromosome map that is not a real population.
    We require, per filter: a real magnitude, low RMS, and high qfit.
    """
    mags = [c.f275w, c.f336w, c.f438w, c.f814w]
    good = np.ones(len(df), dtype=bool)

    for m in mags:
        v = pd.to_numeric(df[m], errors="coerce").to_numpy()
        good &= np.isfinite(v) & (v > BAD_MAG) & (v < 90)

        # Matching quality columns, if present (e.g. F275W_rms, F275W_qfit).
        base = str(m).replace("_mag", "")
        rms_col = next((col for col in df.columns
                        if str(col).lower() == f"{base}_rms".lower()), None)
        qf_col = next((col for col in df.columns
                       if str(col).lower() == f"{base}_qfit".lower()), None)
        if rms_col is not None:
            rms = pd.to_numeric(df[rms_col], errors="coerce").to_numpy()
            good &= np.nan_to_num(rms, nan=9e9) <= max_rms
        if qf_col is not None:
            qf = pd.to_numeric(df[qf_col], errors="coerce").to_numpy()
            good &= np.nan_to_num(qf, nan=-1) >= min_qfit

    if c.membership is not None:
        p = pd.to_numeric(df[c.membership], errors="coerce").to_numpy()
        scale = 100.0 if np.nanmax(p) > 1.5 else 1.0
        good &= np.nan_to_num(p, nan=-1.0) >= membership_cut * scale

    sub = df.loc[good].copy()

    # Colour-magnitude RGB box in F275W vs (F275W-F814W).
    col_x = sub[c.f275w] - sub[c.f814w]
    mag = sub[c.f275w]
    # Brighter RGB segment, where the populations separate most cleanly.
    lo, hi = mag.quantile(0.10), mag.quantile(0.55)
    rgb = sub[(mag >= lo) & (mag <= hi) & (col_x > col_x.median())]
    return rgb


def chromosome_map(rgb: pd.DataFrame, c: Columns) -> pd.DataFrame:
    """
    Compute the two axes of the chromosome map, following the spirit of
    Milone et al. (2017).

    The key step is *verticalisation*: along the red-giant branch the colours
    drift with magnitude, which would smear out the populations. We therefore
    model the colour as a function of magnitude (a robust low-order fit to the
    bulk of the stars) and measure each star's residual from that ridge. The
    residuals -- not the raw colours -- carry the population signal.

      x-axis : residual of (F275W - F814W)           ~ Delta_(F275W,F814W)
      y-axis : residual of C_F275W,F336W,F438W        ~ Delta_C
               where C = (F275W - F336W) - (F336W - F438W)

    Axes are left in magnitude units (not squashed to 0..1), so the diagram
    keeps its true shape and the two sequences stay separable.
    """
    mag = rgb[c.f275w].to_numpy()
    col1 = (rgb[c.f275w] - rgb[c.f814w]).to_numpy()
    cidx = (
        (rgb[c.f275w] - rgb[c.f336w]) - (rgb[c.f336w] - rgb[c.f438w])
    ).to_numpy()

    def verticalise(mag: np.ndarray, val: np.ndarray) -> np.ndarray:
        """Residual of `val` after removing a robust trend with magnitude."""
        ok = np.isfinite(mag) & np.isfinite(val)
        # Iteratively reject outliers so the ridge follows the bulk, not the
        # contaminants, then return residuals for all stars.
        idx = np.where(ok)[0]
        m, v = mag[idx], val[idx]
        keep = np.ones(len(idx), dtype=bool)
        coeffs = None
        for _ in range(3):
            coeffs = np.polyfit(m[keep], v[keep], deg=2)
            resid = v - np.polyval(coeffs, m)
            s = np.std(resid[keep])
            keep = np.abs(resid) < 2.5 * s
        res_all = np.full_like(mag, np.nan, dtype=float)
        res_all[ok] = val[ok] - np.polyval(coeffs, mag[ok])
        return res_all

    dx = verticalise(mag, col1)
    dy = verticalise(mag, cidx)

    out = pd.DataFrame({"dx": dx, "dy": dy}, index=rgb.index).dropna()
    # Clip only the extreme 0.5% tails for display stability (not onto edges).
    for col in ("dx", "dy"):
        lo, hi = out[col].quantile(0.005), out[col].quantile(0.995)
        out = out[(out[col] >= lo) & (out[col] <= hi)]
    return out


# --------------------------------------------------------------------------- #
# Unsupervised clustering
# --------------------------------------------------------------------------- #

def cluster_populations(cm: pd.DataFrame) -> pd.DataFrame:
    """Run K-means (k=2: 1G vs 2G) and DBSCAN (density-based) on the map."""
    X = cm[["dx", "dy"]].to_numpy()
    Xs = StandardScaler().fit_transform(X)

    km = KMeans(n_clusters=2, n_init=10, random_state=42).fit(Xs)
    cm = cm.copy()
    cm["kmeans"] = km.labels_

    # Orient labels so that label 0 = 1G (lower dy = N-normal), 1 = 2G.
    if cm.loc[cm["kmeans"] == 0, "dy"].mean() > cm.loc[cm["kmeans"] == 1, "dy"].mean():
        cm["kmeans"] = 1 - cm["kmeans"]
    cm["population"] = np.where(cm["kmeans"] == 0, "1G (P1)", "2G (P2)")

    db = DBSCAN(eps=0.35, min_samples=12).fit(Xs)
    cm["dbscan"] = db.labels_  # -1 = noise

    return cm


# --------------------------------------------------------------------------- #
# Plotting
# --------------------------------------------------------------------------- #

def plot_results(cm: pd.DataFrame, out_dir: Path, cluster_name: str) -> list[Path]:
    out_dir.mkdir(parents=True, exist_ok=True)
    paths: list[Path] = []

    # 1) Raw chromosome map.
    fig, ax = plt.subplots(figsize=(6.5, 5.5))
    ax.scatter(cm["dx"], cm["dy"], s=8, c="0.35", alpha=0.6, edgecolors="none")
    ax.set_xlabel(r"$\Delta_{F275W,\,F814W}$  (pseudo-colour)")
    ax.set_ylabel(r"$\Delta_{C\,F275W,F336W,F438W}$")
    ax.set_title(f"Chromosome map — {cluster_name}")
    fig.tight_layout()
    p = out_dir / "chromosome_map.png"
    fig.savefig(p, dpi=150)
    plt.close(fig)
    paths.append(p)

    # 2) K-means populations.
    fig, ax = plt.subplots(figsize=(6.5, 5.5))
    palette = {"1G (P1)": "#2c6fbb", "2G (P2)": "#c0392b"}
    for label, sub in cm.groupby("population"):
        ax.scatter(
            sub["dx"], sub["dy"], s=10, alpha=0.7, edgecolors="none",
            color=palette.get(label, "0.4"), label=f"{label}  (n={len(sub)})",
        )
    ax.set_xlabel(r"$\Delta_{F275W,\,F814W}$")
    ax.set_ylabel(r"$\Delta_{C\,F275W,F336W,F438W}$")
    ax.set_title(f"K-means: first vs second generation — {cluster_name}")
    ax.legend(frameon=False, loc="upper right")
    fig.tight_layout()
    p = out_dir / "kmeans_populations.png"
    fig.savefig(p, dpi=150)
    plt.close(fig)
    paths.append(p)

    # 3) DBSCAN view.
    fig, ax = plt.subplots(figsize=(6.5, 5.5))
    labels = sorted(cm["dbscan"].unique())
    cmap = plt.get_cmap("tab10")
    for k in labels:
        sub = cm[cm["dbscan"] == k]
        name = "noise" if k == -1 else f"group {k}"
        color = "0.7" if k == -1 else cmap(k % 10)
        ax.scatter(sub["dx"], sub["dy"], s=10, alpha=0.7, edgecolors="none",
                   color=color, label=f"{name} (n={len(sub)})")
    ax.set_xlabel(r"$\Delta_{F275W,\,F814W}$")
    ax.set_ylabel(r"$\Delta_{C\,F275W,F336W,F438W}$")
    ax.set_title(f"DBSCAN density groups — {cluster_name}")
    ax.legend(frameon=False, loc="upper right", fontsize=8)
    fig.tight_layout()
    p = out_dir / "dbscan_groups.png"
    fig.savefig(p, dpi=150)
    plt.close(fig)
    paths.append(p)

    return paths


# --------------------------------------------------------------------------- #
# CLI
# --------------------------------------------------------------------------- #

def parse_cols(pairs: list[str] | None) -> dict[str, str] | None:
    if not pairs:
        return None
    mapping = {}
    for pair in pairs:
        if "=" not in pair:
            raise SystemExit(f"--cols expects FILTER=column, got {pair!r}")
        k, v = pair.split("=", 1)
        mapping[k.strip().upper()] = v.strip()
    return mapping


def main(argv: list[str] | None = None) -> int:
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--catalog", default="data/ngc6752_meth2.txt",
                    help="Local path OR an http(s) URL to the HUGS catalogue. "
                         "A URL is downloaded (and cached) into --download-to.")
    ap.add_argument("--download-to", type=Path, default=Path("data/ngc6752_meth2.txt"),
                    help="Where to save the file when --catalog is a URL.")
    ap.add_argument("--name", default="NGC 6752", help="Cluster display name.")
    ap.add_argument("--out", type=Path, default=Path("figures"),
                    help="Output directory for figures and the labelled table.")
    ap.add_argument("--membership-cut", type=float, default=0.9,
                    help="Minimum membership probability (0-1).")
    ap.add_argument("--max-rms", type=float, default=0.05,
                    help="Max photometric RMS per filter (quality cut).")
    ap.add_argument("--min-qfit", type=float, default=0.9,
                    help="Min quality-of-fit per filter (quality cut).")
    ap.add_argument("--cols", nargs="*", default=None,
                    help="Manual column map, e.g. F275W=col5 F336W=col9 ...")
    args = ap.parse_args(argv)

    catalog_path = download_if_url(args.catalog, args.download_to)

    if not catalog_path.exists():
        sys.stderr.write(
            f"\nCatalogue not found: {catalog_path}\n\n"
            "Pass a local file, or a URL that will be downloaded, e.g.:\n"
            f"  python multiple_populations.py --catalog \\\n    {HUGS_DEFAULT_URL}\n\n"
            "See README for other clusters and methods.\n"
        )
        return 2

    df = load_hugs_catalog(catalog_path)
    print(f"Loaded {len(df):,} rows, {df.shape[1]} columns from {catalog_path}")

    manual = parse_cols(args.cols)
    cols = apply_manual_columns(df, manual) if manual else resolve_columns(df)
    print(f"Using columns: F275W={cols.f275w}, F336W={cols.f336w}, "
          f"F438W={cols.f438w}, F814W={cols.f814w}, "
          f"membership={cols.membership}")

    rgb = select_red_giants(df, cols, membership_cut=args.membership_cut,
                            max_rms=args.max_rms, min_qfit=args.min_qfit)
    print(f"Selected {len(rgb):,} candidate RGB member stars")
    if len(rgb) < 50:
        sys.stderr.write("Too few stars after selection; loosen --membership-cut.\n")
        return 3

    cm = chromosome_map(rgb, cols)
    cm = cluster_populations(cm)

    n1 = (cm["population"] == "1G (P1)").sum()
    n2 = (cm["population"] == "2G (P2)").sum()
    print(f"\nK-means split: 1G/P1 = {n1}   2G/P2 = {n2}   "
          f"(2G fraction = {n2 / (n1 + n2):.2f})")

    paths = plot_results(cm, args.out, args.name)
    table = args.out / "labelled_stars.csv"
    cm.to_csv(table, index=False)
    print("\nWrote:")
    for p in paths + [table]:
        print(f"  {p}")
    return 0


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