Add FastHydroMap PC regressors

.gitignore

@@ -70,6 +70,7 @@ training/data/rcsb_pdbs
 training/logs/
 training/checkpoints_direct_feat/
 training/checkpoints_delta_feat/
+training/figures/
 training/models/*.pt
 training/models/*.npz
 training/models/*.json

CITATION.cff

@@ -7,8 +7,8 @@ authors:
     given-names: "Samuel"
     email: "samuels.lobo@gmail.com"
 doi: 10.5281/zenodo.19744336
-version: "0.1.2"
-date-released: 2026-04-24
+version: "0.1.3"
+date-released: 2026-05-28
 url: "https://github.com/samlobe/FastHydroMap"
 repository-code: "https://github.com/samlobe/FastHydroMap"
 license: "MIT"

README.md

@@ -5,6 +5,7 @@
 [![DOI](https://zenodo.org/badge/1023802589.svg)](https://doi.org/10.5281/zenodo.19744336)
 
 FastHydroMap predicts per-residue dewetting free energies (`Fdewet`) from protein structures and trajectories.
+It can also predict water structuring  (`PC1`, `PC2`, `PC3` of the [water triplet angle distribution](https://doi.org/10.1021/acs.jpcb.3c00826)).
 
 <p align="center">
   <img
@@ -40,17 +41,20 @@ Typical usage:
 ```bash
 # Single structure
 fasthydromap predict examples/1A1U.pdb -o outputs/1A1U_fdewet
+fasthydromap predict examples/1A1U.pdb --quantity pc1 -o outputs/1A1U_pc1
 
 # Trajectory
 fasthydromap predict-trajectory examples/proteinG.pdb examples/proteinG_short.dcd -o outputs/proteinG_fdewet
 ```
 
+The default `--quantity fdewet` predicts dewetting free energy. Use `--quantity pc1`, `--quantity pc2`, or `--quantity pc3` to predict water-structure principal-component maps instead.
+
 ## Outputs
 
 For a single structure, FastHydroMap writes:
 
-- `*.csv`: one row per residue with `Fdewet`; with `--parts`, intrinsic and context columns are included
-- `*.pdb`: a copy of the input structure with predicted `Fdewet` written to B-factors
+- `*.csv`: one row per residue with the requested quantity; with `--parts`, intrinsic and context columns are included
+- `*.pdb`: a copy of the input structure with the requested quantity written to B-factors
 
 For a trajectory, FastHydroMap writes wide CSV files containing one row per frame and one column per residue.
 Use `--parts` to also write intrinsic, context, and per-frame summary CSVs.
@@ -62,14 +66,22 @@ Predictions for PTMs and other non-canonical chemistries should be treated cauti
 
 ## Visualization
 
-FastHydroMap writes `Fdewet` values to the B-factor column of output PDBs, so you can color structures directly in molecular viewers.
+FastHydroMap writes predicted values to the B-factor column of output PDBs, so you can color structures directly in molecular viewers.
 
 ChimeraX:
 
 ```bash
 color bfactor range 4,6.5 palette ^lipophilicity
 ```
 
+Example PC map coloring in ChimeraX:
+
+```bash
+color bfactor range -8,8 palette red-white-blue # example range and palette for PC1: redder = more tetrahedral angles
+color bfactor range -2,8 palette cyanmaroon # for PC2: maroon = more 90 deg angles (unstructured)
+color bfactor range -2,2 palette ^lipophilicity # for PC3: yellower = fewer 50 deg angles (i.e. fewer highly coordinated waters)
+```
+
 PyMOL:
 
 ```bash
@@ -85,7 +97,7 @@ If you use FastHydroMap in your research, please cite the software release:
 
 [![DOI](https://zenodo.org/badge/1023802589.svg)](https://doi.org/10.5281/zenodo.19744336)
 
-Lobo, S. FastHydroMap (Version 0.1.2) [Computer software]. Zenodo.
+Lobo, S. FastHydroMap (Version 0.1.3) [Computer software]. Zenodo.
 https://doi.org/10.5281/zenodo.19744336
 
 When the manuscript becomes available, please cite that as well.

docs/PYPI_RELEASE.md

@@ -38,21 +38,23 @@ That script builds the project, creates a fresh venv, installs the wheel, and ve
 - `fasthydromap --help`
 - `fasthydromap install-torch --dry-run`
 - `fasthydromap predict examples/1A1U.pdb -o ...`
+- `fasthydromap predict examples/1A1U.pdb --quantity pc1 -o ...`
 
 For a manual end-to-end check in the same environment:
 
 ```bash
 source .wheel-test/bin/activate
 fasthydromap install-torch --variant cpu
 fasthydromap predict examples/1A1U.pdb -o /tmp/1A1U_fdewet
+fasthydromap predict examples/1A1U.pdb --quantity pc1 -o /tmp/1A1U_pc1
 ```
 
 ## 4. Create a Git tag
 
 Example:
 
 ```bash
-VERSION=0.1.1
+VERSION=0.1.3
 git tag "v${VERSION}"
 git push origin "v${VERSION}"
 ```

pyproject.toml

@@ -4,8 +4,8 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name        = "fasthydromap"
-version     = "0.1.2"
-description = "Predict per-residue Fdewet using SASA + MPNN"
+version     = "0.1.3"
+description = "Predict per-residue FastHydroMap Fdewet and water-structure PC maps"
 authors     = [{name="Samuel Lobo", email="samuels.lobo@gmail.com"}]
 license     = "MIT"
 readme      = "README.md"
@@ -50,7 +50,7 @@ Issues = "https://github.com/samlobe/FastHydroMap/issues"
 where = ["src"]
 
 [tool.setuptools.package-data]
-FastHydroMap = ["weights/sasa_feature_stats.npz", "weights/mpnn_latest.pt"]
+FastHydroMap = ["weights/sasa_feature_stats.npz", "weights/*.pt"]
 
 [project.scripts]
 fasthydromap = "FastHydroMap.cli:main"

scripts/smoke_test_install.sh

@@ -13,3 +13,4 @@ mkdir -p "${OUTDIR}"
 fasthydromap --help >/dev/null
 fasthydromap install-torch --variant cpu
 fasthydromap predict examples/1A1U.pdb -o "${OUTDIR}/1A1U_fdewet"
+fasthydromap predict examples/1A1U.pdb --quantity pc1 -o "${OUTDIR}/1A1U_pc1"

scripts/test_wheel_install.sh

@@ -21,9 +21,12 @@ fasthydromap install-torch --dry-run
 fasthydromap install-torch --variant cpu
 OUTROOT="${TEST_ENV}/smoke_1A1U"
 fasthydromap predict "${ROOT_DIR}/examples/1A1U.pdb" -o "${OUTROOT}"
+fasthydromap predict "${ROOT_DIR}/examples/1A1U.pdb" --quantity pc1 -o "${TEST_ENV}/smoke_1A1U_pc1"
 
 test -f "${OUTROOT}.csv"
 test -f "${OUTROOT}.pdb"
+test -f "${TEST_ENV}/smoke_1A1U_pc1.csv"
+test -f "${TEST_ENV}/smoke_1A1U_pc1.pdb"
 
 echo
 echo "Wheel install smoke test passed in ${TEST_ENV}"

src/FastHydroMap/cli.py

@@ -9,13 +9,21 @@
 from .io.pdb            import write_bfactor
 from .install_torch import install_torch, torch_install_command
 
+QUANTITIES = ("fdewet", "pc1", "pc2", "pc3")
+QUANTITY_LABELS = {
+    "fdewet": "Fdewet",
+    "pc1": "PC1",
+    "pc2": "PC2",
+    "pc3": "PC3",
+}
+
 
 # ---------------------------------------------------------------------
 # CLI
 # ---------------------------------------------------------------------
 def _build_parser() -> argparse.ArgumentParser:
     ap = argparse.ArgumentParser(prog="fasthydromap",
-                                 description="FastHydroMap – infer per-residue Fdewet")
+                                 description="FastHydroMap – infer per-residue Fdewet or PC maps")
     sub = ap.add_subparsers(dest="cmd", required=True)
 
     # -------- predict -------------------------------------------------
@@ -31,6 +39,12 @@ def _build_parser() -> argparse.ArgumentParser:
         action="store_true",
         help="for single-structure predictions, include intrinsic and context columns",
     )
+    p.add_argument(
+        "--quantity",
+        choices=QUANTITIES,
+        default="fdewet",
+        help="per-residue quantity to predict (default: fdewet)",
+    )
 
     # -------- predict-trajectory --------------------------------------
     pt = sub.add_parser(
@@ -51,6 +65,12 @@ def _build_parser() -> argparse.ArgumentParser:
         action="store_true",
         help="write intrinsic, context, and per-frame summary CSVs in addition to total",
     )
+    pt.add_argument(
+        "--quantity",
+        choices=QUANTITIES,
+        default="fdewet",
+        help="per-residue quantity to predict (default: fdewet)",
+    )
 
     # -------- install-torch ------------------------------------------
     it = sub.add_parser(
@@ -129,7 +149,9 @@ def main() -> None:
         return
 
     FdewetPredictor = _load_predictor_or_exit()
-    predictor = FdewetPredictor()
+    predictor = FdewetPredictor(quantity=args.quantity)
+    quantity_label = QUANTITY_LABELS[args.quantity]
+    quantity_suffix = args.quantity
 
     if args.cmd == "predict":
         if args.parts and args.dcd is not None:
@@ -161,7 +183,7 @@ def main() -> None:
         # -------------------------------------------------------------
         outroot = args.outroot or args.pdb.with_suffix("")
         if args.outroot is None:
-            outroot = outroot.with_name(outroot.name + "_fdewet")
+            outroot = outroot.with_name(outroot.name + f"_{quantity_suffix}")
 
         csv_path = Path(f"{outroot}.csv")
         pdb_path = Path(f"{outroot}.pdb")
@@ -170,11 +192,11 @@ def main() -> None:
         if args.dcd is None:
             data = {
                 "residue": [str(r) for r in res_ids],
-                "Fdewet": scores.round(2),
+                quantity_label: scores.round(2),
             }
             if args.parts:
-                data["Fdewet_intrinsic"] = parts["intrinsic"].round(2)
-                data["Fdewet_context"] = parts["context"].round(2)
+                data[f"{quantity_label}_intrinsic"] = parts["intrinsic"].round(2)
+                data[f"{quantity_label}_context"] = parts["context"].round(2)
             df = pd.DataFrame(data)
         else:
             col_names = [str(r) for r in res_ids]               # columns = residue numbers
@@ -203,7 +225,7 @@ def main() -> None:
 
     outroot = args.outroot or args.dcd.with_suffix("")
     if args.outroot is None:
-        outroot = outroot.with_name(outroot.name + "_fdewet_traj")
+        outroot = outroot.with_name(outroot.name + f"_{quantity_suffix}_traj")
 
     total_csv = Path(f"{outroot}_total.csv")
     intrinsic_csv = Path(f"{outroot}_intrinsic.csv")

src/FastHydroMap/predictors/fdewet.py

@@ -24,6 +24,52 @@
 from ..utils.atom_names import backbone_alias_priority, canonical_backbone_atom_name
 
 WEIGHT_DIR = Path(__file__).parents[1] / "weights"
+QUANTITY_SPECS = {
+    "fdewet": {
+        "label": "Fdewet",
+        "weight": WEIGHT_DIR / "mpnn_latest.pt",
+        "k_nn": 12,
+        "n_rbf": 3,
+        "rbf_min": 2.0,
+        "rbf_max": 14.0,
+        "rbf_sigma": 4.0,
+    },
+    "pc1": {
+        "label": "PC1",
+        "weight": WEIGHT_DIR / "mpnn_pc1_latest.pt",
+        "k_nn": 12,
+        "n_rbf": 3,
+        "rbf_min": 2.0,
+        "rbf_max": 14.0,
+        "rbf_sigma": 4.0,
+    },
+    "pc2": {
+        "label": "PC2",
+        "weight": WEIGHT_DIR / "mpnn_pc2_latest.pt",
+        "k_nn": 12,
+        "n_rbf": 3,
+        "rbf_min": 2.0,
+        "rbf_max": 14.0,
+        "rbf_sigma": 4.0,
+    },
+    "pc3": {
+        "label": "PC3",
+        "weight": WEIGHT_DIR / "mpnn_pc3_latest.pt",
+        "k_nn": 12,
+        "n_rbf": 3,
+        "rbf_min": 2.0,
+        "rbf_max": 14.0,
+        "rbf_sigma": 4.0,
+    },
+}
+
+
+def normalize_quantity(quantity: str) -> str:
+    q = quantity.lower()
+    if q in QUANTITY_SPECS:
+        return q
+    raise ValueError(f"unknown FastHydroMap quantity {quantity!r}")
+
 
 def _display_residue_label(
     chain_id: str,
@@ -42,15 +88,26 @@ def _display_residue_label(
 class FdewetPredictor:
     def __init__(
         self,
+        quantity: str = "fdewet",
         k_nn: int = 12,
         n_rbf: int = 3,
         rbf_min: float = 2.0,
         rbf_max: float = 14.0,
         rbf_sigma: float = 4.0,
-        mpnn_pt: Path = WEIGHT_DIR / "mpnn_latest.pt",
+        mpnn_pt: Path | None = None,
         sasa_stats_npz: Path = WEIGHT_DIR / "sasa_feature_stats.npz",
         device: str | torch.device | None = None,
     ):
+        self.quantity = normalize_quantity(quantity)
+        spec = QUANTITY_SPECS[self.quantity]
+        if mpnn_pt is None:
+            mpnn_pt = spec["weight"]
+            k_nn = int(spec["k_nn"])
+            n_rbf = int(spec["n_rbf"])
+            rbf_min = float(spec["rbf_min"])
+            rbf_max = float(spec["rbf_max"])
+            rbf_sigma = float(spec["rbf_sigma"])
+        self.output_label = str(spec["label"])
         self.k = k_nn
         self.n_rbf = n_rbf
         self.rbf_min = rbf_min

tests/test_predictor_regression.py

@@ -145,6 +145,22 @@ def test_cli_predict_writes_outputs_for_path_outroot(monkeypatch, tmp_path):
     assert np.max(np.abs(out_scores - ref_scores)) < 0.5
 
 
+def test_cli_predict_can_write_pc_quantity(monkeypatch, tmp_path):
+    outroot = tmp_path / "1A1U_pc1"
+    monkeypatch.setattr(
+        sys,
+        "argv",
+        ["fasthydromap", "predict", str(PDB_PATH), "-o", str(outroot), "--quantity", "pc1"],
+    )
+
+    main()
+
+    df = pd.read_csv(Path(f"{outroot}.csv"))
+    assert list(df.columns) == ["residue", "PC1"]
+    assert len(df) > 0
+    assert np.isfinite(df["PC1"].to_numpy(np.float32)).all()
+
+
 def test_cli_predict_parts_writes_single_structure_decomposition(monkeypatch, tmp_path):
     outroot = tmp_path / "1A1U_fdewet_parts"
     monkeypatch.setattr(