"""

Created on 19.07.2014

@author: Jan-Hendrik Prinz

@author: David W. H. Swenson

"""

import abc

import logging

import numpy as np

import random

import openpathsampling as paths

from openpathsampling.netcdfplus import StorableNamedObject, StorableObject

from openpathsampling.pathmover_inout import InOutSet, InOut

from openpathsampling.rng import default_rng

from .ops_logging import initialization_logging

from .treelogic import TreeMixin

from openpathsampling.deprecations import deprecate, has_deprecations

from openpathsampling.deprecations import (SAMPLE_DETAILS, MOVE_DETAILS,

NEW_SNAPSHOT_KWARG_SELECTOR)

from future.utils import with_metaclass

logger = logging.getLogger(__name__)

init_log = logging.getLogger('openpathsampling.initialization')

# TODO: Remove if really not used anymore

# otherwise might move to utils or tools

def make_list_of_pairs(inlist):

"""

Converts input from several possible formats into a list of pairs: used

to clean input for swap-like moves.

Allowed input formats:

* flat list of length 2N

* list of pairs

* None (returns None)

Anything else will lead to a ValueError or AssertionError

Parameters

----------

inlist : list

input list, either flat list of length 2N, a list of pairs or None

Returns

-------

list of pairs

"""

if inlist is None:

return None

_ = len(inlist) # raises TypeError, avoids everything else

# based on first element, decide whether this should be a list of lists

# or a flat list

try:

_ = len(inlist[0])

list_of_lists = True

except TypeError:

list_of_lists = False

if list_of_lists:

for elem in inlist:

assert len(elem) == 2, "List of lists: inner list length != 2"

outlist = inlist

else:

assert len(inlist) % 2 == 0, \

"Flattened list: length not divisible by 2"

outlist = [

[a, b] for (a, b) in zip(

inlist[slice(0, None, 2)],

inlist[slice(1, None, 2)])

]

# Note that one thing we don't check is whether the items are of the

# same type. That might be worth doing someday; for now, we trust that

# part to work.

return outlist

class SampleNaNError(Exception):

def __init__(self, message, trial_sample, details):

super(SampleNaNError, self).__init__(message)

self.trial_sample = trial_sample

self.details = details

class SampleMaxLengthError(Exception):

def __init__(self, message, trial_sample, details):

super(SampleMaxLengthError, self).__init__(message)

self.trial_sample = trial_sample

self.details = details

class MoveChangeNaNError(Exception):

pass

class PathMover(with_metaclass(abc.ABCMeta, TreeMixin, StorableNamedObject)):

"""

A PathMover is the description of a move in replica space.

Notes

-----

A pathmover takes a SampleSet() and returns MoveChange() that is

used to change the old SampleSet() to the new one.

SampleSet1 + MoveChange1 => SampleSet2

A MoveChange is effectively a list of Samples. The change acts upon

a SampleSet by replacing existing Samples in the same ensemble

sequentially.

SampleSet({samp1(ens1), samp2(ens2), samp3(ens3)}) +

MoveChange([samp4(ens2)])

=> SampleSet({samp1(ens1), samp4(ens2), samp3(ens3)})

Note, that a SampleSet is an unordered list (or a set). Hence the ordering

in the example is arbitrary.

Potential future change: `engine` is not needed for all PathMovers

(replica exchange, ensemble hopping, path reversal, and moves which

combine these [state swap] have no need for the engine). Maybe that

should be moved into only the ensembles that need it? ~~~DWHS

Also, I agree with the separating trial and acceptance. We might choose

to use a different acceptance criterion than Metropolis. For example,

the "waste recycling" approach recently re-discovered by Frenkel (see

also work by Athenes, Jourdain, and old work by Kalos) might be

interesting. I think the best way to do this is to keep the acceptance

in the PathMover, but have it be a separate class ~~~DWHS

"""

# __metaclass__ = abc.ABCMeta

def __init__(self):

StorableNamedObject.__init__(self)

self._rng = default_rng()

self._in_ensembles = None

self._out_ensembles = None

self._len = None

self._inout = None

self._trust_candidate = False

# initialization_logging(logger=init_log, obj=self,

# entries=['ensembles'])

_is_ensemble_change_mover = None

@property

def is_ensemble_change_mover(self):

if self._is_ensemble_change_mover is None:

return False

else:

return self._is_ensemble_change_mover

_is_canonical = None

@property

def is_canonical(self):

return self._is_canonical

@property

def default_name(self):

return self.__class__.__name__[:-5]

# +-------------------------------------------------------------------------

# | tree implementation overrides

# +-------------------------------------------------------------------------

@property

def _subnodes(self):

return self.submovers

@property

def identifier(self):

return self

@staticmethod

def _default_match(original, test):

if isinstance(test, paths.PathMover):

return original is test

elif issubclass(test, paths.PathMover):

return original.__class__ is test

else:

return False

@property

def submovers(self):

"""

Returns a list of submovers

Returns

-------

list of openpathsampling.PathMover

the list of sub-movers

"""

return []

@staticmethod

def _flatten(ensembles):

if type(ensembles) is list:

return [s for ens in ensembles for s in PathMover._flatten(ens)]

else:

return [ensembles]

# +-------------------------------------------------------------------------

# | analyze effects of sample sets

# +-------------------------------------------------------------------------

def move_replica_state(self, replica_states):

return self.in_out.move(replica_states)

def sub_replica_state(self, replica_states):

"""

Return set of replica states that a submover might be called with

Parameters

----------

replica_states : set of `openpathsampling.pathmover_inout.ReplicaState`

Returns

-------

list of set of `ReplicaState`

"""

return [replica_states] * len(self.submovers)

def _generate_in_out(self):

if len(self.output_ensembles) == 0:

return {

InOutSet([])

}

elif (len(self.input_ensembles) == len(self.output_ensembles) == 1):

in_ens = self.input_ensembles[0]

out_ens = self.output_ensembles[0]

return InOutSet([InOut([((in_ens, out_ens, 0), 1)])])

else:

# Fallback could be all possibilities, but for now we ask the user!

raise NotImplementedError(

'Please implement the in-out-matrix for this mover.')

@property

def in_out(self):

"""

List the input -> output relation for ensembles

A mover will pick one or more replicas from specific ensembles.

Alter them (or not) and place these (or additional ones) in specific

ensembles. This relation can be visualized as a mapping of input to

output ensembles. Like

ReplicaExchange

ens1 -> ens2

ens2 -> ens1

EnsembleHop (A sample in ens1 will disappear and appear in ens2)

ens1 -> ens2

DuplicateMover (create a copy with a new replica number) Not used yet!

ens1 -> ens1

None -> ens1

Returns

-------

list of list of tuple : (:obj:`openpathsampling.Ensemble`,

:obj:`openpathsampling.Ensemble`)

a list of possible lists of tuples of ensembles.

Notes

-----

The default implementation will

(1) in case of a single input and output connect the two,

(2) return nothing if there are no out_ensembles and

(3) for more then two require implementation

"""

if self._inout is None:

self._inout = self._generate_in_out()

return self._inout

def _ensemble_signature(self, as_set=False):

"""Return tuple form of (input_ensembles, output_ensembles).

Useful for MoveScheme, e.g., identifying which movers should be

removed as part of a replacement.

"""

inp = tuple(self.input_ensembles)

out = tuple(self.output_ensembles)

if as_set:

inp = set(inp)

out = set(out)

return inp, out

@property

def ensemble_signature(self):

return self._ensemble_signature()

@property

def ensemble_signature_set(self):

return self._ensemble_signature(as_set=True)

@property

def input_ensembles(self):

"""Return a list of possible used ensembles for this mover

This list contains all Ensembles from which this mover might pick

samples. This is very useful to determine on which ensembles a

mover acts for analysis and sanity checking.

Returns

-------

list of :class:`openpathsampling.Ensemble`

the list of input ensembles

"""

if self._in_ensembles is None:

ensembles = self._get_in_ensembles()

self._in_ensembles = list(set(self._flatten(ensembles)))

return self._in_ensembles

@property

def output_ensembles(self):

"""Return a list of possible returned ensembles for this mover

This list contains all Ensembles for which this mover might return

samples. This is very useful to determine on which ensembles a

mover affects in later steps for analysis and sanity checking.

Returns

-------

list of Ensemble

the list of output ensembles

"""

if self._out_ensembles is None:

ensembles = self._get_out_ensembles()

self._out_ensembles = list(set(self._flatten(ensembles)))

return self._out_ensembles

def _get_in_ensembles(self):

"""Function that computes the list of input ensembles

"""

return []

def _get_out_ensembles(self):

"""Function that computes the list of output ensembles

Default is the same as in_ensembles

"""

return self._get_in_ensembles()

@staticmethod

def legal_sample_set(sample_set, ensembles=None, replicas='all'):

"""

This returns all the samples from sample_set which are in both

self.replicas and the parameter ensembles. If ensembles is None, we

use self.ensembles. If you want all ensembles allowed, pass

ensembles='all'.

Parameters

----------

sample_set : `openpathsampling.SampleSet`

the sampleset from which to pick specific samples matching certain

criteria

ensembles : list of `openpathsampling.Ensembles`

the ensembles to pick from

replicas : list of int or `all`

the replicas to pick or `'all'` for all

"""

mover_replicas = sample_set.replica_list()

if replicas == 'all':

selected_replicas = sample_set.replica_list()

else:

selected_replicas = replicas

reps = list(set(mover_replicas) & set(selected_replicas))

rep_samples = []

for rep in reps:

rep_samples.extend(sample_set.all_from_replica(rep))

# logger.debug("ensembles = " + str([ensembles]))

# logger.debug("self.ensembles = " + str(self.ensembles))

if ensembles is None:

ensembles = 'all'

if ensembles == 'all':

legal_samples = rep_samples

else:

ens_samples = []

if type(ensembles) is not list:

ensembles = [ensembles]

for ens in ensembles:

# try:

# ens_samples.extend(sample_set.all_from_ensemble(ens[0]))

# except TypeError:

ens_samples.extend(sample_set.all_from_ensemble(ens))

legal_samples = list(set(rep_samples) & set(ens_samples))

return legal_samples

@staticmethod

def select_sample(sample_set, ensembles=None, replicas=None):

"""

Returns one of the legal samples given self.replica and the ensemble

set in ensembles.

Parameters

----------

sample_set : `openpathsampling.SampleSet`

the sampleset from which to pick specific samples matching certain

criteria

ensembles : list of `openpathsampling.Ensembles` or `None`

the ensembles to pick from or `None` for all

replicas : list of int or None

the replicas to pick or `None` for all

"""

if replicas is None:

replicas = 'all'

logger.debug(

"replicas: " + str(replicas) + " ensembles: " + repr(ensembles))

legal = PathMover.legal_sample_set(sample_set, ensembles, replicas)

for sample in legal:

logger.debug(

"legal: (" + str(sample.replica) +

"," + str(sample.trajectory) +

"," + repr(sample.ensemble) +

")")

# TODO: This can't go through numpy.random as Samples unwrap to

# Snapshots when cast to arrays

selected = random.choice(legal)

logger.debug(

"selected sample: (" + str(selected.replica) +

"," + str(selected.trajectory) +

"," + repr(selected.ensemble) +

")")

return selected

@abc.abstractmethod

def move(self, sample_set):

"""

Run the generation starting with the initial sample_set specified.

Parameters

----------

sample_set : SampleSet

the initially used sampleset

Returns

-------

samples : MoveChange

the MoveChange instance describing the change from the old to

the new SampleSet

"""

return paths.EmptyMoveChange() # pragma: no cover

def __str__(self):

if self.name == self.__class__.__name__:

return self.__repr__()

else:

return self.name

class IdentityPathMover(PathMover):

"""

The simplest Mover that does nothing !

Notes

-----

Since is does nothing it is considered rejected everytime!

It can be used to test function of PathMover

Parameters

----------

counts_as_trial : bool

Whether this mover should count as a trial or not. If `True`, the

`EmptyMoveChange` returned includes this mover, which means it gets

counted as a trial in analysis of acceptance. If `False` (default),

the mover for the returned move change is `None`, which does not get

counted as a trial.

"""

def __init__(self, counts_as_trial=False):

super(IdentityPathMover, self).__init__()

self.counts_as_trial = counts_as_trial

def move(self, sample_set):

mover = self if self.counts_as_trial else None

return paths.EmptyMoveChange(mover=mover)

###############################################################################

# GENERATORS

###############################################################################

class SampleMover(PathMover):

def __init__(self):

super(SampleMover, self).__init__()

def metropolis(self, trials):

"""Implements the Metropolis acceptance for a list of trial samples

The Metropolis uses the .bias for each sample and checks of samples

are valid - are in the proposed ensemble. This will give an acceptance

probability for all samples. If the product is smaller than a random

number the change will be accepted.

Parameters

----------

trials : list of openpathsampling.Sample

the list of all samples to be applied in a change.

Returns

-------

bool

True if the trial is accepted, False otherwise

details : openpathsampling.Details

Returns a Details object that contains information about the

decision, i.e. total acceptance and random number

"""

shoot_str = "MC in {cls} using samples {trials}"

logger.info(shoot_str.format(cls=self.__class__.__name__,

trials=trials))

trial_dict = dict()

for trial in trials:

trial_dict[trial.ensemble] = trial

accepted = True

probability = 1.0

# TODO: This isn't right. `bias` should be associated with the

# change; not with each individual sample. ~~~DWHS

for ens, sample in trial_dict.items():

valid = ens(sample.trajectory, candidate=self._trust_candidate)

if not valid:

# one sample not valid reject

accepted = False

probability = 0.0

break

else:

probability *= sample.bias

rand = self._rng.random()

if rand > probability:

# rejected

accepted = False

details = {

'metropolis_acceptance': probability,

'metropolis_random': rand

}

if accepted:

result_str = "accepted"

else:

result_str = ("rejected. Acceptance probabilty "

+ str(probability))

logger.info("Trial was " + result_str)

return accepted, details

@property

def submovers(self):

# Movers do not have submovers!

return []

def _called_ensembles(self):

"""Function to determine which ensembles to pick samples from

Returns

-------

list of Ensemble

the list of ensembles. Samples can then be selected using

PathMover.select_sample

"""

# Default is that the list of ensembles is in self.ensembles

return []

def get_samples_from_sample_set(self, sample_set):

"""

Select samples to use as input to the move core.

See Also

--------

move_core

move

Parameters

----------

sample_set : :class:`.SampleSet`

current samples to use as potential input

Returns

-------

list of :class:`.Sample`

samples to use as input to the move core

"""

ensembles = self._called_ensembles()

samples = [self.select_sample(sample_set, ens) for ens in ensembles]

return samples

def move(self, sample_set):

samples = self.get_samples_from_sample_set(sample_set)

change = self.move_core(samples)

return change

def move_core(self, samples):

"""Core of the Monte Carlo move. Includes acceptance.

See Also

--------

move

Parameters

----------

samples : list of :class:`.Sample`

input samples from the correct ensembles of this object

Returns

-------

:class:`.MoveChange`

result MoveChange for this move

"""

# this is separated out for reuse and remove dependence core MC move

# dependence on the entire sample set (for parallelization)

try:

# pass these samples to the trial move which might throw

# engine-specific exceptions if something goes wrong.

# Most common should be `EngineNaNError` if nan is detected and

# `EngineMaxLengthError`

trials, call_details = self(*samples)

except SampleNaNError as e:

e.details.update({'rejection_reason': 'nan'})

return paths.RejectedNaNSampleMoveChange(

samples=e.trial_sample,

mover=self,

input_samples=samples,

details=paths.Details(**e.details)

)

except SampleMaxLengthError as e:

e.details.update({'rejection_reason': 'max_length'})

return paths.RejectedMaxLengthSampleMoveChange(

samples=e.trial_sample,

mover=self,

input_samples=samples,

details=paths.Details(**e.details)

)

accepted, acceptance_details = self._accept(trials)

# update details

kwargs = {}

kwargs.update(call_details)

kwargs.update(acceptance_details)

details = Details(**kwargs)

# return change

if accepted:

return paths.AcceptedSampleMoveChange(

samples=trials,

mover=self,

input_samples=samples,

details=details

)

else:

return paths.RejectedSampleMoveChange(

samples=trials,

mover=self,

input_samples=samples,

details=details

)

@abc.abstractmethod

def __call__(self, *args):

"""Generate trial samples directly

PathMovers can also be called directly with a list of samples that are

then used to generate new samples. If the Mover is used as a move

the move will first determine the input samples and then pass these to

this function

"""

# Default is that the original samples are returned

return args

def _accept(self, trials):

"""Function to determine the acceptance of a trial

Defaults to calling the Metropolis acceptance criterion for all

returned trial samples. Means all samples most be valid and accepted.

"""

return self.metropolis(trials)

###############################################################################

# SHOOTING GENERATORS

###############################################################################

class EngineMover(SampleMover):

"""Baseclass for Movers that use an engine

Notes

-----

A few comments for developers working with subclasses of

``EngineMover``: This class is intended to do most of the grunt work for

a wide range of possible engine-based needs. Remember that your

``selector`` can select first or final points, e.g., to extend a move.

In order to help you find your way through the ``EngineMover`` code,

here is an overview of what various private methods do:

* ``__call__``: Creates the trial. Two steps: (1) make the trajectory;

(2) assemble a sample to return

* ``_build_sample``: assembles the final sample

* ``_make_forward_trajectory``/``_make_backward_trajectory``: creates

the actual trajectory, using :class:`.PrefixTrajectoryEnsemble` or

:class:`.SuffixTrajectoryEnsemble` to ensure reasonable behavior (see

below for further discussion)

* ``._run``: this is what is called by ``__call__``, and it in turn

calls the functions to make the trajectories (depending on the nature

of the mover). Frequently, this is the only thing to override (two-way

shooting, shifting).

"""

default_engine = None

reject_max_length = True

# this will store the engine attribute for all subclasses as well

_included_attr = ['_engine']

def __init__(self, ensemble, target_ensemble, selector, engine=None,

modifier=None):

super(EngineMover, self).__init__()

self.selector = selector

self.ensemble = ensemble

self.target_ensemble = target_ensemble

self._engine = engine

self.modifier = modifier or paths.NoModification(as_copy=False)

self._trust_candidate = True # can I safely do that?

# I think that is safe. Note for future: if we come across a bug

# based on this, an alternative would be to have the move strategy

# set _trust_candidate when it builds the movers; that is likely to

# be a little safer (although I think we can trust all candidates

# from engine movers to actually be candidates)

def to_dict(self):

dct = super(EngineMover, self).to_dict()

dct['engine'] = self.engine

return dct

@property

def engine(self):

if self._engine is not None:

return self._engine

else:

return self.default_engine

@engine.setter

def engine(self, engine):

self._engine = engine

def _called_ensembles(self):

return [self.ensemble]

def _get_in_ensembles(self):

return [self.ensemble]

def _get_out_ensembles(self):

return [self.target_ensemble]

def __call__(self, input_sample):

initial_trajectory = input_sample.trajectory

shooting_index = self.selector.pick(initial_trajectory)

try:

trial_trajectory, run_details = self._run(initial_trajectory,

shooting_index)

except paths.engines.EngineNaNError as e:

trial, details = self._build_sample(

input_sample, shooting_index, e.last_trajectory, 'nan')

raise SampleNaNError('Sample with NaN', trial, details)

except paths.engines.EngineMaxLengthError as e:

trial, details = self._build_sample(

input_sample, shooting_index, e.last_trajectory, 'max_length')

if EngineMover.reject_max_length:

raise SampleMaxLengthError('Sample with MaxLength', trial,

details)

else:

trial, details = self._build_sample(

input_sample, shooting_index, trial_trajectory,

run_details=run_details)

trials = [trial]

details.update(run_details)

return trials, details

def _build_sample(

self,

input_sample,

shooting_index,

trial_trajectory,

stopping_reason=None,

run_details={}

):

# TODO OPS 2.0: the passing of run_details is a hack and should be

# properly refactored

initial_trajectory = input_sample.trajectory

if stopping_reason is None:

bias = 1.0

old_snapshot = initial_trajectory[shooting_index]

new_snapshot = run_details.get("modified_shooting_snapshot",

old_snapshot)

# Selector bias

try:

bias *= self.selector.probability_ratio(

initial_trajectory[shooting_index],

initial_trajectory,

trial_trajectory,

new_snapshot=new_snapshot

)

except TypeError:

bias *= self.selector.probability_ratio(

initial_trajectory[shooting_index],

initial_trajectory,

trial_trajectory)

NEW_SNAPSHOT_KWARG_SELECTOR.warn()

# Modifier bias

bias *= self.modifier.probability_ratio(old_snapshot, new_snapshot)

else:

bias = 0.0

# temporary test to make sure nothing went weird

# old_bias = initial_point.sum_bias / trial_point.sum_bias

# assert(abs(bias - old_bias) < 10e-6)

# assert(initial_trajectory[shooting_index] in trial_trajectory)

# we need to save the initial

trial_details = {

'initial_trajectory': initial_trajectory,

'shooting_snapshot': initial_trajectory[shooting_index]

}

if stopping_reason is not None:

trial_details['stopping_reason'] = stopping_reason

trial = paths.Sample(

replica=input_sample.replica,

trajectory=trial_trajectory,

ensemble=self.target_ensemble,

parent=input_sample,

mover=self,

bias=bias

)

return trial, trial_details

def _make_forward_trajectory(self, trajectory, shooting_index):

initial_snapshot = trajectory[shooting_index] # .copy()

run_f = paths.PrefixTrajectoryEnsemble(self.target_ensemble,

trajectory[0:shooting_index]

).can_append

partial_trajectory = self.engine.generate(initial_snapshot,

running=[run_f])

trial_trajectory = (trajectory[0:shooting_index] +

partial_trajectory)

# TODO: this should check for overshoot; only works now if ensemble

# doesn't overshoot

return trial_trajectory

def _make_backward_trajectory(self, trajectory, shooting_index):

initial_snapshot = trajectory[shooting_index].reversed # _copy()

run_f = paths.SuffixTrajectoryEnsemble(self.target_ensemble,

trajectory[shooting_index + 1:]

).can_prepend

partial_trajectory = self.engine.generate(initial_snapshot,

running=[run_f])

trial_trajectory = (partial_trajectory.reversed +

trajectory[shooting_index + 1:])

# TODO: this should check for overshoot; only works now if ensemble

# doesn't overshoot

return trial_trajectory

# direction is an abstract property to disallow instantiation

# of the EngineMover unless we use a concrete subclass that sets this.

# This is not super elegant but is the way to do it with abstract classes

@abc.abstractproperty

def direction(self):

return 'unknown'

def _run(self, trajectory, shooting_index):

"""Takes initial trajectory and shooting point; return trial

trajectory"""

shoot_str = "Running {sh_dir} from frame {fnum} in [0:{maxt}]"

logger.info(shoot_str.format(

fnum=shooting_index,

maxt=len(trajectory) - 1,

sh_dir=self.direction

))

if self.direction == "forward":

trial_trajectory = self._make_forward_trajectory(

trajectory, shooting_index

)

elif self.direction == "backward":

trial_trajectory = self._make_backward_trajectory(

trajectory, shooting_index

)

else:

raise RuntimeError("Unknown direction: " + str(self.direction))

return trial_trajectory, {}

class ForwardShootMover(EngineMover):

"""A forward shooting sample generator

"""

def __init__(self, ensemble, selector, engine=None):

super(ForwardShootMover, self).__init__(

ensemble=ensemble,

target_ensemble=ensemble,

selector=selector,

engine=engine

)

@property

def direction(self):

return 'forward'

class BackwardShootMover(EngineMover):

"""A Backward shooting generator

"""

def __init__(self, ensemble, selector, engine=None):

super(BackwardShootMover, self).__init__(

ensemble=ensemble,

target_ensemble=ensemble,

selector=selector,

engine=engine

)

@property

def direction(self):

return 'backward'

class ForwardExtendMover(EngineMover):

"""

A Sample Mover implementing Forward Extension

"""

_direction = "forward"

def __init__(self, ensemble, target_ensemble, engine=None):

super(ForwardExtendMover, self).__init__(

ensemble=ensemble,

target_ensemble=target_ensemble,

selector=paths.FinalFrameSelector(),

engine=engine

)

@property

def direction(self):

return 'forward'

class BackwardExtendMover(EngineMover):

"""

A Sample Mover implementing Backward Extension

"""

_direction = "backward"

def __init__(self, ensemble, target_ensemble, engine=None):

super(BackwardExtendMover, self).__init__(