Merge pull request #137 from adrn/numpy20

Compatibility with numpy 2.0

.github/workflows/cd.yml

@@ -9,6 +9,7 @@ on:
   release:
     types:
       - published
+      - edited
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.ref }}
@@ -71,7 +72,7 @@ jobs:
     permissions:
       id-token: write
     runs-on: ubuntu-latest
-    if: github.event_name == 'release' && github.event.action == 'published'
+    if: github.event_name == 'release' && (github.event.action == 'published' || github.event.action == 'edited')
 
     steps:
       - name: Download packages built by build-and-inspect-python-package

.github/workflows/ci.yml

@@ -18,25 +18,42 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        python-version: ["3.9", "3.10", "3.11"]
+        python-version: ["3.10", "3.11", "3.12"]
         os: ["ubuntu-latest", "macos-latest"]
+        pymc-version: ["latest"]
+        include:
+          - os: ubuntu-latest
+            python-version: "3.10"
+            pymc-version: "5.10.0"  # Oldest pymc version
     steps:
+      - uses: actions/checkout@v4
 
-    - uses: actions/checkout@v4
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
 
-    - name: Set up Python ${{ matrix.python-version }}
-      uses: actions/setup-python@v5
-      with:
-        python-version: ${{ matrix.python-version }}
+      - name: "Install hdf5 for tables (linux)"
+        if: runner.os == 'ubuntu-latest'
+        run: sudo apt-get install libhdf5-serial-dev
 
-    - name: Install dependencies
-      run: |
-        python -m pip install -U pip
-        python -m pip install -e .[test]
+      - name : "Install hdf5 for tables (mac)"
+        if: runner.os == 'macos-latest'
+        run: brew install hdf5
 
-    - name: Test package
-      run: >-
-        python -m pytest -ra --cov --cov-report=xml --cov-report=term --durations=20
+      - name: Install dependencies
+        run: |
+          python -m pip install -U pip
+          python -m pip install -e .[test]
 
-    - name: Upload coverage report
-      uses: codecov/[email protected]
+      - name: Correct pymc version
+        if: matrix.pymc-version != 'latest'
+        run: |
+          python -m pip install pymc==${{ matrix.pymc-version }}
+
+      - name: Test package
+        run: >-
+          python -m pytest -ra --cov --cov-report=xml --cov-report=term --durations=20 .
+
+      - name: Upload coverage report
+        uses: codecov/[email protected]

docs/run_notebooks.py

@@ -47,6 +47,14 @@ def process_notebook(filename, kernel_name=None):
     nbsphinx_kernel_name = os.environ.get("NBSPHINX_KERNEL", "python3")
 
     for filename in sorted(glob.glob(pattern)):
-        success = process_notebook(filename, kernel_name=nbsphinx_kernel_name)
+        try:
+            success = process_notebook(filename, kernel_name=nbsphinx_kernel_name)
+        except Exception as e:
+            print(f"Error while processing: {filename}")
+            print(e)
+            success = False
+
         if not success:
             sys.exit(1)
+
+        print(f"✅ Processed: {filename}")

pyproject.toml

@@ -2,10 +2,10 @@
 requires = [
     "setuptools>=64",
     "setuptools_scm>=8",
-    "numpy",
+    "numpy>=2.0",
     "scipy",
     "cython",
-    "twobody"
+    "twobody>=0.9.1",
 ]
 build-backend = 'setuptools.build_meta'
 
@@ -14,20 +14,20 @@ name = "thejoker"
 authors = [{name = "Adrian Price-Whelan", email = "[email protected]"}]
 description = "A custom Monte Carlo sampler for the two-body problem."
 readme = "README.rst"
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 license.file = "LICENSE"
 dynamic = ["version"]
 dependencies = [
     "astropy",
-    "numpy",
-    "twobody>=0.9",
+    "numpy>=1.24.4",
+    "twobody>=0.9.1",
     "scipy",
     "h5py",
     "schwimmbad>=0.4",
-    "pymc>=5",
+    "pymc>=5.10.0",
     "pymc_ext>=1.0.1",
     "exoplanet-core[pymc]>=0.3.0",
-    "tables<3.9.2",  # https://github.com/PyTables/PyTables/issues/1093
+    "tables",
     "dill"
 ]
 

thejoker/distributions.py

@@ -3,54 +3,94 @@
 import numpy as np
 import pymc as pm
 import pytensor.tensor as pt
+from packaging.version import Version
 from pymc.distributions.dist_math import check_parameters
 from pytensor.tensor.random.basic import NormalRV, RandomVariable
 
 from thejoker.units import UNIT_ATTR_NAME
 
-__all__ = ["UniformLog", "FixedCompanionMass"]
-
-
-class UniformLogRV(RandomVariable):
-    name = "uniformlog"
-    ndim_supp = 0
-    ndims_params = [0, 0]
-    dtype = "floatX"
-
-    @classmethod
-    def rng_fn(cls, rng, a, b, size):
-        _fac = np.log(b) - np.log(a)
-        uu = rng.uniform(size=size)
-        return np.exp(uu * _fac + np.log(a))
-
-
-uniformlog = UniformLogRV()
+PYMC_GT_516 = Version(pm.__version__) >= Version("5.16.0")
 
+__all__ = ["UniformLog", "FixedCompanionMass"]
 
-class UniformLog(pm.Continuous):
-    rv_op = uniformlog
 
-    @classmethod
-    def dist(cls, a, b, **kwargs):
-        a = pt.as_tensor_variable(a)
-        b = pt.as_tensor_variable(b)
-        return super().dist([a, b], **kwargs)
-
-    def support_point(rv, size, a, b):
-        a, b = pt.broadcast_arrays(a, b)
-        return 0.5 * (a + b)
-
-    # TODO: remove this once new pymc version is released
-    moment = support_point
-
-    def logp(value, a, b):
-        _fac = pt.log(b) - pt.log(a)
-        res = -pt.as_tensor_variable(value) - pt.log(_fac)
-        return check_parameters(
-            res,
-            (a > 0) & (a < b),
-            msg="a > 0 and a < b",
-        )
+if PYMC_GT_516:
+
+    class UniformLogRV(RandomVariable):
+        name: str = "uniformlog"
+        dtype: str = "floatX"
+        signature: str = "(),()->()"
+
+        @classmethod
+        def rng_fn(cls, rng, a, b, size):
+            _fac = np.log(b) - np.log(a)
+            uu = rng.uniform(size=size)
+            return np.exp(uu * _fac + np.log(a))
+
+    uniformlog = UniformLogRV()
+
+    class UniformLog(pm.Continuous):
+        rv_op = uniformlog
+
+        @classmethod
+        def dist(cls, a, b, **kwargs):
+            a = pt.as_tensor_variable(a)
+            b = pt.as_tensor_variable(b)
+            return super().dist([a, b], **kwargs)
+
+        def support_point(rv, size, a, b):
+            a, b = pt.broadcast_arrays(a, b)
+            return 0.5 * (a + b)
+
+        def logp(value, a, b):
+            _fac = pt.log(b) - pt.log(a)
+            res = -pt.as_tensor_variable(value) - pt.log(_fac)
+            return check_parameters(
+                res,
+                (a > 0) & (a < b),
+                msg="a > 0 and a < b",
+            )
+
+else:  # old behavior
+
+    class UniformLogRV(RandomVariable):
+        name: str = "uniformlog"
+        dtype: str = "floatX"
+        ndim_supp: int = 0
+        ndims_params: list[int] = [0, 0]
+
+        @classmethod
+        def rng_fn(cls, rng, a, b, size):
+            _fac = np.log(b) - np.log(a)
+            uu = rng.uniform(size=size)
+            return np.exp(uu * _fac + np.log(a))
+
+    uniformlog = UniformLogRV()
+
+    class UniformLog(pm.Continuous):
+        rv_op = uniformlog
+
+        @classmethod
+        def dist(cls, a, b, **kwargs):
+            a = pt.as_tensor_variable(a)
+            b = pt.as_tensor_variable(b)
+            return super().dist([a, b], **kwargs)
+
+        def support_point(rv, size, a, b):
+            a, b = pt.broadcast_arrays(a, b)
+            return 0.5 * (a + b)
+
+        # TODO: remove this once new pymc version is released
+        moment = support_point
+
+        def logp(value, a, b):
+            _fac = pt.log(b) - pt.log(a)
+            res = -pt.as_tensor_variable(value) - pt.log(_fac)
+            return check_parameters(
+                res,
+                (a > 0) & (a < b),
+                msg="a > 0 and a < b",
+            )
 
 
 class FixedCompanionMassRV(NormalRV):

thejoker/src/fast_likelihood.pyx

@@ -1,10 +1,12 @@
 # coding: utf-8
-# cython: boundscheck=False
-# cython: nonecheck=False
+# cython: boundscheck=True
+# cython: nonecheck=True
+# cython: wraparound=True
+# cython: initializedcheck=True
+# cython: overflowcheck=True
+# cython: linetrace=True
+# cython: profile=True
 # cython: cdivision=True
-# cython: wraparound=False
-# cython: initializedcheck=False
-# cython: profile=False
 # cython: language_level=3
 
 # Third-party
@@ -16,6 +18,7 @@ import cython
 cimport cython
 cimport scipy.linalg.cython_lapack as lapack
 import thejoker.units as xu
+from thejoker.utils import _pytensor_get_mean_std
 
 # from libc.stdio cimport printf
 from libc.math cimport pow, log, fabs, pi
@@ -210,10 +213,7 @@ cdef class CJokerHelper:
             _unit = getattr(dist, xu.UNIT_ATTR_NAME)
             to_unit = self.internal_units[name]
 
-            # mean is par 0, stddev par 1
-            pars = dist.owner.inputs[3:]
-            mu = (pars[0].eval() * _unit).to_value(to_unit)
-            std = (pars[1].eval() * _unit).to_value(to_unit)
+            mu, std = _pytensor_get_mean_std(dist, _unit, to_unit)
 
             # K, v0 = 2 - start at index 2
             self.mu[2+i] = mu
@@ -232,13 +232,7 @@ cdef class CJokerHelper:
             to_unit = self.internal_units[name]
 
             dist = prior.model[name]
-            # first three are (rng, size, dtype) as per
-            # https://github.com/pymc-devs/pymc/blob/main/pymc/printing.py#L43
-            pars = dist.owner.inputs[3:]
-
-            # mean is par 0
-            mu = (pars[0].eval() * _unit).to_value(to_unit)
-            std = (pars[1].eval() * _unit).to_value(to_unit)
+            mu, std = _pytensor_get_mean_std(dist, _unit, to_unit)
 
             if name == 'K' and self.fixed_K_prior == 0:
                 # TODO: here's the major hack

thejoker/src/tests/py_likelihood.py

@@ -8,7 +8,8 @@
 from astroML.utils import log_multivariate_gaussian
 from twobody.wrap import cy_rv_from_elements
 
-from ...samples import JokerSamples
+from thejoker.samples import JokerSamples
+from thejoker.utils import _pytensor_get_mean_std
 
 __all__ = ["get_ivar", "likelihood_worker", "marginal_ln_likelihood"]
 
@@ -141,20 +142,20 @@ def get_M_Lambda_ivar(samples, prior, data):
     n_linear = len(prior._linear_equiv_units)
 
     Lambda = np.zeros(n_linear)
-    for i, k in enumerate(prior._linear_equiv_units.keys()):
-        if k == "K":
+    for i, name in enumerate(prior._linear_equiv_units.keys()):
+        if name == "K":
             continue  # set below
-        pars = prior.pars[k].owner.inputs[3:]
-        Lambda[i] = pars[1].eval() ** 2
+        _, std = _pytensor_get_mean_std(prior.pars[name], u.one, u.one)
+        Lambda[i] = std**2
 
     K_dist = prior.pars["K"]
-    K_pars = K_dist.owner.inputs[3:]
     if K_dist.owner.op._print_name[0] == "FixedCompanionMass":
         sigma_K0 = K_dist._sigma_K0.to_value(v_unit)
         P0 = K_dist._P0.to_value(samples["P"].unit)
         max_K2 = K_dist._max_K.to_value(v_unit) ** 2
     else:
-        Lambda[0] = K_pars[1].eval() ** 2
+        _, std = _pytensor_get_mean_std(K_dist, u.one, u.one)
+        Lambda[0] = std**2
 
     for n in range(n_samples):
         M = design_matrix(packed_samples[n], data, prior)