diff --git a/frustratometer/classes/Frustratometer.py b/frustratometer/classes/Frustratometer.py index 11c22c5..35ac5c1 100644 --- a/frustratometer/classes/Frustratometer.py +++ b/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): """ diff --git a/frustratometer/classes/Structure.py b/frustratometer/classes/Structure.py index 64bd390..ad83746 100644 --- a/frustratometer/classes/Structure.py +++ b/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, diff --git a/frustratometer/frustration/frustration.py b/frustratometer/frustration/frustration.py index 08fd1d8..92a8ffe 100644 --- a/frustratometer/frustration/frustration.py +++ b/frustratometer/frustration/frustration.py @@ -890,10 +890,19 @@ def plot_singleresidue_decoy_energy(decoy_energy, native_energy, method='cluster return g -def write_tcl_script(pdb_file: Union[Path,str], chain: str, mask: np.array, distance_matrix: np.array, distance_cutoff: float, single_frustration: np.array, - pair_frustration: np.array, tcl_script: Union[Path, str] ='frustration.tcl',max_connections: int =None, movie_name: Union[Path, str] =None, still_image_name: Union[Path, str] =None) -> Union[Path, str]: +def write_tcl_script( + pdb_file: Union[Path,str], chain: str, mask: np.array, + distance_matrix: np.array, distance_cutoff: float, + single_frustration: Union[np.array,None], pair_frustration: np.array, + tcl_script: Union[Path, str] ='frustration.tcl',max_connections: int =None, + movie_name: Union[Path, str] =None, still_image_name: Union[Path, str] =None, + ) -> Union[Path, str]: """ Writes a tcl script that can be run with VMD to superimpose the frustration patterns onto the corresponding PDB structure. + + This function is fastest when distance_matrix and mask inputs are dense. + Converting the matrices from sparse to dense or alternatively following the sparse matrix- + based logic adds about 0.4 seconds to this function call for NFKB (615 residues). Parameters ---------- @@ -902,9 +911,9 @@ def write_tcl_script(pdb_file: Union[Path,str], chain: str, mask: np.array, dist chain : str Select chain from pdb mask : np.array - A 2D Boolean array that determines which residue pairs should be considered in the energy computation. The mask should have dimensions (L, L), where L is the length of the sequence. + SparseMatrix or 2D Boolean array that determines which residue pairs should be considered in the energy computation. The mask should have dimensions (L, L), where L is the length of the sequence. distance_matrix : np.array - LxL array for sequence of length L, describing distances between contacts + SparseMatrix or LxL array for sequence of length L, describing distances between contacts distance_cutoff : float Maximum distance at which a contact occurs single_frustration : np.array @@ -926,11 +935,10 @@ def write_tcl_script(pdb_file: Union[Path,str], chain: str, mask: np.array, dist tcl_script : Path or str tcl script file """ - fo = open(tcl_script, 'w+') + to_write = [] single_frustration = np.nan_to_num(single_frustration,nan=0,posinf=0,neginf=0) pair_frustration = np.nan_to_num(pair_frustration,nan=0,posinf=0,neginf=0) - structure = prody.parsePDB(str(pdb_file)) selection = structure.select('protein', chain=chain) residues = np.unique(selection.getResindices()) @@ -942,24 +950,77 @@ def write_tcl_script(pdb_file: Union[Path,str], chain: str, mask: np.array, dist cb_sel = selection.select('name CB or (resname GLY and name CA)') ca_xyz = np.full((len(residues), 3), np.nan) cb_xyz = np.full((len(residues), 3), np.nan) + ca_indices = np.full(len(residues), -1, dtype=int) ca_xyz[np.searchsorted(residues, ca_sel.getResindices())] = ca_sel.getCoords() cb_xyz[np.searchsorted(residues, cb_sel.getResindices())] = cb_sel.getCoords() + ca_indices[np.searchsorted(residues, ca_sel.getResindices())] = ca_sel.getIndices() - fo.write(f'[atomselect top all] set beta 0\n') + to_write.append(f'[atomselect top all] set beta 0\n') # Single residue frustration - for r, f in zip(residues, single_frustration): - fo.write(f'[atomselect top "residue {int(r)}"] set beta {f}\n') - - ii, jj = np.triu_indices(len(residues), k=1) - keep = mask[ii, jj] > 0 - if distance_cutoff: - keep &= distance_matrix[ii, jj] <= distance_cutoff - ii, jj = ii[keep], jj[keep] - frust = pair_frustration[ii, jj] - dist = distance_matrix[ii, jj] - - # Green = minimally frustrated (most negative first); red = frustrated (most positive - # first). For each: sort, cap at max_connections, then draw the in-range contacts. + if not (single_frustration is None): + for r, f in zip(residues, single_frustration): + to_write.append(f'[atomselect top "residue {int(r)}"] set beta {f}\n') + + # if we're really worried about memory, + # we should loop over elements of mask + # and distance_matrix and check their indices + # and values, but distance_matrix shouldn't + # be very big (10^4 residues -> 10^8 + # elements in distance matrix -> 800 MB, + # assuming 8-byte floats) + try: + mask = mask.to_dense(fill=False) + except AttributeError: + pass # it's already a numpy array + try: + distance_matrix = distance_matrix.to_dense(fill=np.inf) + except AttributeError: + pass # it's already a numpy array + def process_dense_arrays(): + ii, jj = np.triu_indices(len(residues), k=1) + keep = mask[ii, jj] > 0 + if distance_cutoff: + keep &= distance_matrix[ii, jj] <= distance_cutoff + ii, jj = ii[keep], jj[keep] + frust = pair_frustration[ii, jj] + dist = distance_matrix[ii, jj] + return frust, dist, ii, jj + if isinstance(distance_matrix,np.ndarray) and isinstance(mask,np.ndarray): + frust, dist, ii, jj = process_dense_arrays() + elif isinstance(distance_matrix,np.ndarray) and not isinstance(mask,np.ndarray): + mask = mask.to_dense(fill=False) + frust, dist, ii, jj = process_dense_arrays() + elif (not isinstance(distance_matrix,np.ndarray)) and isinstance(mask,np.ndarray): + distance_matrix = distance_matrix.to_dense(fill=np.inf) + frust, dist, ii, jj = process_dense_arrays() + else: # This is about the same speed as converting to dense matrix for NFKB (615 residues) + # note that this code is currently unreachable + if distance_cutoff: + mask &= distance_matrix.compute_mask(maximum_contact_distance=distance_cutoff) + upper_triangular_part = np.unique(np.sort(np.array([mask.row, mask.col]),axis=0),axis=1) + ii = upper_triangular_part[0,:] + jj = upper_triangular_part[1,:] + frust = pair_frustration[ii,jj] # pair_frustration is always dense + dist = distance_matrix.lookup(ii,jj) + + # --- START OF HIGH-PERFORMANCE REDRAW Tcl PROCEDURE --- + to_write.append('\nproc redraw_bound_line {} {\n') + to_write.append(' set molid top\n') + to_write.append(' graphics $molid delete all\n\n') + + # 1. Grab all CA atoms AT ONCE + to_write.append(' set sel [atomselect $molid "name CA"]\n') + to_write.append(' set atom_data [$sel get {index x y z}]\n') + to_write.append(' $sel delete\n\n') + + # 2. Build a Tcl array mapping index to {x y z} for O(1) lookups + to_write.append(' array unset coords\n') + to_write.append(' foreach record $atom_data {\n') + to_write.append(' lassign $record idx x y z\n') + to_write.append(' set coords($idx) [list $x $y $z]\n') + to_write.append(' }\n\n') + + # Green = minimally frustrated; red = frustrated. Sort, cap, draw. for color, group, reverse in ( ('green', frust < -0.78, False), ('red', frust > 1, True), @@ -973,90 +1034,92 @@ def write_tcl_script(pdb_file: Union[Path,str], chain: str, mask: np.array, dist gi, gj, gd = gi[:max_connections], gj[:max_connections], gd[:max_connections] draw = (gd >= 3.5) & (gd <= 9.5) - draw &= np.abs(residues[gi] - residues[gj]) != 1 # Skip adjacent residues, which are often connected by a covalent bond and thus always close in space, but not necessarily frustrated. + draw &= np.abs(residues[gi] - residues[gj]) != 1 draw &= ~(np.isnan(ca_xyz[gi]).any(1) | np.isnan(ca_xyz[gj]).any(1) | np.isnan(cb_xyz[gi]).any(1) | np.isnan(cb_xyz[gj]).any(1)) + # Ensure mapping to valid ProDy/VMD indices + draw &= (ca_indices[gi] != -1) & (ca_indices[gj] != -1) gi, gj = gi[draw], gj[draw] - cb_dist = np.linalg.norm(cb_xyz[gi] - cb_xyz[gj], axis=1) # Cbeta-Cbeta -> dashing + cb_dist = np.linalg.norm(cb_xyz[gi] - cb_xyz[gj], axis=1) styles = np.where((cb_dist >= 3.5) & (cb_dist <= 6.5), 'solid', 'dashed') - fo.write(f'draw color {color}\n') - for a, b, style in zip(gi, gj, styles): - p1, p2 = ca_xyz[a], ca_xyz[b] - fo.write("draw line {%.3f %.3f %.3f} {%.3f %.3f %.3f} style %s width 2\n" - % (p1[0], p1[1], p1[2], p2[0], p2[1], p2[2], style)) - - fo.write('''mol delrep top 0 - mol color Beta - mol representation NewCartoon 0.300000 10.000000 4.100000 0 - mol selection all - mol material Opaque - mol addrep top - color scale method GWR - ''') - - if movie_name: - fo.write('''axes location Off - color Display Background white - display resize 800 800 - display projection Orthographic - display depthcue off - display resetview - display resize [expr [lindex [display get size] 0]/2*2] [expr [lindex [display get size] 1]/2*2] ;#Resize display to even height and width - display update ui - - # Set up the movie directory and base file name - mkdir movie_tmp - set workdir "movie_tmp" - ''' + f'set basename "{movie_name}"' + ''' - set numframes 360 - set framerate 25 - - # Function to rotate the molecule and capture frames - proc captureFrames {} { - global workdir basename numframes - for {set i 0} {$i < $numframes} {incr i} { - # Rotate the molecule around the Y-axis - rotate y by 1 - - # Capture the frame - set output [format "%s/$basename.%05d.tga" $workdir $i] - render snapshot $output - } - } - - # Function to convert frames to MP4 - proc convertToMP4 {} { - global workdir basename numframes framerate - - set mybasefilename [format "%s/%s" $workdir $basename] - set outputFile [format "%s.mp4" $basename] - - # Construct and execute the ffmpeg command - - set command "ffmpeg -y -framerate $framerate -i $mybasefilename.%05d.tga -c:v libx264 -profile:v high -crf 20 -pix_fmt yuv420p $outputFile" - puts "Executing: $command" - exec ffmpeg -y -framerate $framerate -i $mybasefilename.%05d.tga -c:v libx264 -profile:v high -crf 20 -pix_fmt yuv420p $outputFile >&@ stdout - } + # 3. Draw lines instantly using the array + if len(gi) > 0: + to_write.append(f' graphics $molid color {color}\n') + for a, b, style in zip(gi, gj, styles): + id1, id2 = ca_indices[a], ca_indices[b] + to_write.append(f' graphics $molid line $coords({id1}) $coords({id2}) style {style} width 2\n') + to_write.append('\n') - # Main script execution - captureFrames - convertToMP4 + to_write.append(' puts "Lines redrawn successfully."\n') + to_write.append('}\n\n') + + # Execute the procedure once initially to render the graphic + to_write.append('redraw_bound_line\n\n') + # --- END OF Tcl PROCEDURE --- + + # Visual representation setup + to_write.append('''mol delrep top 0 + mol color Beta + mol representation NewCartoon 0.300000 10.000000 4.100000 0 + mol selection all + mol material Opaque + mol addrep top + color scale method GWR + ''') - # Cleanup the TGA files if desired + if movie_name: + to_write.append('''axes location Off + color Display Background white + display resize 800 800 + display projection Orthographic + display depthcue off + display resetview + display resize [expr [lindex [display get size] 0]/2*2] [expr [lindex [display get size] 1]/2*2] + display update ui + + mkdir movie_tmp + set workdir "movie_tmp" + ''' + f'set basename "{movie_name}"' + ''' + set numframes 360 + set framerate 25 + + proc captureFrames {} { + global workdir basename numframes for {set i 0} {$i < $numframes} {incr i} { + rotate y by 1 set output [format "%s/$basename.%05d.tga" $workdir $i] - exec rm $output + render snapshot $output } - exit + } + + proc convertToMP4 {} { + global workdir basename numframes framerate + set mybasefilename [format "%s/%s" $workdir $basename] + set outputFile [format "%s.mp4" $basename] + set command "ffmpeg -y -framerate $framerate -i $mybasefilename.%05d.tga -c:v libx264 -profile:v high -crf 20 -pix_fmt yuv420p $outputFile" + puts "Executing: $command" + exec ffmpeg -y -framerate $framerate -i $mybasefilename.%05d.tga -c:v libx264 -profile:v high -crf 20 -pix_fmt yuv420p $outputFile >&@ stdout + } + + captureFrames + convertToMP4 + + for {set i 0} {$i < $numframes} {incr i} { + set output [format "%s/$basename.%05d.tga" $workdir $i] + exec rm $output + } + exit ''') elif still_image_name: - fo.write(f'set output "{still_image_name}"' + ''' - render snapshot $output - exit + to_write.append(f'set output "{still_image_name}"' + ''' + render snapshot $output + exit ''') - fo.close() + + with open(tcl_script, 'w+') as _fo: + _fo.write(''.join(to_write)) return tcl_script