MOD: add unit_index in cgpdb.

cafepy/core/calc_com.py

@@ -51,37 +51,47 @@ def _read(self, inputfile):
     def readIndex(self):
         pass
 
-    def run(self, atom_indexes=[], unit_indexes=[], traj_indexes=[]):
+    def run(self, atom_idx=[], unit_idx=[], traj_idx=[]):
         """ 
         Calculates the Center of mass from DCD-file or PDB-file.
 
-        :Args: atom_indexes (list), traj_indexes (list)
-            :atom_indexes:    You can select Atom for calculating COM with indexes[.ndx,.ninfo]-file 
-            :traj_indexes:    You can extract trajectories for calculating COM.
+        :Args: atom_idx (list), traj_idx (list)
+
+            :atom_idx:    You can select Atoms for calculating COM with idx[.ndx,.ninfo]-file 
+            :unit_idx:    You can select Units in pdb format.
+            :traj_idx:    You can extract trajectories for calculating COM.
 
         """
+
         if self.sfx == 'dcd':
-            return self._dcdrun(atom_indexes, unit_indexes, traj_indexes)
+            return self._dcdrun(atom_idx, unit_idx, traj_idx)
         elif self.sfx == 'pdb':
-            return self._pdbrun(atom_indexes, unit_indexes)
+            return self._pdbrun(atom_idx, unit_idx)
 
-    def _dcdrun(self, atom_indexes=[], unit_indexes=[], traj_indexes=[]):
-        if not atom_indexes:
+    def _dcdrun(self, atom_idx=[], unit_idx=[], traj_idx=[]):
+        if not atom_idx:
             self.com = np.average(self.data[:], axis=1)
         else:
             ndata = np.array(self.data)
-            self.com = np.average(ndata[:, atom_indexes], axis=1)
+            self.com = np.average(ndata[:, atom_idx], axis=1)
         return self.com
 
-    def _pdbrun(self, atom_indexes=[], unit_index=[]):
-        if not atom_indexes:
+    def _pdbrun(self, atom_idx=[], unit_idx=[]):
+        if unit_idx:
+            _data = []
+            for i in unit_idx:
+                i = i-1
+                if i == 0:
+                    s = 0
+                else:
+                    s = sum(self.data.info['aasize'][:i])
+                e = sum(self.data.info['aasize'][:i+1])
+                _data += self.data[:][s: e]
+            ndata = np.array(_data)
+            self.com = np.average(ndata, axis=0)
+        elif not atom_idx:
             self.com = np.average(self.data[:], axis=0)
-        else:
-            ndata = np.array(self.data)
-            self.com = np.average(ndata[:, atom_indexes], axis=1)
         return self.com
-
-
 
     def writeFile(self, outputfile, header= ""):
         np.savetxt(outputfile, self.com, header=header, fmt="%.8e")

cafepy/files/pdbfile.py

@@ -22,8 +22,43 @@
 from ..utils.cafepy_error import CafePyError
 from ..utils.cafepy_base import CafePyBase
 from ..utils.cafepy_math import rotation3D
+
+class BasePDB(CafePyBase, FileIO):
+    def __init__(self, filename):
+        FileIO.__init__(self)
+        self.ftype = 'pdb'
+        self.coard = []
+        self.data = []
+        self.filename = filename
+
+
+    def __getitem__(self,key):
+        if isinstance(key, slice):
+            return [self[i] for i in range(*key.indices(len(self)))]
+        elif isinstance(key, int):
+            if key < 0:
+                key += len(self)
+            if key < 0 or key >= len(self):
+                raise IndexError("The index (%d) is out of range." % key)
+            return self.data[key][5:8]
+        else:
+            raise TypeError("Invalid argument type.")
+
+    def __len__(self):
+        return len(self.data)
+
+    def _read(self):
+        self._file = self.openFile(self.filename)
+        self.readATOM()
+
+    def close(self):
+        self._file.close()
+        return True
 
-class PDB(CafePyBase, FileIO):
+    def write(self, xyz=[], inpsffile=""):
+        pass
+
+class PDB(BasePDB):
     """
     Reading a PDB(Protein Data Bank) file which is an output from CafeMol Software.
     This class returns coordinates as list format. Please check examples below.
@@ -41,20 +76,16 @@ class PDB(CafePyBase, FileIO):
         pdb.close()
     
     """
-    def __init__(self, filename):
-        FileIO.__init__(self)
-        self.filename = filename
-        self.ftype = 'pdb'
-        self.coard = []
-        self.row_data = []
+    def __init__(self, filename, *args, **kwargs):
+        super(PDB, self).__init__(filename, *args, **kwargs)
         self._read()
 
     def get(self):
         """
         :input:   None
         :return:  pdb data as list format
         """
-        return self.row_data()
+        pass
 
     def readATOM(self):
         """
@@ -77,58 +108,55 @@ def readATOM(self):
 
         for line in self._file.readlines():
             if re.match(r"(ATOM)", line):
-                self.row_data.append([slc[2](line[slice(*slc[:2])]) for slc in _format])
-        return self.row_data
+                self.data.append([slc[2](line[slice(*slc[:2])]) for slc in _format])
+        return self.data
 
-    def __getitem__(self,key):
-        if isinstance(key, slice):
-            return [self[i] for i in range(*key.indices(len(self)))]
-        elif isinstance(key, int):
-            if key < 0:
-                key += len(self)
-            if key < 0 or key >= len(self):
-                raise IndexError("The index (%d) is out of range." % key)
-            return self.row_data[key][5:8]
-        else:
-            raise TypeError("Invalid argument type.")
-
-    def __len__(self):
-        return len(self.row_data)
-
-    def _read(self):
-        self._file = self.openFile(self.filename)
-        self.readATOM()
-
-    def close(self):
-        self._file.close()
-        return True
-
-    def write(self, xyz=[], inpsffile=""):
-        pass
 
 
-class CGPDB(CafePyBase, FileIO):
+class CGPDB(BasePDB):
     """
     Reading a Coarse-Grained PDB(Protein Data Bank) file which is an output from CafeMol Software.
     """
-    def __init__(self, filename):
+    def __init__(self, filename, *args, **kwargs):
+        super(CGPDB, self).__init__(filename, *args, **kwargs)        
         self.filename = filename
-        self.ATOM = []
+        self.data = []
         self.coords = []
+        self.info = {}
         self._read()
 
     def _read(self):
         self._file = self.openFile(self.filename)
-        self.readATOM()
+        self._readATOM()
 
-    def readATOM(self):
+    def _readATOM(self):
         _format = [(6,11,int),(12,16,str),(17,20,str),                  #serial,atomName,resName
                    (21,22,str),(22,26,int),                             #chainID,resSeq
                    (30,38,float),(38,46,float),(46,54,float)]
+        usize = 0
+        aasize = []
+        each_aasize = 0
         for line in self._file.readlines():
             if re.match(r"(ATOM)", line):
-                self.ATOM.append([slc[2](line[slice(*slc[:2])]) for slc in _format])
-        return self.ATOM
+                self.data.append([slc[2](line[slice(*slc[:2])]) for slc in _format])
+                each_aasize += 1
+            if re.match(r'TER', line):
+                aasize.append(each_aasize)
+                usize += 1
+                each_aasize = 0
+        ## the number of amino acid                
+        self.info['aasize'] = aasize
+
+        ## the number of amino unit
+        if usize == 0:
+            self.info['usize'] = 1
+        else:
+            self.info['usize'] = usize
+        return self.data
+
+    def getUnitSize(self):
+        pass
+
 
     def rotation(self, alpha, beta, gamma):
         self.coords = rotation3D(self.coords, alpha, beta, gamma)

test/test_calc_com.py

@@ -32,7 +32,6 @@ def test_readPDB(self):
         for x, a in zip(com.com.tolist(), ans):
             self.assertAlmostEqual(x, a)        
 
-
     def test_readBigPDB(self):
         com = CalcCOM(self.bpdb)
 
@@ -60,12 +59,15 @@ def test_calcCOMfromDCD_with_range(self):
         self.assertEqual(com.com.tolist(), self.com_5_atoms)
         com.close()
 
-
     def test_calcCOMfromDCD_with_index_and_slices(self):
         pass
 
-    def test_main(self):
-        pass
+    def test_calc_com_from_cgpdb_with_unitindexes(self):
+        com = CalcCOM(self.bpdb)
+        com.run(unit_idx=[1])
+        ans = [-42.01642589, 15.75263203, 95.42158773]
+        for x, a in zip(com.com.tolist(), ans):
+            self.assertAlmostEqual(x, a)        
 
     def tearDown(self):
         pass

test/test_pdbfile.py

@@ -7,20 +7,19 @@
 
 from unittest.mock import patch
 
-cafepypath = os.path.join(os.path.dirname(__file__), '..')
-sys.path.append(cafepypath)
+from .test_base import CafePyTestBase
+
 from cafepy.files import PDB, CGPDB
 
-class TestPDB(unittest.TestCase):
+class TestPDB(CafePyTestBase):
     @classmethod
     def setUpClass(cls):
-        cls.testdata = os.path.join(os.path.dirname(__file__), 'data/test.pdb')
-        cls.listdata = os.path.join(os.path.dirname(__file__), 'data/')
-        sys.path.append(cls.listdata)
+        pass
+
+    def setUp(self):
         from test_data_in_pdbfile import list_pdb_data, list_pdb_coord
-        cls.pdblist = list_pdb_coord
-
-        cls.pdb = PDB(cls.testdata)
+        self.pdblist = list_pdb_coord
+        self.pdb = PDB(self.data_path+'test.pdb')
 
     def test_readPDB_with_index(self):
         self.assertEqual([15.939, 14.6, -4.238], self.pdb[0])
@@ -49,21 +48,33 @@ def tearDown(self):
     def tearDownClass(cls):
         pass
 
-class TestCGPDB(unittest.TestCase):
+class TestCGPDB(CafePyTestBase):
     @classmethod
     def setUpClass(cls):
-        cls.testdata = os.path.join(os.path.dirname(__file__), 'data/cgtest.pdb')
-        cls.listdata = os.path.join(os.path.dirname(__file__), 'data/')
-        sys.path.append(cls.listdata)
-        cls.cgpdb = CGPDB(cls.testdata)
-        from test_data_in_pdbfile import list_cgpdb_data
-        cls.cgpdb_list = list_cgpdb_data
+        pass
+
+    def setUp(self):
+        self.cgpdb = CGPDB(self.data_path+'cgtest.pdb')
+        self.bigpdb = CGPDB(self.data_path+'big_test.pdb')
+
+    def test_unit_size_from_bigpdb(self):
+        self.assertEqual(self.bigpdb.info['usize'], 60)
+
+    def test_unit_size_from_cgpdb(self):
+        self.assertEqual(self.cgpdb.info['usize'], 1)
 
+    def test_amino_acid_size_from_bigpdb(self):
+        self.assertEqual(self.cgpdb.info['aasize'], 56)
+
+    def test_amino_acid_size_from_bigpdb(self):
+        self.assertEqual(self.bigpdb.info['aasize'], [587 for i in range(60)])
+
     def tearDown(self):
         pass
 
     def test_read_cgpdb(self):
-        self.assertEqual(self.cgpdb.ATOM, self.cgpdb_list)
+        pass
+
 
     @classmethod
     def tearDownClass(cls):