add full_ter kwarg

jamesmkrieger · Nov 23, 2023 · 64f8411 · 64f8411
1 parent 30052ce
commit 64f8411
Showing 1 changed file with 30 additions and 19 deletions.
diff --git a/prody/proteins/pdbfile.py b/prody/proteins/pdbfile.py
@@ -1242,8 +1242,13 @@ def writePDBStream(stream, atoms, csets=None, **kwargs):
         Default is **False**, which means using hexadecimal instead.
         NB: ChimeraX seems to prefer hybrid36 and may have problems with hexadecimal.
     :type hybrid36: bool
+
+    :arg full_ter: whether to write full TER lines with atoms info
+        Default is **True**
+    :type full_ter: bool    
     """    
     renumber = kwargs.get('renumber', True)
+    full_ter = kwargs.get('full_ter', True)
 
     remark = str(atoms)
     try:
@@ -1400,9 +1405,9 @@ def writePDBStream(stream, atoms, csets=None, **kwargs):
                 L = helix_resnums[-1] - helix_resnums[0] + 1
 
                 stream.write(HELIXLINE.format(serNum=i, helixID=helix_secids[0], 
-                            initResName=helix_resnames[0][:3], initChainID=helix_chainids[0][:1], 
+                            initResName=helix_resnames[0][:3], initChainID=helix_chainids[0], 
                             initSeqNum=helix_resnums[0], initICode=helix_icodes[0],
-                            endResName=helix_resnames[-1][:3], endChainID=helix_chainids[-1][:1], 
+                            endResName=helix_resnames[-1][:3], endChainID=helix_chainids[-1], 
                             endSeqNum=helix_resnums[-1], endICode=helix_icodes[-1],
                             helixClass=helix_secclasses[0], length=L))
 
@@ -1426,9 +1431,9 @@ def writePDBStream(stream, atoms, csets=None, **kwargs):
                 strand_icodes = icodes[torf_strand]
 
                 stream.write(SHEETLINE.format(strand=i, sheetID=sheet_id, numStrands=numStrands,
-                            initResName=strand_resnames[0][:3], initChainID=strand_chainids[0][:1], 
+                            initResName=strand_resnames[0][:3], initChainID=strand_chainids[0], 
                             initSeqNum=strand_resnums[0], initICode=strand_icodes[0],
-                            endResName=strand_resnames[-1][:3], endChainID=strand_chainids[-1][:1], 
+                            endResName=strand_resnames[-1][:3], endChainID=strand_chainids[-1], 
                             endSeqNum=strand_resnums[-1], endICode=strand_icodes[-1],
                             sense=strand_secclasses[0]))
         pass
@@ -1571,7 +1576,7 @@ def writePDBStream(stream, atoms, csets=None, **kwargs):
 
                 write(anisouline % ("ANISOU", serial,
                                     atomnames[i], altlocs[i],
-                                    resname, chainids[i][:1], resnum,
+                                    resname, chainids[i], resnum,
                                     icodes[i],
                                     anisou[0], anisou[1], anisou[2],
                                     anisou[3], anisou[4], anisou[5],
@@ -1584,13 +1589,16 @@ def writePDBStream(stream, atoms, csets=None, **kwargs):
                     else:
                         serial += 1
 
-                    false_pdbline = pdbline % ("TER   ", serial,
-                                               "", "",
-                                               resname, chainids[i][:1], resnum,
-                                               icodes[i],
-                                               xyz[0], xyz[1], xyz[2],
-                                               occupancies[i], bfactors[i],
-                                               segments[i], elements[i], charges2[i])
+                    if full_ter:
+                        false_pdbline = pdbline % ("TER   ", serial,
+                                                   "", "",
+                                                   resname, chainids[i], resnum,
+                                                   icodes[i],
+                                                   xyz[0], xyz[1], xyz[2],
+                                                   occupancies[i], bfactors[i],
+                                                   segments[i], elements[i], charges2[i])
+                    else:
+                        false_pdbline = "TER" + " "*23
                     write(false_pdbline[:26] + " "*54 + '\n')
                     num_ter_lines += 1
             else:
@@ -1600,13 +1608,16 @@ def writePDBStream(stream, atoms, csets=None, **kwargs):
                     else:
                         serial += 1
 
-                    false_pdbline = pdbline % ("TER   ", serial,
-                                               "", "",
-                                               resname, chainids[i][:1], resnum,
-                                               icodes[i],
-                                               xyz[0], xyz[1], xyz[2],
-                                               occupancies[i], bfactors[i],
-                                               segments[i], elements[i], charges2[i])
+                    if full_ter:
+                        false_pdbline = pdbline % ("TER   ", serial,
+                                                   "", "",
+                                                   resname, chainids[i], resnum,
+                                                   icodes[i],
+                                                   xyz[0], xyz[1], xyz[2],
+                                                   occupancies[i], bfactors[i],
+                                                   segments[i], elements[i], charges2[i])
+                    else:
+                        false_pdbline = "TER" + " "*23
                     write(false_pdbline[:26] + " "*54 + '\n')
                     num_ter_lines += 1