cavefiller 0.2.0__py3-none-any.whl

cavefiller/__init__.py

@@ -0,0 +1,3 @@
+"""CaveFiller - A tool to find and fill protein cavities with water molecules."""
+
+__version__ = "0.2.0"

cavefiller/cavity_finder.py

@@ -0,0 +1,114 @@
+"""Cavity detection using pyKVFinder."""
+
+import os
+from typing import List, Dict, Tuple, Any
+import numpy as np
+
+# Grid spacing for cavity detection (in Angstroms)
+DEFAULT_GRID_STEP = 0.6
+
+
+def find_cavities(
+    protein_file: str,
+    probe_in: float = 1.4,
+    probe_out: float = 4.0,
+    volume_cutoff: float = 5.0,
+    output_dir: str = "./output",
+) -> Tuple[List[Dict[str, Any]], Any]:
+    """
+    Find cavities in a protein structure using pyKVFinder.
+    
+    Args:
+        protein_file: Path to the protein PDB file
+        probe_in: Probe In radius for cavity detection (Å)
+        probe_out: Probe Out radius for cavity detection (Å)
+        volume_cutoff: Minimum cavity volume to consider (Ų)
+        output_dir: Directory to save cavity detection results
+        
+    Returns:
+        Tuple of (list of cavity dictionaries, cavity_data object)
+    """
+    try:
+        import pyKVFinder
+    except ImportError:
+        raise ImportError(
+            "pyKVFinder is not installed. Please install it with: pip install pykvfinder"
+        )
+    
+    # Run KVFinder to detect cavities
+    cavity_data = pyKVFinder.run_workflow(
+        input=protein_file,
+        probe_in=probe_in,
+        probe_out=probe_out,
+        step=DEFAULT_GRID_STEP,  # Grid step size
+        volume_cutoff=volume_cutoff,
+    )
+    
+    # Extract cavity information
+    cavities = []
+    
+    # Get cavity volumes and areas
+    if hasattr(cavity_data, 'volume') and cavity_data.volume is not None:
+        volumes = cavity_data.volume
+        areas = cavity_data.area if hasattr(cavity_data, 'area') else {}
+        
+        # Create mapping from string IDs to integer IDs
+        # KVFinder uses string IDs like 'KAA', 'KAB', etc., but the grid uses integers
+        cavity_id_map = {}
+        for idx, (cavity_str_id, volume) in enumerate(volumes.items(), start=1):
+            cavity_id_map[cavity_str_id] = idx
+        
+        # Process each cavity
+        for cavity_str_id, volume in volumes.items():
+            if volume >= volume_cutoff:
+                cavity_info = {
+                    "id": cavity_id_map[cavity_str_id],
+                    "string_id": cavity_str_id,
+                    "volume": volume,
+                    "area": areas.get(cavity_str_id, 0.0) if areas else 0.0,
+                }
+                cavities.append(cavity_info)
+    
+    # Sort cavities by volume (largest first)
+    cavities.sort(key=lambda x: x["volume"], reverse=True)
+    
+    return cavities, cavity_data
+
+
+def get_cavity_grid_points(cavity_data: Any, cavity_id: int) -> np.ndarray:
+    """
+    Get the grid points that belong to a specific cavity.
+    
+    Args:
+        cavity_data: The cavity data object from pyKVFinder
+        cavity_id: Integer ID of the cavity (1-indexed)
+        
+    Returns:
+        Array of (x, y, z) coordinates for the cavity grid points
+    """
+    if not hasattr(cavity_data, 'cavities') or cavity_data.cavities is None:
+        return np.array([])
+    
+    # Get cavity grid
+    cavity_grid = cavity_data.cavities
+    
+    # Find all points belonging to this cavity
+    # Note: KVFinder uses 1-indexed cavity IDs in the grid
+    points = np.argwhere(cavity_grid == cavity_id)
+    
+    # Convert grid indices to real coordinates if origin metadata is available.
+    # Different pyKVFinder versions expose metadata on either cavity_data or cavity_data.surface.
+    step = getattr(cavity_data, "step", DEFAULT_GRID_STEP)
+    origin = None
+
+    if hasattr(cavity_data, "surface") and hasattr(cavity_data.surface, "P1"):
+        origin = np.array([cavity_data.surface.P1[i] for i in range(3)], dtype=float)
+    elif hasattr(cavity_data, "P1"):
+        origin = np.array([cavity_data.P1[i] for i in range(3)], dtype=float)
+
+    points = points.astype(float)
+    if origin is not None:
+        return origin + points * float(step)
+
+    # Fallback: return index-space points; downstream code will align to protein frame.
+    return points

cavefiller/cavity_selector.py

@@ -0,0 +1,120 @@
+"""Interactive cavity selection interface."""
+
+from typing import List, Dict, Any, Optional
+import sys
+
+
+def select_cavities(cavities: List[Dict[str, Any]], prompt_for_waters: bool = True) -> tuple:
+    """
+    Allow user to select which cavities to fill and how many waters per cavity.
+    
+    Args:
+        cavities: List of cavity dictionaries with id, volume, and area
+        prompt_for_waters: Whether to prompt for number of waters per cavity
+        
+    Returns:
+        Tuple of (selected cavity dictionaries, waters_per_cavity dict)
+    """
+    print("\n" + "=" * 60)
+    print("Available Cavities:")
+    print("=" * 60)
+    print(f"{'ID':<6} {'Volume (Ų)':<15} {'Area (ų)':<15}")
+    print("-" * 60)
+    
+    for cavity in cavities:
+        print(
+            f"{cavity['id']:<6} {cavity['volume']:<15.2f} {cavity['area']:<15.2f}"
+        )
+    
+    print("=" * 60)
+    print("\nEnter cavity IDs to fill (comma-separated, e.g., '1,2,3')")
+    print("Or enter 'all' to select all cavities")
+    print("Or enter 'q' to quit")
+    
+    selected_cavities = None
+    while selected_cavities is None:
+        try:
+            user_input = input("\nYour selection: ").strip().lower()
+            
+            if user_input == 'q':
+                print("Selection cancelled.")
+                return [], {}
+            
+            if user_input == 'all':
+                print(f"Selected all {len(cavities)} cavities")
+                selected_cavities = cavities
+            else:
+                # Parse comma-separated IDs
+                selected_ids = [int(x.strip()) for x in user_input.split(",")]
+                
+                # Validate IDs
+                valid_ids = {c["id"] for c in cavities}
+                invalid_ids = [sid for sid in selected_ids if sid not in valid_ids]
+                
+                if invalid_ids:
+                    print(f"Invalid cavity IDs: {invalid_ids}")
+                    print(f"Valid IDs are: {sorted(valid_ids)}")
+                    continue
+                
+                # Get selected cavities
+                selected_cavities = [c for c in cavities if c["id"] in selected_ids]
+                
+                if selected_cavities:
+                    print(f"\nSelected {len(selected_cavities)} cavities: {[c['id'] for c in selected_cavities]}")
+                else:
+                    print("No cavities selected. Please try again.")
+                    selected_cavities = None
+                    
+        except ValueError as e:
+            print(f"Invalid input: {e}")
+            print("Please enter comma-separated numbers, 'all', or 'q'")
+        except EOFError:
+            # Handle case when stdin is not available (e.g., in tests)
+            print("\nNo input available. Selecting all cavities by default.")
+            selected_cavities = cavities
+        except KeyboardInterrupt:
+            print("\n\nSelection cancelled.")
+            return [], {}
+    
+    # Prompt for number of waters per cavity
+    waters_per_cavity = {}
+    if prompt_for_waters and selected_cavities:
+        print("\n" + "=" * 60)
+        print("Specify number of water molecules per cavity")
+        print("=" * 60)
+        
+        for cavity in selected_cavities:
+            # Default estimate based on volume
+            default_waters = max(1, int(cavity['volume'] / 30))
+            
+            while True:
+                try:
+                    prompt = f"Cavity {cavity['id']} (volume: {cavity['volume']:.2f} Ų) - waters [default: {default_waters}]: "
+                    user_input = input(prompt).strip()
+                    
+                    if user_input == '':
+                        waters_per_cavity[cavity['id']] = default_waters
+                        break
+                    else:
+                        n_waters = int(user_input)
+                        if n_waters < 0:
+                            print("  Error: Number of waters must be non-negative")
+                            continue
+                        waters_per_cavity[cavity['id']] = n_waters
+                        break
+                        
+                except ValueError:
+                    print("  Error: Please enter a valid number")
+                except EOFError:
+                    # Use default
+                    waters_per_cavity[cavity['id']] = default_waters
+                    break
+                except KeyboardInterrupt:
+                    print("\n\nCancelled. Using defaults for remaining cavities.")
+                    for remaining_cavity in selected_cavities:
+                        if remaining_cavity['id'] not in waters_per_cavity:
+                            default = max(1, int(remaining_cavity['volume'] / 30))
+                            waters_per_cavity[remaining_cavity['id']] = default
+                    return selected_cavities, waters_per_cavity
+    
+    return selected_cavities, waters_per_cavity

cavefiller/cli.py

@@ -0,0 +1,145 @@
+"""Command-line interface for CaveFiller using Typer."""
+
+import typer
+from pathlib import Path
+from typing import Optional, List
+from cavefiller.cavity_finder import find_cavities
+from cavefiller.cavity_selector import select_cavities
+from cavefiller.water_filler import fill_cavities_with_water
+
+app = typer.Typer(help="CaveFiller - Find and fill protein cavities with water molecules")
+
+
+@app.command()
+def run(
+    protein_file: Path = typer.Argument(
+        ...,
+        exists=True,
+        help="Path to the protein PDB file",
+    ),
+    output_dir: Path = typer.Option(
+        Path("./output"),
+        help="Directory to save output files",
+    ),
+    probe_in: float = typer.Option(
+        1.4,
+        help="Probe In radius for cavity detection (Å)",
+    ),
+    probe_out: float = typer.Option(
+        4.0,
+        help="Probe Out radius for cavity detection (Å)",
+    ),
+    volume_cutoff: float = typer.Option(
+        5.0,
+        help="Minimum cavity volume to consider (ų)",
+    ),
+    auto_select: bool = typer.Option(
+        False,
+        help="Automatically select all cavities (no user interaction)",
+    ),
+    cavity_ids: Optional[str] = typer.Option(
+        None,
+        help="Comma-separated list of cavity IDs to fill (e.g., '1,2,3'). If not provided, user will be prompted.",
+    ),
+    waters_per_cavity: Optional[str] = typer.Option(
+        None,
+        help="Comma-separated list of water counts per cavity (e.g., '10,15,20'). Must match cavity_ids order.",
+    ),
+    optimize_mmff94: bool = typer.Option(
+        True,
+        "--optimize-mmff94/--no-optimize-mmff94",
+        help="Run MMFF94 after placement with protein atoms fixed and waters movable.",
+    ),
+    mmff_max_iterations: int = typer.Option(
+        300,
+        help="Maximum MMFF94 iterations when optimization is enabled.",
+    ),
+):
+    """
+    Find cavities in a protein and fill them with explicit water molecules.
+
+    This tool performs the following steps:
+    1. Detects cavities in the protein using KVFinder
+    2. Allows user to select which cavities to fill
+    3. Uses cavity-grid Monte Carlo sampling to place water oxygens
+    4. Optionally runs MMFF94 with protein fixed and waters movable
+    5. Builds explicit RDKit H-O-H waters and writes a combined PDB
+    """
+    typer.echo(f"🔍 Analyzing protein: {protein_file}")
+    
+    # Create output directory
+    output_dir.mkdir(parents=True, exist_ok=True)
+    
+    # Step 1: Find cavities using pykvfinder
+    typer.echo("Step 1: Finding cavities with KVFinder...")
+    cavities, cavity_data = find_cavities(
+        str(protein_file),
+        probe_in=probe_in,
+        probe_out=probe_out,
+        volume_cutoff=volume_cutoff,
+        output_dir=str(output_dir),
+    )
+    
+    if not cavities:
+        typer.echo("❌ No cavities found in the protein.", err=True)
+        raise typer.Exit(code=1)
+    
+    typer.echo(f"✅ Found {len(cavities)} cavities")
+    
+    # Step 2: Select cavities to fill
+    typer.echo("\nStep 2: Selecting cavities to fill...")
+    
+    waters_dict = {}
+    
+    if cavity_ids:
+        # Parse cavity IDs from command line
+        selected_ids = [int(x.strip()) for x in cavity_ids.split(",")]
+        selected_cavities = [c for c in cavities if c["id"] in selected_ids]
+        if not selected_cavities:
+            typer.echo(f"❌ No cavities found with IDs: {cavity_ids}", err=True)
+            raise typer.Exit(code=1)
+        
+        # Parse waters per cavity if provided
+        if waters_per_cavity:
+            water_counts = [int(x.strip()) for x in waters_per_cavity.split(",")]
+            if len(water_counts) != len(selected_ids):
+                typer.echo("❌ Number of water counts must match number of cavity IDs", err=True)
+                raise typer.Exit(code=1)
+            waters_dict = dict(zip(selected_ids, water_counts))
+            
+    elif auto_select:
+        # Auto-select all cavities
+        selected_cavities = cavities
+        typer.echo(f"Auto-selecting all {len(cavities)} cavities")
+        # Use default water counts
+        for cavity in selected_cavities:
+            waters_dict[cavity['id']] = max(1, int(cavity['volume'] / 30))
+    else:
+        # Interactive selection
+        selected_cavities, waters_dict = select_cavities(cavities, prompt_for_waters=True)
+    
+    if not selected_cavities:
+        typer.echo("❌ No cavities selected.", err=True)
+        raise typer.Exit(code=1)
+    
+    typer.echo(f"✅ Selected {len(selected_cavities)} cavities")
+    
+    # Step 3: Fill cavities with water
+    typer.echo("\nStep 3: Filling cavities with water using Monte Carlo sampling...")
+    typer.echo("         (with clash detection, optional fixed-protein MMFF94, and RDKit HOH generation)")
+    output_file = fill_cavities_with_water(
+        str(protein_file),
+        selected_cavities,
+        cavity_data,
+        str(output_dir),
+        waters_per_cavity=waters_dict,
+        optimize_mmff94=optimize_mmff94,
+        mmff_max_iterations=mmff_max_iterations,
+    )
+    
+    typer.echo(f"\n✅ Success! Output saved to: {output_file}")
+    typer.echo("\n🎉 CaveFiller completed successfully!")
+
+
+if __name__ == "__main__":
+    app()