Refactor write_tcl_script

frustratometer/classes/Frustratometer.py

@@ -389,8 +389,9 @@ def auc(self):
         """
         return frustration.compute_auc(self.roc())
 
-    def vmd(self, sequence: str = None, single:Union[str,np.array] = 'singleresidue', pair:Union[str,np.array] = 'mutational',
-             aa_freq:np.array = None, correction:int = 0, max_connections:Union[int,None] = None, movie_name=None, still_image_name=None):
+    def vmd(self, sequence: str = None, single:Union[str,np.array,None] = 'singleresidue', pair:Union[str,np.array] = 'mutational',
+             aa_freq:np.array = None, correction:int = 0, max_connections:Union[int,None] = None, movie_name=None, still_image_name=None,
+             call_vmd: bool = True):
         """
         Calculates frustration indices and superimposes frustration patterns onto PDB structure using the VMD software.
 
@@ -415,12 +416,19 @@ def vmd(self, sequence: str = None, single:Union[str,np.array] = 'singleresidue'
                     from the sequence that was passed to this vmd function. Proceeding further may not\n\
                     perform the computation that you intend to perform.")
 
+        if isinstance(single, str):
+            single = -self.frustration(kind=single, sequence=sequence, aa_freq=aa_freq)
 
-        tcl_script = frustration.write_tcl_script(self.pdb_file, self.chain, self.mask, self.distance_matrix, self.distance_cutoff,
-                                      -self.frustration(kind=single, sequence=sequence, aa_freq=aa_freq),
-                                      -self.frustration(kind=pair, sequence=sequence, aa_freq=aa_freq),
-                                      max_connections=max_connections, movie_name=movie_name, still_image_name=still_image_name)
-        frustration.call_vmd(self.pdb_file, tcl_script)
+        if isinstance(pair, str):
+            pair = -self.frustration(kind=pair, sequence=sequence, aa_freq=aa_freq)
+
+        tcl_script = frustration.write_tcl_script(
+            self.pdb_file, self.chain, self.mask, self.distance_matrix, self.distance_cutoff,
+            single, pair,
+            max_connections=max_connections, movie_name=movie_name, still_image_name=still_image_name)
+
+        if call_vmd:
+            frustration.call_vmd(self.pdb_file, tcl_script)
 
     def view_pair_frustration(self, sequence:str = None, pair:str = 'mutational', aa_freq:np.array = None):
         """

frustratometer/classes/Structure.py

@@ -101,6 +101,65 @@ def __repr__(self):
         kind = 'mask' if self.data is None else 'distance'
         return f"SparseMatrix({kind}, nnz={len(self)}, L={self.shape})"
 
+    def __iter__(self):
+        # check that our attributes all make sense
+        assert self.row.shape == self.col.shape == (self.row.shape[0],) # tuple of length 1
+        if isinstance(self.data, np.ndarray):
+            assert self.data.shape == self.row.shape
+        elif not (self.data is None):
+            raise AssertionError(f'self.data was {self.data}')
+        N = self.row.shape[0]
+        n = 0
+        while n < N:
+            yield (self.row[n], self.col[n], None if self.data is None else self.data[n]) 
+            n += 1
+
+    def __and__(self, other):
+        # create a new SparseMatrix that acts as a boolean mask
+        # for the indices present in self and other
+        if not isinstance(other, SparseMatrix):
+            return NotImplemented
+        else:
+            if self.shape != other.shape:
+                raise ValueError("intersection of two boolean arrays of different shapes is ill-defined")
+            new_row, new_col = self._and_helper(other)
+            return SparseMatrix(new_row, new_col)
+
+    def __rand__(self, other):
+        if not isinstance(other, SparseMatrix):
+            return NotImplemented
+        else:
+            return self.__and__(other)
+
+    def __iand__(self, other):
+        # in-place modification of self
+        if not isinstance(other, SparseMatrix):
+            return NotImplemented
+        if self.shape != other.shape:
+            raise ValueError("intersection of two boolean arrays of different shapes is ill-defined")
+        self.row, self.col = self._and_helper(other) # modify self.row and self.col in place
+        self.data = None
+        return self
+
+    def _and_helper(self, other):
+        #new_row = []
+        #new_col = []
+        #other_set = set(zip(other.row, other.col))
+        #for i, j in zip(self.row, self.col):
+        #    if (i,j) in other_set:
+        #        new_row.append(i)
+        #        new_col.append(j)
+        #return np.array(new_row), np.array(new_col)
+
+        # Vectorized membership using 64-bit integer keys; preserves order of self.
+        # Similar speed to the above option, may be more memory-efficient
+        # if numpy significantly outperforms native python
+        L = self.shape 
+        keys_self = (self.row.astype(np.int64) * L) + self.col.astype(np.int64)
+        keys_other = (other.row.astype(np.int64) * L) + other.col.astype(np.int64)
+        mask = np.in1d(keys_self, keys_other, assume_unique=False)
+        return self.row[mask].astype(np.intp), self.col[mask].astype(np.intp)
+
 class Structure:
 
     def __init__(self, pdb_file: Union[Path,str], chain: Union[str,None]=None, seq_selection: str = None, aligned_sequence: str = None, filtered_aligned_sequence: str = None,