From 9d1c3d53014b09b1d738e859ea79a506d9466625 Mon Sep 17 00:00:00 2001 From: Amruthesh Thirumalaiswamy Date: Mon, 27 Jul 2026 15:59:27 -0700 Subject: [PATCH 1/3] Add domain delta features vs reference structures Compute inter-helix distance and angle changes relative to curated reference PDBs, and document the full feature column manifest for Person 3. --- confostate/features/_structure.py | 5 -- confostate/features/domains.py | 92 +++++++++++++++++++++++++++---- confostate/features/pipeline.py | 7 ++- docs/features.md | 43 ++++++++++++++- tests/test_features.py | 25 ++++++++- 5 files changed, 150 insertions(+), 22 deletions(-) diff --git a/confostate/features/_structure.py b/confostate/features/_structure.py index a0a3e24..a16fb8c 100644 --- a/confostate/features/_structure.py +++ b/confostate/features/_structure.py @@ -86,11 +86,6 @@ def load_structure( return structure -def sort_by_resid(atomgroup: AtomGroup) -> AtomGroup: - """Return atom group sorted by residue number.""" - return atomgroup[np.argsort(atomgroup.resids)] - - def center_of_mass(atomgroup: AtomGroup) -> np.ndarray: if len(atomgroup) == 0: raise ValueError("Cannot compute center of mass for empty atom group") diff --git a/confostate/features/domains.py b/confostate/features/domains.py index efe3948..6d36feb 100644 --- a/confostate/features/domains.py +++ b/confostate/features/domains.py @@ -11,8 +11,10 @@ angle_between_vectors, center_of_mass, helix_axis, + load_structure, pairwise_distance, ) +from confostate.features.rmsd import LEUT_REFERENCE_STRUCTURES, _resolve_reference_path LEUT_TM_HELICES: dict[str, tuple[int, int]] = { "TM1": (22, 52), @@ -36,29 +38,35 @@ ("TM3", "TM10"), ) +_DOMAIN_METRIC_KEYS = ( + "domain_TM1_TM7_distance", + "domain_TM1_TM7_angle", + "domain_TM1_TM6_distance", + "domain_TM1_TM6_angle", + "domain_TM5_TM7_distance", + "domain_TM5_TM7_angle", + "domain_TM3_TM10_distance", + "domain_TM3_TM10_angle", + "domain_gate_TM1_TM6_distance", +) -def extract_domain_features( + +def _domain_geometry( structure: StructureData, - tm_helices: Optional[dict[str, tuple[int, int]]] = None, - domain_pairs: tuple[tuple[str, str], ...] = LEUT_DOMAIN_PAIRS, + tm_helices: dict[str, tuple[int, int]], + domain_pairs: tuple[tuple[str, str], ...], ) -> dict[str, float]: - """ - Compute pairwise helix COM distances and inter-helix angles. - - Uses MDAnalysis selections and ``AtomGroup.center_of_mass()``. - """ - helices = tm_helices or LEUT_TM_HELICES - features: dict[str, float] = {} - + """Compute absolute inter-helix distances and angles.""" coms: dict[str, np.ndarray] = {} axes: dict[str, np.ndarray] = {} - for name, (start, end) in helices.items(): + for name, (start, end) in tm_helices.items(): ag = structure.select_ca_range(start, end) if len(ag) == 0: continue coms[name] = center_of_mass(ag) axes[name] = helix_axis(ag) + features: dict[str, float] = {} for helix_a, helix_b in domain_pairs: key_base = f"domain_{helix_a}_{helix_b}" if helix_a not in coms or helix_b not in coms: @@ -74,3 +82,63 @@ def extract_domain_features( features["domain_gate_TM1_TM6_distance"] = float("nan") return features + + +def _domain_deltas( + base_features: dict[str, float], + reference_dir: str, + reference_structures: dict[str, str], + tm_helices: dict[str, tuple[int, int]], + domain_pairs: tuple[tuple[str, str], ...], +) -> dict[str, float]: + """Compute domain metric deltas vs each unique reference PDB.""" + deltas: dict[str, float] = {} + unique_refs = sorted(set(reference_structures.values())) + + for ref_pdb_id in unique_refs: + try: + ref_path = _resolve_reference_path(ref_pdb_id, reference_dir) + except FileNotFoundError: + continue + + ref_structure = load_structure(str(ref_path), pdb_id=ref_pdb_id) + ref_features = _domain_geometry(ref_structure, tm_helices, domain_pairs) + + for key in _DOMAIN_METRIC_KEYS: + if key not in base_features or key not in ref_features: + continue + base_val = base_features[key] + ref_val = ref_features[key] + if np.isnan(base_val) or np.isnan(ref_val): + deltas[f"{key}_delta_vs_{ref_pdb_id}"] = float("nan") + else: + deltas[f"{key}_delta_vs_{ref_pdb_id}"] = float(base_val - ref_val) + + return deltas + + +def extract_domain_features( + structure: StructureData, + tm_helices: Optional[dict[str, tuple[int, int]]] = None, + domain_pairs: tuple[tuple[str, str], ...] = LEUT_DOMAIN_PAIRS, + reference_dir: Optional[str] = None, + reference_structures: Optional[dict[str, str]] = None, + include_deltas: bool = True, +) -> dict[str, float]: + """ + Compute pairwise helix COM distances, inter-helix angles, and optional deltas. + + When ``reference_dir`` is set, also returns deltas vs each curated reference + PDB (same references as ``rmsd.py``), e.g. + ``domain_TM1_TM7_distance_delta_vs_3F3E``. + """ + helices = tm_helices or LEUT_TM_HELICES + features = _domain_geometry(structure, helices, domain_pairs) + + if include_deltas and reference_dir: + refs = reference_structures or LEUT_REFERENCE_STRUCTURES + features.update( + _domain_deltas(features, reference_dir, refs, helices, domain_pairs) + ) + + return features diff --git a/confostate/features/pipeline.py b/confostate/features/pipeline.py index c311e35..3feb474 100644 --- a/confostate/features/pipeline.py +++ b/confostate/features/pipeline.py @@ -19,7 +19,7 @@ def extract_features( pdb_path: str, pdb_id: Optional[str] = None, - family: str = "LeuT", + family: str = "LeuT", # reserved for multi-family support (LeuT only for now) annotations_row: Optional[dict[str, Any]] = None, reference_dir: Optional[str] = None, include_rmsd: bool = True, @@ -35,7 +35,10 @@ def extract_features( features: dict[str, float] = {} features.update(extract_cavity_features(structure, membrane_normal=membrane_normal)) - features.update(extract_domain_features(structure)) + features.update(extract_domain_features( + structure, + reference_dir=reference_dir or str(Path(pdb_path).parent), + )) features.update(extract_orientation_features(structure, annotations_row=annotations_row)) if include_rmsd: diff --git a/docs/features.md b/docs/features.md index 5b4b091..4d370ed 100644 --- a/docs/features.md +++ b/docs/features.md @@ -55,7 +55,7 @@ Each row = one PDB structure. Columns fall into three categories: | Category | Columns | Status | Source | |----------|---------|--------|--------| -| **Computed features (X)** | `cavity_*`, `domain_*`, `rmsd_*`, `opm_*`, `orientation_principal_axis_*` | **Live computation** | MDAnalysis + SciPy on PDB coordinates from RCSB | +| **Computed features (X)** | `cavity_*`, `domain_*`, `domain_*_delta_vs_*`, `rmsd_*`, `opm_*`, `orientation_principal_axis_*` | **Live computation** | MDAnalysis + SciPy on PDB coordinates from RCSB | | **Label (y)** | `conformation` | **Unverified stub** | Copied from annotations CSV; `conformation_status = literature_estimate` | | **Metadata** | `pdb_id`, `file_path` | **Real** | PDB ID list + local file path | @@ -107,6 +107,7 @@ Person 2 → features (X) extract_features() / feature vectors CS Person 3 → join X + y, train model confostate/data/datasets.py, models/ ``` +## API ### `extract_features(pdb_path, ...)` @@ -151,7 +152,16 @@ Inter-helix distances and angles for LeuT TM helices. | `domain_TM1_TM6_distance` | COM distance between TM1 and TM6 | | `domain_TM5_TM7_distance` | COM distance between TM5 and TM7 | | `domain_TM3_TM10_distance` | COM distance between TM3 and TM10 | -| `domain_gate_TM1_TM6_distance` | Gate-opening distance (TM1–TM6) | +| `domain_gate_TM1_TM6_distance` | Gate-opening distance (TM1–TM6) | Computed | + +For each reference PDB (`3F3A`, `3F3E`, `3F4J`, `3USI`), delta features are also +emitted when reference files are available in `reference_dir`: + +| Feature pattern | Description | Source | +|---|---|---| +| `domain_*_distance_delta_vs_` | Distance change vs reference structure | Computed | +| `domain_*_angle_delta_vs_` | Angle change vs reference structure | Computed | +| `domain_gate_TM1_TM6_distance_delta_vs_` | Gate distance change vs reference | Computed | Helix boundaries are in `LEUT_TM_HELICES`. @@ -184,6 +194,35 @@ available; otherwise estimates tilt/rotation from the structure principal axis. | `opm_depth` | Centroid depth relative to membrane plane (Å) | | `orientation_principal_axis_x/y/z` | Unit vector of first principal component | +## Feature column manifest + +Complete list of columns in `extract_features()` output and +`leu_t_feature_vectors.csv`: + +| Column | Source | Notes | +|--------|--------|-------| +| `cavity_volume` | Computed (PDB) | Convex hull of binding-pocket atoms | +| `cavity_accessibility_in` | Computed (PDB) | Inward membrane-side exposure proxy | +| `cavity_accessibility_out` | Computed (PDB) | Outward membrane-side exposure proxy | +| `domain_*_distance` | Computed (PDB) | Absolute TM helix COM distances | +| `domain_*_angle` | Computed (PDB) | Absolute TM helix axis angles | +| `domain_gate_TM1_TM6_distance` | Computed (PDB) | Gate helix distance | +| `domain_*_delta_vs_` | Computed (PDB) | Change vs reference structure | +| `opm_tilt_angle` | Computed (PDB) or annotations | OPM CSV used when not `N/A` | +| `opm_rotation_angle` | Computed (PDB) or annotations | OPM CSV used when not `N/A` | +| `opm_depth` | Computed (PDB) or annotations | OPM CSV used when not `N/A` | +| `opm_tm_count` | Annotations only | When provided by Person 1 | +| `orientation_principal_axis_x/y/z` | Computed (PDB) | Principal axis components | +| `rmsd_OF_open` | Computed (PDB) | RMSD vs 3F3E | +| `rmsd_IF_open` | Computed (PDB) | RMSD vs 3F3A | +| `rmsd_Occluded` | Computed (PDB) | RMSD vs 3F4J | +| `rmsd_Intermediate` | Computed (PDB) | RMSD vs 3USI | +| `rmsd_min` | Computed (PDB) | Minimum RMSD across references | +| `rmsd_best_state_index` | Computed (PDB) | Index of closest reference state | +| `pdb_id` | Metadata | From filename (batch export only) | +| `file_path` | Metadata | Local PDB path (batch export only) | +| `conformation` | Annotations CSV | **Provisional label** — see provenance section | + ## Dependencies on Person 1 (data) Person 1 will deliver a verified `data/annotations/leu_t_transporters.csv`. diff --git a/tests/test_features.py b/tests/test_features.py index ab6be0d..19a6c5e 100644 --- a/tests/test_features.py +++ b/tests/test_features.py @@ -38,12 +38,34 @@ def test_cavity_features(structure): def test_domain_features(structure): - features = extract_domain_features(structure) + features = extract_domain_features(structure, reference_dir=str(INPUT_DIR)) assert "domain_TM1_TM7_distance" in features assert features["domain_TM1_TM7_distance"] > 0 assert 0 <= features["domain_TM1_TM7_angle"] <= 180 +def test_domain_delta_features_self(structure): + """Delta vs same structure (3F3E) should be approximately zero.""" + if not (INPUT_DIR / "3F3E.pdb").exists(): + pytest.skip("Reference PDB 3F3E not downloaded") + features = extract_domain_features(structure, reference_dir=str(INPUT_DIR)) + key = "domain_TM1_TM7_distance_delta_vs_3F3E" + assert key in features + assert features[key] == pytest.approx(0.0, abs=0.01) + + +def test_domain_delta_features_different_conformation(): + """IF-open 3F3A should have non-zero delta vs OF-open reference 3F3E.""" + path = INPUT_DIR / "3F3A.pdb" + if not path.exists() or not (INPUT_DIR / "3F3E.pdb").exists(): + pytest.skip("Reference PDBs not downloaded") + structure = load_structure(str(path), pdb_id="3F3A") + features = extract_domain_features(structure, reference_dir=str(INPUT_DIR)) + key = "domain_TM1_TM7_distance_delta_vs_3F3E" + assert key in features + assert abs(features[key]) > 0.01 + + def test_orientation_features(structure): features = extract_orientation_features(structure) assert "opm_tilt_angle" in features @@ -67,6 +89,7 @@ def test_extract_features(sample_pdb): features = extract_features(sample_pdb, annotations_row=row, reference_dir=str(INPUT_DIR)) assert "cavity_volume" in features assert "domain_TM1_TM7_distance" in features + assert "domain_TM1_TM7_distance_delta_vs_3F3E" in features assert "opm_tilt_angle" in features # OPM columns are N/A in CSV, so tilt must be computed from coordinates. assert features["opm_tilt_angle"] > 0 From 40811d41647423456922a06b943082d94a5b876f Mon Sep 17 00:00:00 2001 From: Amruthesh Thirumalaiswamy Date: Mon, 10 Aug 2026 15:28:03 -0700 Subject: [PATCH 2/3] black reformat --- confostate/data/loader.py | 8 ++------ confostate/features/_structure.py | 7 ++----- confostate/features/cavity.py | 4 +--- confostate/features/orientation.py | 12 +++--------- confostate/features/pipeline.py | 8 ++++---- confostate/features/rmsd.py | 8 ++------ examples/load_data.py | 4 +--- scripts/download_structures.py | 4 +--- scripts/extract_feature_vectors.py | 4 +--- tests/test_features.py | 7 ++----- 10 files changed, 19 insertions(+), 47 deletions(-) diff --git a/confostate/data/loader.py b/confostate/data/loader.py index 251d524..eca0705 100644 --- a/confostate/data/loader.py +++ b/confostate/data/loader.py @@ -7,9 +7,7 @@ import pandas as pd -def load_annotations( - csv_path: str, family: Optional[str] = None -) -> pd.DataFrame: +def load_annotations(csv_path: str, family: Optional[str] = None) -> pd.DataFrame: """ Load structure annotations from a CSV file. @@ -82,8 +80,6 @@ def load_from_input_dir(input_dir: str = "./input") -> pd.DataFrame: data = [] for pdb_file in sorted(pdb_files): pdb_id = pdb_file.stem.upper() - data.append( - {"pdb_id": pdb_id, "file_path": str(pdb_file), "file_exists": True} - ) + data.append({"pdb_id": pdb_id, "file_path": str(pdb_file), "file_exists": True}) return pd.DataFrame(data) diff --git a/confostate/features/_structure.py b/confostate/features/_structure.py index 73e9e20..0f9e61f 100644 --- a/confostate/features/_structure.py +++ b/confostate/features/_structure.py @@ -34,9 +34,7 @@ def primary_chain_id(self) -> str: def ca_atoms(self) -> AtomGroup: """Alpha-carbon atoms for the primary chain.""" chain = self.primary_chain_id - ag = self.universe.select_atoms( - f"protein and chainID {chain} and name CA" - ) + ag = self.universe.select_atoms(f"protein and chainID {chain} and name CA") if len(ag) == 0: ag = self.universe.select_atoms(f"segid {chain} and name CA") return ag @@ -55,8 +53,7 @@ def select_residues( chain = self.primary_chain_id if heavy_atoms: sel = ( - f"protein and chainID {chain} " - f"and resid {resid_str} and not name H*" + f"protein and chainID {chain} " f"and resid {resid_str} and not name H*" ) return self.universe.select_atoms(sel) return self.ca_atoms.select_atoms(f"resid {resid_str}") diff --git a/confostate/features/cavity.py b/confostate/features/cavity.py index 5640013..8df4ddd 100644 --- a/confostate/features/cavity.py +++ b/confostate/features/cavity.py @@ -82,9 +82,7 @@ def extract_cavity_features( if membrane_normal is None: membrane_normal = np.array([0.0, 0.0, 1.0]) - binding_atoms = structure.select_residues( - binding_residues, heavy_atoms=True - ) + binding_atoms = structure.select_residues(binding_residues, heavy_atoms=True) if len(binding_atoms) == 0: binding_atoms = structure.select_residues(binding_residues) diff --git a/confostate/features/orientation.py b/confostate/features/orientation.py index 5a4fed7..c334796 100644 --- a/confostate/features/orientation.py +++ b/confostate/features/orientation.py @@ -36,15 +36,11 @@ def _maybe_float(value: Any) -> Optional[float]: return None -def _tilt_angle( - protein_axis: np.ndarray, membrane_normal: np.ndarray -) -> float: +def _tilt_angle(protein_axis: np.ndarray, membrane_normal: np.ndarray) -> float: return angle_between_vectors(protein_axis, membrane_normal) -def _rotation_angle( - protein_axis: np.ndarray, membrane_normal: np.ndarray -) -> float: +def _rotation_angle(protein_axis: np.ndarray, membrane_normal: np.ndarray) -> float: normal = membrane_normal / np.linalg.norm(membrane_normal) projected = protein_axis - np.dot(protein_axis, normal) * normal proj_norm = np.linalg.norm(projected) @@ -68,9 +64,7 @@ def _rotation_angle( return angle -def _membrane_depth( - structure: StructureData, membrane_normal: np.ndarray -) -> float: +def _membrane_depth(structure: StructureData, membrane_normal: np.ndarray) -> float: normal = membrane_normal / np.linalg.norm(membrane_normal) centroid = structure.ca_atoms.center_of_mass() return float(abs(np.dot(centroid, normal))) diff --git a/confostate/features/pipeline.py b/confostate/features/pipeline.py index c52d871..3ad9027 100644 --- a/confostate/features/pipeline.py +++ b/confostate/features/pipeline.py @@ -39,7 +39,9 @@ def extract_features( features.update(extract_cavity_features(structure, membrane_normal=membrane_normal)) features.update(extract_domain_features(structure)) - features.update(extract_orientation_features(structure, annotations_row=annotations_row)) + features.update( + extract_orientation_features(structure, annotations_row=annotations_row) + ) if include_rmsd: try: @@ -66,9 +68,7 @@ def extract_features_batch( pdb_id = Path(pdb_path).stem.upper() row_data: Optional[dict[str, Any]] = None if annotations_df is not None and "pdb_id" in annotations_df.columns: - matches = annotations_df[ - annotations_df["pdb_id"].str.upper() == pdb_id - ] + matches = annotations_df[annotations_df["pdb_id"].str.upper() == pdb_id] if len(matches) > 0: row_data = matches.iloc[0].to_dict() diff --git a/confostate/features/rmsd.py b/confostate/features/rmsd.py index 63cc059..05efc08 100644 --- a/confostate/features/rmsd.py +++ b/confostate/features/rmsd.py @@ -91,14 +91,10 @@ def extract_rmsd_features( features[f"rmsd_{state}"] = float(rmsd_val) rmsd_values.append(float(rmsd_val)) - features["rmsd_min"] = ( - float(min(rmsd_values)) if rmsd_values else float("nan") - ) + features["rmsd_min"] = float(min(rmsd_values)) if rmsd_values else float("nan") if rmsd_values: best_state = min(refs.keys(), key=lambda s: features[f"rmsd_{s}"]) - features["rmsd_best_state_index"] = float( - list(refs.keys()).index(best_state) - ) + features["rmsd_best_state_index"] = float(list(refs.keys()).index(best_state)) else: features["rmsd_best_state_index"] = float("nan") diff --git a/examples/load_data.py b/examples/load_data.py index 118d8f8..cca15bc 100644 --- a/examples/load_data.py +++ b/examples/load_data.py @@ -30,9 +30,7 @@ def main(): print("\n First 5 entries:") print( - df_annot[ - ["pdb_id", "conformation", "experimental_method", "year"] - ].head() + df_annot[["pdb_id", "conformation", "experimental_method", "year"]].head() ) else: print(f" ERROR: File not found: {annotations_path}") diff --git a/scripts/download_structures.py b/scripts/download_structures.py index d42cc22..f6284dc 100644 --- a/scripts/download_structures.py +++ b/scripts/download_structures.py @@ -49,9 +49,7 @@ def download_pdb(code: str, output_dir: str, overwrite: bool = False) -> bool: def main() -> None: - parser = argparse.ArgumentParser( - description="Download PDB files from RCSB." - ) + parser = argparse.ArgumentParser(description="Download PDB files from RCSB.") group = parser.add_mutually_exclusive_group() group.add_argument( "--codes-file", diff --git a/scripts/extract_feature_vectors.py b/scripts/extract_feature_vectors.py index aaf48c3..171f2a1 100644 --- a/scripts/extract_feature_vectors.py +++ b/scripts/extract_feature_vectors.py @@ -49,9 +49,7 @@ def main() -> None: merged = annotations.merge(structures, on="pdb_id", how="inner") if len(merged) == 0: - raise SystemExit( - "No overlap between annotations and downloaded PDB files." - ) + raise SystemExit("No overlap between annotations and downloaded PDB files.") print(f"Extracting features for {len(merged)} structures...") df = extract_features_batch( diff --git a/tests/test_features.py b/tests/test_features.py index 770d607..a063127 100644 --- a/tests/test_features.py +++ b/tests/test_features.py @@ -14,8 +14,7 @@ INPUT_DIR = Path(__file__).resolve().parent.parent / "input" ANNOTATIONS = ( - Path(__file__).resolve().parent.parent - / "data/annotations/leu_t_transporters.csv" + Path(__file__).resolve().parent.parent / "data/annotations/leu_t_transporters.csv" ) @@ -23,9 +22,7 @@ def sample_pdb() -> str: path = INPUT_DIR / "3F3E.pdb" if not path.exists(): - pytest.skip( - "Sample PDB not found. Run scripts/download_structures.py first." - ) + pytest.skip("Sample PDB not found. Run scripts/download_structures.py first.") return str(path) From ce7511dc0bcc81c40e9a1e73715446be19cf788c Mon Sep 17 00:00:00 2001 From: Amruthesh Thirumalaiswamy Date: Mon, 10 Aug 2026 15:41:46 -0700 Subject: [PATCH 3/3] ruff reformat --- confostate/data/loader.py | 8 ++++++-- confostate/features/_structure.py | 7 +++++-- confostate/features/cavity.py | 4 +++- confostate/features/domains.py | 23 ++++++++++++++++------- confostate/features/orientation.py | 12 +++++++++--- confostate/features/pipeline.py | 21 ++++++++++++++++----- confostate/features/rmsd.py | 8 ++++++-- examples/load_data.py | 4 +++- scripts/download_structures.py | 4 +++- scripts/extract_feature_vectors.py | 4 +++- tests/test_features.py | 7 +++++-- 11 files changed, 75 insertions(+), 27 deletions(-) diff --git a/confostate/data/loader.py b/confostate/data/loader.py index eca0705..251d524 100644 --- a/confostate/data/loader.py +++ b/confostate/data/loader.py @@ -7,7 +7,9 @@ import pandas as pd -def load_annotations(csv_path: str, family: Optional[str] = None) -> pd.DataFrame: +def load_annotations( + csv_path: str, family: Optional[str] = None +) -> pd.DataFrame: """ Load structure annotations from a CSV file. @@ -80,6 +82,8 @@ def load_from_input_dir(input_dir: str = "./input") -> pd.DataFrame: data = [] for pdb_file in sorted(pdb_files): pdb_id = pdb_file.stem.upper() - data.append({"pdb_id": pdb_id, "file_path": str(pdb_file), "file_exists": True}) + data.append( + {"pdb_id": pdb_id, "file_path": str(pdb_file), "file_exists": True} + ) return pd.DataFrame(data) diff --git a/confostate/features/_structure.py b/confostate/features/_structure.py index 0f9e61f..73e9e20 100644 --- a/confostate/features/_structure.py +++ b/confostate/features/_structure.py @@ -34,7 +34,9 @@ def primary_chain_id(self) -> str: def ca_atoms(self) -> AtomGroup: """Alpha-carbon atoms for the primary chain.""" chain = self.primary_chain_id - ag = self.universe.select_atoms(f"protein and chainID {chain} and name CA") + ag = self.universe.select_atoms( + f"protein and chainID {chain} and name CA" + ) if len(ag) == 0: ag = self.universe.select_atoms(f"segid {chain} and name CA") return ag @@ -53,7 +55,8 @@ def select_residues( chain = self.primary_chain_id if heavy_atoms: sel = ( - f"protein and chainID {chain} " f"and resid {resid_str} and not name H*" + f"protein and chainID {chain} " + f"and resid {resid_str} and not name H*" ) return self.universe.select_atoms(sel) return self.ca_atoms.select_atoms(f"resid {resid_str}") diff --git a/confostate/features/cavity.py b/confostate/features/cavity.py index 8df4ddd..5640013 100644 --- a/confostate/features/cavity.py +++ b/confostate/features/cavity.py @@ -82,7 +82,9 @@ def extract_cavity_features( if membrane_normal is None: membrane_normal = np.array([0.0, 0.0, 1.0]) - binding_atoms = structure.select_residues(binding_residues, heavy_atoms=True) + binding_atoms = structure.select_residues( + binding_residues, heavy_atoms=True + ) if len(binding_atoms) == 0: binding_atoms = structure.select_residues(binding_residues) diff --git a/confostate/features/domains.py b/confostate/features/domains.py index f1ee132..b161a56 100644 --- a/confostate/features/domains.py +++ b/confostate/features/domains.py @@ -14,7 +14,10 @@ load_structure, pairwise_distance, ) -from confostate.features.rmsd import LEUT_REFERENCE_STRUCTURES, _resolve_reference_path +from confostate.features.rmsd import ( + LEUT_REFERENCE_STRUCTURES, + _resolve_reference_path, +) LEUT_TM_HELICES: dict[str, tuple[int, int]] = { "TM1": (22, 52), @@ -108,7 +111,9 @@ def _domain_deltas( continue ref_structure = load_structure(str(ref_path), pdb_id=ref_pdb_id) - ref_features = _domain_geometry(ref_structure, tm_helices, domain_pairs) + ref_features = _domain_geometry( + ref_structure, tm_helices, domain_pairs + ) for key in _DOMAIN_METRIC_KEYS: if key not in base_features or key not in ref_features: @@ -118,7 +123,9 @@ def _domain_deltas( if np.isnan(base_val) or np.isnan(ref_val): deltas[f"{key}_delta_vs_{ref_pdb_id}"] = float("nan") else: - deltas[f"{key}_delta_vs_{ref_pdb_id}"] = float(base_val - ref_val) + deltas[f"{key}_delta_vs_{ref_pdb_id}"] = float( + base_val - ref_val + ) return deltas @@ -132,10 +139,10 @@ def extract_domain_features( include_deltas: bool = True, ) -> dict[str, float]: """ - Compute pairwise helix COM distances, inter-helix angles, and optional deltas. + Compute pairwise helix COM distances, angles, and optional deltas. - When ``reference_dir`` is set, also returns deltas vs each curated reference - PDB (same references as ``rmsd.py``), e.g. + When ``reference_dir`` is set, also returns deltas vs each curated + reference PDB (same references as ``rmsd.py``), e.g. ``domain_TM1_TM7_distance_delta_vs_3F3E``. """ helices = tm_helices or LEUT_TM_HELICES @@ -144,7 +151,9 @@ def extract_domain_features( if include_deltas and reference_dir: refs = reference_structures or LEUT_REFERENCE_STRUCTURES features.update( - _domain_deltas(features, reference_dir, refs, helices, domain_pairs) + _domain_deltas( + features, reference_dir, refs, helices, domain_pairs + ) ) return features diff --git a/confostate/features/orientation.py b/confostate/features/orientation.py index c334796..5a4fed7 100644 --- a/confostate/features/orientation.py +++ b/confostate/features/orientation.py @@ -36,11 +36,15 @@ def _maybe_float(value: Any) -> Optional[float]: return None -def _tilt_angle(protein_axis: np.ndarray, membrane_normal: np.ndarray) -> float: +def _tilt_angle( + protein_axis: np.ndarray, membrane_normal: np.ndarray +) -> float: return angle_between_vectors(protein_axis, membrane_normal) -def _rotation_angle(protein_axis: np.ndarray, membrane_normal: np.ndarray) -> float: +def _rotation_angle( + protein_axis: np.ndarray, membrane_normal: np.ndarray +) -> float: normal = membrane_normal / np.linalg.norm(membrane_normal) projected = protein_axis - np.dot(protein_axis, normal) * normal proj_norm = np.linalg.norm(projected) @@ -64,7 +68,9 @@ def _rotation_angle(protein_axis: np.ndarray, membrane_normal: np.ndarray) -> fl return angle -def _membrane_depth(structure: StructureData, membrane_normal: np.ndarray) -> float: +def _membrane_depth( + structure: StructureData, membrane_normal: np.ndarray +) -> float: normal = membrane_normal / np.linalg.norm(membrane_normal) centroid = structure.ca_atoms.center_of_mass() return float(abs(np.dot(centroid, normal))) diff --git a/confostate/features/pipeline.py b/confostate/features/pipeline.py index 3ad9027..7b0cfa7 100644 --- a/confostate/features/pipeline.py +++ b/confostate/features/pipeline.py @@ -22,7 +22,7 @@ def extract_features( pdb_path: str, pdb_id: Optional[str] = None, - family: str = "LeuT", # reserved for multi-family support (LeuT only for now) + family: str = "LeuT", annotations_row: Optional[dict[str, Any]] = None, reference_dir: Optional[str] = None, include_rmsd: bool = True, @@ -37,10 +37,19 @@ def extract_features( features: dict[str, float] = {} - features.update(extract_cavity_features(structure, membrane_normal=membrane_normal)) - features.update(extract_domain_features(structure)) features.update( - extract_orientation_features(structure, annotations_row=annotations_row) + extract_cavity_features(structure, membrane_normal=membrane_normal) + ) + features.update( + extract_domain_features( + structure, + reference_dir=reference_dir or str(Path(pdb_path).parent), + ) + ) + features.update( + extract_orientation_features( + structure, annotations_row=annotations_row + ) ) if include_rmsd: @@ -68,7 +77,9 @@ def extract_features_batch( pdb_id = Path(pdb_path).stem.upper() row_data: Optional[dict[str, Any]] = None if annotations_df is not None and "pdb_id" in annotations_df.columns: - matches = annotations_df[annotations_df["pdb_id"].str.upper() == pdb_id] + matches = annotations_df[ + annotations_df["pdb_id"].str.upper() == pdb_id + ] if len(matches) > 0: row_data = matches.iloc[0].to_dict() diff --git a/confostate/features/rmsd.py b/confostate/features/rmsd.py index 05efc08..63cc059 100644 --- a/confostate/features/rmsd.py +++ b/confostate/features/rmsd.py @@ -91,10 +91,14 @@ def extract_rmsd_features( features[f"rmsd_{state}"] = float(rmsd_val) rmsd_values.append(float(rmsd_val)) - features["rmsd_min"] = float(min(rmsd_values)) if rmsd_values else float("nan") + features["rmsd_min"] = ( + float(min(rmsd_values)) if rmsd_values else float("nan") + ) if rmsd_values: best_state = min(refs.keys(), key=lambda s: features[f"rmsd_{s}"]) - features["rmsd_best_state_index"] = float(list(refs.keys()).index(best_state)) + features["rmsd_best_state_index"] = float( + list(refs.keys()).index(best_state) + ) else: features["rmsd_best_state_index"] = float("nan") diff --git a/examples/load_data.py b/examples/load_data.py index cca15bc..118d8f8 100644 --- a/examples/load_data.py +++ b/examples/load_data.py @@ -30,7 +30,9 @@ def main(): print("\n First 5 entries:") print( - df_annot[["pdb_id", "conformation", "experimental_method", "year"]].head() + df_annot[ + ["pdb_id", "conformation", "experimental_method", "year"] + ].head() ) else: print(f" ERROR: File not found: {annotations_path}") diff --git a/scripts/download_structures.py b/scripts/download_structures.py index f6284dc..d42cc22 100644 --- a/scripts/download_structures.py +++ b/scripts/download_structures.py @@ -49,7 +49,9 @@ def download_pdb(code: str, output_dir: str, overwrite: bool = False) -> bool: def main() -> None: - parser = argparse.ArgumentParser(description="Download PDB files from RCSB.") + parser = argparse.ArgumentParser( + description="Download PDB files from RCSB." + ) group = parser.add_mutually_exclusive_group() group.add_argument( "--codes-file", diff --git a/scripts/extract_feature_vectors.py b/scripts/extract_feature_vectors.py index 171f2a1..aaf48c3 100644 --- a/scripts/extract_feature_vectors.py +++ b/scripts/extract_feature_vectors.py @@ -49,7 +49,9 @@ def main() -> None: merged = annotations.merge(structures, on="pdb_id", how="inner") if len(merged) == 0: - raise SystemExit("No overlap between annotations and downloaded PDB files.") + raise SystemExit( + "No overlap between annotations and downloaded PDB files." + ) print(f"Extracting features for {len(merged)} structures...") df = extract_features_batch( diff --git a/tests/test_features.py b/tests/test_features.py index a063127..770d607 100644 --- a/tests/test_features.py +++ b/tests/test_features.py @@ -14,7 +14,8 @@ INPUT_DIR = Path(__file__).resolve().parent.parent / "input" ANNOTATIONS = ( - Path(__file__).resolve().parent.parent / "data/annotations/leu_t_transporters.csv" + Path(__file__).resolve().parent.parent + / "data/annotations/leu_t_transporters.csv" ) @@ -22,7 +23,9 @@ def sample_pdb() -> str: path = INPUT_DIR / "3F3E.pdb" if not path.exists(): - pytest.skip("Sample PDB not found. Run scripts/download_structures.py first.") + pytest.skip( + "Sample PDB not found. Run scripts/download_structures.py first." + ) return str(path)