westpa 2022.12__cp313-cp313-macosx_10_13_x86_64.whl

westpa/trajtree/__init__.py

@@ -0,0 +1,4 @@
+from . import trajtree
+from .trajtree import TrajTreeSet
+
+__all__ = ['trajtree', 'TrajTreeSet']

westpa/trajtree/trajtree.py

@@ -0,0 +1,117 @@
+import collections
+
+import numpy as np
+
+import westpa
+from westpa.tools.selected_segs import AllSegmentSelection
+from . import _trajtree
+from ._trajtree import _trajtree_base  # @UnresolvedImport
+
+
+trajnode = collections.namedtuple('trajnode', ('n_iter', 'seg_id'))
+
+
+class TrajTreeSet(_trajtree_base):
+    def __init__(self, segsel=None, data_manager=None):
+        self.data_manager = data_manager or westpa.rc.get_data_manager()
+        self.segsel = segsel or AllSegmentSelection(data_manager=self.data_manager)
+        self._build_table(self.segsel, self.data_manager)
+
+    def __len__(self):
+        return len(self.trajtable)
+
+    def get_roots(self):
+        return self.trajtable[self.trajtable['parent_offset'] == -1]
+        # return [trajnode(root['n_iter'], root['seg_id']) for root in self._get_roots()]
+
+    def get_root_indices(self):
+        return np.squeeze(np.argwhere(self.trajtable['parent_offset'] == -1))
+
+    def trace_trajectories(self, visit, get_visitor_state=None, set_visitor_state=None, vargs=None, vkwargs=None):
+        if (get_visitor_state or set_visitor_state) and not (get_visitor_state and set_visitor_state):
+            raise ValueError('either both or neither of get_visitor_state and set_visitor_state must be specified')
+
+        vargs = vargs or ()
+        vkwargs = vkwargs or {}
+        n_visits = 0
+
+        trajtable = self.trajtable
+
+        roots = collections.deque(self.get_root_indices())
+        print('Examining {:d} roots'.format(len(roots)))
+        state_stack = collections.deque([{'subtrees': roots, 'vstate': get_visitor_state() if get_visitor_state else None}])
+        while state_stack:
+            state = state_stack.pop()
+            subtrees = state['subtrees']
+            if set_visitor_state:
+                set_visitor_state(state['vstate'])
+
+            while subtrees:
+                index = subtrees.popleft()
+                node = trajtable[index]
+
+                state_stack.append({'subtrees': subtrees, 'vstate': get_visitor_state() if get_visitor_state else None})
+
+                subtrees = collections.deque(self.get_child_indices(index))
+
+                n_visits += 1
+                try:
+                    visit(node['n_iter'], node['seg_id'], node['weight'], has_children=(len(subtrees) > 0), *vargs, **vkwargs)
+                except StopIteration:
+                    subtrees = collections.deque()
+                    continue  # to next sibling
+
+        return n_visits
+
+
+class FakeTrajTreeSet(TrajTreeSet):
+    def __init__(self):
+        # _tt_dtype = np.dtype([('n_iter', np.uint32),
+        #                         ('seg_id', np.int64),
+        #                         ('parent_id', np.int64),
+        #                         ('parent_offset', np.int64),  # offset of parent segment into this table
+        #                         ('weight', np.float64)])       # weight of this segment
+
+        self.trajtable = np.array(
+            [
+                (1, 1, -1, -1, 1.0),  # 0
+                (1, 11, -1, -1, 1.0),  # 1
+                (2, 2, 1, 0, 1.0),  # 2
+                (2, 3, 1, 0, 1.0),  # 3
+                (2, 4, 1, 0, 1.0),  # 4
+                (2, 12, 11, 1, 1.0),  # 5
+                (2, 13, 11, 1, 1.0),  # 6
+                (3, 5, 2, 2, 1.0),  # 7
+                (3, 6, 3, 3, 1.0),  # 8
+                (3, 7, 4, 4, 1.0),  # 9
+                (3, 8, 4, 4, 1.0),  # 10
+                (3, 14, 12, 5, 1.0),  # 11
+                (3, 15, 12, 5, 1.0),  # 12
+                (4, 9, 5, 7, 1.0),  # 13
+                (4, 10, 5, 7, 1.0),  # 14
+            ],
+            dtype=_trajtree._tt_dtype,
+        )
+
+        empty_array = np.array([])
+        self.childtable = np.array(
+            [
+                np.array([2, 3, 4]),  # segment 1
+                np.array([5, 6]),  # segment 11
+                np.array([7]),  # segment 2
+                np.array([8]),  # segment 3
+                np.array([9, 10]),  # segment 4
+                np.array([11, 12]),  # segment 12
+                empty_array,  # segment 13
+                np.array([13, 14]),  # segment 5
+                empty_array,  # segment 6
+                empty_array,  # segment 7
+                empty_array,  # segment 8
+                empty_array,  # segment 14
+                empty_array,  # segment 15
+                empty_array,  # segment 9
+                empty_array,  # segment 10
+            ],
+            dtype=np.object_,
+        )
+        self.iter_offsets = {1: 0, 2: 2, 3: 7, 4: 13}

westpa/westext/adaptvoronoi/__init__.py

@@ -0,0 +1,3 @@
+from .adaptVor_driver import AdaptiveVoronoiDriver
+
+__all__ = ['AdaptiveVoronoiDriver']

westpa/westext/adaptvoronoi/adaptVor_driver.py

@@ -0,0 +1,214 @@
+import logging
+
+import numpy as np
+
+import westpa
+from westpa.core import extloader
+from westpa.core.yamlcfg import check_bool, ConfigItemMissing
+from westpa.core.binning import VoronoiBinMapper
+
+log = logging.getLogger(__name__)
+
+
+class AdaptiveVoronoiDriver:
+    '''
+    This plugin implements an adaptive scheme using voronoi bins from
+    Zhang 2010, J Chem Phys, 132. The options exposed to the configuration
+    file are:
+
+      - av_enabled (bool, default False): Enables adaptive binning
+      - max_centers (int, default 10): The maximum number of voronoi centers to be placed
+      - walk_count (integer, default 5): Number of walkers per voronoi center
+      - center_freq (ingeter, default 1): Frequency of center placement
+      - priority (integer, default 1): Priority in the plugin order
+      - dfunc_method (function, non-optional, no default): Non-optional user defined
+          function that will be used to calculate distances between voronoi centers and
+          data points
+      - mapper_func (function, optional): Optional user defined function for building bin
+          mappers for more complicated binning schemes e.g. embedding the voronoi binning
+          in a portion of the state space. If not defined the plugin will build a
+          VoronoiBinMapper with the information it has.
+    '''
+
+    def __init__(self, sim_manager, plugin_config):
+        if not sim_manager.work_manager.is_master:
+            return
+
+        self.sim_manager = sim_manager
+        self.data_manager = sim_manager.data_manager
+        self.system = sim_manager.system
+
+        # Parameters from config file
+        # this enables the adaptive voronoi, allows turning adaptive scheme off
+        self.doAdaptiveVoronoi = check_bool(plugin_config.get('av_enabled', False))
+        # sets maximim number of centers/voronoi bins
+        self.max_centers = plugin_config.get('max_centers', 10)
+        # sets number of walkers per bin/voronoi center
+        self.walk_count = plugin_config.get('walk_count', 5)
+        # center placement frequency in number of iterations
+        self.center_freq = plugin_config.get('center_freq', 1)
+        # priority of the plugin (allows for order of execution)
+        self.priority = plugin_config.get('priority', 0)
+        # pulls the distance function that will be used by the plugin
+        self.dfunc = self.get_dfunc_method(plugin_config)
+        # pulls a user defined function to build the next bin mapper
+        self.mapper_func = self.get_mapper_func(plugin_config)
+
+        # Get initial set of Voronoi centers
+        self.centers = self.get_initial_centers()
+        self.ncenters = len(self.centers)
+
+        # Update the BinMapper
+        self.update_bin_mapper()
+
+        westpa.rc.pstatus('-adaptive voronoi mapping --------------\n')
+        westpa.rc.pstatus('enabled: {}\n'.format(self.doAdaptiveVoronoi))
+        westpa.rc.pstatus('max centers: {}\n'.format(self.max_centers))
+        westpa.rc.pstatus('center adding freq: {}\n'.format(self.center_freq))
+        westpa.rc.pstatus('centers: {}\n'.format(self.centers))
+        westpa.rc.pstatus('----------------------------------------\n')
+        westpa.rc.pflush()
+
+        # Register callback
+        if self.doAdaptiveVoronoi:
+            sim_manager.register_callback(sim_manager.prepare_new_iteration, self.prepare_new_iteration, self.priority)
+
+    def dfunc(self):
+        '''
+        Distance function to be used by the plugin. This function
+        will be used to calculate the distance between each point.
+        '''
+        raise NotImplementedError
+
+    def get_dfunc_method(self, plugin_config):
+        try:
+            methodname = plugin_config['dfunc_method']
+        except KeyError:
+            raise ConfigItemMissing('dfunc_method')
+
+        dfunc_method = extloader.get_object(methodname)
+
+        log.info('loaded adaptive voronoi dfunc method {!r}'.format(dfunc_method))
+
+        return dfunc_method
+
+    def get_mapper_func(self, plugin_config):
+        try:
+            methodname = plugin_config['mapper_func']
+        except KeyError:
+            return False
+
+        mapper_func = extloader.get_object(methodname)
+
+        log.info('loaded adaptive voronoi mapper function {!r}'.format(mapper_func))
+
+        return mapper_func
+
+    def get_initial_centers(self):
+        '''
+        This function pulls from the centers from either the
+        previous bin mapper  or uses the definition from the
+        system to calculate the number of centers
+        '''
+        self.data_manager.open_backing()
+
+        with self.data_manager.lock:
+            n_iter = max(self.data_manager.current_iteration - 1, 1)
+            iter_group = self.data_manager.get_iter_group(n_iter)
+
+            # First attempt to initialize voronoi centers
+            # from data rather than system
+            centers = None
+            try:
+                log.info('Voronoi centers from previous bin mapper')
+                binhash = iter_group.attrs['binhash']
+                bin_mapper = self.data_manager.get_bin_mapper(binhash)
+
+                centers = bin_mapper.centers
+
+            except Exception:
+                log.warning(
+                    'Initializing voronoi centers from data failed; \
+                        Using definition in system instead.'
+                )
+                centers = self.system.bin_mapper.centers
+
+        self.data_manager.close_backing()
+        return centers
+
+    def update_bin_mapper(self):
+        '''Update the bin_mapper using the current set of voronoi centers'''
+
+        westpa.rc.pstatus('westext.adaptvoronoi: Updating bin mapper\n')
+        westpa.rc.pflush()
+
+        # self.mapper_func = plugin_config.get('mapper_func', False)
+        try:
+            dfargs = getattr(self.system, 'dfargs', None)
+            dfkwargs = getattr(self.system, 'dfkwargs', None)
+            if self.mapper_func:
+                # The mapper should take in 1) distance function,
+                # 2) centers, 3) dfargs, 4) dfkwargs and return
+                # the mapper we want
+                self.system.bin_mapper = self.mapper_func(self.dfunc, self.centers, dfargs=dfargs, dfkwargs=dfkwargs)
+            else:
+                self.system.bin_mapper = VoronoiBinMapper(self.dfunc, self.centers, dfargs=dfargs, dfkwargs=dfkwargs)
+            self.ncenters = self.system.bin_mapper.nbins
+            new_target_counts = np.empty((self.ncenters,), int)
+            new_target_counts[...] = self.walk_count
+            self.system.bin_target_counts = new_target_counts
+        except (ValueError, TypeError) as e:
+            log.error(
+                'AdaptiveVoronoiDriver Error: \
+                    Failed updating the bin mapper: {}'.format(
+                    e
+                )
+            )
+            raise
+
+    def update_centers(self, iter_group):
+        '''
+        Update the set of Voronoi centers according to
+        Zhang 2010, J Chem Phys, 132. A short description
+        of the algorithm can be found in the text:
+
+        1) First reference structure is chosen randomly from
+        the first set of given structure
+        2) Given a set of n reference structures, for each
+        configuration in the iteration the distances to each
+        reference structure is calculated and the minimum
+        distance is found
+        3) The configuration with the minimum distance is
+        selected as the next reference
+        '''
+
+        westpa.rc.pstatus('westext.adaptvoronoi: Updating Voronoi centers\n')
+        westpa.rc.pflush()
+
+        # Pull the current coordinates to find distances
+        curr_pcoords = iter_group['pcoord']
+        # Initialize distance array
+        dists = np.zeros(curr_pcoords.shape[0])
+        for iwalk, walk in enumerate(curr_pcoords):
+            # Calculate distances using the provided function
+            # and find the distance to the closest center
+            dists[iwalk] = min(self.dfunc(walk[-1], self.centers))
+        # Find the maximum of the minimum distances
+        max_ind = np.where(dists == dists.max())
+        # Use the maximum progress coordinate as our next center
+        self.centers = np.vstack((self.centers, curr_pcoords[max_ind[0][0]][-1]))
+
+    def prepare_new_iteration(self):
+        n_iter = self.sim_manager.n_iter
+
+        with self.data_manager.lock:
+            iter_group = self.data_manager.get_iter_group(n_iter)
+
+        # Check if we are at the correct frequency for updating the bin mapper
+        if n_iter % self.center_freq == 0:
+            # Check if we still need to add more centers
+            if self.ncenters < self.max_centers:
+                # First find the center to add
+                self.update_centers(iter_group)
+                # Update the bin mapper with the new center
+                self.update_bin_mapper()

westpa/westext/hamsm_restarting/__init__.py

@@ -0,0 +1,3 @@
+from . import restart_driver
+
+__all__ = ['restart_driver']

westpa/westext/hamsm_restarting/example_overrides.py

@@ -0,0 +1,35 @@
+import logging
+
+from rich.logging import RichHandler
+
+FORMAT = "%(message)s"
+logging.basicConfig(level="NOTSET", format=FORMAT, datefmt="[%X]", handlers=[RichHandler()])
+log = logging.getLogger("msm_we")
+log.setLevel(logging.DEBUG)
+
+import numpy as np
+
+
+def processCoordinates(self, coords):
+    log.debug("Processing coordinates")
+
+    if self.dimReduceMethod == "none":
+        nC = np.shape(coords)
+        nC = nC[0]
+        data = coords.reshape(nC, 3 * self.nAtoms)
+        return data
+
+    if self.dimReduceMethod == "pca" or self.dimReduceMethod == "vamp":
+        ### NaCl dimensionality reduction
+        log.warning("Hardcoded selection: Doing dim reduction for Na, Cl. This is only for testing!")
+        indNA = self.reference_structure.topology.select("element Na")
+        indCL = self.reference_structure.topology.select("element Cl")
+
+        diff = np.subtract(coords[:, indNA], coords[:, indCL])
+
+        dist = np.array(np.sqrt(np.mean(np.power(diff, 2), axis=-1)))
+
+        return dist
+
+
+log.info("Loading user-override functions for modelWE")