tinygp 0.2.4__tar.gz → 0.3.0__tar.gz

{tinygp-0.2.4 → tinygp-0.3.0}/.github/workflows/news.yml

@@ -8,11 +8,11 @@ jobs:
     if: ${{ github.actor != 'dependabot[bot]' && github.actor != 'pre-commit-ci[bot]' }}
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
       - name: Setup Python
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: "3.10"
       - name: Install dependencies

{tinygp-0.2.4 → tinygp-0.3.0}/.github/workflows/tests.yml

@@ -15,17 +15,25 @@ jobs:
       matrix:
         python-version: ["3.9", "3.10", "3.11"]
         nox-session: ["test"]
+        x64: ["1"]
         include:
+          - python-version: "3.10"
+            nox-session: "test"
+            x64: "0"
+          - python-version: "3.10"
+            nox-session: "comparison"
+            x64: "1"
           - python-version: "3.10"
             nox-session: "doctest"
+            x64: "1"
 
     steps:
       - name: Checkout
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
         with:
           fetch-depth: 0
       - name: Setup Python
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
@@ -36,14 +44,16 @@ jobs:
         run: |
           python -m nox --non-interactive --error-on-missing-interpreter \
             --session ${{ matrix.nox-session }} --python ${{ matrix.python-version }}
+        env:
+          JAX_ENABLE_X64: ${{ matrix.x64 }}
 
   build:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v3
+      - uses: actions/checkout@v4
         with:
           fetch-depth: 0
-      - uses: actions/setup-python@v4
+      - uses: actions/setup-python@v5
         name: Install Python
         with:
           python-version: "3.10"
@@ -55,7 +65,7 @@ jobs:
         run: python -m build .
       - name: Check the distribution
         run: python -m twine check --strict dist/*
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
         with:
           path: dist/*
 
@@ -69,8 +79,8 @@ jobs:
     runs-on: ubuntu-latest
     if: github.event_name == 'push' && startsWith(github.ref, 'refs/tags/')
     steps:
-      - uses: actions/download-artifact@v3
+      - uses: actions/download-artifact@v4
         with:
           name: artifact
           path: dist
-      - uses: pypa/gh-action-pypi-publish@v1.8.10
+      - uses: pypa/gh-action-pypi-publish@v1.8.11

{tinygp-0.2.4 → tinygp-0.3.0}/.pre-commit-config.yaml

@@ -3,17 +3,17 @@ ci:
 
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: "v4.4.0"
+    rev: "v4.5.0"
     hooks:
       - id: trailing-whitespace
       - id: end-of-file-fixer
         exclude_types: [json, binary]
   - repo: https://github.com/psf/black
-    rev: "23.9.1"
+    rev: "23.12.1"
     hooks:
       - id: black-jupyter
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: "v0.0.291"
+    rev: "v0.1.9"
     hooks:
       - id: ruff
         args: [--fix, --exit-non-zero-on-fix]

{tinygp-0.2.4 → tinygp-0.3.0}/.zenodo.json

@@ -5,6 +5,11 @@
       "affiliation": "Center for Computational Astrophysics, Flatiron Institute, New York, NY, USA",
       "name": "Foreman-Mackey, Daniel"
     },
+    {
+      "orcid": "0000-0003-1262-2897",
+      "affiliation": "Department of Physics and Astronomy, Bishop's University, Canada",
+      "name": "Weixiang Yu"
+    },
     {
       "orcid": "0000-0003-0048-1118",
       "affiliation": "Indian Institute of Technology Gandhinagar: Gandhinagar, Gujarat, IN",
@@ -15,11 +20,26 @@
       "affiliation": "Massachusetts Institute of Technology, Probabilistic Computing Project, Cambridge, MA, USA",
       "name": "Becker, McCoy Reynolds"
     },
+    {
+      "orcid": "0000-0003-3287-5250",
+      "affiliation": "Department of Astronomy and the DiRAC Institute, University of Washington, Seattle, WA, USA",
+      "name": "Caplar, Neven"
+    },
+    {
+      "orcid": "0000-0002-1169-7486",
+      "affiliation": "SRON Netherlands Institute for Space Research, Leiden, The Netherlands",
+      "name": "Huppenkothen, Daniela"
+    },
     {
       "orcid": "0000-0002-0440-9597",
       "name": "Killestein, Thomas",
       "affiliation": "Department of Physics, University of Warwick, Coventry, UK"
     },
+    {
+      "orcid": "0000-0003-1001-0707",
+      "affiliation": "Department of Physics and Astronomy, Aarhus University, DK",
+      "name": "Tronsgaard, René"
+    },
     {
       "affiliation": "School of Public Health, Imperial College London, UK",
       "name": "Rashid, Theo"
@@ -28,16 +48,6 @@
       "orcid": "0000-0003-1354-0578",
       "affiliation": "Helmholtz-Zentrum Dresden-Rossendorf e.V.",
       "name": "Schmerler, Steve"
-    },
-    {
-      "orcid": "0000-0003-1001-0707",
-      "affiliation": "Department of Physics and Astronomy, Aarhus University, DK",
-      "name": "Tronsgaard, René"
-    },
-    {
-      "orcid": "0000-0003-3287-5250",
-      "affiliation": "Department of Astronomy and the DiRAC Institute, University of Washington, Seattle, WA, USA",
-      "name": "Caplar, Neven"
     }
   ],
   "license": "MIT",

{tinygp-0.2.4 → tinygp-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tinygp
-Version: 0.2.4
+Version: 0.3.0
 Summary: The tiniest of Gaussian Process libraries
 Author-email: Dan Foreman-Mackey <foreman.mackey@gmail.com>
 License: MIT
@@ -13,8 +13,12 @@ Classifier: Operating System :: OS Independent
 Classifier: Programming Language :: Python
 Classifier: Programming Language :: Python :: 3
 Requires-Python: >=3.9
+Requires-Dist: equinox
 Requires-Dist: jax
 Requires-Dist: jaxlib
+Provides-Extra: comparison
+Requires-Dist: celerite; extra == 'comparison'
+Requires-Dist: george; extra == 'comparison'
 Provides-Extra: docs
 Requires-Dist: arviz; extra == 'docs'
 Requires-Dist: flax; extra == 'docs'
@@ -27,8 +31,6 @@ Requires-Dist: optax; extra == 'docs'
 Requires-Dist: sphinx-book-theme; extra == 'docs'
 Requires-Dist: statsmodels; extra == 'docs'
 Provides-Extra: test
-Requires-Dist: celerite; extra == 'test'
-Requires-Dist: george; extra == 'test'
 Requires-Dist: pytest; extra == 'test'
 Description-Content-Type: text/markdown
 

{tinygp-0.2.4 → tinygp-0.3.0}/docs/news.rst

@@ -5,6 +5,23 @@ Release Notes
 
 .. towncrier release notes start
 
+tinygp 0.3.0 (2024-01-05)
+-------------------------
+
+Features
+~~~~~~~~
+
+- Added a more robust and better tested implementation of the ``CARMA`` kernel for
+  use with the ``QuasisepSolver``. (`#90 <https://github.com/dfm/tinygp/issues/90>`_)
+- Switched all base classes to `equinox.Module <https://docs.kidger.site/equinox/api/module/module/>`_ objects to simplify dataclass handling. (`#200 <https://github.com/dfm/tinygp/issues/200>`_)
+
+
+Bugfixes
+~~~~~~~~
+
+- Fixed use of `jnp.roots` and `np.roll` to make CARMA kernel jit-compliant. (`#188 <https://github.com/dfm/tinygp/issues/188>`_)
+
+
 tinygp 0.2.4 (2023-09-29)
 -------------------------
 

{tinygp-0.2.4 → tinygp-0.3.0}/docs/tutorials/derivative.ipynb

@@ -105,8 +105,7 @@
     "\n",
     "\n",
     "class DerivativeKernel(tinygp.kernels.Kernel):\n",
-    "    def __init__(self, kernel):\n",
-    "        self.kernel = kernel\n",
+    "    kernel: tinygp.kernels.Kernel\n",
     "\n",
     "    def evaluate(self, X1, X2):\n",
     "        t1, d1 = X1\n",
@@ -301,6 +300,10 @@
     "            shape as ``coeff_prim``.\n",
     "    \"\"\"\n",
     "\n",
+    "    kernel: tinygp.kernels.Kernel\n",
+    "    coeff_prim: jax.Array\n",
+    "    coeff_deriv: jax.Array\n",
+    "\n",
     "    def __init__(self, kernel, coeff_prim, coeff_deriv):\n",
     "        self.kernel = kernel\n",
     "        self.coeff_prim, self.coeff_deriv = jnp.broadcast_arrays(\n",
@@ -497,7 +500,7 @@
    "hash": "d20ea8a315da34b3e8fab0dbd7b542a0ef3c8cf12937343660e6bc10a20768e3"
   },
   "kernelspec": {
-   "display_name": "Python 3.9.9 ('tinygp')",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },
@@ -511,9 +514,9 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.9"
+   "version": "3.10.6"
   }
  },
  "nbformat": 4,
- "nbformat_minor": 2
+ "nbformat_minor": 4
 }

{tinygp-0.2.4 → tinygp-0.3.0}/docs/tutorials/kernels.ipynb

@@ -54,10 +54,9 @@
     "\n",
     "\n",
     "class SpectralMixture(tinygp.kernels.Kernel):\n",
-    "    def __init__(self, weight, scale, freq):\n",
-    "        self.weight = jnp.atleast_1d(weight)\n",
-    "        self.scale = jnp.atleast_1d(scale)\n",
-    "        self.freq = jnp.atleast_1d(freq)\n",
+    "    weight: jax.Array\n",
+    "    scale: jax.Array\n",
+    "    freq: jax.Array\n",
     "\n",
     "    def evaluate(self, X1, X2):\n",
     "        tau = jnp.atleast_1d(jnp.abs(X1 - X2))[..., None]\n",
@@ -210,7 +209,7 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "tinygp",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
    "name": "python3"
   },

{tinygp-0.2.4 → tinygp-0.3.0}/noxfile.py

@@ -9,6 +9,12 @@ PYTHON_VERSIONS = ["3.9", "3.10", "3.11", "3.12"]
 @nox.session(python=PYTHON_VERSIONS)
 def test(session: nox.Session) -> None:
     session.install(".[test]")
+    session.run("pytest", *session.posargs)
+
+
+@nox.session(python=PYTHON_VERSIONS)
+def comparison(session: nox.Session) -> None:
+    session.install(".[test,comparison]")
     session.run("pytest", *session.posargs, env={"JAX_ENABLE_X64": "1"})
 
 

{tinygp-0.2.4 → tinygp-0.3.0}/pyproject.toml

@@ -15,10 +15,11 @@ classifiers = [
     "Programming Language :: Python :: 3",
 ]
 dynamic = ["version"]
-dependencies = ["jax", "jaxlib"]
+dependencies = ["jax", "jaxlib", "equinox"]
 
 [project.optional-dependencies]
-test = ["pytest", "george", "celerite"]
+test = ["pytest"]
+comparison = ["george", "celerite"]
 docs = [
     "sphinx-book-theme",
     "myst-nb",

{tinygp-0.2.4 → tinygp-0.3.0}/src/tinygp/gp.py

@@ -11,8 +11,10 @@ from typing import (
     NamedTuple,
 )
 
+import equinox as eqx
 import jax
 import jax.numpy as jnp
+import numpy as np
 
 from tinygp import kernels, means
 from tinygp.helpers import JAXArray
@@ -20,12 +22,13 @@ from tinygp.kernels.quasisep import Quasisep
 from tinygp.noise import Diagonal, Noise
 from tinygp.solvers import DirectSolver, QuasisepSolver
 from tinygp.solvers.quasisep.core import SymmQSM
+from tinygp.solvers.solver import Solver
 
 if TYPE_CHECKING:
     from tinygp.numpyro_support import TinyDistribution
 
 
-class GaussianProcess:
+class GaussianProcess(eqx.Module):
     """An interface for designing a Gaussian Process regression model
 
     Args:
@@ -50,6 +53,15 @@ class GaussianProcess:
             algebra.
     """
 
+    num_data: int = eqx.field(static=True)
+    dtype: np.dtype = eqx.field(static=True)
+    kernel: kernels.Kernel
+    X: JAXArray
+    mean_function: means.MeanBase
+    mean: JAXArray
+    noise: Noise
+    solver: Solver
+
     def __init__(
         self,
         kernel: kernels.Kernel,
@@ -57,7 +69,7 @@ class GaussianProcess:
         *,
         diag: JAXArray | None = None,
         noise: Noise | None = None,
-        mean: Callable[[JAXArray], JAXArray] | JAXArray | None = None,
+        mean: means.MeanBase | Callable[[JAXArray], JAXArray] | JAXArray | None = None,
         solver: Any | None = None,
         mean_value: JAXArray | None = None,
         covariance_value: Any | None = None,
@@ -66,7 +78,7 @@ class GaussianProcess:
         self.kernel = kernel
         self.X = X
 
-        if callable(mean):
+        if isinstance(mean, means.MeanBase):
             self.mean_function = mean
         elif mean is None:
             self.mean_function = means.Mean(jnp.zeros(()))
@@ -76,7 +88,7 @@ class GaussianProcess:
             mean_value = jax.vmap(self.mean_function)(self.X)
         self.num_data = mean_value.shape[0]
         self.dtype = mean_value.dtype
-        self.loc = self.mean = mean_value
+        self.mean = mean_value
         if self.mean.ndim != 1:
             raise ValueError(
                 "Invalid mean shape: " f"expected ndim = 1, got ndim={self.mean.ndim}"
@@ -92,7 +104,7 @@ class GaussianProcess:
                 solver = QuasisepSolver
             else:
                 solver = DirectSolver
-        self.solver = solver.init(
+        self.solver = solver(
             kernel,
             self.X,
             self.noise,
@@ -100,6 +112,10 @@ class GaussianProcess:
             **solver_kwargs,
         )
 
+    @property
+    def loc(self) -> JAXArray:
+        return self.mean
+
     @property
     def variance(self) -> JAXArray:
         return self.solver.variance()
@@ -209,7 +225,6 @@ class GaussianProcess:
 
     @partial(
         jax.jit,
-        static_argnums=(0,),
         static_argnames=("include_mean", "return_var", "return_cov"),
     )
     def predict(
@@ -281,7 +296,7 @@ class GaussianProcess:
 
         return TinyDistribution(self, **kwargs)
 
-    @partial(jax.jit, static_argnums=(0, 2))
+    @partial(jax.jit, static_argnums=(2,))
     def _sample(
         self,
         key: jax.random.KeyArray,
@@ -296,16 +311,16 @@ class GaussianProcess:
             self.solver.dot_triangular(normal_samples), 0, -1
         )
 
-    @partial(jax.jit, static_argnums=0)
+    @jax.jit
     def _compute_log_prob(self, alpha: JAXArray) -> JAXArray:
         loglike = -0.5 * jnp.sum(jnp.square(alpha)) - self.solver.normalization()
         return jnp.where(jnp.isfinite(loglike), loglike, -jnp.inf)
 
-    @partial(jax.jit, static_argnums=0)
+    @jax.jit
     def _get_alpha(self, y: JAXArray) -> JAXArray:
         return self.solver.solve_triangular(y - self.loc)
 
-    @partial(jax.jit, static_argnums=(0, 3))
+    @partial(jax.jit, static_argnums=(3,))
     def _condition(
         self,
         y: JAXArray,

tinygp-0.3.0/src/tinygp/helpers.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+__all__ = ["JAXArray", "dataclass", "field"]
+
+from typing import Any
+
+import equinox as eqx
+import jax
+
+JAXArray = jax.Array
+
+
+# The following is just for backwards compatibility since tinygp used to provide a
+# custom dataclass implementation
+field = eqx.field
+
+
+def dataclass(clz: type[Any]) -> type[Any]:
+    return clz

{tinygp-0.2.4 → tinygp-0.3.0}/src/tinygp/kernels/base.py

@@ -11,23 +11,23 @@ __all__ = [
     "Polynomial",
 ]
 
-from abc import ABCMeta, abstractmethod
+from abc import abstractmethod
 from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, Callable, Union
 
+import equinox as eqx
 import jax
 import jax.numpy as jnp
 
-from tinygp.helpers import JAXArray, dataclass, field
+from tinygp.helpers import JAXArray
 
 if TYPE_CHECKING:
     from tinygp.solvers.solver import Solver
 
-
 Axis = Union[int, Sequence[int]]
 
 
-class Kernel(metaclass=ABCMeta):
+class Kernel(eqx.Module):
     """The base class for all kernel implementations
 
     This subclass provides default implementations to add and multiply kernels.
@@ -35,11 +35,6 @@ class Kernel(metaclass=ABCMeta):
     :func:`Kernel.evaluate` with custom behavior.
     """
 
-    if TYPE_CHECKING:
-
-        def __init__(self, *args: Any, **kwargs: Any) -> None:
-            pass
-
     @abstractmethod
     def evaluate(self, X1: JAXArray, X2: JAXArray) -> JAXArray:
         """Evaluate the kernel at a pair of input coordinates
@@ -58,6 +53,7 @@ class Kernel(metaclass=ABCMeta):
            ``(n_data, n_dim)``, and you should let the :class:`Kernel` ``vmap``
            magic handle all the broadcasting for you.
         """
+        del X1, X2
         raise NotImplementedError
 
     def evaluate_diag(self, X: JAXArray) -> JAXArray:
@@ -130,7 +126,6 @@ class Kernel(metaclass=ABCMeta):
         return Product(Constant(other), self)
 
 
-@dataclass
 class Conditioned(Kernel):
     """A kernel used when conditioning a process on data
 
@@ -158,7 +153,6 @@ class Conditioned(Kernel):
         return self.kernel.evaluate_diag(X) - K.transpose() @ K
 
 
-@dataclass
 class Custom(Kernel):
     """A custom kernel class implemented as a callable
 
@@ -167,13 +161,12 @@ class Custom(Kernel):
             :func:`Kernel.evaluate`.
     """
 
-    function: Callable[[Any, Any], Any]
+    function: Callable[[Any, Any], Any] = eqx.field(static=True)
 
     def evaluate(self, X1: JAXArray, X2: JAXArray) -> JAXArray:
         return self.function(X1, X2)
 
 
-@dataclass
 class Sum(Kernel):
     """A helper to represent the sum of two kernels"""
 
@@ -184,7 +177,6 @@ class Sum(Kernel):
         return self.kernel1.evaluate(X1, X2) + self.kernel2.evaluate(X1, X2)
 
 
-@dataclass
 class Product(Kernel):
     """A helper to represent the product of two kernels"""
 
@@ -195,7 +187,6 @@ class Product(Kernel):
         return self.kernel1.evaluate(X1, X2) * self.kernel2.evaluate(X1, X2)
 
 
-@dataclass
 class Constant(Kernel):
     r"""This kernel returns the constant
 
@@ -209,15 +200,15 @@ class Constant(Kernel):
         c: The parameter :math:`c` in the above equation.
     """
 
-    value: JAXArray
+    value: JAXArray | float
 
     def evaluate(self, X1: JAXArray, X2: JAXArray) -> JAXArray:
+        del X1, X2
         if jnp.ndim(self.value) != 0:
             raise ValueError("The value of a constant kernel must be a scalar")
-        return self.value
+        return jnp.asarray(self.value)
 
 
-@dataclass
 class DotProduct(Kernel):
     r"""The dot product kernel
 
@@ -234,7 +225,6 @@ class DotProduct(Kernel):
         return X1 @ X2
 
 
-@dataclass
 class Polynomial(Kernel):
     r"""A polynomial kernel
 
@@ -249,9 +239,9 @@ class Polynomial(Kernel):
         sigma: The parameter :math:`\sigma`.
     """
 
-    order: JAXArray
-    scale: JAXArray = field(default_factory=lambda: jnp.ones(()))
-    sigma: JAXArray = field(default_factory=lambda: jnp.zeros(()))
+    order: JAXArray | float
+    scale: JAXArray | float = eqx.field(default_factory=lambda: jnp.ones(()))
+    sigma: JAXArray | float = eqx.field(default_factory=lambda: jnp.zeros(()))
 
     def evaluate(self, X1: JAXArray, X2: JAXArray) -> JAXArray:
         return (

{tinygp-0.2.4 → tinygp-0.3.0}/src/tinygp/kernels/distance.py

@@ -11,14 +11,15 @@ from __future__ import annotations
 
 __all__ = ["Distance", "L1Distance", "L2Distance"]
 
-from abc import ABCMeta, abstractmethod
+from abc import abstractmethod
 
+import equinox as eqx
 import jax.numpy as jnp
 
-from tinygp.helpers import JAXArray, dataclass
+from tinygp.helpers import JAXArray
 
 
-class Distance(metaclass=ABCMeta):
+class Distance(eqx.Module):
     """An abstract base class defining a distance metric interface"""
 
     @abstractmethod
@@ -37,7 +38,6 @@ class Distance(metaclass=ABCMeta):
         return jnp.square(self.distance(X1, X2))
 
 
-@dataclass
 class L1Distance(Distance):
     """The L1 or Manhattan distance between two coordinates"""
 
@@ -45,7 +45,6 @@ class L1Distance(Distance):
         return jnp.sum(jnp.abs(X1 - X2))
 
 
-@dataclass
 class L2Distance(Distance):
     """The L2 or Euclidean distance between two coordinates"""