REFACTOR: Add support for NumPy 2.0 (#362)

* REFACTOR: Adhere to NumPy 2.0 changes * BUILD: update dependencies to support NumPy 2 --------- Co-authored-by: j-ittner <[email protected]>
BCG-X-Official · Jul 4, 2024 · 3fd8072 · 3fd8072
1 parent 8837d09
commit 3fd8072
Show file tree

Hide file tree

Showing 11 changed files with 34 additions and 34 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -50,6 +50,6 @@ repos:
         language_version: python310
         pass_filenames: false
         additional_dependencies:
-          - numpy~=1.24
+          - numpy~=2.0
           - pytest
           - packaging
diff --git a/azure-pipelines.yml b/azure-pipelines.yml
@@ -89,7 +89,7 @@ stages:
           - script: |
               # package dependencies for mypy
               dependencies=(
-                numpy~=1.24
+                numpy~=2.0
                 packaging
                 pytest
               )

diff --git a/environment.yml b/environment.yml
@@ -5,7 +5,7 @@ dependencies:
   # run
   - joblib            ~= 1.2
   - matplotlib        ~= 3.6
-  - numpy             ~= 1.24
+  - numpy             ~= 2.0
   - pandas            ~= 2.0
   - python            ~= 3.10.14
   - scipy             ~= 1.10

diff --git a/pyproject.toml b/pyproject.toml
@@ -16,7 +16,7 @@ license = "Apache Software License v2.0"
 requires = [
     "joblib            ~=1.0",
     "matplotlib        ~=3.6",
-    "numpy             >=1.23,<2a", # cannot use ~= due to conda bug
+    "numpy             >=1.23,<3a", # cannot use ~= due to conda bug
     "pandas            >=1.5",
     "scipy             ~=1.9",
     "typing_inspect    ~=0.7",
@@ -80,7 +80,7 @@ typing_extensions = "~=4.0.0"
 # maximum requirements of gamma-pytools
 joblib = "~=1.3"
 matplotlib = "~=3.8"
-numpy = ">=1.26,<2a"  # cannot use ~= due to conda bug
+numpy = ">=2,<3a"  # cannot use ~= due to conda bug
 pandas = "~=2.0"
 python = ">=3.12,<4a"   # cannot use ~= due to conda bug
 scipy = "~=1.12"

diff --git a/src/pytools/data/_linkage.py b/src/pytools/data/_linkage.py
@@ -32,7 +32,7 @@
 # Type variables
 #
 
-LinkageMatrix = npt.NDArray[np.float_]
+LinkageMatrix = npt.NDArray[np.float64]
 
 
 #

diff --git a/src/pytools/data/_matrix.py b/src/pytools/data/_matrix.py
@@ -32,7 +32,7 @@
 #
 
 T_Matrix = TypeVar("T_Matrix", bound="Matrix[Any]")
-T_Number = TypeVar("T_Number", bound="np.number[npt.NBitBase]")
+T_Number = TypeVar("T_Number", bound="np.number[Any]")
 
 #
 # Ensure all symbols introduced below are included in __all__
@@ -59,7 +59,7 @@ class Matrix(HasExpressionRepr, Generic[T_Number]):
     names: tuple[npt.NDArray[Any] | None, npt.NDArray[Any] | None]
 
     #: the weights of the rows and columns
-    weights: tuple[npt.NDArray[np.float_] | None, npt.NDArray[np.float_] | None]
+    weights: tuple[npt.NDArray[np.float64] | None, npt.NDArray[np.float64] | None]
 
     #: the labels for the row and column axes
     name_labels: tuple[str | None, str | None]
@@ -155,8 +155,8 @@ def _arg_to_array(
         else:
 
             def _ensure_positive(
-                w: npt.NDArray[np.float_] | None, axis: int
-            ) -> npt.NDArray[np.float_] | None:
+                w: npt.NDArray[np.float64] | None, axis: int
+            ) -> npt.NDArray[np.float64] | None:
                 if w is not None and (w < 0).any():
                     raise ValueError(
                         f"arg weights[{axis}] should be all positive, "
@@ -352,7 +352,7 @@ def _message(error: str) -> str:
 
 
 def _top_items_mask(
-    weights: npt.NDArray[np.float_] | None,
+    weights: npt.NDArray[np.float64] | None,
     current_size: int,
     target_size: tuple[int | None, float | None],
 ) -> npt.NDArray[np.bool_]:
@@ -385,7 +385,7 @@ def _top_items_mask(
         # THe target weight is expressed as a ratio of total weight
         # (0 < target_ratio <= 1).
 
-        weights_sorted_cumsum: npt.NDArray[np.float_] = weights[
+        weights_sorted_cumsum: npt.NDArray[np.float64] = weights[
             ix_weights_descending_stable
         ].cumsum()
         mask[
@@ -401,12 +401,12 @@ def _top_items_mask(
 
 def _resize_rows(
     values: npt.NDArray[T_Number],
-    weights: npt.NDArray[np.float_] | None,
+    weights: npt.NDArray[np.float64] | None,
     names: npt.NDArray[Any] | None,
     current_size: int,
     target_size: tuple[int | None, float | None],
 ) -> tuple[
-    npt.NDArray[T_Number], npt.NDArray[np.float_] | None, npt.NDArray[Any] | None
+    npt.NDArray[T_Number], npt.NDArray[np.float64] | None, npt.NDArray[Any] | None
 ]:
     mask = _top_items_mask(
         weights=weights, current_size=current_size, target_size=target_size

diff --git a/src/pytools/viz/_matplot.py b/src/pytools/viz/_matplot.py
@@ -304,13 +304,13 @@ def color_for_value(self, z: int | float) -> RgbaColor:
         pass
 
     @overload
-    def color_for_value(self, z: npt.NDArray[np.float_]) -> npt.NDArray[np.float_]:
+    def color_for_value(self, z: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]:
         """[overload]"""
         pass
 
     def color_for_value(
-        self, z: int | float | npt.NDArray[np.float_]
-    ) -> RgbaColor | npt.NDArray[np.float_]:
+        self, z: int | float | npt.NDArray[np.float64]
+    ) -> RgbaColor | npt.NDArray[np.float64]:
         """
         Get the color(s) associated with the given value(s), based on the color map and
         normalization defined for this style.

diff --git a/src/pytools/viz/dendrogram/_style.py b/src/pytools/viz/dendrogram/_style.py
@@ -222,7 +222,7 @@ def draw_leaf_labels(
 
     def _get_ytick_locations(
         self, *, weights: Sequence[float]
-    ) -> npt.NDArray[np.float_]:
+    ) -> npt.NDArray[np.float64]:
         """
         Get the tick locations for the y axis.
 
@@ -231,7 +231,7 @@ def _get_ytick_locations(
         """
         weights_array = np.array(weights)
         # noinspection PyTypeChecker
-        ytick_locations: npt.NDArray[np.float_] = -(
+        ytick_locations: npt.NDArray[np.float64] = -(
             np.arange(len(weights)) * self.padding
             + weights_array.cumsum()
             - weights_array / 2

diff --git a/src/pytools/viz/matrix/_matrix.py b/src/pytools/viz/matrix/_matrix.py
@@ -155,25 +155,25 @@ def draw_matrix(
             npt.NDArray[Any] | None,
         ],
         weights: tuple[
-            npt.NDArray[np.float_] | None,
-            npt.NDArray[np.float_] | None,
+            npt.NDArray[np.float64] | None,
+            npt.NDArray[np.float64] | None,
         ],
     ) -> None:
         """[see superclass]"""
         ax: Axes = self.ax
         colors = self.colors
 
-        weights_rows: npt.NDArray[np.float_]
-        weights_columns: npt.NDArray[np.float_]
+        weights_rows: npt.NDArray[np.float64]
+        weights_columns: npt.NDArray[np.float64]
         # replace undefined weights with all ones
         weights_rows, weights_columns = tuple(
             np.ones(n) if w is None else w for w, n in zip(weights, data.shape)
         )
 
         # calculate the horizontal and vertical matrix cell bounds based on the
         # cumulative sums of the axis weights; default all weights to 1 if not defined
-        column_bounds: npt.NDArray[np.float_]
-        row_bounds: npt.NDArray[np.float_]
+        column_bounds: npt.NDArray[np.float64]
+        row_bounds: npt.NDArray[np.float64]
 
         row_bounds = -np.array([0, *weights_rows]).cumsum()
         column_bounds = np.array([0, *weights_columns]).cumsum()
@@ -200,7 +200,7 @@ def draw_matrix(
         # draw the matrix cells
         for c, (x0, x1) in enumerate(zip(column_bounds, column_bounds[1:])):
             for r, (y1, y0) in enumerate(zip(row_bounds, row_bounds[1:])):
-                color: npt.NDArray[np.float_] = cell_colors[r, c]
+                color: npt.NDArray[np.float64] = cell_colors[r, c]
                 ax.add_patch(
                     Rectangle(
                         (
@@ -224,7 +224,7 @@ def draw_matrix(
         y_tick_locations = (row_bounds[:-1] + row_bounds[1:]) / 2
 
         def _set_ticks(
-            tick_locations: npt.NDArray[np.float_],
+            tick_locations: npt.NDArray[np.float64],
             tick_labels: npt.NDArray[Any],
             axis: Axis,
             tick_params: dict[str, Any],
@@ -461,15 +461,15 @@ def draw_matrix(
             npt.NDArray[Any] | None,
         ],
         weights: tuple[
-            npt.NDArray[np.float_] | None,
-            npt.NDArray[np.float_] | None,
+            npt.NDArray[np.float64] | None,
+            npt.NDArray[np.float64] | None,
         ],
     ) -> None:
         """[see superclass]"""
 
         def _axis_marks(
             axis_names: npt.NDArray[Any] | None,
-            axis_weights: npt.NDArray[np.float_] | None,
+            axis_weights: npt.NDArray[np.float64] | None,
         ) -> Iterable[str] | None:
             axis_names_iter: Iterable[Any]
 

diff --git a/src/pytools/viz/matrix/base/_base.py b/src/pytools/viz/matrix/base/_base.py
@@ -67,8 +67,8 @@ def draw_matrix(
             npt.NDArray[Any] | None,
         ],
         weights: tuple[
-            npt.NDArray[np.float_] | None,
-            npt.NDArray[np.float_] | None,
+            npt.NDArray[np.float64] | None,
+            npt.NDArray[np.float64] | None,
         ],
     ) -> None:
         """

diff --git a/test/test/pytools/viz/dendrogram/test_dendrogram.py b/test/test/pytools/viz/dendrogram/test_dendrogram.py
@@ -25,7 +25,7 @@ def linkage_matrix() -> npt.NDArray[np.int_]:
 
 
 @pytest.fixture
-def linkage_tree(linkage_matrix: npt.NDArray[np.float_]) -> LinkageTree:
+def linkage_tree(linkage_matrix: npt.NDArray[np.float64]) -> LinkageTree:
     """Create a linkage tree for drawing tests."""
     return LinkageTree(
         scipy_linkage_matrix=linkage_matrix,
@@ -34,7 +34,7 @@ def linkage_tree(linkage_matrix: npt.NDArray[np.float_]) -> LinkageTree:
     )
 
 
-def test_dendrogram_drawer_text(linkage_matrix: npt.NDArray[np.float_]) -> None:
+def test_dendrogram_drawer_text(linkage_matrix: npt.NDArray[np.float64]) -> None:
     leaf_names = list("ABCDEFGH")
     leaf_weights = [(w + 1) / 36 for w in range(8)]