"""Shared tools for the head-start study. * paths and physical constants * 1-D vertical Galactic potential and orbit integrals (column-density weights) * EM (Richardson-Lucy type) deconvolution of an age distribution from stars with heteroscedastic, asymmetric age errors * summary statistics of a head-start distribution * compact loader for MIST isochrones """ import os import io import tarfile import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) STUDY = os.path.dirname(HERE) DATA = os.path.join(STUDY, "data") FIGS = os.path.join(STUDY, "figures") RAW = os.environ.get("HEADSTART_RAW", "fermi_raw/headstart/") for _d in (DATA, FIGS): os.makedirs(_d, exist_ok=True) # ----------------------------------------------------------------- constants T_SUN = 4.567 # Gyr, age of the Solar System (CAI age; Connelly+ 2012) T_MAX = 13.8 # Gyr, age of the Universe (Planck 2018: 13.787) G_PC = 4.30091e-3 # pc (km/s)^2 / Msun KMS_TO_PC_PER_MYR = 1.0227 Z_SUN = 20.8 # pc, Sun's height above the plane (Bennett & Bovy 2019) W_SUN = 7.25 # km/s, solar vertical motion wrt LSR (Schoenrich+ 2010) U_SUN, V_SUN = 11.1, 12.24 IMF_SLOPE = 2.3 # Kroupa (2001) slope above 0.5 Msun, dN/dM ~ M^-2.3 # age grid used everywhere (Gyr); bins of 0.25 Gyr from 0 to 14 AGE_EDGES = np.arange(0.0, 14.0001, 0.25) AGE_MID = 0.5 * (AGE_EDGES[1:] + AGE_EDGES[:-1]) # ------------------------------------------------------- vertical potential class VerticalPotential: """One-dimensional vertical potential Phi(z) at the solar radius and the orbit integrals needed to turn a local sample into column densities. Units: z in pc, velocities in km/s, Phi in (km/s)^2, times in pc/(km/s). """ def __init__(self, phi_func, name=""): self._phi = phi_func self.name = name self._build() def phi(self, z): return self._phi(np.abs(np.asarray(z, float))) def _build(self, n=600, ntheta=400): # grid in zmax (pc); E = Phi(zmax) self.zmax_grid = np.logspace(-0.5, np.log10(3.0e4), n) self.E_grid = self.phi(self.zmax_grid) th = (np.arange(ntheta) + 0.5) * (np.pi / 2) / ntheta # midpoints dth = (np.pi / 2) / ntheta z = self.zmax_grid[:, None] * np.sin(th)[None, :] v = np.sqrt(np.maximum(2 * (self.E_grid[:, None] - self.phi(z)), 1e-30)) integrand = self.zmax_grid[:, None] * np.cos(th)[None, :] / v # period (pc / (km/s)) and action (pc km/s) self.T_grid = 4 * integrand.sum(axis=1) * dth self.J_grid = (2 / np.pi) * (v * self.zmax_grid[:, None] * np.cos(th)[None, :]).sum(axis=1) * dth # -- interpolators (all in log space, monotonic) -- def _clipE(self, E): return np.clip(E, self.E_grid[0], self.E_grid[-1]) def zmax_of_E(self, E): return np.exp(np.interp(np.log(self._clipE(E)), np.log(self.E_grid), np.log(self.zmax_grid))) def T_of_E(self, E): return np.exp(np.interp(np.log(self._clipE(E)), np.log(self.E_grid), np.log(self.T_grid))) def E_of_J(self, J): J = np.clip(J, self.J_grid[0], self.J_grid[-1]) return np.exp(np.interp(np.log(J), np.log(self.J_grid), np.log(self.E_grid))) def energy(self, z, w): return 0.5 * np.asarray(w, float) ** 2 + self.phi(z) def frac_weighted(self, E, gfunc, ntheta=400, chunk=20000): """Time average over the vertical orbit of energy E. gfunc(z) must accept arrays of z in pc. Works for array E. """ E = np.atleast_1d(np.asarray(E, float)) out = np.empty(len(E)) th = -np.pi / 2 + (np.arange(ntheta) + 0.5) * np.pi / ntheta sin, cos = np.sin(th)[None, :], np.cos(th)[None, :] for a in range(0, len(E), chunk): zmax = self.zmax_of_E(E[a:a + chunk]) Ez = self.phi(zmax) # snap E to the interpolated zmax z = zmax[:, None] * sin v = np.sqrt(np.maximum(2 * (Ez[:, None] - self.phi(z)), 1e-30)) dt = zmax[:, None] * cos / v out[a:a + chunk] = (gfunc(z) * dt).sum(axis=1) / dt.sum(axis=1) return out def sphere_fraction(self, E, R, zsun=Z_SUN): """Fraction of a vertical column of orbits with energy E (uniform in the plane) that lies inside a sphere of radius R centred on the Sun, relative to a cylinder of cross-section pi R^2: f = < max(0, 1 - (z - zsun)^2 / R^2) >_orbit A star in the sphere therefore stands for 1/(pi R^2 f) stars per unit area of the Galactic plane. """ return self.frac_weighted( E, lambda z: np.clip(1 - ((z - zsun) / R) ** 2, 0, None)) def slab_fraction(self, E, z0): """Fraction of time an orbit of energy E spends at |z| < z0.""" return self.frac_weighted(E, lambda z: (np.abs(z) < z0).astype(float)) def kg_phi(sigma0, D, rho_h): """Kuijken & Gilmore (1989) form Phi(z) = 2 pi G [ Sigma0 (sqrt(z^2 + D^2) - D) + rho_h z^2 ] Sigma0 [Msun/pc^2] disc surface density, D [pc] disc scale, rho_h [Msun/pc^3] effective halo density.""" def phi(z): return 2 * np.pi * G_PC * (sigma0 * (np.sqrt(z * z + D * D) - D) + rho_h * z * z) return phi def mw14_phi(R0=8178.0, v0=220.0): """Vertical cut Phi(R0, z) - Phi(R0, 0) of galpy's MWPotential2014 (Bovy 2015): Miyamoto-Nagai disc (a = 3/8 R0, b = 0.28/8 R0, 60% of v_c^2 at R0), NFW halo (a = 2 R0, 35%), and a bulge (5%) treated as a point mass (its cut-off radius, 1.9 kpc, is far inside R0). This is the potential in which Xiang & Rix (2022) computed their actions.""" a, b, ah = 3.0 / 8.0 * R0, 0.28 / 8.0 * R0, 2.0 * R0 GMd = 0.60 * v0 ** 2 * (R0 ** 2 + (a + b) ** 2) ** 1.5 / R0 ** 2 GMh = 0.35 * v0 ** 2 * R0 / (np.log(1 + R0 / ah) - (R0 / ah) / (1 + R0 / ah)) GMb = 0.05 * v0 ** 2 * R0 def raw(z): r = np.sqrt(R0 ** 2 + z * z) return (-GMd / np.sqrt(R0 ** 2 + (a + np.sqrt(z * z + b * b)) ** 2) - GMh * np.log(1 + r / ah) / r - GMb / r) p0 = raw(np.array(0.0)) return lambda z: raw(z) - p0 # Fiducial values reproduce rho_0 = Sigma0/(2D) + rho_h = 0.096 Msun/pc^3 # (McKee+ 2015: 0.097) and Sigma(|z|<1.1 kpc) = 72 Msun/pc^2 (cf. 68-74: # Kuijken & Gilmore 1991, Bovy & Rix 2013, Zhang+ 2013). POTENTIALS = { "fiducial": lambda: kg_phi(47.0, 280.0, 0.012), "KG89": lambda: kg_phi(46.0, 180.0, 0.010), # Kuijken & Gilmore 1989 "light": lambda: kg_phi(38.0, 300.0, 0.008), # low-mass disc "heavy": lambda: kg_phi(56.0, 260.0, 0.016), # high-mass disc "MW14": lambda: mw14_phi(), # galpy MWPotential2014 } _POT_CACHE = {} def get_potential(name="fiducial"): if name not in _POT_CACHE: _POT_CACHE[name] = VerticalPotential(POTENTIALS[name](), name=name) return _POT_CACHE[name] # ---------------------------------------------------------------- kinematics def galactic_uvw(l_deg, b_deg, plx_mas, pmra, pmdec, rv, ra_deg, dec_deg): """Heliocentric Galactic (U, V, W) in km/s (U towards the Galactic centre) and heliocentric Galactic z in pc, from Gaia ICRS astrometry.""" from astropy.coordinates import SkyCoord import astropy.units as u c = SkyCoord(ra=np.asarray(ra_deg) * u.deg, dec=np.asarray(dec_deg) * u.deg, distance=(1000.0 / np.asarray(plx_mas)) * u.pc, pm_ra_cosdec=np.asarray(pmra) * u.mas / u.yr, pm_dec=np.asarray(pmdec) * u.mas / u.yr, radial_velocity=np.asarray(rv) * u.km / u.s, frame="icrs") g = c.galactic g.representation_type = "cartesian" return (g.velocity.d_x.to_value(u.km / u.s), g.velocity.d_y.to_value(u.km / u.s), g.velocity.d_z.to_value(u.km / u.s), g.cartesian.z.to_value(u.pc)) # ------------------------------------------------------------ deconvolution def likelihood_matrix(age, s_lo, s_hi, grid=AGE_MID, kind="splitnormal", floor=0.05): """L[i, k] = p(reported age_i | true age grid_k), up to a per-star constant. A catalogue posterior obtained with a flat age prior is proportional to the likelihood as a function of true age. We approximate it by a split normal with the reported lower/upper 1-sigma widths ("splitnormal"), or by a normal in ln(age) ("lognormal"). Widths are floored at `floor` Gyr. """ age = np.asarray(age, float)[:, None] s_lo = np.maximum(np.asarray(s_lo, float), floor)[:, None] s_hi = np.maximum(np.asarray(s_hi, float), floor)[:, None] t = np.asarray(grid, float)[None, :] if kind == "splitnormal": s = np.where(t >= age, s_hi, s_lo) L = np.exp(-0.5 * ((t - age) / s) ** 2) elif kind == "lognormal": sl = 0.5 * (s_lo + s_hi) / np.maximum(age, 0.1) L = np.exp(-0.5 * (np.log(np.maximum(t, 1e-3) / np.maximum(age, 0.1)) / sl) ** 2) else: raise ValueError(kind) L = L + 1e-300 return L / L.sum(axis=1, keepdims=True) _TN_CACHE = {} def truncnorm_fit(p16, p50, p84, lo=0.0, hi=14.0): """Parameters (mu, sigma) of a normal truncated to [lo, hi] whose 16th, 50th and 84th percentiles best match the reported ones. Isochrone-age posteriors of unevolved stars are cut off by the edges of the age grid: a star whose data say only "younger than ~6 Gyr" is reported as e.g. 3 (+2.5, -2) Gyr. Treating that as a Gaussian centred on 3 Gyr invents information; the truncated-normal fit recovers a likelihood that is flat or peaks at the edge when the data require it. Solved by nearest-neighbour look-up in a pre-computed (mu, sigma) table. """ from scipy.stats import norm from scipy.spatial import cKDTree key = (lo, hi) if key not in _TN_CACHE: span = hi - lo mu = lo + span * np.concatenate([np.linspace(-3, -0.02, 150), np.linspace(0, 1, 201), np.linspace(1.02, 4, 150)]) sg = span * np.logspace(-2.6, 1.0, 180) M, Sg = np.meshgrid(mu, sg, indexing="ij") a, b = norm.cdf((lo - M) / Sg), norm.cdf((hi - M) / Sg) q = [M + Sg * norm.ppf(np.clip(a + f * (b - a), 1e-300, 1 - 1e-16)) for f in (0.16, 0.50, 0.84)] pts = np.clip(np.stack([x.ravel() for x in q], axis=1), lo, hi) good = np.isfinite(pts).all(axis=1) & ((b - a).ravel() > 1e-12) _TN_CACHE[key] = (cKDTree(pts[good]), M.ravel()[good], Sg.ravel()[good]) tree, Mv, Sv = _TN_CACHE[key] x = np.clip(np.stack([np.asarray(p16, float), np.asarray(p50, float), np.asarray(p84, float)], axis=1), lo, hi) _, j = tree.query(x) return Mv[j], Sv[j] def truncnorm_likelihood(mu, sigma, t, lo=0.0, hi=14.0): """L[i, k] proportional to N(t_k; mu_i, sigma_i) for lo <= t_k <= hi.""" t = np.asarray(t, float)[None, :] L = np.exp(-0.5 * ((t - mu[:, None]) / sigma[:, None]) ** 2) L = L * ((t >= lo) & (t <= hi)) + 1e-300 return L / L.sum(axis=1, keepdims=True) def em_deconvolve(L, weights=None, n_iter=200, phi0=None, smooth=0.0): """Weighted EM for the mixture weights phi_k of an age histogram. phi_k <- sum_i w_i phi_k L_ik / (sum_j phi_j L_ij) / sum_i w_i `smooth` > 0 applies a light (1, s, 1)-type smoothing after every step (a roughness penalty that stabilises the ill-posed deconvolution). Returns phi (sums to 1). """ n, K = L.shape w = np.ones(n) if weights is None else np.asarray(weights, float) w = w / w.sum() phi = np.full(K, 1.0 / K) if phi0 is None else np.asarray(phi0, float) for _ in range(n_iter): r = L * phi[None, :] r /= r.sum(axis=1, keepdims=True) phi = (w[:, None] * r).sum(axis=0) if smooth > 0: p = np.pad(phi, 1, mode="edge") phi = (smooth * p[:-2] + phi + smooth * p[2:]) / (1 + 2 * smooth) phi /= phi.sum() return phi def ml_mixture(L, weights=None, n_em=30): """Maximum-likelihood mixture weights for a small number of components: a few EM steps followed by L-BFGS on softmax parameters (EM alone converges very slowly when the per-star likelihoods are broad).""" from scipy.optimize import minimize n, K = L.shape w = np.ones(n) if weights is None else np.asarray(weights, float) w = w / w.sum() m0 = em_deconvolve(L, w, n_iter=n_em) def nll(theta): e = np.exp(theta - theta.max()) m = e / e.sum() mix = L @ m f = -(w * np.log(mix)).sum() dm = -(L * (w / mix)[:, None]).sum(axis=0) # d f / d m_k grad = m * (dm - (dm * m).sum()) return f, grad res = minimize(nll, np.log(np.maximum(m0, 1e-6)), jac=True, method="L-BFGS-B", options=dict(maxiter=500)) e = np.exp(res.x - res.x.max()) return e / e.sum() def weighted_hist(age, weights=None, edges=AGE_EDGES): h, _ = np.histogram(np.clip(age, edges[0], edges[-1] - 1e-9), bins=edges, weights=weights) return h / h.sum() # -------------------------------------------------- head-start statistics HEADSTART_THRESHOLDS = (0.0, 0.1, 1.0, 2.0, 3.0, 5.0) def frac_older(phi, x, edges=AGE_EDGES, t_sun=T_SUN): """Fraction of the (binned, piecewise-uniform) age distribution phi that is older than t_sun + x Gyr.""" t = t_sun + x lo, hi = edges[:-1], edges[1:] part = np.clip((hi - t) / (hi - lo), 0, 1) return float((phi * part).sum() / phi.sum()) def headstart_stats(phi, edges=AGE_EDGES, t_sun=T_SUN, t_cap=T_MAX): """Summary of the head-start distribution h = age - t_sun. Ages above t_cap (unphysical, error-driven) are counted at t_cap when computing means/medians (they still count as 'older').""" phi = np.asarray(phi, float) / np.sum(phi) out = {f"f_gt_{x:g}": frac_older(phi, x, edges, t_sun) for x in HEADSTART_THRESHOLDS} # fine sampling of the piecewise-uniform density nsub = 50 lo, hi = edges[:-1], edges[1:] tt = (lo[:, None] + (np.arange(nsub)[None, :] + 0.5) / nsub * (hi - lo)[:, None]).ravel() ww = np.repeat(phi / nsub, nsub) tt = np.minimum(tt, t_cap) out["mean_age"] = float((tt * ww).sum()) order = np.argsort(tt) cs = np.cumsum(ww[order]) out["median_age"] = float(tt[order][np.searchsorted(cs, 0.5)]) older = tt > t_sun h = tt[older] - t_sun wh = ww[older] / ww[older].sum() out["mean_headstart_older"] = float((h * wh).sum()) o2 = np.argsort(h) out["median_headstart_older"] = float(h[o2][np.searchsorted( np.cumsum(wh[o2]), 0.5)]) out["mean_offset_all"] = out["mean_age"] - t_sun return out def bootstrap(func, n, n_boot=200, seed=1): """Run func(idx) on bootstrap resamples idx of range(n); returns list.""" rng = np.random.default_rng(seed) return [func(rng.integers(0, n, n)) for _ in range(n_boot)] def pct(a, q=(16, 50, 84)): return np.percentile(np.asarray(a, float), q, axis=0) # ------------------------------------------------------------------- MIST MIST_TARBALL = os.path.join(RAW, "mist", "MIST_v1.2_vvcrit0.0_UBVRIplus.txz") MIST_CACHE = os.path.join(DATA, "mist_compact.npz") MIST_FEH = (-2.00, -1.75, -1.50, -1.25, -1.00, -0.75, -0.50, -0.25, 0.00, 0.25, 0.50) def _feh_tag(feh): return ("p" if feh >= 0 else "m") + f"{abs(feh):.2f}" def load_mist(rebuild=False): """Return dict feh -> array with columns [log_age, EEP, initial_mass, log_Teff, log_g, log_L, 2MASS_Ks, phase, Gaia_G_EDR3, Gaia_BP_EDR3, Gaia_RP_EDR3] restricted to 8.9 <= log_age <= 10.3 and EEP <= 605 (tip of the RGB). A compact cache is written to data/mist_compact.npz (2.6 MB) so that the 157 MB tarball is only needed to rebuild it. """ if os.path.exists(MIST_CACHE) and not rebuild: z = np.load(MIST_CACHE) return {float(k[1:]): z[k] for k in z.files} out = {} with tarfile.open(MIST_TARBALL, "r:xz") as tar: members = {m.name: m for m in tar.getmembers() if m.isfile()} for feh in MIST_FEH: name = [n for n in members if f"feh_{_feh_tag(feh)}_" in n][0] raw = tar.extractfile(members[name]).read().decode() rows = [ln for ln in raw.splitlines() if ln and not ln.startswith("#")] arr = np.loadtxt(io.StringIO("\n".join(rows)), usecols=(1, 0, 2, 4, 5, 6, 16, 33, 30, 31, 32)) # 0-based columns: 16 = 2MASS_Ks, 33 = phase, 30-32 = Gaia EDR3 keep = (arr[:, 0] >= 8.9) & (arr[:, 0] <= 10.31) & (arr[:, 1] <= 605) out[feh] = arr[keep].astype(np.float32) np.savez_compressed(MIST_CACHE, **{f"f{feh:+.2f}": a for feh, a in out.items()}) return out def mh_from_feh_alpha(feh, afe): """Salaris, Chieffi & Straniero (1993) total metallicity.""" afe = np.nan_to_num(np.asarray(afe, float), nan=0.0) return np.asarray(feh, float) + np.log10(0.638 * 10 ** afe + 0.362)