diff --git a/frustratometer/classes/AWSEM.py b/frustratometer/classes/AWSEM.py index c25c259..a9f3c5c 100644 --- a/frustratometer/classes/AWSEM.py +++ b/frustratometer/classes/AWSEM.py @@ -64,17 +64,21 @@ class AWSEM(Frustratometer): q = 20 aa_map_awsem_list = [0, 0, 4, 3, 6, 13, 7, 8, 9, 11, 10, 12, 2, 14, 5, 1, 15, 16, 19, 17, 18] #A gap has no energy aa_map_awsem_x, aa_map_awsem_y = np.meshgrid(aa_map_awsem_list, aa_map_awsem_list, indexing='ij') - - def __init__(self, + + def __init__(self, pdb_structure: Structure, - sequence: Union[str, None] =None, - expose_indicator_functions: bool=False, - sparse: bool=True, - fast: bool=False, - backend: str='numba', + sequence: Union[str, None] = None, + expose_indicator_functions: bool = False, + sparse: bool = True, + fast: bool = False, + backend: str = 'numba', **parameters): """ - Generate AWSEM object + Generate AWSEM object. + + This method simply delegates to :meth:`AWSEM.from_structure`. + New code is encouraged to call :meth:`AWSEM.from_structure` + or :meth:`AWSEM.from_distance_matrix` directly. Parameters ---------- @@ -92,7 +96,215 @@ def __init__(self, ------- AWSEM object """ + built = type(self).from_structure( + pdb_structure, + sequence=sequence, + expose_indicator_functions=expose_indicator_functions, + sparse=sparse, + fast=fast, + backend=backend, + **parameters) + self.__dict__.update(built.__dict__) + + # ------------------------------------------------------------------ + # Factory classmethods (one per input source) + # ------------------------------------------------------------------ + + @classmethod + def from_structure(cls, + pdb_structure: Structure, + sequence: Union[str, None] = None, + expose_indicator_functions: bool = False, + sparse: bool = True, + fast: bool = False, + backend: str = 'numba', + **parameters) -> "AWSEM": + """ + Build an AWSEM object from a Structure object (the common case: a PDB + file plus everything the Structure class derives from it, including + atom coordinates, chain/residue bookkeeping, and — for membrane + calculations — the z-coordinates used to place residues relative to + the membrane). + + Parameters + ---------- + pdb_structure : Structure + Structure object generated by the Structure class. + sequence : str, optional + One-letter-code amino acid sequence. Defaults to the structure's sequence. + expose_indicator_functions : bool + If True, indicator functions of the contact and burial energy terms are exposed. + sparse : bool + If True (default), build a sparse Potts model. If False, build the dense (N, N, Q, Q) model. + fast : bool + If True, skip Potts model construction and use the fast numba/cuda backend. + Requires a sparse Structure (Structure built with sparse=True). + backend : str + Backend to use when fast=True. One of 'numba' or 'cuda'. + **parameters + Any AWSEMParameters field overrides (k_contact, gamma, membrane, etc). + + Returns + ------- + AWSEM + """ + is_sparse_structure = pdb_structure._is_sparse + return cls._construct( + distance_matrix=pdb_structure.distance_matrix, + full_pdb_distance_matrix=pdb_structure.full_pdb_distance_matrix, + sequence=sequence if sequence is not None else pdb_structure.sequence, + chain_breaks=pdb_structure.chain_breaks, + sparse_distance_matrix=pdb_structure._sparse_distance_matrix if is_sparse_structure else None, + sparse_distance_matrix_elec=pdb_structure._sparse_distance_matrix_elec if is_sparse_structure else None, + structure=pdb_structure.structure, + chain=pdb_structure.chain, + pdb_file=pdb_structure.pdb_file, + full_to_aligned_index_dict=pdb_structure.full_to_aligned_index_dict, + init_index_shift=pdb_structure.init_index_shift, + fin_index_shift=getattr(pdb_structure, 'fin_index_shift', None), + distance_matrix_method=getattr(pdb_structure, 'distance_matrix_method', None), + seq_selection=getattr(pdb_structure, 'seq_selection', None), + expose_indicator_functions=expose_indicator_functions, + sparse=sparse, + fast=fast, + backend=backend, + **parameters) + + @classmethod + def from_distance_matrix(cls, + distance_matrix, + sequence: str, + chain_breaks=None, + full_pdb_distance_matrix=None, + expose_indicator_functions: bool = False, + sparse: bool = False, + fast: bool = False, + backend: str = 'numba', + z_coords=None, + **parameters) -> "AWSEM": + """ + Build an AWSEM object directly from a residue-residue distance + matrix, without a Structure/PDB file. Useful when the distance matrix + was computed or obtained independently (e.g. predicted contacts, a + coarse-grained model, or distances loaded from another pipeline). + + Bookkeeping fields that only make sense for an actual PDB (atom + coordinates, chain id, residue numbering, the on-disk file path) are + left as None on the returned object; nothing in the energy or + frustration calculations depends on them. + + Parameters + ---------- + distance_matrix : np.ndarray (N, N) + Residue-residue distance matrix (Angstrom). + sequence : str + One-letter-code amino acid sequence, length N. + chain_breaks : optional + Chain break specification, as consumed by `frustration.compute_mask`. + Defaults to None (no chain breaks). + full_pdb_distance_matrix : np.ndarray (N, N), optional + Distance matrix for burial calculations if it differs from + `distance_matrix` (e.g. a substructure embedded in a larger + complex). Defaults to `distance_matrix`. + expose_indicator_functions : bool + If True, indicator functions of the contact and burial energy terms are exposed. + sparse : bool + If True, build a sparse Potts model. If False (default here, since a dense + matrix was supplied), build the dense (N, N, Q, Q) model. + fast : bool + If True, skip Potts model construction and use the fast numba/cuda backend. + Requires a sparse distance_matrix and full_pdb_distance_matrix (SparseMatrix) + z_coords : np.ndarray (N,), optional + Per-residue z-coordinate relative to the membrane, only required + if `membrane=True` is also passed in `**parameters` (there is no + Structure here to derive it from). Raises if membrane mode is + requested without z_coords. + **parameters + Any AWSEMParameters field overrides (k_contact, gamma, membrane, etc). + + Returns + ------- + AWSEM + """ + # TODO: decide how to handle cases where one matrix is sparse and another is dense + if full_pdb_distance_matrix: + if (isinstance(distance_matrix, SparseMatrix) and not isinstance(full_pdb_distance_matrix, SparseMatrix)) or \ + (not isinstance(distance_matrix, SparseMatrix) and isinstance(full_pdb_distance_matrix, SparseMatrix)): + raise TypeError("not sure what to do when distance_matrix and full_pdb_distance_matrix are different types") + # check user input + if fast and not isinstance(distance_matrix, SparseMatrix): + # helpful warning + print('\ndistance_matrix seems to be a dense matrix. \n' + f'Dense matrices cannot use the True-like argument fast={fast}. \n' + 'This will probably raise an Exception later on in the initialization.\n') + #'Reverting to fast=False.\n') + #fast = False + #with open('before_construct_called.txt','a') as file: + # file.write('construct called --------') + # file.write(str(distance_matrix)) + # file.write('-------------------------') + return cls._construct( + distance_matrix=distance_matrix, + full_pdb_distance_matrix=full_pdb_distance_matrix if full_pdb_distance_matrix is not None else distance_matrix, + sequence=sequence, + chain_breaks=chain_breaks, + sparse_distance_matrix=distance_matrix, + sparse_distance_matrix_elec=distance_matrix, + z_coords=z_coords, + expose_indicator_functions=expose_indicator_functions, + sparse=sparse, + fast=fast, + backend=backend, + **parameters) + + # ------------------------------------------------------------------ + # Shared construction logic + # ------------------------------------------------------------------ + + @classmethod + def _construct(cls, + distance_matrix, + full_pdb_distance_matrix, + sequence: str, + chain_breaks, + sparse_distance_matrix=None, + sparse_distance_matrix_elec=None, + z_coords=None, + structure=None, + chain=None, + pdb_file=None, + full_to_aligned_index_dict=None, + resid=None, + init_index_shift: int = 0, + fin_index_shift=None, + distance_matrix_method=None, + seq_selection=None, + expose_indicator_functions: bool = False, + sparse: bool = True, + fast: bool = False, + backend: str = 'numba', + **parameters) -> "AWSEM": + """Shared implementation behind all factory classmethods and __init__. + + Takes the primitive fields the physics actually needs (distance + matrices, sequence, chain_breaks) plus optional PDB-provenance fields + (structure/chain/pdb_file/resid/...) that are only used for + bookkeeping or for rebuilding a dense matrix from a sparse Structure. + Builds a fresh instance via ``cls.__new__`` and dispatches to the + appropriate mode-specific initializer (`_init_fast`, + `_init_from_sparse_distances`, or `_init_from_dense_distances`). + """ + # We don't want to do a simple cls() here + # because the implicit __init__() would lock us + # in an infinite loop by calling this method + # (indirectly, by way of from_structure). + # The __init__() must be written in this way + # for backward compatibility. + self = cls.__new__(cls) + # So, here, we just call the bare-bones __new__() + # and initialize our object using the below code. + #Set attributes p = AWSEMParameters(**parameters) if p.min_sequence_separation_contact is None: @@ -145,71 +357,86 @@ def __init__(self, self.membrane_protein_gamma = membrane_gamma_obj['Protein'][0] * p.k_relative_mem self.membrane_water_gamma = membrane_gamma_obj['Water'][0] * p.k_relative_mem - #Structure details - self.full_to_aligned_index_dict=pdb_structure.full_to_aligned_index_dict - if sequence is None: - self.sequence=pdb_structure.sequence - else: - self.sequence=sequence - self.structure=pdb_structure.structure - self.chain=pdb_structure.chain - self.pdb_file=pdb_structure.pdb_file - self.init_index_shift=pdb_structure.init_index_shift - self.distance_matrix=pdb_structure.distance_matrix - self.full_pdb_distance_matrix=pdb_structure.full_pdb_distance_matrix - self.chain_breaks=pdb_structure.chain_breaks - self._distance_is_sparse = pdb_structure._is_sparse + #Structure/provenance details (None-able when there is no PDB Structure) + self.full_to_aligned_index_dict = full_to_aligned_index_dict + self.sequence = sequence + self.structure = structure + self.chain = chain + self.pdb_file = pdb_file + self.init_index_shift = init_index_shift + self.fin_index_shift = fin_index_shift + self.distance_matrix=distance_matrix + self.full_pdb_distance_matrix=full_pdb_distance_matrix + self.chain_breaks=chain_breaks + self._distance_is_sparse = sparse_distance_matrix is not None #Sparse matrices - self._sparse_distance_matrix = None - self._sparse_distance_matrix_elec = None - if self._distance_is_sparse: - self._sparse_distance_matrix = pdb_structure._sparse_distance_matrix - self._sparse_distance_matrix_elec = pdb_structure._sparse_distance_matrix_elec - if not sparse: - # Dense Potts model from sparse distances. - from frustratometer.pdb import distance as _pdb_dist - dense_dm = _pdb_dist.get_dense_distance_matrix( - pdb_file=pdb_structure.pdb_file, - chain=pdb_structure.chain, - method=pdb_structure.distance_matrix_method) - if pdb_structure.seq_selection is not None: - dense_dm = dense_dm[ - pdb_structure.init_index_shift:pdb_structure.fin_index_shift, - pdb_structure.init_index_shift:pdb_structure.fin_index_shift] - self.distance_matrix = dense_dm - self.full_pdb_distance_matrix = dense_dm - self._distance_is_sparse = False - else: - self._sparse_distance_matrix = None - self._sparse_distance_matrix_elec = None - selection_CA = self.structure.select('name CA') - selection_CB = self.structure.select('name CB or (resname GLY IGL and name CA)') - - resid = selection_CB.getResindices() - self.resid=resid - self.N=len(self.resid) + self._sparse_distance_matrix = sparse_distance_matrix + self._sparse_distance_matrix_elec = sparse_distance_matrix_elec + #with open('construct_called.txt','a') as file: + # file.write('construct called --------') + # file.write(str(sparse_distance_matrix_elec)) + # file.write('-------------------------') + if self._distance_is_sparse and not sparse: + # Dense Potts model requested from a sparse distance matrix: this can + # only be rebuilt from the original PDB file (only available when + # constructed via `from_structure`). + if pdb_file is None or distance_matrix_method is None: + raise ValueError( + "Building a dense Potts model (sparse=False) from a sparse " + "distance matrix requires the original pdb_file and " + "distance_matrix_method, which are only available when " + "constructing via AWSEM.from_structure(...).") + from frustratometer.pdb import distance as _pdb_dist + dense_dm = _pdb_dist.get_dense_distance_matrix( + pdb_file=pdb_file, + chain=chain, + method=distance_matrix_method) + if seq_selection is not None: + dense_dm = dense_dm[init_index_shift:fin_index_shift, + init_index_shift:fin_index_shift] + self.distance_matrix = dense_dm + self.full_pdb_distance_matrix = dense_dm + self._distance_is_sparse = False + + # Residue bookkeeping: derived from atom selections when a Structure is + # available, otherwise falls back to the sequence length. + selection_CA = selection_CB = None + if structure is not None: + selection_CA = structure.select('name CA') + selection_CB = structure.select('name CB or (resname GLY IGL and name CA)') + resid = selection_CB.getResindices() + self.resid = resid + self.N = len(self.resid) if self.resid is not None else len(self.sequence) assert self.N == len(self.sequence), "The pdb is incomplete. Try setting 'repair_pdb=True' when constructing the Structure object." # Membrane alpha and phi_z computation if p.membrane: - # OpenAWSEM has two membrane Zim behaviors: - # - membrane_preassigned_term: CB (or CA fallback) - # - membrane_term (simple, no preassigned file): CA - # May need to adjust in the future to set custom z_source - if p.z_source == "Auto": - z_source = selection_CB if p.zim is not None else selection_CA - elif p.z_source == "CB": - z_source = selection_CB - elif p.z_source == "CA": - z_source = selection_CA + if structure is not None: + # OpenAWSEM has two membrane Zim behaviors: + # - membrane_preassigned_term: CB (or CA fallback) + # - membrane_term (simple, no preassigned file): CA + # May need to adjust in the future to set custom z_source + if p.z_source == "Auto": + z_source = selection_CB if p.zim is not None else selection_CA + elif p.z_source == "CB": + z_source = selection_CB + elif p.z_source == "CA": + z_source = selection_CA + else: + raise ValueError(f"Invalid z_source value: {p.z_source}. Must be 'Auto', 'CB', or 'CA'.") + z_coords_arr = z_source.getCoords()[:, 2] - p.membrane_center + elif z_coords is not None: + z_coords_arr = np.asarray(z_coords, dtype=np.float64) - p.membrane_center else: - raise ValueError(f"Invalid z_source value: {p.z_source}. Must be 'Auto', 'CB', or 'CA'.") - z_coords = z_source.getCoords()[:, 2] - p.membrane_center - self.z_coords = z_coords + raise ValueError( + "membrane=True requires either constructing from a Structure " + "(AWSEM.from_structure) or passing explicit z_coords to " + "AWSEM.from_distance_matrix.") + self.z_coords = z_coords_arr # alpha_i ≈ 1 inside membrane, ≈ 0 outside - self.alpha = 0.5 * np.tanh(p.eta_switching * (z_coords + p.z_m)) + 0.5 * np.tanh(p.eta_switching * (p.z_m - z_coords)) + self.alpha = 0.5 * np.tanh(p.eta_switching * (z_coords_arr + p.z_m)) + 0.5 * np.tanh(p.eta_switching * (p.z_m - z_coords_arr)) # phi_z for Zim potential (same shape but sharper boundary via k_m) - self.phi_z = 0.5 * np.tanh(p.k_m * (z_coords + p.z_m)) + 0.5 * np.tanh(p.k_m * (p.z_m - z_coords)) + self.phi_z = 0.5 * np.tanh(p.k_m * (z_coords_arr + p.z_m)) + 0.5 * np.tanh(p.k_m * (p.z_m - z_coords_arr)) # Per-residue DGwoct values # Load Wimley-White DGwoct scale (20 values in AWSEM aa order) _AA = 'ARNDCQEGHILKMFPSTWYV' @@ -246,16 +473,19 @@ def __init__(self, if fast: if not self._distance_is_sparse: - raise ValueError("fast=True requires a sparse Structure (sparse=True)") + raise ValueError("\nAWSEM(*args, **kwargs, fast=True) requires a sparse\n" + "Structure(*args, **kwargs, sparse=True) or SparseMatrix distance matrix\n") if backend not in ('numba', 'cuda'): raise ValueError(f"backend must be 'numba' or 'cuda', got {backend!r}") - self._init_fast(p, pdb_structure) + self._init_fast(p) elif self._distance_is_sparse: - self._init_from_sparse_distances(p, sparse, expose_indicator_functions, pdb_structure) + self._init_from_sparse_distances(p, sparse, expose_indicator_functions) else: - self._init_from_dense_distances(p, sparse, expose_indicator_functions, pdb_structure) + self._init_from_dense_distances(p, sparse, expose_indicator_functions) - def _init_from_sparse_distances(self, p, sparse, expose, pdb_structure): + return self + + def _init_from_sparse_distances(self, p, sparse, expose): """Initialize AWSEM physics from sparse COO distance data (no L×L arrays).""" dm = self._sparse_distance_matrix @@ -283,8 +513,6 @@ def _init_from_sparse_distances(self, p, sparse, expose, pdb_structure): # Handle burial_in_context substructure slicing if sel_dm.shape != dm.shape and self.burial_in_context: - self.init_index_shift = pdb_structure.init_index_shift - self.fin_index_shift = pdb_structure.fin_index_shift rho_r = rho_r[self.init_index_shift:self.fin_index_shift] self.rho = None # No dense rho matrix in sparse mode @@ -341,7 +569,7 @@ def _init_from_sparse_distances(self, p, sparse, expose, pdb_structure): self._build_sparse_from_contacts(p, contact_mask.row, contact_mask.col, theta_c, thetaII_c, sigma_water_c, sigma_protein_c, burial_energy, expose) - def _init_from_dense_distances(self, p, sparse, expose, pdb_structure): + def _init_from_dense_distances(self, p, sparse, expose): """Initialize AWSEM physics from dense distance matrix (original path).""" if self.burial_in_context==True: selected_matrix=self.full_pdb_distance_matrix @@ -367,8 +595,6 @@ def _init_from_dense_distances(self, p, sparse, expose, pdb_structure): rho_r = (rho).sum(axis=1) if self.full_pdb_distance_matrix.shape!=self.distance_matrix.shape: if self.burial_in_context==True: - self.init_index_shift=pdb_structure.init_index_shift - self.fin_index_shift=pdb_structure.fin_index_shift rho_r=rho_r[self.init_index_shift:self.fin_index_shift] self.rho_r=rho_r rho_b = np.expand_dims(rho_r, 1) @@ -409,16 +635,15 @@ def _init_from_dense_distances(self, p, sparse, expose, pdb_structure): else: self._build_dense(p, sequence_mask_contact, theta, thetaII, burial_energy, expose) - def _init_fast(self, p, pdb_structure): + def _init_fast(self, p): """Initialize fast mode: build FrustrationData, skip Potts model construction. - ``p`` and ``pdb_structure`` are retained so the sparse Potts model can be - materialized lazily (``_ensure_potts_model``) if an inherited method needs it. + ``p`` is retained so the sparse Potts model can be materialized lazily + (``_ensure_potts_model``) if an inherited method needs it. """ from ..frustration.data import FrustrationData self._init_params = p - self._init_pdb_structure = pdb_structure dm = self._sparse_distance_matrix self._frustration_data = FrustrationData.from_awsem(self) @@ -466,6 +691,18 @@ def _init_fast(self, p, pdb_structure): f"initialized: {e}. Use backend='numba' to run on CPU." ) from e + + def change_conformation(self, + alt_distance_matrix, alt_full_pdb_distance_matrix, + _sparse_distance_matrix, _sparse_distance_matrix_elec): + # WARNING: assumes sparse/fast mode, and distance matrices must be sparse! + self.distance_matrix = alt_distance_matrix + self.full_pdb_distance_matrix = alt_full_pdb_distance_matrix + self._sparse_distance_matrix = _sparse_distance_matrix + self._sparse_distance_matrix_elec = _sparse_distance_matrix_elec + self._init_fast(self._init_params) + + def _ensure_potts_model(self): """Lazily build the sparse Potts model for a fast-mode model. @@ -481,8 +718,7 @@ def _ensure_potts_model(self): return self._building_potts_model = True try: - self._init_from_sparse_distances(self._init_params, sparse=True, expose=False, - pdb_structure=self._init_pdb_structure) + self._init_from_sparse_distances(self._init_params, sparse=True, expose=False) finally: self._building_potts_model = False @@ -1079,4 +1315,4 @@ def zim_energy(self) -> float: if self._zim_h is None: return 0.0 seq_index = np.array([_AA_21.index(aa) for aa in self.sequence]) - return float(self._zim_h[np.arange(self.N), seq_index].sum()) + return float(self._zim_h[np.arange(self.N), seq_index].sum()) \ No newline at end of file