From b9e4b6412a3e6c390b028d640858be12977535c7 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 28 Jan 2024 16:09:35 -0500 Subject: [PATCH 1/4] add size and replace arguments to deepmd.utils.random.choice Fix https://github.com/deepmodeling/deepmd-kit/security/code-scanning/2096 Signed-off-by: Jinzhe Zeng --- deepmd/utils/random.py | 27 +++++++++++++++++++++------ 1 file changed, 21 insertions(+), 6 deletions(-) diff --git a/deepmd/utils/random.py b/deepmd/utils/random.py index 8944419412..5832978fbc 100644 --- a/deepmd/utils/random.py +++ b/deepmd/utils/random.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: LGPL-3.0-or-later from typing import ( Optional, + Tuple, + Union, ) import numpy as np @@ -8,22 +10,35 @@ _RANDOM_GENERATOR = np.random.RandomState() -def choice(a: np.ndarray, p: Optional[np.ndarray] = None): +def choice( + a: Union[np.ndarray, int], + size: Union[int, Tuple[int, ...]] = None, + replace: bool = True, + p: Optional[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. + a: 1-D array-like or int + If an ndarray, a random sample is generated from its elements. If an int, + the random sample is generated as if it were np.arange(a) + size: int or tuple of ints, optional + Output shape. If the given shape is, e.g., (m, n, k), then m * n * k samples + are drawn. Default is None, in which case a single value is returned. + replace: boolean, optional + Whether the sample is with or without replacement. Default is True, meaning + that a value of a can be selected multiple times. + p: 1-D array-like, optional + The probabilities associated with each entry in a. If not given, the sample + assumes a uniform distribution over all entries in a. Returns ------- np.ndarray arrays with results and their shapes """ - return _RANDOM_GENERATOR.choice(a, p=p) + return _RANDOM_GENERATOR.choice(a, size=size, replace=replace, p=p) def random(size=None): From a3671d7dab827cb9a1516524036213ba5458d844 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sun, 28 Jan 2024 21:10:11 +0000 Subject: [PATCH 2/4] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- deepmd/utils/random.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/deepmd/utils/random.py b/deepmd/utils/random.py index 5832978fbc..ecc0787d38 100644 --- a/deepmd/utils/random.py +++ b/deepmd/utils/random.py @@ -11,25 +11,25 @@ def choice( - a: Union[np.ndarray, int], - size: Union[int, Tuple[int, ...]] = None, - replace: bool = True, - p: Optional[np.ndarray] = None, - ): + a: Union[np.ndarray, int], + size: Union[int, Tuple[int, ...]] = None, + replace: bool = True, + p: Optional[np.ndarray] = None, +): """Generates a random sample from a given 1-D array. Parameters ---------- - a: 1-D array-like or int + a : 1-D array-like or int If an ndarray, a random sample is generated from its elements. If an int, the random sample is generated as if it were np.arange(a) - size: int or tuple of ints, optional + size : int or tuple of ints, optional Output shape. If the given shape is, e.g., (m, n, k), then m * n * k samples are drawn. Default is None, in which case a single value is returned. - replace: boolean, optional + replace : boolean, optional Whether the sample is with or without replacement. Default is True, meaning that a value of a can be selected multiple times. - p: 1-D array-like, optional + p : 1-D array-like, optional The probabilities associated with each entry in a. If not given, the sample assumes a uniform distribution over all entries in a. From b3f0b8d2e685320f816f0e36f23b0bf47a56bc40 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 28 Jan 2024 16:11:27 -0500 Subject: [PATCH 3/4] add Optional Signed-off-by: Jinzhe Zeng --- deepmd/utils/random.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepmd/utils/random.py b/deepmd/utils/random.py index ecc0787d38..44ea6a1dac 100644 --- a/deepmd/utils/random.py +++ b/deepmd/utils/random.py @@ -12,7 +12,7 @@ def choice( a: Union[np.ndarray, int], - size: Union[int, Tuple[int, ...]] = None, + size: Optional[Union[int, Tuple[int, ...]]] = None, replace: bool = True, p: Optional[np.ndarray] = None, ): From 893d9043a6622ab0acfd24e791ed42db1be5b680 Mon Sep 17 00:00:00 2001 From: Jinzhe Zeng Date: Sun, 28 Jan 2024 16:29:50 -0500 Subject: [PATCH 4/4] fix p argument calling --- deepmd/pt/utils/dataset.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/deepmd/pt/utils/dataset.py b/deepmd/pt/utils/dataset.py index 24daa6e37e..3920499d3a 100644 --- a/deepmd/pt/utils/dataset.py +++ b/deepmd/pt/utils/dataset.py @@ -879,7 +879,7 @@ def __len__(self): def __getitem__(self, index=None): """Get a batch of frames from the selected system.""" if index is None: - index = dp_random.choice(np.arange(self.nsystems), self.probs) + index = dp_random.choice(np.arange(self.nsystems), p=self.probs) b_data = self._data_systems[index].get_batch(self._batch_size) b_data["natoms"] = torch.tensor( self._natoms_vec[index], device=env.PREPROCESS_DEVICE @@ -892,7 +892,7 @@ def __getitem__(self, index=None): def get_training_batch(self, index=None): """Get a batch of frames from the selected system.""" if index is None: - index = dp_random.choice(np.arange(self.nsystems), self.probs) + index = dp_random.choice(np.arange(self.nsystems), p=self.probs) b_data = self._data_systems[index].get_batch_for_train(self._batch_size) b_data["natoms"] = torch.tensor( self._natoms_vec[index], device=env.PREPROCESS_DEVICE