<jupyter_start><jupyter_text>Run: `scripts/compute_actions.py` before this notebook!<jupyter_code>import os
import pathlib
import itertools
import pickle

import astropy.coordinates as coord
from astropy.convolution import convolve, Gaussian2DKernel
from astropy.io import fits
import astropy.table as at
import astropy.units as u
import matplotlib as mpl
import matplotlib.pyplot as plt
%matplotlib inline
import numpy as np
from scipy.stats import binned_statistic
from scipy.spatial import cKDTree
from scipy.interpolate import interp1d
from tqdm.notebook import trange, tqdm

# gala
import gala.coordinates as gc
import gala.dynamics as gd
import gala.integrate as gi
import gala.potential as gp
from gala.units import galactic
from gala.mpl_style import hesperia, laguna

from totoro.config import vcirc, rsun, plot_path, fig_path, cache_path, fiducial_mdisk
from totoro.config import plot_config as pc
from totoro.config import galcen_frame, elem_names as apogee_elem_names
from totoro.data import datasets
from totoro.potentials import potentials
from totoro.abundance_helpers import elem_to_label
from totoro.atm import (AbundanceTorusMaschine, 
                        run_bootstrap_coeffs, 
                        get_cos2th_zerocross)
data_name = 'apogee-rgb-loalpha'
d = datasets[data_name]
this_cache_path = cache_path / data_name
all_aafs = {}
for name in potentials:
    filename = this_cache_path / f'aaf-{name}.fits'
    if not filename.exists():
        continue
    all_aafs[name] = at.join(
        at.QTable.read(filename), 
        d.t, 
        keys='APOGEE_ID')  # NOTE: could be d._id_column
with open(cache_path / 'w0s.pkl', 'rb') as f:
    w0s = pickle.load(f)
    
orbits = {}
for k, w0 in w0s.items():
    orbits[k] = potentials[k].integrate_orbit(
        w0, dt=0.5*u.Myr, t1=0, t2=6*u.Gyr
    )
sorted_keys = sorted(
    orbits.keys(), 
    key=lambda k: potentials[k]['disk'].parameters['m'])
sorted_keys = [x for x in sorted_keys if x in all_aafs]
all_finite_mask = None
for name in sorted_keys:
    finite_mask = np.all(np.isfinite(all_aafs[name]['actions']), axis=1)
    finite_mask &= np.all(all_aafs[name]['angles'] < 2*np.pi*u.rad, axis=1)
    if all_finite_mask is None:
        all_finite_mask = finite_mask
    else:
        all_finite_mask &= finite_mask
        
print(all_finite_mask.sum())
for elem_name in apogee_elem_names:
    print(elem_name, (all_aafs['1.0'][elem_name][all_finite_mask] > -3).sum())
main_elem = 'O_FE'
# main_elem = 'N_FE'

main_elem_mask = None
for name in sorted_keys:
    elem = all_aafs[name][main_elem]
    _mask = np.isfinite(elem) & (elem > -10) & (elem < 10)
    if main_elem_mask is None:
        main_elem_mask = _mask
    else:
        main_elem_mask &= _mask
        
print(main_elem_mask.sum())
aafs = {}
for name in all_aafs:
    aafs[name] = all_aafs[name][all_finite_mask & main_elem_mask]<jupyter_output><empty_output><jupyter_text>Visualize results from fitting the full sample for each potential model<jupyter_code>fig, axes = plt.subplots(3, 5, figsize=(15, 10), 
                         sharex=True, sharey=True)
for i, name in enumerate(sorted_keys):
    atm = AbundanceTorusMaschine(aafs[name])
    
    coeffs, coeff_cov = atm.get_coeffs_for_elem(main_elem)
    angz, d_elem, d_elem_errs = atm.get_theta_z_anomaly(main_elem)
    binx, bin_anom, bin_anom_err = atm.get_binned_anomaly(main_elem)

    # ---
    
    ax = axes.flat[i]
    ax.plot(angz, d_elem, 
            marker='o', ls='none', mew=0, ms=2, alpha=0.3)

    ax.plot(binx, bin_anom, 
            marker='', drawstyle='steps-mid', 
            zorder=10, color='tab:red', alpha=1)
    ax.errorbar(binx, bin_anom, bin_anom_err,
                marker='o', ls='none', ecolor='tab:red', 
                zorder=5, alpha=0.6)
    
    plot_x = np.linspace(0, 2*np.pi, 1024)
    plot_M = atm.get_M(plot_x)
    plot_y = plot_M @ coeffs
    ax.plot(plot_x, plot_y, marker='', lw=2, color='tab:blue', zorder=100)
    
    ax.set_title(name, fontsize=18, pad=11)
    
ylim = np.max(np.abs(np.nanpercentile(d_elem, [25, 75])))
    
ax.set_xlim(0, 2*np.pi)
# ax.set_ylim(-0.025, 0.025)
ax.set_ylim(-ylim, ylim)
all_bs_coeffs = run_bootstrap_coeffs(aafs, main_elem, bootstrap_K=128,
                                     cache_path=this_cache_path)
fig, axes = plt.subplots(1, 3, figsize=(12, 4.4), sharex=True, sharey=True)

plot_x = np.linspace(0, 2*np.pi, 1024)
for i, name in enumerate(['0.4', '1.0', '1.6']):
    atm = AbundanceTorusMaschine(aafs[name])
    bin_x, bin_anom, bin_anom_err = atm.get_binned_anomaly(main_elem)
    
    ax = axes[i]
    
    mean_c = np.mean(all_bs_coeffs[name], axis=0)
    if i == 0:
        label = 'mean abundance deviation'
    else:
        label = None
    ax.plot(bin_x, bin_anom - mean_c[0], 
            marker='', drawstyle='steps-mid', color='k', lw=2,
            label=label)
    
    for j, c in enumerate(all_bs_coeffs[name]):
        plot_y = atm.get_M(plot_x) @ c
        if j == 0 and i == 1:
            label = 'bootstrap samples'
        else:
            label = None
        ax.plot(plot_x, plot_y - mean_c[0], 
                alpha=0.25, color='tab:blue', marker='', lw=1.,
                label=label)
        
    tmp = np.zeros_like(c)
    tmp[3] = c[3]
    plot_y = atm.get_M(plot_x) @ tmp
    if i == 2:
        label = r'mean $\cos(2\,\theta_z)$ term'
    else:
        label = None
    ax.plot(plot_x, plot_y, 
            alpha=0.8, color='tab:red', marker='', lw=2, zorder=100, 
            label=label)
    
    title = (r'${\rm M}_{\rm disk} / {\rm M}_{\rm disk}^\star =' 
             + f' {float(name):.1f}$')
    ax.set_title(title, pad=11, fontsize=26)

    ax.axhline(0, zorder=-100, color='#aaaaaa')
    ax.set_xlabel(r'vertical angle, $\theta_z$ [rad]')

ax = axes[0]
ax.set_ylim(-0.012, 0.012)
ax.set_ylabel(f'mean {elem_to_label(main_elem)}\ndeviation, ' + 
              r'$\Delta^{[{\rm O}/{\rm Fe}]}$')

ax.set_xlim(0, 2*np.pi)
ax.set_xticks(np.arange(0, 2+1e-3, 1) * np.pi)
ax.set_xticklabels(['0'] + [rf'${x:.0f}\,\pi$' for x in np.arange(1, 2+1e-3, 1)])
ax.set_xticks(np.arange(0, 2+1e-3, 0.5) * np.pi, minor=True)


for ax in axes:
    ax.legend(loc='lower center', fontsize=13)

fig.tight_layout()
fig.set_facecolor('w')

fig.savefig(fig_path / 'sinusoid-fits.pdf')<jupyter_output><empty_output><jupyter_text>---<jupyter_code>summary, zero_cross, zero_cross_err = get_cos2th_zerocross(all_bs_coeffs)
zero_cross_err
fig, axes = plt.subplots(2, 2, figsize=(12*0.9, 8*0.9), 
                         sharex=True, sharey=True)

for j, i in enumerate(range(1, 5)):
    ax = axes.flat[j]
    
    ax.errorbar(summary[i]['mdisk'], 
                summary[i]['y'], 
                yerr=summary[i]['y_err'],
                marker='o', ecolor='#666666', ls='')
    
    ax.plot(summary[i]['mdisk'], 
            summary[i]['y'], 
            marker='', ls='-', color='#aaaaaa', zorder=-10)
    
    ax.axhline(0, marker='', ls='--', zorder=-100, color='tab:green', alpha=0.5)

    
axes[1, 0].errorbar(zero_cross, [0], 
                    xerr=np.array(zero_cross_err)[None].T, 
                    ls=None, marker='s', color='tab:red',
                    alpha=0.75, mew=0, ms=8)
    
axes[0, 0].set_ylabel('$a_1$')
axes[0, 1].set_ylabel('$b_1$')
axes[1, 0].set_ylabel('$a_2$')
axes[1, 1].set_ylabel('$b_2$')
axes[1, 0].set_xlabel(r'${\rm M}_{\rm disk} / {\rm M}_{\rm disk}^\star$')
axes[1, 1].set_xlabel(r'${\rm M}_{\rm disk} / {\rm M}_{\rm disk}^\star$')

ax = axes[1, 0]
xlim = (min(summary[0]['mdisk'])-0.1, 
        max(summary[0]['mdisk'])+0.1)
ax.set_xlim(xlim)
ax.set_xticks(np.arange(min(summary[0]['mdisk']), max(summary[0]['mdisk'])+1e-3, 0.2))
ax.set_xticks(np.arange(min(summary[0]['mdisk']), max(summary[0]['mdisk'])+1e-3, 0.1), 
              minor=True)
    
subtitles = ['(sensitive to solar motion in $v_z$)',
             '(sensitive to solar position in $z$)',
             '(sensitive to local density)',
             '(verboten)']
for i, text in enumerate(subtitles):    
    axes.flat[i].text(
        1.825, -0.005, text,
        fontsize=20, ha='right', va='bottom', 
        color='#555555', fontstyle='italic')
    
axes[0, 0].set_title(r'$\cos\,\theta_z$ amplitude', pad=11, fontsize=24)
axes[0, 1].set_title(r'$\sin\,\theta_z$ amplitude', pad=11, fontsize=24)
axes[1, 0].set_title(r'$\cos\,2\theta_z$ amplitude', pad=11, fontsize=24)
axes[1, 1].set_title(r'$\sin\,2\theta_z$ amplitude', pad=11, fontsize=24)
axes[0, 0].set_ylim(-0.0055, 0.0022)

fig.tight_layout()
fig.subplots_adjust(hspace=0.3)

fig.savefig(fig_path / 'coeff-vs-mdisk.pdf')<jupyter_output><empty_output><jupyter_text>Anomaly plot for all elements:<jupyter_code>for elem_name in apogee_elem_names:
    fig, axes = plt.subplots(3, 5, figsize=(15, 10), 
                             sharex=True, sharey=True)
    for i, name in enumerate(sorted_keys):
        atm = AbundanceTorusMaschine(aafs[name])

        coeffs, coeff_cov = atm.get_coeffs_for_elem(elem_name)
        angz, d_elem, d_elem_errs = atm.get_theta_z_anomaly(elem_name)
        binx, bin_anom, bin_anom_err = atm.get_binned_anomaly(elem_name)

        # ---

        ax = axes.flat[i]
        ax.plot(angz, d_elem, 
                marker='o', ls='none', mew=0, ms=2, alpha=0.3)

        ax.plot(binx, bin_anom, 
                marker='', drawstyle='steps-mid', 
                zorder=10, color='tab:red', alpha=1)
        ax.errorbar(binx, bin_anom, bin_anom_err,
                    marker='o', ls='none', ecolor='tab:red', 
                    zorder=5, alpha=0.6)

        plot_x = np.linspace(0, 2*np.pi, 1024)
        plot_M = atm.get_M(plot_x)
        plot_y = plot_M @ coeffs
        ax.plot(plot_x, plot_y, marker='', lw=2, color='tab:blue', zorder=100)

        ax.set_title(name, fontsize=18, pad=11)

    ax.set_xlim(0, 2*np.pi)
    ax.set_ylim(-0.025, 0.025)
    
    fig.savefig(plot_path / data_name / elem_name / f'atm-anomaly-{elem_name}.png', dpi=250)
    plt.close(fig)<jupyter_output><empty_output><jupyter_text>Run zero-cross finder for all elements:<jupyter_code>from totoro.atm import ZeroCrossWorker
from schwimmbad import MultiPool, SerialPool
zerocross_worker = ZeroCrossWorker(
    aafs, cache_path=this_cache_path, 
    bootstrap_K=32)
results = []
with MultiPool() as pool:
# with SerialPool() as pool:
    for r in pool.map(zerocross_worker, apogee_elem_names):
        results.append(r)
results
labels = []
xs = []
ys = []
yerrs = []

for i, r in enumerate(results):
    labels.append(r[0])
    xs.append(i)
    ys.append(r[1][0])
    yerrs.append(r[1][1])

ys = np.array(ys)
yerrs = np.array(yerrs)
    
fig, ax = plt.subplots(1, 1, figsize=(8.5, 6))
ax.errorbar(xs, ys, yerrs.T, marker='s', ms=6,
            ls='none', ecolor='#666666', zorder=10)
ax.errorbar(xs[4], ys[4], yerrs.T[:, 4:5], marker='s', ms=6,
            ls='none', color='tab:blue', zorder=100)
# ax.axhline(1., zorder=-10, color='#aaaaaa', alpha=0.5)

dumb_yerrs = np.mean(yerrs, axis=1)
ivars = 1 / dumb_yerrs**2
comb_y = np.sum(ys * ivars) / np.sum(ivars)
comb_err = np.sqrt(1 / np.sum(ivars))
ax.axhline(comb_y, color='tab:purple', zorder=-5, alpha=1)
ax.axhspan(comb_y - comb_err, comb_y + comb_err, 
           color='tab:purple', lw=0, alpha=0.5, zorder=-5)

ax.set_xlim(-0.75, len(results)-1 + 0.75)
ax.set_ylim(0, 2)

ax.set_ylabel(r'inferred ${\rm M}_{\rm disk} / {\rm M}_{\rm disk}^\star$')

ax.set_xticks([])
ax.set_xticklabels([])
# ax.set_xticklabels([elem_to_label(x) for x in labels])
for x, y, yerr, label in zip(xs, ys, yerrs, labels):
    ax.text(x, y - yerr[0] - 0.12, elem_to_label(label),
            ha='center', va='top', fontsize=18, color='#555555')

fig.tight_layout()
fig.set_facecolor('w')

# fig.savefig(fig_path / 'mdisk-vs-elem.pdf')
for elem_name in apogee_elem_names:
    _all_bs_coeffs = run_bootstrap_coeffs(aafs, elem_name, 
                                          bootstrap_K=32,
                                          cache_path=cache_path / data_name)
    elem_summary, *_ = get_cos2th_zerocross(_all_bs_coeffs)

    fig, axes = plt.subplots(2, 2, figsize=(12, 9), 
                             sharex=True, sharey='row')

    for j, i in enumerate(range(1, 5)):
        ax = axes.flat[j]

        ax.errorbar(elem_summary[i]['mdisk'], 
                    elem_summary[i]['y'], 
                    elem_summary[i]['y_err'],
                    marker='o', ecolor='#666666', ls='')
        
        _idx = np.argsort(elem_summary[i]['mdisk'])
        ax.plot(np.array(elem_summary[i]['mdisk'])[_idx], 
                np.array(elem_summary[i]['y'])[_idx], 
                marker='', ls='-', color='#aaaaaa', zorder=-10)

        ax.axhline(0, marker='', zorder=-100, color='tab:green', alpha=0.5)

    axes[0, 0].set_ylabel(r'$\cos\,\theta_z$ amplitude')
    axes[0, 1].set_ylabel(r'$\sin\,\theta_z$ amplitude')
    axes[1, 0].set_ylabel(r'$\cos\,2\theta_z$ amplitude')
    axes[1, 1].set_ylabel(r'$\sin\,2\theta_z$ amplitude')
    axes[1, 0].set_xlabel(r'${\rm factor} \times {\rm M}_{\rm disk}$')
    axes[1, 1].set_xlabel(r'${\rm factor} \times {\rm M}_{\rm disk}$')

    ax = axes[1, 0]
    xlim = (min(elem_summary[0]['mdisk'])-0.1, 
            max(elem_summary[0]['mdisk'])+0.1)
    ax.set_xlim(xlim)
    ax.set_xticks(np.arange(min(elem_summary[0]['mdisk']), max(elem_summary[0]['mdisk'])+1e-3, 0.2))
    ax.set_xticks(np.arange(min(elem_summary[0]['mdisk']), max(elem_summary[0]['mdisk'])+1e-3, 0.1), 
                  minor=True)

    axes[0, 0].set_title('(sensitive to solar motion in $v_z$)', pad=11, fontsize=20)
    axes[0, 1].set_title('(sensitive to solar position in $z$)', pad=11, fontsize=20)
    axes[1, 0].set_title('(sensitive to local density)', pad=11, fontsize=20)
    axes[1, 1].set_title('(verboten)', pad=11, fontsize=20)

    fig.tight_layout()
    fig.subplots_adjust(hspace=0.3)
    
    fig.savefig(plot_path / data_name / elem_name / f'coeffs-vs-mdisk.png', dpi=250)
    plt.close(fig)<jupyter_output><empty_output><jupyter_text>---  Note: below shows results from a test comparing bootstrap estimated coefficient errors to propagated uncertainty Showing that the d_elem distribution has an intrinsic width of ~0.04 dex<jupyter_code>np.mean(d_elem), np.median(d_elem), np.std(d_elem)
plt.hist(d_elem, bins=np.linspace(-0.2, 0.2, 64));<jupyter_output><empty_output><jupyter_text>If we run the below, but replacing d_elem_errs with 0.04, the ratio goes to 1<jupyter_code>np.mean(bs_coeffs[:, 3]), coeffs[3]
for i in range(5):
    print(np.std(bs_coeffs[:, i]), np.sqrt(coeffs_cov[i, i]),
          np.std(bs_coeffs[:, i]) / np.sqrt(coeffs_cov[i, i]))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16399372142223437, 'fraction_numerical': 0.033235514911622195, 'mean_word_length': 3.170176437108708, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 31, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 5, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '414594', 'license': 'MIT', 'n_tokens_mistral': 6143, 'n_tokens_neox': 5766, 'n_words': 1069, 'path': 'notebooks/figures/4-Abundance-anomaly-bootstrap-1d.ipynb', 'repo_name': 'adrn/thrift-shop', 'size': '14653'}