use insertion codes for water chain

modules/mol/alg/src/pdbize.cc

@@ -106,13 +106,20 @@ void PDBize::Add(EntityView asu, const geom::Mat4List& transforms,
       }
       if (chain.GetType() == CHAINTYPE_WATER) {
         if (!water_chain_) {
+          last_water_rnum_ = ResNum(0, 'Z');
           water_chain_ = edi.InsertChain(WATER_CHAIN_NAME);
           edi.SetChainDescription(water_chain_, chain.GetDescription());
           edi.SetChainType(water_chain_, chain.GetType());
         }
         for (ResidueViewList::const_iterator k = chain.GetResidueList().begin(),
              e3 = chain.GetResidueList().end(); k != e3; ++k) {
-          ResidueHandle new_res = edi.AppendResidue(water_chain_, k->GetName());
+          if (last_water_rnum_.GetInsCode() == 'Z') {
+            last_water_rnum_ = ResNum(last_water_rnum_.GetNum()+1, 'A');
+          } else {
+            last_water_rnum_.SetInsCode(last_water_rnum_.GetInsCode()+1);
+          }
+          ResidueHandle new_res = edi.AppendResidue(water_chain_, k->GetName(), 
+                                                    last_water_rnum_);
           transfer_residue_properties(*k, new_res);
           new_res.SetStringProp("type", StringFromChainType(chain.GetType()));
           new_res.SetStringProp("description", chain.GetDescription());

modules/mol/alg/src/pdbize.hh

@@ -52,6 +52,7 @@ private:
   const char*                           curr_chain_name_;
   bool                                  needs_adjustment_;
   int                                   last_rnum_;
+  ResNum                                last_water_rnum_;
   std::map<ResidueHandle,ResidueHandle> dst_to_src_map_;
 };
 

modules/mol/alg/tests/CMakeLists.txt

@@ -2,6 +2,7 @@ set(OST_MOL_ALG_UNIT_TESTS
   test_superposition.cc
   tests.cc
   test_consistency_checks.cc
+  test_pdbize.py
   test_convenient_superpose.py
 )
 

modules/mol/alg/tests/test_pdbize.py

+from ost import io, mol, geom, seq
+import unittest
+import os
+import random
+
+class TestPDBize(unittest.TestCase):
+  
+  def test_numbers_water_molecules_with_ins_codes(self):
+    m = mol.CreateEntity()
+    e = m.EditXCS(mol.BUFFERED_EDIT)
+    c = e.InsertChain("A");
+    e.SetChainType(c, mol.CHAINTYPE_WATER)
+    for i in range(27):
+      e.AppendResidue(c, "HOH")
+    pdbizer = mol.alg.PDBize()
+    transformations = geom.Mat4List()
+    transformations.append(geom.Mat4())
+    seqs = seq.CreateSequenceList()
+    pdbizer.Add(m.Select(''), transformations, seqs)
+    pdbized = pdbizer.Finish()
+    self.assertEquals([c.name for c in pdbized.chains], ["-"])
+    residues = pdbized.residues
+    for i in range(26):
+      self.assertEquals(residues[i].number.num, 1)
+      self.assertEquals(residues[i].number.ins_code, chr(ord('A')+i))
+    self.assertEquals(residues[26].number.num, 2)
+    self.assertEquals(residues[26].number.ins_code, 'A')
+  def test_starts_from_last_water_rnum(self):
+    m = mol.CreateEntity()
+    e = m.EditXCS(mol.BUFFERED_EDIT)
+    c = e.InsertChain("A");
+    e.SetChainType(c, mol.CHAINTYPE_WATER)
+    e.AppendResidue(c, "HOH")
+    pdbizer = mol.alg.PDBize()
+    transformations = geom.Mat4List()
+    transformations.append(geom.Mat4())
+    seqs = seq.CreateSequenceList()
+    pdbizer.Add(m.Select(''), transformations,seqs)
+    pdbizer.Add(m.Select(''), transformations,seqs)
+    pdbized = pdbizer.Finish()
+    self.assertEquals([c.name for c in pdbized.chains], ["-"])
+    residues = pdbized.residues
+    self.assertEquals([r.number for r in residues],
+                      [mol.ResNum(1, 'A'), mol.ResNum(1, 'B')])
+
+
+
+
+if __name__ == "__main__":
+  from ost import testutils
+  testutils.RunTests()