Feature/Add JAX-accelerated interaction fingerprinting

prolif/interactions/__init__.py

@@ -1,4 +1,3 @@
-# ruff: noqa: F401
 from prolif.interactions.base import (
     BasePiStacking,
     Distance,

prolif/interactions/_jax/README.md

@@ -0,0 +1,101 @@
+# JAX Interaction Fingerprinting
+
+JAX-backed interaction fingerprinting for ProLIF trajectory and single-frame workflows.
+
+## Requirements
+
+- `jax`
+- `jaxlib`
+
+CPU install:
+
+```bash
+pip install "jax[cpu]"
+```
+
+GPU install depends on your CUDA setup. Use the JAX install guide for the correct wheel.
+
+## Quick Start
+
+### Trajectory analysis
+
+```python
+import MDAnalysis as mda
+from prolif.interactions._jax import analyze_trajectory
+
+u = mda.Universe("topology.pdb", "trajectory.xtc")
+
+results = analyze_trajectory(
+    u,
+    ligand_selection="resname LIG",
+    protein_selection="protein",
+    device="cpu",  # or "gpu"
+)
+
+df = results.to_dataframe()
+```
+
+### Single frame analysis
+
+```python
+from prolif.interactions._jax import analyze_frame
+
+frame_result = analyze_frame(ligand_mol, protein_mol, cutoff=6.0)
+```
+
+## Main API
+
+```python
+analyze_trajectory(
+    universe,
+    ligand_selection="resname LIG",
+    protein_selection="protein",
+    *,
+    cutoff=6.0,
+    max_frames=None,
+    device="cpu",
+    chunk_size=None,
+    residue_mode="all",
+    scan_stride=1,
+)
+```
+
+Key parameters:
+
+- `device`: `"cpu"` or `"gpu"`
+- `chunk_size`: frames per batch; if `None`, CPU uses `256`, GPU uses auto size capped at `256`
+- `residue_mode`:
+  - `"all"`: select residues seen near ligand across trajectory
+  - `"first"`: select residues from first frame only
+- `scan_stride`: only used with `residue_mode="all"`
+
+Returns `InteractionResult` with:
+
+- `interactions`: `dict[str, np.ndarray]` of shape `(n_frames, n_residues)`
+- `residue_ids`
+- `n_frames`
+- `n_residues`
+
+## Interaction Types
+
+The high-level trajectory API returns these nine interaction maps:
+
+- `Hydrophobic`
+- `Cationic`
+- `Anionic`
+- `VdWContact`
+- `HBAcceptor`
+- `HBDonor`
+- `PiStacking`
+- `CationPi`
+- `PiCation`
+
+## Notes
+
+- First call may be slower due to JIT compilation.
+- If GPU execution fails due memory pressure, set a smaller `chunk_size` (for example `128`, `64`, or `32`).
+- Check availability with:
+
+```python
+from prolif.interactions._jax import JAX_AVAILABLE
+```

prolif/interactions/_jax/__init__.py

@@ -0,0 +1,66 @@
+"""Public JAX interaction API exports."""
+
+try:
+    import jax
+
+    JAX_AVAILABLE = True
+except ImportError:
+    JAX_AVAILABLE = False
+
+if JAX_AVAILABLE:
+    # High-level API (recommended entry points)
+    from .api import InteractionResult, analyze_frame, analyze_trajectory
+    from .framebatch import (
+        build_actor_masks,
+        build_angle_indices,
+        build_ring_cation_indices,
+        build_vdw_radii,
+        calculate_chunk_size,
+        cationpi_frames,
+        chunked_has_interactions_frames,
+        estimate_memory_per_frame,
+        get_gpu_device,
+        get_gpu_memory_info,
+        has_interactions_frames,
+        hbacceptor_frames,
+        hbdonor_frames,
+        pairwise_distances_frames,
+        pistacking_frames,
+        prepare_for_device,
+        xbacceptor_frames,
+        xbdonor_frames,
+    )
+    from .integration import compute_distances_batch, has_interaction_batch
+
+    # Low-level primitives
+    from .primitives import pairwise_distances
+
+    __all__ = [
+        "JAX_AVAILABLE",
+        "InteractionResult",
+        "analyze_frame",
+        "analyze_trajectory",
+        "build_actor_masks",
+        "build_angle_indices",
+        "build_ring_cation_indices",
+        "build_vdw_radii",
+        "calculate_chunk_size",
+        "cationpi_frames",
+        "chunked_has_interactions_frames",
+        "compute_distances_batch",
+        "estimate_memory_per_frame",
+        "get_gpu_device",
+        "get_gpu_memory_info",
+        "has_interaction_batch",
+        "has_interactions_frames",
+        "hbacceptor_frames",
+        "hbdonor_frames",
+        "pairwise_distances",
+        "pairwise_distances_frames",
+        "pistacking_frames",
+        "prepare_for_device",
+        "xbacceptor_frames",
+        "xbdonor_frames",
+    ]
+else:
+    __all__ = ["JAX_AVAILABLE"]