{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dfe0fd3a-8c35-4763-af55-cbbcd3793ea8",
   "metadata": {},
   "source": [
    "# 8. Coloring - BETA -\n",
    "This notebook takes a look at coloring (fitting population data). Upon its first start, it will create a model for the example galaxy FCC167. Based on that model and observational data, the notebook will fit age and metallicity of the populations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9deac14-d89b-48e5-8eb8-db7d2a0d8f91",
   "metadata": {},
   "source": [
    "Please note that this notebook is still in beta, as is the DYNAMITE coloring functionality. For more information on the options and functionality of the methods described here, please refer to the API docs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f463a057-9e09-4597-9d11-b2147d466642",
   "metadata": {},
   "source": [
    "## 8.1. Prerequisites"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00ac1a55-5570-4d57-93d8-b56cba9d6914",
   "metadata": {},
   "source": [
    "Import the required modules."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91267ba7-8793-4910-996a-10bd5442be3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import pymc as pm  # only necessary for posterior plots\n",
    "from vorbin.voronoi_2d_binning import voronoi_2d_binning\n",
    "import dynamite as dyn\n",
    "\n",
    "print('DYNAMITE version', dyn.__version__)\n",
    "# print('    installed at ', dyn.__path__)  # Uncomment to print the complete DYNAMITE installation path"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4691fa8a-40a8-4c37-a95d-216dca08fd28",
   "metadata": {},
   "source": [
    "Run the DYNAMITE model. The configuration file `FCC167_config.yaml` fixes all parameter to a specific value, resulting in a single model. For performance, the orbit library is rather small, so the calculation time should only take a few minutes. Note that the model will not be run if already existing on the disk (e.g., if the notebook is re-run)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80c78a71-b1a7-4ae8-a488-67b64ea5b4a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "fname = 'FCC167_config.yaml'\n",
    "c = dyn.config_reader.Configuration(fname, reset_logging=True)\n",
    "# c = dyn.config_reader.Configuration(fname, reset_logging=True, reset_existing_output=True)  # uncomment to force fresh run\n",
    "_ = dyn.model_iterator.ModelIterator(c)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f6e05ba-cb0c-4b06-92d6-cc881a242794",
   "metadata": {},
   "source": [
    "## 8.2. Voronoi phase space binning\n",
    "The model's orbits in the phase space of circularity $\\lambda_z$ versus time-averaged orbital radius $r$ are binned into $N_\\mathrm{bundle}$ Voronoi orbit bundles. Each Voronoi orbit bundle amounts to a certain minimum total weight."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a508b446-a1e6-4b17-b479-1cc06ed923c2",
   "metadata": {},
   "source": [
    "First, we need to choose the underlying resolution in $r$ and $\\lambda_z$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01952fa0-4fd1-4b3a-939d-6b756539e381",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Number of desired r and lambda_z bins\n",
    "nr = 6  # 20\n",
    "nl = 7  # 41"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22391b5c-ba9a-4dd9-8b00-3b29954d469d",
   "metadata": {},
   "source": [
    "Optional step: Before binning the orbits, let's have a look at the $r$ - $\\lambda_z$ phase space. The standard DYNAMITE orbit distribution plot with the parameter `force_lambda_z=True` will plot the distribution of all orbits. Note that the plot title is incorrect (all orbits are included in the plot, not only the short axis tubes)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a88b30f-5a8a-4f5b-8ff8-c870c1168d8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "plotter = dyn.plotter.Plotter(c)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c4c9901-818c-4a52-83a3-fdf9f40734cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "# NOTE: when using force_lambda_z=True, then the title of the orbit-distribution plot is incorrect: all orbits are shown in this distribution - not only short-axis tubes!\n",
    "#       If force_lambda_z is True, then the entire orbit distribution is collapsed onto (r, lambda_z) space.\n",
    "#       This is done by forcing all orbits to be classified as short axis-tube orbits.\n",
    "fig2 = plotter.orbit_distribution(model=None,\n",
    "                                  minr=None,\n",
    "                                  maxr=None,\n",
    "                                  r_scale='linear',\n",
    "                                  nr=nr,\n",
    "                                  nl=nl,\n",
    "                                  orientation='vertical',\n",
    "                                  subset='short',\n",
    "                                  force_lambda_z=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb008326-4d34-41fa-905a-d6a679b87053",
   "metadata": {},
   "source": [
    "The Voronoi binning of orbits in the radius-circularity phase space is done by a method in the `dyn.coloring.Coloring` class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34f6c8d4-e203-4a99-bfd7-dfe18a21c383",
   "metadata": {},
   "outputs": [],
   "source": [
    "coloring = dyn.coloring.Coloring(c, nr=nr, nl=nl)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03ed52f9-b83f-4b2d-9686-4989bc69f6d1",
   "metadata": {},
   "source": [
    "The original `n_orbits` orbit bundles are binned into fewer $N_\\mathrm{bundle}=$`n_bundle` Voronoi orbit bundles with each of these Voronoi bundles accounting for a weight of at least `vor_weight`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "054d6600-ba6d-49c3-90ae-932806b2c080",
   "metadata": {},
   "outputs": [],
   "source": [
    "vor_weight = 0.05  # define the desired (minimum) total orbital weight in each Voronoi bin"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ea2315e-944c-41bc-afa4-616b017e1f42",
   "metadata": {},
   "source": [
    "The method `coloring.bin_phase_space()` performs the binning of orbits, based on the best-fitting model so far (in our case this is the only model). The result of the binning is a tuple `(vor_bundle_mapping, phase_space_binning)`:\n",
    "```\n",
    "vor_bundle_mapping :    np.array of shape (n_bundle, n_orbits)\n",
    "                        Mapping between the \"original\" orbit bundles and the Voronoi\n",
    "                        orbit bundles: vor_bundle_mapping(i_bundle, i_orbit) is the\n",
    "                        fraction of i_orbit assigned to i_bundle, multiplied by i_orbit's weight.\n",
    "phase_space_binning :   dict\n",
    "                        'out':  np.array of shape (3, n_bundle), the Voronoi binning output:\n",
    "                                weighted Voronoi bin centroid coordinates r_bar, lambda_bar\n",
    "                                and Voronoi bin total weights\n",
    "                        'map':  np.array of shape (nr*nl,) the phase space mapping:\n",
    "                                Voronoi bin numbers for each input bin\n",
    "```\n",
    "\n",
    "As the binning can be time-consuming, `coloring.bin_phase_space()` will write the binning result to the model directory so that subsequent calls with the same parameters for the same model will read the existing binning from disk if `use_cache=True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c87996d-d461-4f61-b6bc-3c646ace389a",
   "metadata": {},
   "outputs": [],
   "source": [
    "vor_bundle_mapping, phase_space_binning = coloring.bin_phase_space(model=None,\n",
    "                                                                   vor_weight=vor_weight,\n",
    "                                                                   vor_ignore_zeros=True,\n",
    "                                                                   make_diagnostic_plots=True,\n",
    "                                                                   extra_diagnostic_output=True,\n",
    "                                                                   cvt=False,\n",
    "                                                                   wvt=False,\n",
    "                                                                   use_noise=True,\n",
    "                                                                   use_cache=False)  # set use_cache=False to force fresh binning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe0b8c61-a541-4acd-9e4e-ef47b4c3834b",
   "metadata": {},
   "source": [
    "Let's visualize the Voronoi bundle mapping, starting with the probability density of stellar orbits in the circularity-radius phase space, overlayed by the Voronoi binning scheme:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85b1c71c-fba4-46c4-8477-c67ed460cbe2",
   "metadata": {},
   "outputs": [],
   "source": [
    "coloring.orbit_bundle_plot(phase_space_mapping=phase_space_binning['map']);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33374c87-c4d6-47e1-90e6-8edc6e34962e",
   "metadata": {},
   "source": [
    "As a quick check, let's plot the orbit weights in each $(r, \\lambda_z)$ Voronoi bin. Note that most likely the exact desired weight of the Voronoi bins will not be met, especially for small orbit libraries."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df613cac-1fb9-4de4-8dfa-4cfe203fbd00",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.bar([str(i) for i in range(phase_space_binning['out'].shape[1])], phase_space_binning['out'][2])\n",
    "plt.xlabel('Voronoi bin ID')\n",
    "plt.ylabel('Weight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b00fee4-0749-48c4-aeb3-4189869aa534",
   "metadata": {},
   "source": [
    "Now, let's visualize how much weight each original orbit bundle contributes to each Voronoi orbit bundle.\n",
    "- Original orbit bundles with zero weight in the specific model will not contribute to any of the Voronoi orbit bundles.\n",
    "- Each original orbit bundle consists of one actual orbit if `dithering=1` in the configuration file's `orblib_settings`, but will comprise multiple orbits if ``dithering > 1``. In the latter case, one original orbit bundle can be split among neighboring bins in the radius-circularity phase space and can hence contribut to multiple (neighboring) Voronoi orbit bundles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac6e86d1-f552-49d9-acc3-602600c5bcdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(24,4))\n",
    "plt.gca().set_title('Weight that each orbit bundle contributes to Voronoi orbit bundles')\n",
    "bundle_map = np.copy(vor_bundle_mapping)\n",
    "bundle_map[bundle_map == 0] = np.nan\n",
    "plt.pcolormesh(np.log10(bundle_map, where=bundle_map is not np.nan), shading='flat', cmap='Greys')\n",
    "plt.xlabel('Original orbit bundle id')\n",
    "plt.ylabel('Voronoi orbit bundle id')\n",
    "plt.colorbar(label='log Weight')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8cba5c1b-a2e1-4183-959d-f098d280f46f",
   "metadata": {},
   "source": [
    "For the subsequent analysis, we will need to know how much mass (for mass-weighted models) or flux (for light-weighted models) each Voronoi orbit bundle contributes to each spatial bin. As the concept of orbit bundles per se is independent of coloring (fitting population data), the associated method is part of the `dyn.analysis.Analysis` class:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c153d06c-29c2-4e74-a590-d762a8a4b34e",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = dyn.analysis.Analysis(c)  # orbit bundle maps are residing in the Analysis class"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c4a333b-9492-446a-b8e6-19177531c867",
   "metadata": {},
   "source": [
    "The method ``get_orbit_bundle_maps()`` returns an astropy table with as many rows as there are spatial bins and the number of columns is the number of Voronoi orbit bundles plus 1: Each Voronoi orbit bundle's contribution corresponds to one column and at the end there is the combined bundle's contribution in column 'flux_all'."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9a9206a-b7b6-4457-b871-47662caf01d8",
   "metadata": {},
   "source": [
    "Setting the parameter `create_figure=True` will create a figure of the individual orbit bundles' contributions and return it along with the bundle maps (see the commented-out example below - try it...). Setting `normalize=True` will normalize the mass (flux) contributions so that in each spatial bin the sum of all orbit bundles' contributions is 1. We activate this option because we will need normalized flux in the Bayesian statistical analysis further down."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6cce8eeb-ea8e-40fd-aa42-00f95b67627a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# bundle_maps, bundle_figure = a.get_orbit_bundle_maps(pop_set=0,\n",
    "#                                                      bundle_mapping=vor_bundle_mapping,\n",
    "#                                                      normalize=True,\n",
    "#                                                      create_figure=True)  # comment-in to view the plots\n",
    "bundle_maps = a.get_orbit_bundle_maps(pop_set=0,\n",
    "                                      bundle_mapping=vor_bundle_mapping,\n",
    "                                      normalize=True,\n",
    "                                      create_figure=False)  # calculation only, won't display the plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1fba4862-b5b8-43a6-b28b-5cd6fac524f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'{type(bundle_maps) = }\\n{len(bundle_maps) = }\\n{bundle_maps.colnames = }')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15bc5b53-8d3d-4509-a564-4ab0ed4fe18a",
   "metadata": {},
   "source": [
    "As you see, the `bundle_maps` contain the Voronoi orbit bundles' contributions along with the aggregate map. For the subsequent calculations, we will only need the data in columns corresponding to the individual Voronoi bundles which we store in `flux_data_norm`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "325881d9-9c0d-47ec-89ed-2f7a4cb7dc30",
   "metadata": {},
   "outputs": [],
   "source": [
    "flux_data_norm = np.array([bundle_maps[a] for a in bundle_maps.columns if a != 'flux_all']).T\n",
    "print(f'{flux_data_norm.shape = }') # (n_spatial_bins, n_bundle)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b1d9f47-4815-4b76-9703-242c3bd305e1",
   "metadata": {},
   "source": [
    "## 8.3. Bayesian statistical analysis"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "acb3f50e-52a8-4df3-adae-2b4e1f5fae28",
   "metadata": {},
   "source": [
    "Fitting age and metallicity essentially follows the procedure described in **Zhu et al., 2020, MNRAS, 496, 1579**. For brevity, only the priors \"R1\" for both age and metallicity will be used. In the following sections, equation numbers refer to the corresponding equations in that paper."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb6201da-6b43-483d-a268-f2d8052cf3f1",
   "metadata": {},
   "source": [
    "The fitting is done via Bayesian inference as provided by the Python package PyMC. The model uses a truncated normal or lognormal distribution for the prior of the observed quantity and a Student's t distribution with fixed $\\sigma$ (Half-Cauchy distributed with $\\beta=5$) and $\\nu$ (Exponential distributed with parameter 1/30) parameters for the likelihood of the observed data. The solution method uses the Markov chain Monte Carlo (MCMC) sampling algorithm NUTS (No-U-Turn Sampler), initialized with the ADVI(Automatic Differentiation Variational Inference) method.\n",
    "\n",
    "There is no need to directly interact with PyMC. DYNAMITE users can call the `coloring.fit_bayesian()` method that has many settings predefined already (see below)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02203c67-dce0-433c-a9fc-deec9547e090",
   "metadata": {},
   "source": [
    "The chain is initialized with the ADVI (automatic differentiation variational inference) with 200000 draws and 2500 tuning steps. The last 500 MCMC steps are used to get the mean and standard deviation of the age and metallicity, respectively:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "985b6062-ebb2-4900-8324-6d91365ead4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample = {'n_draws': 500,\n",
    "          'n_tune': 2500,\n",
    "          'advi_init': 200000}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ee5f199-1dde-4a29-8f72-bb8bbf11bb92",
   "metadata": {},
   "source": [
    "Finally, we need to assign some DYNAMITE data structures (note that currently, we only support one population dataset `stars.population_data[0]`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61caabba-06c2-4fa5-80dc-9112199cb1b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "stars = c.system.get_unique_triaxial_visible_component()\n",
    "pops = stars.population_data[0]\n",
    "age, dage, met, dmet = [pops.get_data()[i] for i in ('age', 'dage', 'met', 'dmet')]  # dage and dmet will not be used"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "207d4324-ea9e-43f8-bed2-b7e0da7b9f7c",
   "metadata": {},
   "source": [
    "### 8.3.1. Age"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c8b07f6-bc1d-4f17-bacb-9e4b4467f146",
   "metadata": {},
   "source": [
    "The age prior is a bounded normal distribution $f(t_k|\\mu_k,\\sigma_k)$ with a lower boundary of 0, an upper boundary of 20, and $\\mu_k=\\mathrm{Randn}(<t_\\mathrm{obs}>,2\\sigma(t_\\mathrm{obs}))$ and $\\sigma_k=2\\sigma(t_\\mathrm{obs})$, see Eq. (11)-(13):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bba6b418-9ae0-43d2-a914-1d68fadbd6fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "prior_t = {'mu': np.random.normal(age.mean(), 2 * age.std(), size=len(vor_bundle_mapping)),  # Eq. (12)\n",
    "           'sigma': 2 * age.std(),                                                           # Eq. (13)\n",
    "           'lower': 0,\n",
    "           'upper': 20}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16c16e45-71a4-461d-ad22-096bfe7e3d3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_t, trace_t = coloring.fit_bayesian(prior_dist='normal',\n",
    "                                         prior_par=prior_t,\n",
    "                                         flux_data_norm=flux_data_norm,\n",
    "                                         obs_data=age,\n",
    "                                         sample=sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "407ed96b-6ebd-4a7c-a04e-b0c2dbf9399c",
   "metadata": {},
   "source": [
    "Optional: inspect the Bayesian model details via"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5792f2c4-3919-4c65-9659-1e35b30ec9fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_t"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a99890b4-8ee6-4777-a17b-234c22214701",
   "metadata": {},
   "source": [
    "Note that per default PyMC per default uses as many chains as there are physical CPU cores available (e.g., 4). There are 3000 draws in each chain, corresponding to 2500 tuning steps and 500 draws for the results.\n",
    "\n",
    "The resulting age values are accessible via `trace_t.posterior['pop']`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9a8fdcd-cc86-45c4-9739-5331f9f2c8c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "trace_t.posterior['pop']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a1131e26-ef57-4d23-9bfe-48ed7a0168e9",
   "metadata": {},
   "source": [
    "We store it in symbol `age_posterior` for later. It is a data structure compatible with an array of shape (< number of chains >, < number of draws >, <$N_\\mathrm{bundle}$>). Consequently, the mean and error of each Voronoi orbit bundle's age are given by `age_mean = age_posterior.mean(axis=(0,1))` and `age_err = age_posterior.std(axis=(0,1))`, respectively:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5d8f844-1fd8-4887-a4cd-e0c34151a311",
   "metadata": {},
   "outputs": [],
   "source": [
    "age_posterior = trace_t.posterior['pop']\n",
    "age_mean = age_posterior.mean(axis=(0,1))\n",
    "age_err = age_posterior.std(axis=(0,1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "933fb150-2dfb-411c-a11c-5042e498a879",
   "metadata": {},
   "outputs": [],
   "source": [
    "age_mean"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cdbc2777-8df5-401a-abd8-afa142f1e990",
   "metadata": {},
   "source": [
    "Optional: display the posterior plot (requires the PyMC module to be imported). To better format the figure, we use `tight_layout()`, but as `plot_trace()` returns an array of `Axes` objects, we first must extract the figure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55c893b0-70af-4769-9e9b-33fe63c4f13d",
   "metadata": {},
   "outputs": [],
   "source": [
    "axes_array = pm.plot_trace(trace_t, combined=True)\n",
    "figure = axes_array[0,0].figure\n",
    "figure.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99123962-2f8e-4402-bc98-43e3e26381af",
   "metadata": {},
   "source": [
    "### 8.3.2. Metallicity"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5a38aff-b709-408f-8516-e5a1110ae3cd",
   "metadata": {},
   "source": [
    "The metallicity prior is a bounded lognormal distribution $f(Z_k|\\mu_k,\\sigma_k)$ with a lower boundary of 0, an upper boundary of 10, and $\\mu_k=\\ln(\\mathrm{Randn}(<Z_\\mathrm{obs}>,\\sigma(Z_\\mathrm{obs})))$ and $\\sigma_k=\\sigma(Z_\\mathrm{obs})$, see Eq. (16)-(18):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c86807b4-4cc3-4f4e-b0e8-0923d80638f7",
   "metadata": {},
   "outputs": [],
   "source": [
    "prior_z = {'mu': np.log(np.random.normal(met.mean(), met.std(), size=len(vor_bundle_mapping))),  # (17)\n",
    "           'sigma': met.std(),  # (18)\n",
    "           'lower': 0,\n",
    "           'upper': 10}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1108199-0149-4281-a2a1-f33dae6714dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_z, trace_z = coloring.fit_bayesian(prior_dist='lognormal',\n",
    "                                         prior_par=prior_z,\n",
    "                                         flux_data_norm=flux_data_norm,\n",
    "                                         obs_data=met,\n",
    "                                         sample=sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a3f3f04-caa2-4637-9149-5cacc003c3e7",
   "metadata": {},
   "source": [
    "Optional: inspect the Bayesian model details via"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "196d449d-f88a-4107-8b99-c80712d9765d",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_z"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d147875b-a1d7-46e7-b69b-4fca5633ee1f",
   "metadata": {},
   "source": [
    "In analogy to above, the resulting metallicity values are accessible via `trace_z.posterior['pop']`, a data structure compatible with an array of shape (< number of chains >, < number of draws >, <$N_\\mathrm{bundle}$>). After storing it in `met_posterior`, the mean and error of each Voronoi orbit bundle's metallicity are given by `met_mean = met_posterior.mean(axis=(0,1))` and `met_err = met_posterior.std(axis=(0,1))`, respectively:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a345131b-90b5-4409-b8af-54045c852e0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "met_posterior = trace_z.posterior['pop']\n",
    "met_mean = met_posterior.mean(axis=(0,1))\n",
    "met_err = met_posterior.std(axis=(0,1))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e7a29456-8723-4bd7-968e-bc8e91a20c62",
   "metadata": {},
   "source": [
    "Optional: display the posterior plot (requires the PyMC module to be imported):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "608a07f3-ffa5-4dbf-828d-49cc1bd92e67",
   "metadata": {},
   "outputs": [],
   "source": [
    "axes_array = pm.plot_trace(trace_z, combined=True)\n",
    "figure = axes_array[0,0].figure\n",
    "figure.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65097809-2b11-4479-a66d-affabc1cd0a0",
   "metadata": {},
   "source": [
    "## 8.4. Results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45c913f4-6f95-4b7a-8810-e537473c0ae3",
   "metadata": {},
   "source": [
    "### 8.4.1. Check how the model matches the data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "065d223a-0f96-4e1d-bc07-9da4d4b43fb8",
   "metadata": {},
   "source": [
    "Here, we plot the observed age and metallicity maps along with the maps of the fitted age and metallicity data. Note that for the observed data (first row), the errors refer to the read-in observation errors and for the fitted data (second row), the error columns refer to the standard deviations of the posteriors for age and metallicity, respectively. The population maps are consistent with those used in the kinematic map plots. Also, the residuals are defined as `residual = (model - data) / data_error`, consistent with the kinematic maps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4fd90af-8f88-4f9b-a6d8-c391b10b0b3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "coloring.pop_maps(pop_data={'age': 'Stellar age [Gyr]', 'met': r'Metallicity $Z/Z_\\odot$'},  # choose the population datasets to be plotted\n",
    "                  model_data=[age_mean, age_err, met_mean, met_err],                         # calculated from the posteriors\n",
    "                  flux_norm=flux_data_norm,                                                  # resulting from get_orbit_bundle_maps() above\n",
    "                  cbar_lims='data');                                                         # 'data', 'model', or 'auto'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0be7579-11ce-4232-8a1e-0b9073b3175b",
   "metadata": {},
   "source": [
    "Including the error columns is optional. In the function call above, model_data has twice the number of entries in pop_data, which results in every other column in pop_data being interpreted as the error of the preceding dataset. If model_data and pop_data have the same lengths, just the datasets, but no errors are plotted:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77fc64f6-2dab-42e3-8599-dca61d27ae5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "coloring.pop_maps(pop_data={'age': 'Stellar age [Gyr]', 'met': r'Metallicity $Z/Z_\\odot$'},\n",
    "                  model_data=[age_mean, met_mean],\n",
    "                  flux_norm=flux_data_norm,\n",
    "                  cbar_lims='data');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "588cfb4c-f1f5-42b4-81b5-c0cc60366ed8",
   "metadata": {},
   "source": [
    "### 8.4.2. Visualize the relation between population datasets (here: age-metallicity relation, AMR)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1ce39b3-a72c-45fa-92b8-454735718074",
   "metadata": {},
   "source": [
    "Let's plot the relation of two population datasets (e.g. age vs metallicity, AMR) for a set of stellar bundles, averaged over multiple MCMC chains and draws. The plot includes the probability density distribution in the population data values and a scatter plot with diamonds indicating the average population data values of each orbit bundle. The diamond sizes are proportional to the orbit bundles' weights. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73ff68e4-235f-4714-b5fa-d030f67584ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f'We will smooth over {age_posterior.shape[1]} posterior values.')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e44d55ea-d13d-4a9f-9c81-b01badae01dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "coloring.pop_pop_plot(age_posterior, met_posterior, phase_space_binning['out'][2],  # the third parameter is the bundle weights, see API docs\n",
    "                      x_label='Stellar age [Gyr]', y_label='$Z/Z_\\\\odot$',\n",
    "                      x_scale='linear', y_scale='linear',\n",
    "                      n_smooth=500);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf1c1a75-4c12-4c07-bf38-d3e5280c756f",
   "metadata": {},
   "source": [
    "### 8.4.3. Create an orbital decomposition plot for the population datasets"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "662814d4-4f15-4f18-8be5-c0c077472b5a",
   "metadata": {},
   "source": [
    "The orbital decomposition plot can deal with multiple models. The orbits' probability distribution and the population data in the phase space bins are then averaged over say, all models within a 1$\\sigma$ confidence level of the model hyperparameters. Here, we have only one model and will use that to demonstrate how to create the orbital decomposition plot using the (only) model in row 0 of the all_models table, along with the just calculated Voronoi orbit bundles and the estimates for the age and metallicity:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a5f7e1e-ef97-408b-be50-39b90f061617",
   "metadata": {},
   "outputs": [],
   "source": [
    "distr = coloring.get_pop_orbital_decomp(models=[c.all_models.get_model_from_row(0)],\n",
    "                                        vor_bundle_mappings=[vor_bundle_mapping],\n",
    "                                        pop_data=[[age_mean, met_mean]])\n",
    "print(f'{distr.shape = }, {{weights, age, metallicity}} x (number of r bins) x (number of lambda_z bins) x (number of models)')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26a1084f-6f91-4af5-8a38-ee33e42f975b",
   "metadata": {},
   "source": [
    "The intermediate result ``distr`` is a 4-dimensional numpy array. Its first dimension has three entries, corresponding to the orbital weight distribution plus the two population data distributions (age and metallicity). The second and third indices are the phase space bins in $r$ and $\\lambda_z$. The last index enumerates the number of models for which the weight and population data distributions are available (here, only one model)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bc4d4d6-b9c7-4bda-85d2-3637e338d351",
   "metadata": {},
   "source": [
    "The next step is to average over the models and to plot the data. For this, ``distr`` can directly be passed to the plotting method, along with the desired labels for the individual plots:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b331ddbf-fcef-4fba-962e-7b423ffa70d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "coloring.pop_decomp_plot(distr,\n",
    "                         plot_labels=[r'Probability density [$M_*$/unit]', 'Stellar age [Gyr]', r'Metallicity $Z/Z_\\odot$'],\n",
    "                         colorbar_scale=['linear','linear','linear']);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f7a2cb5-5795-4cb9-bbc0-406ea5d75067",
   "metadata": {},
   "source": [
    "### 8.4.4. Create a plot that shows the orbit probability distribution in (age, circularity) and the disk ratio vs age"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fd6ef945-71a0-48d8-a483-8e1013297148",
   "metadata": {},
   "source": [
    "We use `distr`, the result from the last section to extract the weight and age data and plot the orbit distribution in the (age, circularity) phase space, averaged over multiple DYNAMITE models. On top of that plot, we display the disk fraction as a function of the age and identify the the age at which the cold orbit fraction crosses 50% for the first time.\n",
    "The method ``circularity_pop_plot()`` expects the orbit bundle weights as the first parameter and the color distribution in the second. We can directly use ``distr[0]`` for the weights and ``distr[1]`` for the age data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d55fa4c4-7c1f-4886-b2f4-ceb0c3f38d56",
   "metadata": {},
   "outputs": [],
   "source": [
    "coloring.circularity_pop_plot(distr[0],\n",
    "                              distr[1],\n",
    "                              pop_label='Stellar age [Gyr]');  # note that per default pop_label='Stellar age [Gyr]' to account for the often-used circularity-age-plot"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6acb8b1e-5818-4507-b159-78b027a254a9",
   "metadata": {},
   "source": [
    "The method ``circularity_pop_plot()`` has quite a few arguments. For details, please see its docstring. As an illustration, here is an example that plots metallicity (`distr[2]`) instead of age. Feel free to experiment with a few more settings..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41c41bae-a459-452d-bdda-d48dc0e4adab",
   "metadata": {},
   "outputs": [],
   "source": [
    "coloring.circularity_pop_plot(distr[0],\n",
    "                              distr[2],\n",
    "                              pop_label=r'Metallicity $Z/Z_\\odot$',\n",
    "                              pop_scale='linear',\n",
    "                              prob_scale='linear',\n",
    "                              n_pop_bins=14,\n",
    "                              interpolation='spline16',  # try 'none', 'spline16',...\n",
    "                              disk_fraction=True);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a39a742-316d-4a28-a33b-c250879f41ca",
   "metadata": {},
   "source": [
    "## Appendix: experiments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "557826cc-ed06-4462-bc4f-3f451e83d386",
   "metadata": {},
   "outputs": [],
   "source": [
    "plotter.orbit_plot(model=c.all_models.get_model_from_row(0), Rmax_arcs=316);"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}