<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
"""
import numpy as np
import glob
import sys
import argparse
import cPickle
from pyemma.coordinates.clustering import AssignCenters
from AdaptivePELE.testing import runMarkovChainModel as run
import itertools

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"

    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) as MSMfile:
        MSM_object = cPickle.load(MSMfile)
    return MSM_object

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
        counts += 1/float(counts.shape[0])
        trans = run.buildRevTransitionMatrix(counts)
        _, eic = run.getSortedEigen(trans)
        pi = run.getStationaryDistr(eic[:, 0])
        clusters = allClusters
    return pi, clusters

def gather_coordinates(originalFilenames):
    originalCoordinates = []
    for originalFilename in originalFilenames:
        trajOriginalCoordinates = list(np.loadtxt(originalFilename, ndmin=2)[:, 1:])
        if np.random.random() < 0.0:
            # Add artificial points nearby to improve volume estimation, set
            # randomly since its very slow
            sys.stderr.write("Introducing artificial neighbours\n")
            newCoords = map(expandTrajs, trajOriginalCoordinates)
            trajOriginalCoordinates.extend(list(itertools.chain.from_iterable(newCoords)))
        originalCoordinates.append(np.array(trajOriginalCoordinates))
    return originalCoordinates

def create_box(clusters, originalCoordinates, d):
    """
        Create the box discretization for the estimation of volumes
    """
    dimensions = clusters.shape[1]
    maxval = dimensions*[-np.inf]
    minval = dimensions*[np.inf]
    for coord in originalCoordinates:
        cmaxval = coord.max(axis=0)
        cminval = coord.min(axis=0)
        maxval = np.maximum(cmaxval, maxval)
        minval = np.minimum(cminval, minval)

    print "Maximum bounds", maxval, "Minimum bounds", minval

    # Rounded floor and ceiling in intervals of "d" (e.g., floor of 1.73 with d = 0.5, will be 1.5 instead of 1.0, in order to optimize box creation.
    # An extra box is included in the ceiling, so that all the points are contained in the range given by arange
    bins = np.array([np.arange(np.floor(minval[i]) + d*int((minval[i] - np.floor(minval[i]))/d),
                               np.ceil(maxval[i]) + d*(int((maxval[i] - np.ceil(maxval[i]))/d) + 1),
                               d) for i in range(3)])
    return bins

def calculate_microstate_volumes(clusters, originalCoordinates, bins, d):
    """
        Estimate the clusters volumes using a cubic discretization of volumes
    """
    numberOfClusters = clusters.shape[0]
    print "Number of clusters", numberOfClusters
    histogram = np.array([])
    histograms = []
    microstateVolume = np.zeros(numberOfClusters)

    # dtrajs = clusteringObject.assign(originalCoordinates)
    dtrajs = assignNewTrajectories(originalCoordinates, clusters)
    for i in range(numberOfClusters):
        allCoords = []
        for trajOriginalCoordinates, dtraj in zip(originalCoordinates, dtrajs):
            assert dtraj.shape[0] == trajOriginalCoordinates.shape[0]
            belongingFrames = np.argwhere(dtraj == i)
            trajCoords = trajOriginalCoordinates[belongingFrames, :3]
            trajCoords = trajCoords.flatten().tolist()

            allCoords.extend(trajCoords)

        allCoords = np.reshape(allCoords, (-1, 3))

        current_hist, _ = np.histogramdd(allCoords, bins=bins)
        histograms.append(current_hist)

        if histogram.size == 0:
            histogram = np.copy(current_hist)
        else:
            histogram += current_hist

    nRows, nCols, nDepth = histogram.shape
    pseudo = False
    for i in range(numberOfClusters):
        histogramCluster = histograms[i]
        if pseudo:
            # Add "pseudocounts" to try to fill the holes that lead to volume
            # underestimation compared to Matlab script for free energies
            histogramTotal = histogram.copy()
            for x, y, z in zip(*np.where(histogramCluster)):
                upBound = max(x-1, 0)
                lowBound = min(x+2, nRows)
                leftBound = max(0, y-1)
                rightBound = min(y+2, nCols)
                topBound = max(z-1, 0)
                botBound = min(z+2, nDepth)
                signsCluster = np.sign(histogramCluster[upBound:lowBound, leftBound:rightBound, topBound:botBound])
                # signs = np.sign(histogramTotal[upBound:lowBound, leftBound:rightBound, topBound:botBound])
                histogramCluster[upBound:lowBound, leftBound:rightBound, topBound:botBound] += (1-signsCluster)*d/8  # + signsCluster*d/2
                histogramTotal[upBound:lowBound, leftBound:rightBound, topBound:botBound] += (1-signsCluster)  # + signs * d/2
            histogramTotal = histogramTotal[histogramCluster > 0]
        else:
            histogramTotal = histogram[histogramCluster > 0]
        histogramCluster = histogramCluster[histogramCluster > 0]
        microstateVolume[i] = (histogramCluster/histogramTotal).sum() * d**3
    return microstateVolume

def calculate_microstate_volumes_new(clusters, originalCoordinates, bins, d):
    """
        Estimate the clusters volumes using a cubic discretization of volumes
    """
    numberOfClusters = clusters.shape[0]
    print "Number of clusters", numberOfClusters

    allCoords = []
    # The coordinates array is built through lists in order to be able to
    # process trajectories of different lenght
    for coord in originalCoordinates:
        allCoords.extend(coord.tolist())
    allCoords = np.array(allCoords)

    histogram, _ = np.histogramdd(allCoords, bins=bins)

    centers_trajs = []
    for i, matrix in enumerate(histogram):
        centers_x = []
        x_val = bins[0][i]
        for j, row in enumerate(matrix):
            y_val = bins[1][j]
            indices, = np.nonzero(row)
            if indices.size > 0:
                init = indices[0]
                end = indices[-1]
                centers_x.extend([[x_val, y_val, bins[2][zi]] for zi in xrange(init, end+1)])
        if len(centers_x):
            centers_trajs.append(np.array(centers_x))

    dtrajs = assignNewTrajectories(centers_trajs, clusters)
    microstateVolume = np.zeros(numberOfClusters)
    for traj in dtrajs:
        counts, _ = np.histogram(traj, bins=np.arange(-0.5, numberOfClusters))
        microstateVolume += counts
    microstateVolume *= d**3
    # provisional return for debugging and testing script
    # return microstateVolume, centers_trajs, dtrajs
    return microstateVolume

def calculate_pmf(microstateVolume, pi):
    """
        Compute a potential of mean force given a stationary distribution
        (probabilities) and cluster volumes
    """
    kb = 0.0019872041
    T = 300
    beta = 1 / (kb * T)
    newDist = pi/microstateVolume
    newDist /= newDist[newDist != np.inf].sum()
    gpmf = -kb*T*np.log(newDist)
    print gpmf[gpmf == -np.inf]
    print gpmf[gpmf == np.inf]
    gpmf[gpmf == -np.inf] = np.inf  # to avoid contribution later
    gpmf -= gpmf.min()

    deltaW = -gpmf[gpmf != np.inf].max()
    print "bound    Delta G     Delta W     Binding Volume:     Binding Volume contribution"

    upperGpmfValues = np.arange(0, -deltaW, 0.5)

    # Initialize string variable in case loop is not accessed
    string = ""

    for upperGpmfValue in upperGpmfValues:
        bindingVolume = 0
        for g, volume in zip(gpmf, microstateVolume):
            if g <= upperGpmfValue:
                bindingVolume += np.exp(-beta * g) * volume
        deltaG = deltaW - kb*T*np.log(bindingVolume/1661)
        string = "%.1f\t%.3f\t%.3f\t%.3f\t%.3f" % (upperGpmfValue, deltaG, deltaW, bindingVolume, -kb*T*np.log(bindingVolume/1661))
        print string
    return gpmf, string

def main(trajWildcard, reweightingT=1000):
    allClusters = np.loadtxt("discretized/clusterCenters.dat")
    MSMObject = loadMSM('MSM_object.pkl')

    pi, clusters = ensure_connectivity(MSMObject, allClusters)
    # radius of the cube for volume determination
    d = 0.75

    originalFilenames = glob.glob(trajWildcard)
    # originalFilenames = glob.glob("rawData/"+trajWildcard)
    originalCoordinates = gather_coordinates(originalFilenames)

    bins = create_box(clusters, originalCoordinates, d)
    method = "new"
    if method == "new":
        print "Using new volume estimation with radius %.2f" % d
        microstateVolume = calculate_microstate_volumes_new(clusters, originalCoordinates, bins, d)
    else:
        print "Using old volume estimation with radius %.2f" % d
        microstateVolume = calculate_microstate_volumes(clusters, originalCoordinates, bins, d)
    np.savetxt("volumeOfClusters.dat", microstateVolume)

    gpmf, string = calculate_pmf(microstateVolume, pi)

    pmf_xyzg = np.hstack((clusters, np.expand_dims(gpmf, axis=1)))
    np.savetxt("pmf_xyzg.dat", pmf_xyzg)

    writePDB(pmf_xyzg)
    return string

if __name__ == "__main__":
    trajWildcard_name, reweight = parseArgs()
    main(trajWildcard_name, reweight)
<commit_msg>Minor fix in writePDB function 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
"""
import numpy as np
import glob
import sys
import argparse
import cPickle
from pyemma.coordinates.clustering import AssignCenters
from AdaptivePELE.testing import runMarkovChainModel as run
import itertools

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) as MSMfile:
        MSM_object = cPickle.load(MSMfile)
    return MSM_object

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
        counts += 1/float(counts.shape[0])
        trans = run.buildRevTransitionMatrix(counts)
        _, eic = run.getSortedEigen(trans)
        pi = run.getStationaryDistr(eic[:, 0])
        clusters = allClusters
    return pi, clusters

def gather_coordinates(originalFilenames):
    originalCoordinates = []
    for originalFilename in originalFilenames:
        trajOriginalCoordinates = list(np.loadtxt(originalFilename, ndmin=2)[:, 1:])
        if np.random.random() < 0.0:
            # Add artificial points nearby to improve volume estimation, set
            # randomly since its very slow
            sys.stderr.write("Introducing artificial neighbours\n")
            newCoords = map(expandTrajs, trajOriginalCoordinates)
            trajOriginalCoordinates.extend(list(itertools.chain.from_iterable(newCoords)))
        originalCoordinates.append(np.array(trajOriginalCoordinates))
    return originalCoordinates

def create_box(clusters, originalCoordinates, d):
    """
        Create the box discretization for the estimation of volumes
    """
    dimensions = clusters.shape[1]
    maxval = dimensions*[-np.inf]
    minval = dimensions*[np.inf]
    for coord in originalCoordinates:
        cmaxval = coord.max(axis=0)
        cminval = coord.min(axis=0)
        maxval = np.maximum(cmaxval, maxval)
        minval = np.minimum(cminval, minval)

    print "Maximum bounds", maxval, "Minimum bounds", minval

    # Rounded floor and ceiling in intervals of "d" (e.g., floor of 1.73 with d = 0.5, will be 1.5 instead of 1.0, in order to optimize box creation.
    # An extra box is included in the ceiling, so that all the points are contained in the range given by arange
    bins = np.array([np.arange(np.floor(minval[i]) + d*int((minval[i] - np.floor(minval[i]))/d),
                               np.ceil(maxval[i]) + d*(int((maxval[i] - np.ceil(maxval[i]))/d) + 1),
                               d) for i in range(3)])
    return bins

def calculate_microstate_volumes(clusters, originalCoordinates, bins, d):
    """
        Estimate the clusters volumes using a cubic discretization of volumes
    """
    numberOfClusters = clusters.shape[0]
    print "Number of clusters", numberOfClusters
    histogram = np.array([])
    histograms = []
    microstateVolume = np.zeros(numberOfClusters)

    # dtrajs = clusteringObject.assign(originalCoordinates)
    dtrajs = assignNewTrajectories(originalCoordinates, clusters)
    for i in range(numberOfClusters):
        allCoords = []
        for trajOriginalCoordinates, dtraj in zip(originalCoordinates, dtrajs):
            assert dtraj.shape[0] == trajOriginalCoordinates.shape[0]
            belongingFrames = np.argwhere(dtraj == i)
            trajCoords = trajOriginalCoordinates[belongingFrames, :3]
            trajCoords = trajCoords.flatten().tolist()

            allCoords.extend(trajCoords)

        allCoords = np.reshape(allCoords, (-1, 3))

        current_hist, _ = np.histogramdd(allCoords, bins=bins)
        histograms.append(current_hist)

        if histogram.size == 0:
            histogram = np.copy(current_hist)
        else:
            histogram += current_hist

    nRows, nCols, nDepth = histogram.shape
    pseudo = False
    for i in range(numberOfClusters):
        histogramCluster = histograms[i]
        if pseudo:
            # Add "pseudocounts" to try to fill the holes that lead to volume
            # underestimation compared to Matlab script for free energies
            histogramTotal = histogram.copy()
            for x, y, z in zip(*np.where(histogramCluster)):
                upBound = max(x-1, 0)
                lowBound = min(x+2, nRows)
                leftBound = max(0, y-1)
                rightBound = min(y+2, nCols)
                topBound = max(z-1, 0)
                botBound = min(z+2, nDepth)
                signsCluster = np.sign(histogramCluster[upBound:lowBound, leftBound:rightBound, topBound:botBound])
                # signs = np.sign(histogramTotal[upBound:lowBound, leftBound:rightBound, topBound:botBound])
                histogramCluster[upBound:lowBound, leftBound:rightBound, topBound:botBound] += (1-signsCluster)*d/8  # + signsCluster*d/2
                histogramTotal[upBound:lowBound, leftBound:rightBound, topBound:botBound] += (1-signsCluster)  # + signs * d/2
            histogramTotal = histogramTotal[histogramCluster > 0]
        else:
            histogramTotal = histogram[histogramCluster > 0]
        histogramCluster = histogramCluster[histogramCluster > 0]
        microstateVolume[i] = (histogramCluster/histogramTotal).sum() * d**3
    return microstateVolume

def calculate_microstate_volumes_new(clusters, originalCoordinates, bins, d):
    """
        Estimate the clusters volumes using a cubic discretization of volumes
    """
    numberOfClusters = clusters.shape[0]
    print "Number of clusters", numberOfClusters

    allCoords = []
    # The coordinates array is built through lists in order to be able to
    # process trajectories of different lenght
    for coord in originalCoordinates:
        allCoords.extend(coord.tolist())
    allCoords = np.array(allCoords)

    histogram, _ = np.histogramdd(allCoords, bins=bins)

    centers_trajs = []
    for i, matrix in enumerate(histogram):
        centers_x = []
        x_val = bins[0][i]
        for j, row in enumerate(matrix):
            y_val = bins[1][j]
            indices, = np.nonzero(row)
            if indices.size > 0:
                init = indices[0]
                end = indices[-1]
                centers_x.extend([[x_val, y_val, bins[2][zi]] for zi in xrange(init, end+1)])
        if len(centers_x):
            centers_trajs.append(np.array(centers_x))

    dtrajs = assignNewTrajectories(centers_trajs, clusters)
    microstateVolume = np.zeros(numberOfClusters)
    for traj in dtrajs:
        counts, _ = np.histogram(traj, bins=np.arange(-0.5, numberOfClusters))
        microstateVolume += counts
    microstateVolume *= d**3
    # provisional return for debugging and testing script
    # return microstateVolume, centers_trajs, dtrajs
    return microstateVolume

def calculate_pmf(microstateVolume, pi):
    """
        Compute a potential of mean force given a stationary distribution
        (probabilities) and cluster volumes
    """
    kb = 0.0019872041
    T = 300
    beta = 1 / (kb * T)
    newDist = pi/microstateVolume
    newDist /= newDist[newDist != np.inf].sum()
    gpmf = -kb*T*np.log(newDist)
    print gpmf[gpmf == -np.inf]
    print gpmf[gpmf == np.inf]
    gpmf[gpmf == -np.inf] = np.inf  # to avoid contribution later
    gpmf -= gpmf.min()

    deltaW = -gpmf[gpmf != np.inf].max()
    print "bound    Delta G     Delta W     Binding Volume:     Binding Volume contribution"

    upperGpmfValues = np.arange(0, -deltaW, 0.5)

    # Initialize string variable in case loop is not accessed
    string = ""

    for upperGpmfValue in upperGpmfValues:
        bindingVolume = 0
        for g, volume in zip(gpmf, microstateVolume):
            if g <= upperGpmfValue:
                bindingVolume += np.exp(-beta * g) * volume
        deltaG = deltaW - kb*T*np.log(bindingVolume/1661)
        string = "%.1f\t%.3f\t%.3f\t%.3f\t%.3f" % (upperGpmfValue, deltaG, deltaW, bindingVolume, -kb*T*np.log(bindingVolume/1661))
        print string
    return gpmf, string

def main(trajWildcard, reweightingT=1000):
    allClusters = np.loadtxt("discretized/clusterCenters.dat")
    MSMObject = loadMSM('MSM_object.pkl')

    pi, clusters = ensure_connectivity(MSMObject, allClusters)
    # radius of the cube for volume determination
    d = 0.75

    originalFilenames = glob.glob(trajWildcard)
    # originalFilenames = glob.glob("rawData/"+trajWildcard)
    originalCoordinates = gather_coordinates(originalFilenames)

    bins = create_box(clusters, originalCoordinates, d)
    method = "new"
    if method == "new":
        print "Using new volume estimation with radius %.2f" % d
        microstateVolume = calculate_microstate_volumes_new(clusters, originalCoordinates, bins, d)
    else:
        print "Using old volume estimation with radius %.2f" % d
        microstateVolume = calculate_microstate_volumes(clusters, originalCoordinates, bins, d)
    np.savetxt("volumeOfClusters.dat", microstateVolume)

    gpmf, string = calculate_pmf(microstateVolume, pi)

    pmf_xyzg = np.hstack((clusters, np.expand_dims(gpmf, axis=1)))
    np.savetxt("pmf_xyzg.dat", pmf_xyzg)

    writePDB(pmf_xyzg)
    return string

if __name__ == "__main__":
    trajWildcard_name, reweight = parseArgs()
    main(trajWildcard_name, reweight)
<|endoftext|>


<commit_before>from oemof.core.network import Entity
from oemof.core.network.entities.components import Transformer
from oemof.core.network.entities.components import Transport
from oemof.core.network.entities.components import Source
from oemof.core.network.entities.buses import Bus

#from oemof.core.network.entities.buses import BusPypo
#from oemof.core.network.entities.components.transports import BranchPypo
#from oemof.core.network.entities.components.sources import GenPypo

import networkx as nx
import matplotlib.pyplot as plt

class GridDingo:
    """ DINGO grid

    Parameters
    ----------
    id_db : id according to database table
    """

    def __init__(self, **kwargs):
        self.id_db = kwargs.get('id_db', None)
        self.region = kwargs.get('region', None)
        #self.geo_data = kwargs.get('geo_data', None)

        self._graph = nx.Graph()

    def graph_add_node(self, node_object):
        """Adds a station or cable distributor object to grid graph if not already existing"""
        if node_object not in self._graph.nodes()\
                and (isinstance(node_object, StationDingo) or isinstance(node_object, CableDistributorDingo)):
            self._graph.add_node(node_object)

    # TODO: UPDATE DRAW FUNCTION -> make draw method work for both MV and LV regions!
    def graph_draw(self):
        """ Draws grid graph using networkx

        caution: The geo coords (for used crs see database import in class `NetworkDingo`) are used as positions for
                 drawing but networkx uses cartesian crs. Since no coordinate transformation is performed, the drawn
                 graph representation is falsified!
        """

        g = self._graph

        # get draw params from nodes and edges (coordinates, colors, demands, etc.)
        nodes_pos = {}; demands = {}; demands_pos = {}
        nodes_color = []
        for node in g.nodes():
            if isinstance(node, StationDingo) or isinstance(node, CableDistributorDingo):
                nodes_pos[node] = (node.geo_data.x, node.geo_data.y)
                # TODO: MOVE draw/color settings to config
            if node == self.station():
                nodes_color.append((1, 0.5, 0.5))
            else:
                #demands[node] = 'd=' + '{:.3f}'.format(node.grid.region.peak_load_sum)
                #demands_pos[node] = tuple([a+b for a, b in zip(nodes_pos[node], [0.003]*len(nodes_pos[node]))])
                nodes_color.append((0.5, 0.5, 1))

        plt.figure()
        nx.draw_networkx(g, nodes_pos, node_color=nodes_color, font_size=10)
        nx.draw_networkx_labels(g, demands_pos, labels=demands, font_size=8)
        plt.show()

    def graph_edges(self):
        """ Returns a generator for iterating over graph edges

        The edge of a graph is described by the to adjacent node and the branch
        object itself. Whereas the branch object is used to hold all relevant
        power system parameters.

        Note
        ----

        There are generator functions for nodes (`Graph.nodes()`) and edges
        (`Graph.edges()`) in NetworkX but unlike graph nodes, which can be
        represented by objects, branch objects can only be accessed by using an
        edge attribute ('branch' is used here)

        To make access to attributes of the branch objects simplier and more
        intuitive for the user, this generator yields a dictionary for each edge
        that contains information about adjacent nodes and the branch object.

        Note, the construction of the dictionary highly depends on the structure
        of the in-going tuple (which is defined by the needs of networkX). If
        this changes, the code will break.
        """
        for edge in nx.get_edge_attributes(self._graph, 'branch').items():
            yield {'adj_nodes': edge[0], 'branch': edge[1]}

class StationDingo():
    """
    Defines a MV/LVstation in DINGO
    -------------------------------

    id_db: id according to database table
    """
    # TODO: add method remove_transformer()

    def __init__(self, **kwargs):
        self.id_db = kwargs.get('id_db', None)
        self.geo_data = kwargs.get('geo_data', None)
        self.grid = kwargs.get('grid', None)
        self._transformers = []
        self.busbar = None
        self.peak_load = kwargs.get('peak_load', None)

    def transformers(self):
        """Returns a generator for iterating over transformers"""
        for trans in self._transformers:
            yield trans

    def add_transformer(self, transformer):
        """Adds a transformer to _transformers if not already existing"""
        # TODO: check arg
        if transformer not in self.transformers() and isinstance(transformer, TransformerDingo):
            self._transformers.append(transformer)
        # TODO: what if it exists? -> error message

class BusDingo(Bus):
    """ Create new pypower Bus class as child from oemof Bus used to define
    busses and generators data
    """

    def __init__(self, **kwargs):
        """Assigned minimal required pypower input parameters of the bus and
        generator as arguments

        Keyword description of bus arguments:
        bus_id -- the bus number (also used as GEN_BUS parameter for generator)
        bus_type -- the bus type (1 = PQ, 2 = PV, 3 = ref, 4 = Isolated)
        PD -- the real power demand in MW
        QD -- the reactive power demand in MVAr
        GS -- the shunt conductance (demanded at V = 1.0 p.u.) in MW
        BS -- the shunt susceptance (injected at V = 1.0 p.u.) in MVAr
        bus_area -- area number (positive integer)
        VM -- the voltage magnitude in p.u.
        VA -- the voltage angle in degrees
        base_kv -- the base voltage in kV
        zone -- loss zone (positive integer)
        vmax -- the maximum allowed voltage magnitude in p.u.
        vmin -- the minimum allowed voltage magnitude in p.u.
        """

        super().__init__(**kwargs)
        # Bus Data parameters

class BranchDingo():
    """
    Cables and lines
    ----------------
    geo_data : shapely.geometry object
        Geo-spatial data with informations for location/region-shape. The
        geometry can be a polygon/multi-polygon for regions, a line for
        transport objects or a point for objects such as transformer sources.
    equip_line_id : int
        ID of cable/line type according to DB table 'equip_line'
    out_max : float
        Maximum output which can possibly be obtained when using the transport,
        in $MW$.
    """

    def __init__(self, **kwargs):
        # inherit parameters from oemof's Transport
        #super().__init__(**kwargs)

        #

        # more params (OLD)
        self.equip_line_id = kwargs.get('equip_line_id', None)
        self.v_level = kwargs.get('v_level', None)
        self.type = kwargs.get('type', None)
        self.cable_cnt = kwargs.get('cable_cnt', None)
        self.wire_cnt = kwargs.get('wire_cnt', None)
        self.cs_area = kwargs.get('cs_area', None)
        self.r = kwargs.get('r', None)
        self.x = kwargs.get('x', None)
        self.c = kwargs.get('c', None)
        self.i_max_th = kwargs.get('i_max_th', None)
        self.s_max_a = kwargs.get('s_max_a', None)
        self.s_max_b = kwargs.get('s_max_b', None)
        self.s_max_c = kwargs.get('s_max_c', None)

class TransformerDingo():
    """
    Transformers
    ------------
    geo_data : shapely.geometry object
        Geo-spatial data with informations for location/region-shape. The
        geometry can be a polygon/multi-polygon for regions, a line for
        transport objects or a point for objects such as transformer sources.
    equip_trans_id : int
        ID of transformer type according to DB table 'equip_trans'
    v_level :
        voltage level
    s_max_a : float
        rated power (long term)
    s_max_b : float
        rated power (short term)
    s_max_c : float
        rated power (emergency)
    phase_angle : float
        phase shift angle
    tap_ratio: float
        off nominal turns ratio
    """

    def __init__(self, **kwargs):
        #inherit parameters from oemof's Transformer
        # super().__init__(**kwargs)
        #more params
        self.equip_trans_id = kwargs.get('equip_trans_id', None)
        self.v_level = kwargs.get('v_level', None)
        self.s_max_a = kwargs.get('s_max_longterm', None)
        self.s_max_b = kwargs.get('s_max_shortterm', None)
        self.s_max_c = kwargs.get('s_max_emergency', None)
        self.phase_angle = kwargs.get('phase_angle', None)
        self.tap_ratio = kwargs.get('tap_ratio', None)

class SourceDingo(Source):
    """
    Generators
    """

    def __init__(self, **kwargs):
        #inherit parameters from oemof's Transformer
        super().__init__(**kwargs)

class CableDistributorDingo():
    """ Cable distributor (connection point) """

    def __init__(self, **kwargs):
        self.id_db = kwargs.get('id_db', None)
        self.geo_data = kwargs.get('geo_data', None)
        self.grid = kwargs.get('grid', None)

    def __repr__(self):
        return 'cable_dist_' + str(self.id_db)
<commit_msg>add method to grid to get path length between 2 nodes<commit_after>from oemof.core.network import Entity
from oemof.core.network.entities.components import Transformer
from oemof.core.network.entities.components import Transport
from oemof.core.network.entities.components import Source
from oemof.core.network.entities.buses import Bus

#from oemof.core.network.entities.buses import BusPypo
#from oemof.core.network.entities.components.transports import BranchPypo
#from oemof.core.network.entities.components.sources import GenPypo

import networkx as nx
import matplotlib.pyplot as plt

class GridDingo:
    """ DINGO grid

    Parameters
    ----------
    id_db : id according to database table
    """

    def __init__(self, **kwargs):
        self.id_db = kwargs.get('id_db', None)
        self.region = kwargs.get('region', None)
        #self.geo_data = kwargs.get('geo_data', None)

        self._graph = nx.Graph()

    def graph_add_node(self, node_object):
        """Adds a station or cable distributor object to grid graph if not already existing"""
        if node_object not in self._graph.nodes()\
                and (isinstance(node_object, StationDingo) or isinstance(node_object, CableDistributorDingo)):
            self._graph.add_node(node_object)

    # TODO: UPDATE DRAW FUNCTION -> make draw method work for both MV and LV regions!
    def graph_draw(self):
        """ Draws grid graph using networkx

        caution: The geo coords (for used crs see database import in class `NetworkDingo`) are used as positions for
                 drawing but networkx uses cartesian crs. Since no coordinate transformation is performed, the drawn
                 graph representation is falsified!
        """

        g = self._graph

        # get draw params from nodes and edges (coordinates, colors, demands, etc.)
        nodes_pos = {}; demands = {}; demands_pos = {}
        nodes_color = []
        for node in g.nodes():
            if isinstance(node, StationDingo) or isinstance(node, CableDistributorDingo):
                nodes_pos[node] = (node.geo_data.x, node.geo_data.y)
                # TODO: MOVE draw/color settings to config
            if node == self.station():
                nodes_color.append((1, 0.5, 0.5))
            else:
                #demands[node] = 'd=' + '{:.3f}'.format(node.grid.region.peak_load_sum)
                #demands_pos[node] = tuple([a+b for a, b in zip(nodes_pos[node], [0.003]*len(nodes_pos[node]))])
                nodes_color.append((0.5, 0.5, 1))

        plt.figure()
        nx.draw_networkx(g, nodes_pos, node_color=nodes_color, font_size=10)
        nx.draw_networkx_labels(g, demands_pos, labels=demands, font_size=8)
        plt.show()

    def graph_edges(self):
        """ Returns a generator for iterating over graph edges

        The edge of a graph is described by the to adjacent node and the branch
        object itself. Whereas the branch object is used to hold all relevant
        power system parameters.

        Note
        ----

        There are generator functions for nodes (`Graph.nodes()`) and edges
        (`Graph.edges()`) in NetworkX but unlike graph nodes, which can be
        represented by objects, branch objects can only be accessed by using an
        edge attribute ('branch' is used here)

        To make access to attributes of the branch objects simplier and more
        intuitive for the user, this generator yields a dictionary for each edge
        that contains information about adjacent nodes and the branch object.

        Note, the construction of the dictionary highly depends on the structure
        of the in-going tuple (which is defined by the needs of networkX). If
        this changes, the code will break.
        """
        for edge in nx.get_edge_attributes(self._graph, 'branch').items():
            yield {'adj_nodes': edge[0], 'branch': edge[1]}

    def graph_path_length(self, node_source, node_target):
        """ Calculates the absolute distance between `node_source` and `node_target` in meters.
        Args:
            node_source:
            node_target:

        Returns:

        """

        length = 0
        path = nx.shortest_path(self._graph, node_source, node_target)
        node_pairs = list(zip(path[0:len(path)-1], path[1:len(path)]))

        for n1, n2 in node_pairs:
            length += self._graph.edge[n1][n2]['branch'].length

        return length

class StationDingo():
    """
    Defines a MV/LVstation in DINGO
    -------------------------------

    id_db: id according to database table
    """
    # TODO: add method remove_transformer()

    def __init__(self, **kwargs):
        self.id_db = kwargs.get('id_db', None)
        self.geo_data = kwargs.get('geo_data', None)
        self.grid = kwargs.get('grid', None)
        self._transformers = []
        self.busbar = None
        self.peak_load = kwargs.get('peak_load', None)

    def transformers(self):
        """Returns a generator for iterating over transformers"""
        for trans in self._transformers:
            yield trans

    def add_transformer(self, transformer):
        """Adds a transformer to _transformers if not already existing"""
        # TODO: check arg
        if transformer not in self.transformers() and isinstance(transformer, TransformerDingo):
            self._transformers.append(transformer)
        # TODO: what if it exists? -> error message

class BusDingo(Bus):
    """ Create new pypower Bus class as child from oemof Bus used to define
    busses and generators data
    """

    def __init__(self, **kwargs):
        """Assigned minimal required pypower input parameters of the bus and
        generator as arguments

        Keyword description of bus arguments:
        bus_id -- the bus number (also used as GEN_BUS parameter for generator)
        bus_type -- the bus type (1 = PQ, 2 = PV, 3 = ref, 4 = Isolated)
        PD -- the real power demand in MW
        QD -- the reactive power demand in MVAr
        GS -- the shunt conductance (demanded at V = 1.0 p.u.) in MW
        BS -- the shunt susceptance (injected at V = 1.0 p.u.) in MVAr
        bus_area -- area number (positive integer)
        VM -- the voltage magnitude in p.u.
        VA -- the voltage angle in degrees
        base_kv -- the base voltage in kV
        zone -- loss zone (positive integer)
        vmax -- the maximum allowed voltage magnitude in p.u.
        vmin -- the minimum allowed voltage magnitude in p.u.
        """

        super().__init__(**kwargs)
        # Bus Data parameters

class BranchDingo():
    """
    Cables and lines
    ----------------
    geo_data : shapely.geometry object
        Geo-spatial data with informations for location/region-shape. The
        geometry can be a polygon/multi-polygon for regions, a line for
        transport objects or a point for objects such as transformer sources.
    equip_line_id : int
        ID of cable/line type according to DB table 'equip_line'
    out_max : float
        Maximum output which can possibly be obtained when using the transport,
        in $MW$.
    """

    def __init__(self, **kwargs):
        # inherit parameters from oemof's Transport
        #super().__init__(**kwargs)

        #

        # more params (OLD)
        self.equip_line_id = kwargs.get('equip_line_id', None)
        self.v_level = kwargs.get('v_level', None)
        self.type = kwargs.get('type', None)
        self.cable_cnt = kwargs.get('cable_cnt', None)
        self.wire_cnt = kwargs.get('wire_cnt', None)
        self.cs_area = kwargs.get('cs_area', None)
        self.r = kwargs.get('r', None)
        self.x = kwargs.get('x', None)
        self.c = kwargs.get('c', None)
        self.i_max_th = kwargs.get('i_max_th', None)
        self.s_max_a = kwargs.get('s_max_a', None)
        self.s_max_b = kwargs.get('s_max_b', None)
        self.s_max_c = kwargs.get('s_max_c', None)

class TransformerDingo():
    """
    Transformers
    ------------
    geo_data : shapely.geometry object
        Geo-spatial data with informations for location/region-shape. The
        geometry can be a polygon/multi-polygon for regions, a line for
        transport objects or a point for objects such as transformer sources.
    equip_trans_id : int
        ID of transformer type according to DB table 'equip_trans'
    v_level :
        voltage level
    s_max_a : float
        rated power (long term)
    s_max_b : float
        rated power (short term)
    s_max_c : float
        rated power (emergency)
    phase_angle : float
        phase shift angle
    tap_ratio: float
        off nominal turns ratio
    """

    def __init__(self, **kwargs):
        #inherit parameters from oemof's Transformer
        # super().__init__(**kwargs)
        #more params
        self.equip_trans_id = kwargs.get('equip_trans_id', None)
        self.v_level = kwargs.get('v_level', None)
        self.s_max_a = kwargs.get('s_max_longterm', None)
        self.s_max_b = kwargs.get('s_max_shortterm', None)
        self.s_max_c = kwargs.get('s_max_emergency', None)
        self.phase_angle = kwargs.get('phase_angle', None)
        self.tap_ratio = kwargs.get('tap_ratio', None)

class SourceDingo(Source):
    """
    Generators
    """

    def __init__(self, **kwargs):
        #inherit parameters from oemof's Transformer
        super().__init__(**kwargs)

class CableDistributorDingo():
    """ Cable distributor (connection point) """

    def __init__(self, **kwargs):
        self.id_db = kwargs.get('id_db', None)
        self.geo_data = kwargs.get('geo_data', None)
        self.grid = kwargs.get('grid', None)

    def __repr__(self):
        return 'cable_dist_' + str(self.id_db)
<|endoftext|>


<commit_before>import numpy as np
from dipy.direction.peaks import peak_directions, default_sphere
from dipy.reconst.shm import order_from_ncoef, sph_harm_lookup
from dipy.tracking.local.direction_getter import DirectionGetter
from dipy.tracking.local.interpolation import trilinear_interpolate4d

def closest_peak(peak_dirs, prev_step, cos_similarity):
    """Return the closest direction to prev_step from peak_dirs.

    All directions should be unit vectors. Antipodal symmetry is assumed, ie
    direction x is the same as -x.

    Parameters
    ----------
    peak_dirs : array (N, 3)
        N unit vectors.
    prev_step : array (3,) or None
        Previous direction.
    cos_similarity : float
        `cos(max_angle)` where `max_angle` is the maximum allowed angle between
        prev_step and the returned direction.

    Returns
    -------
    direction : array or None
        If prev_step is None, returns peak_dirs. Otherwise returns the
        closest direction to prev_step. If no directions are close enough to
        prev_step, returns None
    """
    peak_dots = np.dot(peak_dirs, prev_step)
    closest_peak = abs(peak_dots).argmax()
    dot_closest_peak = peak_dots[closest_peak]
    if dot_closest_peak >= cos_similarity:
        return peak_dirs[closest_peak]
    elif dot_closest_peak <= -cos_similarity:
        return -peak_dirs[closest_peak]
    else:
        return None

class PmfGen(object):
    pass

class SimplePmfGen(PmfGen):

    def __init__(self, pmf_array):
        if pmf_array.min() < 0:
            raise ValueError("pmf should not have negative values")
        self.pmf_array = pmf_array

    def get_pmf(self, point):
        return trilinear_interpolate4d(self.pmf_array, point)

class SHCoeffPmfGen(PmfGen):

    def __init__(self, shcoeff, sphere, basis_type):
        self.shcoeff = shcoeff
        self.sphere = sphere
        sh_order = order_from_ncoef(shcoeff.shape[-1])
        try:
            basis = sph_harm_lookup[basis_type]
        except KeyError:
            raise ValueError("%s is not a known basis type." % basis_type)
        self._B, m, n = basis(sh_order, sphere.theta, sphere.phi)

    def get_pmf(self, point):
        coeff = trilinear_interpolate4d(self.shcoeff, point)
        pmf = np.dot(self._B, coeff)
        return pmf

class BaseDirectionGetter(DirectionGetter):
    """A base class for dynamic direction getters"""

    def __init__(self, pmf_gen, max_angle, sphere, pmf_threshold=.1, **kwargs):
        self.sphere = sphere
        self._pf_kwargs = kwargs
        self.pmf_gen = pmf_gen
        if pmf_threshold < 0:
            raise ValueError("pmf threshould must be >= 0.")
        self.pmf_threshold = pmf_threshold
        self.cos_similarity = np.cos(np.deg2rad(max_angle))

    def _get_peak_directions(self, blob):
        """Gets directions using parameters provided at init.

        Blob can be any function defined on ``self.sphere``, ie an ODF, PMF,
        FOD.
        """
        return peak_directions(blob, self.sphere, **self._pf_kwargs)[0]

    def initial_direction(self, point):
        """Returns best directions at seed location to start tracking.

        Parameters
        ----------
        point : ndarray, shape (3,)
            The point in an image at which to lookup tracking directions.

        Returns
        -------
        directions : ndarray, shape (N, 3)
            Possible tracking directions from point. ``N`` may be 0, all
            directions should be unique.

        """
        pmf = self.pmf_gen.get_pmf(point)
        pmf.clip(min=self.pmf_threshold, out=pmf)
        return self._get_peak_directions(pmf)

class PmfGenDirectionGetter(BaseDirectionGetter):
    """A direction getter that returns the closest odf peak to previous tracking
    direction.

    """
    @classmethod
    def from_pmf(klass, pmf, max_angle, sphere=default_sphere,
                 pmf_threshold=0.1, **kwargs):
        """Constructor for making a DirectionGetter from an array of Pmfs

        Parameters
        ----------
        pmf : array, 4d
            The pmf to be used for tracking at each voxel.
        max_angle : float, [0, 90]
            The maximum allowed angle between incoming direction and new
            direction.
        sphere : Sphere
            The set of directions to be used for tracking.
        pmf_threshold : float [0., 1.]
            Used to remove direction from the probability mass function for
            selecting the tracking direction.
        relative_peak_threshold : float in [0., 1.]
            Used for extracting initial tracking directions. Passed to
            peak_directions.
        min_separation_angle : float in [0, 90]
            Used for extracting initial tracking directions. Passed to
            peak_directions.

        See also
        --------
        dipy.direction.peaks.peak_directions

        """
        pmf = np.asarray(pmf, dtype=float)
        if pmf.ndim != 4:
            raise ValueError("pmf should be a 4d array.")
        if pmf.shape[3] != len(sphere.theta):
            msg = ("The last dimension of pmf should match the number of "
                   "points in sphere.")
            raise ValueError(msg)
        pmf_gen = SimplePmfGen(pmf)
        return klass(pmf_gen, max_angle, sphere, pmf_threshold, **kwargs)

    @classmethod
    def from_shcoeff(klass, shcoeff, max_angle, sphere=default_sphere,
                     pmf_threshold=0.1,
                     basis_type=None, **kwargs):
        """Probabilistic direction getter from a distribution of directions
        on the sphere.

        Parameters
        ----------
        shcoeff : array
            The distribution of tracking directions at each voxel represented
            as a function on the sphere using the real spherical harmonic
            basis. For example the FOD of the Constrained Spherical
            Deconvolution model can be used this way. This distribution will
            be discretized using ``sphere`` and tracking directions will be
            chosen from the vertices of ``sphere`` based on the distribution.
        max_angle : float, [0, 90]
            The maximum allowed angle between incoming direction and new
            direction.
        sphere : Sphere
            The set of directions to be used for tracking.
        pmf_threshold : float [0., 1.]
            Used to remove direction from the probability mass function for
            selecting the tracking direction.
        basis_type : name of basis
            The basis that ``shcoeff`` are associated with.
            ``dipy.reconst.shm.real_sym_sh_basis`` is used by default.
        relative_peak_threshold : float in [0., 1.]
            Used for extracting initial tracking directions. Passed to
            peak_directions.
        min_separation_angle : float in [0, 90]
            Used for extracting initial tracking directions. Passed to
            peak_directions.

        See also
        --------
        dipy.direction.peaks.peak_directions

        """
        pmf_gen = SHCoeffPmfGen(shcoeff, sphere, basis_type)
        return klass(pmf_gen, max_angle, sphere, pmf_threshold, **kwargs)

class ClosestDirectionGetter(PmfGenDirectionGetter):
    """A direction getter that returns the closest odf peak to previous tracking
    direction.
    """

    def get_direction(self, point, direction):
        pmf = self.pmf_gen.get_pmf(point)
        pmf.clip(min=self.pmf_threshold, out=pmf)
        peaks = self._get_peak_directions(pmf)
        if len(peaks) == 0:
            return 1
        new_dir = closest_peak(peaks, direction, self.cos_similarity)
        if new_dir is None:
            return 1
        else:
            direction[:] = new_dir
            return 0
<commit_msg>DOC - PmfGenDirectionGetter<commit_after>import numpy as np
from dipy.direction.peaks import peak_directions, default_sphere
from dipy.reconst.shm import order_from_ncoef, sph_harm_lookup
from dipy.tracking.local.direction_getter import DirectionGetter
from dipy.tracking.local.interpolation import trilinear_interpolate4d

def closest_peak(peak_dirs, prev_step, cos_similarity):
    """Return the closest direction to prev_step from peak_dirs.

    All directions should be unit vectors. Antipodal symmetry is assumed, ie
    direction x is the same as -x.

    Parameters
    ----------
    peak_dirs : array (N, 3)
        N unit vectors.
    prev_step : array (3,) or None
        Previous direction.
    cos_similarity : float
        `cos(max_angle)` where `max_angle` is the maximum allowed angle between
        prev_step and the returned direction.

    Returns
    -------
    direction : array or None
        If prev_step is None, returns peak_dirs. Otherwise returns the
        closest direction to prev_step. If no directions are close enough to
        prev_step, returns None
    """
    peak_dots = np.dot(peak_dirs, prev_step)
    closest_peak = abs(peak_dots).argmax()
    dot_closest_peak = peak_dots[closest_peak]
    if dot_closest_peak >= cos_similarity:
        return peak_dirs[closest_peak]
    elif dot_closest_peak <= -cos_similarity:
        return -peak_dirs[closest_peak]
    else:
        return None

class PmfGen(object):
    pass

class SimplePmfGen(PmfGen):

    def __init__(self, pmf_array):
        if pmf_array.min() < 0:
            raise ValueError("pmf should not have negative values")
        self.pmf_array = pmf_array

    def get_pmf(self, point):
        return trilinear_interpolate4d(self.pmf_array, point)

class SHCoeffPmfGen(PmfGen):

    def __init__(self, shcoeff, sphere, basis_type):
        self.shcoeff = shcoeff
        self.sphere = sphere
        sh_order = order_from_ncoef(shcoeff.shape[-1])
        try:
            basis = sph_harm_lookup[basis_type]
        except KeyError:
            raise ValueError("%s is not a known basis type." % basis_type)
        self._B, m, n = basis(sh_order, sphere.theta, sphere.phi)

    def get_pmf(self, point):
        coeff = trilinear_interpolate4d(self.shcoeff, point)
        pmf = np.dot(self._B, coeff)
        return pmf

class BaseDirectionGetter(DirectionGetter):
    """A base class for dynamic direction getters"""

    def __init__(self, pmf_gen, max_angle, sphere, pmf_threshold=.1, **kwargs):
        self.sphere = sphere
        self._pf_kwargs = kwargs
        self.pmf_gen = pmf_gen
        if pmf_threshold < 0:
            raise ValueError("pmf threshold must be >= 0.")
        self.pmf_threshold = pmf_threshold
        self.cos_similarity = np.cos(np.deg2rad(max_angle))

    def _get_peak_directions(self, blob):
        """Gets directions using parameters provided at init.

        Blob can be any function defined on ``self.sphere``, ie an ODF, PMF,
        FOD.
        """
        return peak_directions(blob, self.sphere, **self._pf_kwargs)[0]

    def initial_direction(self, point):
        """Returns best directions at seed location to start tracking.

        Parameters
        ----------
        point : ndarray, shape (3,)
            The point in an image at which to lookup tracking directions.

        Returns
        -------
        directions : ndarray, shape (N, 3)
            Possible tracking directions from point. ``N`` may be 0, all
            directions should be unique.

        """
        pmf = self.pmf_gen.get_pmf(point)
        pmf.clip(min=self.pmf_threshold, out=pmf)
        return self._get_peak_directions(pmf)

class PmfGenDirectionGetter(BaseDirectionGetter):
    """A base class for direction getter using a pmf"""
    @classmethod
    def from_pmf(klass, pmf, max_angle, sphere=default_sphere,
                 pmf_threshold=0.1, **kwargs):
        """Constructor for making a DirectionGetter from an array of Pmfs

        Parameters
        ----------
        pmf : array, 4d
            The pmf to be used for tracking at each voxel.
        max_angle : float, [0, 90]
            The maximum allowed angle between incoming direction and new
            direction.
        sphere : Sphere
            The set of directions to be used for tracking.
        pmf_threshold : float [0., 1.]
            Used to remove direction from the probability mass function for
            selecting the tracking direction.
        relative_peak_threshold : float in [0., 1.]
            Used for extracting initial tracking directions. Passed to
            peak_directions.
        min_separation_angle : float in [0, 90]
            Used for extracting initial tracking directions. Passed to
            peak_directions.

        See also
        --------
        dipy.direction.peaks.peak_directions

        """
        pmf = np.asarray(pmf, dtype=float)
        if pmf.ndim != 4:
            raise ValueError("pmf should be a 4d array.")
        if pmf.shape[3] != len(sphere.theta):
            msg = ("The last dimension of pmf should match the number of "
                   "points in sphere.")
            raise ValueError(msg)
        pmf_gen = SimplePmfGen(pmf)
        return klass(pmf_gen, max_angle, sphere, pmf_threshold, **kwargs)

    @classmethod
    def from_shcoeff(klass, shcoeff, max_angle, sphere=default_sphere,
                     pmf_threshold=0.1,
                     basis_type=None, **kwargs):
        """Probabilistic direction getter from a distribution of directions
        on the sphere.

        Parameters
        ----------
        shcoeff : array
            The distribution of tracking directions at each voxel represented
            as a function on the sphere using the real spherical harmonic
            basis. For example the FOD of the Constrained Spherical
            Deconvolution model can be used this way. This distribution will
            be discretized using ``sphere`` and tracking directions will be
            chosen from the vertices of ``sphere`` based on the distribution.
        max_angle : float, [0, 90]
            The maximum allowed angle between incoming direction and new
            direction.
        sphere : Sphere
            The set of directions to be used for tracking.
        pmf_threshold : float [0., 1.]
            Used to remove direction from the probability mass function for
            selecting the tracking direction.
        basis_type : name of basis
            The basis that ``shcoeff`` are associated with.
            ``dipy.reconst.shm.real_sym_sh_basis`` is used by default.
        relative_peak_threshold : float in [0., 1.]
            Used for extracting initial tracking directions. Passed to
            peak_directions.
        min_separation_angle : float in [0, 90]
            Used for extracting initial tracking directions. Passed to
            peak_directions.

        See also
        --------
        dipy.direction.peaks.peak_directions

        """
        pmf_gen = SHCoeffPmfGen(shcoeff, sphere, basis_type)
        return klass(pmf_gen, max_angle, sphere, pmf_threshold, **kwargs)

class ClosestDirectionGetter(PmfGenDirectionGetter):
    """A direction getter that returns the closest odf peak to previous tracking
    direction.
    """

    def get_direction(self, point, direction):
        pmf = self.pmf_gen.get_pmf(point)
        pmf.clip(min=self.pmf_threshold, out=pmf)
        peaks = self._get_peak_directions(pmf)
        if len(peaks) == 0:
            return 1
        new_dir = closest_peak(peaks, direction, self.cos_similarity)
        if new_dir is None:
            return 1
        else:
            direction[:] = new_dir
            return 0
<|endoftext|>


<commit_before>from DirectGuiGlobals import *
from DirectFrame import *
from DirectButton import *

def findDialog(uniqueName):
    """findPanel(string uniqueName)

    Returns the panel whose uniqueName is given.  This is mainly
    useful for debugging, to get a pointer to the current onscreen
    panel of a particular type.

    """
    if DirectDialog.AllDialogs.has_key(uniqueName):
        return DirectDialog.AllDialogs[uniqueName]
    return None

def cleanupDialog(uniqueName):
    """cleanupPanel(string uniqueName)

    Cleans up (removes) the panel with the given uniqueName.  This
    may be useful when some panels know about each other and know
    that opening panel A should automatically close panel B, for
    instance.
    """

    if DirectDialog.AllDialogs.has_key(uniqueName):
        # calling cleanup() will remove it out of the AllDialogs dict
        # This way it will get removed from the dict even it we did
        # not clean it up using this interface (ie somebody called
        # self.cleanup() directly
        DirectDialog.AllDialogs[uniqueName].cleanup()

class DirectDialog(DirectFrame):

    AllDialogs = {}
    PanelIndex = 0

    def __init__(self, parent = aspect2d, **kw):
        """
        DirectDialog(kw)

        Creates a popup dialog to alert and/or interact with user.
        Some of the main keywords that can be used to customize the dialog:
            Keyword              Definition
            -------              ----------
            text                 Text message/query displayed to user
            geom                 Geometry to be displayed in dialog
            buttonTextList       List of text to show on each button
            buttonGeomList       List of geometry to show on each button
            buttonImageList      List of images to show on each button
            buttonValueList      List of values sent to dialog command for
                                 each button.  If value is [] then the
                                 ordinal rank of the button is used as
                                 its value
            buttonHotKeyList     List of hotkeys to bind to each button.
                                 Typing hotkey is equivalent to pressing
                                 the corresponding button.
            supressKeys          Set to true if you wish to supress keys
                                 (i.e. Dialog eats key event), false if
                                 you wish Dialog to pass along key event
            buttonSize           4-tuple used to specify custom size for
                                 each button (to make bigger then geom/text
                                 for example)
            pad                  Space between border and interior graphics
            topPad               Extra space added above text/geom/image
            midPad               Extra space added between text/buttons
            buttonPadSF          Scale factor used to expand/contract
                                 button horizontal spacing
            command              Callback command used when a button is
                                 pressed.  Value supplied to command
                                 depends on values in buttonValueList

         Note: Number of buttons on the dialog depends upon the maximum
               length of any button[Text|Geom|Image|Value]List specified.
               Values of None are substituted for lists that are shorter
               than the max length
         """

        # Inherits from DirectFrame
        optiondefs = (
            # Define type of DirectGuiWidget
            ('dialogName',        None,          INITOPT),
            # Default position is slightly forward in Y, so as not to
            # intersect the near plane, which is incorrectly set to 0
            # in DX for some reason.
            ('pos',               (0, 0.1, 0),   None),
            ('pad',               (0.1, 0.1),    None),
            ('text',              '',            None),
            ('text_align',        TextNode.ALeft,   None),
            ('text_scale',        0.06,          None),
            ('image',  getDefaultDialogGeom(),   None),
            ('relief',            None,          None),
            ('buttonTextList',    [],            INITOPT),
            ('buttonGeomList',    [],            INITOPT),
            ('buttonImageList',   [],            INITOPT),
            ('buttonValueList',   [],            INITOPT),
            ('buttonHotKeyList',  [],            INITOPT),
            ('button_borderWidth',(.01,.01),     None),
            ('button_pad',        (.01,.01),     None),
            ('button_relief',     RAISED,        None),
            ('button_text_scale'  , 0.06,        None),
            ('buttonSize',        None,          INITOPT),
            ('topPad',            0.06,          INITOPT),
            ('midPad',            0.12,          INITOPT),
            ('buttonPadSF',       1.05,          INITOPT),
            # Alpha of fade screen behind dialog
            ('fadeScreen',        0,             None),
            ('command',           None,          None),
            ('extraArgs',         [],            None),
            ('sortOrder',    NO_FADE_SORT_INDEX, None),
            )
        # Merge keyword options with default options
        self.defineoptions(kw, optiondefs, dynamicGroups = ("button",))

        # Initialize superclasses
        DirectFrame.__init__(self, parent)

        if not self['dialogName']:
            self['dialogName'] = 'DirectDialog_' + `DirectDialog.PanelIndex`
        # Clean up any previously existing panel with the same unique
        # name.  We don't allow any two panels with the same name to
        # coexist.
        cleanupDialog(self['dialogName'])
        # Store this panel in our map of all open panels.
        DirectDialog.AllDialogs[self['dialogName']] = self
        DirectDialog.PanelIndex += 1

        # Determine number of buttons
        self.numButtons = max(len(self['buttonTextList']),
                              len(self['buttonGeomList']),
                              len(self['buttonImageList']),
                              len(self['buttonValueList']))
        # Create buttons
        self.buttonList = []
        index = 0
        for i in range(self.numButtons):
            name = 'Button' + `i`
            try:
                text = self['buttonTextList'][i]
            except IndexError:
                text = None
            try:
                geom = self['buttonGeomList'][i]
            except IndexError:
                geom = None
            try:
                image = self['buttonImageList'][i]
            except IndexError:
                image = None
            try:
                value = self['buttonValueList'][i]
            except IndexError:
                value = i
                self['buttonValueList'].append(i)
            try:
                hotKey = self['buttonHotKeyList'][i]
            except IndexError:
                hotKey = None
            button = self.createcomponent(
                name, (), "button",
                DirectButton, (self,),
                text = text,
                geom = geom,
                image = image,
                suppressKeys = self['suppressKeys'],
                frameSize = self['buttonSize'],
                command = lambda s = self, v = value: s.buttonCommand(v)
                )
            self.buttonList.append(button)

        # Update dialog when everything has been initialised
        self.postInitialiseFuncList.append(self.configureDialog)
        self.initialiseoptions(DirectDialog)

    def configureDialog(self):
        # Set up hot key bindings
        bindList = zip(self.buttonList, self['buttonHotKeyList'],
                       self['buttonValueList'])
        for button, hotKey, value in bindList:
            if ((type(hotKey) == types.ListType) or
                (type(hotKey) == types.TupleType)):
                for key in hotKey:
                    button.bind('press-' + key + '-', self.buttonCommand,
                                extraArgs = [value])
                    self.bind('press-' + key + '-', self.buttonCommand,
                              extraArgs = [value])

            else:
                button.bind('press-' + hotKey + '-',self.buttonCommand,
                            extraArgs = [value])
                self.bind('press-' + hotKey + '-', self.buttonCommand,
                          extraArgs = [value])
        # Position buttons and text
        pad = self['pad']
        image = self.component('image0')
        # Get size of text/geom without image (for state 0)
        if image:
            image.reparentTo(hidden)
        bounds = self.stateNodePath[0].getTightBounds()
        if image:
            image.reparentTo(self.stateNodePath[0])
        l = bounds[0][0]
        r = bounds[1][0]
        b = bounds[0][2]
        t = bounds[1][2]
        # Center text and geom around origin
        # How far is center of text from origin?
        xOffset = -(l+r)*0.5
        zOffset = -(b+t)*0.5
        # Update bounds to reflect text movement
        l += xOffset
        r += xOffset
        b += zOffset
        t += zOffset
        # Offset text and geom to center
        if self['text']:
            self['text_pos'] = (self['text_pos'][0] + xOffset,
                                self['text_pos'][1] + zOffset)
        if self['geom']:
            self['geom_pos'] = Point3(self['geom_pos'][0] + xOffset,
                                      self['geom_pos'][1],
                                      self['geom_pos'][2] + zOffset)
        if self.numButtons != 0:
            bpad = self['button_pad']
            # Get button size
            if self['buttonSize']:
                # Either use given size
                buttonSize = self['buttonSize']
                bl = buttonSize[0]
                br = buttonSize[1]
                bb = buttonSize[2]
                bt = buttonSize[3]
            else:
                # Or get bounds of union of buttons
                bl = br = bb = bt = 0
                for button in self.buttonList:
                    bounds = button.stateNodePath[0].getTightBounds()
                    bl = min(bl, bounds[0][0])
                    br = max(br, bounds[1][0])
                    bb = min(bb, bounds[0][2])
                    bt = max(bt, bounds[1][2])
                bl -= bpad[0]
                br += bpad[0]
                bb -= bpad[1]
                bt += bpad[1]
                # Now resize buttons to match largest
                for button in self.buttonList:
                    button['frameSize'] = (bl,br,bb,bt)
            # Must compensate for scale
            scale = self['button_scale']
            # Can either be a Vec3 or a tuple of 3 values
            if (isinstance(scale, Vec3) or
                (type(scale) == types.ListType) or
                (type(scale) == types.TupleType)):
                sx = scale[0]
                sz = scale[2]
            elif ((type(scale) == types.IntType) or
                  (type(scale) == types.FloatType)):
                sx = sz = scale
            else:
                sx = sz = 1
            bl *= sx
            br *= sx
            bb *= sz
            bt *= sz
            # Position buttons
            # Calc button width and height
            bHeight = bt - bb
            bWidth = br - bl
            # Add pad between buttons
            bSpacing = self['buttonPadSF'] * bWidth
            bPos = -bSpacing * (self.numButtons - 1)*0.5
            index = 0
            for button in self.buttonList:
                button.setPos(bPos + index * bSpacing, 0,
                              b - self['midPad'] - bpad[1] - bt)
                index += 1
            bMax = bPos + bSpacing * (self.numButtons - 1)
        else:
            bpad = 0
            bl = br = bb = bt = 0
            bPos = 0
            bMax = 0
            bpad = (0,0)
            bHeight = bWidth = 0
        # Resize frame to fit text and buttons
        l = min(bPos + bl, l) - pad[0]
        r = max(bMax + br, r) + pad[0]
        # reduce bottom by pad, button height and 2*button pad
        b = min(b - self['midPad'] - bpad[1] - bHeight - bpad[1], b) - pad[1]
        t = t + self['topPad'] + pad[1]
        self['image_scale'] = (r - l, 1, t - b)
        # Center frame about text and buttons
        self['image_pos'] = ((l+r)*0.5, 0.0,(b+t)*0.5)
        self.resetFrameSize()

    def show(self):
        """show(self)
        """
        if self['fadeScreen']:
            base.transitions.fadeScreen(self['fadeScreen'])
        NodePath.show(self)

    def hide(self):
        """hide(self)
        """
        if self['fadeScreen']:
            base.transitions.noTransitions()
        NodePath.hide(self)

    def buttonCommand(self, value, event = None):
        if self['command']:
            self['command'](value)

    def setMessage(self, message):
        self['text'] = message
        self.configureDialog()

    def cleanup(self):
        """unload(self)
        """
        # Remove this panel out of the AllDialogs list
        uniqueName = self['dialogName']
        if DirectDialog.AllDialogs.has_key(uniqueName):
            del DirectDialog.AllDialogs[uniqueName]
        self.destroy()

    def destroy(self):
        if self['fadeScreen']:
            base.transitions.noTransitions()
        DirectFrame.destroy(self)

class OkDialog(DirectDialog):
    def __init__(self, parent = aspect2d, **kw):
        # Inherits from DirectFrame
        optiondefs = (
            # Define type of DirectGuiWidget
            ('buttonTextList',  ['OK'],       INITOPT),
            ('buttonValueList', [DIALOG_OK],          INITOPT),
            )
        # Merge keyword options with default options
        self.defineoptions(kw, optiondefs)
        DirectDialog.__init__(self, parent)
        self.initialiseoptions(OkDialog)

class OkCancelDialog(DirectDialog):
    def __init__(self, parent = aspect2d, **kw):
        # Inherits from DirectFrame
        optiondefs = (
            # Define type of DirectGuiWidget
            ('buttonTextList',  ['OK','Cancel'],       INITOPT),
            ('buttonValueList', [DIALOG_OK,DIALOG_CANCEL], INITOPT),
            )
        # Merge keyword options with default options
        self.defineoptions(kw, optiondefs)
        DirectDialog.__init__(self, parent)
        self.initialiseoptions(OkCancelDialog)

class YesNoDialog(DirectDialog):
    def __init__(self, parent = aspect2d, **kw):
        # Inherits from DirectFrame
        optiondefs = (
            # Define type of DirectGuiWidget
            ('buttonTextList',  ['Yes', 'No'],       INITOPT),
            ('buttonValueList', [DIALOG_YES,DIALOG_NO], INITOPT),
            )
        # Merge keyword options with default options
        self.defineoptions(kw, optiondefs)
        DirectDialog.__init__(self, parent)
        self.initialiseoptions(YesNoDialog)

class YesNoCancelDialog(DirectDialog):
    def __init__(self, parent = aspect2d, **kw):
        # Inherits from DirectFrame
        optiondefs = (
            # Define type of DirectGuiWidget
            ('buttonTextList',  ['Yes', 'No', 'Cancel'],  INITOPT),
            ('buttonValueList', [DIALOG_YES,DIALOG_NO,DIALOG_CANCEL],
             INITOPT),
            )
        # Merge keyword options with default options
        self.defineoptions(kw, optiondefs)
        DirectDialog.__init__(self, parent)
        self.initialiseoptions(YesNoCancelDialog)

class RetryCancelDialog(DirectDialog):
    def __init__(self, parent = aspect2d, **kw):
        # Inherits from DirectFrame
        optiondefs = (
            # Define type of DirectGuiWidget
            ('buttonTextList',  ['Retry','Cancel'],   INITOPT),
            ('buttonValueList', [DIALOG_RETRY,DIALOG_CANCEL], INITOPT),
            )
        # Merge keyword options with default options
        self.defineoptions(kw, optiondefs)
        DirectDialog.__init__(self, parent)
        self.initialiseoptions(RetryCancelDialog)

<commit_msg>spelling fix<commit_after>from DirectGuiGlobals import *
from DirectFrame import *
from DirectButton import *

def findDialog(uniqueName):
    """findPanel(string uniqueName)

    Returns the panel whose uniqueName is given.  This is mainly
    useful for debugging, to get a pointer to the current onscreen
    panel of a particular type.

    """
    if DirectDialog.AllDialogs.has_key(uniqueName):
        return DirectDialog.AllDialogs[uniqueName]
    return None

def cleanupDialog(uniqueName):
    """cleanupPanel(string uniqueName)

    Cleans up (removes) the panel with the given uniqueName.  This
    may be useful when some panels know about each other and know
    that opening panel A should automatically close panel B, for
    instance.
    """

    if DirectDialog.AllDialogs.has_key(uniqueName):
        # calling cleanup() will remove it out of the AllDialogs dict
        # This way it will get removed from the dict even it we did
        # not clean it up using this interface (ie somebody called
        # self.cleanup() directly
        DirectDialog.AllDialogs[uniqueName].cleanup()

class DirectDialog(DirectFrame):

    AllDialogs = {}
    PanelIndex = 0

    def __init__(self, parent = aspect2d, **kw):
        """
        DirectDialog(kw)

        Creates a popup dialog to alert and/or interact with user.
        Some of the main keywords that can be used to customize the dialog:
            Keyword              Definition
            -------              ----------
            text                 Text message/query displayed to user
            geom                 Geometry to be displayed in dialog
            buttonTextList       List of text to show on each button
            buttonGeomList       List of geometry to show on each button
            buttonImageList      List of images to show on each button
            buttonValueList      List of values sent to dialog command for
                                 each button.  If value is [] then the
                                 ordinal rank of the button is used as
                                 its value
            buttonHotKeyList     List of hotkeys to bind to each button.
                                 Typing hotkey is equivalent to pressing
                                 the corresponding button.
            suppressKeys         Set to true if you wish to suppress keys
                                 (i.e. Dialog eats key event), false if
                                 you wish Dialog to pass along key event
            buttonSize           4-tuple used to specify custom size for
                                 each button (to make bigger then geom/text
                                 for example)
            pad                  Space between border and interior graphics
            topPad               Extra space added above text/geom/image
            midPad               Extra space added between text/buttons
            buttonPadSF          Scale factor used to expand/contract
                                 button horizontal spacing
            command              Callback command used when a button is
                                 pressed.  Value supplied to command
                                 depends on values in buttonValueList

         Note: Number of buttons on the dialog depends upon the maximum
               length of any button[Text|Geom|Image|Value]List specified.
               Values of None are substituted for lists that are shorter
               than the max length
         """

        # Inherits from DirectFrame
        optiondefs = (
            # Define type of DirectGuiWidget
            ('dialogName',        None,          INITOPT),
            # Default position is slightly forward in Y, so as not to
            # intersect the near plane, which is incorrectly set to 0
            # in DX for some reason.
            ('pos',               (0, 0.1, 0),   None),
            ('pad',               (0.1, 0.1),    None),
            ('text',              '',            None),
            ('text_align',        TextNode.ALeft,   None),
            ('text_scale',        0.06,          None),
            ('image',  getDefaultDialogGeom(),   None),
            ('relief',            None,          None),
            ('buttonTextList',    [],            INITOPT),
            ('buttonGeomList',    [],            INITOPT),
            ('buttonImageList',   [],            INITOPT),
            ('buttonValueList',   [],            INITOPT),
            ('buttonHotKeyList',  [],            INITOPT),
            ('button_borderWidth',(.01,.01),     None),
            ('button_pad',        (.01,.01),     None),
            ('button_relief',     RAISED,        None),
            ('button_text_scale'  , 0.06,        None),
            ('buttonSize',        None,          INITOPT),
            ('topPad',            0.06,          INITOPT),
            ('midPad',            0.12,          INITOPT),
            ('buttonPadSF',       1.05,          INITOPT),
            # Alpha of fade screen behind dialog
            ('fadeScreen',        0,             None),
            ('command',           None,          None),
            ('extraArgs',         [],            None),
            ('sortOrder',    NO_FADE_SORT_INDEX, None),
            )
        # Merge keyword options with default options
        self.defineoptions(kw, optiondefs, dynamicGroups = ("button",))

        # Initialize superclasses
        DirectFrame.__init__(self, parent)

        if not self['dialogName']:
            self['dialogName'] = 'DirectDialog_' + `DirectDialog.PanelIndex`
        # Clean up any previously existing panel with the same unique
        # name.  We don't allow any two panels with the same name to
        # coexist.
        cleanupDialog(self['dialogName'])
        # Store this panel in our map of all open panels.
        DirectDialog.AllDialogs[self['dialogName']] = self
        DirectDialog.PanelIndex += 1

        # Determine number of buttons
        self.numButtons = max(len(self['buttonTextList']),
                              len(self['buttonGeomList']),
                              len(self['buttonImageList']),
                              len(self['buttonValueList']))
        # Create buttons
        self.buttonList = []
        index = 0
        for i in range(self.numButtons):
            name = 'Button' + `i`
            try:
                text = self['buttonTextList'][i]
            except IndexError:
                text = None
            try:
                geom = self['buttonGeomList'][i]
            except IndexError:
                geom = None
            try:
                image = self['buttonImageList'][i]
            except IndexError:
                image = None
            try:
                value = self['buttonValueList'][i]
            except IndexError:
                value = i
                self['buttonValueList'].append(i)
            try:
                hotKey = self['buttonHotKeyList'][i]
            except IndexError:
                hotKey = None
            button = self.createcomponent(
                name, (), "button",
                DirectButton, (self,),
                text = text,
                geom = geom,
                image = image,
                suppressKeys = self['suppressKeys'],
                frameSize = self['buttonSize'],
                command = lambda s = self, v = value: s.buttonCommand(v)
                )
            self.buttonList.append(button)

        # Update dialog when everything has been initialised
        self.postInitialiseFuncList.append(self.configureDialog)
        self.initialiseoptions(DirectDialog)

    def configureDialog(self):
        # Set up hot key bindings
        bindList = zip(self.buttonList, self['buttonHotKeyList'],
                       self['buttonValueList'])
        for button, hotKey, value in bindList:
            if ((type(hotKey) == types.ListType) or
                (type(hotKey) == types.TupleType)):
                for key in hotKey:
                    button.bind('press-' + key + '-', self.buttonCommand,
                                extraArgs = [value])
                    self.bind('press-' + key + '-', self.buttonCommand,
                              extraArgs = [value])

            else:
                button.bind('press-' + hotKey + '-',self.buttonCommand,
                            extraArgs = [value])
                self.bind('press-' + hotKey + '-', self.buttonCommand,
                          extraArgs = [value])
        # Position buttons and text
        pad = self['pad']
        image = self.component('image0')
        # Get size of text/geom without image (for state 0)
        if image:
            image.reparentTo(hidden)
        bounds = self.stateNodePath[0].getTightBounds()
        if image:
            image.reparentTo(self.stateNodePath[0])
        l = bounds[0][0]
        r = bounds[1][0]
        b = bounds[0][2]
        t = bounds[1][2]
        # Center text and geom around origin
        # How far is center of text from origin?
        xOffset = -(l+r)*0.5
        zOffset = -(b+t)*0.5
        # Update bounds to reflect text movement
        l += xOffset
        r += xOffset
        b += zOffset
        t += zOffset
        # Offset text and geom to center
        if self['text']:
            self['text_pos'] = (self['text_pos'][0] + xOffset,
                                self['text_pos'][1] + zOffset)
        if self['geom']:
            self['geom_pos'] = Point3(self['geom_pos'][0] + xOffset,
                                      self['geom_pos'][1],
                                      self['geom_pos'][2] + zOffset)
        if self.numButtons != 0:
            bpad = self['button_pad']
            # Get button size
            if self['buttonSize']:
                # Either use given size
                buttonSize = self['buttonSize']
                bl = buttonSize[0]
                br = buttonSize[1]
                bb = buttonSize[2]
                bt = buttonSize[3]
            else:
                # Or get bounds of union of buttons
                bl = br = bb = bt = 0
                for button in self.buttonList:
                    bounds = button.stateNodePath[0].getTightBounds()
                    bl = min(bl, bounds[0][0])
                    br = max(br, bounds[1][0])
                    bb = min(bb, bounds[0][2])
                    bt = max(bt, bounds[1][2])
                bl -= bpad[0]
                br += bpad[0]
                bb -= bpad[1]
                bt += bpad[1]
                # Now resize buttons to match largest
                for button in self.buttonList:
                    button['frameSize'] = (bl,br,bb,bt)
            # Must compensate for scale
            scale = self['button_scale']
            # Can either be a Vec3 or a tuple of 3 values
            if (isinstance(scale, Vec3) or
                (type(scale) == types.ListType) or
                (type(scale) == types.TupleType)):
                sx = scale[0]
                sz = scale[2]
            elif ((type(scale) == types.IntType) or
                  (type(scale) == types.FloatType)):
                sx = sz = scale
            else:
                sx = sz = 1
            bl *= sx
            br *= sx
            bb *= sz
            bt *= sz
            # Position buttons
            # Calc button width and height
            bHeight = bt - bb
            bWidth = br - bl
            # Add pad between buttons
            bSpacing = self['buttonPadSF'] * bWidth
            bPos = -bSpacing * (self.numButtons - 1)*0.5
            index = 0
            for button in self.buttonList:
                button.setPos(bPos + index * bSpacing, 0,
                              b - self['midPad'] - bpad[1] - bt)
                index += 1
            bMax = bPos + bSpacing * (self.numButtons - 1)
        else:
            bpad = 0
            bl = br = bb = bt = 0
            bPos = 0
            bMax = 0
            bpad = (0,0)
            bHeight = bWidth = 0
        # Resize frame to fit text and buttons
        l = min(bPos + bl, l) - pad[0]
        r = max(bMax + br, r) + pad[0]
        # reduce bottom by pad, button height and 2*button pad
        b = min(b - self['midPad'] - bpad[1] - bHeight - bpad[1], b) - pad[1]
        t = t + self['topPad'] + pad[1]
        self['image_scale'] = (r - l, 1, t - b)
        # Center frame about text and buttons
        self['image_pos'] = ((l+r)*0.5, 0.0,(b+t)*0.5)
        self.resetFrameSize()

    def show(self):
        """show(self)
        """
        if self['fadeScreen']:
            base.transitions.fadeScreen(self['fadeScreen'])
        NodePath.show(self)

    def hide(self):
        """hide(self)
        """
        if self['fadeScreen']:
            base.transitions.noTransitions()
        NodePath.hide(self)

    def buttonCommand(self, value, event = None):
        if self['command']:
            self['command'](value)

    def setMessage(self, message):
        self['text'] = message
        self.configureDialog()

    def cleanup(self):
        """unload(self)
        """
        # Remove this panel out of the AllDialogs list
        uniqueName = self['dialogName']
        if DirectDialog.AllDialogs.has_key(uniqueName):
            del DirectDialog.AllDialogs[uniqueName]
        self.destroy()

    def destroy(self):
        if self['fadeScreen']:
            base.transitions.noTransitions()
        DirectFrame.destroy(self)

class OkDialog(DirectDialog):
    def __init__(self, parent = aspect2d, **kw):
        # Inherits from DirectFrame
        optiondefs = (
            # Define type of DirectGuiWidget
            ('buttonTextList',  ['OK'],       INITOPT),
            ('buttonValueList', [DIALOG_OK],          INITOPT),
            )
        # Merge keyword options with default options
        self.defineoptions(kw, optiondefs)
        DirectDialog.__init__(self, parent)
        self.initialiseoptions(OkDialog)

class OkCancelDialog(DirectDialog):
    def __init__(self, parent = aspect2d, **kw):
        # Inherits from DirectFrame
        optiondefs = (
            # Define type of DirectGuiWidget
            ('buttonTextList',  ['OK','Cancel'],       INITOPT),
            ('buttonValueList', [DIALOG_OK,DIALOG_CANCEL], INITOPT),
            )
        # Merge keyword options with default options
        self.defineoptions(kw, optiondefs)
        DirectDialog.__init__(self, parent)
        self.initialiseoptions(OkCancelDialog)

class YesNoDialog(DirectDialog):
    def __init__(self, parent = aspect2d, **kw):
        # Inherits from DirectFrame
        optiondefs = (
            # Define type of DirectGuiWidget
            ('buttonTextList',  ['Yes', 'No'],       INITOPT),
            ('buttonValueList', [DIALOG_YES,DIALOG_NO], INITOPT),
            )
        # Merge keyword options with default options
        self.defineoptions(kw, optiondefs)
        DirectDialog.__init__(self, parent)
        self.initialiseoptions(YesNoDialog)

class YesNoCancelDialog(DirectDialog):
    def __init__(self, parent = aspect2d, **kw):
        # Inherits from DirectFrame
        optiondefs = (
            # Define type of DirectGuiWidget
            ('buttonTextList',  ['Yes', 'No', 'Cancel'],  INITOPT),
            ('buttonValueList', [DIALOG_YES,DIALOG_NO,DIALOG_CANCEL],
             INITOPT),
            )
        # Merge keyword options with default options
        self.defineoptions(kw, optiondefs)
        DirectDialog.__init__(self, parent)
        self.initialiseoptions(YesNoCancelDialog)

class RetryCancelDialog(DirectDialog):
    def __init__(self, parent = aspect2d, **kw):
        # Inherits from DirectFrame
        optiondefs = (
            # Define type of DirectGuiWidget
            ('buttonTextList',  ['Retry','Cancel'],   INITOPT),
            ('buttonValueList', [DIALOG_RETRY,DIALOG_CANCEL], INITOPT),
            )
        # Merge keyword options with default options
        self.defineoptions(kw, optiondefs)
        DirectDialog.__init__(self, parent)
        self.initialiseoptions(RetryCancelDialog)

<|endoftext|>


<commit_before>import numpy as np
from .matutils import mkvc

def cyl2cart(grid):
    """
    Take a grid defined in cylindrical coordinates :math:`(r, \theta, z)` and
    transform it to cartesian coordinates.
    """
    return np.hstack([
        mkvc(grid[:, 0]*np.cos(grid[:, 1]), 2),
        mkvc(grid[:, 0]*np.sin(grid[:, 1]), 2),
        mkvc(grid[:, 2], 2)
    ])

def rotationMatrixFromNormals(v0, v1, tol=1e-20):
    """
        Performs the minimum number of rotations to define a rotation from the direction indicated by the vector n0 to the direction indicated by n1.
        The axis of rotation is n0 x n1
        https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula

        :param numpy.array v0: vector of length 3
        :param numpy.array v1: vector of length 3
        :param tol = 1e-20: tolerance. If the norm of the cross product between the two vectors is below this, no rotation is performed
        :rtype: numpy.array, 3x3
        :return: rotation matrix which rotates the frame so that n0 is aligned with n1

    """

    # ensure both n0, n1 are vectors of length 1
    assert len(v0) == 3, "Length of n0 should be 3"
    assert len(v1) == 3, "Length of n1 should be 3"

    # ensure both are true normals
    n0 = v0*1./np.linalg.norm(v0)
    n1 = v1*1./np.linalg.norm(v1)

    n0dotn1 = n0.dot(n1)

    # define the rotation axis, which is the cross product of the two vectors
    rotAx = np.cross(n0, n1)

    if np.linalg.norm(rotAx) < tol:
        return np.eye(3, dtype=float)

    rotAx *= 1./np.linalg.norm(rotAx)

    cosT = n0dotn1/(np.linalg.norm(n0)*np.linalg.norm(n1))
    sinT = np.sqrt(1.-n0dotn1**2)

    ux = np.array([[0., -rotAx[2], rotAx[1]], [rotAx[2], 0., -rotAx[0]], [-rotAx[1], rotAx[0], 0.]], dtype=float)

    return np.eye(3, dtype=float) + sinT*ux + (1.-cosT)*(ux.dot(ux))

def rotatePointsFromNormals(XYZ, n0, n1, x0=np.r_[0., 0., 0.]):
    """
        rotates a grid so that the vector n0 is aligned with the vector n1

        :param numpy.array n0: vector of length 3, should have norm 1
        :param numpy.array n1: vector of length 3, should have norm 1
        :param numpy.array x0: vector of length 3, point about which we perform the rotation
        :rtype: numpy.array, 3x3
        :return: rotation matrix which rotates the frame so that n0 is aligned with n1
    """

    R = rotationMatrixFromNormals(n0, n1)

    assert XYZ.shape[1] == 3, "Grid XYZ should be 3 wide"
    assert len(x0) == 3, "x0 should have length 3"

    X0 = np.ones([XYZ.shape[0], 1])*mkvc(x0)

    return (XYZ - X0).dot(R.T) + X0 # equivalent to (R*(XYZ - X0)).T + X0
<commit_msg>allow a single point to be input to the coordutils cyl to cartesian conversion<commit_after>import numpy as np
from .matutils import mkvc

def cyl2cart(grid):
    """
    Take a grid defined in cylindrical coordinates :math:`(r, \theta, z)` and
    transform it to cartesian coordinates.
    """
    grid = np.atleast_2d(grid)
    return np.hstack([
        mkvc(grid[:, 0]*np.cos(grid[:, 1]), 2),
        mkvc(grid[:, 0]*np.sin(grid[:, 1]), 2),
        mkvc(grid[:, 2], 2)
    ])

def rotationMatrixFromNormals(v0, v1, tol=1e-20):
    """
        Performs the minimum number of rotations to define a rotation from the direction indicated by the vector n0 to the direction indicated by n1.
        The axis of rotation is n0 x n1
        https://en.wikipedia.org/wiki/Rodrigues%27_rotation_formula

        :param numpy.array v0: vector of length 3
        :param numpy.array v1: vector of length 3
        :param tol = 1e-20: tolerance. If the norm of the cross product between the two vectors is below this, no rotation is performed
        :rtype: numpy.array, 3x3
        :return: rotation matrix which rotates the frame so that n0 is aligned with n1

    """

    # ensure both n0, n1 are vectors of length 1
    assert len(v0) == 3, "Length of n0 should be 3"
    assert len(v1) == 3, "Length of n1 should be 3"

    # ensure both are true normals
    n0 = v0*1./np.linalg.norm(v0)
    n1 = v1*1./np.linalg.norm(v1)

    n0dotn1 = n0.dot(n1)

    # define the rotation axis, which is the cross product of the two vectors
    rotAx = np.cross(n0, n1)

    if np.linalg.norm(rotAx) < tol:
        return np.eye(3, dtype=float)

    rotAx *= 1./np.linalg.norm(rotAx)

    cosT = n0dotn1/(np.linalg.norm(n0)*np.linalg.norm(n1))
    sinT = np.sqrt(1.-n0dotn1**2)

    ux = np.array([[0., -rotAx[2], rotAx[1]], [rotAx[2], 0., -rotAx[0]], [-rotAx[1], rotAx[0], 0.]], dtype=float)

    return np.eye(3, dtype=float) + sinT*ux + (1.-cosT)*(ux.dot(ux))

def rotatePointsFromNormals(XYZ, n0, n1, x0=np.r_[0., 0., 0.]):
    """
        rotates a grid so that the vector n0 is aligned with the vector n1

        :param numpy.array n0: vector of length 3, should have norm 1
        :param numpy.array n1: vector of length 3, should have norm 1
        :param numpy.array x0: vector of length 3, point about which we perform the rotation
        :rtype: numpy.array, 3x3
        :return: rotation matrix which rotates the frame so that n0 is aligned with n1
    """

    R = rotationMatrixFromNormals(n0, n1)

    assert XYZ.shape[1] == 3, "Grid XYZ should be 3 wide"
    assert len(x0) == 3, "x0 should have length 3"

    X0 = np.ones([XYZ.shape[0], 1])*mkvc(x0)

    return (XYZ - X0).dot(R.T) + X0 # equivalent to (R*(XYZ - X0)).T + X0
<|endoftext|>