westpa 2022.13__cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl

westpa/core/_rc.py

@@ -0,0 +1,673 @@
+"""WEST run control and configuration routines"""
+
+import math
+import os
+import sys
+import errno
+import logging
+import warnings
+from copy import deepcopy
+from datetime import timedelta
+
+import numpy as np
+
+import westpa
+import westpa.core.data_manager
+from westpa.core.binning.assign import BinMapper
+from westpa.core.binning import RectilinearBinMapper, RecursiveBinMapper, MABBinMapper, BinlessMapper
+from .yamlcfg import YAMLConfig
+from .systems import WESTSystem
+from . import extloader
+from ..work_managers import SerialWorkManager
+
+log = logging.getLogger('westpa.rc')
+
+
+def bins_from_yaml_dict(bin_dict):
+    kwargs = deepcopy(bin_dict)
+    typename = kwargs.pop('type')
+
+    try:
+        mapper_type = getattr(sys.modules['westpa.core.binning'], typename)
+    except AttributeError:
+        try:
+            mapper_type = extloader.get_object(typename)
+        except AttributeError:
+            raise KeyError('unknown bin mapper type {!r}'.format(typename))
+
+    if not issubclass(mapper_type, BinMapper):
+        raise ValueError('{} is not a BinMapper'.format(mapper_type.__name__))
+
+    if mapper_type is RectilinearBinMapper:
+        boundary_lists = kwargs.pop('boundaries', None)
+        if boundary_lists is None:
+            raise KeyError('RectilinearBinMapper: missing boundaries')
+        parsed_lists = boundary_lists[:]
+        for iboundary, boundary in enumerate(boundary_lists):
+            if boundary.__class__ == str:
+                parsed_lists[iboundary] = parsePCV(boundary)[0]
+            else:
+                parsed_lists[iboundary] = list(map((lambda x: float(x) if isinstance(x, str) else x), boundary))
+        return RectilinearBinMapper(parsed_lists)
+    elif mapper_type is RecursiveBinMapper:
+        base_mapper_config = kwargs.pop('base', None)
+        base_mapper_config = kwargs.pop('base_mapper', base_mapper_config)
+        if base_mapper_config is None:
+            raise KeyError('RecursiveBinMapper: missing base_mapper')
+
+        base_mapper = bins_from_yaml_dict(base_mapper_config)
+        start_index = kwargs.pop('start_index', 0)
+
+        rec_mapper = RecursiveBinMapper(base_mapper, start_index)
+        mapper_configs = kwargs.pop('mappers')
+        mappers = []
+        if mapper_configs is not None:
+            for config in mapper_configs:
+                replaced_bin = config.pop('replaces_bin_at', None)
+                replaced_bin = config.pop('at', replaced_bin)
+                if replaced_bin is None:
+                    raise KeyError('RecursiveBinMapper: missing replaces_bin_at for ' 'at least one of the child mappers')
+                mapper = bins_from_yaml_dict(config)
+                mappers.append((mapper, replaced_bin))
+
+        for mapper, replaced_bin in mappers:
+            rec_mapper.add_mapper(mapper, replaced_bin)
+
+        return rec_mapper
+    else:
+        try:
+            return mapper_type(**kwargs)
+        except Exception:
+            log.exception('exception instantiating mapper')
+            raise
+
+
+def detect_mab_mapper(mapper):
+    if isinstance(mapper, MABBinMapper):
+        return True
+    elif isinstance(mapper, RecursiveBinMapper):
+        if detect_mab_mapper(mapper.base_mapper):
+            return True
+
+        for ibin in mapper._recursion_targets:
+            rec_mapper = mapper._recursion_targets[ibin]
+            if detect_mab_mapper(rec_mapper):
+                return True
+    else:
+        return False
+
+
+def detect_binless_mapper(mapper):
+    if isinstance(mapper, BinlessMapper):
+        return True
+    elif isinstance(mapper, RecursiveBinMapper):
+        if detect_binless_mapper(mapper.base_mapper):
+            return True
+
+        for ibin in mapper._recursion_targets:
+            rec_mapper = mapper._recursion_targets[ibin]
+            if detect_binless_mapper(rec_mapper):
+                return True
+    else:
+        return False
+
+
+def parsePCV(pc_str):
+    # Execute arbitrary code within a limited
+    # scope to avoid nastyness. Stolen fully from
+    # other parts of the WESTPA code.
+    namespace = {'math': math, 'numpy': np, 'np': np, 'inf': float('inf')}
+
+    arr = np.array(eval(pc_str, namespace))
+    if arr.ndim == 0:
+        arr.shape = (1, 1)
+    elif arr.ndim == 1:
+        arr.shape = (1,) + arr.shape
+    else:
+        raise ValueError('too many dimensions')
+    # return list(arr[...])
+    return arr[...]
+
+
+def lazy_loaded(backing_name, loader, docstring=None):
+    def getter(self):
+        obj = getattr(self, backing_name, None)
+        if obj is None:
+            obj = loader()
+            setattr(self, backing_name, obj)
+        return obj
+
+    def setter(self, val):
+        setattr(self, backing_name, val)
+
+    def deleter(self):
+        delattr(self, backing_name)
+        setattr(self, backing_name, None)
+
+    return property(getter, setter, deleter, docstring)
+
+
+class WESTRC:
+    '''A class, an instance of which is accessible as ``westpa.rc``, to handle global issues for WEST-PA code,
+    such as loading modules and plugins, writing output based on verbosity level, adding default command line options,
+    and so on.'''
+
+    # Runtime config file management
+    ENV_RUNTIME_CONFIG = 'WESTRC'
+    RC_DEFAULT_FILENAME = 'west.cfg'
+
+    def __init__(self):
+        self.verbosity = None
+        self.rcfile = os.environ.get(self.ENV_RUNTIME_CONFIG) or self.RC_DEFAULT_FILENAME
+
+        self.config = YAMLConfig()
+        try:
+            self.process_name = os.path.splitext(os.path.basename(sys.argv[0]))[0]
+        except (TypeError, IndexError):
+            self.process_name = "unknown"
+
+        # Crucial simulation and analysis drivers
+        self._system = None
+        self._data_manager = None
+        self._sim_manager = None
+        self._we_driver = None
+        self._propagator = None
+
+        self.work_manager = SerialWorkManager()
+
+        self.status_stream = sys.stdout
+
+    def add_args(self, parser):
+        group = parser.add_argument_group('general options')
+        group.add_argument(
+            '-r',
+            '--rcfile',
+            metavar='RCFILE',
+            dest='rcfile',
+            default=(os.environ.get(self.ENV_RUNTIME_CONFIG) or self.RC_DEFAULT_FILENAME),
+            help='use RCFILE as the WEST run-time configuration file (default: %(default)s)',
+        )
+
+        egroup = group.add_mutually_exclusive_group()
+        egroup.add_argument(
+            '--quiet', dest='verbosity', action='store_const', const='quiet', help='emit only essential information'
+        )
+        egroup.add_argument('--verbose', dest='verbosity', action='store_const', const='verbose', help='emit extra information')
+        egroup.add_argument(
+            '--debug',
+            dest='verbosity',
+            action='store_const',
+            const='debug',
+            help='enable extra checks and emit copious information',
+        )
+
+        group.add_argument('--version', action='version', version='WEST version %s' % westpa.__version__)
+
+    @property
+    def verbose_mode(self):
+        return self.verbosity in ('verbose', 'debug')
+
+    @property
+    def debug_mode(self):
+        return self.verbosity == 'debug'
+
+    @property
+    def quiet_mode(self):
+        return self.verbosity == 'quiet'
+
+    def process_args(self, args, config_required=True):
+        self.cmdline_args = args
+        self.verbosity = args.verbosity
+
+        if args.rcfile:
+            self.rcfile = args.rcfile
+
+        try:
+            self.read_config()
+        except IOError as e:
+            if e.errno == errno.ENOENT and not config_required:
+                pass
+            else:
+                raise
+        self.config_logging()
+        self.config['args'] = {k: v for k, v in args.__dict__.items() if not k.startswith('_')}
+        self.process_config()
+        self.time_config()
+
+    def time_config(self):
+        '''Convert non-pyYAML accepted time formats into seconds'''
+        max_run_wallclock = self.config.get(['west', 'propagation', 'max_run_wallclock'], None)
+        if isinstance(max_run_wallclock, str):
+            time_formats = ['days', 'hours', 'minutes', 'seconds']
+            wallclock_str = max_run_wallclock.split(':')
+            kwargs = {key: int(val) for key, val in zip(time_formats[-len(wallclock_str) :], wallclock_str)}
+            try:
+                diff = int(timedelta(**kwargs).total_seconds())
+                self.config['west']['propagation']['max_run_wallclock'] = diff
+                log.debug(f'Automatically converted {max_run_wallclock=} to {diff}')
+            except ValueError:
+                log.debug(f'Failed to convert {max_run_wallclock=} to seconds.')
+
+    def process_config(self):
+        log.debug('config: {!r}'.format(self.config))
+        sys.path.extend(self.config.get_pathlist(['west', 'drivers', 'module_path'], []))
+        try:
+            sys.path.append(os.environ['WEST_SIM_ROOT'])
+        except KeyError:
+            pass
+
+    def read_config(self, filename=None):
+        if filename:
+            self.rcfile = filename
+
+        if 'WEST_SIM_ROOT' not in os.environ:
+            # sys.stderr.write('-- WARNING -- setting $WEST_SIM_ROOT to current directory ({})\n'.format(os.getcwd()))
+            os.environ['WEST_SIM_ROOT'] = os.getcwd()
+
+        self.config.update_from_file(self.rcfile)
+
+    def config_logging(self):
+        import logging.config
+
+        logging_config = {
+            'version': 1,
+            'incremental': False,
+            'formatters': {
+                'standard': {'format': '-- %(levelname)-8s [%(name)s] -- %(message)s'},
+                'debug': {
+                    'format': '''\
+-- %(levelname)-8s %(asctime)24s PID %(process)-12d TID %(thread)-20d
+   from logger "%(name)s"
+   at location %(pathname)s:%(lineno)d [%(funcName)s()]
+   ::
+   %(message)s
+'''
+                },
+            },
+            'handlers': {'console': {'class': 'logging.StreamHandler', 'stream': 'ext://sys.stdout', 'formatter': 'standard'}},
+            'loggers': {
+                'west': {'handlers': ['console'], 'propagate': False},
+                'westpa': {'handlers': ['console'], 'propagate': False},
+                'oldtools': {'handlers': ['console'], 'propagate': False},
+                'westtools': {'handlers': ['console'], 'propagate': False},
+                'westext': {'handlers': ['console'], 'propagate': False},
+                'work_managers': {'handlers': ['console'], 'propagate': False},
+                'py.warnings': {'handlers': ['console'], 'propagate': False},
+            },
+            'root': {'handlers': ['console']},
+        }
+
+        logging_config['loggers'][self.process_name] = {'handlers': ['console'], 'propagate': False}
+
+        if self.verbosity == 'debug':
+            logging_config['root']['level'] = 5  # 'DEBUG'
+            logging_config['handlers']['console']['formatter'] = 'debug'
+        elif self.verbosity == 'verbose':
+            logging_config['root']['level'] = 'INFO'
+        else:
+            logging_config['root']['level'] = 'WARNING'
+
+        logging.config.dictConfig(logging_config)
+        logging_config['incremental'] = True
+
+        if self.verbosity == 'debug':
+            warnings.resetwarnings()
+            warnings.simplefilter("default")
+            logging.captureWarnings(True)
+        else:
+            if not sys.warnoptions:
+                warnings.simplefilter("ignore")
+            logging.captureWarnings(False)
+
+    def pstatus(self, *args, **kwargs):
+        fileobj = kwargs.pop('file', self.status_stream)
+        if kwargs.get('termonly', False) and not fileobj.isatty():
+            return
+        if self.verbosity != 'quiet':
+            print(*args, file=fileobj, **kwargs)
+
+    def pstatus_term(self, *args, **kwargs):
+        fileobj = kwargs.pop('file', self.status_stream)
+        if fileobj.isatty() and self.verbosity != 'quiet':
+            print(*args, file=fileobj, **kwargs)
+
+    def pflush(self):
+        for stream in (self.status_stream, sys.stdout, sys.stderr):
+            try:
+                stream.flush()
+            except AttributeError:
+                pass
+
+    def detect_mab_mapper(self):
+        bin_dict = self.config.get(['west', 'system', 'system_options', 'bins'])
+        use_mab = False
+        if bin_dict is not None:
+            mapper = bins_from_yaml_dict(bin_dict)
+            use_mab = detect_mab_mapper(mapper)
+
+        return use_mab
+
+    def detect_binless_mapper(self):
+        bin_dict = self.config.get(['west', 'system', 'system_options', 'bins'])
+        use_binless = False
+        if bin_dict is not None:
+            mapper = bins_from_yaml_dict(bin_dict)
+            use_binless = detect_binless_mapper(mapper)
+
+        return use_binless
+
+    def new_sim_manager(self):
+        drivername = self.config.get(['west', 'drivers', 'sim_manager'], 'default')
+
+        if drivername.lower() == 'default':
+            use_mab = self.detect_mab_mapper()
+            use_binless = self.detect_binless_mapper()
+
+            if use_mab:
+                from .binning.mab_manager import MABSimManager
+
+                sim_manager = MABSimManager(rc=self)
+            elif use_binless:
+                from .binning.binless_manager import BinlessSimManager
+
+                sim_manager = BinlessSimManager(rc=self)
+            else:
+                from .sim_manager import WESimManager
+
+                sim_manager = WESimManager(rc=self)
+        else:
+            sim_manager = extloader.get_object(drivername)(rc=self)
+
+        log.debug('loaded simulation manager {!r}'.format(sim_manager))
+        return sim_manager
+
+    def get_sim_manager(self):
+        if self._sim_manager is None:
+            self._sim_manager = self.new_sim_manager()
+        return self._sim_manager
+
+    def new_data_manager(self):
+        import westpa
+
+        drivername = self.config.get(['west', 'drivers', 'data_manager'], 'hdf5')
+        if drivername.lower() in ('hdf5', 'default'):
+            data_manager = westpa.core.data_manager.WESTDataManager()
+        else:
+            data_manager = extloader.get_object(drivername)(rc=self)
+        log.debug('loaded data manager: {!r}'.format(data_manager))
+        return data_manager
+
+    def get_data_manager(self):
+        if self._data_manager is None:
+            self._data_manager = self.new_data_manager()
+        return self._data_manager
+
+    def new_we_driver(self):
+        import westpa
+
+        drivername = self.config.get(['west', 'drivers', 'we_driver'], 'default')
+
+        if drivername.lower() == 'default':
+            use_mab = self.detect_mab_mapper()
+            use_binless = self.detect_binless_mapper()
+            if use_mab:
+                from .binning.mab_driver import MABDriver
+
+                we_driver = MABDriver()
+            elif use_binless:
+                from .binning.binless_driver import BinlessDriver
+
+                we_driver = BinlessDriver()
+            else:
+                from .we_driver import WEDriver
+
+                we_driver = WEDriver()
+        else:
+            we_driver = extloader.get_object(drivername)(rc=self)
+
+        log.debug('loaded WE algorithm driver: {!r}'.format(we_driver))
+
+        subgroup_function = self.config.get(['west', 'drivers', 'subgroup_function'], 'default')
+        if subgroup_function.lower() == 'default':
+            try:
+                subgroup_function = 'westpa.core.we_driver._group_walkers_identity'
+                we_driver.subgroup_function = westpa.core.we_driver._group_walkers_identity
+            except Exception:
+                pass
+        else:
+            we_driver.subgroup_function = extloader.get_object(subgroup_function)
+        we_driver.subgroup_function_kwargs = self.config.get(['west', 'drivers', 'subgroup_arguments'])
+        # Necessary if the user hasn't specified any options.
+        if we_driver.subgroup_function_kwargs is None:
+            we_driver.subgroup_function_kwargs = {}
+        log.debug('loaded WE algorithm driver subgrouping function {!r}'.format(subgroup_function))
+        log.debug('WE algorithm driver subgrouping function kwargs: {!r}'.format(we_driver.subgroup_function_kwargs))
+
+        return we_driver
+
+    def get_we_driver(self):
+        if self._we_driver is None:
+            self._we_driver = self.new_we_driver()
+        return self._we_driver
+
+    def new_propagator(self):
+        drivername = self.config.require(['west', 'propagation', 'propagator'])
+        if drivername.lower() == 'executable':
+            from westpa.core.propagators.executable import ExecutablePropagator
+
+            propagator = ExecutablePropagator()
+        else:
+            propagator = extloader.get_object(drivername)(rc=self)
+        log.debug('loaded propagator {!r}'.format(propagator))
+        return propagator
+
+    def get_propagator(self):
+        if self._propagator is None:
+            self._propagator = self.new_propagator()
+        return self._propagator
+
+    def new_system_driver(self):
+        '''
+        Returns a new system object either from the driver OR from the YAML
+        file. Currently builds off of system then updates with YAML
+        overwriting the previous settings if both are specified.
+
+        There are no default settings, all settings MUST be specified
+        in the config YAML file OR the system driver.
+
+        Settings that are specified here are:
+          Progress coordinate settings:
+            Progress coordinate dimensionality
+            Progress coordinate length/number of data points in each tau
+            Progress coordinate data type
+          Bin settings:
+            Bin mapper type
+            Bins
+            Target simulation counts for each bin
+          Generic setting for flexibility:
+            Both paths allow for generic attribute setting for
+            future flexibility.
+        '''
+
+        # First step is to see if we have a driver specified
+        # if so load that one first, setting the defaults for
+        # YAML to possibly modify.
+
+        # I will keep this as is for now since I like the idea
+        # of being able to set some defaults from the system
+        # driver and then just modify the YAML instead for basic
+        # stuff. This might make it easier to set up systems since
+        # playing around with the YAML format is easier.
+
+        # Still still has the end-user issue of possibly confusing
+        # people, maybe KISS is better, not sure. I'll keep this
+        # for development purposes if nothing else.
+
+        system = None
+        # Get method checks for us
+        sysdrivername = self.config.get(['west', 'system', 'driver'])
+        if not sysdrivername:
+            # Warn user that driver is not specified
+            log.info("System driver not specified")
+        else:
+            log.info('loading system driver %r' % sysdrivername)
+            system = extloader.get_object(sysdrivername)(rc=self)
+            log.debug('loaded system driver {!r}'.format(system))
+            system.initialize()
+        # Second let's see if we have info in the YAML file
+        yamloptions = self.config.get(['west', 'system', 'system_options'])
+        if not yamloptions:
+            # Same here, yaml file doesn't have sys info
+            log.info("Config file doesn't contain any system info")
+        else:
+            log.info("Loading system options from configuration file")
+            if system:
+                system = self.update_from_yaml(system, yamloptions)
+            else:
+                system = self.system_from_yaml(yamloptions)
+
+        if system:
+            if not yamloptions:
+                print("System is being built only off of the system driver")
+            return system
+        else:
+            log.info("No system specified! Exiting program.")
+            # Gracefully exit
+            raise ValueError("No system defined!")
+
+    def system_from_yaml(self, system_dict):
+        """
+        System builder directly from the config YAML file.
+
+        Arguments:
+          system_dict (dict): Parsed YAML file as a dictionary, parsed by
+            PyYAML by default.
+
+        Returns:
+          A modified WESTSystem object as defined in yamlcfg.py with
+          the parsed settings from the config file.
+        """
+
+        yamlSystem = WESTSystem()
+        print("System building only off of the configuration file")
+        # Now for the building of the system from YAML we need to use
+        # require for these settings since they are musts.
+
+        # First basic pcoord settings
+        ndim = self.config.require(['west', 'system', 'system_options', 'pcoord_ndim'])
+        plen = self.config.require(['west', 'system', 'system_options', 'pcoord_len'])
+        # Dtype needs to be ran as code from YAML file, document YAML code execution syntax
+        # somewhere
+        ptype = self.config.require(['west', 'system', 'system_options', 'pcoord_dtype'])
+        # Bins
+        bins_obj = self.config.require(['west', 'system', 'system_options', 'bins'])
+        trgt_cnt = self.config.require(['west', 'system', 'system_options', 'bin_target_counts'])
+        # Now add the parsed settings to the system
+        mapper = bins_from_yaml_dict(bins_obj)
+        setattr(yamlSystem, 'pcoord_ndim', ndim)
+        setattr(yamlSystem, 'pcoord_len', plen)
+        setattr(yamlSystem, 'pcoord_dtype', ptype)
+        setattr(yamlSystem, 'bin_mapper', mapper)
+        # Check if the supplied target count object is
+        # an iterable or not,
+
+        # This one I designed in a way that you can either
+        # directly supply an iterable that has the correct
+        # size OR an integer. Anything else will fail.
+
+        # Main issue: For higher bin dimensions this
+        # is tough since the bins might not correspond
+        # properly to the flattened array you are supplying.
+
+        # I might just scrap the iterable later and only allow
+        # integers.
+        if hasattr(trgt_cnt, "__iter__"):
+            assert len(trgt_cnt) == mapper.nbins, "Count iterable size doesn't match the number of bins"
+            trgt_cnt_arr = trgt_cnt
+        else:
+            assert trgt_cnt == int(trgt_cnt), "Counts are not integer valued, ambiguous input"
+            trgt_cnt_arr = np.zeros(mapper.nbins)
+            trgt_cnt_arr[:] = trgt_cnt
+        setattr(yamlSystem, 'bin_target_counts', trgt_cnt_arr)
+
+        # Attach generic attribute to system
+        for attr in system_dict.keys():
+            if not hasattr(yamlSystem, attr):
+                setattr(yamlSystem, attr, system_dict[attr])
+
+        # Return complete system
+        return yamlSystem
+
+    def update_from_yaml(self, init_system, system_dict):
+        """
+        Updates the system built from the driver with the options
+        from the YAML file. For now it overwrites everything specified
+        in the system driver.
+
+        Arguments:
+          system_dict (dict): Parsed YAML file as a dictionary, parsed by
+            PyYAML by default.
+          init_system (WESTSystem): System returned by the driver.
+
+        Returns:
+          A modified WESTSystem object with settings from the system
+          driver and the config YAML file.
+        """
+
+        # First we want to overwrite whatever we have from the YAML
+        # file.
+        print("Updating system with the options from the configuration file")
+        for key, value in system_dict.items():
+            if key == 'pcoord_ndim':
+                self.overwrite_option(init_system, key, value)
+            elif key == 'pcoord_len':
+                self.overwrite_option(init_system, key, value)
+            elif key == 'pcoord_dtype':
+                self.overwrite_option(init_system, key, value)
+            elif key == "bins":
+                self.overwrite_option(init_system, 'bin_mapper', bins_from_yaml_dict(value))
+        # Target counts have to be parsed after we have a mapper in
+        # place
+        try:
+            trgt_cnt = system_dict['bin_target_counts']
+            if hasattr(trgt_cnt, "__iter__"):
+                assert len(trgt_cnt) == init_system.bin_mapper.nbins, "Count iterable size doesn't match the number of bins"
+                trgt_cnt_arr = trgt_cnt
+            else:
+                assert trgt_cnt == int(trgt_cnt), "Counts are not integer valued, ambiguous input"
+                trgt_cnt_arr = np.zeros(init_system.bin_mapper.nbins)
+                trgt_cnt_arr[:] = int(trgt_cnt)
+            self.overwrite_option(init_system, 'bin_target_counts', trgt_cnt_arr)
+        except KeyError:
+            pass
+        # The generic attribute settings added here
+        for attr in system_dict.keys():
+            if not hasattr(init_system, attr):
+                setattr(init_system, attr, system_dict[attr])
+        return init_system
+
+    def overwrite_option(self, system, key, value):
+        if hasattr(system, key):
+            log.info("Overwriting system option: %s" % key)
+        setattr(system, key, value)
+
+    def get_system_driver(self):
+        if self._system is None:
+            self._system = self.new_system_driver()
+        return self._system
+
+    def get_work_manager(self):
+        return self.work_manager
+
+    def clear_state(self):
+        self._sim_manager = None
+        self._system = None
+        self._data_manager = None
+        self._we_driver = None
+        self._propagator = None
+
+    propagator = property(get_propagator)
+    we_driver = property(get_we_driver)
+    system = property(get_system_driver)
+    data_manager = property(get_data_manager)
+    sim_manager = property(get_sim_manager)