Use a local Numpy random generator. (#1030)

* Use a local Numpy random generator.

* Fix import error.

* Replace random generator in test_gen_stat_data.py.

Author: shishaochen

deepmd/entrypoints/test.py

@@ -6,6 +6,7 @@
 import numpy as np
 from deepmd import DeepPotential
 from deepmd.common import expand_sys_str
+from deepmd.utils import random as dp_random
 from deepmd.utils.data import DeepmdData
 from deepmd.utils.weight_avg import weighted_average
 
@@ -64,7 +65,7 @@ def test(
 
     # init random seed
     if rand_seed is not None:
-        np.random.seed(rand_seed % (2 ** 32))
+        dp_random.seed(rand_seed % (2 ** 32))
 
     # init model
     dp = DeepPotential(model)

deepmd/entrypoints/train.py

@@ -9,12 +9,12 @@
 import os
 from typing import Dict, List, Optional, Any
 
-import numpy as np
 from deepmd.common import data_requirement, expand_sys_str, j_loader, j_must_have
 from deepmd.env import tf, reset_default_tf_session_config
 from deepmd.infer.data_modifier import DipoleChargeModifier
 from deepmd.train.run_options import BUILD, CITATION, WELCOME, RunOptions
 from deepmd.train.trainer import DPTrainer
+from deepmd.utils import random as dp_random
 from deepmd.utils.argcheck import normalize
 from deepmd.utils.compat import updata_deepmd_input
 from deepmd.utils.data_system import DeepmdDataSystem
@@ -135,7 +135,7 @@ def _do_work(jdata: Dict[str, Any], run_opt: RunOptions, is_compress: bool = Fal
     seed = jdata["training"].get("seed", None)
     if seed is not None:
         seed = seed % (2 ** 32)
-    np.random.seed(seed)
+    dp_random.seed(seed)
 
     # setup data modifier
     modifier = get_modifier(jdata["model"].get("modifier", None))

deepmd/utils/data.py

@@ -9,6 +9,7 @@
 
 from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
 from deepmd.env import GLOBAL_ENER_FLOAT_PRECISION
+from deepmd.utils import random as dp_random
 
 log = logging.getLogger(__name__)
 
@@ -396,7 +397,7 @@ def _shuffle_data (self,
         ret = {}
         nframes = data['coord'].shape[0]
         idx = np.arange (nframes)
-        np.random.shuffle (idx)
+        dp_random.shuffle(idx)
         for kk in data :
             if type(data[kk]) == np.ndarray and \
                len(data[kk].shape) == 2 and \
@@ -676,7 +677,7 @@ def load_set(self, set_name, shuffle = True):
         # shuffle data
         if shuffle:
             idx = np.arange (nframe)
-            np.random.shuffle (idx)
+            dp_random.shuffle(idx)
             for ii in data:
                 if ii != "prop_c":
                     data[ii] = data[ii][idx]

deepmd/utils/data_system.py

@@ -7,6 +7,7 @@
 import numpy as np
 from typing import Tuple, List
 
+from deepmd.utils import random as dp_random
 from deepmd.utils.data import DataSets
 from deepmd.utils.data import DeepmdData
 
@@ -330,7 +331,7 @@ def get_batch(self, sys_idx : int = None):
             self.pick_idx = sys_idx
         else :
             # prob = self._get_sys_probs(sys_probs, auto_prob_style)
-            self.pick_idx = np.random.choice(np.arange(self.nsystems), p=self.sys_probs)
+            self.pick_idx = dp_random.choice(np.arange(self.nsystems), p=self.sys_probs)
         b_data = self.data_systems[self.pick_idx].get_batch(self.batch_size[self.pick_idx])
         b_data["natoms_vec"] = self.natoms_vec[self.pick_idx]
         b_data["default_mesh"] = self.default_mesh[self.pick_idx]
@@ -671,7 +672,7 @@ def get_batch (self,
                     raise RuntimeError("unkown get_batch style")
             else :
                 prob = self.process_sys_weights(sys_weights)
-            self.pick_idx = np.random.choice(np.arange(self.nsystems), p = prob)
+            self.pick_idx = dp_random.choice(np.arange(self.nsystems), p=prob)
         b_data = self.data_systems[self.pick_idx].get_batch(self.batch_size[self.pick_idx])
         b_data["natoms_vec"] = self.natoms_vec[self.pick_idx]
         b_data["default_mesh"] = self.default_mesh[self.pick_idx]

deepmd/utils/random.py

@@ -0,0 +1,63 @@
+import numpy as np
+
+
+_RANDOM_GENERATOR = np.random.RandomState()
+
+
+def choice(a: np.ndarray, p: np.ndarray = None):
+    """Generates a random sample from a given 1-D array.
+
+    Parameters
+    ----------
+    a : np.ndarray
+        A random sample is generated from its elements.
+    p : np.ndarray
+        The probabilities associated with each entry in a.
+
+    Returns
+    -------
+    np.ndarray
+        arrays with results and their shapes
+    """
+    return _RANDOM_GENERATOR.choice(a, p=p)
+
+
+def random(size=None):
+    """Return random floats in the half-open interval [0.0, 1.0).
+
+    Parameters
+    ----------
+    size
+        Output shape.
+
+    Returns
+    -------
+    np.ndarray
+        Arrays with results and their shapes.
+    """
+    return _RANDOM_GENERATOR.random_sample(size)
+
+
+def seed(val: int = None):
+    """Seed the generator.
+
+    Parameters
+    ----------
+    val : int
+        Seed.
+    """
+    _RANDOM_GENERATOR.seed(val)
+
+
+def shuffle(x: np.ndarray):
+    """Modify a sequence in-place by shuffling its contents.
+
+    Parameters
+    ----------
+    x : np.ndarray
+        The array or list to be shuffled.
+    """
+    _RANDOM_GENERATOR.shuffle(x)
+
+
+__all__ = ['choice', 'random', 'seed', 'shuffle']

source/tests/test_gen_stat_data.py

@@ -1,8 +1,9 @@
-import os,sys,platform,json,shutil
+import shutil
 import numpy as np
 import unittest
 import dpdata
 
+from deepmd.utils import random as dp_random
 from deepmd.utils.data_system import DeepmdDataSystem
 from deepmd.fit import EnerFitting
 from deepmd.model.model_stat import make_stat_input, merge_sys_stat, _make_all_stat_ref
@@ -46,31 +47,31 @@ def _comp_data(self, d0, d1) :
                     self.assertAlmostEqual(d0[ii][jj][kk], d1[ii][jj][kk])
 
     def test_merge_all_stat(self):
-        np.random.seed(0)
+        dp_random.seed(0)
         data0 = DeepmdDataSystem(['system_0', 'system_1'], 
                                 5, 
                                 10, 
                                 1.0)
         data0.add('energy', 1, must = True)
-        np.random.seed(0)
+        dp_random.seed(0)
         data1 = DeepmdDataSystem(['system_0', 'system_1'], 
                                 5, 
                                 10, 
                                 1.0)
         data1.add('force', 3, atomic = True, must = True)
-        np.random.seed(0)
+        dp_random.seed(0)
         data2 = DeepmdDataSystem(['system_0', 'system_1'], 
                                 5, 
                                 10, 
                                 1.0)
         data2.add('force', 3, atomic = True, must = True)
 
-        np.random.seed(0)
+        dp_random.seed(0)
         all_stat_0 = make_stat_input(data0, 10, merge_sys = False)
-        np.random.seed(0)
+        dp_random.seed(0)
         all_stat_1 = make_stat_input(data1, 10, merge_sys = True)
         all_stat_2 = merge_sys_stat(all_stat_0)
-        np.random.seed(0)
+        dp_random.seed(0)
         all_stat_3 = _make_all_stat_ref(data2, 10)
 
         ####################################
@@ -109,7 +110,7 @@ def tearDown(self):
         shutil.rmtree('system_1')
 
     def test_ener_shift(self):
-        np.random.seed(0)
+        dp_random.seed(0)
         data = DeepmdDataSystem(['system_0', 'system_1'], 
                                 5, 
                                 10,