Source code for schrodinger.application.desmond.ana

"""
Desmond analyses

Copyright Schrodinger, LLC. All rights reserved.
"""
from bisect import bisect
from bisect import bisect_left

from schrodinger.application.desmond.packages import analysis
from schrodinger.application.desmond.packages import topo
from schrodinger.application.desmond.packages import traj

QUANTITY_CLASS_MAP = {
    "angle": analysis.Angle,
    "dihedral": analysis.Torsion,
    "distance": analysis.Distance,
}


def calc_time_series(requests, model_fname, traj_fname):
    """
    :type requests: `list` of requests where each request is a `list` of
        quantity name, arguments, and (optional) time selection.
        Examples::

          [['dihedral', 1, 2, 3, 4],
           ['dihedral', 1, 2, 3, 4, {'begin':0, 'end':24.001}],]

    :return: `list` of analysis results
    """
    if not requests:
        return []

    models = list(topo.read_cms(model_fname))
    tr = traj.read_traj(traj_fname)
    sim_times = [fr.time for fr in tr]

    analyzers, times = [], []
    for req in requests:
        has_opt = isinstance(req[-1], dict)
        args = models + (req[1:-1] if has_opt else req[1:])
        times.append(req[-1] if has_opt else None)
        analyzers.append(QUANTITY_CLASS_MAP[req[0]](*args))
    results = analysis.analyze(tr, *analyzers)

    # FIXME: this is to undo the special treatment of analysis.analyze() when
    # there is only one analyzer
    if len(analyzers) == 1:
        results = [results]

    answer = []
    for res, time_sel in zip(results, times):
        if time_sel:
            left = bisect_left(sim_times, time_sel['begin'])
            right = bisect(sim_times, time_sel['end'])
            answer.append(list(zip(sim_times[left:right], res[left:right])))
        else:
            answer.append(list(zip(sim_times, res)))
    return answer


def calc_prob_profile(data, bin_width, min, max, is_periodic=False):
    """
    FIXME: To be added.
    """
    num_bin = int(float(max - min) / bin_width)
    bin = [0] * num_bin
    range = max - min
    for e in data:
        while (e > max):
            e -= range
        while (e < min):
            e += range
        i_bin = int((e - min) / bin_width + 0.5)
        if (i_bin >= num_bin):
            i_bin = i_bin - num_bin if (is_periodic) else num_bin - 1
        elif (i_bin < 0):
            i_bin = i_bin + num_bin if (is_periodic) else 0
        bin[i_bin] += 1
    result = []
    num_data = len(data)
    for i, e in enumerate(bin):
        result.append((i * bin_width + min, float(e) / num_data))
    return result