Skip to content

Commit

Permalink
merge prob_sys_size with prob_sys_size;0:nsys:1.0 (#2963)
Browse files Browse the repository at this point in the history
to be consistent with Pytorch version

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
  • Loading branch information
CaRoLZhangxy and pre-commit-ci[bot] authored Oct 31, 2023
1 parent 5be8fc9 commit 389d403
Showing 1 changed file with 7 additions and 6 deletions.
13 changes: 7 additions & 6 deletions deepmd/utils/data_system.py
Original file line number Diff line number Diff line change
Expand Up @@ -195,8 +195,7 @@ def __init__(
assert isinstance(self.test_size, (list, np.ndarray))
assert len(self.test_size) == self.nsystems

# prob of batch, init pick idx
self.prob_nbatches = [float(i) for i in self.nbatches] / np.sum(self.nbatches)
# init pick idx
self.pick_idx = 0

# derive system probabilities
Expand Down Expand Up @@ -350,11 +349,13 @@ def set_sys_probs(self, sys_probs=None, auto_prob_style: str = "prob_sys_size"):
if auto_prob_style == "prob_uniform":
prob_v = 1.0 / float(self.nsystems)
probs = [prob_v for ii in range(self.nsystems)]
elif auto_prob_style == "prob_sys_size":
probs = self.prob_nbatches
elif auto_prob_style[:14] == "prob_sys_size;":
elif auto_prob_style[:13] == "prob_sys_size":
if auto_prob_style == "prob_sys_size":
prob_style = f"prob_sys_size;0:{self.get_nsystems()}:1.0"
else:
prob_style = auto_prob_style
probs = prob_sys_size_ext(
auto_prob_style, self.get_nsystems(), self.nbatches
prob_style, self.get_nsystems(), self.nbatches
)
else:
raise RuntimeError("Unknown auto prob style: " + auto_prob_style)
Expand Down

0 comments on commit 389d403

Please sign in to comment.