<jupyter_start><jupyter_text>z_max vs Jz for fiducial model:$J_z \propto \sqrt{z_{\rm max}}$<jupyter_code>from totoro.action_helpers import get_staeckel_aaf
elem_name = 'O_FE'

path = pathlib.Path(this_cache_path / f"optimize-results-{elem_name}.csv")
tbl = at.Table.read(path)

xs = []
for col in tbl.colnames:
    xs.append(np.nanmean(tbl[col]))

obj = TorusImagingObjective(d, elem_name)
atm, w0, pot = obj.get_atm_w0(*xs)
Jz = atm.aaf['actions'][:, 2].to_value(u.kpc*u.km/u.s)
orbits = pot.integrate_orbit(w0[:1000], dt=1., n_steps=1000)
zmax = orbits.zmax(approximate=True)
coeffs = np.polyfit(np.log(Jz[:1000]), np.log(zmax.to_value(u.kpc)), deg=2)
logjz_to_logzmax_poly = np.poly1d(coeffs)
plt.scatter(np.log(Jz[:1000]), np.log(zmax.to_value(u.kpc)))<jupyter_output><empty_output><jupyter_text>---<jupyter_code>fig, axes = plt.subplots(2, 4, figsize=(14, 7.),
                         sharex=True, 
                         constrained_layout=True)

for i, (ax, elem_name) in enumerate(zip(axes.flat, elem_names)):
    path = pathlib.Path(this_cache_path / f"optimize-results-{elem_name}.csv")
    tbl = at.Table.read(path)

    xs = []
    for col in tbl.colnames:
        xs.append(np.nanmean(tbl[col]))

    obj = TorusImagingObjective(d, elem_name)
    atm, *_ = obj.get_atm_w0(*xs)

    Jz = atm.aaf['actions'][:, 2].to_value(u.kpc*u.km/u.s)
    elem = atm.aaf[elem_name]  # nan-safe

    mask = (elem > -3) & (elem < 3)
    elem = elem - np.nanmedian(elem[mask])

    elem = elem[mask]
    Jz = Jz[mask]

    pcl = np.nanpercentile(elem, [1, 99])
    lim = 2 * np.max(np.abs(pcl))

    ax.hist2d(Jz, elem, 
              bins=(10 ** np.arange(-3, 3., 0.02),
                    np.linspace(-lim, lim, 128)),
              norm=mpl.colors.LogNorm(vmin=5e-1),
              cmap='Greys', rasterized=True)
    
    stat = binned_statistic(Jz, elem, 
                            bins=10 ** np.arange(-3, 3, 0.1))
    ctr = 0.5 * (stat.bin_edges[:-1] + stat.bin_edges[1:])
    ax.plot(ctr, stat.statistic, 
            marker='', drawstyle='steps-mid', 
            lw=1.5, color='tab:blue')
    
    for val in [16, 84]:
        stat = binned_statistic(Jz, elem, 
                                bins=stat.bin_edges, 
                                statistic=lambda x: np.percentile(x, val))
        ax.plot(ctr, stat.statistic, 
                marker='', drawstyle='steps-mid', 
                lw=1., color='tab:purple')
    
    ax.set_xscale('log')
    
    if elem_name == 'FE_H':
        lim = 0.6
        ax.set_yticks(np.arange(-0.6, 0.6+1e-3, 0.2))
    else:
        lim = 0.3
        if elem_name not in ['C_FE', 'MG_FE']:
            ax.set_yticklabels([])
    
    ax.set_yticks(np.arange(-0.6, 0.6+1e-3, 0.1), minor=True)
    ax.set_ylim(-lim, lim)
    
    ax.errorbar(1e2, lim-0.1*lim, 
                np.median(d.g.data[elem_name + '_ERR']),
                marker='o', ms=4, color='tab:red', alpha=0.75)

    lbl = elem_to_label(elem_name)
    ax.text(1.5e-2, lim-0.05*lim, lbl, 
            ha='left', va='top', fontsize=18)
    
ax.set_xticks(10. ** np.arange(-2, 2+1, 1))
    
for i in range(axes.shape[0]):
    axes[i, 0].set_ylabel(rf'$X - \langle X \rangle$')
    
for j in range(axes.shape[1]):
    axes[-1, j].set_xlabel(f'$J_z$ [{u.kpc*u.km/u.s:latex_inline}]')
    
fig.set_facecolor('w')
# fig.savefig(fig_path / 'elem-Jz-gradients.pdf', dpi=250)
fig, axes = plt.subplots(2, 4, figsize=(14, 7.),
                         sharex=True, 
                         constrained_layout=True)

for i, (ax, elem_name) in enumerate(zip(axes.flat, elem_names)):
    path = pathlib.Path(this_cache_path / f"optimize-results-{elem_name}.csv")
    tbl = at.Table.read(path)

    xs = []
    for col in tbl.colnames:
        xs.append(np.nanmean(tbl[col]))

    obj = TorusImagingObjective(d, elem_name)
    atm, *_ = obj.get_atm_w0(*xs)

    Jz = atm.aaf['actions'][:, 2].to_value(u.kpc*u.km/u.s)
    zmax = np.exp(logjz_to_logzmax_poly(np.log(Jz)))
    elem = atm.aaf[elem_name]  # nan-safe

    mask = (elem > -3) & (elem < 3)
    elem = elem - np.nanmedian(elem[mask])

    elem = elem[mask]
    zmax = zmax[mask]

    pcl = np.nanpercentile(elem, [1, 99])
    lim = 2 * np.max(np.abs(pcl))

    ax.hist2d(zmax, elem, 
              bins=(10 ** np.arange(-2, 1, 0.02),
                    np.linspace(-lim, lim, 128)),
              norm=mpl.colors.LogNorm(vmin=5e-1),
              cmap='Greys', rasterized=True)
    
    stat = binned_statistic(zmax, elem, 
                            bins=10 ** np.arange(-2, 1, 0.1))
    ctr = 0.5 * (stat.bin_edges[:-1] + stat.bin_edges[1:])
    ax.plot(ctr, stat.statistic, 
            marker='', drawstyle='steps-mid', 
            lw=1.5, color='tab:blue')
    
    for val in [16, 84]:
        stat = binned_statistic(zmax, elem, 
                                bins=stat.bin_edges, 
                                statistic=lambda x: np.percentile(x, val))
        ax.plot(ctr, stat.statistic, 
                marker='', drawstyle='steps-mid', 
                lw=1., color='tab:purple')
    
    ax.set_xscale('log')
    
    if elem_name == 'FE_H':
        lim = 0.6
        ax.set_yticks(np.arange(-0.6, 0.6+1e-3, 0.2))
    else:
        lim = 0.3
        if elem_name not in ['C_FE', 'MG_FE']:
            ax.set_yticklabels([])
    
    ax.set_yticks(np.arange(-0.6, 0.6+1e-3, 0.1), minor=True)
    ax.set_ylim(-lim, lim)
    
    ax.errorbar(1e2, lim-0.1*lim, 
                np.median(d.g.data[elem_name + '_ERR']),
                marker='o', ms=4, color='tab:red', alpha=0.75)

    lbl = elem_to_label(elem_name)
    ax.text(1.5e-2, lim-0.05*lim, lbl, 
            ha='left', va='top', fontsize=18)
    
    ax.axvline(0.5, color='tab:red', alpha=0.5, ls='--')
    
ax.set_xticks(10. ** np.arange(-2, 1+1, 1))
ax.set_xlim(1e-2, 1e1)
    
for i in range(axes.shape[0]):
    axes[i, 0].set_ylabel(rf'$X - \langle X \rangle$')
    
for j in range(axes.shape[1]):
    axes[-1, j].set_xlabel(r'$z_{\rm max}$ ' + f'[{u.kpc:latex_inline}]')
    
fig.set_facecolor('w')
fig.savefig(fig_path / 'elem-zmax-gradients.pdf', dpi=250)<jupyter_output><empty_output><jupyter_text>---<jupyter_code>polys = {}

fig, axes = plt.subplots(2, 4, figsize=(14, 7.),
                         sharex=True, 
                         constrained_layout=True)

for i, (ax, elem_name) in enumerate(zip(axes.flat, elem_names)):
    path = pathlib.Path(this_cache_path / f"optimize-results-{elem_name}.csv")
    tbl = at.Table.read(path)

    xs = []
    for col in tbl.colnames:
        xs.append(np.nanmean(tbl[col]))

    obj = TorusImagingObjective(d, elem_name)
    atm = obj.get_atm(*xs)

    Jz = atm.aaf['actions'][:, 2].to_value(u.kpc*u.km/u.s)
    elem = atm.aaf[elem_name]  # nan-safe
    elem_err = atm.aaf[elem_name + '_ERR']

    mask = (elem > -3) & (elem < 3)
    elem = elem - np.nanmedian(elem[mask])

    elem = elem[mask]
    Jz = Jz[mask]
    elem_err = elem_err[mask]

    pcl = np.nanpercentile(elem, [1, 99])
    lim = 2 * np.max(np.abs(pcl))

    ax.hist2d(Jz, elem, 
              bins=(10 ** np.arange(-2, 2.5, 0.02),
                    np.linspace(-lim, lim, 128)),
              norm=mpl.colors.LogNorm(vmin=5e-1),
              cmap='Greys', rasterized=True)
    
#     stat = binned_statistic(Jz, elem, 
#                             bins=10 ** np.arange(-3, 3, 0.1))
#     y = stat.statistic

    stat1 = binned_statistic(Jz, elem / elem_err**2, statistic='sum', 
                             bins=10 ** np.arange(-3, 3, 0.1))
    stat2 = binned_statistic(Jz, 1 / elem_err**2, statistic='sum', 
                             bins=10 ** np.arange(-3, 3, 0.1))
    y = stat1.statistic / stat2.statistic

    ctr = 0.5 * (stat1.bin_edges[:-1] + stat1.bin_edges[1:])
    ax.plot(ctr, y, 
            marker='', drawstyle='steps-mid', 
            lw=2, color='tab:blue')
    
    logJz = np.log10(Jz)
    coeffs = np.polyfit(logJz, elem, w=1/elem_err**2, deg=7)
    poly = np.poly1d(coeffs)
    polys[elem_name] = poly
    
    x = np.linspace(logJz.min(), logJz.max(), 1024)
    ax.plot(10**x, poly(x))
    
    ax.set_xscale('log')
    ax.set_ylim(-lim, lim)

    lbl = elem_to_label(elem_name)
    ax.text(1.5e-2, lim-0.05*lim, lbl, 
            ha='left', va='top', fontsize=18)
    
for i in range(axes.shape[0]):
    axes[i, 0].set_ylabel(rf'$X - \langle X \rangle$')
    
for j in range(axes.shape[1]):
    axes[-1, j].set_xlabel(f'$J_z$ [{u.kpc*u.km/u.s:latex_inline}]')
    
fig.set_facecolor('w')
len(Jz), len(mask), len(elem)
polys = {}

fig, axes = plt.subplots(2, 4, figsize=(14, 7.),
                         sharex=True, 
                         constrained_layout=True)

for i, (ax, elem_name) in enumerate(zip(axes.flat, elem_names)):
    path = pathlib.Path(this_cache_path / f"optimize-results-{elem_name}.csv")
    tbl = at.Table.read(path)

    xs = []
    for col in tbl.colnames:
        xs.append(np.nanmean(tbl[col]))
    
    # only used to get actions
    tmpd = d[np.isfinite(d.t[elem_name]) &
             (d.t[elem_name] > -3) & 
             (d.t[elem_name] < 3)]
    obj = TorusImagingObjective(tmpd, elem_name)
    atm = obj.get_atm(*xs)
    Jz = atm.aaf['actions'][:, 2].to_value(u.kpc*u.km/u.s)
        
    if elem_name != 'FE_H':
        elem_name = '_'.join([elem_name.split('_')[0], 'H'])

    elem = tmpd.get_elem_ratio(elem_name)
    elem = elem - np.nanmedian(elem)

    pcl = np.nanpercentile(elem, [1, 99])
    lim = 2 * np.max(np.abs(pcl))

    ax.hist2d(Jz, elem, 
              bins=(10 ** np.arange(-2, 2.5, 0.02),
                    np.linspace(-lim, lim, 128)),
              norm=mpl.colors.LogNorm(vmin=5e-1),
              cmap='Greys', rasterized=True)
    
    logJz = np.log10(Jz)
    coeffs = np.polyfit(logJz, elem, deg=5)
    poly = np.poly1d(coeffs)
    polys[elem_name] = poly
    
    x = np.linspace(logJz.min(), logJz.max(), 1024)
    ax.plot(10**x, poly(x))
    
    ax.set_xscale('log')
    ax.set_ylim(-lim, lim)

    lbl = elem_to_label(elem_name)
    ax.text(1.5e-2, lim-0.05*lim, lbl, 
            ha='left', va='top', fontsize=18)
    
for i in range(axes.shape[0]):
    axes[i, 0].set_ylabel(rf'$X - \langle X \rangle$')
    
for j in range(axes.shape[1]):
    axes[-1, j].set_xlabel(f'$J_z$ [{u.kpc*u.km/u.s:latex_inline}]')
    
fig.set_facecolor('w')
fig = plt.figure(figsize=(6, 6))

grid = np.linspace(0, 2.2, 256)
for i, elem_name in enumerate(polys.keys()):
    poly = polys[elem_name]
    deriv = poly.deriv()(grid)
    
    plt.plot(grid, np.abs(deriv), marker='', label=elem_name)

plt.legend()

plt.xlim(0, 2.2)
plt.ylim(-0.05, 0.5)

plt.xlabel(r'$\log J_z$')
plt.ylabel(r'$\left| \frac{{\rm d} [{\rm X}/{\rm H}]}{{\rm d} \log J_z} \right|$')

fig.set_facecolor('w')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.18097361575622445, 'fraction_numerical': 0.03363062058714233, 'mean_word_length': 2.283013113754193, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '48238', 'license': 'MIT', 'n_tokens_mistral': 4713, 'n_tokens_neox': 4394, 'n_words': 714, 'path': 'notebooks/figures/6-APOGEE-elem-Jz-gradients.ipynb', 'repo_name': 'adrn/thrift-shop', 'size': '10764'}