Skip to content

CU-86ad1nbm0 Fix instability in multi-worker Wepy setups by replacing omm.State with custom State class #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
edeustua wants to merge 5 commits into
base: master
Choose a base branch
from
Open

CU-86ad1nbm0 Fix instability in multi-worker Wepy setups by replacing omm.State with custom State class #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
bbf12bf
fix: Fix bugs
edeustua Jul 15, 2025
bb96d01
format: Format code with ruff
edeustua Aug 4, 2025
1e2398b
feat: Add types, format, and cleanup unused symbols
edeustua Oct 14, 2025
a02f108
feat: PlainWalkerState
edeustua Oct 14, 2025
a79321e
fix: Clean up test code
edeustua Nov 5, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion pyproject.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -37,7 +37,10 @@ dependencies = [
"jinja2",
"pint",
"multiprocessing_logging",
"openmm[cuda12]",
"openmm",
"mdtraj>=1.11.0",
"openmm-cuda-12>=8.3.1",
"rich>=14.1.0",
]

[project.optional-dependencies]
Expand Down
30 changes: 13 additions & 17 deletions src/wepy/analysis/contig_tree.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -175,23 +175,19 @@ def __init__(
if continuations is Ellipsis:
# add continuations involving both ends of the continuation

self._continuations.update(
[
(a, b)
for a, b in wepy_h5.continuations
if a in self._run_idxs and b in self._run_idxs
]
)
self._continuations.update([
(a, b)
for a, b in wepy_h5.continuations
if a in self._run_idxs and b in self._run_idxs
])

# if a subset of continuations was given use only those
elif continuations is not None:
self._continuations.update(
[
(a, b)
for a, b in continuations
if a in self._run_idxs and b in self._run_idxs
]
)
self._continuations.update([
(a, b)
for a, b in continuations
if a in self._run_idxs and b in self._run_idxs
])

# using the wepy_h5 create a tree of the cycles
self._create_tree(wepy_h5)
Expand Down Expand Up @@ -1564,9 +1560,9 @@ def __init__(self, wepy_h5, **kwargs):

# check that the result is a single contig
spanning_contig_traces = self.spanning_contig_traces()
assert (
len(spanning_contig_traces) == 1
), "continuations given do not form a single contig"
assert len(spanning_contig_traces) == 1, (
"continuations given do not form a single contig"
)

# if so we add some useful attributes valid for only a
# standalone contig
Expand Down
170 changes: 92 additions & 78 deletions src/wepy/hdf5.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -892,19 +892,21 @@ def __init__(
# read only mode
elif self._wepy_mode == "r":
# if any data was given, warn the user
if any([
kwarg is not None
for kwarg in [
topology,
units,
sparse_fields,
feature_shapes,
feature_dtypes,
n_dims,
alt_reps,
main_rep_idxs,
if any(
[
kwarg is not None
for kwarg in [
topology,
units,
sparse_fields,
feature_shapes,
feature_dtypes,
n_dims,
alt_reps,
main_rep_idxs,
]
]
]):
):
warn("Data was given but opening in read-only mode", RuntimeWarning)

# then run the initialization process
Expand Down Expand Up @@ -973,9 +975,9 @@ def _create_init(self):
and set with the new ones if given.
"""

assert self._topology is not None, (
"Topology must be given for a creation constructor"
)
assert (
self._topology is not None
), "Topology must be given for a creation constructor"

# initialize the runs group
runs_grp = self._h5.create_group(RUNS)
Expand Down Expand Up @@ -1617,9 +1619,9 @@ def _extend_contiguous_traj_field(self, run_idx, traj_idx, field_path, field_dat
field = traj_grp[field_path]

# make sure this is a feature vector
assert len(field_data.shape) > 1, (
"field_data must be a feature vector with the same number of dimensions as the number"
)
assert (
len(field_data.shape) > 1
), "field_data must be a feature vector with the same number of dimensions as the number"

# of datase new frames
n_new_frames = field_data.shape[0]
Expand All @@ -1628,9 +1630,9 @@ def _extend_contiguous_traj_field(self, run_idx, traj_idx, field_path, field_dat
if all([i == 0 for i in field.shape]):
# check the feature shape against the maxshape which gives
# the feature dimensions for an empty dataset
assert field_data.shape[1:] == field.maxshape[1:], (
"field feature dimensions must be the same, i.e. all but the first dimension"
)
assert (
field_data.shape[1:] == field.maxshape[1:]
), "field feature dimensions must be the same, i.e. all but the first dimension"

# if it is empty resize it to make an array the size of
# the new field_data with the maxshape for the feature
Expand All @@ -1644,9 +1646,9 @@ def _extend_contiguous_traj_field(self, run_idx, traj_idx, field_path, field_dat
else:
# make sure the new data has the right dimensions against
# the shape it already has
assert field_data.shape[1:] == field.shape[1:], (
"field feature dimensions must be the same, i.e. all but the first dimension"
)
assert (
field_data.shape[1:] == field.shape[1:]
), "field feature dimensions must be the same, i.e. all but the first dimension"

# append to the dataset on the first dimension, keeping the
# others the same, these must be feature vectors and therefore
Expand Down Expand Up @@ -1688,10 +1690,10 @@ def _extend_sparse_traj_field(
if all([i == 0 for i in field_data.shape]):
# check the feature shape against the maxshape which gives
# the feature dimensions for an empty dataset
assert values.shape[1:] == field_data.maxshape[1:], (
"input value features have shape {}, expected {}".format(
values.shape[1:], field_data.maxshape[1:]
)
assert (
values.shape[1:] == field_data.maxshape[1:]
), "input value features have shape {}, expected {}".format(
values.shape[1:], field_data.maxshape[1:]
)

# if it is empty resize it to make an array the size of
Expand All @@ -1705,25 +1707,29 @@ def _extend_sparse_traj_field(

else:
# make sure the new data has the right dimensions
assert values.shape[1:] == field_data.shape[1:], (
"field feature dimensions must be the same, i.e. all but the first dimension"
)
assert (
values.shape[1:] == field_data.shape[1:]
), "field feature dimensions must be the same, i.e. all but the first dimension"

# append to the dataset on the first dimension, keeping the
# others the same, these must be feature vectors and therefore
# must exist
field_data.resize((
field_data.shape[0] + n_new_frames,
*field_data.shape[1:],
))
field_data.resize(
(
field_data.shape[0] + n_new_frames,
*field_data.shape[1:],
)
)
# add the new data
field_data[-n_new_frames:, ...] = values

# add the sparse idxs in the same way
field_sparse_idxs.resize((
field_sparse_idxs.shape[0] + n_new_frames,
*field_sparse_idxs.shape[1:],
))
field_sparse_idxs.resize(
(
field_sparse_idxs.shape[0] + n_new_frames,
*field_sparse_idxs.shape[1:],
)
)
# add the new data
field_sparse_idxs[-n_new_frames:, ...] = sparse_idxs

Expand Down Expand Up @@ -1871,9 +1877,9 @@ def _extend_run_record_data_field(
field = records_grp[field_name]

# make sure this is a feature vector
assert len(field_data.shape) > 1, (
"field_data must be a feature vector with the same number of dimensions as the number"
)
assert (
len(field_data.shape) > 1
), "field_data must be a feature vector with the same number of dimensions as the number"

# of datase new frames
n_new_frames = field_data.shape[0]
Expand Down Expand Up @@ -1914,9 +1920,9 @@ def _extend_run_record_data_field(
if all([i == 0 for i in field.shape]):
# check the feature shape against the maxshape which gives
# the feature dimensions for an empty dataset
assert field_data.shape[1:] == field.maxshape[1:], (
"field feature dimensions must be the same, i.e. all but the first dimension"
)
assert (
field_data.shape[1:] == field.maxshape[1:]
), "field feature dimensions must be the same, i.e. all but the first dimension"

# if it is empty resize it to make an array the size of
# the new field_data with the maxshape for the feature
Expand Down Expand Up @@ -3940,14 +3946,18 @@ def add_continuation(self, continuation_run, base_run):
"""

continuations_dset = self.settings_grp[CONTINUATIONS]
continuations_dset.resize((
continuations_dset.shape[0] + 1,
continuations_dset.shape[1],
))
continuations_dset[continuations_dset.shape[0] - 1] = np.array([
continuation_run,
base_run,
])
continuations_dset.resize(
(
continuations_dset.shape[0] + 1,
continuations_dset.shape[1],
)
)
continuations_dset[continuations_dset.shape[0] - 1] = np.array(
[
continuation_run,
base_run,
]
)

def new_run(self, init_walkers, continue_run=None, **kwargs):
"""Initialize a new run.
Expand Down Expand Up @@ -4346,9 +4356,9 @@ def add_traj(self, run_idx, data, weights=None, sparse_idxs=None, metadata=None)
weights = np.ones((n_frames, 1), dtype=float)
else:
assert isinstance(weights, np.ndarray), "weights must be a numpy.ndarray"
assert weights.shape[0] == n_frames, (
"weights and the number of frames must be the same length"
)
assert (
weights.shape[0] == n_frames
), "weights and the number of frames must be the same length"

# current traj_idx
traj_idx = self.next_run_traj_idx(run_idx)
Expand All @@ -4374,15 +4384,15 @@ def add_traj(self, run_idx, data, weights=None, sparse_idxs=None, metadata=None)
)

# check to make sure the positions are the right shape
assert traj_data[POSITIONS].shape[1] == self.num_atoms, (
"positions given have different number of atoms: {}, should be {}".format(
traj_data[POSITIONS].shape[1], self.num_atoms
)
assert (
traj_data[POSITIONS].shape[1] == self.num_atoms
), "positions given have different number of atoms: {}, should be {}".format(
traj_data[POSITIONS].shape[1], self.num_atoms
)
assert traj_data[POSITIONS].shape[2] == self.num_dims, (
"positions given have different number of dims: {}, should be {}".format(
traj_data[POSITIONS].shape[2], self.num_dims
)
assert (
traj_data[POSITIONS].shape[2] == self.num_dims
), "positions given have different number of dims: {}, should be {}".format(
traj_data[POSITIONS].shape[2], self.num_dims
)

# add datasets to the traj group
Expand Down Expand Up @@ -4453,9 +4463,9 @@ def extend_traj(self, run_idx, traj_idx, data, weights=None):
"""

if self._wepy_mode == "c-":
assert self._append_flags[dataset_key], (
"dataset is not available for appending to"
)
assert self._append_flags[
dataset_key
], "dataset is not available for appending to"

# convenient alias
traj_data = data
Expand All @@ -4479,20 +4489,22 @@ def extend_traj(self, run_idx, traj_idx, data, weights=None):
weights = np.ones((n_new_frames, 1), dtype=float)
else:
assert isinstance(weights, np.ndarray), "weights must be a numpy.ndarray"
assert weights.shape[0] == n_new_frames, (
"weights and the number of frames must be the same length"
)
assert (
weights.shape[0] == n_new_frames
), "weights and the number of frames must be the same length"

# add the weights
weights_ds = traj_grp[WEIGHTS]

# append to the dataset on the first dimension, keeping the
# others the same, if they exist
if len(weights_ds.shape) > 1:
weights_ds.resize((
weights_ds.shape[0] + n_new_frames,
*weights_ds.shape[1:],
))
weights_ds.resize(
(
weights_ds.shape[0] + n_new_frames,
*weights_ds.shape[1:],
)
)
else:
weights_ds.resize((weights_ds.shape[0] + n_new_frames,))

Expand Down Expand Up @@ -5517,10 +5529,12 @@ def get_contig_trace_fields(self, contig_trace, fields):

# (there must be the same number of trajectories in each run)
n_trajs_test = self.num_run_trajs(run_idxs[0])
assert all([
True if n_trajs_test == self.num_run_trajs(run_idx) else False
for run_idx in run_idxs
])
assert all(
[
True if n_trajs_test == self.num_run_trajs(run_idx) else False
for run_idx in run_idxs
]
)

# then using this we go run by run and get all the
# trajectories
Expand Down
16 changes: 7 additions & 9 deletions src/wepy/orchestration/configuration.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -134,9 +134,9 @@ def __init__(
### Monitor options

# get the names of the reporters in the order they are
reporter_order = tuple(
[str(reporter_class.__name__) for reporter_class in self._reporter_classes]
)
reporter_order = tuple([
str(reporter_class.__name__) for reporter_class in self._reporter_classes
])

# init the kwargs for the monitor
if monitor_partial_kwargs is None:
Expand Down Expand Up @@ -242,12 +242,10 @@ def _gen_reporters(self):

# the number of filenames
all_exts = list(
it.chain(
*[
[ext for ext in rep.SUGGESTED_EXTENSIONS]
for rep in self.reporter_classes
]
)
it.chain(*[
[ext for ext in rep.SUGGESTED_EXTENSIONS]
for rep in self.reporter_classes
])
)
n_exts = len(all_exts)

Expand Down
Loading