make spatial, spectral and spectro-glacial sampling reproducible if seed set in Formulae (#1026)

slayoo · abulenok · web-flow · commit eb4c52b51590 · 2023-03-29T16:21:14.000Z
Co-authored-by: Oleksii Bulenok &lt;63914185+abulenok@users.noreply.github.com&gt;
diff --git a/PySDM/backends/impl_numba/storage_impl.py b/PySDM/backends/impl_numba/storage_impl.py
@@ -70,9 +70,3 @@ def power(output, exponent):
 @numba.njit(**{**conf.JIT_FLAGS, **{"parallel": False}})
 def subtract(output, subtrahend):
     output[:] -= subtrahend[:]
-
-
-# @numba.njit(void(f8[:]), **conf.JIT_FLAGS)
-def urand(output):
-    # TODO #986: seed
-    output.data[:] = np.random.uniform(0, 1, output.shape)
diff --git a/PySDM/environments/box.py b/PySDM/environments/box.py
@@ -30,6 +30,6 @@ def register(self, builder):
     def init_attributes(self, *, spectral_discretisation):
         attributes = {}
         attributes["volume"], attributes["n"] = spectral_discretisation.sample(
-            self.particulator.n_sd
+            backend=self.particulator.backend, n_sd=self.particulator.n_sd
         )
         return attributes
diff --git a/PySDM/environments/kinematic_1d.py b/PySDM/environments/kinematic_1d.py
@@ -40,15 +40,19 @@ def init_attributes(
         attributes = {}
         with np.errstate(all="raise"):
             positions = spatial_discretisation.sample(
-                self.mesh.grid, self.particulator.n_sd
+                backend=self.particulator.backend,
+                grid=self.mesh.grid,
+                n_sd=self.particulator.n_sd,
             )
             (
                 attributes["cell id"],
                 attributes["cell origin"],
                 attributes["position in cell"],
             ) = self.mesh.cellular_attributes(positions)
 
-            r_dry, n_per_kg = spectral_discretisation.sample(self.particulator.n_sd)
+            r_dry, n_per_kg = spectral_discretisation.sample(
+                backend=self.particulator.backend, n_sd=self.particulator.n_sd
+            )
             attributes["dry volume"] = self.formulae.trivia.volume(radius=r_dry)
             attributes["kappa times dry volume"] = attributes["dry volume"] * kappa
             r_wet = equilibrate_wet_radii(
diff --git a/PySDM/environments/kinematic_2d.py b/PySDM/environments/kinematic_2d.py
@@ -11,7 +11,7 @@
     default_rtol,
     equilibrate_wet_radii,
 )
-from PySDM.initialisation.sampling import spectral_sampling
+from PySDM.initialisation.sampling.spectral_sampling import ConstantMultiplicity
 
 from ..impl import arakawa_c
 
@@ -42,30 +42,26 @@ def init_attributes(
         kappa,
         dry_radius_spectrum,
         rtol=default_rtol,
-        n_sd=None
+        n_sd=None,
+        spectral_sampling=ConstantMultiplicity
     ):
         super().sync()
         self.notify()
         n_sd = n_sd or self.particulator.n_sd
         attributes = {}
         with np.errstate(all="raise"):
-            positions = spatial_discretisation.sample(self.mesh.grid, n_sd)
-            (
-                attributes["cell id"],
-                attributes["cell origin"],
-                attributes["position in cell"],
-            ) = self.mesh.cellular_attributes(positions)
-        with np.errstate(all="raise"):
-            positions = spatial_discretisation.sample(self.mesh.grid, n_sd)
+            positions = spatial_discretisation.sample(
+                backend=self.particulator.backend, grid=self.mesh.grid, n_sd=n_sd
+            )
             (
                 attributes["cell id"],
                 attributes["cell origin"],
                 attributes["position in cell"],
             ) = self.mesh.cellular_attributes(positions)
 
-            r_dry, n_per_kg = spectral_sampling.ConstantMultiplicity(
-                spectrum=dry_radius_spectrum
-            ).sample(n_sd)
+            r_dry, n_per_kg = spectral_sampling(spectrum=dry_radius_spectrum).sample(
+                n_sd=n_sd, backend=self.particulator.backend
+            )
 
             attributes["dry volume"] = self.formulae.trivia.volume(radius=r_dry)
             attributes["kappa times dry volume"] = kappa * attributes["dry volume"]
diff --git a/PySDM/initialisation/sampling/spatial_sampling.py b/PySDM/initialisation/sampling/spatial_sampling.py
@@ -1,17 +1,20 @@
 """
 spatial sampling logic (i.e., physical x-y-z coordinates)
 """
-import numpy as np
-
 # TODO #305 QUASIRANDOM & GRID
 #  http://www.ii.uj.edu.pl/~arabas/workshop_2019/files/talk_Shima.pdf
 
 
 class Pseudorandom:  # pylint: disable=too-few-public-methods
     @staticmethod
-    def sample(grid, n_sd):
+    def sample(*, backend, grid, n_sd):
         dimension = len(grid)
-        positions = np.random.rand(dimension, n_sd)
+        n_elements = dimension * n_sd
+
+        storage = backend.Storage.empty(n_elements, dtype=float)
+        backend.Random(seed=backend.formulae.seed, size=n_elements)(storage)
+        positions = storage.to_ndarray().reshape(dimension, n_sd)
+
         for dim in range(dimension):
             positions[dim, :] *= grid[dim]
         return positions
diff --git a/PySDM/initialisation/sampling/spectral_sampling.py b/PySDM/initialisation/sampling/spectral_sampling.py
@@ -7,20 +7,12 @@
 import numpy as np
 from scipy import optimize
 
-from PySDM.physics import constants as const
-
 default_cdf_range = (0.00001, 0.99999)
 
 
 class SpectralSampling:  # pylint: disable=too-few-public-methods
-    def __init__(
-        self,
-        spectrum,
-        size_range: [None, Tuple[float, float]] = None,
-        error_threshold: Optional[float] = None,
-    ):
+    def __init__(self, spectrum, size_range: Optional[Tuple[float, float]] = None):
         self.spectrum = spectrum
-        self.error_threshold = error_threshold or 0.01
 
         if size_range is None:
             if hasattr(spectrum, "percentiles"):
@@ -40,6 +32,20 @@ def __init__(
             assert size_range[1] > size_range[0]
             self.size_range = size_range
 
+
+class DeterministicSpectralSampling(
+    SpectralSampling
+):  # pylint: disable=too-few-public-methods
+    # TODO #1031 - error_threshold will be also used in non-deterministic sampling
+    def __init__(
+        self,
+        spectrum,
+        size_range: Optional[Tuple[float, float]] = None,
+        error_threshold: Optional[float] = None,
+    ):
+        super().__init__(spectrum, size_range)
+        self.error_threshold = error_threshold or 0.01
+
     def _sample(self, grid, spectrum):
         x = grid[1:-1:2]
         cdf = spectrum.cumulative(grid[0::2])
@@ -54,13 +60,15 @@ def _sample(self, grid, spectrum):
         return x, y_float
 
 
-class Linear(SpectralSampling):  # pylint: disable=too-few-public-methods
-    def sample(self, n_sd):
+class Linear(DeterministicSpectralSampling):  # pylint: disable=too-few-public-methods
+    def sample(self, n_sd, *, backend=None):  # pylint: disable=unused-argument
         grid = np.linspace(*self.size_range, num=2 * n_sd + 1)
         return self._sample(grid, self.spectrum)
 
 
-class Logarithmic(SpectralSampling):  # pylint: disable=too-few-public-methods
+class Logarithmic(
+    DeterministicSpectralSampling
+):  # pylint: disable=too-few-public-methods
     def __init__(
         self,
         spectrum,
@@ -71,12 +79,14 @@ def __init__(
         self.start = np.log10(self.size_range[0])
         self.stop = np.log10(self.size_range[1])
 
-    def sample(self, n_sd):
+    def sample(self, n_sd, *, backend=None):  # pylint: disable=unused-argument
         grid = np.logspace(self.start, self.stop, num=2 * n_sd + 1)
         return self._sample(grid, self.spectrum)
 
 
-class ConstantMultiplicity(SpectralSampling):  # pylint: disable=too-few-public-methods
+class ConstantMultiplicity(
+    DeterministicSpectralSampling
+):  # pylint: disable=too-few-public-methods
     def __init__(self, spectrum, size_range=None):
         super().__init__(spectrum, size_range)
 
@@ -86,7 +96,7 @@ def __init__(self, spectrum, size_range=None):
         )
         assert 0 < self.cdf_range[0] < self.cdf_range[1]
 
-    def sample(self, n_sd):
+    def sample(self, n_sd, *, backend=None):  # pylint: disable=unused-argument
         cdf_arg = np.linspace(self.cdf_range[0], self.cdf_range[1], num=2 * n_sd + 1)
         cdf_arg /= self.spectrum.norm_factor
         percentiles = self.spectrum.percentiles(cdf_arg)
@@ -97,11 +107,13 @@ def sample(self, n_sd):
 
 
 class UniformRandom(SpectralSampling):  # pylint: disable=too-few-public-methods
-    def __init__(self, spectrum, size_range=None, seed=const.default_random_seed):
-        super().__init__(spectrum, size_range)
-        self.rng = np.random.default_rng(seed)
+    def sample(self, n_sd, *, backend):
+        n_elements = n_sd
+        storage = backend.Storage.empty(n_elements, dtype=float)
+        backend.Random(seed=backend.formulae.seed, size=n_elements)(storage)
+        u01 = storage.to_ndarray()
 
-    def sample(self, n_sd):
-        pdf_arg = self.rng.uniform(*self.size_range, n_sd)
+        pdf_arg = self.size_range[0] + u01 * (self.size_range[1] - self.size_range[0])
         dr = abs(self.size_range[1] - self.size_range[0]) / n_sd
+        # TODO #1031 - should also handle error_threshold check
         return pdf_arg, dr * self.spectrum.size_distribution(pdf_arg)
diff --git a/PySDM/initialisation/sampling/spectro_glacial_sampling.py b/PySDM/initialisation/sampling/spectro_glacial_sampling.py
@@ -5,21 +5,14 @@
 import numpy as np
 
 from PySDM.initialisation.sampling.spectral_sampling import default_cdf_range
-from PySDM.physics import constants as const
 
 DIM_TEMP = 0
 DIM_SURF = 1
 N_DIMS = 2
 
 
 class SpectroGlacialSampling:  # pylint: disable=too-few-public-methods
-    def __init__(
-        self,
-        *,
-        freezing_temperature_spectrum,
-        insoluble_surface_spectrum,
-        seed=const.default_random_seed
-    ):
+    def __init__(self, *, freezing_temperature_spectrum, insoluble_surface_spectrum):
         self.insoluble_surface_spectrum = insoluble_surface_spectrum
         self.freezing_temperature_spectrum = freezing_temperature_spectrum
 
@@ -29,15 +22,20 @@ def __init__(
         self.temperature_range = freezing_temperature_spectrum.invcdf(
             np.asarray(default_cdf_range), insoluble_surface_spectrum.median
         )
-        self.seed = seed
 
-    def sample(self, n_sd):
+    def sample(self, *, backend, n_sd):
         simulated = np.empty((n_sd, N_DIMS))
+
+        n_elements = n_sd * N_DIMS
+        storage = backend.Storage.empty(n_elements, dtype=float)
+        backend.Random(seed=backend.formulae.seed, size=n_elements)(storage)
+        random_numbers = storage.to_ndarray().reshape(n_sd, N_DIMS)
+
         simulated[:, DIM_SURF] = self.insoluble_surface_spectrum.percentiles(
-            np.random.random(n_sd)
+            random_numbers[:, 0]
         )
         simulated[:, DIM_TEMP] = self.freezing_temperature_spectrum.invcdf(
-            np.random.random(n_sd), simulated[:, DIM_SURF]
+            random_numbers[:, 1], simulated[:, DIM_SURF]
         )
 
         return (
diff --git a/test-time-requirements.txt b/test-time-requirements.txt
@@ -6,7 +6,7 @@ ghapi
 pytest
 
 # note: if cloning both PySDM and PySDM examples, consider "pip install -e"
-PySDM-examples @ git+https://github.com/open-atmos/PySDM-examples@9142101#egg=PySDM-examples
+PySDM-examples @ git+https://github.com/open-atmos/PySDM-examples@4900000#egg=PySDM-examples
 PyMPDATA @ git+https://github.com/open-atmos/PyMPDATA@63f1c11#egg=PyMPDATA
 
 PyPartMC==0.0.26
diff --git a/tests/unit_tests/dynamics/collisions/test_sdm_multi_cell.py b/tests/unit_tests/dynamics/collisions/test_sdm_multi_cell.py
@@ -33,7 +33,9 @@ def test_coalescence_call(n_sd, backend_class, adaptive):
         particulator, sut = get_dummy_particulator_and_coalescence(
             backend_class, len(n), environment=env
         )
-        cell_id, _, _ = env.mesh.cellular_attributes(Pseudorandom.sample(grid, len(n)))
+        cell_id, _, _ = env.mesh.cellular_attributes(
+            Pseudorandom.sample(backend=particulator.backend, grid=grid, n_sd=len(n))
+        )
         attributes = {"n": n, "volume": v, "cell id": cell_id}
         particulator.build(attributes)
         sut.actual_length = particulator.attributes._ParticleAttributes__idx.length
diff --git a/tests/unit_tests/initialisation/test_aerosol_init.py b/tests/unit_tests/initialisation/test_aerosol_init.py
@@ -7,8 +7,6 @@
 from PySDM.initialisation.aerosol_composition import DryAerosolMixture
 from PySDM.physics import si
 
-# spectrum = spectra.Lognormal(norm_factor=100.0 / si.cm**3, m_mode=50.0 * si.nm, s_geom=2.0)
-
 
 @pytest.mark.parametrize(
     "mass_fractions",
diff --git a/tests/unit_tests/initialisation/test_spatial_discretisation.py b/tests/unit_tests/initialisation/test_spatial_discretisation.py
@@ -0,0 +1,38 @@
+# pylint: disable=missing-module-docstring,missing-class-docstring,missing-function-docstring
+import numpy as np
+import pytest
+
+from PySDM import Formulae
+from PySDM.initialisation.sampling.spatial_sampling import Pseudorandom
+
+from ...backends_fixture import backend_class
+
+assert hasattr(backend_class, "_pytestfixturefunction")
+
+
+@pytest.mark.parametrize(
+    "seeds",
+    (
+        pytest.param(
+            (44, 33), marks=pytest.mark.xfail(strict=True), id="different seeds"
+        ),
+        pytest.param((44, 44), id="same seeds"),
+    ),
+)
+def test_pseudorandom_reproducible(
+    seeds, backend_class
+):  # pylint: disable=redefined-outer-name
+    # arrange
+    assert len(seeds) == 2
+    backends = [backend_class(Formulae(seed=seed)) for seed in seeds]
+    grid = (12, 13)
+    n_sd = 14
+
+    # act
+    positions = [
+        Pseudorandom.sample(backend=backend, grid=grid, n_sd=n_sd)
+        for backend in backends
+    ]
+
+    # assert
+    np.testing.assert_array_equal(positions[0], positions[1])
diff --git a/tests/unit_tests/initialisation/test_spectral_discretisation.py b/tests/unit_tests/initialisation/test_spectral_discretisation.py
@@ -3,39 +3,39 @@
 import pytest
 
 from PySDM import Formulae
-from PySDM.initialisation.sampling import spectral_sampling, spectro_glacial_sampling
+from PySDM.initialisation.sampling import spectral_sampling
 from PySDM.initialisation.spectra.lognormal import Lognormal
+from PySDM.physics import si
+
+from ...backends_fixture import backend_class
+
+assert hasattr(backend_class, "_pytestfixturefunction")
 
 m_mode = 0.5e-5
 n_part = 256 * 16
 s_geom = 1.5
 spectrum = Lognormal(n_part, m_mode, s_geom)
-m_range = (0.1e-6, 100e-6)
-formulae = Formulae(
-    freezing_temperature_spectrum="Niemand_et_al_2012",
-    constants={"NIEMAND_A": -0.517, "NIEMAND_B": 8.934},
-)
+m_range = (0.1 * si.um, 100 * si.um)
+formulae = Formulae()
 
 
 @pytest.mark.parametrize(
     "discretisation",
-    [
+    (
         pytest.param(spectral_sampling.Linear(spectrum, m_range)),
         pytest.param(spectral_sampling.Logarithmic(spectrum, m_range)),
         pytest.param(spectral_sampling.ConstantMultiplicity(spectrum, m_range)),
         pytest.param(spectral_sampling.UniformRandom(spectrum, m_range)),
-        # TODO #599
-    ],
+    ),
 )
-def test_spectral_discretisation(discretisation):
+# pylint: disable=redefined-outer-name
+def test_spectral_discretisation(discretisation, backend_class):
     # Arrange
     n_sd = 100000
+    backend = backend_class()
 
     # Act
-    if isinstance(discretisation, spectro_glacial_sampling.SpectroGlacialSampling):
-        m, _, __, n = discretisation.sample(n_sd)
-    else:
-        m, n = discretisation.sample(n_sd)
+    m, n = discretisation.sample(n_sd, backend=backend)
 
     # Assert
     assert m.shape == n.shape
diff --git a/tests/unit_tests/initialisation/test_spectro_glacial_discretisation.py b/tests/unit_tests/initialisation/test_spectro_glacial_discretisation.py

Original file line number	Diff line number	Diff line change
`@@ -30,6 +30,6 @@ def register(self, builder):`
`30`	`30`	`def init_attributes(self, *, spectral_discretisation):`
`31`	`31`	`attributes = {}`
`32`	`32`	`attributes["volume"], attributes["n"] = spectral_discretisation.sample(`
`33`		`- self.particulator.n_sd`
	`33`	`+ backend=self.particulator.backend, n_sd=self.particulator.n_sd`
`34`	`34`	`)`
`35`	`35`	`return attributes`