protect agains residue count (#17)

Co-authored-by: pegerto-fernandez-ck <pegerto.fernandez@creditkarma.com>
pegerto · Nov 22, 2023 · 54c53ec · 54c53ec
1 parent c18431c
commit 54c53ec
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 7 deletions.
diff --git a/mdakit_sasa/analysis/sasaanalysis.py b/mdakit_sasa/analysis/sasaanalysis.py
@@ -14,10 +14,12 @@
 import numpy as np
 import freesasa
 import os
+import logging
 
 if TYPE_CHECKING:
     from MDAnalysis.core.universe import Universe, AtomGroup
 
+logger = logging.getLogger(__name__)
 
 class SASAAnalysis(AnalysisBase):
     """SASAAnalysis class.
@@ -76,7 +78,7 @@ def _prepare(self):
             dtype=float,
         )
         self.results.residue_area = np.zeros(
-            (self.n_frames, len(self.universe.residues)),
+            (self.n_frames, len(self.universe.residues.resids)),
             dtype=float,
         )
 
@@ -87,19 +89,25 @@ def _single_frame(self):
         structure = freesasa.Structure()  
         # FreeSasa structure accepts PDBS if not available requires to reconstruct the structure using `addAtom`
         for a in self.atomgroup:
-            x,y,z = a.position            
-            structure.addAtom(a.name, a.resname, a.resnum.item(), a.segid, x, y, z)
+            x,y,z = a.position
+            structure.addAtom(a.type.rjust(2), a.resname, a.resnum.item(), a.segid, x, y, z)
 
         # Define 1 cpu for windows avoid freesasa code to calculate it.
         parametes =  freesasa.Parameters()
         if self._is_windows():
             parametes.setNThreads(1)
 
         result = freesasa.calc(structure, parametes)
+
         residue_areas = [result.residueAreas()[s][r] for s in sorted(list(result.residueAreas().keys())) for r in sorted(list(result.residueAreas()[s].keys()))]
 
         self.results.total_area[self._frame_index] = result.totalArea()
-        self.results.residue_area[self._frame_index] = [r.total for r in residue_areas]
+
+        # Defend agains residue counts mismatch
+        if  len(self.universe.residues.resids)!= len(residue_areas):
+            logger.error(f'Residude count do not match the expectation, residue SASA not in results { len(self.universe.residues.resids)} != {len(residue_areas)}')
+        else:
+            self.results.residue_area[self._frame_index] = [r.total for r in residue_areas]
 
     def _conclude(self):
         self.results.mean_total_area= self.results.total_area.mean()

diff --git a/mdakit_sasa/tests/analysis/test_sasaanalysis.py b/mdakit_sasa/tests/analysis/test_sasaanalysis.py
@@ -5,7 +5,7 @@
 
 from mdakit_sasa.analysis.sasaanalysis import SASAAnalysis
 from mdakit_sasa.tests.utils import make_Universe
-from MDAnalysis.core.topologyattrs import Atomnames, Resnames, Resids, Resnums, Segids
+from MDAnalysis.core.topologyattrs import Atomnames, Resnames, Resids, Resnums, Segids, Atomtypes
 
 class TestSASAAnalysis:
 
@@ -25,6 +25,7 @@ def universe(self):
         u.add_TopologyAttr(Resids(list(range(0, len(u.residues)))))
         u.add_TopologyAttr(Resnums(list(range(0, len(u.residues)))))
         u.add_TopologyAttr(Segids(["A"] * len(u.segments)))
+        u.add_TopologyAttr(Atomtypes(["O"]* len(u.atoms)))
         return u
 
     @pytest.fixture
@@ -60,5 +61,4 @@ def test_residue_sasa_calculation_results(self, analysis):
         assert analysis.n_frames == 3
         assert analysis.results['residue_area'].dtype ==  np.dtype('float64')
         assert np.all(analysis.results['residue_area'] >= 0)
-        assert analysis.results['residue_area'].shape == (3,25)
-        0
+        assert analysis.results['residue_area'].shape == (3,25)