Skip to content

Commit 4f3bfc9

Browse files
dfmdfm
authored
Adding support for PyMC v4 (#67)
1 parent c6874e2 commit 4f3bfc9

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

60 files changed

+3437
-1527
lines changed

‎.github/workflows/python.yml‎

Lines changed: 22 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -15,15 +15,28 @@ concurrency:
1515

1616
jobs:
1717
tests:
18-
runs-on: ubuntu-latest
18+
runs-on: ${{ matrix.os }}
1919
strategy:
2020
fail-fast: false
2121
matrix:
22-
python-version: ["3.9"]
23-
extras: ["core", "jax", "pymc3"]
22+
python-version: ["3.10"]
23+
os: ["ubuntu-latest"]
24+
session:
25+
- "core"
26+
- "jax"
27+
- "pymc3"
28+
- "pymc4"
29+
- "pymc4_jax"
2430
include:
25-
- python-version: "3.10"
26-
extras: "core"
31+
- os: "ubuntu-latest"
32+
python-version: "3.9"
33+
session: "full"
34+
- os: "macos-latest"
35+
python-version: "3.9"
36+
session: "full"
37+
- os: "macos-latest"
38+
python-version: "3.10"
39+
session: "full"
2740

2841
steps:
2942
- name: Checkout
@@ -40,12 +53,12 @@ jobs:
4053
- name: Install dependencies
4154
run: |
4255
python -m pip install -U pip
43-
python -m pip install -U coveralls coverage[toml] tox tox-gh-actions pybind11
56+
python -m pip install -U coveralls coverage[toml] nox
4457
4558
- name: Run tests
46-
run: python -m tox
47-
env:
48-
EXTRAS: ${{ matrix.extras }}
59+
run: |
60+
python -m nox --non-interactive --error-on-missing-interpreter \
61+
--session ${{ matrix.session }}-${{ matrix.python-version }}
4962
5063
- name: Combine and upload coverage
5164
run: |

‎.github/workflows/tutorials.yml‎ renamed to ‎.github/workflows/tutorials.yml.off‎

Lines changed: 12 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -31,22 +31,21 @@ jobs:
3131
- name: Install dependencies
3232
run: |
3333
python -m pip install -U pip
34-
python -m pip install "numpy<1.22"
3534
python -m pip install ".[tutorials]"
3635

37-
- name: Get theano compiledir
38-
id: compiledir
39-
run: |
40-
python -c "import theano; print('::set-output name=compiledir::' + theano.config.compiledir.split('/')[-1])"
36+
# - name: Get theano compiledir
37+
# id: compiledir
38+
# run: |
39+
# python -c "import theano; print('::set-output name=compiledir::' + theano.config.compiledir.split('/')[-1])"
4140

42-
- name: "Cache ~/.theano"
43-
uses: actions/cache@v2
44-
with:
45-
path: ~/.theano
46-
key: tutorials-${{ steps.compiledir.outputs.compiledir }}-${{ hashFiles('docs/tutorials/*.py') }}
47-
restore-keys: |
48-
tutorials-${{ steps.compiledir.outputs.compiledir }}-
49-
tutorials-
41+
# - name: "Cache ~/.theano"
42+
# uses: actions/cache@v2
43+
# with:
44+
# path: ~/.theano
45+
# key: tutorials-${{ steps.compiledir.outputs.compiledir }}-${{ hashFiles('docs/tutorials/*.py') }}
46+
# restore-keys: |
47+
# tutorials-${{ steps.compiledir.outputs.compiledir }}-
48+
# tutorials-
5049

5150
- name: Execute the notebooks
5251
run: |

‎.gitignore‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,9 @@ build
3434
dist
3535
.envrc
3636
.coverage
37-
*.ipynb
37+
/*.ipynb
3838
.coverage*
3939
.eggs
4040
.tox
41+
.nox
42+
.coverage.*

‎.pre-commit-config.yaml‎

Lines changed: 14 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -4,22 +4,23 @@ repos:
44
hooks:
55
- id: trailing-whitespace
66
- id: end-of-file-fixer
7-
- id: debug-statements
8-
7+
exclude_types: [json, binary]
98
- repo: https://github.com/PyCQA/isort
109
rev: "5.10.1"
1110
hooks:
12-
- id: isort
13-
args: []
14-
additional_dependencies: [toml]
15-
exclude: docs/tutorials
16-
11+
- id: isort
12+
additional_dependencies: [toml]
13+
exclude: docs/tutorials
1714
- repo: https://github.com/psf/black
18-
rev: 22.3.0
15+
rev: "22.3.0"
1916
hooks:
20-
- id: black
21-
22-
- repo: https://github.com/dfm/black_nbconvert
23-
rev: v0.4.0
17+
- id: black-jupyter
18+
- repo: https://github.com/kynan/nbstripout
19+
rev: "0.5.0"
2420
hooks:
25-
- id: black_nbconvert
21+
- id: nbstripout
22+
- repo: https://github.com/hadialqattan/pycln
23+
rev: "v1.3.1"
24+
hooks:
25+
- id: pycln
26+
additional_dependencies: ["click<8.1.0"]

‎docs/api/jax.rst‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
.. _jax-api:
2+
3+
JAX interface
4+
=============
5+
6+
This ``celerite2.jax`` submodule provides an interface to *celerite2* models
7+
that can be used from `JAX <https://jax.readthedocs.io>`_.
8+
9+
The :ref:`first` tutorial demonstrates the use of this interface, while this
10+
page provides the details for the :class:`celerite2.jax.GaussianProcess` class
11+
which provides all this functionality. This page does not include documentation
12+
for the term models defined in JAX, but you can refer to the
13+
:ref:`python-terms` section of the :ref:`python-api` documentation. All of those
14+
models are implemented in JAX and you can access them using something like
15+
the following:
16+
17+
.. code-block:: python
18+
19+
import jax
20+
from celerite2.jax import GaussianProcess, terms
21+
22+
@jax.jit
23+
def log_likelihood(params, x, diag, y):
24+
term = terms.SHOTerm(S0=params["S0"], w0=params["w0"], Q=params["Q"])
25+
gp = GaussianProcess(term)
26+
gp.compute(x, diag=diag)
27+
return gp.log_likelihood(y)
28+
29+
The :class:`celerite2.jax.GaussianProcess` class is detailed below:
30+
31+
.. autoclass:: celerite2.jax.GaussianProcess
32+
:inherited-members:
33+
:exclude-members: sample, sample_conditional, recompute
34+
35+
36+
numpyro support
37+
---------------
38+
39+
This implementation comes with a custom `numpyro <https://num.pyro.ai>`_
40+
``Distribution`` that represents a multivariate normal with a *celerite*
41+
covariance matrix. This is used by the
42+
:func:`celerite2.jax.GaussianProcess.numpyro_dist` method documented above which
43+
adds a marginal likelihood node to a numpyro model.
44+
45+
.. autoclass:: celerite2.jax.distribution.CeleriteNormal

‎docs/api/pymc3.rst‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
.. _pymc3-api:
2+
3+
PyMC3 interface
4+
===============
5+
6+
This ``celerite2.pymc3`` submodule provides access to the *celerite2* models
7+
within the `Theano <http://deeplearning.net/software/theano/>`_ framework. Of
8+
special interest, this adds support for probabilistic model building using
9+
`PyMC3 <https://docs.pymc.io/>`_.
10+
11+
This page does not include documentation for the term models defined in Theano,
12+
but you can refer to the :ref:`python-terms` section of the :ref:`python-api`
13+
documentation. All of those models are implemented in Theano and you can access
14+
them using something like the following:
15+
16+
.. code-block:: python
17+
18+
import theano.tensor as tt
19+
from celerite2.pymc3 import GaussianProcess, terms
20+
21+
term = terms.SHOTerm(S0=tt.dscalar(), w0=tt.dscalar(), Q=tt.dscalar())
22+
gp = GaussianProcess(term)
23+
24+
The :class:`celerite2.pymc3.GaussianProcess` class is detailed below:
25+
26+
.. autoclass:: celerite2.pymc3.GaussianProcess
27+
:inherited-members:
28+
:exclude-members: sample, sample_conditional, recompute
29+
30+
31+
PyMC3 support
32+
-------------
33+
34+
This implementation comes with a custom PyMC3 ``Distribution`` that represents a
35+
multivariate normal with a *celerite* covariance matrix. This is used by the
36+
:func:`celerite2.pymc3.GaussianProcess.marginal` method documented above which
37+
adds a marginal likelihood node to a PyMC3 model.
38+
39+
.. autoclass:: celerite2.pymc3.distribution.CeleriteNormal

‎docs/api/pymc4.rst‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
.. _pymc4-api:
2+
3+
PyMC4 interface
4+
===============
5+
6+
This ``celerite2.pymc4`` submodule provides access to the *celerite2* models
7+
within the `Aesara <https://aesara.readthedocs.io>`_ framework. Of special
8+
interest, this adds support for probabilistic model building using `PyMC
9+
<https://docs.pymc.io/>`_ v4 or later.
10+
11+
The :ref:`first` tutorial demonstrates the use of this interface, while this
12+
page provides the details for the :class:`celerite2.pymc4.GaussianProcess` class
13+
which provides all this functionality. This page does not include documentation
14+
for the term models defined in Aesara, but you can refer to the
15+
:ref:`python-terms` section of the :ref:`python-api` documentation. All of those
16+
models are implemented in Aesara and you can access them using something like
17+
the following:
18+
19+
.. code-block:: python
20+
21+
import aesara.tensor as at
22+
from celerite2.pymc4 import GaussianProcess, terms
23+
24+
term = terms.SHOTerm(S0=at.scalar(), w0=at.scalar(), Q=at.scalar())
25+
gp = GaussianProcess(term)
26+
27+
The :class:`celerite2.pymc4.GaussianProcess` class is detailed below:
28+
29+
.. autoclass:: celerite2.pymc4.GaussianProcess
30+
:inherited-members:
31+
:exclude-members: sample, sample_conditional, recompute
32+
33+
34+
PyMC (v4) support
35+
-----------------
36+
37+
This implementation comes with a custom PyMC ``Distribution`` that represents a
38+
multivariate normal with a *celerite* covariance matrix. This is used by the
39+
:func:`celerite2.pymc4.GaussianProcess.marginal` method documented above which
40+
adds a marginal likelihood node to a PyMC model.
41+
42+
.. autoclass:: celerite2.pymc4.distribution.CeleriteNormal

‎docs/api/theano.rst‎

Lines changed: 0 additions & 43 deletions
This file was deleted.

‎docs/conf.py‎

Lines changed: 9 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,5 @@
11
#!/usr/bin/env python
22
# -*- coding: utf-8 -*-
3-
import os
43
import subprocess
54

65
import sphinx_material
@@ -9,7 +8,7 @@
98
try:
109
__version__ = get_distribution("celerite2").version
1110
except DistributionNotFound:
12-
__version__ = "unknown version"
11+
__version__ = "dev"
1312

1413

1514
subprocess.call("doxygen Doxyfile.in", shell=True)
@@ -19,26 +18,13 @@
1918
"sphinx.ext.intersphinx",
2019
"sphinx.ext.napoleon",
2120
"sphinx.ext.mathjax",
21+
"sphinx.ext.viewcode",
2222
"sphinx_copybutton",
23-
"nbsphinx",
23+
"myst_nb",
2424
"breathe",
2525
]
2626

27-
autodoc_mock_imports = [
28-
"numpy",
29-
"theano",
30-
"pymc3",
31-
"celerite2.driver",
32-
"celerite2.backprop",
33-
]
34-
35-
intersphinx_mapping = {
36-
"python": ("https://docs.python.org/3/", None),
37-
"numpy": ("https://docs.scipy.org/doc/numpy/", None),
38-
"scipy": ("https://docs.scipy.org/doc/scipy/reference/", None),
39-
"pymc3": ("https://docs.pymc.io/", None),
40-
"celerite": ("https://celerite.readthedocs.io/en/stable/", None),
41-
}
27+
autodoc_mock_imports = ["theano", "pymc3"]
4228

4329
exclude_patterns = ["_build"]
4430
templates_path = ["_templates"]
@@ -48,19 +34,10 @@
4834
# General information about the project.
4935
project = "celerite2"
5036
author = "Dan Foreman-Mackey"
51-
copyright = "2020, " + author
37+
copyright = "2020-2022" + author
5238
version = __version__
5339
release = __version__
5440

55-
# RTDs-action
56-
if "GITHUB_TOKEN" in os.environ:
57-
extensions.append("rtds_action")
58-
59-
rtds_action_github_repo = "exoplanet-dev/celerite2"
60-
rtds_action_path = "tutorials"
61-
rtds_action_artifact_prefix = "notebooks-for-"
62-
rtds_action_github_token = os.environ["GITHUB_TOKEN"]
63-
6441
# HTML theme
6542
html_show_sourcelink = False
6643
html_sidebars = {
@@ -107,3 +84,7 @@
10784
breathe_projects = {"celerite2": "c++/xml/"}
10885
breathe_default_project = "celerite2"
10986
breathe_default_members = ("members", "undoc-members")
87+
88+
# Myst notebooks
89+
nb_execution_mode = "auto"
90+
nb_execution_timeout = -1

0 commit comments

Comments
 (0)