westpa 2022.13__cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- westpa/__init__.py +14 -0
- westpa/_version.py +21 -0
- westpa/analysis/__init__.py +5 -0
- westpa/analysis/core.py +749 -0
- westpa/analysis/statistics.py +27 -0
- westpa/analysis/trajectories.py +369 -0
- westpa/cli/__init__.py +0 -0
- westpa/cli/core/__init__.py +0 -0
- westpa/cli/core/w_fork.py +152 -0
- westpa/cli/core/w_init.py +230 -0
- westpa/cli/core/w_run.py +77 -0
- westpa/cli/core/w_states.py +212 -0
- westpa/cli/core/w_succ.py +99 -0
- westpa/cli/core/w_truncate.py +68 -0
- westpa/cli/tools/__init__.py +0 -0
- westpa/cli/tools/ploterr.py +506 -0
- westpa/cli/tools/plothist.py +706 -0
- westpa/cli/tools/w_assign.py +597 -0
- westpa/cli/tools/w_bins.py +166 -0
- westpa/cli/tools/w_crawl.py +119 -0
- westpa/cli/tools/w_direct.py +557 -0
- westpa/cli/tools/w_dumpsegs.py +94 -0
- westpa/cli/tools/w_eddist.py +506 -0
- westpa/cli/tools/w_fluxanl.py +376 -0
- westpa/cli/tools/w_ipa.py +832 -0
- westpa/cli/tools/w_kinavg.py +127 -0
- westpa/cli/tools/w_kinetics.py +96 -0
- westpa/cli/tools/w_multi_west.py +414 -0
- westpa/cli/tools/w_ntop.py +213 -0
- westpa/cli/tools/w_pdist.py +515 -0
- westpa/cli/tools/w_postanalysis_matrix.py +82 -0
- westpa/cli/tools/w_postanalysis_reweight.py +53 -0
- westpa/cli/tools/w_red.py +491 -0
- westpa/cli/tools/w_reweight.py +780 -0
- westpa/cli/tools/w_select.py +226 -0
- westpa/cli/tools/w_stateprobs.py +111 -0
- westpa/cli/tools/w_timings.py +113 -0
- westpa/cli/tools/w_trace.py +599 -0
- westpa/core/__init__.py +0 -0
- westpa/core/_rc.py +673 -0
- westpa/core/binning/__init__.py +55 -0
- westpa/core/binning/_assign.c +36018 -0
- westpa/core/binning/_assign.cpython-312-aarch64-linux-gnu.so +0 -0
- westpa/core/binning/_assign.pyx +370 -0
- westpa/core/binning/assign.py +454 -0
- westpa/core/binning/binless.py +96 -0
- westpa/core/binning/binless_driver.py +54 -0
- westpa/core/binning/binless_manager.py +189 -0
- westpa/core/binning/bins.py +47 -0
- westpa/core/binning/mab.py +506 -0
- westpa/core/binning/mab_driver.py +54 -0
- westpa/core/binning/mab_manager.py +197 -0
- westpa/core/data_manager.py +1761 -0
- westpa/core/extloader.py +74 -0
- westpa/core/h5io.py +1079 -0
- westpa/core/kinetics/__init__.py +24 -0
- westpa/core/kinetics/_kinetics.c +45174 -0
- westpa/core/kinetics/_kinetics.cpython-312-aarch64-linux-gnu.so +0 -0
- westpa/core/kinetics/_kinetics.pyx +815 -0
- westpa/core/kinetics/events.py +147 -0
- westpa/core/kinetics/matrates.py +156 -0
- westpa/core/kinetics/rate_averaging.py +266 -0
- westpa/core/progress.py +218 -0
- westpa/core/propagators/__init__.py +54 -0
- westpa/core/propagators/executable.py +592 -0
- westpa/core/propagators/loaders.py +196 -0
- westpa/core/reweight/__init__.py +14 -0
- westpa/core/reweight/_reweight.c +36899 -0
- westpa/core/reweight/_reweight.cpython-312-aarch64-linux-gnu.so +0 -0
- westpa/core/reweight/_reweight.pyx +439 -0
- westpa/core/reweight/matrix.py +126 -0
- westpa/core/segment.py +119 -0
- westpa/core/sim_manager.py +839 -0
- westpa/core/states.py +359 -0
- westpa/core/systems.py +93 -0
- westpa/core/textio.py +74 -0
- westpa/core/trajectory.py +603 -0
- westpa/core/we_driver.py +910 -0
- westpa/core/wm_ops.py +43 -0
- westpa/core/yamlcfg.py +298 -0
- westpa/fasthist/__init__.py +34 -0
- westpa/fasthist/_fasthist.c +38755 -0
- westpa/fasthist/_fasthist.cpython-312-aarch64-linux-gnu.so +0 -0
- westpa/fasthist/_fasthist.pyx +222 -0
- westpa/mclib/__init__.py +271 -0
- westpa/mclib/__main__.py +28 -0
- westpa/mclib/_mclib.c +34610 -0
- westpa/mclib/_mclib.cpython-312-aarch64-linux-gnu.so +0 -0
- westpa/mclib/_mclib.pyx +226 -0
- westpa/oldtools/__init__.py +4 -0
- westpa/oldtools/aframe/__init__.py +35 -0
- westpa/oldtools/aframe/atool.py +75 -0
- westpa/oldtools/aframe/base_mixin.py +26 -0
- westpa/oldtools/aframe/binning.py +178 -0
- westpa/oldtools/aframe/data_reader.py +560 -0
- westpa/oldtools/aframe/iter_range.py +200 -0
- westpa/oldtools/aframe/kinetics.py +117 -0
- westpa/oldtools/aframe/mcbs.py +153 -0
- westpa/oldtools/aframe/output.py +39 -0
- westpa/oldtools/aframe/plotting.py +88 -0
- westpa/oldtools/aframe/trajwalker.py +126 -0
- westpa/oldtools/aframe/transitions.py +469 -0
- westpa/oldtools/cmds/__init__.py +0 -0
- westpa/oldtools/cmds/w_ttimes.py +361 -0
- westpa/oldtools/files.py +34 -0
- westpa/oldtools/miscfn.py +23 -0
- westpa/oldtools/stats/__init__.py +4 -0
- westpa/oldtools/stats/accumulator.py +35 -0
- westpa/oldtools/stats/edfs.py +129 -0
- westpa/oldtools/stats/mcbs.py +96 -0
- westpa/tools/__init__.py +33 -0
- westpa/tools/binning.py +472 -0
- westpa/tools/core.py +340 -0
- westpa/tools/data_reader.py +159 -0
- westpa/tools/dtypes.py +31 -0
- westpa/tools/iter_range.py +198 -0
- westpa/tools/kinetics_tool.py +343 -0
- westpa/tools/plot.py +283 -0
- westpa/tools/progress.py +17 -0
- westpa/tools/selected_segs.py +154 -0
- westpa/tools/wipi.py +751 -0
- westpa/trajtree/__init__.py +4 -0
- westpa/trajtree/_trajtree.c +17829 -0
- westpa/trajtree/_trajtree.cpython-312-aarch64-linux-gnu.so +0 -0
- westpa/trajtree/_trajtree.pyx +130 -0
- westpa/trajtree/trajtree.py +117 -0
- westpa/westext/__init__.py +0 -0
- westpa/westext/adaptvoronoi/__init__.py +3 -0
- westpa/westext/adaptvoronoi/adaptVor_driver.py +214 -0
- westpa/westext/hamsm_restarting/__init__.py +3 -0
- westpa/westext/hamsm_restarting/example_overrides.py +35 -0
- westpa/westext/hamsm_restarting/restart_driver.py +1165 -0
- westpa/westext/stringmethod/__init__.py +11 -0
- westpa/westext/stringmethod/fourier_fitting.py +69 -0
- westpa/westext/stringmethod/string_driver.py +253 -0
- westpa/westext/stringmethod/string_method.py +306 -0
- westpa/westext/weed/BinCluster.py +180 -0
- westpa/westext/weed/ProbAdjustEquil.py +100 -0
- westpa/westext/weed/UncertMath.py +247 -0
- westpa/westext/weed/__init__.py +10 -0
- westpa/westext/weed/weed_driver.py +192 -0
- westpa/westext/wess/ProbAdjust.py +101 -0
- westpa/westext/wess/__init__.py +6 -0
- westpa/westext/wess/wess_driver.py +217 -0
- westpa/work_managers/__init__.py +57 -0
- westpa/work_managers/core.py +396 -0
- westpa/work_managers/environment.py +134 -0
- westpa/work_managers/mpi.py +318 -0
- westpa/work_managers/processes.py +201 -0
- westpa/work_managers/serial.py +28 -0
- westpa/work_managers/threads.py +79 -0
- westpa/work_managers/zeromq/__init__.py +20 -0
- westpa/work_managers/zeromq/core.py +635 -0
- westpa/work_managers/zeromq/node.py +131 -0
- westpa/work_managers/zeromq/work_manager.py +526 -0
- westpa/work_managers/zeromq/worker.py +320 -0
- westpa-2022.13.dist-info/METADATA +179 -0
- westpa-2022.13.dist-info/RECORD +162 -0
- westpa-2022.13.dist-info/WHEEL +7 -0
- westpa-2022.13.dist-info/entry_points.txt +30 -0
- westpa-2022.13.dist-info/licenses/LICENSE +21 -0
- westpa-2022.13.dist-info/top_level.txt +1 -0
|
@@ -0,0 +1,592 @@
|
|
|
1
|
+
import logging
|
|
2
|
+
import os
|
|
3
|
+
import shutil
|
|
4
|
+
import signal
|
|
5
|
+
import subprocess
|
|
6
|
+
import sys
|
|
7
|
+
import tempfile
|
|
8
|
+
import time
|
|
9
|
+
|
|
10
|
+
import numpy as np
|
|
11
|
+
from numpy.random import MT19937, Generator
|
|
12
|
+
|
|
13
|
+
from westpa.core.data_manager import makepath
|
|
14
|
+
from westpa.core.extloader import get_object
|
|
15
|
+
from westpa.core.propagators import WESTPropagator
|
|
16
|
+
from westpa.core.propagators.loaders import (
|
|
17
|
+
data_loaders,
|
|
18
|
+
trajectory_loaders,
|
|
19
|
+
pcoord_loader,
|
|
20
|
+
mdtraj_trajectory_loader,
|
|
21
|
+
restart_loader,
|
|
22
|
+
seglog_loader,
|
|
23
|
+
restart_writer,
|
|
24
|
+
aux_data_loader,
|
|
25
|
+
)
|
|
26
|
+
from westpa.core.propagators.loaders import * # noqa
|
|
27
|
+
from westpa.core.states import BasisState, InitialState, return_state_type
|
|
28
|
+
from westpa.core.segment import Segment
|
|
29
|
+
from westpa.core.yamlcfg import check_bool
|
|
30
|
+
|
|
31
|
+
log = logging.getLogger(__name__)
|
|
32
|
+
|
|
33
|
+
# Get a list of user-friendly signal names
|
|
34
|
+
SIGNAL_NAMES = {getattr(signal, name): name for name in dir(signal) if name.startswith('SIG') and not name.startswith('SIG_')}
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
class ExecutablePropagator(WESTPropagator):
|
|
38
|
+
ENV_CURRENT_ITER = 'WEST_CURRENT_ITER'
|
|
39
|
+
|
|
40
|
+
# Environment variables set during propagation
|
|
41
|
+
ENV_CURRENT_SEG_ID = 'WEST_CURRENT_SEG_ID'
|
|
42
|
+
ENV_CURRENT_SEG_DATA_REF = 'WEST_CURRENT_SEG_DATA_REF'
|
|
43
|
+
ENV_CURRENT_SEG_INITPOINT = 'WEST_CURRENT_SEG_INITPOINT_TYPE'
|
|
44
|
+
ENV_PARENT_SEG_ID = 'WEST_PARENT_ID'
|
|
45
|
+
ENV_PARENT_DATA_REF = 'WEST_PARENT_DATA_REF'
|
|
46
|
+
|
|
47
|
+
# Environment variables set during propagation and state generation
|
|
48
|
+
ENV_BSTATE_ID = 'WEST_BSTATE_ID'
|
|
49
|
+
ENV_BSTATE_DATA_REF = 'WEST_BSTATE_DATA_REF'
|
|
50
|
+
ENV_ISTATE_ID = 'WEST_ISTATE_ID'
|
|
51
|
+
ENV_ISTATE_DATA_REF = 'WEST_ISTATE_DATA_REF'
|
|
52
|
+
|
|
53
|
+
# Environment variables for progress coordinate calculation
|
|
54
|
+
ENV_STRUCT_DATA_REF = 'WEST_STRUCT_DATA_REF'
|
|
55
|
+
|
|
56
|
+
ENV_RAND16 = 'WEST_RAND16'
|
|
57
|
+
ENV_RAND32 = 'WEST_RAND32'
|
|
58
|
+
ENV_RAND64 = 'WEST_RAND64'
|
|
59
|
+
ENV_RAND128 = 'WEST_RAND128'
|
|
60
|
+
ENV_RANDFLOAT = 'WEST_RANDFLOAT'
|
|
61
|
+
|
|
62
|
+
makepath = staticmethod(makepath)
|
|
63
|
+
|
|
64
|
+
def __init__(self, rc=None):
|
|
65
|
+
super().__init__(rc)
|
|
66
|
+
|
|
67
|
+
# A mapping of environment variables to template strings which will be
|
|
68
|
+
# added to the environment of all children launched.
|
|
69
|
+
self.addtl_child_environ = dict()
|
|
70
|
+
|
|
71
|
+
# A mapping of executable name ('propagator', 'pre_iteration', 'post_iteration') to
|
|
72
|
+
# a dictionary of attributes like 'executable', 'stdout', 'stderr', 'environ', etc.
|
|
73
|
+
self.exe_info = {}
|
|
74
|
+
self.exe_info['propagator'] = {}
|
|
75
|
+
self.exe_info['pre_iteration'] = {}
|
|
76
|
+
self.exe_info['post_iteration'] = {}
|
|
77
|
+
self.exe_info['get_pcoord'] = {}
|
|
78
|
+
self.exe_info['gen_istate'] = {}
|
|
79
|
+
|
|
80
|
+
# A mapping of data set name ('pcoord', 'coord', 'com', etc) to a dictionary of
|
|
81
|
+
# attributes like 'loader', 'dtype', etc
|
|
82
|
+
self.data_info = {'pcoord': {}}
|
|
83
|
+
|
|
84
|
+
# Validate configuration
|
|
85
|
+
config = self.rc.config
|
|
86
|
+
|
|
87
|
+
for key in [
|
|
88
|
+
('west', 'executable', 'propagator', 'executable'),
|
|
89
|
+
('west', 'data', 'data_refs', 'segment'),
|
|
90
|
+
('west', 'data', 'data_refs', 'basis_state'),
|
|
91
|
+
('west', 'data', 'data_refs', 'initial_state'),
|
|
92
|
+
]:
|
|
93
|
+
config.require(key)
|
|
94
|
+
|
|
95
|
+
self.segment_ref_template = config['west', 'data', 'data_refs', 'segment']
|
|
96
|
+
self.basis_state_ref_template = config['west', 'data', 'data_refs', 'basis_state']
|
|
97
|
+
self.initial_state_ref_template = config['west', 'data', 'data_refs', 'initial_state']
|
|
98
|
+
store_h5 = config.get(['west', 'data', 'data_refs', 'iteration']) is not None
|
|
99
|
+
|
|
100
|
+
# Create a persistent RNG for each worker
|
|
101
|
+
self.rng = Generator(MT19937())
|
|
102
|
+
|
|
103
|
+
# Load additional environment variables for all child processes
|
|
104
|
+
self.addtl_child_environ.update({k: str(v) for k, v in (config['west', 'executable', 'environ'] or {}).items()})
|
|
105
|
+
|
|
106
|
+
# Load configuration items relating to child processes
|
|
107
|
+
for child_type in ('propagator', 'pre_iteration', 'post_iteration', 'get_pcoord', 'gen_istate', 'subgroup_walkers'):
|
|
108
|
+
child_info = config.get(['west', 'executable', child_type])
|
|
109
|
+
if not child_info:
|
|
110
|
+
continue
|
|
111
|
+
|
|
112
|
+
info_prefix = ['west', 'executable', child_type]
|
|
113
|
+
|
|
114
|
+
# require executable to be specified if anything is specified at all
|
|
115
|
+
config.require(info_prefix + ['executable'])
|
|
116
|
+
|
|
117
|
+
self.exe_info[child_type]['executable'] = child_info['executable']
|
|
118
|
+
self.exe_info[child_type]['stdin'] = child_info.get('stdin', os.devnull)
|
|
119
|
+
self.exe_info[child_type]['stdout'] = child_info.get('stdout', None)
|
|
120
|
+
self.exe_info[child_type]['stderr'] = child_info.get('stderr', None)
|
|
121
|
+
self.exe_info[child_type]['cwd'] = child_info.get('cwd', None)
|
|
122
|
+
|
|
123
|
+
if child_type not in ('propagator', 'get_pcoord', 'gen_istate'):
|
|
124
|
+
self.exe_info[child_type]['enabled'] = child_info.get('enabled', True)
|
|
125
|
+
else:
|
|
126
|
+
# for consistency, propagator, get_pcoord, and gen_istate can never be disabled
|
|
127
|
+
self.exe_info[child_type]['enabled'] = True
|
|
128
|
+
|
|
129
|
+
# apply environment modifications specific to this executable
|
|
130
|
+
self.exe_info[child_type]['environ'] = {k: str(v) for k, v in (child_info.get('environ') or {}).items()}
|
|
131
|
+
|
|
132
|
+
log.debug('exe_info: {!r}'.format(self.exe_info))
|
|
133
|
+
|
|
134
|
+
# Load configuration items relating to dataset input
|
|
135
|
+
self.data_info['pcoord'] = {'name': 'pcoord', 'loader': pcoord_loader, 'enabled': True, 'filename': None, 'dir': False}
|
|
136
|
+
|
|
137
|
+
self.data_info['trajectory'] = {
|
|
138
|
+
'name': 'trajectory',
|
|
139
|
+
'loader': mdtraj_trajectory_loader,
|
|
140
|
+
'enabled': store_h5,
|
|
141
|
+
'filename': None,
|
|
142
|
+
'dir': True,
|
|
143
|
+
}
|
|
144
|
+
self.data_info['restart'] = {
|
|
145
|
+
'name': 'restart',
|
|
146
|
+
'loader': restart_loader,
|
|
147
|
+
'enabled': store_h5,
|
|
148
|
+
'filename': None,
|
|
149
|
+
'dir': True,
|
|
150
|
+
}
|
|
151
|
+
self.data_info['log'] = {'name': 'seglog', 'loader': seglog_loader, 'enabled': store_h5, 'filename': None, 'dir': False}
|
|
152
|
+
|
|
153
|
+
# Grab config from west.executable.datasets, else fallback to west.data.datasets.
|
|
154
|
+
dataset_configs = config.get(["west", "executable", "datasets"]) or config.get(['west', 'data', 'datasets'], {})
|
|
155
|
+
for dsinfo in dataset_configs:
|
|
156
|
+
try:
|
|
157
|
+
dsname = dsinfo['name']
|
|
158
|
+
except KeyError:
|
|
159
|
+
raise ValueError('dataset specifications require a ``name`` field')
|
|
160
|
+
|
|
161
|
+
if dsname == 'pcoord':
|
|
162
|
+
# can never disable pcoord collection
|
|
163
|
+
dsinfo['enabled'] = True
|
|
164
|
+
else:
|
|
165
|
+
check_bool(dsinfo.setdefault('enabled', True))
|
|
166
|
+
|
|
167
|
+
loader_directive = dsinfo.get('loader', None)
|
|
168
|
+
if 'module_path' in dsinfo:
|
|
169
|
+
dspath = self.makepath(dsinfo['module_path'])
|
|
170
|
+
else:
|
|
171
|
+
dspath = None
|
|
172
|
+
|
|
173
|
+
match loader_directive:
|
|
174
|
+
case loader_directive if callable(loader_directive):
|
|
175
|
+
# If directly callable, then use it
|
|
176
|
+
loader = loader_directive
|
|
177
|
+
case 'pcoord' | 'seglog' | 'restart':
|
|
178
|
+
# These are "protected" dataset names
|
|
179
|
+
if loader_directive in data_loaders:
|
|
180
|
+
loader = data_loaders[loader_directive]
|
|
181
|
+
else:
|
|
182
|
+
loader = get_object(loader_directive)
|
|
183
|
+
case 'trajectory':
|
|
184
|
+
# Special dataset for saving trajectory coordinates in HDF5 Framework
|
|
185
|
+
if loader_directive in trajectory_loaders:
|
|
186
|
+
loader = trajectory_loaders[loader_directive]
|
|
187
|
+
else:
|
|
188
|
+
loader = get_object(loader_directive, path=dspath)
|
|
189
|
+
case _:
|
|
190
|
+
# All other dataset names
|
|
191
|
+
if loader_directive in data_loaders:
|
|
192
|
+
loader = data_loaders[loader_directive]
|
|
193
|
+
elif isinstance(loader_directive, str):
|
|
194
|
+
loader = get_object(loader_directive, path=dspath)
|
|
195
|
+
else:
|
|
196
|
+
# Assumed aux dataset, defaulting to aux_data_loader
|
|
197
|
+
loader = aux_data_loader
|
|
198
|
+
|
|
199
|
+
if loader:
|
|
200
|
+
dsinfo['loader'] = loader
|
|
201
|
+
self.data_info.setdefault(dsname, {}).update(dsinfo)
|
|
202
|
+
|
|
203
|
+
log.debug('data_info: {!r}'.format(self.data_info))
|
|
204
|
+
|
|
205
|
+
def random_val_env_vars(self):
|
|
206
|
+
'''Return a set of environment variables containing random seeds. These are returned
|
|
207
|
+
as a dictionary, suitable for use in ``os.environ.update()`` or as the ``env`` argument to
|
|
208
|
+
``subprocess.Popen()``. Every child process executed by ``exec_child()`` gets these.'''
|
|
209
|
+
|
|
210
|
+
return {
|
|
211
|
+
self.ENV_RAND16: str(self.rng.integers(2**16, dtype=np.uint16)),
|
|
212
|
+
self.ENV_RAND32: str(self.rng.integers(2**32, dtype=np.uint32)),
|
|
213
|
+
self.ENV_RAND64: str(self.rng.integers(2**64, dtype=np.uint64)),
|
|
214
|
+
self.ENV_RAND128: str(int(self.rng.integers(2**64, dtype=np.uint64)) + int(self.rng.integers(2**64, dtype=np.uint64))),
|
|
215
|
+
self.ENV_RANDFLOAT: str(self.rng.random()),
|
|
216
|
+
}
|
|
217
|
+
|
|
218
|
+
def exec_child(self, executable, environ=None, stdin=None, stdout=None, stderr=None, cwd=None):
|
|
219
|
+
'''Execute a child process with the environment set from the current environment, the
|
|
220
|
+
values of self.addtl_child_environ, the random numbers returned by self.random_val_env_vars, and
|
|
221
|
+
the given ``environ`` (applied in that order). stdin/stdout/stderr are optionally redirected.
|
|
222
|
+
|
|
223
|
+
This function waits on the child process to finish, then returns
|
|
224
|
+
(rc, rusage), where rc is the child's return code and rusage is the resource usage tuple from os.wait4()'''
|
|
225
|
+
|
|
226
|
+
all_environ = dict(os.environ)
|
|
227
|
+
all_environ.update(self.addtl_child_environ)
|
|
228
|
+
all_environ.update(self.random_val_env_vars())
|
|
229
|
+
all_environ.update(environ or {})
|
|
230
|
+
|
|
231
|
+
stdin = open(stdin, 'rb') if stdin else sys.stdin
|
|
232
|
+
stdout = open(stdout, 'wb') if stdout else sys.stdout
|
|
233
|
+
if stderr == 'stdout':
|
|
234
|
+
stderr = stdout
|
|
235
|
+
else:
|
|
236
|
+
stderr = open(stderr, 'wb') if stderr else sys.stderr
|
|
237
|
+
|
|
238
|
+
# close_fds is critical for preventing out-of-file errors
|
|
239
|
+
proc = subprocess.Popen(
|
|
240
|
+
[executable],
|
|
241
|
+
cwd=cwd,
|
|
242
|
+
stdin=stdin,
|
|
243
|
+
stdout=stdout,
|
|
244
|
+
stderr=stderr if stderr != stdout else subprocess.STDOUT,
|
|
245
|
+
close_fds=True,
|
|
246
|
+
env=all_environ,
|
|
247
|
+
)
|
|
248
|
+
|
|
249
|
+
# Wait on child and get resource usage
|
|
250
|
+
(_pid, _status, rusage) = os.wait4(proc.pid, 0)
|
|
251
|
+
# Do a subprocess.Popen.wait() to let the Popen instance (and subprocess module) know that
|
|
252
|
+
# we are done with the process, and to get a more friendly return code
|
|
253
|
+
rc = proc.wait()
|
|
254
|
+
|
|
255
|
+
return (rc, rusage)
|
|
256
|
+
|
|
257
|
+
def exec_child_from_child_info(self, child_info, template_args, environ):
|
|
258
|
+
for key, value in child_info.get('environ', {}).items():
|
|
259
|
+
environ[key] = self.makepath(value)
|
|
260
|
+
return self.exec_child(
|
|
261
|
+
executable=self.makepath(child_info['executable'], template_args),
|
|
262
|
+
environ=environ,
|
|
263
|
+
cwd=self.makepath(child_info['cwd'], template_args) if child_info['cwd'] else None,
|
|
264
|
+
stdin=self.makepath(child_info['stdin'], template_args) if child_info['stdin'] else os.devnull,
|
|
265
|
+
stdout=self.makepath(child_info['stdout'], template_args) if child_info['stdout'] else None,
|
|
266
|
+
stderr=self.makepath(child_info['stderr'], template_args) if child_info['stderr'] else None,
|
|
267
|
+
)
|
|
268
|
+
|
|
269
|
+
# Functions to create template arguments and environment values for child processes
|
|
270
|
+
def update_args_env_basis_state(self, template_args, environ, basis_state):
|
|
271
|
+
new_template_args = {'basis_state': basis_state}
|
|
272
|
+
new_env = {
|
|
273
|
+
self.ENV_BSTATE_ID: str(basis_state.state_id if basis_state.state_id is not None else -1),
|
|
274
|
+
self.ENV_BSTATE_DATA_REF: self.makepath(self.basis_state_ref_template, new_template_args),
|
|
275
|
+
}
|
|
276
|
+
template_args.update(new_template_args)
|
|
277
|
+
environ.update(new_env)
|
|
278
|
+
return template_args, environ
|
|
279
|
+
|
|
280
|
+
def update_args_env_initial_state(self, template_args, environ, initial_state):
|
|
281
|
+
new_template_args = {'initial_state': initial_state}
|
|
282
|
+
new_env = {
|
|
283
|
+
self.ENV_ISTATE_ID: str(initial_state.state_id if initial_state.state_id is not None else -1),
|
|
284
|
+
self.ENV_ISTATE_DATA_REF: self.makepath(self.initial_state_ref_template, new_template_args),
|
|
285
|
+
}
|
|
286
|
+
|
|
287
|
+
if initial_state.basis_state is not None:
|
|
288
|
+
basis_state = initial_state.basis_state
|
|
289
|
+
elif initial_state.istate_type == InitialState.ISTATE_TYPE_START:
|
|
290
|
+
basis_state = BasisState(
|
|
291
|
+
label=f"sstate_{initial_state.state_id}", pcoord=initial_state.pcoord, probability=0.0, auxref=""
|
|
292
|
+
)
|
|
293
|
+
else:
|
|
294
|
+
basis_state = self.basis_states[initial_state.basis_state_id]
|
|
295
|
+
|
|
296
|
+
self.update_args_env_basis_state(new_template_args, new_env, basis_state)
|
|
297
|
+
|
|
298
|
+
template_args.update(new_template_args)
|
|
299
|
+
environ.update(new_env)
|
|
300
|
+
return template_args, environ
|
|
301
|
+
|
|
302
|
+
def update_args_env_iter(self, template_args, environ, n_iter):
|
|
303
|
+
environ[self.ENV_CURRENT_ITER] = str(n_iter if n_iter is not None else -1)
|
|
304
|
+
template_args['n_iter'] = int(n_iter)
|
|
305
|
+
return template_args, n_iter
|
|
306
|
+
|
|
307
|
+
def update_args_env_segment(self, template_args, environ, segment):
|
|
308
|
+
template_args['segment'] = segment
|
|
309
|
+
|
|
310
|
+
environ[self.ENV_CURRENT_SEG_INITPOINT] = Segment.initpoint_type_names[segment.initpoint_type]
|
|
311
|
+
|
|
312
|
+
if segment.initpoint_type == Segment.SEG_INITPOINT_CONTINUES:
|
|
313
|
+
# Could use actual parent object here if the work manager cared to pass that much data
|
|
314
|
+
# to us (we'd need at least the subset of parents for all segments sent in the call to propagate)
|
|
315
|
+
# that may make a good west.cfg option for future crazy extensibility, but for now,
|
|
316
|
+
# just populate the bare minimum
|
|
317
|
+
parent = Segment(n_iter=segment.n_iter - 1, seg_id=segment.parent_id)
|
|
318
|
+
parent_template_args = dict(template_args)
|
|
319
|
+
parent_template_args['segment'] = parent
|
|
320
|
+
|
|
321
|
+
environ[self.ENV_PARENT_SEG_ID] = str(segment.parent_id if segment.parent_id is not None else -1)
|
|
322
|
+
environ[self.ENV_PARENT_DATA_REF] = self.makepath(self.segment_ref_template, parent_template_args)
|
|
323
|
+
elif segment.initpoint_type == Segment.SEG_INITPOINT_NEWTRAJ:
|
|
324
|
+
# This segment is initiated from a basis state; WEST_PARENT_SEG_ID and WEST_PARENT_DATA_REF are
|
|
325
|
+
# set to the basis state ID and data ref
|
|
326
|
+
initial_state = self.initial_states[segment.initial_state_id]
|
|
327
|
+
|
|
328
|
+
if initial_state.istate_type == InitialState.ISTATE_TYPE_START:
|
|
329
|
+
basis_state = BasisState(
|
|
330
|
+
label=f"sstate_{initial_state.state_id}", pcoord=initial_state.pcoord, probability=0.0, auxref=""
|
|
331
|
+
)
|
|
332
|
+
|
|
333
|
+
else:
|
|
334
|
+
basis_state = self.basis_states[initial_state.basis_state_id]
|
|
335
|
+
|
|
336
|
+
if self.ENV_BSTATE_ID not in environ:
|
|
337
|
+
self.update_args_env_basis_state(template_args, environ, basis_state)
|
|
338
|
+
if self.ENV_ISTATE_ID not in environ:
|
|
339
|
+
self.update_args_env_initial_state(template_args, environ, initial_state)
|
|
340
|
+
|
|
341
|
+
assert initial_state.istate_type in (
|
|
342
|
+
InitialState.ISTATE_TYPE_BASIS,
|
|
343
|
+
InitialState.ISTATE_TYPE_GENERATED,
|
|
344
|
+
InitialState.ISTATE_TYPE_START,
|
|
345
|
+
)
|
|
346
|
+
if initial_state.istate_type == InitialState.ISTATE_TYPE_BASIS:
|
|
347
|
+
environ[self.ENV_PARENT_DATA_REF] = environ[self.ENV_BSTATE_DATA_REF]
|
|
348
|
+
|
|
349
|
+
elif initial_state.istate_type == InitialState.ISTATE_TYPE_START:
|
|
350
|
+
# This points to the start-state PDB
|
|
351
|
+
environ[self.ENV_PARENT_DATA_REF] = environ[self.ENV_BSTATE_DATA_REF] + '/' + initial_state.basis_auxref
|
|
352
|
+
else: # initial_state.type == InitialState.ISTATE_TYPE_GENERATED
|
|
353
|
+
environ[self.ENV_PARENT_DATA_REF] = environ[self.ENV_ISTATE_DATA_REF]
|
|
354
|
+
|
|
355
|
+
environ[self.ENV_CURRENT_SEG_ID] = str(segment.seg_id if segment.seg_id is not None else -1)
|
|
356
|
+
environ[self.ENV_CURRENT_SEG_DATA_REF] = self.makepath(self.segment_ref_template, template_args)
|
|
357
|
+
return template_args, environ
|
|
358
|
+
|
|
359
|
+
def template_args_for_segment(self, segment):
|
|
360
|
+
template_args, environ = {}, {}
|
|
361
|
+
self.update_args_env_iter(template_args, environ, segment.n_iter)
|
|
362
|
+
self.update_args_env_segment(template_args, environ, segment)
|
|
363
|
+
return template_args
|
|
364
|
+
|
|
365
|
+
def exec_for_segment(self, child_info, segment, addtl_env=None):
|
|
366
|
+
'''Execute a child process with environment and template expansion from the given
|
|
367
|
+
segment.'''
|
|
368
|
+
template_args, environ = {}, {}
|
|
369
|
+
self.update_args_env_iter(template_args, environ, segment.n_iter)
|
|
370
|
+
self.update_args_env_segment(template_args, environ, segment)
|
|
371
|
+
environ.update(addtl_env or {})
|
|
372
|
+
self.prepare_file_system(segment, environ)
|
|
373
|
+
child_info['cwd'] = environ[self.ENV_CURRENT_SEG_DATA_REF]
|
|
374
|
+
return self.exec_child_from_child_info(child_info, template_args, environ)
|
|
375
|
+
|
|
376
|
+
def exec_for_iteration(self, child_info, n_iter, addtl_env=None):
|
|
377
|
+
'''Execute a child process with environment and template expansion from the given
|
|
378
|
+
iteration number.'''
|
|
379
|
+
template_args, environ = {}, {}
|
|
380
|
+
self.update_args_env_iter(template_args, environ, n_iter)
|
|
381
|
+
environ.update(addtl_env or {})
|
|
382
|
+
return self.exec_child_from_child_info(child_info, template_args, environ)
|
|
383
|
+
|
|
384
|
+
def exec_for_basis_state(self, child_info, basis_state, addtl_env=None):
|
|
385
|
+
'''Execute a child process with environment and template expansion from the
|
|
386
|
+
given basis state'''
|
|
387
|
+
template_args, environ = {}, {}
|
|
388
|
+
self.update_args_env_basis_state(template_args, environ, basis_state)
|
|
389
|
+
environ.update(addtl_env or {})
|
|
390
|
+
return self.exec_child_from_child_info(child_info, template_args, environ)
|
|
391
|
+
|
|
392
|
+
def exec_for_initial_state(self, child_info, initial_state, addtl_env=None):
|
|
393
|
+
'''Execute a child process with environment and template expansion from the given
|
|
394
|
+
initial state.'''
|
|
395
|
+
template_args, environ = {}, {}
|
|
396
|
+
self.update_args_env_initial_state(template_args, environ, initial_state)
|
|
397
|
+
environ.update(addtl_env or {})
|
|
398
|
+
return self.exec_child_from_child_info(child_info, template_args, environ)
|
|
399
|
+
|
|
400
|
+
def prepare_file_system(self, segment, environ):
|
|
401
|
+
try:
|
|
402
|
+
# If the filesystem is properly clean.
|
|
403
|
+
os.makedirs(environ[self.ENV_CURRENT_SEG_DATA_REF])
|
|
404
|
+
except Exception:
|
|
405
|
+
# If the filesystem is NOT properly clean.
|
|
406
|
+
shutil.rmtree(environ[self.ENV_CURRENT_SEG_DATA_REF])
|
|
407
|
+
os.makedirs(environ[self.ENV_CURRENT_SEG_DATA_REF])
|
|
408
|
+
if self.data_info['restart']['enabled']:
|
|
409
|
+
restart_writer(environ[self.ENV_CURRENT_SEG_DATA_REF], segment=segment)
|
|
410
|
+
|
|
411
|
+
def setup_dataset_return(self, segment=None, subset_keys=None):
|
|
412
|
+
'''Set up temporary files and environment variables that point to them for segment
|
|
413
|
+
runners to return data. ``segment`` is the ``Segment`` object that the return data
|
|
414
|
+
is associated with. ``subset_keys`` specifies the names of a subset of data to be
|
|
415
|
+
returned.'''
|
|
416
|
+
if subset_keys is None:
|
|
417
|
+
subset_keys = self.data_info.keys()
|
|
418
|
+
|
|
419
|
+
addtl_env = {}
|
|
420
|
+
return_files = {}
|
|
421
|
+
del_return_files = {}
|
|
422
|
+
|
|
423
|
+
for dataset in self.data_info:
|
|
424
|
+
if dataset not in subset_keys:
|
|
425
|
+
continue
|
|
426
|
+
|
|
427
|
+
if not self.data_info[dataset].get('enabled', False):
|
|
428
|
+
continue
|
|
429
|
+
|
|
430
|
+
return_template = self.data_info[dataset].get('filename')
|
|
431
|
+
if return_template:
|
|
432
|
+
if segment is None:
|
|
433
|
+
raise ValueError('segment needs to be provided for dataset return')
|
|
434
|
+
return_files[dataset] = self.makepath(return_template, self.template_args_for_segment(segment))
|
|
435
|
+
del_return_files[dataset] = False
|
|
436
|
+
else:
|
|
437
|
+
isdir = self.data_info[dataset].get('dir', False)
|
|
438
|
+
if isdir:
|
|
439
|
+
rfname = tempfile.mkdtemp()
|
|
440
|
+
else:
|
|
441
|
+
(fd, rfname) = tempfile.mkstemp()
|
|
442
|
+
os.close(fd)
|
|
443
|
+
return_files[dataset] = rfname
|
|
444
|
+
del_return_files[dataset] = True
|
|
445
|
+
|
|
446
|
+
addtl_env['WEST_{}_RETURN'.format(dataset.upper())] = return_files[dataset]
|
|
447
|
+
|
|
448
|
+
return addtl_env, return_files, del_return_files
|
|
449
|
+
|
|
450
|
+
def retrieve_dataset_return(self, state, return_files, del_return_files, single_point):
|
|
451
|
+
'''Retrieve returned data from the temporary locations directed by the environment variables.
|
|
452
|
+
``state`` is a ``Segment``, ``BasisState`` , or ``InitialState``object that the return data is
|
|
453
|
+
associated with. ``return_files`` is a ``dict`` where the keys are the dataset names and
|
|
454
|
+
the values are the paths to the temporarily files that contain the returned data.
|
|
455
|
+
``del_return_files`` is a ``dict`` where the keys are the names of datasets to be deleted
|
|
456
|
+
(if the corresponding value is set to ``True``) once the data is retrieved.'''
|
|
457
|
+
|
|
458
|
+
state_name, state_id = return_state_type(state)
|
|
459
|
+
|
|
460
|
+
for dataset in self.data_info:
|
|
461
|
+
if dataset not in return_files:
|
|
462
|
+
continue
|
|
463
|
+
|
|
464
|
+
# pcoord is always enabled (see __init__)
|
|
465
|
+
if not self.data_info[dataset].get('enabled', False):
|
|
466
|
+
continue
|
|
467
|
+
|
|
468
|
+
filename = return_files[dataset]
|
|
469
|
+
loader = self.data_info[dataset]['loader']
|
|
470
|
+
try:
|
|
471
|
+
loader(dataset, filename, state, single_point=single_point)
|
|
472
|
+
except Exception as e:
|
|
473
|
+
log.error('could not read {} for {} {} from {!r}: {!r}'.format(dataset, state_name, state_id, filename, e))
|
|
474
|
+
if isinstance(state, Segment):
|
|
475
|
+
state.status = state.SEG_STATUS_FAILED
|
|
476
|
+
break
|
|
477
|
+
else:
|
|
478
|
+
if del_return_files.get(dataset, False):
|
|
479
|
+
try:
|
|
480
|
+
if os.path.isfile(filename):
|
|
481
|
+
os.unlink(filename)
|
|
482
|
+
else:
|
|
483
|
+
shutil.rmtree(filename)
|
|
484
|
+
except Exception as e:
|
|
485
|
+
log.warning(
|
|
486
|
+
'could not delete {} file {!r} for {} {}: {!r}'.format(dataset, filename, state_name, state_id, e)
|
|
487
|
+
)
|
|
488
|
+
else:
|
|
489
|
+
log.debug('deleted {} file {!r} for {} {}'.format(dataset, filename, state_name, state_id))
|
|
490
|
+
|
|
491
|
+
# Specific functions required by the WEST framework
|
|
492
|
+
def get_pcoord(self, state):
|
|
493
|
+
'''Get the progress coordinate of the given basis or initial state.'''
|
|
494
|
+
|
|
495
|
+
template_args, environ = {}, {}
|
|
496
|
+
|
|
497
|
+
if isinstance(state, BasisState):
|
|
498
|
+
execfn = self.exec_for_basis_state
|
|
499
|
+
self.update_args_env_basis_state(template_args, environ, state)
|
|
500
|
+
struct_ref = environ[self.ENV_BSTATE_DATA_REF]
|
|
501
|
+
elif isinstance(state, InitialState):
|
|
502
|
+
execfn = self.exec_for_initial_state
|
|
503
|
+
self.update_args_env_initial_state(template_args, environ, state)
|
|
504
|
+
struct_ref = environ[self.ENV_ISTATE_DATA_REF]
|
|
505
|
+
else:
|
|
506
|
+
raise TypeError('state must be a BasisState or InitialState')
|
|
507
|
+
|
|
508
|
+
child_info = self.exe_info.get('get_pcoord')
|
|
509
|
+
addtl_env, return_files, del_return_files = self.setup_dataset_return(
|
|
510
|
+
subset_keys=['pcoord', 'trajectory', 'restart', 'log']
|
|
511
|
+
)
|
|
512
|
+
addtl_env[self.ENV_STRUCT_DATA_REF] = struct_ref
|
|
513
|
+
|
|
514
|
+
rc, rusage = execfn(child_info, state, addtl_env)
|
|
515
|
+
if rc != 0:
|
|
516
|
+
log.error('get_pcoord executable {!r} returned {}'.format(child_info['executable'], rc))
|
|
517
|
+
|
|
518
|
+
self.retrieve_dataset_return(state, return_files, del_return_files, True)
|
|
519
|
+
|
|
520
|
+
def gen_istate(self, basis_state, initial_state):
|
|
521
|
+
'''Generate a new initial state from the given basis state.'''
|
|
522
|
+
child_info = self.exe_info.get('gen_istate')
|
|
523
|
+
rc, rusage = self.exec_for_initial_state(child_info, initial_state)
|
|
524
|
+
if rc != 0:
|
|
525
|
+
log.error('gen_istate executable {!r} returned {}'.format(child_info['executable'], rc))
|
|
526
|
+
initial_state.istate_status = InitialState.ISTATE_STATUS_FAILED
|
|
527
|
+
return
|
|
528
|
+
|
|
529
|
+
# Determine and load the progress coordinate value for this state
|
|
530
|
+
try:
|
|
531
|
+
self.get_pcoord(initial_state)
|
|
532
|
+
except Exception:
|
|
533
|
+
log.exception('could not get progress coordinate for initial state {!r}'.format(initial_state))
|
|
534
|
+
initial_state.istate_status = InitialState.ISTATE_STATUS_FAILED
|
|
535
|
+
raise
|
|
536
|
+
else:
|
|
537
|
+
initial_state.istate_status = InitialState.ISTATE_STATUS_PREPARED
|
|
538
|
+
|
|
539
|
+
def prepare_iteration(self, n_iter, segments):
|
|
540
|
+
child_info = self.exe_info.get('pre_iteration')
|
|
541
|
+
if child_info and child_info['enabled']:
|
|
542
|
+
try:
|
|
543
|
+
rc, rusage = self.exec_for_iteration(child_info, n_iter)
|
|
544
|
+
except OSError as e:
|
|
545
|
+
log.warning('could not execute pre-iteration program {!r}: {}'.format(child_info['executable'], e))
|
|
546
|
+
else:
|
|
547
|
+
if rc != 0:
|
|
548
|
+
log.warning('pre-iteration executable {!r} returned {}'.format(child_info['executable'], rc))
|
|
549
|
+
|
|
550
|
+
def finalize_iteration(self, n_iter, segments):
|
|
551
|
+
child_info = self.exe_info.get('post_iteration')
|
|
552
|
+
if child_info and child_info['enabled']:
|
|
553
|
+
try:
|
|
554
|
+
rc, rusage = self.exec_for_iteration(child_info, n_iter)
|
|
555
|
+
except OSError as e:
|
|
556
|
+
log.warning('could not execute post-iteration program {!r}: {}'.format(child_info['executable'], e))
|
|
557
|
+
else:
|
|
558
|
+
if rc != 0:
|
|
559
|
+
log.warning('post-iteration executable {!r} returned {}'.format(child_info['executable'], rc))
|
|
560
|
+
|
|
561
|
+
def propagate(self, segments):
|
|
562
|
+
child_info = self.exe_info['propagator']
|
|
563
|
+
|
|
564
|
+
for segment in segments:
|
|
565
|
+
starttime = time.time()
|
|
566
|
+
|
|
567
|
+
addtl_env, return_files, del_return_files = self.setup_dataset_return(segment)
|
|
568
|
+
|
|
569
|
+
# Spawn propagator and wait for its completion
|
|
570
|
+
rc, rusage = self.exec_for_segment(child_info, segment, addtl_env)
|
|
571
|
+
|
|
572
|
+
if rc == 0:
|
|
573
|
+
segment.status = Segment.SEG_STATUS_COMPLETE
|
|
574
|
+
elif rc < 0:
|
|
575
|
+
log.error('child process for segment %d exited on signal %d (%s)' % (segment.seg_id, -rc, SIGNAL_NAMES[-rc]))
|
|
576
|
+
segment.status = Segment.SEG_STATUS_FAILED
|
|
577
|
+
continue
|
|
578
|
+
else:
|
|
579
|
+
log.error('child process for segment %d exited with code %d' % (segment.seg_id, rc))
|
|
580
|
+
segment.status = Segment.SEG_STATUS_FAILED
|
|
581
|
+
continue
|
|
582
|
+
|
|
583
|
+
# Extract data and store on segment for recording in the master thread/process/node
|
|
584
|
+
self.retrieve_dataset_return(segment, return_files, del_return_files, False)
|
|
585
|
+
|
|
586
|
+
if segment.status == Segment.SEG_STATUS_FAILED:
|
|
587
|
+
continue
|
|
588
|
+
|
|
589
|
+
# Record timing info
|
|
590
|
+
segment.walltime = time.time() - starttime
|
|
591
|
+
segment.cputime = rusage.ru_utime
|
|
592
|
+
return segments
|