feat: major rewrite

lgrcia · Oct 28, 2024 · 798c5c9 · 798c5c9
1 parent 4f5331d
commit 798c5c9
Show file tree

Hide file tree

Showing 23 changed files with 1,410 additions and 2,473 deletions.
diff --git a/docs/_static/test_2.gif b/docs/_static/test_2.gif
diff --git a/docs/_templates/class.rst b/docs/_templates/class.rst
diff --git a/docs/api.rst b/docs/api.rst
diff --git a/docs/conf.py b/docs/conf.py
@@ -9,7 +9,7 @@
     "sphinx.ext.napoleon",
     "sphinx.ext.autosummary",
     "sphinx.ext.viewcode",
-    "matplotlib.sphinxext.plot_directive",
+    "autoapi.extension",
 ]
 
 templates_path = ["_templates"]
@@ -37,10 +37,24 @@
     "dollarmath",
 ]
 
-nb_execution_mode = "off"
 html_logo = "_static/spotter.png"
-html_css_files = ["style.css"]
 myst_url_schemes = ("http", "https")
 
 plot_html_show_formats = False
 plot_html_show_source_link = False
+
+autoapi_dirs = ["../spotter"]
+autoapi_ignore = ["*_version*", "*/types*"]
+autoapi_options = [
+    "members",
+    "undoc-members",
+    # "private-members",
+    "show-inheritance",
+    "show-module-summary",
+    "special-members",
+    # "imported-members",
+]
+# autoapi_add_toctree_entry = False
+autoapi_template_dir = "_autoapi_templates"
+
+suppress_warnings = ["autoapi.python_import_resolution"]
diff --git a/docs/index.md b/docs/index.md
@@ -11,12 +11,6 @@ Use at your own risk as the code is completely untested and its API subject to c
 *spotter* uses the [HEALPix](https://healpix.sourceforge.io/) subdivision scheme and is powered by the high-performance numerical package [JAX](https://jax.readthedocs.io/en/latest/notebooks/quickstart.html), enabling its use on GPUs.
 
 
-## Features
-
-- Small-scale surface features (e.g. beyond limitations of [starry]()) <span style="color:grey">
-- Modeling of any active regions with their limb laws (e.g. limb-brightened faculae)
-- GPU compatible <span style="color:grey">
-- Possibility to input any stellar spectra model
 
 ```{toctree}
 :maxdepth: 1
@@ -25,10 +19,17 @@ Use at your own risk as the code is completely untested and its API subject to c
 notebooks/introduction
 ```
 
+```{toctree}
+:maxdepth: 1
+:caption: Tutorials
+
+notebooks/surface_gp
+notebooks/ensemble
+```
+
 ```{toctree}
 :maxdepth: 1
 :caption: Reference
 
 notebooks/rotation.ipynb
-api
 ```
diff --git a/docs/notebooks/amplitude_constraints.ipynb b/docs/notebooks/amplitude_constraints.ipynb
diff --git a/docs/notebooks/ensemble.ipynb b/docs/notebooks/ensemble.ipynb
diff --git a/docs/notebooks/experiments.ipynb b/docs/notebooks/experiments.ipynb
diff --git a/docs/notebooks/introduction.ipynb b/docs/notebooks/introduction.ipynb
diff --git a/docs/notebooks/jax_features.ipynb b/docs/notebooks/jax_features.ipynb
diff --git a/docs/notebooks/rotation.ipynb b/docs/notebooks/rotation.ipynb
diff --git a/docs/notebooks/simple_example.ipynb b/docs/notebooks/simple_example.ipynb
diff --git a/docs/notebooks/surface_gp.ipynb b/docs/notebooks/surface_gp.ipynb
@@ -0,0 +1,282 @@
+{
+    "cells": [
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "# Surface Gaussian Process\n",
+                "\n",
+                "In this tutorial we show how to draw a non-uniform surface using a Gaussian Process (GP) on the sphere. Here, we will use this GP to model a star covered in [starspots](https://en.wikipedia.org/wiki/Sunspot), but of course spotter can be used to model any spherical surface such that of a planet.\n",
+                "\n",
+                "```{note}\n",
+                "GPs in spotter are defined and computed using the [tinygp](https://tinygp.readthedocs.io/en/stable/) Python package. For more details about GPs on the sphere, check out the [Custom Geometry](https://tinygp.readthedocs.io/en/stable/tutorials/geometry.html) `tinygp` tutorial.\n",
+                "```\n",
+                "\n",
+                "The main idea is that a 1D Gaussian Process can be easily defined on the sphere, using the [great circle distance](https://en.wikipedia.org/wiki/Great-circle_distance) as the distance metric in the GP kernel."
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "## Isotropic kernel\n",
+                "\n",
+                "In this first example, we define a kernel to draw a stellar surface uniformly covered with active regions, such as starspots.\n",
+                "\n",
+                "We start by defining the properties of our stellar surface with few variables"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "from spotter.core import vec\n",
+                "\n",
+                "N = 2**4  # number of sides\n",
+                "u = (0.4, 0.1)  # limb darkening coefficient\n",
+                "X = vec(N)  # pixels coordinates"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "We can then use `tinygp` to define an isotropic kernel on the sphere"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "import tinygp\n",
+                "from spotter.kernels import GreatCircleDistance\n",
+                "\n",
+                "kernel = 0.1 * tinygp.kernels.Matern52(0.4, distance=GreatCircleDistance())\n",
+                "gp = tinygp.GaussianProcess(kernel, X)"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "and draw a surface from it"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "import jax\n",
+                "import matplotlib.pyplot as plt\n",
+                "from spotter import viz\n",
+                "\n",
+                "y = gp.sample(jax.random.PRNGKey(4), shape=(1,))[0]\n",
+                "y = 1.0 - y.clip(0.0, 1.0)\n",
+                "\n",
+                "plt.figure(figsize=(2, 2))\n",
+                "viz.show(y, u=u)"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "```{note}\n",
+                "Any GP library allowing a custom distance metric can be used in combination with *spotter* to draw surfaces on the sphere.\n",
+                "``` \n",
+                "\n",
+                "\n",
+                "The GP we just built model active regions on a star with sizes related to the kernel length scale. Let's show surfaces drawn from kernels with different length scales"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "from matplotlib import gridspec\n",
+                "\n",
+                "lenght_scales = [0.1, 0.2, 0.3, 0.4, 0.5]\n",
+                "gridspec.GridSpec(1, 6)\n",
+                "plt.figure(figsize=(12, 2))\n",
+                "\n",
+                "for i, l in enumerate(lenght_scales):\n",
+                "    kernel = 0.1 * tinygp.kernels.Matern52(l, distance=GreatCircleDistance())\n",
+                "    gp = tinygp.GaussianProcess(kernel, X)\n",
+                "    y = gp.sample(jax.random.PRNGKey(i + 1), shape=(1,))[0]\n",
+                "    y = 1.0 - y.clip(0.0, 1.0)\n",
+                "    plt.subplot(1, 6, i + 1)\n",
+                "    viz.show(y, u=u)\n",
+                "    plt.title(f\"l={l}\")"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "## Active latitudes\n",
+                "\n",
+                "As [seen on the sun](https://en.wikipedia.org/wiki/Solar_cycle), active regions can be preferentially located at certain latitudes. In order to draw surfaces with such properties, *spotter* features the `ActiveLatitude` kernel (non-isotropic!).\n",
+                "\n",
+                "This kernel takes an isotropic kernel as input, which sets the active regions properties, as well as their preferred latitude and the standard deviation of the latitude band where they are located."
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "from spotter import kernels\n",
+                "\n",
+                "kernel = kernels.ActiveLatitude(\n",
+                "    kernel=0.1 * tinygp.kernels.Matern52(0.3, distance=kernels.GreatCircleDistance()),\n",
+                "    latitude=0.5,\n",
+                "    sigma=0.2,\n",
+                "    symetric=True,\n",
+                ")\n",
+                "\n",
+                "gp = tinygp.GaussianProcess(kernel, X)"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "We can plot the distribution of active regions along $\\cos{\\theta}$"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "import jax.numpy as jnp\n",
+                "\n",
+                "cos_theta = X.T[2] / jnp.linalg.norm(X.T, axis=0)\n",
+                "_ = plt.plot(cos_theta, jax.vmap(kernel.amplitude)(vec(N)))"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "and draw a sample from the GP"
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "import jax\n",
+                "import matplotlib.pyplot as plt\n",
+                "from spotter import viz\n",
+                "\n",
+                "y = gp.sample(jax.random.PRNGKey(6), shape=(1,))[0]\n",
+                "y = 1.0 - y.clip(0.0, 1.0)\n",
+                "\n",
+                "plt.figure(figsize=(2, 2))\n",
+                "viz.show(y, u=u)"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "As before, let's see how surfaces look for different values of the length scale."
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "from matplotlib import gridspec\n",
+                "\n",
+                "lenght_scales = [0.1, 0.2, 0.3, 0.4, 0.5]\n",
+                "gridspec.GridSpec(1, 6)\n",
+                "plt.figure(figsize=(12, 2))\n",
+                "\n",
+                "for i, l in enumerate(lenght_scales):\n",
+                "    kernel = kernels.ActiveLatitude(\n",
+                "        kernel=0.1 * tinygp.kernels.Matern52(l, distance=kernels.GreatCircleDistance()),\n",
+                "        latitude=0.5,\n",
+                "        sigma=0.2,\n",
+                "        symetric=True,\n",
+                "    )\n",
+                "    gp = tinygp.GaussianProcess(kernel, X)\n",
+                "    y = gp.sample(jax.random.PRNGKey(i), shape=(1,))[0]\n",
+                "    y = 1.0 - y.clip(0.0, 1.0)\n",
+                "    plt.subplot(1, 6, i + 1)\n",
+                "    viz.show(y, u=u)\n",
+                "    plt.title(f\"l={l}\")"
+            ]
+        },
+        {
+            "cell_type": "markdown",
+            "metadata": {},
+            "source": [
+                "We note that the kernel doesn't have to be symmetric in latitudes."
+            ]
+        },
+        {
+            "cell_type": "code",
+            "execution_count": null,
+            "metadata": {},
+            "outputs": [],
+            "source": [
+                "kernel = kernels.ActiveLatitude(\n",
+                "    kernel=0.1 * tinygp.kernels.Matern52(0.1, distance=kernels.GreatCircleDistance()),\n",
+                "    latitude=jnp.pi / 2,\n",
+                "    sigma=0.2,\n",
+                "    symetric=False,\n",
+                ")\n",
+                "gp = tinygp.GaussianProcess(kernel, X)\n",
+                "\n",
+                "y = gp.sample(jax.random.PRNGKey(0), shape=(1,))[0]\n",
+                "y = 1.0 - y.clip(0.0, 1.0)\n",
+                "\n",
+                "plt.figure(figsize=(4, 2))\n",
+                "ax = plt.subplot(121)\n",
+                "viz.show(y, inc=0.8, u=u, ax=ax, vmin=0, vmax=1)\n",
+                "ax.set_title(\"north pole\")\n",
+                "\n",
+                "ax = plt.subplot(122)\n",
+                "viz.show(y, inc=jnp.pi / 2 + 0.8, u=u, ax=ax, vmin=0, vmax=1)\n",
+                "_ = ax.set_title(\"south pole\")"
+            ]
+        }
+    ],
+    "metadata": {
+        "kernelspec": {
+            "display_name": "spotter",
+            "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.10.14"
+        }
+    },
+    "nbformat": 4,
+    "nbformat_minor": 2
+}