<commit_before>import logging

import numpy as np

import openpathsampling as paths
from openpathsampling.numerics import (
    Histogram, histograms_to_pandas_dataframe, LookupFunction, Histogrammer
)
from openpathsampling.numerics import WHAM
from openpathsampling.netcdfplus import StorableNamedObject

from openpathsampling.analysis.tools import (
    pathlength, max_lambdas, guess_interface_lambda, minus_sides_summary,
    sampleset_sample_generator
)

logger = logging.getLogger(__name__)


class Transition(StorableNamedObject):
    """
    Describes (in general) a transition between two states.
    """
    def __init__(self, stateA, stateB):
        super(Transition, self).__init__()
        self.stateA = stateA
        self.stateB = stateB
        # whoops: can't be set here, but must be set in subclass
        # TODO: make that work in a more sensible way
        #self.ensembles = []

    @property
    def all_ensembles(self):
        return self.ensembles

    def to_dict(self):
        return {
            'stateA' : self.stateA,
            'stateB' : self.stateB,
        }

    @classmethod
    def from_dict(cls, dct):
        return Transition(
            stateA=dct['stateA'],
            stateB=dct['stateB']
        )


class TPSTransition(Transition):
    """
    Transition using TPS ensembles
    """
    def __init__(self, stateA, stateB, name=None):
        super(TPSTransition, self).__init__(stateA, stateB)
        if name is not None:
            self.name = name
        if not hasattr(self, "ensembles"):
            self.ensembles = [self._tps_ensemble(stateA, stateB)]

    def to_dict(self):
        return {
            'stateA' : self.stateA,
            'stateB' : self.stateB,
            'ensembles' : self.ensembles,
            'name' : self.name
        }

    @classmethod
    def from_dict(cls, dct):
        mytrans = TPSTransition(dct['stateA'], dct['stateB'], dct['name'])
        mytrans.ensembles = dct['ensembles']
        return mytrans

    def _tps_ensemble(self, stateA, stateB):
        return paths.SequentialEnsemble([
            paths.AllInXEnsemble(stateA) & paths.LengthEnsemble(1),
            paths.AllOutXEnsemble(stateA | stateB),
            paths.AllInXEnsemble(stateB) & paths.LengthEnsemble(1)
        ])

    def add_transition(self, stateA, stateB):
        new_ens = self._tps_ensemble(stateA, stateB)
        try:
            self.ensembles[0] = self.ensembles[0] | new_ens
        except AttributeError:
            self.ensembles = [new_ens]


class FixedLengthTPSTransition(TPSTransition):
    """Transition using fixed length TPS ensembles"""
    def __init__(self, stateA, stateB, length, name=None):
        self.length = length
        super(FixedLengthTPSTransition, self).__init__(stateA, stateB, name)

    def to_dict(self):
        dct = super(FixedLengthTPSTransition, self).to_dict()
        dct['length'] = self.length
        return dct

    @classmethod
    def from_dict(cls, dct):
        mytrans = super(FixedLengthTPSTransition, cls).from_dict(dct)
        mytrans.length = dct['length']
        return mytrans

    def _tps_ensemble(self, stateA, stateB):
        return paths.SequentialEnsemble([
            paths.LengthEnsemble(1) & paths.AllInXEnsemble(stateA),
            paths.LengthEnsemble(self.length - 2),
            paths.LengthEnsemble(1) & paths.AllInXEnsemble(stateB)
        ])


class TISTransition(Transition):
    """
    Transition using TIS ensembles.

    The additional information from the TIS ensembles allows us to set up
    all the analysis (assuming we built these are proper TIS ensembles,
    which we DO in the intitialization!)

    Parameters
    ----------
    stateA : Volume
        Volume for the state from which the transition begins
    stateB : Volume
        Volume for the state in which the transition ends
    interfaces : list of Volume
        Volumes for the interfaces
    orderparameter : CollectiveVariable
        order parameter to be used in the analysis (does not need to be the
        parameter which defines the interfaces, although it usually is)
    name : string
        name for the transition

    """

    def __init__(self, stateA, stateB, interfaces, orderparameter=None, name=None):
        super(TISTransition, self).__init__(stateA, stateB)

        self.stateA = stateA
        self.stateB = stateB
        self.interfaces = interfaces
        if name is not None:
            self.name = name

        # If we reload from a storage file, we want to use the
        # ensembles from the file, not the automatically generated
        # ones here

        # build ensembles if we don't already have them
        self.orderparameter = orderparameter
        if not hasattr(self, "ensembles"):
            self.build_ensembles(self.stateA, self.stateB,
                                 self.interfaces, self.orderparameter)

        self.default_orderparameter = self.orderparameter

        self.total_crossing_probability_method = "wham"
        self.histograms = {}
        # caches for the results of our calculation
        self._flux = None
        self._rate = None

        self.hist_args = {} # shortcut to ensemble_histogram_info[].hist_args
        self.ensemble_histogram_info = {
            'max_lambda' : Histogrammer(
                f=max_lambdas,
                f_args={'orderparameter' : self.orderparameter},
                hist_args={}
            ),
            'pathlength' : Histogrammer(
                f=pathlength,
                f_args={},
                hist_args={}
            )
        }

        self.minus_ensemble = paths.MinusInterfaceEnsemble(
            state_vol=stateA,
            innermost_vols=interfaces[0]
        ).named("Out " + stateA.name + " minus")

    def copy(self, with_results=True):
        copy = self.from_dict(self.to_dict())
        copy.copy_analysis_from(self)
        return copy

    def copy_analysis_from(self, other):
        self.default_orderparameter = other.default_orderparameter
        self.total_crossing_probability_method = other.total_crossing_probability_method
        self.hist_args = other.hist_args
        self.ensemble_histogram_info = other.ensemble_histogram_info
        self.histograms = other.histograms
        self._flux = other._flux
        self._rate = other._rate
        try:
            self.tcp = other.tcp
        except AttributeError:
            pass
        try:
            self.ctp = other.ctp
        except AttributeError:
            pass


    def __str__(self):
        mystr = str(self.__class__.__name__) + ": " + str(self.name) + "\n"
        mystr += (str(self.stateA.name) + " -> " + str(self.stateA.name)
                  + " or " + str(self.stateB.name) + "\n")
        for iface in self.interfaces:
            mystr += "Interface: " + str(iface.name) + "\n"
        return mystr


    def build_ensembles(self, stateA, stateB, interfaces, orderparameter):
        self.ensembles = paths.EnsembleFactory.TISEnsembleSet(
            stateA, stateB, self.interfaces, orderparameter
        )
        for ensemble in self.ensembles:
            ensemble.named(self.name + " " +
                           str(self.ensembles.index(ensemble)))


    # parameters for different types of output
    def ensemble_statistics(self, ensemble, samples, weights=None, force=False):
        """Calculate stats for a given ensemble: path length, crossing prob

        In general we do all of these at once because the extra cost of
        running through the samples twice is worse than doing the extra
        calculations.

        Parameters
        ----------
        ensemble: Ensemble
        samples : iterator over samples
        """
        # figure out which histograms need to updated for this ensemble
        run_it = []
        if not force:
            # TODO figure out which need to be rerun
            pass
        else:
            run_it = self.ensemble_histogram_info.keys()

        for hist in run_it:
            hist_info = self.ensemble_histogram_info[hist]
            if hist_info.hist_args == {} and self.hist_args[hist] != {}:
                hist_info.hist_args = self.hist_args[hist]

            if hist not in self.histograms.keys():
                self.histograms[hist] = {}
            self.histograms[hist][ensemble] = Histogram(**(hist_info.hist_args))

        in_ens_samples = (s for s in samples if s.ensemble is ensemble)
        hist_data = {}
        buflen = -1
        sample_buf = []
        for sample in in_ens_samples:
            for hist in run_it:
                hist_info = self.ensemble_histogram_info[hist]
                hist_data_sample = hist_info.f(sample, **hist_info.f_args)
                try:
                    hist_data[hist].append(hist_data_sample)
                except KeyError:
                    hist_data[hist] = [hist_data_sample]


        for hist in run_it:<commit_msg>Cache previous to speed up histograms in TIS<commit_after>import logging

import numpy as np

import openpathsampling as paths
from openpathsampling.numerics import (
    Histogram, histograms_to_pandas_dataframe, LookupFunction, Histogrammer
)
from openpathsampling.numerics import WHAM
from openpathsampling.netcdfplus import StorableNamedObject

from openpathsampling.analysis.tools import (
    pathlength, max_lambdas, guess_interface_lambda, minus_sides_summary,
    sampleset_sample_generator
)

logger = logging.getLogger(__name__)


class Transition(StorableNamedObject):
    """
    Describes (in general) a transition between two states.
    """
    def __init__(self, stateA, stateB):
        super(Transition, self).__init__()
        self.stateA = stateA
        self.stateB = stateB
        # whoops: can't be set here, but must be set in subclass
        # TODO: make that work in a more sensible way
        #self.ensembles = []

    @property
    def all_ensembles(self):
        return self.ensembles

    def to_dict(self):
        return {
            'stateA' : self.stateA,
            'stateB' : self.stateB,
        }

    @classmethod
    def from_dict(cls, dct):
        return Transition(
            stateA=dct['stateA'],
            stateB=dct['stateB']
        )


class TPSTransition(Transition):
    """
    Transition using TPS ensembles
    """
    def __init__(self, stateA, stateB, name=None):
        super(TPSTransition, self).__init__(stateA, stateB)
        if name is not None:
            self.name = name
        if not hasattr(self, "ensembles"):
            self.ensembles = [self._tps_ensemble(stateA, stateB)]

    def to_dict(self):
        return {
            'stateA' : self.stateA,
            'stateB' : self.stateB,
            'ensembles' : self.ensembles,
            'name' : self.name
        }

    @classmethod
    def from_dict(cls, dct):
        mytrans = TPSTransition(dct['stateA'], dct['stateB'], dct['name'])
        mytrans.ensembles = dct['ensembles']
        return mytrans

    def _tps_ensemble(self, stateA, stateB):
        return paths.SequentialEnsemble([
            paths.AllInXEnsemble(stateA) & paths.LengthEnsemble(1),
            paths.AllOutXEnsemble(stateA | stateB),
            paths.AllInXEnsemble(stateB) & paths.LengthEnsemble(1)
        ])

    def add_transition(self, stateA, stateB):
        new_ens = self._tps_ensemble(stateA, stateB)
        try:
            self.ensembles[0] = self.ensembles[0] | new_ens
        except AttributeError:
            self.ensembles = [new_ens]


class FixedLengthTPSTransition(TPSTransition):
    """Transition using fixed length TPS ensembles"""
    def __init__(self, stateA, stateB, length, name=None):
        self.length = length
        super(FixedLengthTPSTransition, self).__init__(stateA, stateB, name)

    def to_dict(self):
        dct = super(FixedLengthTPSTransition, self).to_dict()
        dct['length'] = self.length
        return dct

    @classmethod
    def from_dict(cls, dct):
        mytrans = super(FixedLengthTPSTransition, cls).from_dict(dct)
        mytrans.length = dct['length']
        return mytrans

    def _tps_ensemble(self, stateA, stateB):
        return paths.SequentialEnsemble([
            paths.LengthEnsemble(1) & paths.AllInXEnsemble(stateA),
            paths.LengthEnsemble(self.length - 2),
            paths.LengthEnsemble(1) & paths.AllInXEnsemble(stateB)
        ])


class TISTransition(Transition):
    """
    Transition using TIS ensembles.

    The additional information from the TIS ensembles allows us to set up
    all the analysis (assuming we built these are proper TIS ensembles,
    which we DO in the intitialization!)

    Parameters
    ----------
    stateA : Volume
        Volume for the state from which the transition begins
    stateB : Volume
        Volume for the state in which the transition ends
    interfaces : list of Volume
        Volumes for the interfaces
    orderparameter : CollectiveVariable
        order parameter to be used in the analysis (does not need to be the
        parameter which defines the interfaces, although it usually is)
    name : string
        name for the transition

    """

    def __init__(self, stateA, stateB, interfaces, orderparameter=None, name=None):
        super(TISTransition, self).__init__(stateA, stateB)

        self.stateA = stateA
        self.stateB = stateB
        self.interfaces = interfaces
        if name is not None:
            self.name = name

        # If we reload from a storage file, we want to use the
        # ensembles from the file, not the automatically generated
        # ones here

        # build ensembles if we don't already have them
        self.orderparameter = orderparameter
        if not hasattr(self, "ensembles"):
            self.build_ensembles(self.stateA, self.stateB,
                                 self.interfaces, self.orderparameter)

        self.default_orderparameter = self.orderparameter

        self.total_crossing_probability_method = "wham"
        self.histograms = {}
        # caches for the results of our calculation
        self._flux = None
        self._rate = None

        self.hist_args = {} # shortcut to ensemble_histogram_info[].hist_args
        self.ensemble_histogram_info = {
            'max_lambda' : Histogrammer(
                f=max_lambdas,
                f_args={'orderparameter' : self.orderparameter},
                hist_args={}
            ),
            'pathlength' : Histogrammer(
                f=pathlength,
                f_args={},
                hist_args={}
            )
        }

        self.minus_ensemble = paths.MinusInterfaceEnsemble(
            state_vol=stateA,
            innermost_vols=interfaces[0]
        ).named("Out " + stateA.name + " minus")

    def copy(self, with_results=True):
        copy = self.from_dict(self.to_dict())
        copy.copy_analysis_from(self)
        return copy

    def copy_analysis_from(self, other):
        self.default_orderparameter = other.default_orderparameter
        self.total_crossing_probability_method = other.total_crossing_probability_method
        self.hist_args = other.hist_args
        self.ensemble_histogram_info = other.ensemble_histogram_info
        self.histograms = other.histograms
        self._flux = other._flux
        self._rate = other._rate
        try:
            self.tcp = other.tcp
        except AttributeError:
            pass
        try:
            self.ctp = other.ctp
        except AttributeError:
            pass


    def __str__(self):
        mystr = str(self.__class__.__name__) + ": " + str(self.name) + "\n"
        mystr += (str(self.stateA.name) + " -> " + str(self.stateA.name)
                  + " or " + str(self.stateB.name) + "\n")
        for iface in self.interfaces:
            mystr += "Interface: " + str(iface.name) + "\n"
        return mystr


    def build_ensembles(self, stateA, stateB, interfaces, orderparameter):
        self.ensembles = paths.EnsembleFactory.TISEnsembleSet(
            stateA, stateB, self.interfaces, orderparameter
        )
        for ensemble in self.ensembles:
            ensemble.named(self.name + " " +
                           str(self.ensembles.index(ensemble)))


    # parameters for different types of output
    def ensemble_statistics(self, ensemble, samples, weights=None, force=False):
        """Calculate stats for a given ensemble: path length, crossing prob

        In general we do all of these at once because the extra cost of
        running through the samples twice is worse than doing the extra
        calculations.

        Parameters
        ----------
        ensemble: Ensemble
        samples : iterator over samples
        """
        # figure out which histograms need to updated for this ensemble
        run_it = []
        if not force:
            # TODO figure out which need to be rerun
            pass
        else:
            run_it = self.ensemble_histogram_info.keys()

        for hist in run_it:
            hist_info = self.ensemble_histogram_info[hist]
            if hist_info.hist_args == {} and self.hist_args[hist] != {}:
                hist_info.hist_args = self.hist_args[hist]

            if hist not in self.histograms.keys():
                self.histograms[hist] = {}
            self.histograms[hist][ensemble] = Histogram(**(hist_info.hist_args))

        in_ens_samples = (s for s in samples if s.ensemble is ensemble)
        hist_data = {}
        buflen = -1
        sample_buf = []
        prev_sample = {h: None for h in run_it}
        prev_result = {h: None for h in run_it}
        for sample in in_ens_samples:
            for hist in run_it:
                hist_info = self.ensemble_histogram_info[hist]
                if sample is prev_sample[hist]:
                    hist_data_sample = prev_result[hist]
                else:
                    hist_data_sample = hist_info.f(sample,
                                                   **hist_info.f_args)
                prev_result[hist] = hist_data_sample
                prev_sample[hist] = sample
                try:
                    hist_data[hist].append(hist_data_sample)
                except KeyError:
                    hist_data[hist] = [hist_data_sample]


        for hist in run_it:
{'fraction_non_alphanumeric': 0.08619841486745013, 'fraction_numerical': 0.0009838753757857339, 'mean_word_length': 2.148511443813457, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1202231', 'n_tokens_mistral': 5406, 'n_tokens_neox': 5178, 'n_words': 1311}