example.py

"""
ClusterForge - worked example.

Runs the full pipeline (excluding mass completeness, which is a separate
tool - see the README) on a small synthetic cluster catalog, so this script
works immediately after cloning the repo with `python make_mock_catalog.py`
already run once (see below).

To use your own data instead of the bundled mock catalog: point `data_path`
at your own CSV and update the column-name arguments to `cluster_data(...)`
to match your file's actual column names. The columns listed on the right
of each keyword below are the ones `cluster_data` expects to find.
"""
#%% Setup
import pandas as pd
import numpy as np
import os

from packages import cluster_main
from packages.kinematics import get_kinematics_analysis
from packages.cmd_mass_analysis import get_cmd_mass_analysis

for d in ['Figs', 'output']:
    os.makedirs(d, exist_ok=True)

#%% Load the catalog and wrap it in a cluster_data object
# Regenerate the bundled mock catalog with: python make_mock_catalog.py
data_path = 'input/mock_cluster.csv'
if not os.path.exists(data_path):
    raise FileNotFoundError(
        f"{data_path} not found - run `python make_mock_catalog.py` once to generate it, "
        "or point data_path at your own catalog."
    )
data = pd.read_csv(data_path)

my_cluster = cluster_main.cluster_data(
    df=data,
    ra='ra', dec='dec',
    pmra='pmra', pmra_e='pmra_error',
    pmdec='pmdec', pmdec_e='pmdec_error',
    plx='PLXcorr', plx_e='parallax_error',
    rv='RVbest', rv_e='e_RVbest',
    BP='Bpmag', BP_e='Bperr',
    RP='Rpmag', RP_e='Rperr',
    Gmag='Gmag', Gmag_e='Gperr',
)

#%% Quick look: median position and proper motion
cluster_main.centroid(my_cluster)

#%% Structural characterization: radial density profile (King + EFF models)
cluster_main.get_density_profiles(my_cluster.ra, my_cluster.dec, plot=True)

#%% Velocity distributions: RV, V_RA, V_DEC (Gaussian fits)
cluster_main.get_velocity_analysis(my_cluster, plot=True)

#%% Age and extinction from isochrone fitting
# Single fit on the full sample (fast) - swap method='repeat' for bootstrap
# age/A_V uncertainties (slower: n_iter refits).
#
# NOTE: iso_test contains multiple metallicities per (log_age, A_V) pair,
# but load_isochrones() currently keys its grid only on (log_age, A_V) - see
# the README's "Known limitations" section. best_iso/best_params below are
# internally consistent (returned directly from the fit), but reconstructing
# the file path for the next step can't know which metallicity that was, so
# we just take whichever file matches - fine for this demo, but worth fixing
# in estimate_age_isochrones.py before relying on iso_test for real science.
import glob
color, mag, best_iso, best_params = cluster_main.get_age_estimation(
    my_cluster, iso_folder='iso_test', method='single', binning=10
)
best_log_age, best_av = best_params
iso_path = glob.glob(f'input/isochrones/iso_test/iso_{best_log_age:.4f}_*_{best_av:.2f}.csv')[0]

#%% Kinematics: expansion pattern and kinematic (expansion) age
# med_rv should be the cluster's median/fitted RV - here we just use the
# catalog's own median as a stand-in; use the velocity_analysis RV fit's
# result in a real analysis.
med_rv = float(np.median(my_cluster.rv))
df_pos_vel, kin_results = get_kinematics_analysis(my_cluster, med_rv=med_rv, plot=True)

#%% Mass estimation from the CMD (observed stars only - no completeness correction)
color, mag, mass_stars, params, df_mass = get_cmd_mass_analysis(my_cluster, iso=iso_path, plot=True)

#%% Physical parameters (size, virial mass, dynamical state, ...)
# Needs the posterior CSVs written by the density-profile and velocity-fit
# steps above.
king = 'output/King_Posteriors.csv'
eff = 'output/EFF_Posteriors.csv'
vel = 'output/Velocity_Posteriors.csv'
cluster_main.get_cluster_parameters(my_cluster, iso_path, king, eff, vel, mass='None')

#%% Re-running with an externally-supplied mass
# If you've separately run the mass-completeness/mock-catalog tool
# (https://github.com/<you>/ClusterForge-MassCompleteness) and it reported a
# completeness-corrected total mass, feed it back in here instead of
# re-estimating from the observed CMD alone:
#
# cluster_main.get_cluster_parameters(my_cluster, iso_path, king, eff, vel, mass=(total_mass, total_mass_error))