<jupyter_start><jupyter_text>Compute Staeckel delta on a grid:<jupyter_code>Rz_grids = (np.arange(8-2.5, 8+2.5 + 1e-3, 0.05),
            np.arange(-2.5, 2.5 + 1e-3, 0.05))
Rz_grid = np.stack(list(map(np.ravel, np.meshgrid(*Rz_grids)))).T
for pot_name in ['0.4', '1.0', '1.6']:
    pot = galpy_potentials[pot_name]

    delta_staeckels = []
    for i in range(Rz_grid.shape[0]):
        R = (Rz_grid[i, 0] * u.kpc).to_value(ro)
        z = (Rz_grid[i, 1] * u.kpc).to_value(ro)
        delta_staeckels.append(estimateDeltaStaeckel(
            pot, R, z))

    plt.figure()
    plt.title(pot_name)
    plt.scatter(Rz_grid[:, 0], Rz_grid[:, 1], 
                c=delta_staeckels,
                vmin=2, vmax=6, s=8, marker='s')
from scipy.interpolate import NearestNDInterpolator
pot = galpy_potentials['1.0']

delta_staeckels = []
for i in range(Rz_grid.shape[0]):
    R = (Rz_grid[i, 0] * u.kpc).to_value(ro)
    z = (Rz_grid[i, 1] * u.kpc).to_value(ro)
    delta_staeckels.append(estimateDeltaStaeckel(
        pot, R, z))
    
delta_interp = NearestNDInterpolator(Rz_grid, 
                                     delta_staeckels)
# def fast_actions():

deltas = delta_interp(w0s.cylindrical.rho.to_value(u.kpc),
                      w0s.z.to_value(u.kpc))

o = gala_to_galpy_orbit(w0s)
aAS = actionAngleStaeckel(pot=pot, delta=deltas)
actions = np.squeeze(aAS(o)).T * ro * vo<jupyter_output><empty_output><jupyter_text>Compare to Sanders & Binney actions<jupyter_code>sb_aaf = at.Table.read('../cache_new_zsun/aaf-1.0.fits')
sb_aaf = sb_aaf[np.isin(sb_aaf['APOGEE_ID'], t['APOGEE_ID'])]

assert len(sb_aaf) == len(t)
sb_aaf = sb_aaf[np.argsort(sb_aaf['APOGEE_ID'])] 
assert np.all(t['APOGEE_ID'] == sb_aaf['APOGEE_ID'])
sb_actions = sb_aaf['actions']
actions.shape, sb_actions.shape
from scipy.stats import binned_statistic
from astropy.stats import median_absolute_deviation
k = 0
for k in [0, 2]:
    sb_J = sb_actions[:, k]
    J = actions[:, k]
    mask = np.isfinite(sb_J) & np.isfinite(J)
    sb_J = sb_J[mask]
    J = J[mask]
    stat = binned_statistic(np.log10(sb_J), 
                            (J - sb_J) / sb_J,
                            statistic=lambda x: 1.5 * median_absolute_deviation(x),
                            bins=np.arange(-1, 3, 0.1))
    bincen = 0.5 * (10 ** stat.bin_edges[:-1] + 10 ** stat.bin_edges[1:])
    
    fig = plt.figure()
    plt.plot(sb_J, (J - sb_J) / sb_J,
             alpha=0.1, ls='none', ms=2, mew=0)

    plt.plot(bincen, stat.statistic)

    plt.xscale('log')
    plt.xlim(0.1, 2000)
    plt.ylim(-1, 1)
    
    fig.set_facecolor('w')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.16170861937452327, 'fraction_numerical': 0.028604118993135013, 'mean_word_length': 2.8630338733431517, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '2', 'dir': 'jupyter-structured-clean-dedup', 'id': '5890', 'license': 'MIT', 'n_tokens_mistral': 1184, 'n_tokens_neox': 1134, 'n_words': 177, 'path': 'notebooks/Staeckel-Delta-grid.ipynb', 'repo_name': 'adrn/thrift-shop', 'size': '2622'}