Fix seed for random test data.

Also switch to using default_rng instead of RandomState.
pydata · dcherian · Dec 3, 2024 · Dec 2, 2024 · Dec 2, 2024 · Dec 2, 2024
commit d474b67c5c5346fb86ba18cea5b98665e1cfe1e0
diff --git a/asv_bench/benchmarks/__init__.py b/asv_bench/benchmarks/__init__.py
@@ -30,13 +30,13 @@ def requires_sparse():
 
 
 def randn(shape, frac_nan=None, chunks=None, seed=0):
-    rng = np.random.RandomState(seed)
+    rng = np.random.default_rng(seed)
     if chunks is None:
         x = rng.standard_normal(shape)
     else:
         import dask.array as da
 
-        rng = da.random.RandomState(seed)
+        rng = da.random.default_rng(seed)
         x = rng.standard_normal(shape, chunks=chunks)
 
     if frac_nan is not None:
@@ -47,7 +47,7 @@ def randn(shape, frac_nan=None, chunks=None, seed=0):
 
 
 def randint(low, high=None, size=None, frac_minus=None, seed=0):
-    rng = np.random.RandomState(seed)
+    rng = np.random.default_rng(seed)
     x = rng.randint(low, high, size)
     if frac_minus is not None:
         inds = rng.choice(range(x.size), int(x.size * frac_minus))

diff --git a/asv_bench/benchmarks/reindexing.py b/asv_bench/benchmarks/reindexing.py
@@ -11,7 +11,7 @@
 
 class Reindex:
     def setup(self):
-        data = np.random.RandomState(0).randn(ntime, nx, ny)
+        data = np.random.default_rng(0).randn(ntime, nx, ny)
         self.ds = xr.Dataset(
             {"temperature": (("time", "x", "y"), data)},
             coords={"time": np.arange(ntime), "x": np.arange(nx), "y": np.arange(ny)},

diff --git a/asv_bench/benchmarks/unstacking.py b/asv_bench/benchmarks/unstacking.py
@@ -8,7 +8,7 @@
 
 class Unstacking:
     def setup(self):
-        data = np.random.RandomState(0).randn(250, 500)
+        data = np.random.default_rng(0).randn(250, 500)
         self.da_full = xr.DataArray(data, dims=list("ab")).stack(flat_dim=[...])
         self.da_missing = self.da_full[:-1]
         self.df_missing = self.da_missing.to_pandas()

diff --git a/doc/user-guide/computation.rst b/doc/user-guide/computation.rst
@@ -30,7 +30,7 @@ numpy) over all array values:
 .. ipython:: python
 
     arr = xr.DataArray(
-        np.random.RandomState(0).randn(2, 3), [("x", ["a", "b"]), ("y", [10, 20, 30])]
+        np.random.default_rng(0).randn(2, 3), [("x", ["a", "b"]), ("y", [10, 20, 30])]
     )
     arr - 3
     abs(arr)

diff --git a/doc/user-guide/dask.rst b/doc/user-guide/dask.rst
@@ -292,7 +292,7 @@ work as a streaming operation, when run on arrays loaded from disk:
 .. ipython::
     :verbatim:
 
-    In [56]: rs = np.random.RandomState(0)
+    In [56]: rs = np.random.default_rng(0)
 
     In [57]: array1 = xr.DataArray(rs.randn(1000, 100000), dims=["place", "time"])  # 800MB
 

diff --git a/doc/user-guide/pandas.rst b/doc/user-guide/pandas.rst
@@ -202,7 +202,7 @@ Let's take a look:
 
 .. ipython:: python
 
-    data = np.random.RandomState(0).rand(2, 3, 4)
+    data = np.random.default_rng(0).rand(2, 3, 4)
     items = list("ab")
     major_axis = list("mno")
     minor_axis = pd.date_range(start="2000", periods=4, name="date")

diff --git a/xarray/tests/__init__.py b/xarray/tests/__init__.py
@@ -298,12 +298,12 @@ def assert_allclose(a, b, check_default_indexes=True, **kwargs):
 
 
 def create_test_data(
-    seed: int | None = None,
+    seed: int = 12345,
     add_attrs: bool = True,
     dim_sizes: tuple[int, int, int] = _DEFAULT_TEST_DIM_SIZES,
     use_extension_array: bool = False,
 ) -> Dataset:
-    rs = np.random.RandomState(seed)
+    rs = np.random.default_rng(seed)
     _vars = {
         "var1": ["dim1", "dim2"],
         "var2": ["dim1", "dim2"],

diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py
@@ -2082,7 +2082,7 @@
    @requires_netCDF4_1_6_2_or_above
    @pytest.mark.xfail(ON_WINDOWS, reason="new compression not yet implemented")
    def test_compression_encoding(self, compression: str | None) -> None:
        data = create_test_data(dim_sizes=(20, 80, 10))
        encoding_params: dict[str, Any] = dict(compression=compression, blosc_shuffle=1)
        data["var2"].encoding.update(encoding_params)
        data["var2"].encoding.update(
@@ -2634,7 +2634,7 @@
            assert_identical(original, actual)

        # check append mode for append write
        ds, ds_to_append, _ = create_append_test_data()
        with self.create_zarr_target() as store_target:
            ds.to_zarr(store_target, mode="w", group=group, **self.version_kwargs)
            ds_to_append.to_zarr(
@@ -2739,7 +2739,7 @@
                modified.to_zarr(store, mode="r+", **self.version_kwargs)

    def test_append_with_invalid_dim_raises(self) -> None:
        ds, ds_to_append, _ = create_append_test_data()
        with self.create_zarr_target() as store_target:
            ds.to_zarr(store_target, mode="w", **self.version_kwargs)
            with pytest.raises(
@@ -2760,14 +2760,14 @@
                )

    def test_append_with_append_dim_not_set_raises(self) -> None:
        ds, ds_to_append, _ = create_append_test_data()
        with self.create_zarr_target() as store_target:
            ds.to_zarr(store_target, mode="w", **self.version_kwargs)
            with pytest.raises(ValueError, match="different dimension sizes"):
                ds_to_append.to_zarr(store_target, mode="a", **self.version_kwargs)

    def test_append_with_mode_not_a_raises(self) -> None:
        ds, ds_to_append, _ = create_append_test_data()
        with self.create_zarr_target() as store_target:
            ds.to_zarr(store_target, mode="w", **self.version_kwargs)
            with pytest.raises(ValueError, match="cannot set append_dim unless"):
@@ -2776,7 +2776,7 @@
                )

    def test_append_with_existing_encoding_raises(self) -> None:
        ds, ds_to_append, _ = create_append_test_data()
        with self.create_zarr_target() as store_target:
            ds.to_zarr(store_target, mode="w", **self.version_kwargs)
            with pytest.raises(ValueError, match="but encoding was provided"):
@@ -2811,7 +2811,7 @@
            xr.testing.assert_identical(expected, actual)

    def test_check_encoding_is_consistent_after_append(self) -> None:
        ds, ds_to_append, _ = create_append_test_data()

        # check encoding consistency
        with self.create_zarr_target() as store_target:
@@ -2848,7 +2848,7 @@
            )

    def test_append_with_new_variable(self) -> None:
        ds, ds_to_append, ds_with_new_var = create_append_test_data()

        # check append mode for new variable
        with self.create_zarr_target() as store_target:
@@ -2864,7 +2864,7 @@
            )

    def test_append_with_append_dim_no_overwrite(self) -> None:
        ds, ds_to_append, _ = create_append_test_data()
        with self.create_zarr_target() as store_target:
            ds.to_zarr(store_target, mode="w", **self.version_kwargs)
            original = xr.concat([ds, ds_to_append], dim="time")
@@ -2914,7 +2914,7 @@
    def test_to_zarr_append_compute_false_roundtrip(self) -> None:
        from dask.delayed import Delayed

        ds, ds_to_append, _ = create_append_test_data()
        ds, ds_to_append = ds.chunk(), ds_to_append.chunk()

        with pytest.warns(SerializationWarning):
@@ -6478,7 +6478,7 @@
         arr.isel(a=slice(5, -1)).chunk(a=5).to_zarr(store, region="auto", mode="r+")
 
     # Test if the code is detecting the last chunk correctly
-    data = np.random.RandomState(0).randn(2920, 25, 53)
+    data = np.random.default_rng(0).randn(2920, 25, 53)
     ds = xr.Dataset({"temperature": (("time", "lat", "lon"), data)})
     chunks = {"time": 1000, "lat": 25, "lon": 53}
     ds.chunk(chunks).to_zarr(store, compute=False, mode="w")

diff --git a/xarray/tests/test_computation.py b/xarray/tests/test_computation.py
@@ -1293,7 +1293,7 @@ def covariance(x, y):
             (x - x.mean(axis=-1, keepdims=True)) * (y - y.mean(axis=-1, keepdims=True))
         ).mean(axis=-1)
 
-    rs = np.random.RandomState(42)
+    rs = np.random.default_rng(42)
     array1 = da.from_array(rs.randn(4, 4), chunks=(2, 4))
     array2 = da.from_array(rs.randn(4, 4), chunks=(2, 4))
     data_array_1 = xr.DataArray(array1, dims=("x", "z"))

diff --git a/xarray/tests/test_dask.py b/xarray/tests/test_dask.py
@@ -38,7 +38,7 @@
 
 
 def test_raise_if_dask_computes():
-    data = da.from_array(np.random.RandomState(0).randn(4, 6), chunks=(2, 2))
+    data = da.from_array(np.random.default_rng(0).randn(4, 6), chunks=(2, 2))
     with pytest.raises(RuntimeError, match=r"Too many computes"):
         with raise_if_dask_computes():
             data.compute()
@@ -77,7 +77,7 @@
 
     @pytest.fixture(autouse=True)
     def setUp(self):
-        self.values = np.random.RandomState(0).randn(4, 6)
+        self.values = np.random.default_rng(0).randn(4, 6)
         self.data = da.from_array(self.values, chunks=(2, 2))
 
         self.eager_var = Variable(("x", "y"), self.values)

diff --git a/xarray/tests/test_dataarray.py b/xarray/tests/test_dataarray.py
@@ -2284,7 +2284,7 @@
         assert isinstance(converted_subok.data, NdArraySubclass)
 
     def test_is_null(self) -> None:
-        x = np.random.RandomState(42).randn(5, 6)
+        x = np.random.default_rng(42).randn(5, 6)
         x[x < 0] = np.nan
         original = DataArray(x, [-np.arange(5), np.arange(6)], ["x", "y"])
         expected = DataArray(pd.isnull(x), [-np.arange(5), np.arange(6)], ["x", "y"])
@@ -3528,7 +3528,7 @@
 
         idx = pd.MultiIndex.from_product([np.arange(3), np.arange(5)], names=["a", "b"])
         series: pd.Series = pd.Series(
-            np.random.RandomState(0).random(len(idx)), index=idx
+            np.random.default_rng(0).random(len(idx)), index=idx
         ).sample(n=5, random_state=3)
 
         dense = DataArray.from_series(series, sparse=False)
@@ -3703,8 +3703,8 @@
         assert expected_attrs == actual["attrs"]
 
     def test_to_masked_array(self) -> None:
-        rs = np.random.RandomState(44)
+        rs = np.random.default_rng(44)
         x = rs.random_sample(size=(10, 20))
         x_masked = np.ma.masked_where(x < 0.5, x)
         da = DataArray(x_masked)


diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py
@@ -99,7 +99,7 @@
 
 
 def create_append_test_data(seed=None) -> tuple[Dataset, Dataset, Dataset]:
-    rs = np.random.RandomState(seed)
+    rs = np.random.default_rng(seed)
 
     lat = [2, 1, 0]
     lon = [0, 1, 2]
@@ -5650,7 +5650,7 @@
        assert list(actual) == expected

    def test_reduce_non_numeric(self) -> None:
        data1 = create_test_data(seed=44, use_extension_array=True)
        data2 = create_test_data(seed=44)
        add_vars = {"var5": ["dim1", "dim2"], "var6": ["dim1"]}
        for v, dims in sorted(add_vars.items()):
@@ -7161,7 +7161,7 @@
 @pytest.mark.parametrize("dask", [True, False])
 @pytest.mark.parametrize("edge_order", [1, 2])
 def test_differentiate(dask, edge_order) -> None:
-    rs = np.random.RandomState(42)
+    rs = np.random.default_rng(42)
     coord = [0.2, 0.35, 0.4, 0.6, 0.7, 0.75, 0.76, 0.8]
 
     da = xr.DataArray(
@@ -7210,7 +7210,7 @@
 @pytest.mark.filterwarnings("ignore:Converting non-nanosecond")
 @pytest.mark.parametrize("dask", [True, False])
 def test_differentiate_datetime(dask) -> None:
-    rs = np.random.RandomState(42)
+    rs = np.random.default_rng(42)
     coord = np.array(
         [
             "2004-07-13",
@@ -7260,7 +7260,7 @@
 @requires_cftime
 @pytest.mark.parametrize("dask", [True, False])
 def test_differentiate_cftime(dask) -> None:
-    rs = np.random.RandomState(42)
+    rs = np.random.default_rng(42)
     coord = xr.cftime_range("2000", periods=8, freq="2ME")
 
     da = xr.DataArray(
@@ -7289,7 +7289,7 @@
 
 @pytest.mark.parametrize("dask", [True, False])
 def test_integrate(dask) -> None:
-    rs = np.random.RandomState(42)
+    rs = np.random.default_rng(42)
     coord = [0.2, 0.35, 0.4, 0.6, 0.7, 0.75, 0.76, 0.8]
 
     da = xr.DataArray(
@@ -7343,7 +7343,7 @@
 @requires_scipy
 @pytest.mark.parametrize("dask", [True, False])
 def test_cumulative_integrate(dask) -> None:
-    rs = np.random.RandomState(43)
+    rs = np.random.default_rng(43)
     coord = [0.2, 0.35, 0.4, 0.6, 0.7, 0.75, 0.76, 0.8]
 
     da = xr.DataArray(
@@ -7406,7 +7406,7 @@
 @pytest.mark.parametrize("dask", [True, False])
 @pytest.mark.parametrize("which_datetime", ["np", "cftime"])
 def test_trapezoid_datetime(dask, which_datetime) -> None:
-    rs = np.random.RandomState(42)
+    rs = np.random.default_rng(42)
     coord: ArrayLike
     if which_datetime == "np":
         coord = np.array(

diff --git a/xarray/tests/test_duck_array_ops.py b/xarray/tests/test_duck_array_ops.py
@@ -341,7 +341,7 @@
 
 def construct_dataarray(dim_num, dtype, contains_nan, dask):
     # dimnum <= 3
-    rng = np.random.RandomState(0)
+    rng = np.random.default_rng(0)
     shapes = [16, 8, 4][:dim_num]
     dims = ("x", "y", "z")[:dim_num]
 
@@ -649,7 +649,7 @@
            pytest.skip("numpy's argmin (not nanargmin) does not handle object-dtype")
        if skipna and np.dtype(dtype).kind in "iufc":
            pytest.skip("numpy's nanargmin raises ValueError for all nan axis")
    da = construct_dataarray(dim_num, dtype, contains_nan=contains_nan, dask=dask)

    with warnings.catch_warnings():
        warnings.filterwarnings("ignore", "All-NaN slice")

diff --git a/xarray/tests/test_formatting_html.py b/xarray/tests/test_formatting_html.py
@@ -11,7 +11,7 @@
 
 @pytest.fixture
 def dataarray() -> xr.DataArray:
-    return xr.DataArray(np.random.RandomState(0).randn(4, 6))
+    return xr.DataArray(np.random.default_rng(0).randn(4, 6))
 
 
 @pytest.fixture

diff --git a/xarray/tests/test_groupby.py b/xarray/tests/test_groupby.py
@@ -2879,7 +2879,7 @@ def test_multiple_groupers(use_flox: bool, shuffle: bool) -> None:
         )
 
     b = xr.DataArray(
-        np.random.RandomState(0).randn(2, 3, 4),
+        np.random.default_rng(0).randn(2, 3, 4),
         coords={"xy": (("x", "y"), [["a", "b", "c"], ["b", "c", "c"]], {"foo": "bar"})},
         dims=["x", "y", "z"],
     )

diff --git a/xarray/tests/test_missing.py b/xarray/tests/test_missing.py
@@ -62,7 +62,7 @@ def ds():
 
 
 def make_interpolate_example_data(shape, frac_nan, seed=12345, non_uniform=False):
-    rs = np.random.RandomState(seed)
+    rs = np.random.default_rng(seed)
     vals = rs.normal(size=shape)
     if frac_nan == 1:
         vals[:] = np.nan