Merge branch 'main' into Issue/176-Amplitude_dif

.github/dependabot.yml

@@ -0,0 +1,6 @@
+version: 2
+updates:
+  - package-ecosystem: "github-actions"
+    directory: "/"
+    schedule:
+      interval: "monthly"

.github/workflows/news.yml

@@ -9,11 +9,11 @@ jobs:
   news:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
         with:
           fetch-depth: 0
       - name: Setup Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v4
         with:
           python-version: "3.9"
       - name: Install dependencies

.github/workflows/tests.yml

@@ -17,11 +17,11 @@ jobs:
 
     steps:
       - name: Checkout
-        uses: actions/checkout@v2
+        uses: actions/checkout@v3
         with:
           fetch-depth: 0
       - name: Setup Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v4
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
@@ -45,7 +45,7 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Setup Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v4
         with:
           python-version: "3.9"
       - name: Finish coverage collection
@@ -59,11 +59,11 @@ jobs:
   doctest:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
         with:
           fetch-depth: 0
       - name: Setup Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v4
         with:
           python-version: "3.9"
       - name: Install dependencies
@@ -76,11 +76,11 @@ jobs:
   lint:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
         with:
           fetch-depth: 0
       - name: Setup Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v4
         with:
           python-version: "3.9"
       - name: Install dependencies
@@ -93,10 +93,10 @@ jobs:
   build:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v3
         with:
           fetch-depth: 0
-      - uses: actions/setup-python@v2
+      - uses: actions/setup-python@v4
         name: Install Python
         with:
           python-version: "3.9"
@@ -105,7 +105,7 @@ jobs:
           python -m pip install -U pip
           python -m pip install -U build
           python -m build .
-      - uses: actions/upload-artifact@v2
+      - uses: actions/upload-artifact@v3
         with:
           path: dist/*
 
@@ -114,12 +114,12 @@ jobs:
     runs-on: ubuntu-latest
     if: startsWith(github.ref, 'refs/tags/')
     steps:
-      - uses: actions/download-artifact@v2
+      - uses: actions/download-artifact@v3
         with:
           name: artifact
           path: dist
 
-      - uses: pypa/gh-action-pypi-publish@v1.4.2
+      - uses: pypa/gh-action-pypi-publish@v1.8.8
         with:
           user: __token__
           password: ${{ secrets.pypi_password }}

.pre-commit-config.yaml

@@ -1,31 +1,33 @@
+ci:
+  autoupdate_schedule: monthly
+
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.2.0
+    rev: v4.4.0
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
         exclude_types: [json, binary]
   - repo: https://github.com/PyCQA/isort
-    rev: "5.10.1"
+    rev: "5.12.0"
     hooks:
       - id: isort
         additional_dependencies: [toml]
         exclude: docs/tutorials
   - repo: https://github.com/psf/black
-    rev: "22.3.0"
+    rev: "23.7.0"
     hooks:
       - id: black-jupyter
   - repo: https://github.com/kynan/nbstripout
-    rev: "0.5.0"
+    rev: "0.6.1"
     hooks:
       - id: nbstripout
         exclude: docs/benchmarks.ipynb
   - repo: https://github.com/hadialqattan/pycln
-    rev: "v1.3.1"
+    rev: "v2.2.0"
     hooks:
       - id: pycln
-        additional_dependencies: ["click<8.1.0"]
   - repo: https://github.com/pre-commit/mirrors-mypy
-    rev: "v0.942"
+    rev: "v1.4.1"
     hooks:
       - id: mypy

README.md

@@ -5,7 +5,7 @@
   <br>
   <br>
   <a href="https://github.com/dfm/tinygp/actions/workflows/tests.yml">
-    <img alt="GitHub Workflow Status" src="https://img.shields.io/github/workflow/status/dfm/tinygp/Tests">
+    <img alt="GitHub Workflow Status" src="https://img.shields.io/github/actions/workflow/status/dfm/tinygp/tests.yml?branch=main">
   </a>
   <a href="https://tinygp.readthedocs.io">
     <img alt="Read the Docs" src="https://img.shields.io/readthedocs/tinygp">

docs/conf.py

@@ -9,6 +9,7 @@
     "sphinx.ext.napoleon",
     "sphinx.ext.viewcode",
     "myst_nb",
+    "IPython.sphinxext.ipython_console_highlighting",
 ]
 
 myst_enable_extensions = ["dollarmath", "colon_fence"]
@@ -47,9 +48,9 @@
     "use_download_button": True,
 }
 html_baseurl = "https://tinygp.readthedocs.io/en/latest/"
-jupyter_execute_notebooks = "auto"
-execution_excludepatterns = ["benchmarks.ipynb"]
-execution_timeout = -1
+nb_execution_mode = "auto"
+nb_execution_excludepatterns = ["benchmarks.ipynb"]
+nb_execution_timeout = -1
 
 autodoc_type_aliases = {
     "JAXArray": "tinygp.helpers.JAXArray",

docs/news.rst

@@ -5,6 +5,30 @@ Release Notes
 
 .. towncrier release notes start
 
+tinygp 0.2.3 (2022-10-31)
+-------------------------
+
+Features
+~~~~~~~~
+
+- Removed deprecation warning from ``predict`` method and wrapped it in a
+  ``jax.jit`` in order to support interactive use. (`#120 <https://github.com/dfm/tinygp/issues/120>`_)
+- Added check for sorted input coordinates when using the ``QuasisepSolver``;
+  a ``ValueError`` is thrown if they are not. (`#123 <https://github.com/dfm/tinygp/issues/123>`_)
+
+
+Bugfixes
+~~~~~~~~
+
+- Fixed incorrect definition of ``observation_model`` for ``Celerite`` kernel. (`#88 <https://github.com/dfm/tinygp/issues/88>`_)
+- Fixed ``FutureWarning`` by updating ``tree_map`` to ``tree_util.tree_map``. (`#114 <https://github.com/dfm/tinygp/issues/114>`_)
+- Fixed issue when tree structure and shape of ``X_test`` input to ``condition``
+  was incompatible with the initial input. (`#119 <https://github.com/dfm/tinygp/issues/119>`_)
+- Fixed bug where the gradient of the L2 distance would return NaN when the
+  distance was zero. (`#121 <https://github.com/dfm/tinygp/issues/121>`_)
+- Fixed behavior of DotProduct kernel on scalar inputs. (`#124 <https://github.com/dfm/tinygp/issues/124>`_)
+
+
 tinygp 0.2.2 (2022-04-20)
 -------------------------
 

docs/tutorials/intro.ipynb

@@ -232,6 +232,7 @@
     "gp = GaussianProcess(kernel, X, diag=1e-5, mean=2.0)\n",
     "y_const = gp.sample(jax.random.PRNGKey(4), shape=(5,))\n",
     "\n",
+    "\n",
     "# And a GP with a general mean function\n",
     "def mean_function(x):\n",
     "    return 5 * jax.numpy.sin(x)\n",

docs/tutorials/kernels.ipynb

@@ -23,6 +23,7 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "efc5f580-005f-43d5-9056-28f1885873eb",
    "metadata": {},
@@ -32,13 +33,8 @@
     "# Custom Kernels\n",
     "\n",
     "One of the goals of the `tinygp` interface design was to make the kernel building framework flexible and easily extensible.\n",
-    "In this tutorial, we demontrate this interface with two examples; one simple, and one more complicated.\n",
-    "Besides describing this interface, we also show how `tinygp` can support arbitrary [JAX pytrees](https://jax.readthedocs.io/en/latest/pytrees.html) as input.\n",
-    "\n",
-    "## Example: Spectral mixture kernel\n",
-    "\n",
-    "In this section, we will implement the \"spectral mixture kernel\" proposed by [Gordon Wilson & Adams (2013)](https://arxiv.org/abs/1302.4245).\n",
-    "It would be possible to implement this using sums of built in kernels, but the interface seems better if we implement a custom kernel and I expect that we'd get somewhat better performance for mixtures with many components.\n",
+    "In this tutorial, we demonstrate this interface using the \"spectral mixture kernel\" proposed by [Gordon Wilson & Adams (2013)](https://arxiv.org/abs/1302.4245).\n",
+    "It would be possible to implement this using sums of built-in kernels, but the interface seems better if we implement a custom kernel and I expect that we'd get somewhat better performance for mixtures with many components.\n",
     "\n",
     "Now, let's implement this kernel in a way that `tinygp` understands.\n",
     "When doing this, you will subclass {class}`tinygp.kernels.Kernel` and implement the {func}`tinygp.kernels.Kernel.evaluate` method.\n",
@@ -70,7 +66,7 @@
     "            * jnp.prod(\n",
     "                jnp.exp(-2 * jnp.pi**2 * tau**2 / self.scale**2)\n",
     "                * jnp.cos(2 * jnp.pi * self.freq * tau),\n",
-    "                axis=-1,\n",
+    "                axis=0,\n",
     "            )\n",
     "        )"
    ]
@@ -108,7 +104,7 @@
     "params = {\n",
     "    \"log_weight\": np.log([1.0, 1.0]),\n",
     "    \"log_scale\": np.log([10.0, 20.0]),\n",
-    "    \"log_freq\": np.log([1.0, 1.0 / 3.0]),\n",
+    "    \"log_freq\": np.log([1.0, 1.0 / 2.0]),\n",
     "    \"log_diag\": np.log(0.1),\n",
     "    \"mean\": 0.0,\n",
     "}\n",
@@ -126,11 +122,12 @@
    ]
   },
   {
+   "attachments": {},
    "cell_type": "markdown",
    "id": "0a8b4279-f681-49c2-8735-23ed5b8890c8",
    "metadata": {},
    "source": [
-    "One thing to note here is that we've used named parameters in a dictionary, instead of an array of parameters as in some of the other examples.\n",
+    "One thing to note here is that we've used named parameters in a dictionary, instead of an array of parameters as in some other examples.\n",
     "This would be awkward (but not impossible) to fit using `scipy`, so instead we'll use [`optax`](https://github.com/deepmind/optax) for optimization:"
    ]
   },
@@ -213,7 +210,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "tinygp",
    "language": "python",
    "name": "python3"
   },
@@ -227,7 +224,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.9"
+   "version": "3.10.6"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "d20ea8a315da34b3e8fab0dbd7b542a0ef3c8cf12937343660e6bc10a20768e3"
+   }
   }
  },
  "nbformat": 4,

docs/tutorials/means.ipynb

@@ -222,7 +222,7 @@
     "import jaxopt\n",
     "\n",
     "solver = jaxopt.ScipyMinimize(fun=loss)\n",
-    "soln = solver.run(jax.tree_map(jnp.asarray, params))\n",
+    "soln = solver.run(jax.tree_util.tree_map(jnp.asarray, params))\n",
     "print(f\"Final negative log likelihood: {soln.state.fun_val}\")"
    ]
   },

docs/tutorials/modeling.ipynb

@@ -311,7 +311,7 @@
    "id": "c2f9c49e-9eea-4971-a939-7390a31368fc",
    "metadata": {},
    "source": [
-    "And, finally we can plot our posterior inferences of the comditional process, compared to the true model:"
+    "And, finally we can plot our posterior inferences of the conditional process, compared to the true model:"
    ]
   },
   {

docs/tutorials/quasisep-custom.ipynb

@@ -670,7 +670,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "Python 3 (ipykernel)",
+   "display_name": "Python 3.10.6 ('tinygp')",
    "language": "python",
    "name": "python3"
   },
@@ -684,7 +684,12 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.9"
+   "version": "3.10.6"
+  },
+  "vscode": {
+   "interpreter": {
+    "hash": "d20ea8a315da34b3e8fab0dbd7b542a0ef3c8cf12937343660e6bc10a20768e3"
+   }
   }
  },
  "nbformat": 4,

docs/tutorials/quasisep.ipynb

@@ -185,7 +185,7 @@
     "import jaxopt\n",
     "\n",
     "solver = jaxopt.ScipyMinimize(fun=loss)\n",
-    "soln = solver.run(jax.tree_map(jnp.asarray, params))\n",
+    "soln = solver.run(jax.tree_util.tree_map(jnp.asarray, params))\n",
     "print(f\"Final negative log likelihood: {soln.state.fun_val}\")"
    ]
   },

docs/tutorials/quickstart.ipynb

@@ -109,7 +109,6 @@
     "\n",
     "\n",
     "def build_gp(theta, X):\n",
-    "\n",
     "    # We want most of our parameters to be positive so we take the `exp` here\n",
     "    # Note that we're using `jnp` instead of `np`\n",
     "    amps = jnp.exp(theta[\"log_amps\"])\n",

docs/tutorials/transforms.ipynb

@@ -170,8 +170,14 @@
     "\n",
     "        # Define a custom transform to pass the input coordinates through our\n",
     "        # `Transformer` network from above\n",
+    "        # Note: with recent version of flax, you can't directly vmap modules,\n",
+    "        # but we can get around that by explicitly constructing the init and\n",
+    "        # apply functions. Ref:\n",
+    "        # https://flax.readthedocs.io/en/latest/advanced_topics/lift.html\n",
     "        transform = Transformer()\n",
-    "        kernel = transforms.Transform(transform, base_kernel)\n",
+    "        transform_params = self.param(\"transform\", transform.init, x[:1])\n",
+    "        apply_fn = lambda x: transform.apply(transform_params, x)\n",
+    "        kernel = transforms.Transform(apply_fn, base_kernel)\n",
     "\n",
     "        # Evaluate and return the GP negative log likelihood as usual with the\n",
     "        # transformed features\n",
@@ -333,7 +339,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.9"
+   "version": "3.10.6"
   }
  },
  "nbformat": 4,

news/143.doc

@@ -0,0 +1,2 @@
+Fixed incorrect definition of "spectral mixture kernel" in the custom kernels
+tutorial.

news/148.feature

@@ -0,0 +1 @@
+Removed `__post_init__` checks after kernel construction to avoid extraneous errors when returning kernels out of `jax.vmap`'d functions.

news/151.feature

@@ -0,0 +1 @@
+Added Zenodo data to improve citation tracking.

news/153.doc

@@ -0,0 +1,2 @@
+Unpinned the docs theme version to fix release compatibility with recent
+versions of setuptools.

news/154.doc

@@ -0,0 +1 @@
+Added past contributor metadata to `.zenodo.json`.

news/159.bugfix

@@ -0,0 +1 @@
+Fixed syntax for `vmap` of `flax` modules in `transforms` tutorial.