dipolar_segid.py

from math import acos, pi, sqrt
import re
import string
import sys


class Dipole:
    def __init__(self, id, D, segid, Dxx, Dyy, Dzz):
        Dmin = min(min(Dxx, Dyy), Dzz)
        Dmax = max(max(Dxx, Dyy), Dzz)

        if D < Dmin:
            print("changing dipole id", id, segid, "from", D, "to", Dmin)
            D = Dmin

        elif D > Dmax:
            print("changing dipole id", id, segid, "from", D, "to", Dmax)
            D = Dmax

        self.id = id
        self.D = D
        self.segid = segid

        x = (D - Dzz) / (Dyy - Dzz)  # yz plane
        y = (D - Dxx) / (Dzz - Dxx)  # xz plane
        z = (D - Dyy) / (Dxx - Dyy)  # xy plane

        # print 'x, y, z =', x, y, z

        if (x < 0) or (x > 1):
            s1 = sqrt(y)
            c1 = sqrt(1 - y)
            s2 = sqrt(z)
            c2 = sqrt(1 - z)

            self.excluded = 0
            self.circle1 = [(c1, 0, s1), (c1, 0, -s1), (s2, c2, 0), (s2, -c2, 0)]
            self.circle2 = [(-c1, 0, s1), (-c1, 0, -s1), (-s2, c2, 0), (-s2, -c2, 0)]

        elif (y < 0) or (y > 1):
            s1 = sqrt(z)
            c1 = sqrt(1 - z)
            s2 = sqrt(x)
            c2 = sqrt(1 - x)

            self.excluded = 1
            self.circle1 = [(s1, c1, 0), (-s1, c1, 0), (0, s2, c2), (0, s2, -c2)]
            self.circle2 = [(s1, -c1, 0), (-s1, -c1, 0), (0, -s2, c2), (0, -s2, -c2)]

        else:  # z <= 0  or z >= 1
            s1 = sqrt(x)
            c1 = sqrt(1 - x)
            s2 = sqrt(y)
            c2 = sqrt(1 - y)

            self.excluded = 2
            self.circle1 = [(0, s1, c1), (0, -s1, c1), (c2, 0, s2), (-c2, 0, s2)]
            self.circle2 = [(0, s1, -c1), (0, -s1, -c1), (c2, 0, -s2), (-c2, 0, -s2)]


class DipolePair:
    def __init__(self, id1, id2, segid1, segid2, phimin1, phimax1, phimin2, phimax2):
        if phimin2 < phimin1:
            (phimin1, phimin2) = (phimin2, phimin1)
            (phimax1, phimax2) = (phimax2, phimax1)

        if phimin2 < phimax1:  # one range
            phimax1 = 90
            phimin2 = 90

        included = phimax2 - phimin2 + phimax1 - phimin1

        self.id1 = id1
        self.id2 = id2
        self.segid1 = segid1
        self.segid2 = segid2
        self.phimin1 = phimin1
        self.phimax1 = phimax1
        self.phimin2 = phimin2
        self.phimax2 = phimax2
        self.phimid1 = 0.5 * (self.phimin1 + self.phimax1)
        self.phiwidth1 = 0.5 * (self.phimax1 - self.phimin1)
        self.phimid2 = 0.5 * (self.phimin2 + self.phimax2)
        self.phiwidth2 = 0.5 * (self.phimax2 - self.phimin2)

        self.excluded = 1 - included / 180

    def write(self, file):
        file.write("assign (resid ")
        file.write(self.id1)
        file.write(" and name N and segid ")
        file.write(self.segid1)
        file.write(" ) (resid ")
        file.write(self.id1)
        file.write(" and name HN and segid ")
        file.write(self.segid1)
        file.write(" ) (resid ")
        file.write(self.id2)
        file.write(" and name N and segid ")
        file.write(self.segid2)
        file.write(" ) (resid ")
        file.write(self.id2)
        file.write(" and name HN and segid ")
        file.write(self.segid2)
        file.write(" ) ")
        file.write("%2.1f" % self.phimid1)
        file.write(" ")
        file.write("%2.1f" % self.phiwidth1)
        file.write(" ")
        file.write("%2.1f" % self.phimid2)
        file.write(" ")
        file.write("%2.1f" % self.phiwidth2)
        file.write(" ! excluded ")
        file.write("%4.3f" % self.excluded)
        file.write("\n")


def compare_dipoles(d1, d2):
    dmin = 2
    dmax = -2

    for point1 in d1.circle1:
        (x1, y1, z1) = point1
        # print '1:', x1, y1, z1

        for point2 in d2.circle1:
            (x2, y2, z2) = point2
            # print '  2:', x2, y2, z2

            d = x1 * x2 + y1 * y2 + z1 * z2

            dmin = min(d, dmin)
            dmax = max(d, dmax)
            # print '  3:', d, dmin, dmax

    phimin1 = acos(dmax) * 180 / pi
    phimax1 = acos(dmin) * 180 / pi

    dmin = 2
    dmax = -2

    for point1 in d1.circle2:
        (x1, y1, z1) = point1
        # print '4:', x1, y1, z1

        for point2 in d2.circle1:
            (x2, y2, z2) = point2
            # print '  5:', x2, y2, z2

            d = x1 * x2 + y1 * y2 + z1 * z2

            dmin = min(d, dmin)
            dmax = max(d, dmax)
            # print '  6:', d, dmin, dmax

    phimin2 = acos(dmax) * 180 / pi
    phimax2 = acos(dmin) * 180 / pi

    dipole_pair = DipolePair(
        d1.id, d2.id, d1.segid, d2.segid, phimin1, phimax1, phimin2, phimax2
    )

    return dipole_pair


def write_dipole_pairs(file, dipole_pairs):
    for dipole_pair in dipole_pairs:
        dipole_pair.write(file)


def read_dipoles(input_file, Dxx, Dyy, Dzz):
    file = open(input_file, "r")

    for i in range(1):  # header information
        line = file.readline()

    lines = file.readlines()

    file.close()

    dipoles = []
    # count = [ 0, 0, 0 ]

    for line in lines:
        line = string.strip(line)

        if line:
            # [ id, junk1, junk2, value ] = re.split('\t', line)
            [id, value, segid] = re.split("\t", line)

            dipole = Dipole(id, string.atof(value), segid, Dxx, Dyy, Dzz)
            dipoles.append(dipole)

            # count[dipole.excluded] = count[dipole.excluded] + 1

    print("found", len(dipoles), "dipoles")
    # print 'exclusion cases:', count[0], count[1], count[2]

    return dipoles


def compare_exclusions(dipole_pair1, dipole_pair2):
    if dipole_pair1.excluded > dipole_pair2.excluded:
        return -1

    elif dipole_pair1.excluded < dipole_pair2.excluded:
        return 1

    else:
        return 0


def process_dipoles(
    input_file, output_file, Dxx, Dyy, Dzz, correct_trace=1, min_excluded=0, max_pairs=0
):
    if correct_trace:
        t = -(Dxx + Dyy + Dzz) / 3

        if t != 0:
            print("adding", t, "to Dxx, Dyy, Dzz")

            Dxx = Dxx + t
            Dyy = Dyy + t
            Dzz = Dzz + t

    dipoles = read_dipoles(input_file, Dxx, Dyy, Dzz)
    file = open(output_file, "w")

    dipole_pairs = []

    for n1 in range(len(dipoles) - 1):
        print("working on dipole id", dipoles[n1].id)

        for n2 in range(n1 + 1, len(dipoles)):
            dipole_pair = compare_dipoles(dipoles[n1], dipoles[n2])

            if dipole_pair.excluded >= min_excluded:
                dipole_pairs.append(dipole_pair)

    dipole_pairs.sort(compare_exclusions)

    if max_pairs:
        dipole_pairs = dipole_pairs[:max_pairs]

    write_dipole_pairs(file, dipole_pairs)

    file.close()

    print("output", len(dipole_pairs), "dipole pairs")


if __name__ == "__main__":
    if len(sys.argv) != 6:
        print("args required: input file, output file, Dxx, Dyy, Dzz")

        sys.exit(1)

    input_file = sys.argv[1]
    output_file = sys.argv[2]
    Dxx = string.atof(sys.argv[3])
    Dyy = string.atof(sys.argv[4])
    Dzz = string.atof(sys.argv[5])

    process_dipoles(input_file, output_file, Dxx, Dyy, Dzz)