diff spring_package/Molecule.py @ 39:172398348efd draft

"planemo upload commit 26b4018c88041ee0ca7c2976e0a012015173d7b6-dirty"

--- a/spring_package/Molecule.py	Wed Nov 25 17:38:24 2020 +0000
+++ b/spring_package/Molecule.py	Fri Jan 22 15:50:27 2021 +0000
@@ -86,26 +86,31 @@
         matLine = list(map(lambda x: self.toFloat(x), matLine))
         return matId, matLine
 
-    def createUnit(self, biomolNumber=1):
-        molecule = Molecule()
-        chainCount = 0
-        for matrixDict in self.biomol[biomolNumber]:
-            for chain in matrixDict["chains"]:
-                if chain in self.calpha:
-                    chainCopy = dict()
-                    for residue in self.calpha[chain]:
-                        chainCopy[residue] = self.calpha[chain][residue].copy()
-                    for atomNumber in chainCopy:
-                        atom = chainCopy[atomNumber]
-                        rotmat = matrixDict["matrix"]
-                        self.applyMatrix(atom, rotmat)
-                    if chain in molecule.calpha:
-                        chainName = "%s%d" % (chain, chainCount)
-                    else:
-                        chainName = chain
-                    molecule.calpha[chainName] = chainCopy
-                    chainCount = chainCount + 1
-        return molecule
+    def createUnit(self, biomolNumber=0):
+        if biomolNumber == 0:
+            return self
+        else:
+            molecule = Molecule()
+            chainCount = dict()
+            for matrixDict in self.biomol[biomolNumber]:
+                for chain in matrixDict["chains"]:
+                    if chain in self.calpha:
+                        chainCopy = dict()
+                        for residue in self.calpha[chain]:
+                            chainCopy[residue] = self.calpha[chain][residue].copy()
+                        for atomNumber in chainCopy:
+                            atom = chainCopy[atomNumber]
+                            rotmat = matrixDict["matrix"]
+                            self.applyMatrix(atom, rotmat)
+                        if chain in chainCount:
+                            chainCount = chainCount[chain]
+                            chainName = "%s_%d" % (chain, chainCount)
+                            chainCount[chain] = chainCount + 1
+                        else:
+                            chainName = chain
+                            chainCount[chain] = 0
+                        molecule.calpha[chainName] = chainCopy
+            return molecule
 
     def getSequence(self, chainName):
         seq = ""
@@ -145,7 +150,6 @@
         return code[seq] if seq in code else "X"
 
     def saveChain(self, chainName, outputName):
-        print("Writing PDB file to %s." % outputName)
         f = open(outputName, "w")
         for residueNumber in sorted(self.calpha[chainName].keys()):  
             ca = self.calpha[chainName][residueNumber]
@@ -153,10 +157,11 @@
                 f.write(self.atomString(ca))
         f.close()
 
-    def save(self, outputName, append=False, chainName=None):
-        print("Writing atoms to PDB file to %s." % outputName)
+    def save(self, outputName, append=False, chainName=None, payload=None):
         fileFlag = "+a" if append else "w"
         f = open(outputName, fileFlag)
+        if payload:
+            f.write("%s\n" % payload)
         for atom in self.atoms:
             atom["chainName"] = chainName if chainName else atom["chainName"]
             f.write(self.atomString(atom))