atomicDomino.py

from __future__ import print_function
import IMP
import IMP.domino
import IMP.core
import IMP.rmf
import RMF
import time
import IMP.algebra
import re
import sys
import os
import resource
from . import atomicDominoUtilities


class AtomicDomino:

    def __init__(self, model, protein, parameterFileName):
        self.model = model
        self.protein = protein
        self.namesToParticles = atomicDominoUtilities.makeNamesToParticles(
            protein)
        self.readParameters(parameterFileName)
        self.wroteNativeProtein = 0
        self.maxMem = 0

    def readParameters(self, parameterFileName):
        self.parameters = atomicDominoUtilities.readParameters(
            parameterFileName)

    def logMemory(self):
        currentMem = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
        currentMem *= 10 ** -6
        print("log memory: %s" % currentMem)
        if (currentMem > self.maxMem):
            self.maxMem = currentMem

    # Get parameter value. Throws dictionary exception if not found
    def getParam(self, paramName):
        paramValue = self.parameters[paramName]
        return paramValue

    def loadDominoHelpers(self):
        self.subsetStateScores = {}
        self.subsetRestraintScores = {}
        self.subsetStateFailures = {}
        self.nodesToAssignmentCount = {}

        self.totalAssignments = 0
        outputDir = self.getParam("output_directory")
        mtreeCytoscapeAssignmentFile = self.getParam(
            "mtree_cytoscape_assignment_file")
        self.mtreeCytoscapeAssignmentFh = open(
            os.path.join(outputDir, mtreeCytoscapeAssignmentFile), 'w')
        self.mtreeCytoscapeAssignmentFh.write("Assignments\n")

        self.mtNamesToIndices = {}

        self.subsetNamesToAssignmentFiles = {}

    # Simple way to calculate run-time for certain methods and procedures.
    # When reset is 0, compares the previous saved time to the current one
    # and returns the difference (saving the current time too).
    # When reset is 1, just saves the current time.
    def logTimePoint(self, reset):
        newTime = time.time()
        if (reset == 0):
            timeDifference = newTime - self.timePoint
            timeDifference = round(timeDifference, 0)
            self.timePoint = newTime
            return timeDifference
        else:

            self.timePoint = newTime
            return newTime

    def createParticleStatesTable(self):

        for particleName in self.particleInfo.keys():

            statesToCenters = {}
            dataArray = self.particleInfo[particleName]
            for i in range(len(dataArray)):
                [state, center] = dataArray[i]
                statesToCenters[state] = center

            statesList = []
            for stateIndex in sorted(statesToCenters.keys()):
                vector3d = IMP.algebra.Vector3D(statesToCenters[stateIndex])
                statesList.append(vector3d)
                # print "appending particle %s state %s center %s" %
                # (particleName, stateIndex, vector3d)
            xyzStates = IMP.domino.XYZStates(statesList)
            self.dominoPst.set_particle_states(
                self.namesToParticles[particleName],
                xyzStates)

    def quickParticleName(self, particle):
        return atomicDominoUtilities.quickParticleName(particle)

    def filterAssignments(self, assignments, subset, nodeIndex, rssft):

        filteredSubsets = []
        restraintList = []

        # make dependency graph for stats
        for r in IMP.get_restraints([self.model.get_root_restraint_set()]):
            restraintList.append(r)

        stateCounter = 0
        passedCounter = 0

        # create hdf5AssignmentContainer

        sFilter = rssft.get_subset_filter(subset, filteredSubsets)
        # check each unique state to see if passes filter
        filteredAssignments = []
        for assignment in assignments:

            if (sFilter is None or sFilter.get_is_ok(assignment)):
                # add to assignment container if it passes
                passedCounter += 1
                filteredAssignments.append(assignment)

            stateCounter += 1
        fraction = (passedCounter * 1.0) / (stateCounter * 1.0)
        print("%s states passed out of %s total for this subset (fraction %s)"
              % (passedCounter, stateCounter, fraction))
        if (passedCounter == 0):
            print("subset %s had 0 assignments (out of %s) pass. Exiting..."
                  % (subset, stateCounter))
            sys.exit()

        return filteredAssignments

    # Convert the name of the subset to something more readable. Currently
    # returning name as sorted list of atoms
    def quickSubsetName(self, subset):
        cleanName = self.cleanVertexName(subset)
        atomNameList = cleanName.split(" ")
        sortedList = sorted(atomNameList)

        name = " ".join(sortedList)
        return name

    def getNodeIndexList(self):
        return self.mt.get_vertices()

    def getLeafNodeIndexList(self):
        leafNodeIndexList = []
        for subset in self.subsets:
            index = self.getMtIndexForSubset(self.quickSubsetName(subset))
            leafNodeIndexList.append(index)
        return leafNodeIndexList

    # Create Domino sampler and give it all the other domino objects it needs
    def createSampler(self):
        s = IMP.domino.DominoSampler(self.model, self.dominoPst)

        s.set_merge_tree(self.mt)
        filterTables = []
        filterTables.append(self.rssft)
        s.set_subset_filter_tables(filterTables)

        s.set_assignments_table(self.lat)

        if (self.getParam("cross_subset_filtering") == 1):
            s.set_use_cross_subset_filtering(1)

        self.sampler = s

    # Use the model restraint set to get the interaction graph, junction tree,
    # and merge tree, and also get subsets
    # from the junction tree and return them
    def createSubsets(self):
        self.initializeParticleStatesTable()
        ig = IMP.domino.get_interaction_graph(
            [self.model.get_root_restraint_set()],
            self.dominoPst)
        print("interaction graph:")
        ig.show()
        jt = IMP.domino.get_junction_tree(ig)
        print("junction tree:")
        jt.show()
        self.subsetNamesToSubsets = {}
        mt = IMP.domino.get_balanced_merge_tree(jt)

        # make map of vertex indices to atoms in subsets
        for index in mt.get_vertices():
            subset = mt.get_vertex_name(index)
            subsetName = self.cleanVertexName(subset)
            self.mtNamesToIndices[subsetName] = index

        subsets = IMP.domino.get_subsets(jt)
        for subset in subsets:
            print("created subset %s" % subset)
            subsetName = self.quickSubsetName(subset)
            self.subsetNamesToSubsets[subsetName] = subset

        print("merge tree:")
        mt.show()

        self.ig = ig
        self.jt = jt
        self.mt = mt
        self.subsets = subsets

        self.parentSiblingMap = {}
        self.parentSiblingMap[self.mt.get_vertices()[-1]] = {}
        self.createSiblingMap(self.mt.get_vertices()[-1])

    def getMtIndexForSubset(self, subsetName):
        for index in self.mt.get_vertices():
            mtNode = self.mt.get_vertex_name(index)
            mtName = self.cleanVertexName(mtNode)
            mtName = mtName.rstrip()
            mtName = mtName.lstrip()

            mtNameList = mtName.split(" ")
            subsetNameList = subsetName.split(" ")
            mtNameListSorted = sorted(mtNameList)
            subsetNameListSorted = sorted(subsetNameList)
            if (" ".join(subsetNameListSorted) == " ".join(mtNameListSorted)):

                return index
        print("did not find merge tree index for subset name %s" % subsetName)
        sys.exit()

    def getDominoParticleNames(self):
        particles = self.dominoPst.get_particles()
        particleNameList = []
        for particle in particles:
            pName = self.quickParticleName(particle)
            particleNameList.append(pName)

        return particleNameList

    def getSubsetNameList(self):
        subsetNameList = []
        for subset in self.subsets:
            name = self.quickSubsetName(subset)
            subsetNameList.append(name)
        return subsetNameList

    def getMtIndexToParticles(self):

        mtIndexToParticles = {}
        allVertices = self.mt.get_vertices()
        for nodeIndex in allVertices:
            subset = self.mt.get_vertex_name(nodeIndex)
            particleList = self.quickSubsetName(subset)
            mtIndexToParticles[nodeIndex] = particleList
            print("createtMtIndexToParticles: index %s is particleList %s "
                  % (nodeIndex, particleList))

        return mtIndexToParticles

    def readTrajectoryFile(
        self,
        atomList,
        rh,
        frames,
        scoreOutputFile,
        skipDomino,
            flexibleAtoms):

        bestScore = 100000
        bestRmsd = 10000
        bestScoreFrame = 0
        bestRmsdFrame = 0
        scoreOutputFh = open(scoreOutputFile, 'w')
        for i in frames:
            try:
                IMP.rmf.load_frame(rh, i, self.protein)
            except Exception:
                print("exception in loading frame %s" % i)
                continue
            score = self.model.evaluate(False)
            # for leaf in leaves:
            #    xyzD = IMP.core.XYZ.decorate_particle(leaf)
            # print "read trajectory file: coordinates for frame %s particle %s
            # are %s" % (i, self.quickParticleName(leaf),
            # xyzD.get_coordinates())

            rmsd = self.calculateNativeRmsd(flexibleAtoms)
            scoreOutputFh.write("%s\t%s\t%s\n" % (i, score, rmsd))
            if (score < bestScore):
                bestScore = score
                bestScoreFrame = i

            if (rmsd < bestRmsd):
                bestRmsd = rmsd
                bestRmsdFrame = i
            if (skipDomino == 0):
                for atomName in atomList:

                    particle = self.namesToParticles[atomName]

                    # get grid index and coordinates
                    gridIndex = self.snapToGrid(particle)
                    center = self.grid.get_center(gridIndex)
                    pythonCenter = []
                    for coordinate in center:
                        pythonCenter.append(coordinate)

                    # grid indices are mapped to coordinate states. Check if
                    # we've seen this grid index
                    if ((gridIndex in self.particleStatesSeen[atomName]) == 0):
                        # set this particle state index to size of the
                        # dictionary, which effectively increments the index
                        # with each new state
                        currentSize = len(
                            list(self.particleStatesSeen[atomName].keys()))
                        self.particleStatesSeen[
                            atomName][
                            gridIndex] = currentSize
                    state = self.particleStatesSeen[atomName][gridIndex]
                    self.particleInfo[atomName].append([state, pythonCenter])
            else:
                print("didn't add domino states due to skip parameters")
        return [bestScore, bestScoreFrame, bestRmsd, bestRmsdFrame]

    # Get the grid index for the given particle. Returns an integer that
    # can be used to get the center of the grid space for discretizing
    # the particle
    def getParticleGridIndex(self, leaf):
        xyzDecorator = IMP.core.XYZ(leaf)
        coordinates = xyzDecorator.get_coordinates()
        extendedIndex = 0
        extendedIndex = self.grid.get_extended_index(coordinates)
        if (self.grid.get_has_index(extendedIndex) == 0):
            # mostly working with sparse grids now
            if (self.getParam("grid_type") == "sparse"):
                self.grid.add_voxel(extendedIndex, 1)
            else:
                self.grid.add_voxel(extendedIndex)

        index = self.grid.get_index(extendedIndex)
        return index

    # Set coordinates of this atom to be the center of the grid space
    # containing it (effectiely discretizing the system)
    def snapToGrid(self, particle):

        index = self.getParticleGridIndex(particle)
        center = self.grid.get_center(index)
        xyzDecorator = IMP.core.XYZ(particle)
        xyzDecorator.set_coordinates(center)

        return index

    # Take initial protein and snap every atom to the center of its gridpoint
    def discretizeNativeProtein(self):

        outputDir = self.getParam("output_directory")
        nativeSnappedFile = self.getParam("native_protein_snapped_output_file")

        leaves = IMP.atom.get_leaves(self.protein)
        for leaf in leaves:
            self.snapToGrid(leaf)

        IMP.atom.write_pdb(
            self.protein,
            os.path.join(
                outputDir,
                nativeSnappedFile))

    # Get particle representing the alpha carbon at the center of the peptide
    def getPeptideCa(self):

        # get all residue indices in the peptide
        residues = IMP.atom.get_by_type(self.protein, IMP.atom.RESIDUE_TYPE)
        peptideIndicesToResidues = {}
        for residueH in residues:
            chain = IMP.atom.get_chain(residueH)
            chainId = chain.get_id()
            residue = IMP.atom.Residue(residueH)
            if (chainId == self.getParam("peptide_chain")):
                peptideIndicesToResidues[residue.get_index()] = residue

        # use the min and max residue indices to get the residue in the middle
        # (rounding down)
        minPeptide = min(sorted(peptideIndicesToResidues.keys()))
        maxPeptide = max(sorted(peptideIndicesToResidues.keys()))
        centerPeptide = round(((maxPeptide - minPeptide) / 2 + minPeptide), 0)

        # get the particle corresponding to the ca atom at the middle residue
        # and return it,
        centerName = atomicDominoUtilities.makeParticleName(
            self.getParam("peptide_chain"),
            int(centerPeptide),
            "CA")
        centerParticle = self.namesToParticles[centerName]
        return centerParticle

    # Create grid object that will be used to create discrete states for each
    # particle
    def createGrid(self):

        protBb = IMP.atom.get_bounding_box(self.protein)

        gridSpacing = float(self.getParam("grid_spacing"))
        bufferSpace = float(self.getParam("grid_buffer_space"))

        protBb += bufferSpace  # add buffer around grid
        g = 0
        if (self.getParam("grid_type") == "sparse"):
            ca = self.getPeptideCa()
            xyzCa = IMP.core.XYZ(ca)
            g = IMP.algebra.SparseUnboundedIntGrid3D(
                gridSpacing, xyzCa.get_coordinates())
        else:
            g = IMP.algebra.DenseDoubleGrid3D(gridSpacing, protBb)

        self.grid = g

    # Create Particle States Table and for each particle in the system,
    # add XYZStates with that particle's initial location
    def initializeParticleStatesTable(self):

        dominoPst = IMP.domino.ParticleStatesTable()
        restrainedParticles = atomicDominoUtilities.getRestrainedParticles(
            self.protein, self.model, self.namesToParticles)

        for p in restrainedParticles:
            xyz = IMP.core.XYZ(p).get_coordinates()
            xyzStates = IMP.domino.XYZStates([xyz])
            dominoPst.set_particle_states(p, xyzStates)

        self.dominoPst = dominoPst

    def createUniqueLeafAssignments(self, particleNameList, particleInfo):

        size = len(particleInfo[particleNameList[0]])
        allAssignments = []

        for i in range(size):
            nextAssignment = []
            for particleName in particleNameList:
                dataArray = particleInfo[particleName]
                [state, center] = dataArray[i]
                nextAssignment.append(state)
            allAssignments.append(nextAssignment)

        # make unique assignments to avoid duplicates
        uniqueAssignments = {}
        for assignment in allAssignments:

            stateString = ""
            for index in assignment:
                stateString = stateString + str(index) + "_"
            uniqueAssignments[stateString] = assignment
        finalAssignments = []

        # add all assignments to final list
        for stateString in uniqueAssignments.keys():
            assignmentList = uniqueAssignments[stateString]
            assignment = IMP.domino.Assignment(assignmentList)
            finalAssignments.append(assignment)
        return finalAssignments

    # Read MD trajectory; for each particle, save all unique states, and for
    # each subset, save all assignments
    def readMdTrajectory(self, atomList, flexibleAtoms):

        # open trajectory file
        outputDir = self.getParam("output_directory")
        trajectoryFile = self.getParam("md_trajectory_output_file")
        fullFile = os.path.join(outputDir, trajectoryFile)
        rh = RMF.open_rmf_file(fullFile)
        IMP.rmf.set_hierarchies(rh, [self.protein])
        framesToRead = atomicDominoUtilities.getMdIntervalFrames(
            rh, int(self.getParam("md_interval")), self.protein)
        print("preparing to read md frames %s" % framesToRead)
        # prepare data structures for tracking
        particleInfo = {}
        # for each particle, list where each entry corresponds to an md step,
        # and its value [domino state, coordinates]
        # for each particle, dictionary where the key is the grid index and
        # value is domino state
        particleStatesSeen = {}
        for atomName in atomList:
            particleInfo[atomName] = []
            particleStatesSeen[atomName] = {}

        self.particleStatesSeen = particleStatesSeen
        self.particleInfo = particleInfo

        # read trajectory file
        mdScoreOutputFile = os.path.join(
            outputDir, "%s" %
            self.getParam("md_score_output_file"))
        [bestMdScore,
         bestScoreFrame,
         bestRmsd,
         bestRmsdFrame] = self.readTrajectoryFile(
            atomList,
            rh,
            framesToRead,
            mdScoreOutputFile,
            0,
            flexibleAtoms)

        # output best score information
        # print "try loading bad frame"
        self.singlePdbResults(
            fullFile,
            bestScoreFrame,
            self.getParam("best_md_score_output_file"))

        self.singlePdbResults(
            fullFile,
            bestRmsdFrame,
            self.getParam("best_md_rmsd_output_file"))
        self.singlePdbResults(
            fullFile,
            -1,
            self.getParam("final_md_frame_output_file"))

        self.bestMdScore = round(bestMdScore, 2)
        self.bestMdRmsd = round(bestRmsd, 2)
        self.bestMdScoreFrame = bestScoreFrame
        self.bestMdRmsdFrame = bestRmsdFrame

    def readCgTrajectories(self, atomList, flexibleAtoms):

        cgFileName = self.getParam("cg_output_file")
        bestCgScore = 10000000
        bestCgScoreFile = ""
        bestCgRmsd = 10000000
        bestCgRmsdFile = ""

        outputDir = self.getParam("output_directory")
        trajectoryFile = self.getParam("md_trajectory_output_file")
        fullFile = os.path.join(outputDir, trajectoryFile)
        rh = RMF.open_rmf_file(fullFile)
        IMP.rmf.set_hierarchies(rh, [self.protein])
        framesToRead = atomicDominoUtilities.getMdIntervalFrames(
            rh, int(self.getParam("cg_interval")), self.protein)

        skipCgDomino = int(self.getParam("skip_cg_domino"))

        if (len(framesToRead) > 0):
            for cgNumber in framesToRead:
                # Open next cg trajectory
                outputDir = self.getParam("output_directory")
                fullCgFileName = os.path.join(
                    outputDir, "%s%s" %
                    (cgFileName, cgNumber))
                rh = RMF.open_rmf_file(fullCgFileName)
                IMP.rmf.set_hierarchies(rh, [self.protein])

                # Only look at the bottom 20 frames
                frameCount = IMP.rmf.get_number_of_frames(rh, self.protein)
                cgFrames = []
                startFrameCount = 0
                if (frameCount > 20):
                    startFrameCount = frameCount - 20

                for i in range(startFrameCount, frameCount):
                    cgFrames.append(i)

                # Process trajectory
                cgScoreOutputFile = os.path.join(
                    outputDir, "%s%s" %
                    (self.getParam("cg_score_output_file"), cgNumber))
                [cgScore, cgScoreFrame, cgRmsd, cgRmsdFrame] = \
                    self.readTrajectoryFile(atomList, rh, cgFrames,
                                            cgScoreOutputFile, skipCgDomino,
                                            flexibleAtoms)
                print("cg number %s rmsd %s score %s"
                      % (cgNumber, cgRmsd, cgScore))
                # Update best score
                if (cgScore < bestCgScore):
                    bestCgScore = cgScore
                    bestCgScoreFile = fullCgFileName
                if (cgRmsd < bestCgRmsd):
                    bestCgRmsd = cgRmsd
                    bestCgRmsdFile = fullCgFileName

            # output best score information
            self.singlePdbResults(
                bestCgScoreFile,
                -1,
                self.getParam("best_cg_score_output_file"))
            self.singlePdbResults(
                bestCgRmsdFile,
                -1,
                self.getParam("best_cg_rmsd_output_file"))
            self.singlePdbResults(
                "%s%s" % (cgFileName, framesToRead[-1]),
                -1, self.getParam("final_cg_frame_output_file"))
            finalCgRmsd = self.calculateNativeRmsd(flexibleAtoms)
            print("final cg rmsd is %s " % finalCgRmsd)
        self.bestCgScore = round(bestCgScore, 2)
        self.bestCgRmsd = round(bestCgRmsd, 2)
        self.bestCgScoreFile = bestCgScoreFile
        self.bestCgRmsdFile = bestCgRmsdFile

    def singlePdbResults(self, trajectoryFile, frame, outputPdbFile):

        fullTrajectoryFile = os.path.join(
            self.getParam("output_directory"),
            trajectoryFile)
        fullOutputFile = os.path.join(
            self.getParam("output_directory"),
            outputPdbFile)
        rh = RMF.open_rmf_file(fullTrajectoryFile)
        IMP.rmf.set_hierarchies(rh, [self.protein])
        if (frame == -1):
            frame = IMP.rmf.get_number_of_frames(rh, self.protein) - 1
        IMP.rmf.load_frame(rh, frame, self.protein)
        IMP.atom.write_pdb(self.protein, fullOutputFile)

    def calculateRmsd(self, otherProtein, flexibleAtoms):
        otherNamesToParticles = atomicDominoUtilities.makeNamesToParticles(
            otherProtein)
        otherVector = []
        modelVector = []
        for pName in otherNamesToParticles.keys():
            if ((pName in flexibleAtoms) == 0):
                continue
            otherParticle = otherNamesToParticles[pName]
            modelParticle = self.namesToParticles[pName]
            otherVector.append(IMP.core.XYZ(otherParticle).get_coordinates())
            modelVector.append(IMP.core.XYZ(modelParticle).get_coordinates())
        rmsd = IMP.atom.get_rmsd(otherVector, modelVector)
        return rmsd

    def calculateNativeRmsd(self, flexibleAtoms):

        if (self.wroteNativeProtein == 0):
            pdbName = self.getParam("native_pdb_input_file")
            self.nativeModel = IMP.Model()
            self.nativeProtein = IMP.atom.read_pdb(
                pdbName,
                self.nativeModel,
                IMP.atom.ATOMPDBSelector())
            self.wroteNativeProtein = 1

        return self.calculateRmsd(self.nativeProtein, flexibleAtoms)

    def calculateTrajectoryRmsd(
        self,
        trajectoryFile,
        trajectoryFrame,
            flexibleAtoms):
        pdbName = self.getParam("native_pdb_input_file")
        otherProtein = IMP.atom.read_pdb(
            pdbName,
            self.nativeModel,
            IMP.atom.ATOMPDBSelector())
        outputDir = self.getParam("output_directory")
        fullFile = os.path.join(outputDir, trajectoryFile)
        print("open calculate traj rmf %s" % fullFile)
        rh = RMF.open_rmf_file(fullFile)
        IMP.rmf.set_hierarchies(rh, [otherProtein])
        if (trajectoryFrame == -1):
            trajectoryFrame = IMP.rmf.get_number_of_frames(
                rh, otherProtein) - 1
        IMP.rmf.load_frame(rh, trajectoryFrame, otherProtein)
        return self.calculateRmsd(otherProtein, flexibleAtoms)

    def createAllSubsetAssignments(self):

        lat = IMP.domino.ListAssignmentsTable()
        rssft = IMP.domino.RestraintScoreSubsetFilterTable(
            self.model.get_root_restraint_set(),
            self.dominoPst)

        leafNodeIndexList = self.getLeafNodeIndexList()

        for nodeIndex in leafNodeIndexList:
            # get subset for this leaf
            subset = self.mt.get_vertex_name(nodeIndex)
            particleNameList = []
            for particle in subset:
                particleNameList.append(self.quickParticleName(particle))
            print("creating initial assignments for leaf %s" % nodeIndex)
            # use particleInfo to create assignments and filter them
            assignments = self.createUniqueLeafAssignments(
                particleNameList, self.particleInfo)
            filteredAssignments = self.filterAssignments(
                assignments, subset, nodeIndex, rssft)

            # add assignemtns to container and listAssignmentTable
            packedAssignmentContainer = IMP.domino.PackedAssignmentContainer()
            for assignment in filteredAssignments:
                packedAssignmentContainer.add_assignment(assignment)
            lat.set_assignments(subset, packedAssignmentContainer)

        self.lat = lat
        self.rssft = rssft

    def runDomino(self):
        root = self.mt.get_vertices()[-1]
        completeAc = self.loadAssignments(root)
        self.completeAc = completeAc

    def loadAssignments(self, nodeIndex):

        children = self.mt.get_out_neighbors(nodeIndex)
        subset = self.mt.get_vertex_name(nodeIndex)
        heapCount = int(self.getParam("heap_count"))
        mine = IMP.domino.HeapAssignmentContainer(
            heapCount, self.rssft.get_subset_filter(subset, []))
        if len(children) == 0:
            print("computing assignments for leaf %s" % nodeIndex)

            self.sampler.load_vertex_assignments(nodeIndex, mine)
            print("leaf node %s has %s leaf assignments"
                  % (nodeIndex, mine.get_number_of_assignments()))
        else:
            if (children[0] > children[1]):
                children = [children[1], children[0]]
            # recurse on the two children
            firstAc = self.loadAssignments(children[0])
            secondAc = self.loadAssignments(children[1])
            self.logTimePoint(1)
            self.sampler.load_vertex_assignments(
                nodeIndex, firstAc, secondAc, mine)

            timeDifference = self.logTimePoint(0)
            print("Done Parent %s Assignments %s first child %s "
                  "second child %s time %s"
                  % (nodeIndex, mine.get_number_of_assignments(),
                     firstAc.get_number_of_assignments(),
                     secondAc.get_number_of_assignments(), timeDifference))
        self.totalAssignments += mine.get_number_of_assignments()
        self.logMemory()
        return mine

    def writeOutput(self, flexibleAtoms, startTime):
        bestAssignment = -1
        bestDominoScore = 100000
        bestAssignment = 0
        finalAssignments = self.completeAc.get_assignments()
        for assignment in finalAssignments:
            IMP.domino.load_particle_states(
                self.dominoPst.get_subset(),
                assignment,
                self.dominoPst)
            score = self.model.evaluate(False)
            if (score < bestDominoScore):
                bestAssignment = assignment
                bestDominoScore = round(score, 2)
        print("best domino score is %s " % bestDominoScore)
        print("best md score is %s" % self.bestMdScore)
        print("best md rmsd is %s" % self.bestMdRmsd)
        print("best cg score is %s" % self.bestCgScore)
        print("best cg rmsd is %s" % self.bestCgRmsd)
        print("merge tree contained %s total assignments"
              % self.totalAssignments)

        IMP.domino.load_particle_states(
            self.dominoPst.get_subset(),
            bestAssignment,
            self.dominoPst)
        dominoVsMdRmsd = round(
            self.calculateTrajectoryRmsd(
                self.getParam(
                    "md_trajectory_output_file"),
                self.bestMdScoreFrame,
                flexibleAtoms),
            2)
        cg = IMP.core.ConjugateGradients(self.model)
        cg.optimize(100)
        IMP.atom.write_pdb(
            self.protein,
            os.path.join(self.getParam("output_directory"),
                         self.getParam("minimum_domino_score_pdb")))

        dominoVsCgRmsd = round(
            self.calculateTrajectoryRmsd(
                self.bestCgScoreFile,
                -1,
                flexibleAtoms),
            2)
        dominoMinimizedScore = round(self.model.evaluate(False), 2)
        dominoRmsd = round(self.calculateNativeRmsd(flexibleAtoms), 2)

        runTime = round(time.time() - startTime, 2)
        print("final domino score (after cg): %s" % dominoMinimizedScore)
        print("final domino rmsd: %s" % dominoRmsd)
        print("best domino rmsd with best md score: %s" % dominoVsMdRmsd)
        print("domino rmsd with best cg score: %s" % dominoVsCgRmsd)
        print("Final Results\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s"
              % (self.bestMdScore, self.bestMdRmsd, self.bestCgScore,
                 self.bestCgRmsd, bestDominoScore, dominoRmsd,
                 dominoMinimizedScore, dominoVsCgRmsd, self.totalAssignments,
                 self.maxMem, runTime))

    #
    # Parallel methods
    #
    def createSubsetFromParticles(self, particleNames):
        particleNameList = particleNames.split(" ")
        particleList = []
        for particleName in particleNameList:
            particle = self.namesToParticles[particleName]
            particleList.append(particle)

        subset = IMP.domino.Subset(particleList)
        return [subset, particleList]

    def createHdf5AssignmentContainer(self, index, particleNames, read):
        root = self.getAssignmentContainerRoot(index, read)
        print("got root for index %s" % index)
        dataset = 0
        if (read == 1):
            dataset = root.get_child_index_data_set_2d(str(index))
        else:
            dataset = root.add_child_index_data_set_2d(str(index))
        print("added child index dataset")
        # firstDataset.set_size([
        [subset, particleOrder] = self.createSubsetFromParticles(particleNames)
        print("created subset for index %s" % index)
        hdf5Ac = IMP.domino.create_assignments_container(
            dataset, subset, particleOrder)
        print("returning from create")
        order = IMP.domino.get_order(subset, particleOrder)
        for nextInt in order:
            print("next int is %s" % nextInt)
        return [subset, hdf5Ac]

    def loadAssignmentsParallel(
        self,
        nodeIndex,
        particleInfo,
        mtIndexToNodeInfo,
        mtIndexToSubsetOrder,
            mtIndexToParticles):
        IMP.set_log_level(IMP.WARNING)

        if (("firstChild" in mtIndexToNodeInfo[nodeIndex]) == 0):
            print("writing file for leaf index %s" % nodeIndex)
            return (
                self.createAssignmentsParallel(
                    particleInfo,
                    nodeIndex,
                    mtIndexToParticles)
            )

        else:
            beginTime = self.logTimePoint(1)
            firstChildIndex = mtIndexToNodeInfo[nodeIndex]["firstChild"]
            [firstSubset, firstAc] = self.createHdf5AssignmentContainer(
                firstChildIndex, mtIndexToSubsetOrder[firstChildIndex], 1)
            firstAcCreateTime = self.logTimePoint(0)

            secondChildIndex = mtIndexToNodeInfo[nodeIndex]["secondChild"]
            [secondSubset, secondAc] = self.createHdf5AssignmentContainer(
                secondChildIndex, mtIndexToSubsetOrder[secondChildIndex], 1)
            secondAcCreateTime = self.logTimePoint(0)

            print("getting assignments for nodeIndex %s first child %s "
                  "second child %s"
                  % (nodeIndex, firstChildIndex, secondChildIndex))
            firstChildParticles = mtIndexToSubsetOrder[firstChildIndex]
            secondChildParticles = mtIndexToSubsetOrder[secondChildIndex]
            myParticles = mtIndexToParticles[nodeIndex]

            print("first child particles %s\nsecond child particles %s\n"
                  "my particles; %s\n"
                  % (firstChildParticles, secondChildParticles, myParticles))
            for p in firstSubset:
                print("next particle in first subset: %s"
                      % self.quickParticleName(p))

            for p in secondSubset:
                print("next particle in second subset: %s"
                      % self.quickParticleName(p))

            for assignment in firstAc.get_assignments():
                print("next assignment for first child %s: %s"
                      % (firstChildIndex, assignment))

            for assignment in secondAc.get_assignments():
                print("next assignment for second child %s: %s"
                      % (secondChildIndex, assignment))

            [mySubset, myAc] = self.createHdf5AssignmentContainer(
                nodeIndex, mtIndexToParticles[nodeIndex], 0)
            print("done creating hdf5")
            prepTime = self.logTimePoint(0)
            rssft = IMP.domino.RestraintScoreSubsetFilterTable(
                self.model.get_root_restraint_set(),
                self.dominoPst)

            IMP.domino.load_merged_assignments(
                firstSubset,
                firstAc,
                secondSubset,
                secondAc,
                [rssft],
                myAc)

            heapTime = self.logTimePoint(0)
            # myAc.add_assignments(heapAc.get_assignments())
            addTime = self.logTimePoint(0)
            for assignment in myAc.get_assignments():
                print("loadAssignmentsParallel next assignment for %s: %s"
                      % (nodeIndex, assignment))
            doneTime = self.logTimePoint(0)
            firstChildCount = firstAc.get_number_of_assignments()
            secondChildCount = secondAc.get_number_of_assignments()
            print("first count: %s second count: %s begin: %s firstAc: %s "
                  "secondAc: %s  prep: %s heap: %s add: %s done: %s"
                  % (firstChildCount, secondChildCount, beginTime,
                     firstAcCreateTime, secondAcCreateTime, prepTime,
                     heapTime, addTime, doneTime))
            subsetOrder = self.getSubsetOrderList(mySubset)
            return subsetOrder

    def createAssignmentsParallel(self, particleInfo, nodeIndex,
                                  mtIndexToParticles):

        subsetName = mtIndexToParticles[nodeIndex]
        print("starting assignments parallel leaf index %s subset name %s"
              % (nodeIndex, subsetName))
        [subset, particleList] = self.createSubsetFromParticles(subsetName)

        particleNameList = subsetName.split(" ")

        # create assignment by reading states
        finalAssignments = self.createUniqueLeafAssignments(
            particleNameList, particleInfo)

        rssft = IMP.domino.RestraintScoreSubsetFilterTable(
            self.model.get_root_restraint_set(),
            self.dominoPst)
        filteredAssignments = self.filterAssignments(
            finalAssignments, subset, nodeIndex, rssft)
        root = self.getAssignmentContainerRoot(nodeIndex, 0)
        dataset = root.add_child_index_data_set_2d(str(nodeIndex))
        dataset.set_size([0, len(subset)])

        hdf5AssignmentContainer = IMP.domino.create_assignments_container(
                dataset, subset, particleOrder)   # noqa: F821
        for assignment in filteredAssignments:
            hdf5AssignmentContainer.add_assignment(assignment)
        for assignment in hdf5AssignmentContainer.get_assignments():
            print("hdf5 assignment container node %s next assignment %s"
                  % (nodeIndex, assignment))

        # self.checkAssignments(subset, nodeIndex, particleOrder)
        subsetOrder = self.getSubsetOrderList(subset)
        print("leaf node returning with order %s" % subsetOrder)
        return subsetOrder

    # Write pymol session files for the interactions across atoms and all
    # subsets
    def writePymolData(self):

        outputDir = self.getParam("output_directory")
        geometry = IMP.domino.get_interaction_graph_geometry(self.ig)
        pymolInteractions = self.getParam("pymol_interactions_file")
        w = IMP.display.PymolWriter(os.path.join(outputDir, pymolInteractions))

        for gg in geometry:
            w.add_geometry(gg)

        pymolSubsets = self.getParam("pymol_subsets_file")
        geometry = IMP.domino.get_subset_graph_geometry(self.jt)
        w = IMP.display.PymolWriter(os.path.join(outputDir, pymolSubsets))
        for gg in geometry:
            w.add_geometry(gg)

    # Clean the default name of the vertex (in brackets and with each atom
    # contained in quotes) and return a string where [] and " are removed
    def cleanVertexName(self, vertexName):

        # not sure if any vertices still have a Subset prefix but keeping
        # anyway
        nodeRe = re.compile(r'Subset\(\[(.*?)\s*\]')
        secondNodeRe = re.compile(r'\[(.*?)\s*\]')  # atom name
        node = nodeRe.search(str(vertexName))
        secondNode = secondNodeRe.search(str(vertexName))
        vertexNameFinal = ""
        foundName = 0
        if node:
            foundName = node.group(1)
        if secondNode:
            foundName = secondNode.group(1)
        vertexNameFinal = foundName.replace('"', '')
        return vertexNameFinal

    def getSubsetOrderList(self, subset):
        subsetOrderList = []
        for particle in subset:
            name = self.quickParticleName(particle)
            subsetOrderList.append(name)
        subsetOrder = " ".join(subsetOrderList)
        return subsetOrder

    def checkAssignments(self, subset, nodeIndex, particleOrder):
        print("reading back in assignments for leaf index %s" % nodeIndex)
        root = self.getAssignmentContainerRoot(nodeIndex, 1)
        dataset = root.get_child_index_data_set_2d(str(nodeIndex))
        hdf5 = IMP.domino.create_assignments_container(
            dataset, subset, particleOrder)
        for assignment in hdf5.get_assignments():
            print("leaf index %s read back in %s" % (nodeIndex, assignment))

    def createSamplerLite(self):
        s = IMP.domino.DominoSampler(self.model, self.dominoPst)

        if (self.getParam("cross_subset_filtering") == 1):
            s.set_use_cross_subset_filtering(1)

        self.sampler = s

    def createSiblingMap(self, parentIndex):

        children = self.mt.get_out_neighbors(parentIndex)
        if (len(children) > 0):
            firstChild = children[0]
            secondChild = children[1]
            self.parentSiblingMap[firstChild] = {}
            self.parentSiblingMap[firstChild]["sibling"] = secondChild
            self.parentSiblingMap[firstChild]["parent"] = parentIndex

            self.parentSiblingMap[secondChild] = {}
            self.parentSiblingMap[secondChild]["sibling"] = firstChild
            self.parentSiblingMap[secondChild]["parent"] = parentIndex
            print("created map for parent %s first child %s second child %s"
                  % (parentIndex, firstChild, secondChild))

            self.parentSiblingMap[parentIndex]["firstChild"] = firstChild
            self.parentSiblingMap[parentIndex]["secondChild"] = secondChild

            self.createSiblingMap(firstChild)
            self.createSiblingMap(secondChild)

    def getMtIndexToNodeInfo(self):
        return self.parentSiblingMap

    def getLeafParentNodeIndexList(self):
        leafParentNodeIndexList = {}
        leafNodeIndexList = self.getLeafNodeIndexList()
        for leafIndex in leafNodeIndexList:
            parent = self.parentSiblingMap[leafIndex]["parent"]
            leafParentNodeIndexList[parent] = 1
        return leafParentNodeIndexList

    def getMtIndexToNameList(self):
        mtIndexToNames = {}
        for index in self.mt.get_vertices():
            name = self.mt.get_vertex_name(index)
            mtIndexToNames[index] = name
        return mtIndexToNames

    def getAssignmentContainerRoot(self, subsetIndex, read):
        outputDir = self.getParam("output_directory")
        filePrefix = self.getParam("subset_assignment_output_file")
        assignmentFileName = os.path.join(
            outputDir, "%s%s" %
            (filePrefix, subsetIndex))
        print("creating hdf5 file with name %s" % assignmentFileName)
        root = 0
        if (read == 1):
            root = RMF.open_hdf5_file(assignmentFileName)
        else:
            root = RMF.create_hdf5_file(assignmentFileName)
        return root

    #
    # Begin Cytoscape Methods
    #

    def writeVisualization(self):

        self.writeCytoscapeIgInput()
        self.writeCytoscapeJtInput()
        self.writeCytoscapeMtInput()
        self.writeCytoscapeScripts()
        self.writePymolData()

    def writeCytoscapeScripts(self):
        outputDir = self.getParam("output_directory")
        mTreeCytoscapeInput = self.getParam("mtree_cytoscape_input_file")
        mTreeCytoscapeAssignments = self.getParam(
            "mtree_cytoscape_assignment_file")
        mTreeCytoscapeAtomChains = self.getParam(
            "mtree_cytoscape_atom_chain_file")
        mTreeCytoscapeAtomSummary = self.getParam(
            "mtree_cytoscape_atom_summary_file")

        mTreeCytoscapeScript = self.getParam("mtree_cytoscape_script")
        mTreeFh = open(os.path.join(outputDir, mTreeCytoscapeScript), 'w')
        mTreeFh.write(
            "network import file=%s\n" %
            os.path.join(outputDir, mTreeCytoscapeInput))
        mTreeFh.write(
            "node import attributes file=\"%s\"\n" %
            os.path.join(outputDir, mTreeCytoscapeAssignments))
        mTreeFh.write(
            "node import attributes file=\"%s\"\n" %
            os.path.join(outputDir, mTreeCytoscapeAtomSummary))
        mTreeFh.write(
            "node import attributes file=\"%s\"\n" %
            os.path.join(outputDir, mTreeCytoscapeAtomChains))
        mTreeFh.write("layout jgraph-tree\n")

    def getGraphStructure(self, graph, fileName, separator):
        # write junction tree to file
        graphLogWrite = open(fileName, 'w')
        graph.show(graphLogWrite)
        graphLogWrite.close()

        # read file
        graphLogRead = open(fileName, 'r')

        nodeRe = re.compile(r'^(\d+)\[label\=\"\[*(.*?)\s*\]*\"')  # atom name
        separatorEscape = "\" + separator
        edgeString = r"^(\d+)\-%s(\d+)" % separatorEscape
        edgeRe = re.compile(edgeString)

        nodesToNodes = {}
        # keys: source node, value: target node (track edges)
        nodesToNames = {}  # keys: node number, value; string parsed from file

        for line in graphLogRead:

            # search for nodes
            node = nodeRe.search(line)
            if node:
                nodeNumber = node.group(1)
                atomString = node.group(2)
                nodesToNames[nodeNumber] = atomString
                continue

            # search for edges
            edge = edgeRe.search(line)
            if edge:
                firstNode = edge.group(1)
                secondNode = edge.group(2)
                firstNodeDict = {}
                if (firstNode in nodesToNodes):
                    firstNodeDict = nodesToNodes[firstNode]
                firstNodeDict[secondNode] = 1
                nodesToNodes[firstNode] = firstNodeDict

        return [nodesToNames, nodesToNodes]

    def writeEdgeFile(self, nodesToNodes, edgeFileName):
        # write edge file
        outputDir = self.getParam("output_directory")
        graphInputFile = open(os.path.join(outputDir, edgeFileName), 'w')
        for firstNode in nodesToNodes.keys():
            nodeDict = nodesToNodes[firstNode]
            for secondNode in nodeDict.keys():
                graphInputFile.write("%s ttt %s\n" % (firstNode, secondNode))
        graphInputFile.close()

    def writeCytoscapeIgInput(self):
        outputDir = self.getParam("output_directory")
        igOutputFile = self.getParam("ig_output_file")
        [nodesToNames, nodesToNodes] = self.getGraphStructure(
            self.ig, os.path.join(outputDir, igOutputFile), "-")

        self.writeEdgeFile(
            nodesToNodes,
            self.getParam("ig_cytoscape_input_file"))

        # write residue numbers for each node
        igResiduesFile = self.getParam("ig_cytoscape_residues_file")
        peptideChain = self.getParam("peptide_chain")

        subsetResidueFile = open(os.path.join(outputDir, igResiduesFile), 'w')
        subsetResidueFile.write("ResidueNumber\n")
        for nodeNumber in nodesToNames.keys():
            nodeName = nodesToNames[nodeNumber]

            [nodeChain, residueNumber,
                nodeAtom] = atomicDominoUtilities.getAtomInfoFromName(nodeName)
            if (nodeChain == peptideChain):  # peptideAtom
                subsetResidueFile.write(
                    "%s = %s\n" %
                    (nodeNumber, residueNumber))
            else:
                # for now just write arbitrary number to designate as protein
                # atom
                subsetResidueFile.write("%s = 100\n" % nodeNumber)

    def writeCytoscapeMtInput(self):
        outputDir = self.getParam("output_directory")
        mTreeOutputFile = self.getParam("mtree_output_file")
        [nodesToNames, nodesToNodes] = self.getGraphStructure(
            self.mt, os.path.join(outputDir, mTreeOutputFile), ">")

        self.writeEdgeFile(
            nodesToNodes,
            self.getParam("mtree_cytoscape_input_file"))
        self.writeNodeNameAttributes(
            nodesToNames, self.getParam("mtree_cytoscape_atom_name_file"),
            self.getParam("mtree_cytoscape_atom_summary_file"),
            self.getParam("mtree_cytoscape_atom_chain_file"))

    def writeCytoscapeJtInput(self):

        outputDir = self.getParam("output_directory")
        jTreeOutputFile = self.getParam("jtree_output_file")
        [nodesToNames, nodesToNodes] = self.getGraphStructure(
            self.jt, os.path.join(outputDir, jTreeOutputFile), "-")

        self.writeEdgeFile(
            nodesToNodes,
            self.getParam("jtree_cytoscape_input_file"))

        self.writeNodeNameAttributes(
            nodesToNames, self.getParam("jtree_cytoscape_atom_name_file"),
            self.getParam("jtree_cytoscape_atom_summary_file"),
            self.getParam("jtree_cytoscape_atom_chain_file"))

        # write edge weight file -- weights are number of shared particles
        # across nodes
        jtreeCytoscapeEdgeFile = self.getParam("jtree_cytoscape_edge_file")
        edgeWeightFile = open(
            os.path.join(outputDir,
                         jtreeCytoscapeEdgeFile),
            'w')
        edgeWeightFile.write("SubsetOverlap (class=Integer)\n")
        for firstNode in nodesToNodes.keys():
            nodeDict = nodesToNodes[firstNode]
            for secondNode in nodeDict.keys():
                firstNodeAtoms = nodesToNames[firstNode].split(" ")
                secondNodeAtoms = nodesToNames[secondNode].split(" ")
                intersection = [
                    val for val in firstNodeAtoms if val in secondNodeAtoms]
                edgeWeightFile.write(
                    "%s (pp) %s = %s\n" %
                    (firstNode, secondNode, len(intersection)))
        edgeWeightFile.close()

    def getAtomTypeCounts(self, atomNames):

        atomNames = atomNames.lstrip('[')
        atomNames = atomNames.rstrip(']')
        atomNames = atomNames.rstrip()

        atomList = atomNames.split(" ")  # atom names
        peptideAtomCount = 0
        proteinAtomCount = 0
        for atom in atomList:
            [chain, residue, atom] = atomicDominoUtilities.getAtomInfoFromName(
                atom)
            if (chain == self.getParam("peptide_chain")):
                peptideAtomCount += 1
            else:
                proteinAtomCount += 1
        return [peptideAtomCount, proteinAtomCount]

    def writeNodeNameAttributes(
        self,
        nodesToNames,
        atomNameFile,
        atomSummaryFile,
            atomChainFile):
        # write attribute file (atom names for each node)
        outputDir = self.getParam("output_directory")
        subsetAtomNameFile = open(os.path.join(outputDir, atomNameFile), 'w')
        subsetAtomSummaryFile = open(
            os.path.join(outputDir, atomSummaryFile), 'w')
        subsetAtomChainFile = open(os.path.join(outputDir, atomChainFile), 'w')

        subsetAtomNameFile.write("Atom names\n")
        subsetAtomSummaryFile.write("Atom Summary\n")
        subsetAtomChainFile.write("Atom chain\n")
        for node in nodesToNames.keys():
            atomNames = nodesToNames[node]
            subsetAtomNameFile.write("%s = %s\n" % (node, atomNames))
            [peptideAtomCount, proteinAtomCount] = self.getAtomTypeCounts(
                atomNames)
            # Number of protein and peptide atoms in each subset
            subsetAtomSummaryFile.write(
                "%s = %sp %sl\n" %
                (node, proteinAtomCount, peptideAtomCount))

            # whether each subset has protein, peptide, or a mix of atoms
            if (proteinAtomCount == 0):
                subsetAtomChainFile.write("%s = 1\n" % node)
            elif (peptideAtomCount == 0):
                subsetAtomChainFile.write("%s = 2\n" % node)
            else:
                subsetAtomChainFile.write("%s = 3\n" % node)

        subsetAtomChainFile.close()
        subsetAtomSummaryFile.close()
        subsetAtomNameFile.close()
    #
    # End Cytoscape Methods
    #