westpa 2022.10__cp312-cp312-macosx_11_0_arm64.whl

westpa/core/states.py

@@ -0,0 +1,359 @@
+import numpy as np
+
+from westpa.core.segment import Segment
+
+
+class BasisState:
+    '''Describes an basis (micro)state. These basis states are used to generate
+    initial states for new trajectories, either at the beginning of the simulation
+    (i.e. at w_init) or due to recycling.
+
+    :ivar state_id:     Integer identifier of this state, usually set by the
+                        data manager.
+    :ivar label:        A descriptive label for this microstate (may be empty)
+    :ivar probability:  Probability of this state to be selected when creating a
+                        new trajectory.
+    :ivar pcoord:       The representative progress coordinate of this state.
+
+    :ivar auxref:       A user-provided (string) reference for locating data associated
+                        with this state (usually a filesystem path).
+    '''
+
+    def __init__(self, label, probability, pcoord=None, auxref=None, state_id=None):
+        self.label = str(label, encoding="UTF-8") if type(label) is bytes else label
+        self.probability = probability
+        self.pcoord = np.atleast_1d(pcoord)
+        self.auxref = auxref
+        self.state_id = state_id
+        self.data = {}
+
+    def __repr__(self):
+        return '{} state_id={self.state_id!r} label={self.label!r} prob={self.probability!r} pcoord={self.pcoord!r}>'.format(
+            object.__repr__(self)[:-1], self=self
+        )
+
+    @classmethod
+    def states_to_file(cls, states, fileobj):
+        '''Write a file defining basis states, which may then be read by `states_from_file()`.'''
+
+        if isinstance(fileobj, str):
+            fileobj = open(fileobj, 'wt')
+
+        max_label_len = max(8, max(len(state.label or '') for state in states))
+        max_auxref_len = max(8, max(len(state.auxref or '') for state in states))
+        fmt = (
+            '{state.label:<{max_label_len}s}    {state.probability:12.7g}    {auxref_str:<{max_auxref_len}s}'
+            '    # state_id={state_id_str:s}    pcoord={pcoord_str}\n'
+        )
+        fileobj.write(
+            '# {:{max_label_len}s}    {:>12s}    {:{max_auxref_len}s}\n'.format(
+                'Label', 'Probability', 'Auxref', max_label_len=max_label_len - 2, max_auxref_len=max_auxref_len
+            )
+        )
+        for state in states:
+            state_id_str = str(state.state_id) if state.state_id is not None else 'None'
+            pcoord_str = str(list(state.pcoord))
+            auxref_str = 'None' if state.auxref == b'' or state.auxref == "b''" else str(state.auxref)
+            fileobj.write(
+                fmt.format(
+                    state=state,
+                    pcoord_str=pcoord_str,
+                    auxref_str=auxref_str,
+                    state_id_str=state_id_str,
+                    max_label_len=max_label_len,
+                    max_auxref_len=max_auxref_len,
+                )
+            )
+
+    @classmethod
+    def states_from_file(cls, statefile):
+        '''Read a file defining basis states.  Each line defines a state, and contains a label, the probability,
+        and optionally a data reference, separated by whitespace, as in::
+
+            unbound    1.0
+
+        or::
+
+            unbound_0    0.6        state0.pdb
+            unbound_1    0.4        state1.pdb
+
+        '''
+        states = []
+        lineno = 0
+
+        try:
+            open_statefile = open(statefile, 'rt')
+        except TypeError:
+            open_statefile = statefile
+
+        for line in open_statefile:
+            lineno += 1
+
+            # remove comment portion
+            line = line.partition('#')[0].strip()
+            if not line:
+                continue
+
+            fields = line.split()
+            label = fields[0]
+            try:
+                probability = float(fields[1])
+            except ValueError:
+                raise ValueError('invalid probability ({!r}) {} line {:d}'.format(fields[1], statefile, lineno))
+
+            try:
+                auxref = fields[2].strip()
+            except IndexError:
+                auxref = None
+
+            states.append(cls(state_id=None, probability=probability, label=label, auxref=auxref))
+
+        try:
+            open_statefile.close()
+        except Exception:
+            pass
+
+        return states
+
+    def as_numpy_record(self):
+        '''Return the data for this state as a numpy record array.'''
+
+        from westpa.core.data_manager import vstr_dtype, weight_dtype, seg_id_dtype
+
+        bstate_dtype = np.dtype(
+            [
+                ('state_id', seg_id_dtype),
+                ('probability', weight_dtype),
+                ('pcoord', self.pcoord.dtype, (len(self.pcoord),)),
+                ('label', vstr_dtype),
+                ('auxref', vstr_dtype),
+            ]
+        )
+        bstaterec = np.array(
+            [(self.state_id, self.probability, self.pcoord, self.label or '', self.auxref or '')], dtype=bstate_dtype
+        )[0]
+        return bstaterec
+
+
+class InitialState:
+    '''Describes an initial state for a new trajectory. These are generally constructed by
+    appropriate modification of a basis state.
+
+    :ivar state_id:         Integer identifier of this state, usually set by the
+                            data manager.
+    :ivar basis_state_id:   Identifier of the basis state from which this state was
+                            generated, or None.
+    :ivar basis_state:      The `BasisState` from which this state was generated, or None.
+    :ivar iter_created:     Iteration in which this state was generated (0 for
+                            simulation initialization).
+    :ivar iter_used:        Iteration in which this state was used to initiate a
+                            trajectory (None for unused).
+    :ivar istate_type:      Integer describing the type of this initial state
+                            (ISTATE_TYPE_BASIS for direct use of a basis state,
+                            ISTATE_TYPE_GENERATED for a state generated from a basis state,
+                            ISTATE_TYPE_RESTART for a state corresponding to the endpoint
+                            of a segment in another simulation, or
+                            ISTATE_TYPE_START for a state generated from a start state).
+    :ivar istate_status:    Integer describing whether this initial state has been properly
+                            prepared.
+    :ivar pcoord:           The representative progress coordinate of this state.
+    '''
+
+    ISTATE_TYPE_UNSET = 0
+    ISTATE_TYPE_BASIS = 1
+    ISTATE_TYPE_GENERATED = 2
+    ISTATE_TYPE_RESTART = 3
+    ISTATE_TYPE_START = 4
+
+    ISTATE_UNUSED = 0
+
+    ISTATE_STATUS_PENDING = 0
+    ISTATE_STATUS_PREPARED = 1
+    ISTATE_STATUS_FAILED = 2
+
+    istate_types = {}
+    istate_type_names = {}
+
+    istate_statuses = {}
+    istate_status_names = {}
+
+    def __init__(
+        self,
+        state_id,
+        basis_state_id,
+        iter_created,
+        iter_used=None,
+        istate_type=None,
+        istate_status=None,
+        pcoord=None,
+        basis_state=None,
+        basis_auxref=None,
+    ):
+        self.state_id = state_id
+        self.basis_state_id = basis_state_id
+        self.basis_state = basis_state
+        self.istate_type = istate_type
+        self.istate_status = istate_status
+        self.iter_created = iter_created
+        self.iter_used = iter_used
+        self.pcoord = np.atleast_1d(pcoord)
+        self.basis_auxref = basis_auxref
+        self.data = {}
+
+    def __repr__(self):
+        return '{} state_id={self.state_id!r} istate_type={self.istate_type!r} basis_state_id={self.basis_state_id!r} iter_created={self.iter_created!r} pcoord={self.pcoord!r}>'.format(
+            object.__repr__(self)[:-1], self=self
+        )
+
+    def as_numpy_record(self):
+        from westpa.core.data_manager import seg_id_dtype, istate_type_dtype, istate_status_dtype, vstr_dtype
+
+        istate_dtype = np.dtype(
+            [
+                ('state_id', seg_id_dtype),
+                ('basis_state_id', seg_id_dtype),
+                ('iter_created', np.uint),
+                ('iter_used', np.uint),
+                ('istate_type', istate_type_dtype),
+                ('istate_status', istate_status_dtype),
+                ('basis_auxref', vstr_dtype),
+                ('pcoord', self.pcoord.dtype, (len(self.pcoord),)),
+            ]
+        )
+        return np.array(
+            [
+                (
+                    self.state_id,
+                    self.basis_state_id or 0,
+                    self.iter_created or 0,
+                    self.iter_used or 0,
+                    self.istate_type or 0,
+                    self.istate_status or 0,
+                    self.basis_auxref or '',
+                    self.pcoord,
+                )
+            ],
+            dtype=istate_dtype,
+        )[0]
+
+
+InitialState.istate_statuses.update(
+    {_attr: getattr(InitialState, _attr) for _attr in dir(InitialState) if _attr.startswith('ISTATE_STATUS_')}
+)
+InitialState.istate_types.update(
+    {_attr: getattr(InitialState, _attr) for _attr in dir(InitialState) if _attr.startswith('ISTATE_TYPE_')}
+)
+InitialState.istate_status_names.update({istate_status: name for name, istate_status in InitialState.istate_statuses.items()})
+InitialState.istate_type_names.update({istate_type: name for name, istate_type in InitialState.istate_types.items()})
+
+
+class TargetState:
+    '''Describes a target state.
+
+    :ivar state_id:     Integer identifier of this state, usually set by the
+                        data manager.
+    :ivar label:        A descriptive label for this microstate (may be empty)
+    :ivar pcoord: The representative progress coordinate of this state.
+
+    '''
+
+    def __init__(self, label, pcoord, state_id=None):
+        self.label = label
+        self.pcoord = np.atleast_1d(pcoord)
+        self.state_id = state_id
+
+    def __repr__(self):
+        return '{} state_id={self.state_id!r} label={self.label!r} pcoord={self.pcoord!r}>'.format(
+            object.__repr__(self)[:-1], self=self
+        )
+
+    @classmethod
+    def states_to_file(cls, states, fileobj):
+        '''Write a file defining basis states, which may then be read by `states_from_file()`.'''
+
+        if isinstance(fileobj, str):
+            fileobj = open(fileobj, 'wt')
+
+        max_label_len = max(8, max(len(state.label or '') for state in states))
+
+        fileobj.write('# {:{max_label_len}s}    {:s}\n'.format('Label', 'Pcoord', max_label_len=max_label_len - 2))
+        for state in states:
+            pcoord_str = '    '.join(str(field) for field in state.pcoord)
+            fileobj.write('{:{max_label_len}s}    {:s}\n'.format(state.label, pcoord_str, max_label_len=max_label_len))
+
+    @classmethod
+    def states_from_file(cls, statefile, dtype):
+        '''Read a file defining target states.  Each line defines a state, and contains a label followed
+        by a representative progress coordinate value, separated by whitespace, as in::
+
+            bound     0.02
+
+        for a single target and one-dimensional progress coordinates or::
+
+            bound    2.7    0.0
+            drift    100    50.0
+
+        for two targets and a two-dimensional progress coordinate.
+        '''
+
+        labels = []
+        pcoord_values = []
+
+        try:
+            open_statefile = open(statefile, 'r')
+        except TypeError:
+            open_statefile = statefile
+
+        for line in open_statefile:
+            line = line.partition('#')[0].strip()
+            if not line:
+                continue
+
+            fields = line.split()
+            labels.append(fields[0])
+            pcoord_values.append(np.array(list(map(dtype, fields[1:])), dtype=dtype))
+
+        try:
+            open_statefile.close()
+        except Exception:
+            pass
+
+        return [cls(label=label, pcoord=pcoord) for label, pcoord in zip(labels, pcoord_values)]
+
+
+def pare_basis_initial_states(basis_states, initial_states, segments=None):
+    '''Given iterables of basis and initial states (and optionally segments that use them),
+    return minimal sets (as in __builtins__.set) of states needed to describe the history of the given
+    segments an initial states.'''
+
+    bstatemap = {state.state_id: state for state in basis_states}
+    istatemap = {state.state_id: state for state in initial_states}
+
+    if segments is not None:
+        segments = list(segments)
+        return_istates = set(
+            istatemap[segment.initial_state_id] for segment in segments if segment.initpoint_type == Segment.SEG_INITPOINT_NEWTRAJ
+        )
+    else:
+        return_istates = set(initial_states)
+
+    return_bstates = set(
+        bstatemap[istate.basis_state_id]
+        for istate in return_istates
+        if istate.istate_type != InitialState.ISTATE_TYPE_RESTART and istate.istate_type != InitialState.ISTATE_TYPE_START
+    )
+
+    return return_bstates, return_istates
+
+
+def return_state_type(state_obj):
+    '''Convinience function for returning the state ID and type of the state_obj pointer'''
+
+    if isinstance(state_obj, Segment):
+        return type(state_obj).__name__, state_obj.seg_id
+    elif isinstance(state_obj, InitialState):
+        return type(state_obj).__name__, state_obj.state_id
+    elif isinstance(state_obj, BasisState):
+        return type(state_obj).__name__, state_obj.state_id
+    else:
+        return 'Unknown', float('inf')

westpa/core/systems.py

@@ -0,0 +1,93 @@
+import logging
+
+import numpy as np
+
+import westpa
+from westpa.core.binning import NopMapper
+
+log = logging.getLogger(__name__)
+
+
+class WESTSystem:
+    '''A description of the system being simulated, including the dimensionality and
+    data type of the progress coordinate, the number of progress coordinate entries
+    expected from each segment, and binning. To construct a simulation, the user must
+    subclass WESTSystem and set several instance variables.
+
+    At a minimum, the user must subclass ``WESTSystem`` and override
+    :method:`initialize` to set the data type and dimensionality of progress
+    coordinate data and define a bin mapper.
+
+    :ivar pcoord_ndim:    The number of dimensions in the progress coordinate.
+                          Defaults to 1 (i.e. a one-dimensional progress
+                          coordinate).
+    :ivar pcoord_dtype:   The data type of the progress coordinate, which must be
+                          callable (e.g. ``np.float32`` and ``long`` will work,
+                          but ``'<f4'`` and ``'<i8'`` will not).  Defaults to
+                          ``np.float64``.
+    :ivar pcoord_len:     The length of the progress coordinate time series
+                          generated by each segment, including *both* the initial
+                          and final values.  Defaults to 2 (i.e. only the initial
+                          and final progress coordinate values for a segment are
+                          returned from propagation).
+    :ivar bin_mapper:     A bin mapper describing the progress coordinate space.
+    :ivar bin_target_counts: A vector of target counts, one per bin.
+    '''
+
+    def __init__(self, rc=None):
+        self.rc = rc or westpa.rc
+
+        # Number of dimentions in progress coordinate data
+        self.pcoord_ndim = 1
+
+        # Length of progress coordinate data for each segment
+        self.pcoord_len = 2
+
+        # Data type of progress coordinate
+        self.pcoord_dtype = np.float32
+
+        # Mapper
+        self.bin_mapper = NopMapper()
+        self._bin_target_counts = None
+
+        self.bin_target_counts = [1]
+
+    @property
+    def bin_target_counts(self):
+        return self._bin_target_counts
+
+    @bin_target_counts.setter
+    def bin_target_counts(self, target_counts):
+        maxcount = max(target_counts)
+        self._bin_target_counts = np.array(target_counts, dtype=np.min_scalar_type(maxcount))
+
+    def initialize(self):
+        '''Prepare this system object for use in simulation or analysis,
+        creating a bin space, setting replicas per bin, and so on. This
+        function is called whenever a WEST tool creates an instance of the
+        system driver.
+        '''
+        pass
+
+    def prepare_run(self):
+        '''Prepare this system for use in a simulation run. Called by w_run in
+        all worker processes.'''
+        pass
+
+    def finalize_run(self):
+        '''A hook for system-specific processing for the end of a simulation run
+        (as defined by such things as maximum wallclock time, rather than perhaps
+        more scientifically-significant definitions of "the end of a simulation run")'''
+        pass
+
+    def new_pcoord_array(self, pcoord_len=None):
+        '''Return an appropriately-sized and -typed pcoord array for a timepoint, segment,
+        or number of segments. If ``pcoord_len`` is not specified (or None), then
+        a length appropriate for a segment is returned.'''
+
+        if pcoord_len is None:
+            pcoord_len = self.pcoord_len
+        return np.zeros((pcoord_len, self.pcoord_ndim), self.pcoord_dtype)
+
+    def new_region_set(self):
+        raise NotImplementedError('This method has been removed.')

westpa/core/textio.py

@@ -0,0 +1,74 @@
+'''Miscellaneous routines to help with input and output of WEST-related data in text format'''
+
+import getpass
+import socket
+import sys
+import time
+
+
+class NumericTextOutputFormatter:
+    comment_string = '# '
+    emit_header = True
+
+    def __init__(self, output_file, mode='wt', emit_header=None):
+        if hasattr(output_file, 'write'):
+            self._file = output_file
+        else:
+            self._file = open(output_file, mode)
+
+        self._header_written = False
+        self._header_lines = []
+
+        self.write_header('created by program: {}'.format(sys.argv[0] or 'unknown program'))
+        self.write_header('created by user:    {}'.format(getpass.getuser()))
+        self.write_header('created on host:    {}'.format(socket.gethostname()))
+        self.write_header('created at:         {}'.format(time.strftime('%c', time.localtime(time.time()))))
+
+        if emit_header is not None:
+            self.emit_header = emit_header
+        else:
+            self.emit_header = self.__class__.emit_header
+
+    def __getattr__(self, attr):
+        return getattr(self._file, attr)
+
+    def close(self):
+        if not self._header_written:
+            self._write_header()
+        self._file.close()
+
+    def write(self, str):
+        if not self._header_written:
+            self._write_header()
+        self._file.write(str)
+
+    def writelines(self, sequence):
+        if not self._header_written:
+            self._write_header()
+        self._file.writelines(sequence)
+
+    def write_comment(self, line):
+        '''Writes a line beginning with the comment string'''
+        self._file.write('{}{}\n'.format(self.comment_string, line))
+
+    def write_header(self, line):
+        '''Appends a line to those written when the file header is written. The
+        appropriate comment string will be prepended, so ``line`` should not include
+        a comment character.'''
+        if self._header_written:
+            raise EnvironmentError('cannot append header lines after header has been written')
+        else:
+            self._header_lines.append(line)
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self._file.close()
+        return False
+
+    def _write_header(self):
+        if self.emit_header and not self._header_written:
+            for line in self._header_lines:
+                self.write_comment(line)
+        self._header_written = True