feat: Combined crn

README.md

@@ -1,6 +1,8 @@
 # Discovery
 [![DOI](https://zenodo.org/badge/DOI/10.5281/zenodo.17711453.svg)](https://doi.org/10.5281/zenodo.17711453)
 
+[Sphinx Documentation is here](https://nanograv.github.io/discovery/). The old README user guide is below as well.
+
 <img src="discovery.png" alt="Logo" width="300" align="left" style="margin-right: 20px;"/>
 
 _Discovery_ is a next-generation pulsar-timing-array data-analysis package, _built for speed_ on a [JAX](https://jax.readthedocs.io/en/latest/) backend that supports GPU execution and autodifferentiation.

docs/conf.py

@@ -27,6 +27,7 @@
     'sphinx.ext.intersphinx',
     'sphinx_autodoc_typehints',
     'numpydoc',
+    'myst_nb',
 ]
 
 # Napoleon settings for numpy-style docstrings
@@ -49,6 +50,9 @@
 numpydoc_show_class_members = False
 numpydoc_class_members_toctree = False
 
+# MyST-NB settings
+nb_execution_mode = 'off'
+
 # Autodoc settings
 autodoc_typehints = 'description'
 autodoc_member_order = 'bysource'
@@ -68,20 +72,17 @@
 # -- Options for HTML output -------------------------------------------------
 # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output
 
-html_theme = 'sphinx_rtd_theme'
+html_theme = 'pydata_sphinx_theme'
 html_static_path = ['_static']
 
 # Theme options
 html_theme_options = {
-    'logo_only': False,
-    'display_version': True,
-    'prev_next_buttons_location': 'bottom',
-    'style_external_links': False,
-    'collapse_navigation': False,
-    'sticky_navigation': True,
     'navigation_depth': 4,
-    'includehidden': True,
-    'titles_only': False
+}
+
+# Remove empty left sidebar from all pages
+html_sidebars = {
+    '**': [],
 }
 
 # If you want to add a custom logo, uncomment and specify:

docs/index.rst

@@ -3,42 +3,48 @@
 Welcome to Discovery's documentation!
 ======================================
 
-Discovery is a next-generation pulsar-timing-array data analysis package built on JAX.
+Discovery is a next-generation PTA data analysis package built on JAX.
 
 .. toctree::
-   :maxdepth: 2
+   :maxdepth: 1
    :caption: User Guide
 
+   tutorials/curn_example
    guide/overview
-   guide/data_model
-   guide/pulsar_data
    installation
 
 .. toctree::
-   :maxdepth: 2
+   :maxdepth: 1
    :caption: Tutorials
 
    quickstart
    tutorials/basic_likelihood
-   tutorials/simulations
    tutorials/optimal_statistic
+   tutorials/simulations
 
 .. toctree::
-   :maxdepth: 2
+   :maxdepth: 1
    :caption: Component Reference
 
    components/noise_signals
    components/priors_spectra
    components/delays
 
 .. toctree::
-   :maxdepth: 2
+   :maxdepth: 1
+   :caption: Other Useful Information
+
+   guide/data_model
+   guide/pulsar_data
+
+.. toctree::
+   :maxdepth: 1
    :caption: Advanced Topics
 
    advanced/conditional_sampling
 
 .. toctree::
-   :maxdepth: 2
+   :maxdepth: 1
    :caption: API Reference
 
    api/index

docs/tutorials/curn_example.ipynb

@@ -0,0 +1,215 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "ce3d00c3-eb23-4686-80fc-090913a839da",
+   "metadata": {},
+   "source": [
+    "# How best to run a CURN model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "ff14aca8-d77b-4ae1-baea-3ab4429efa51",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import discovery as ds\n",
+    "import jax\n",
+    "import numpy as np\n",
+    "import jax.numpy as jnp\n",
+    "import glob\n",
+    "from pathlib import Path\n",
+    "import discovery as ds\n",
+    "import discovery.samplers.numpyro as ds_numpyro\n",
+    "import matplotlib.pyplot as plt"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "7216e4a9-b649-483d-b9ed-d9332b52a418",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "datapath = Path(ds.__path__[0]) / '../../data'"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "43e38887-eba5-4791-9376-2e40ac43ca38",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "psrs = [ds.Pulsar.read_feather(f) for f in sorted(datapath.glob('*v1p1*.feather'))][:10] # only 10 pulsars for now"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "b1e4c150-d086-43b8-a421-76455936b3c3",
+   "metadata": {},
+   "source": [
+    "## Making the model \n",
+    "{func}`~discovery.signals.make_combined_crn` will combine intrinsic red noise and common nosie\n",
+    "into a single GP that can use a single Fourier basis.\n",
+    "since this fixes the names of the common process when it's created, we also return those parameters\n",
+    "so you can pass them later on.\n",
+    "\n",
+    "If you want different Fourier bases (different T-spans) for different pulsars, then\n",
+    "you will have to do something different."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6dc5fd0b-fe92-41b9-b619-209ba7cc06ab",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "\n",
+    "# common_parnames are the names of parameters\n",
+    "# that are shared for all pulsars.\n",
+    "mypl, common_parnames = ds.make_combined_crn(14, ds.powerlaw, ds.powerlaw, crn_prefix='gw_')\n",
+    "\n",
+    "psls = [ds.PulsarLikelihood([psr.residuals,\n",
+    "                             ds.makenoise_measurement(psr, psr.noisedict),\n",
+    "                             ds.makegp_ecorr(psr, psr.noisedict),\n",
+    "                             ds.makegp_timing(psr, svd=True)]) for psr in psrs]\n",
+    "\n",
+    "commongp = ds.makecommongp_fourier(psrs, mypl, 30, T=ds.getspan(psrs), name='red_noise',\n",
+    "                                   common=common_parnames)\n",
+    "\n",
+    "array_likelihood = ds.ArrayLikelihood(psls, commongp=commongp)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "881cc1d5-3fe1-4a05-b415-452e37a99b36",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# array_likelihood.logL.params"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c80e2067-7a65-4283-bb86-1582673755a6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "test_params = ds.sample_uniform(array_likelihood.logL.params)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "95ab3572-45ef-4391-ae5c-7b1d6dcc7997",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "jlogl = jax.jit(array_likelihood.logL)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "73c26aa8-11f4-49be-8f4d-0b55a9d44945",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "jlogl(test_params)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "e73229c5-b6fa-4a2a-9916-6399f6d014f9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "%%timeit\n",
+    "jlogl(test_params)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "e212fdca-2048-47e2-9c9e-ae7b8de33bfa",
+   "metadata": {},
+   "source": [
+    "## Variable transformations and performance\n",
+    "This does a transformation so that the parameters that get sampled\n",
+    "live on the full real line instead of uniform in a fixed range\n",
+    "this helps with NUTS sampling\n",
+    "\n",
+    "ATTENTION!!! In creating this transformed likelihood, JAX\n",
+    "actually bypasses the parameter dictionary completely\n",
+    "when it is compiled. This seems to give a large performance benefit on GPUs, where the dictionary\n",
+    "rolling and unrolling seems to cause significant overhead.\n",
+    "\n",
+    "For both the sampling reason, and this performance reason, I'd recommend using these transformations if possible. "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "a5cff402-90f7-478a-a3bd-6da2ec764aae",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "npmodel = ds_numpyro.makemodel_transformed(jlogl)\n",
+    "npsampler = ds_numpyro.makesampler_nuts(npmodel,\n",
+    "                                        num_warmup=100,\n",
+    "                                        num_samples=100)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1b7ca6d5-34b6-47b0-b65a-f14296bd20c2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "npsampler.run(jax.random.key(0))\n",
+    "chain = npsampler.to_df()\n",
+    "# chain.to_csv('chain.feather', index=False)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "52cdf3c4-bf1b-4e9e-b068-e8687249bcb2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "plt.hist(chain['gw_gamma'])\n",
+    "plt.xlabel(\"$\\gamma_{gw}$\")\n",
+    "plt.show()"
+   ]
+  }
+ ],
+ "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.13.7"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

pyproject.toml

@@ -69,9 +69,10 @@ flow = [
 
 docs = [
     "sphinx>=7.0",
-    "sphinx-rtd-theme>=2.0",
+    "pydata-sphinx-theme>=0.15",
     "sphinx-autodoc-typehints>=1.25",
     "numpydoc>=1.6",
+    "myst-nb>=1.0",
 ]
 
 [project.urls]