Adding support for beta version of PyMC v4 and Python 3.10 (#27)

.github/workflows/tests.yml

@@ -12,21 +12,18 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ["3.9"]
+        python-version: ["3.9", "3.10"]
         os: ["ubuntu-latest"]
         extras:
+          - 'starry batman-package "numpy<1.22"'
           - "jax jaxlib"
-          - 'aesara-theano-fallback "pymc3~=3.9.0" "arviz<0.11"'
-          - 'aesara-theano-fallback "pymc3~=3.10.0" "arviz<0.11"'
-          - 'aesara-theano-fallback "pymc3~=3.11.0"'
-          - 'aesara-theano-fallback "https://github.com/pymc-devs/pymc3/archive/v3.zip"'
+          - 'pymc3 "numpy<1.22" aesara-theano-fallback'
+          - '"https://github.com/pymc-devs/pymc/archive/v3.zip" "numpy<1.22" aesara-theano-fallback'
+          - '"https://github.com/pymc-devs/pymc/archive/main.zip" aesara-theano-fallback'
         include:
-          - os: "windows-latest"
-            python-version: "3.9"
-            extras: "aesara-theano-fallback pymc3"
           - os: "macos-latest"
-            python-version: "3.9"
-            extras: "aesara-theano-fallback pymc3"
+            python-version: "3.10"
+            extras: 'pymc3 "numpy<1.22" aesara-theano-fallback'
 
     steps:
       - uses: actions/checkout@v2
@@ -44,52 +41,5 @@ jobs:
           python -m pip install ${{ matrix.extras }}
           python -m pip install -e .[test]
 
-      - name: Install starry and batman for comparisons on Linux
-        if: ${{ matrix.os == 'ubuntu-latest' }}
-        run: |
-          python -m pip install starry batman-package
-
       - name: Run tests
         run: python -m pytest -v tests
-
-  benchmarks:
-    name: Benchmark ops
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v2
-        with:
-          fetch-depth: 0
-
-      - uses: actions/setup-python@v2
-        name: Install Python
-        with:
-          python-version: 3.9
-
-      - name: Install dependencies
-        run: |
-          python -m pip install -U pip
-          python -m pip install cython h5py
-          python -m pip install -e .[benchmark]
-
-      - name: Get dependency versions
-        run: |
-          python -c 'import radvel;print(radvel.__version__)'
-          python -c 'import batman;print(batman.__version__)'
-          python -c 'import starry;print(starry.__version__)'
-          python -c 'import kepler;print(kepler.__version__)'
-
-      - name: Run tests
-        run: python -m pytest -v tests/benchmark.py --benchmark-json benchmark.json
-
-      - name: Store benchmark result
-        uses: rhysd/github-action-benchmark@v1
-        with:
-          name: Op performance
-          tool: "pytest"
-          output-file-path: benchmark.json
-          github-token: ${{ secrets.PERSONAL_GITHUB_TOKEN }}
-          auto-push: true
-          alert-threshold: "200%"
-          comment-on-alert: true
-          fail-on-alert: true
-          alert-comment-cc-users: "@dfm"

.github/workflows/wheels.yml

@@ -14,14 +14,11 @@ jobs:
         with:
           submodules: true
           fetch-depth: 0
-      - uses: ilammy/msvc-dev-cmd@v1
-      - uses: joerick/cibuildwheel@v2.3.1
+      - uses: pypa/cibuildwheel@v2.3.1
         env:
-          CIBW_BUILD: "cp3?-*"
-          CIBW_SKIP: "cp35-* *-win32 *-manylinux_i686"
+          CIBW_SKIP: "pp* *-win32 *-manylinux_i686 *-musllinux*"
           CIBW_MANYLINUX_X86_64_IMAGE: manylinux2014
-          DISTUTILS_USE_SDK: 1
-          MSSdk: 1
+          CIBW_ARCHS_MACOS: "x86_64 universal2 arm64"
       - uses: actions/upload-artifact@v2
         with:
           path: ./wheelhouse/*.whl
@@ -56,7 +53,7 @@ jobs:
           name: artifact
           path: dist
 
-      - uses: pypa/gh-action-pypi-publish@v1.4.2
+      - uses: pypa/gh-action-pypi-publish@v1.5.0
         with:
           user: __token__
           password: ${{ secrets.pypi_password }}

src/exoplanet_core/driver.cpp

@@ -14,10 +14,10 @@ struct flat_unchecked_array {
     data = (Scalar *)info.ptr;
   }
 
-  inline Scalar &operator()(ssize_t index) { return data[index]; }
-  inline ssize_t shape(ssize_t index) const { return info.shape[index]; }
-  inline ssize_t size() const { return info.size; }
-  inline ssize_t ndim() const { return info.ndim; }
+  inline Scalar &operator()(py::ssize_t index) { return data[index]; }
+  inline py::ssize_t shape(py::ssize_t index) const { return info.shape[index]; }
+  inline py::ssize_t size() const { return info.size; }
+  inline py::ssize_t ndim() const { return info.ndim; }
 
   py::buffer_info info;
   Scalar *data;
@@ -29,10 +29,10 @@ void solve_kepler(py::array_t<double, py::array::c_style> M_in,
                   py::array_t<double, py::array::c_style> cosf_out) {
   flat_unchecked_array<double, py::array::c_style> M(M_in), ecc(ecc_in);
   flat_unchecked_array<double, py::array::c_style> sinf(sinf_out, true), cosf(cosf_out, true);
-  ssize_t N = M.size();
+  py::ssize_t N = M.size();
   if (ecc.size() != N || cosf.size() != N || sinf.size() != N)
     throw std::invalid_argument("dimension mismatch");
-  for (ssize_t n = 0; n < N; ++n) {
+  for (py::ssize_t n = 0; n < N; ++n) {
     if (ecc(n) < 0 || ecc(n) > 1)
       throw std::invalid_argument("eccentricity must be in the range [0, 1)");
     exoplanet::kepler::solve_kepler(M(n), ecc(n), &(sinf(n)), &(cosf(n)));
@@ -47,9 +47,9 @@ void quad_solution_vector(py::array_t<double, py::array::c_style> b_in,
   flat_unchecked_array<double, py::array::c_style> s(s_out, true);
   const double eps = std::numeric_limits<double>::epsilon();
 
-  ssize_t N = b.size();
+  py::ssize_t N = b.size();
   if (r.size() != N || s.size() != 3 * N) throw std::invalid_argument("dimension mismatch");
-  for (ssize_t n = 0; n < N; ++n) {
+  for (py::ssize_t n = 0; n < N; ++n) {
     exoplanet::limbdark::quad_solution_vector<false>(eps, std::abs(b(n)), r(n), &(s(3 * n)),
                                                      (double *)NULL, (double *)NULL);
   }
@@ -65,11 +65,11 @@ void quad_solution_vector_with_grad(py::array_t<double, py::array::c_style> b_in
   flat_unchecked_array<double, py::array::c_style> s(s_out, true), dsdb(dsdb_out, true),
       dsdr(dsdr_out, true);
   const double eps = std::numeric_limits<double>::epsilon();
-  ssize_t N = b.size();
+  py::ssize_t N = b.size();
   if (r.size() != N || s.size() != 3 * N || dsdb.size() != 3 * N || dsdr.size() != 3 * N)
     throw std::invalid_argument("dimension mismatch");
-  for (ssize_t n = 0; n < N; ++n) {
-    ssize_t ind = 3 * n;
+  for (py::ssize_t n = 0; n < N; ++n) {
+    py::ssize_t ind = 3 * n;
     int sgn = exoplanet::sgn(b(n));
     exoplanet::limbdark::quad_solution_vector<true>(eps, std::abs(b(n)), r(n), &(s(ind)),
                                                     &(dsdb(ind)), &(dsdr(ind)));
@@ -96,12 +96,12 @@ void contact_points(py::array_t<double, py::array::c_style> a_in,
   flat_unchecked_array<double, py::array::c_style> M_left(M_left_out, true),
       M_right(M_right_out, true);
   flat_unchecked_array<int, py::array::c_style> flag(flag_out, true);
-  ssize_t N = a.size();
+  py::ssize_t N = a.size();
   if (e.size() != N || cosw.size() != N || sinw.size() != N || cosi.size() != N ||
       sini.size() != N || L.size() != N || M_left.size() != N || M_right.size() != N ||
       flag.size() != N)
     throw std::invalid_argument("dimension mismatch");
-  for (ssize_t n = 0; n < N; ++n) {
+  for (py::ssize_t n = 0; n < N; ++n) {
     auto const solver = exoplanet::contact_points::ContactPointSolver<double>(
         a(n), e(n), cosw(n), sinw(n), cosi(n), sini(n));
     auto const roots = solver.find_roots(L(n), tol);

tests/core_test.py

@@ -5,19 +5,9 @@
 
 from exoplanet_core import core
 
-try:
-    import starry
-except ImportError:
-    starry = None
 
-try:
-    import batman
-except ImportError:
-    batman = None
-
-
-@pytest.mark.skipif(starry is None, reason="starry is not installed")
 def test_compare_starry():
+    starry = pytest.importorskip("starry")
     u1 = 0.2
     u2 = 0.3
     b = np.linspace(-1.5, 1.5, 100)
@@ -31,8 +21,8 @@ def test_compare_starry():
     np.testing.assert_allclose(computed, expected, atol=1e-12)
 
 
-@pytest.mark.skipif(batman is None, reason="batman is not installed")
 def test_compare_batman():
+    batman = pytest.importorskip("batman")
     u1 = 0.2
     u2 = 0.3
     b = np.linspace(-1.5, 1.5, 100)

tests/pymc/pymc_ops_test.py

@@ -9,9 +9,11 @@
     get_mean_and_true_anomaly,
 )
 
-aesara = pytest.importorskip("aesara_theano_fallback.aesara")
+fallback = pytest.importorskip("aesara_theano_fallback")
 ops = pytest.importorskip("exoplanet_core.pymc.ops")
 
+aesara = fallback.aesara
+
 
 def compare_to_numpy(nop, op, *args):
     expected = nop(*args)
@@ -39,16 +41,21 @@ def test_kepler(kepler_data):
 
 
 def test_kepler_grad(kepler_data):
+    try:
+        verify_grad = aesara.tensor.verify_grad
+    except AttributeError:
+        from aesara.gradient import verify_grad
+
     M, e, f = kepler_data
     np.random.seed(1324)
     ind = 100
-    aesara.tensor.verify_grad(
+    verify_grad(
         lambda *x: ops.kepler(*x)[0],
         [M[:, ind], e[:, ind]],
         rng=np.random,
         eps=1e-8,
     )
-    aesara.tensor.verify_grad(
+    verify_grad(
         lambda *x: ops.kepler(*x)[1],
         [M[:, ind], e[:, ind]],
         rng=np.random,
@@ -71,16 +78,19 @@ def test_quad_solution_vector(limbdark_data):
 
 
 def test_quad_solution_vector_grad(limbdark_data):
+    try:
+        verify_grad = aesara.tensor.verify_grad
+    except AttributeError:
+        from aesara.gradient import verify_grad
+
     # The numerical estimate is bad at discontinuities
     eps = 1e-7
     b, r = limbdark_data
     m = np.abs(b) > 2 * eps
     m &= np.abs(np.abs(b) - (1 - r)) > 2 * eps
     m &= np.abs(np.abs(b) - (1 + r)) > 2 * eps
 
-    aesara.tensor.verify_grad(
-        ops.quad_solution_vector, (b[m], r[m]), rng=np.random, eps=eps
-    )
+    verify_grad(ops.quad_solution_vector, (b[m], r[m]), rng=np.random, eps=eps)
 
 
 @pytest.mark.parametrize("a", [5.0, 12.1234, 20000.0])