Implementing compatibility with Aesara (#150)

* adding aesara theano shim * don't push tutorials on PR * updating pre-commit hooks * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * pymc3 version for tutorials * upddatiing pymc3-ext * deprecating * reboundx installation issues * rebound install issues * rebound install issues * updating tutorial notebooks * one more notebook issues * verify grad has moved * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
exoplanet-dev · Mar 11, 2021 · dadf549 · dadf549
1 parent d46b9fd
commit dadf549
Show file tree

Hide file tree

Showing 55 changed files with 259 additions and 1,301 deletions.
diff --git a/.github/ISSUE_TEMPLATE/bug_report.md b/.github/ISSUE_TEMPLATE/bug_report.md
@@ -19,7 +19,7 @@ A clear and concise description of what you expected to happen.
 **Your setup (please complete the following information):**
  - Version of exoplanet:
  - Operating system:
- - Python version & installation method (pip, conda, etc.): 
+ - Python version & installation method (pip, conda, etc.):
 
 **Additional context**
 Add any other context about the problem here.
diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
@@ -12,52 +12,51 @@ on:
 
 jobs:
   tests:
-    name: "${{ matrix.pymc-version }} on py${{ matrix.python-version }}"
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.8"]
-        pymc-version: ["pymc3==3.9.3", "pymc3==3.10", "pymc3==3.11"]
+        python-version: ["3.9"]
+        arviz-version: ["arviz"]
+        pymc-version:
+          - "pymc3~=3.11.0"
+          - "https://github.com/pymc-devs/pymc3/archive/master.zip"
         include:
-          - python-version: "3.8"
-            pymc-version: "https://github.com/pymc-devs/pymc3/archive/master.zip"
+          - python-version: "3.9"
+            arviz-version: "arviz<0.11"
+            pymc-version: "pymc3~=3.9.0"
+          - python-version: "3.9"
+            arviz-version: "arviz<0.11"
+            pymc-version: "pymc3~=3.10.0"
 
     steps:
       - uses: actions/checkout@v2
         with:
           fetch-depth: 0
           submodules: true
 
-      - uses: actions/cache@v2
+      - uses: actions/setup-python@v2
+        name: Install Python
         with:
-          path: ~/.theano
-          key: xo-${{ runner.os }}-py${{ matrix.python-version }}-${{ matrix.pymc-version }}-${{ github.sha }}
-          restore-keys: |
-            xo-${{ runner.os }}-py${{ matrix.python-version }}-${{ matrix.pymc-version }}-
-
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: conda-incubator/setup-miniconda@v2
-        with:
-          auto-update-conda: true
           python-version: ${{ matrix.python-version }}
 
       # Setup
       - name: Install dependencies
         shell: bash -l {0}
         run: |
           python -m pip install -U pip
-          python -m pip install "${{ matrix.pymc-version }}"
+          python -m pip install "${{ matrix.pymc-version }}" "${{ matrix.arviz-version }}"
+          python -m pip install rebound
+          python -m pip install "reboundx>=3.2.0"
           python -m pip install -e .[test]
           python -m pip install --no-deps starry rebound-pymc3
-          python -c 'import theano; print(theano.config.__str__(print_doc=False))'
 
       - name: Run tests
         shell: bash -l {0}
         run: python -m pytest --cov=exoplanet -v tests
 
       - name: Coveralls
-        uses: AndreMiras/coveralls-python-action@develop
+        uses: AndreMiras/coveralls-python-action@v20201129
         with:
           parallel: true
           flag-name: Unit Tests
@@ -68,6 +67,6 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - name: Coveralls Finished
-        uses: AndreMiras/coveralls-python-action@develop
+        uses: AndreMiras/coveralls-python-action@v20201129
         with:
           parallel-finished: true
diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml
@@ -5,6 +5,7 @@ on:
       - main
   pull_request:
     branches:
+      - main
   release:
     types:
       - published
@@ -28,12 +29,12 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v2
         with:
-          python-version: 3.8
+          python-version: 3.9
 
       - name: Install dependencies
         run: |
           python -m pip install -U pip
-          python -m pip install "pymc3<3.10"
+          python -m pip install "pymc3~=3.10"
           python -m pip install ".[tutorials]"
 
       - name: Execute the notebooks
@@ -46,6 +47,7 @@ jobs:
           path: docs/tutorials
 
       - name: Trigger RTDs build
+        if: ${{ github.event_name != 'pull_request' }}
         uses: dfm/rtds-action@v1.0.0
         with:
           webhook_url: ${{ secrets.RTDS_WEBHOOK_URL }}

diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml
@@ -79,7 +79,7 @@ jobs:
           name: artifact
           path: dist
 
-      - uses: pypa/gh-action-pypi-publish@master
+      - uses: pypa/gh-action-pypi-publish@v1.4.2
         with:
           user: __token__
           password: ${{ secrets.pypi_password }}

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -1,6 +1,6 @@
 repos:
 - repo: https://github.com/pre-commit/pre-commit-hooks
-  rev: v2.3.0
+  rev: v3.4.0
   hooks:
   - id: trailing-whitespace
     exclude: src/exoplanet/theano_ops/lib/vendor
@@ -14,7 +14,7 @@ repos:
     exclude: docs/tutorials
 
 - repo: https://github.com/PyCQA/isort
-  rev: "5.4.2"
+  rev: "5.7.0"
   hooks:
   - id: isort
     args: []
@@ -27,6 +27,6 @@ repos:
   - id: black
 
 - repo: https://github.com/dfm/black_nbconvert
-  rev: stable
+  rev: v0.2.0
   hooks:
   - id: black_nbconvert
diff --git a/docs/_themes/exoplanet/globaltoc.html b/docs/_themes/exoplanet/globaltoc.html
@@ -3,4 +3,4 @@
     collapse=theme_collapse_navigation|tobool,
     includehidden=theme_includehidden|tobool,
     titles_only=theme_titles_only|tobool) }}
-</div>
+</div>
diff --git a/docs/tutorials/astrometric.py b/docs/tutorials/astrometric.py
@@ -5,7 +5,7 @@
 #       extension: .py
 #       format_name: light
 #       format_version: '1.5'
-#       jupytext_version: 1.6.0
+#       jupytext_version: 1.7.1
 #   kernelspec:
 #     display_name: Python 3
 #     language: python
@@ -14,7 +14,6 @@
 
 # + nbsphinx="hidden"
 # %matplotlib inline
-# -
 
 # + nbsphinx="hidden"
 # %run notebook_setup
@@ -127,11 +126,12 @@
 
 # +
 import pymc3 as pm
-import theano
-import theano.tensor as tt
+import pymc3_ext as pmx
+
+from aesara_theano_fallback import aesara as theano
+import aesara_theano_fallback.tensor as tt
 
 import exoplanet as xo
-from exoplanet.distributions import Angle
 
 from astropy import constants
 
@@ -235,15 +235,15 @@ def get_model(parallax=None):
         # directly
         omega0 = 251.6 * deg - np.pi
         Omega0 = 159.6 * deg
-        p = Angle("p", testval=0.5 * (Omega0 + omega0))
-        m = Angle("m", testval=0.5 * (Omega0 - omega0))
+        p = pmx.Angle("p", testval=0.5 * (Omega0 + omega0))
+        m = pmx.Angle("m", testval=0.5 * (Omega0 - omega0))
         omega = pm.Deterministic("omega", p - m)
         Omega = pm.Deterministic("Omega", p + m)
 
         # For these orbits, it can also be better to fit for a phase angle
         # (relative to a reference time) instead of the time of periasteron
         # passage directly
-        phase = Angle("phase", testval=0.0)
+        phase = pmx.Angle("phase", testval=0.0)
         tperi = pm.Deterministic("tperi", T0 + P * phase / (2 * np.pi))
 
         # Geometric uiform prior on cos(incl)
@@ -299,11 +299,11 @@ def get_model(parallax=None):
 
         # Optimize to find the initial parameters
         map_soln = model.test_point
-        map_soln = xo.optimize(map_soln, vars=[log_rho_s, log_theta_s])
-        map_soln = xo.optimize(map_soln, vars=[phase])
-        map_soln = xo.optimize(map_soln, vars=[p, m, ecc])
-        map_soln = xo.optimize(map_soln, vars=[logP, a_ang, phase])
-        map_soln = xo.optimize(map_soln)
+        map_soln = pmx.optimize(map_soln, vars=[log_rho_s, log_theta_s])
+        map_soln = pmx.optimize(map_soln, vars=[phase])
+        map_soln = pmx.optimize(map_soln, vars=[p, m, ecc])
+        map_soln = pmx.optimize(map_soln, vars=[logP, a_ang, phase])
+        map_soln = pmx.optimize(map_soln)
 
     return model, map_soln
 
@@ -350,41 +350,37 @@ def get_model(parallax=None):
 
 np.random.seed(1234)
 with model:
-    trace = pm.sample(
+    trace = pmx.sample(
         tune=5000,
         draws=4000,
         start=map_soln,
         cores=2,
         chains=2,
-        init="adapt_full",
         target_accept=0.9,
     )
 
 # First we can check the convergence for some of the key parameters.
 
+# +
+import arviz as az
+
 with model:
-    summary = pm.summary(
+    summary = az.summary(
         trace,
-        varnames=["P", "tperi", "a_ang", "omega", "Omega", "incl", "ecc"],
+        var_names=["P", "tperi", "a_ang", "omega", "Omega", "incl", "ecc"],
     )
 summary
+# -
 
 # That looks pretty good.
 # Now here's a corner plot showing the covariances between parameters.
 
 # +
 import corner
 
-samples = pm.trace_to_dataframe(trace, varnames=["ecc"])
-samples["$P$ [yr]"] = trace["P"] / yr
-samples["$T_\mathrm{peri} - T_0$ [day]"] = trace["tperi"] - T0
-samples["$a$ [arcsec]"] = trace["a_ang"]
-samples["$\omega$ [deg]"] = (trace["omega"] / deg) % 360
-samples["$\Omega$ [deg]"] = (trace["Omega"] / deg) % 360
-samples["$i$ [deg]"] = (trace["incl"] / deg) % 360
-samples["$e$"] = samples["ecc"]
-del samples["ecc"]
-_ = corner.corner(samples)
+_ = corner.corner(
+    trace, var_names=["P", "tperi", "a_ang", "omega", "Omega", "incl", "ecc"]
+)
 # -
 
 # Finally, we can plot the posterior constraints on $\rho$ and $\theta$.
@@ -432,22 +428,21 @@ def get_model(parallax=None):
 
 np.random.seed(5432)
 with plx_model:
-    plx_trace = pm.sample(
+    plx_trace = pmx.sample(
         tune=5000,
         draws=4000,
         start=plx_map_soln,
         cores=2,
         chains=2,
-        init="adapt_full",
         target_accept=0.9,
     )
 
 # Check the convergence diagnostics.
 
 with model:
-    summary = pm.summary(
+    summary = az.summary(
         plx_trace,
-        varnames=[
+        var_names=[
             "P",
             "tperi",
             "a_ang",
@@ -463,14 +458,7 @@ def get_model(parallax=None):
 
 # And make the corner plot for the physical parameters.
 
-samples = pm.trace_to_dataframe(plx_trace, varnames=["ecc"])
-samples["$P$ [yr]"] = plx_trace["P"] / yr
-samples["$T_\mathrm{peri} - T_0$ [day]"] = plx_trace["tperi"] - T0
-samples["$a$ [au]"] = plx_trace["a"]
-samples["$M_\mathrm{tot}$ [$M_\odot$]"] = plx_trace["M_tot"]
-samples["$e$"] = plx_trace["ecc"]
-del samples["ecc"]
-_ = corner.corner(samples)
+_ = corner.corner(plx_trace, var_names=["P", "tperi", "a", "ecc", "M_tot"])
 
 # ## Citations
 #

diff --git a/docs/tutorials/citation.py b/docs/tutorials/citation.py
@@ -5,7 +5,7 @@
 #       extension: .py
 #       format_name: light
 #       format_version: '1.5'
-#       jupytext_version: 1.6.0
+#       jupytext_version: 1.7.1
 #   kernelspec:
 #     display_name: Python 3
 #     language: python