Repository: polsys/ennemi

Issue:
Add support for NumPy 2.0
- [ ] Modify `setup.py` to accept NumPy 2.0
- [ ] Add prerelease NumPy 2.0 feed to the CI test matrix
- [ ] Create followup item for testing with official NumPy 2.0 releases

Discussion:

Patch:
diff --git a/.github/workflows/continuous-integration.yml b/.github/workflows/continuous-integration.yml
index 6fb901e..ac9e4d8 100644
--- a/.github/workflows/continuous-integration.yml
+++ b/.github/workflows/continuous-integration.yml
@@ -14,32 +14,31 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: [ ubuntu-latest, macos-12, windows-latest ]
-        python-version: [ '3.9', '3.10', '3.11' ]
         # Save a bit of resources by running only some combinations:
-        # Windows 3.9; macOS 3.10; Ubuntu 3.11
-        # Python 3.8 is only used in integration tests
-        # Note that some steps are only run on Ubuntu/3.11 (must update when 3.12 comes)
-        exclude:
-          - os: windows-latest
-            python-version: '3.10'
+        # Windows 3.11; macOS 3.10; Ubuntu 3.12
+        # Additionally, macOS on ARM processor is tested
+        # Code coverage, SonarCloud and test package build are only done once
+        include:
+          - os: ubuntu-latest
+            python-version: '3.12'
+            analysis: true
           - os: windows-latest
             python-version: '3.11'
-          - os: macos-12
-            python-version: '3.9'
-          - os: macos-12
-            python-version: '3.11'
-          - os: ubuntu-latest
-            python-version: '3.9'
-          - os: ubuntu-latest
+            analysis: false
+          - os: macos-latest
             python-version: '3.10'
+            analysis: false
+          - os: macos-14
+            python-version: '3.12'
+            analysis: false
+
 
     steps:
       - name: Check out code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
 
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
 
@@ -48,7 +47,7 @@ jobs:
           pip install -e .[dev]
       
       - name: Run unit tests (with coverage)
-        if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.11'
+        if: matrix.analysis == true
         # We have to rename the path in coverage.xml since Sonar is run in Docker
         # and has a different view of the file system
         run: |
@@ -57,12 +56,12 @@ jobs:
           sed 's/\/home\/runner\/work\/ennemi\/ennemi/\/github\/workspace/g' coverage.xml > coverage-mod.xml
       
       - name: Run unit tests (no coverage)
-        if: matrix.os != 'ubuntu-latest' || matrix.python-version != '3.11'
+        if: matrix.analysis != true
         run: |
           python -m pytest tests/unit/ --junit-xml=test-results.xml
 
       - name: Try building the package
-        if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.11'
+        if: matrix.analysis == true
         run: |
           pip install setuptools wheel
           python setup.py sdist bdist_wheel
@@ -72,7 +71,7 @@ jobs:
           python -m mypy ennemi/ tests/unit tests/integration tests/pandas
       
       - name: Run Sonar code analysis
-        if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.11'
+        if: matrix.analysis == true
         uses: sonarsource/sonarcloud-github-action@master
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
diff --git a/.github/workflows/nightly-deps.yml b/.github/workflows/nightly-deps.yml
new file mode 100644
index 0000000..4744cc4
--- /dev/null
+++ b/.github/workflows/nightly-deps.yml
@@ -0,0 +1,39 @@
+# Run the unit tests with prerelease dependencies
+
+name: Nightly dependencies
+
+on:
+  pull_request:
+    branches: [ master ]
+    paths-ignore:
+      - 'docs/**'
+      - '**.md'
+
+jobs:
+  build:
+
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [ ubuntu-latest ]
+
+    steps:
+      - name: Check out code
+        uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.12'
+
+      - name: Install dependencies with custom versions
+        run: |
+          pip install -U --pre --only-binary :all: -i https://pypi.anaconda.org/scientific-python-nightly-wheels/simple --extra-index-url https://pypi.org/simple numpy scipy pandas
+          pip install pytest~=8.0
+          pip install --no-deps -e .
+
+      - name: Run unit and integration tests
+        run: |
+          python -m pytest .
+
diff --git a/.github/workflows/release-pypi.yml b/.github/workflows/release-pypi.yml
index a7d8b76..529eb58 100644
--- a/.github/workflows/release-pypi.yml
+++ b/.github/workflows/release-pypi.yml
@@ -10,10 +10,10 @@ jobs:
 
     steps:
       - name: Check out code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
 
       - name: Set up Python
-        uses: actions/setup-python@v4
+        uses: actions/setup-python@v5
         with:
           python-version: '3.11'
 
diff --git a/DESCRIPTION.md b/DESCRIPTION.md
index 5c1b6b3..76e36e3 100644
--- a/DESCRIPTION.md
+++ b/DESCRIPTION.md
@@ -16,8 +16,10 @@ with no theoretical background required.
 
 This package depends only on NumPy and SciPy;
 Pandas is suggested for more enjoyable data analysis.
-Python 3.8+ on the latest macOS, Ubuntu and Windows versions
+Recent versions of NumPy 1.x and 2.x are supported.
+Python 3.10+ on the latest macOS, Ubuntu and Windows versions
 is officially supported.
+Older `ennemi` versions have generally identical behavior if you need to run on older Python.
 
 For more information on theoretical background and usage, please see the
 [documentation](https://polsys.github.io/ennemi).
diff --git a/ennemi/_driver.py b/ennemi/_driver.py
index 605d8df..5b1119d 100644
--- a/ennemi/_driver.py
+++ b/ennemi/_driver.py
@@ -34,6 +34,12 @@ DISCRETE_NORMALIZATION_WARNING = \
     "If you really want to calculate correlation coefficients, you can suppress " +\
     "this warning by setting normalize=False and calling normalize_mi() on the results."
 
+PREPROCESS_CONSTANT_DATA_WARNING = \
+    "A variable not marked as discrete takes only a single value, " +\
+    "or values in a very small numerical range. " +\
+    "If this is intentional, you can suppress this warning by passing preprocess=False. " +\
+    "Note that this disables rescaling on all other variables as well."
+
 def normalize_mi(mi: Union[float, GenArrayLike]) -> GenArrayLike:
     """Normalize mutual information values to the unit interval.
 
@@ -61,10 +67,16 @@ def normalize_mi(mi: Union[float, GenArrayLike]) -> GenArrayLike:
     """
 
     # If the parameter is a pandas type, preserve the columns and indices
+    # In pandas 2.1 applymap was renamed to map,
+    # so use that if possible, and fall back to older function if that fails
+    # (since our support goes all the way to pandas 1.0).
     if "pandas" in sys.modules:
         import pandas
         if isinstance(mi, (pandas.DataFrame, pandas.Series)):
-            return mi.applymap(_normalize)
+            if "map" in pandas.DataFrame.__dict__:
+                return mi.map(_normalize)
+            else:
+                return mi.applymap(_normalize)
     
     return np.vectorize(_normalize, otypes=[float])(mi)
 
@@ -207,7 +219,7 @@ def _estimate_conditional_entropy(x: FloatArray, cond: FloatArray, k: int, multi
         return np.asarray(_call_entropy_func(xs, k, discrete) - marginal)
     else:
         nvar = x.shape[1]
-        joint = np.empty(nvar) # type: npt.NDArray[np.float64]
+        joint = np.empty(nvar)
         for i in range(nvar):
             xs = _mask_and_validate_entropy(np.column_stack((x[:,i], cond)), mask, drop_nan, discrete, k)
             joint[i] = _call_entropy_func(xs, k, discrete)
@@ -328,7 +340,7 @@ def estimate_mi(y: ArrayLike, x: ArrayLike,
     else:
         cond_arr = None
         ncond = 1
-    mask_arr = None # type: Optional[npt.NDArray[np.float64]]
+    mask_arr = None
     if mask is not None: mask_arr = np.asarray(mask)
 
     # Broadcast cond_lag to be (#lags, #cond vars) in shape
@@ -590,8 +602,8 @@ def pairwise_mi(data: ArrayLike,
 
     # Convert arrays to consistent type; _lagged_mi assumes cond to be 2D
     data_arr = np.asarray(data)
-    cond_arr = None # type: Optional[npt.NDArray[np.float64]]
-    mask_arr = None # type: Optional[npt.NDArray[np.float64]]
+    cond_arr = None
+    mask_arr = None
     if cond is not None: cond_arr = np.column_stack((np.asarray(cond),))
     if mask is not None: mask_arr = np.asarray(mask)
 
@@ -868,17 +880,29 @@ def _rescale_data(xs: FloatArray, ys: FloatArray, zs: Optional[FloatArray],
     rng = np.random.default_rng(2_718281828)
 
     if not discrete_x:
-        # This warns if the standard deviation is zero
-        xs = (xs - xs.mean()) / xs.std()
-        xs += rng.normal(0.0, 1e-10, xs.shape)
+        # Do not try to divide by zero
+        std = xs.std()
+        if np.abs(std) < 1e-20:
+            warn(PREPROCESS_CONSTANT_DATA_WARNING)
+        else:
+            xs = (xs - xs.mean()) / std
+            xs += rng.normal(0.0, 1e-10, xs.shape)
 
     if not discrete_y:
-        ys = (ys - ys.mean()) / ys.std()
-        ys += rng.normal(0.0, 1e-10, ys.shape)
+        std = ys.std()
+        if np.abs(std) < 1e-20:
+            warn(PREPROCESS_CONSTANT_DATA_WARNING)
+        else:
+            ys = (ys - ys.mean()) / std
+            ys += rng.normal(0.0, 1e-10, ys.shape)
     
     # If both X and Y are discrete, we assume the condition to be discrete too
     if zs is not None and not (discrete_x and discrete_y):
-        zs = (zs - zs.mean(axis=0)) / zs.std(axis=0)
-        zs += rng.normal(0.0, 1e-10, zs.shape) # type: ignore # mypy does not realize this is ndarray
+        std = zs.std(axis=0)
+        if np.any(np.abs(std) < 1e-20):
+            warn(PREPROCESS_CONSTANT_DATA_WARNING)
+        else:
+            zs = (zs - zs.mean(axis=0)) / std
+            zs += rng.normal(0.0, 1e-10, zs.shape) # type: ignore # mypy does not realize this is ndarray
 
     return xs, ys, zs
diff --git a/ennemi/_entropy_estimators.py b/ennemi/_entropy_estimators.py
index eb73176..79560af 100644
--- a/ennemi/_entropy_estimators.py
+++ b/ennemi/_entropy_estimators.py
@@ -303,6 +303,10 @@ def _estimate_conditional_discrete_mi(x: FloatArray, y: FloatArray, cond: FloatA
     # Determine probabilities of the conditioning variable
     cond_vals, cond_inverses, cond_counts = np.unique(cond,
         axis=0, return_inverse=True, return_counts=True)
+    
+    # Starting from NumPy 2.0, cond_inverses is multidimensional, so we need to flatten it
+    # This is a no-op for NumPy 1.x
+    cond_inverses = np.ravel(cond_inverses)
 
     # For each condition, compute the conditional probability (given by basic MI on subset of data)
     cond_probs = np.zeros(len(cond_vals))
diff --git a/setup.py b/setup.py
index b4bcaf7..14d5acd 100644
--- a/setup.py
+++ b/setup.py
@@ -10,7 +10,7 @@ with open(description_file, encoding="utf-8") as f:
 
 setup(
     name = "ennemi",
-    version = "1.3.0",
+    version = "1.4.0",
     description = "Non-linear correlation detection with mutual information",
     long_description = long_description,
     long_description_content_type = "text/markdown",
@@ -25,10 +25,9 @@ setup(
         "License :: OSI Approved :: MIT License",
         "Operating System :: OS Independent",
         "Programming Language :: Python :: 3",
-        "Programming Language :: Python :: 3.8",
-        "Programming Language :: Python :: 3.9",
         "Programming Language :: Python :: 3.10",
         "Programming Language :: Python :: 3.11",
+        "Programming Language :: Python :: 3.12",
         "Topic :: Scientific/Engineering",
         "Topic :: Scientific/Engineering :: Information Analysis",
         "Topic :: Scientific/Engineering :: Mathematics",
@@ -46,9 +45,9 @@ setup(
 
     packages = [ "ennemi" ],
     package_data = { "ennemi": ["py.typed"] },
-    python_requires = "~=3.8",
-    install_requires = [ "numpy~=1.21", "scipy~=1.7" ],
+    python_requires = "~=3.10",
+    install_requires = [ "numpy>=1.23", "scipy~=1.9" ],
     extras_require = {
-        "dev": [ "pandas>1.0", "pytest~=6.2", "mypy~=0.991" ]
+        "dev": [ "pandas>=1.5", "pytest~=8.0", "mypy~=1.9" ]
     }
 )
\ No newline at end of file


Test patch:
diff --git a/.github/workflows/integration-tests.yml b/.github/workflows/integration-tests.yml
index 463678c..37921e0 100644
--- a/.github/workflows/integration-tests.yml
+++ b/.github/workflows/integration-tests.yml
@@ -17,32 +17,33 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        os: [ ubuntu-20.04, windows-2019, macos-11 ]
+        os: [ windows-2019, macos-11 ]
 
     steps:
       - name: Check out code
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
 
-      - name: Set up Python 3.8
-        uses: actions/setup-python@v4
+      - name: Set up Python 3.10
+        uses: actions/setup-python@v5
         with:
-          python-version: '3.8'
+          python-version: '3.10'
 
       - name: Install dependencies with custom versions
         # Note: This must be kept in sync with the minimum versions in setup.py
         run: |
-          pip install numpy==1.21.0 scipy==1.7 pytest~=6.2
+          pip install numpy==1.23.0 scipy==1.9.0 pytest~=8.0
           pip install --no-deps -e .
 
       - name: Run integration tests (without pandas)
         run: |
           python -m pytest tests/integration/
 
-      - name: Install pandas 1.0
+      - name: Install pandas 1.x
+        # This too needs to be updated at some cadence (mainly due to NumPy support of pandas)
         run: |
-          pip install pandas==1.0
+          pip install pandas==1.5.0
 
-      - name: Run integration tests (with pandas 1.0)
+      - name: Run integration tests (with pandas 1.x)
         run: |
           python -m pytest tests/pandas/
 
diff --git a/tests/unit/test_driver.py b/tests/unit/test_driver.py
index ce1e652..6a1f887 100644
--- a/tests/unit/test_driver.py
+++ b/tests/unit/test_driver.py
@@ -937,15 +937,16 @@ class TestEstimateMi(unittest.TestCase):
             estimate_mi(cont, disc2d, discrete_x=True, normalize=True)
 
         # No warning should be emitted if data is continuous...
+        # (Need to disable rescaling since it has its own discreteness check)
         with catch_warnings(record=True) as w:
-            estimate_mi(cont, cont, discrete_x=False, normalize=True)
+            estimate_mi(cont, cont, discrete_x=False, normalize=True, preprocess=False)
 
             matching_warns = [msg for msg in w if "discrete" in str(msg.message)]
             self.assertEqual(len(matching_warns), 0, "should not warn")
         
         # ...or normalization disabled
         with catch_warnings(record=True) as w:
-            estimate_mi(cont, disc2d, discrete_x=True, normalize=False)
+            estimate_mi(cont, disc2d, discrete_x=True, normalize=False, preprocess=False)
 
             matching_warns = [msg for msg in w if "discrete" in str(msg.message)]
             self.assertEqual(len(matching_warns), 0, "should not warn")
@@ -1181,6 +1182,18 @@ class TestEstimateMi(unittest.TestCase):
         self.assertNotAlmostEqual(mi_unscaled.item(), 0.6, delta=0.1)
         self.assertAlmostEqual(mi_scaled.item(), 0.6, delta=0.03)
 
+    def test_preprocess_zero_variance_warns(self) -> None:
+        zeros = np.zeros(10)
+        ones = np.ones(10)
+
+        with self.assertWarnsRegex(UserWarning, "only a single value"):
+            estimate_mi(zeros, ones, preprocess=True)
+        with self.assertWarnsRegex(UserWarning, "only a single value"):
+            estimate_mi(ones, zeros, preprocess=True)
+        with self.assertWarnsRegex(UserWarning, "only a single value"):
+            estimate_mi(ones, ones, cond=zeros, preprocess=True)
+        with self.assertWarnsRegex(UserWarning, "only a single value"):
+            estimate_mi(ones, ones, cond=np.asarray([ones, zeros]).T, preprocess=True)
 
     def test_normalization(self) -> None:
         rng = np.random.default_rng(17)
@@ -1320,21 +1333,22 @@ class TestPairwiseMi(unittest.TestCase):
         data = np.zeros((10, 2))
 
         # Warnings emitted both in scalar and array case
+        # Note: need to disable rescaling since it has its own discreteness check
         with self.assertWarnsRegex(UserWarning, "discrete"):
-            pairwise_mi(data, discrete=True, normalize=True)
+            pairwise_mi(data, discrete=True, normalize=True, preprocess=False)
         with self.assertWarnsRegex(UserWarning, "discrete"):
-            pairwise_mi(data, discrete=[False, True], normalize=True)
+            pairwise_mi(data, discrete=[False, True], normalize=True, preprocess=False)
 
         # No warning should be emitted if data is continuous...
         with catch_warnings(record=True) as w:
-            pairwise_mi(data, discrete=False, normalize=True)
+            pairwise_mi(data, discrete=False, normalize=True, preprocess=False)
 
             matching_warns = [msg for msg in w if "discrete" in str(msg.message)]
             self.assertEqual(len(matching_warns), 0, "should not warn")
         
         # ...or normalization is disabled
         with catch_warnings(record=True) as w:
-            pairwise_mi(data, discrete=True, normalize=False)
+            pairwise_mi(data, discrete=True, normalize=False, preprocess=False)
 
             matching_warns = [msg for msg in w if "discrete" in str(msg.message)]
             self.assertEqual(len(matching_warns), 0, "should not warn")