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.
Files changed (162) hide show
  1. westpa/__init__.py +14 -0
  2. westpa/_version.py +21 -0
  3. westpa/analysis/__init__.py +5 -0
  4. westpa/analysis/core.py +749 -0
  5. westpa/analysis/statistics.py +27 -0
  6. westpa/analysis/trajectories.py +369 -0
  7. westpa/cli/__init__.py +0 -0
  8. westpa/cli/core/__init__.py +0 -0
  9. westpa/cli/core/w_fork.py +152 -0
  10. westpa/cli/core/w_init.py +230 -0
  11. westpa/cli/core/w_run.py +77 -0
  12. westpa/cli/core/w_states.py +212 -0
  13. westpa/cli/core/w_succ.py +99 -0
  14. westpa/cli/core/w_truncate.py +68 -0
  15. westpa/cli/tools/__init__.py +0 -0
  16. westpa/cli/tools/ploterr.py +506 -0
  17. westpa/cli/tools/plothist.py +706 -0
  18. westpa/cli/tools/w_assign.py +597 -0
  19. westpa/cli/tools/w_bins.py +166 -0
  20. westpa/cli/tools/w_crawl.py +119 -0
  21. westpa/cli/tools/w_direct.py +557 -0
  22. westpa/cli/tools/w_dumpsegs.py +94 -0
  23. westpa/cli/tools/w_eddist.py +506 -0
  24. westpa/cli/tools/w_fluxanl.py +376 -0
  25. westpa/cli/tools/w_ipa.py +832 -0
  26. westpa/cli/tools/w_kinavg.py +127 -0
  27. westpa/cli/tools/w_kinetics.py +96 -0
  28. westpa/cli/tools/w_multi_west.py +414 -0
  29. westpa/cli/tools/w_ntop.py +213 -0
  30. westpa/cli/tools/w_pdist.py +515 -0
  31. westpa/cli/tools/w_postanalysis_matrix.py +82 -0
  32. westpa/cli/tools/w_postanalysis_reweight.py +53 -0
  33. westpa/cli/tools/w_red.py +491 -0
  34. westpa/cli/tools/w_reweight.py +780 -0
  35. westpa/cli/tools/w_select.py +226 -0
  36. westpa/cli/tools/w_stateprobs.py +111 -0
  37. westpa/cli/tools/w_timings.py +113 -0
  38. westpa/cli/tools/w_trace.py +599 -0
  39. westpa/core/__init__.py +0 -0
  40. westpa/core/_rc.py +673 -0
  41. westpa/core/binning/__init__.py +55 -0
  42. westpa/core/binning/_assign.c +36018 -0
  43. westpa/core/binning/_assign.cpython-312-aarch64-linux-gnu.so +0 -0
  44. westpa/core/binning/_assign.pyx +370 -0
  45. westpa/core/binning/assign.py +454 -0
  46. westpa/core/binning/binless.py +96 -0
  47. westpa/core/binning/binless_driver.py +54 -0
  48. westpa/core/binning/binless_manager.py +189 -0
  49. westpa/core/binning/bins.py +47 -0
  50. westpa/core/binning/mab.py +506 -0
  51. westpa/core/binning/mab_driver.py +54 -0
  52. westpa/core/binning/mab_manager.py +197 -0
  53. westpa/core/data_manager.py +1761 -0
  54. westpa/core/extloader.py +74 -0
  55. westpa/core/h5io.py +1079 -0
  56. westpa/core/kinetics/__init__.py +24 -0
  57. westpa/core/kinetics/_kinetics.c +45174 -0
  58. westpa/core/kinetics/_kinetics.cpython-312-aarch64-linux-gnu.so +0 -0
  59. westpa/core/kinetics/_kinetics.pyx +815 -0
  60. westpa/core/kinetics/events.py +147 -0
  61. westpa/core/kinetics/matrates.py +156 -0
  62. westpa/core/kinetics/rate_averaging.py +266 -0
  63. westpa/core/progress.py +218 -0
  64. westpa/core/propagators/__init__.py +54 -0
  65. westpa/core/propagators/executable.py +592 -0
  66. westpa/core/propagators/loaders.py +196 -0
  67. westpa/core/reweight/__init__.py +14 -0
  68. westpa/core/reweight/_reweight.c +36899 -0
  69. westpa/core/reweight/_reweight.cpython-312-aarch64-linux-gnu.so +0 -0
  70. westpa/core/reweight/_reweight.pyx +439 -0
  71. westpa/core/reweight/matrix.py +126 -0
  72. westpa/core/segment.py +119 -0
  73. westpa/core/sim_manager.py +839 -0
  74. westpa/core/states.py +359 -0
  75. westpa/core/systems.py +93 -0
  76. westpa/core/textio.py +74 -0
  77. westpa/core/trajectory.py +603 -0
  78. westpa/core/we_driver.py +910 -0
  79. westpa/core/wm_ops.py +43 -0
  80. westpa/core/yamlcfg.py +298 -0
  81. westpa/fasthist/__init__.py +34 -0
  82. westpa/fasthist/_fasthist.c +38755 -0
  83. westpa/fasthist/_fasthist.cpython-312-aarch64-linux-gnu.so +0 -0
  84. westpa/fasthist/_fasthist.pyx +222 -0
  85. westpa/mclib/__init__.py +271 -0
  86. westpa/mclib/__main__.py +28 -0
  87. westpa/mclib/_mclib.c +34610 -0
  88. westpa/mclib/_mclib.cpython-312-aarch64-linux-gnu.so +0 -0
  89. westpa/mclib/_mclib.pyx +226 -0
  90. westpa/oldtools/__init__.py +4 -0
  91. westpa/oldtools/aframe/__init__.py +35 -0
  92. westpa/oldtools/aframe/atool.py +75 -0
  93. westpa/oldtools/aframe/base_mixin.py +26 -0
  94. westpa/oldtools/aframe/binning.py +178 -0
  95. westpa/oldtools/aframe/data_reader.py +560 -0
  96. westpa/oldtools/aframe/iter_range.py +200 -0
  97. westpa/oldtools/aframe/kinetics.py +117 -0
  98. westpa/oldtools/aframe/mcbs.py +153 -0
  99. westpa/oldtools/aframe/output.py +39 -0
  100. westpa/oldtools/aframe/plotting.py +88 -0
  101. westpa/oldtools/aframe/trajwalker.py +126 -0
  102. westpa/oldtools/aframe/transitions.py +469 -0
  103. westpa/oldtools/cmds/__init__.py +0 -0
  104. westpa/oldtools/cmds/w_ttimes.py +361 -0
  105. westpa/oldtools/files.py +34 -0
  106. westpa/oldtools/miscfn.py +23 -0
  107. westpa/oldtools/stats/__init__.py +4 -0
  108. westpa/oldtools/stats/accumulator.py +35 -0
  109. westpa/oldtools/stats/edfs.py +129 -0
  110. westpa/oldtools/stats/mcbs.py +96 -0
  111. westpa/tools/__init__.py +33 -0
  112. westpa/tools/binning.py +472 -0
  113. westpa/tools/core.py +340 -0
  114. westpa/tools/data_reader.py +159 -0
  115. westpa/tools/dtypes.py +31 -0
  116. westpa/tools/iter_range.py +198 -0
  117. westpa/tools/kinetics_tool.py +343 -0
  118. westpa/tools/plot.py +283 -0
  119. westpa/tools/progress.py +17 -0
  120. westpa/tools/selected_segs.py +154 -0
  121. westpa/tools/wipi.py +751 -0
  122. westpa/trajtree/__init__.py +4 -0
  123. westpa/trajtree/_trajtree.c +17829 -0
  124. westpa/trajtree/_trajtree.cpython-312-aarch64-linux-gnu.so +0 -0
  125. westpa/trajtree/_trajtree.pyx +130 -0
  126. westpa/trajtree/trajtree.py +117 -0
  127. westpa/westext/__init__.py +0 -0
  128. westpa/westext/adaptvoronoi/__init__.py +3 -0
  129. westpa/westext/adaptvoronoi/adaptVor_driver.py +214 -0
  130. westpa/westext/hamsm_restarting/__init__.py +3 -0
  131. westpa/westext/hamsm_restarting/example_overrides.py +35 -0
  132. westpa/westext/hamsm_restarting/restart_driver.py +1165 -0
  133. westpa/westext/stringmethod/__init__.py +11 -0
  134. westpa/westext/stringmethod/fourier_fitting.py +69 -0
  135. westpa/westext/stringmethod/string_driver.py +253 -0
  136. westpa/westext/stringmethod/string_method.py +306 -0
  137. westpa/westext/weed/BinCluster.py +180 -0
  138. westpa/westext/weed/ProbAdjustEquil.py +100 -0
  139. westpa/westext/weed/UncertMath.py +247 -0
  140. westpa/westext/weed/__init__.py +10 -0
  141. westpa/westext/weed/weed_driver.py +192 -0
  142. westpa/westext/wess/ProbAdjust.py +101 -0
  143. westpa/westext/wess/__init__.py +6 -0
  144. westpa/westext/wess/wess_driver.py +217 -0
  145. westpa/work_managers/__init__.py +57 -0
  146. westpa/work_managers/core.py +396 -0
  147. westpa/work_managers/environment.py +134 -0
  148. westpa/work_managers/mpi.py +318 -0
  149. westpa/work_managers/processes.py +201 -0
  150. westpa/work_managers/serial.py +28 -0
  151. westpa/work_managers/threads.py +79 -0
  152. westpa/work_managers/zeromq/__init__.py +20 -0
  153. westpa/work_managers/zeromq/core.py +635 -0
  154. westpa/work_managers/zeromq/node.py +131 -0
  155. westpa/work_managers/zeromq/work_manager.py +526 -0
  156. westpa/work_managers/zeromq/worker.py +320 -0
  157. westpa-2022.13.dist-info/METADATA +179 -0
  158. westpa-2022.13.dist-info/RECORD +162 -0
  159. westpa-2022.13.dist-info/WHEEL +7 -0
  160. westpa-2022.13.dist-info/entry_points.txt +30 -0
  161. westpa-2022.13.dist-info/licenses/LICENSE +21 -0
  162. 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