<jupyter_start><jupyter_text>Getting started with KinesisIn this notebook, we will generate a mock cluster, observe (add noise) and fit to see if we recover the true parameters of its internal kinematics.Notes on dependences- [gapipes](https://github.com/smoh/gapipes): for custom pandas accessor `g` to get covariance matrices, astropy coordinate objects..- [arviz](https://github.com/arviz-devs/arviz): for visualization<jupyter_code>%matplotlib inline
import matplotlib.pyplot as plt
from matplotlib import colors
import pandas as pd
import numpy as np
import astropy.units as u
import astropy.coordinates as coords
import arviz as az

# project dependencies
import gapipes as gp
import kinesis as kn
np.random.seed(18324)<jupyter_output><empty_output><jupyter_text>Making and observing a mock clusterFirst, let's make a cluster with some fiducial parameters.<jupyter_code>N = 150  # number of sources
b0 = np.array([17.7, 41.2, 13.3])    # pc
v0 = np.array([-6.32, 45.24, 5.30])  # [vx, vy, vz] in km/s
sigv = 1.  # dispersion, km/s

cl = kn.Cluster(v0, sigv, b0=b0)<jupyter_output><empty_output><jupyter_text>The cluster currently has no members. We need to 'sample' the cluster in order to populate it. There are two methods for sampling a cluster:- sample_sphere: sampling uniformly within a maximum radius- sample_at: sample at exactly the positions you want by giving it astropy coordinates.Let's do uniform sphere for now. All normal methods of `Cluster` will return itself with modified attributes for method chaining.Sampling will modify its `members` from `None` to an instance of `ClusterMembers`.<jupyter_code>cl.sample_sphere(N=N, Rmax=5)
cl.members<jupyter_output><empty_output><jupyter_text>There are two Gaia-like DataFrame associated with `ClusterMembers`:- `truths`: this is the noise-free Gaia data consistent with the `Cluster`'s internal motions- `observed`: this is the noise-added Gaia data. If the `ClusterMembers` has not been `observe`d, it will be `None`.Now let's `observe` the cluster. In order to observe, you need to specify the 3x3 covariance matrix for (parallax, pmra, pmdec) with `cov` keyword. If `cov` is (3,3) array, this will be assumed for all stars. You can also specify different covariance matrix for different stars by specifying `cov` to be (N, 3, 3) array. The noise model is assumed to be Gaussian.It is also possible to specify covariance by giving another Gaia-like data frame using `error_from` keyword. The following columns needed to construct to covariance matrix are expected: `parallax_error`, `pmra_error`, `pmdec_error`, `parallax_pmra_corr`, `parallax_pmdec_corr`, `pmra_pmdec_corr`. This is useful, for example, if you want to simulate data as noisy as the actual Gaia data you want to model.Radial velocity errors should be specified separately with `rv_error` keyword, which should be 1-d array with `N` elements, each specifying the Gaussian standard deviation in km/s. If any element in `rv_error` is `np.nan`, it is assumed that the star does not have RV measurement.This will add noise to the astrometry and radial velocities of the cluster members, and the noise-added Gaia-like data frame is now accessible with `cl.members.observed`. Let's specify simply 0.01 mas or mas/yr errors for parallax and pmra/pmdec and no available RVs.<jupyter_code>cl.observe(cov=np.eye(3)*0.01)
cl.members.observed.head()<jupyter_output><empty_output><jupyter_text>Fitting the mock dataNow let's try to fit the mock data that we generated.Under the hood, kinesis uses [stan](https://mc-stan.org/) using its python interface [pystan](https://pystan.readthedocs.io/en/latest/).<jupyter_code>fitter = kn.Fitter(include_T=False, recompile=True)
df = cl.members.observed.copy()
fit = fitter.fit(df, sample=False)
print(f"v0, sigv = {fit['v0']}, {fit['sigv']:.4f}")
print(f"diff from truth: {fit['v0']-v0}, {fit['sigv']-sigv:.4f}")
stanfit = fitter.fit(df)
stanfit.azfit
azfit=stanfit.azfit
az.summary(azfit, var_names=['v0','sigv'])
az.plot_forest(azfit, kind='ridgeplot', var_names=['sigv']);
az.plot_pair(azfit, var_names=['v0','sigv']);
v0_diff = stanfit['v0'] - cl.v0[None,:]

plt.figure(figsize=(4,3))
plt.hist(v0_diff, bins=64, histtype='step', density=True, label=['x','y','z']);
plt.xlabel("diff from truth [km/s]");
plt.legend(loc='upper left');
plt.axvline(0, c='gray', lw=.5);
az.plot_posterior(azfit, var_names=['sigv'], ref_val=cl.sigmav, figsize=(4,3));
# plt.figure(figsize=(4,3))
# g = np.einsum('ni,nij,nj->n', a-r['a_model'], np.linalg.inv(C), a-r['a_model'])

# plt.hist(g, 32, density=True);
# plt.hist(g, np.logspace(0, 4, 32), density=True);
# plt.xscale('log');
# plt.yscale('log')<jupyter_output><empty_output><jupyter_text>Basic cluster model with partial RVs<jupyter_code>N = 150
b0 = np.array([17.7, 41.2, 13.3])    # pc
v0 = np.array([-6.32, 45.24, 5.30])
sigv = 1.

cl = kn.Cluster(v0, sigv, b0=b0)\
       .sample_sphere(N=N, Rmax=5)\
       .observe(cov=np.eye(3)*0.01)

# Give random half of the stars RV with 0.5 km/s uncertainty
Nrv = int(N*0.5)
rv_error = 0.5
irand = np.random.choice(np.arange(N), size=Nrv, replace=False)

df = cl.members.observed.copy()
df['radial_velocity'] = np.random.normal(cl.members.truth['radial_velocity'].values, rv_error)
df['radial_velocity_error'] = rv_error
df = cl.members.observed.copy()
fit = fitter.fit(df, sample=False)
print(f"v0, sigv = {fit['v0']}, {fit['sigv']:.4f}")
print(f"diff from truth: {fit['v0']-v0}, {fit['sigv']-sigv:.4f}")
stanfit_partialrv = fitter.fit(df)
# workaround for a bug in arviz instantiating InferenceData
d = stanfit_partialrv.extract(permuted=False, pars=['v0', 'sigv'])
d = {k:np.swapaxes(v, 0, 1) for k,v in d.items()}
azfit_partialrv = az.from_dict(d)
v0_diff0 = stanfit['v0'] - cl.v0[None,:]
v0_diff = stanfit_partialrv['v0'] - cl.v0[None,:]

plt.figure(figsize=(4,3))
color = ['tab:blue', 'tab:green', 'tab:orange']
plt.hist(v0_diff0, bins=64, histtype='step', density=True, color=color, lw=.5);
plt.hist(v0_diff, bins=64, histtype='step', density=True, color=color, label=['x','y','z'], lw=2);
plt.xlabel("diff from truth [km/s]");
plt.legend(loc='upper left');
plt.axvline(0, c='gray', lw=.5);
axs = az.plot_posterior(azfit, var_names=['sigv'], ref_val=cl.sigmav, figsize=(4,3))
az.plot_posterior(azfit_partialrv, var_names=['sigv'], ax=axs);<jupyter_output><empty_output><jupyter_text>Cluster with rotation<jupyter_code>N = 150
b0 = np.array([17.7, 41.2, 13.3])    # pc
v0 = np.array([-6.32, 45.24, 5.30])
sigv = 1.
omegas = [40., 20., 50.]

cl = kn.Cluster(v0, sigv, b0=b0, omegas=omegas)\
       .sample_sphere(N=N, Rmax=15)\
       .observe(cov=np.eye(3)*0.001)
cl0 = kn.Cluster(v0, sigv, b0=b0).sample_at(cl.members.truth.g.icrs)

# Give random half of the stars RV with 0.5 km/s uncertainty
Nrv = int(N*0.5)
rv_error = 0.5
irand = np.random.choice(np.arange(N), size=Nrv, replace=False)

df = cl.members.observed.copy()
df['radial_velocity'] = np.random.normal(cl.members.truth['radial_velocity'].values, rv_error)
df['radial_velocity_error'] = rv_error
m = kn.Fitter(include_T=True)
fit = m.fit(df, sample=False, b0=b0)
print(f"v0, sigv = {fit['v0']}, {fit['sigv']:.4f}")
print(f"diff from truth: {fit['v0']-v0}, {fit['sigv']-sigv:.4f}")
print(f"{fit['T_param']}")
print(f"{cl.T}")
# omegax = 0.5*(r['T_param'][2, 1] - r['T_param'][1, 2])
# omegay = 0.5*(r['T_param'][0, 2] - r['T_param'][2, 0])
# omegaz = 0.5*(r['T_param'][1, 0] - r['T_param'][0, 1])

# w1 = 0.5*(r['T_param'][2, 1] + r['T_param'][1, 2])
# w2 = 0.5*(r['T_param'][0, 2] + r['T_param'][2, 0])
# w3 = 0.5*(r['T_param'][1, 0] + r['T_param'][0, 1])
# w4 = r['T_param'][0, 0]
# w5 = r['T_param'][1, 1]
# kappa = w4 + w5 + r['T_param'][2, 2]
# print(omegax, omegay, omegaz)
# print(w1, w2, w3)
# print(w4, w5)
# print(kappa)
stanfit = m.fit(df, b0=b0)
azfit = az.from_pystan(stanfit)
v0_diff = stanfit['v0'] - cl.v0[None,:]

plt.figure(figsize=(4,3))
color = ['tab:blue', 'tab:green', 'tab:orange']
plt.hist(v0_diff, bins=64, histtype='step', density=True, color=color, label=['x','y','z'], lw=2);
plt.xlabel("diff from truth [km/s]");
plt.legend(loc='upper left');
plt.axvline(0, c='gray', lw=.5);
fig, ax = plt.subplots(3, 3, figsize=(8, 8), sharex=True, sharey=True)
fig.subplots_adjust(bottom=0.1, top=0.95, right=0.95, left=0.15)
ax = ax.ravel()
for cax, cT, truth in zip(ax, stanfit["T_param"].reshape((-1, 9)).T, cl.T.ravel()):
    cax.hist(cT, bins=32, density=True, histtype="step")
    cax.axvline(truth, c="k")
    cax.axvline(0, c="gray", lw=0.5)
fig.text(0.55, 0.05, "m/s/pc", ha="center", va="center", size=20)
fig.text(0.05, 0.55, "Density", ha="center", va="center", rotation=90, size=20)
fig.savefig("mock_posterior_T.png")<jupyter_output><empty_output><jupyter_text>Cluster with rotation - ideal caseTest the most ideal case when the cluster has small dispersion and all velocities are measured to high precision.<jupyter_code>N = 150
b0 = np.array([17.7, 41.2, 13.3])    # pc
v0 = np.array([-6.32, 45.24, 5.30])
sigv = 0.1 # small dispersion
omegas = [40., 20., 50.]

cl = kn.Cluster(v0, sigv, b0=b0, omegas=omegas)\
       .sample_sphere(N=N, Rmax=15)\
       .observe(cov=np.eye(3)*0.001)
cl0 = kn.Cluster(v0, sigv, b0=b0).sample_at(cl.members.truth.g.icrs)

# Give all stars observed RV with small uncertainty
Nrv = int(N)
rv_error = 0.1
irand = np.random.choice(np.arange(N), size=Nrv, replace=False)

df = cl.members.observed.copy()
df['radial_velocity'] = np.random.normal(cl.members.truth['radial_velocity'].values, rv_error)
df['radial_velocity_error'] = rv_error
m = kn.Fitter(include_T=True)
fit = m.fit(df, sample=False, b0=b0)
print(f"v0, sigv = {fit['v0']}, {fit['sigv']:.4f}")
print(f"diff from truth: {fit['v0']-v0}, {fit['sigv']-sigv:.4f}")
print(f"{fit['T_param']}")
print(f"{cl.T}")
stanfit = m.fit(df, b0=b0)
azfit = az.from_pystan(stanfit, coords={'x':['v1','v2','v3']}, dims={'v0':['x']})
azfit.posterior
azfit.sample_stats
v0_diff = stanfit['v0'] - cl.v0[None,:]

plt.figure(figsize=(4,3))
color = ['tab:blue', 'tab:green', 'tab:orange']
plt.hist(v0_diff, bins=64, histtype='step', density=True, color=color, label=['x','y','z'], lw=2);
plt.xlabel("diff from truth [km/s]");
plt.legend(loc='upper left');
plt.axvline(0, c='gray', lw=.5);
az.plot_posterior(azfit, var_names=['sigv']);
az.plot_posterior(azfit, var_names=['T_param'], coords={'T_param_dim_0':[0], 'T_param_dim_1':[0]});
fig, ax = plt.subplots(3, 3, figsize=(8, 8), sharex=True, sharey=True)
fig.subplots_adjust(bottom=0.1, top=0.95, right=0.95, left=0.15)
ax = ax.ravel()
for cax, cT, truth in zip(ax, stanfit["T_param"].reshape((-1, 9)).T, cl.T.ravel()):
    cax.hist(cT, bins=32, density=True, histtype="step")
    cax.axvline(truth, c="k")
    cax.axvline(0, c="gray", lw=0.5)
fig.text(0.55, 0.05, "m/s/pc", ha="center", va="center", size=20)
fig.text(0.05, 0.55, "Density", ha="center", va="center", rotation=90, size=20)
fig.savefig("mock_posterior_T.png")
az.summary(azfit, var_names=['v0', 'sigv', 'T_param'])<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.18596810933940774, 'fraction_numerical': 0.03917995444191344, 'mean_word_length': 5.67639902676399, 'pattern_counts': {'":': 0, '<': 33, '<?xml version=': 0, '>': 34, 'https://': 4, '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': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '90820', 'license': 'MIT', 'n_tokens_mistral': 4664, 'n_tokens_neox': 4388, 'n_words': 1036, 'path': 'docs/tutorial.ipynb', 'repo_name': 'smoh/kinesis', 'size': '10975'}