<commit_before>"""
    Script that computes the absolute binding free energy of an MSM.
    It needs:
        1) Trajectory wildcard in order to compute cluster volume
        2) A discretized/clusterCenters.dat file with the cluster center
        3) A MSM_object.pkl obtained with pyemma in order to obtain the stationary distribution
        4) For the moment, it needs of a reweightingT, in order to do a histogram reweighting, but does not seem to work that well
"""
from __future__ import absolute_import, division, print_function, unicode_literals
from builtins import range
import glob
import sys
import itertools
import argparse
import numpy as np
try:
    import cPickle
except ImportError:
    import pickle as cPickle
from pyemma.coordinates.clustering import AssignCenters
from AdaptivePELE.freeEnergies import runMarkovChainModel as run
from AdaptivePELE.freeEnergies import utils


def assignNewTrajectories(trajs, clusterCenters):
    assign = AssignCenters(clusterCenters)
    dTrajs = assign.assign(trajs)
    return dTrajs


def expandTrajs(trajList):
    d = 0.5
    combinations = np.array([[0, 1, 0], [1, 0, 0], [0, 0, 1], [0, -1, 0],
                             [-1, 0, 0], [0, 0, -1]])
    return list(trajList+combinations*d)


def parseArgs():
    parser = argparse.ArgumentParser(description="Script that computes delta G")
    parser.add_argument('trajectory', type=str)
    parser.add_argument('-r', '--reweight', type=float, default=1000)
    args = parser.parse_args()
    return args.trajectory, args.reweight


def writePDB(pmf_xyzg, title="clusters.pdb"):
    # templateLine = "HETATM%s  H%sCLT L 502    %s%s%s  0.75%s           H\n"
    templateLine = "HETATM%s H%sCLT L 502    %s%s%s  0.75%s                H\n"

    content = ""
    for i, line in enumerate(pmf_xyzg):
        number = str(i).rjust(5)
        number3 = str(i).ljust(3)
        x = ("%.3f" % line[0]).rjust(8)
        y = ("%.3f" % line[1]).rjust(8)
        z = ("%.3f" % line[2]).rjust(8)
        g = ("%.3f" % line[-1]).rjust(8)

        content += templateLine % (number, number3, x, y, z, g)
    f = open(title, 'w')
    f.write(content)
    f.close()


def calcQ(T, gpmf):
    kb = 0.0019872041
    Q = np.sum(np.exp(-gpmf/(kb*T)))
    return Q


def reweightProbabilities(T, Torig, origProb):
    """
        Torig: temperature at which simulations were run
    """
    kb = 0.0019872041
    gpmf = -kb*Torig*np.log(origProb)
    gpmf[gpmf == -np.inf] = np.inf
    gpmf -= gpmf.min()  # It does not make any difference here

    print(gpmf[gpmf == np.inf])
    print(gpmf[gpmf == -np.inf])
    a = (T-Torig) / (kb * T * Torig)
    correction1 = np.exp(a * gpmf)
    beta = 1/kb/T
    betaOrig = 1/kb/Torig
    Qorig = np.sum(np.exp(-betaOrig*gpmf))
    Qnew = np.sum(np.exp(-beta*gpmf))
    correction2 = Qorig / Qnew

    return correction1 * correction2 * origProb


def loadMSM(MSMFile):
    with open(MSMFile, "rb") as MSMfile:
        MSM_object = cPickle.load(MSMfile)
    return MSM_object


def reestimate_transition_matrix(count_matrix):
    count_matrix += 1/float(count_matrix.shape[0])
    return utils.buildRevTransitionMatrix(count_matrix)


def ensure_connectivity(MSMObject, allClusters):
    if len(allClusters) == MSMObject.stationary_distribution.size:
        pi = MSMObject.stationary_distribution
        clusters = allClusters[MSMObject.connected_sets[0]]
    else:
        ######################
        # Reconstruct stationary distribution with pseudocounts to ensure
        # connectivity
        print("Adding pseudocounts to enforce connectivity")
        counts = MSMObject.count_matrix_full
        trans = reestimate_transition_matrix(counts)
        _, eic = run.getSortedEigen(trans)
        pi = run.getStationaryDistr(eic[:, 0])<commit_msg>Make import more precise in computeDeltaG<commit_after>"""
    Script that computes the absolute binding free energy of an MSM.
    It needs:
        1) Trajectory wildcard in order to compute cluster volume
        2) A discretized/clusterCenters.dat file with the cluster center
        3) A MSM_object.pkl obtained with pyemma in order to obtain the stationary distribution
        4) For the moment, it needs of a reweightingT, in order to do a histogram reweighting, but does not seem to work that well
"""
from __future__ import absolute_import, division, print_function, unicode_literals
from builtins import range
import glob
import sys
import itertools
import argparse
import numpy as np
try:
    import cPickle
except ImportError:
    import pickle as cPickle
from pyemma.coordinates.clustering import AssignCenters
from AdaptivePELE.freeEnergies import utils
from AdaptivePELE.freeEnergies.runMarkovChainModel import getStationaryDistr, getSortedEigen


def assignNewTrajectories(trajs, clusterCenters):
    assign = AssignCenters(clusterCenters)
    dTrajs = assign.assign(trajs)
    return dTrajs


def expandTrajs(trajList):
    d = 0.5
    combinations = np.array([[0, 1, 0], [1, 0, 0], [0, 0, 1], [0, -1, 0],
                             [-1, 0, 0], [0, 0, -1]])
    return list(trajList+combinations*d)


def parseArgs():
    parser = argparse.ArgumentParser(description="Script that computes delta G")
    parser.add_argument('trajectory', type=str)
    parser.add_argument('-r', '--reweight', type=float, default=1000)
    args = parser.parse_args()
    return args.trajectory, args.reweight


def writePDB(pmf_xyzg, title="clusters.pdb"):
    # templateLine = "HETATM%s  H%sCLT L 502    %s%s%s  0.75%s           H\n"
    templateLine = "HETATM%s H%sCLT L 502    %s%s%s  0.75%s                H\n"

    content = ""
    for i, line in enumerate(pmf_xyzg):
        number = str(i).rjust(5)
        number3 = str(i).ljust(3)
        x = ("%.3f" % line[0]).rjust(8)
        y = ("%.3f" % line[1]).rjust(8)
        z = ("%.3f" % line[2]).rjust(8)
        g = ("%.3f" % line[-1]).rjust(8)

        content += templateLine % (number, number3, x, y, z, g)
    f = open(title, 'w')
    f.write(content)
    f.close()


def calcQ(T, gpmf):
    kb = 0.0019872041
    Q = np.sum(np.exp(-gpmf/(kb*T)))
    return Q


def reweightProbabilities(T, Torig, origProb):
    """
        Torig: temperature at which simulations were run
    """
    kb = 0.0019872041
    gpmf = -kb*Torig*np.log(origProb)
    gpmf[gpmf == -np.inf] = np.inf
    gpmf -= gpmf.min()  # It does not make any difference here

    print(gpmf[gpmf == np.inf])
    print(gpmf[gpmf == -np.inf])
    a = (T-Torig) / (kb * T * Torig)
    correction1 = np.exp(a * gpmf)
    beta = 1/kb/T
    betaOrig = 1/kb/Torig
    Qorig = np.sum(np.exp(-betaOrig*gpmf))
    Qnew = np.sum(np.exp(-beta*gpmf))
    correction2 = Qorig / Qnew

    return correction1 * correction2 * origProb


def loadMSM(MSMFile):
    with open(MSMFile, "rb") as MSMfile:
        MSM_object = cPickle.load(MSMfile)
    return MSM_object


def reestimate_transition_matrix(count_matrix):
    count_matrix += 1/float(count_matrix.shape[0])
    return utils.buildRevTransitionMatrix(count_matrix)


def ensure_connectivity(MSMObject, allClusters):
    if len(allClusters) == MSMObject.stationary_distribution.size:
        pi = MSMObject.stationary_distribution
        clusters = allClusters[MSMObject.connected_sets[0]]
    else:
        ######################
        # Reconstruct stationary distribution with pseudocounts to ensure
        # connectivity
        print("Adding pseudocounts to enforce connectivity")
        counts = MSMObject.count_matrix_full
        trans = reestimate_transition_matrix(counts)
        _, eic = getSortedEigen(trans)
        pi = getStationaryDistr(eic[:, 0])
{'fraction_non_alphanumeric': 0.11939123589609027, 'fraction_numerical': 0.02309105221726581, 'mean_word_length': 3.059105431309904, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '120742', 'n_tokens_mistral': 2754, 'n_tokens_neox': 2601, 'n_words': 671}