Merge pull request #6 from Becksteinlab:person2-feature-engg

Feature Engineering second take with new addtions

Author: amruthesht

confostate/features/_structure.py

@@ -90,11 +90,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")

confostate/features/domains.py

@@ -11,8 +11,13 @@
     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 +41,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:
@@ -80,3 +91,69 @@ 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, 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

confostate/features/pipeline.py

@@ -40,7 +40,12 @@ def extract_features(
     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

docs/features/features.md

@@ -56,7 +56,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 |
 
@@ -108,6 +108,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, ...)`
 
@@ -152,7 +153,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_<PDB>` | Distance change vs reference structure | Computed |
+| `domain_*_angle_delta_vs_<PDB>` | Angle change vs reference structure | Computed |
+| `domain_gate_TM1_TM6_distance_delta_vs_<PDB>` | Gate distance change vs reference | Computed |
 
 Helix boundaries are in `LEUT_TM_HELICES`.
 
@@ -185,6 +195,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_<PDB>` | 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`.

tests/test_features.py

@@ -43,12 +43,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
@@ -74,6 +96,7 @@ def test_extract_features(sample_pdb):
     )
     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