Creating a Dockerfile

.dockerignore

@@ -0,0 +1,88 @@
+# =============================================================================
+# WaterFlow Docker Ignore
+# Exclude files that shouldn't be in the Docker build context
+# =============================================================================
+
+# Git
+.git
+.gitignore
+.gitattributes
+
+# Python
+__pycache__
+*.py[cod]
+*$py.class
+*.so
+.Python
+.venv
+venv
+ENV
+env
+.eggs
+*.egg-info
+*.egg
+.mypy_cache
+.pytest_cache
+.ruff_cache
+.coverage
+htmlcov
+
+# IDE
+.idea
+.vscode
+*.swp
+*.swo
+*~
+
+# Jupyter
+.ipynb_checkpoints
+notebooks/
+*.ipynb
+
+# Data and model files (mount these as volumes instead)
+data/
+*.pt
+*.ckpt
+*.pth
+*.safetensors
+*.h5
+*.hdf5
+*.pkl
+*.pickle
+
+# W&B
+wandb/
+*.wandb
+
+# Logs
+logs/
+*.log
+
+# Docker
+Dockerfile*
+docker-compose*.yml
+.docker
+
+# Documentation
+docs/
+*.md
+!README.md
+
+# Tests (not needed in production image)
+tests/
+test_*.py
+*_test.py
+conftest.py
+
+# CI/CD
+.github/
+.gitlab-ci.yml
+.travis.yml
+Makefile
+
+# Misc
+.DS_Store
+Thumbs.db
+*.bak
+*.tmp
+*.temp

Dockerfile

@@ -0,0 +1,100 @@
+# =============================================================================
+# WaterFlow Docker image for GPU workflows (CUDA 12.6)
+# =============================================================================
+
+FROM nvidia/cuda:12.6.3-devel-ubuntu22.04
+
+# Prevent interactive prompts during package installation
+ENV DEBIAN_FRONTEND=noninteractive
+
+# Install system dependencies
+# Note: rm -rf /var/lib/apt/lists/* removes apt cache to reduce image size (~30MB savings)
+RUN apt-get update && apt-get install -y --no-install-recommends \
+    software-properties-common \
+    curl \
+    git \
+    build-essential \
+    && add-apt-repository ppa:deadsnakes/ppa \
+    && apt-get update \
+    && apt-get install -y --no-install-recommends \
+    python3.12 \
+    python3.12-dev \
+    python3.12-venv \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set Python 3.12 as default
+RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.12 1 \
+    && update-alternatives --install /usr/bin/python3 python3 /usr/bin/python3.12 1
+
+# Install uv package manager (pinned for reproducible builds)
+COPY --from=ghcr.io/astral-sh/uv:0.7.3 /uv /usr/local/bin/uv
+
+# Set working directory
+WORKDIR /app
+
+# Copy dependency files first for better layer caching
+COPY pyproject.toml uv.lock ./
+
+# Create virtual environment and install dependencies
+ENV UV_COMPILE_BYTECODE=1
+ENV UV_LINK_MODE=copy
+RUN uv sync --frozen --no-install-project
+
+# Copy source code
+COPY src/ ./src/
+COPY scripts/ ./scripts/
+
+# Install the project itself
+RUN uv sync --frozen
+
+# Pre-download ESM3 model to bake it into the image.
+# The generate script loads esm3-open without authentication, so no HF token is needed.
+ENV HF_HOME=/app/.cache/huggingface
+RUN . .venv/bin/activate && python -c "\
+from esm.models.esm3 import ESM3; \
+ESM3.from_pretrained('esm3-open'); \
+print('ESM3 model downloaded successfully')"
+
+# Compile Python bytecode for faster startup
+RUN . .venv/bin/activate && python -m compileall -q src/ scripts/
+
+# Verify core GPU and preprocessing imports work
+RUN . .venv/bin/activate && python <<'EOF'
+import torch
+print(f'PyTorch {torch.__version__}, CUDA {torch.version.cuda}')
+from torch_scatter import scatter_add
+print('torch-scatter OK')
+from torch_cluster import radius_graph
+print('torch-cluster OK')
+import pymol2
+print('pymol2 OK')
+EOF
+
+# Copy entrypoint script
+COPY docker/entrypoint.sh /app/entrypoint.sh
+RUN chmod +x /app/entrypoint.sh
+
+# Create mount points for data volumes
+RUN mkdir -p /data/pdb /data/cache /data/checkpoints /data/outputs /data/logs /data/splits
+
+# Environment variables
+ENV VIRTUAL_ENV=/app/.venv
+ENV PATH="/app/.venv/bin:$PATH"
+ENV PYTHONPATH=/app
+ENV PYTHONDONTWRITEBYTECODE=1
+ENV PYTHONUNBUFFERED=1
+
+# CUDA configuration for H100 GPUs
+ENV CUDA_HOME=/usr/local/cuda
+ENV TORCH_CUDA_ARCH_LIST="9.0"
+
+# Default data paths (can be overridden via docker run -e)
+ENV WATERFLOW_PDB_DIR=/data/pdb
+ENV WATERFLOW_CACHE_DIR=/data/cache
+ENV WATERFLOW_CHECKPOINT_DIR=/data/checkpoints
+ENV WATERFLOW_OUTPUT_DIR=/data/outputs
+ENV WATERFLOW_LOG_DIR=/data/logs
+ENV WATERFLOW_SPLITS_DIR=/data/splits
+
+ENTRYPOINT ["/app/entrypoint.sh"]
+CMD ["--help"]

README.md

@@ -241,7 +241,6 @@ To resume training from a checkpoint, you can load the model weights and optimiz
 | `--save_dir` | `../flow_checkpoints` | Directory to save checkpoints |
 | `--save_every` | `10` | Save checkpoint every N epochs |
 | `--eval_every` | `5` | Run evaluation every N epochs |
-| `--edia_dir` | (none) | Root directory for EDIA CSV files |
 | `--min_edia` | `0.4` | Minimum EDIA score threshold for waters |
 | `--no_filter_by_edia` | - | Disable EDIA-based water filtering |
 
@@ -286,12 +285,10 @@ These filters remove individual low-quality waters (can be toggled):
 EDIA measures how well an atom's position is supported by the experimental electron density map. Higher EDIA scores indicate more reliable atomic positions.
 
 **Configuration:**
-- EDIA filtering is enabled by default but only activates if `--edia_dir` is provided
-- If `--edia_dir` is not set, EDIA filtering is skipped (with a warning logged)
+- EDIA filtering is enabled by default 
+- The EDIA data lives in the `json` file of the format `<pdb_id>_final.json` in the same directory as the `pdb` file, and is obtained from PDB-REDO.
 - Use `--no_filter_by_edia` to explicitly disable EDIA filtering
 
-**Directory structure:** `{edia_dir}/{pdb_id}/{pdb_id}_residue_stats.csv`
-
 </details>
 
 ## Inference

docker-compose.yml

@@ -0,0 +1,71 @@
+# =============================================================================
+# WaterFlow Docker Compose Configuration
+# Local development with GPU passthrough
+# =============================================================================
+
+services:
+  waterflow:
+    build:
+      context: .
+      dockerfile: Dockerfile
+    image: waterflow:latest
+    runtime: nvidia
+
+    # GPU configuration - expose all GPUs
+    deploy:
+      resources:
+        reservations:
+          devices:
+            - driver: nvidia
+              count: all
+              capabilities: [gpu]
+
+    # Shared memory size for PyTorch DataLoader
+    shm_size: 16gb
+
+    # Environment variables
+    environment:
+      - NVIDIA_VISIBLE_DEVICES=all
+      - WANDB_API_KEY=${WANDB_API_KEY:-}
+      - WANDB_PROJECT=${WANDB_PROJECT:-waterflow}
+      - WATERFLOW_PDB_DIR=/data/pdb
+      - WATERFLOW_CACHE_DIR=/data/cache
+      - WATERFLOW_CHECKPOINT_DIR=/data/checkpoints
+      - WATERFLOW_OUTPUT_DIR=/data/outputs
+      - WATERFLOW_LOG_DIR=/data/logs
+      - WATERFLOW_SPLITS_DIR=/data/splits
+
+    # Volume mounts
+    volumes:
+      # Input PDB files (read-only)
+      - ${PDB_DIR:-./data/pdb}:/data/pdb:ro
+      # Cache for embeddings and preprocessed data
+      - ${CACHE_DIR:-./data/cache}:/data/cache
+      # Model checkpoints
+      - ${CHECKPOINT_DIR:-./data/checkpoints}:/data/checkpoints
+      # Inference outputs
+      - ${OUTPUT_DIR:-./data/outputs}:/data/outputs
+      # W&B and training logs
+      - ${LOG_DIR:-./data/logs}:/data/logs
+      # Train/val/test split files
+      - ${SPLITS_DIR:-./splits}:/data/splits:ro
+
+    # Interactive mode for development
+    stdin_open: true
+    tty: true
+
+  # Development profile with source code mounted
+  waterflow-dev:
+    extends:
+      service: waterflow
+    volumes:
+      # Mount source for live development
+      - ./src:/app/src:ro
+      - ./scripts:/app/scripts:ro
+      # Data volumes
+      - ${PDB_DIR:-./data/pdb}:/data/pdb:ro
+      - ${CACHE_DIR:-./data/cache}:/data/cache
+      - ${CHECKPOINT_DIR:-./data/checkpoints}:/data/checkpoints
+      - ${OUTPUT_DIR:-./data/outputs}:/data/outputs
+      - ${LOG_DIR:-./data/logs}:/data/logs
+      - ${SPLITS_DIR:-./splits}:/data/splits:ro

docker/entrypoint.sh

@@ -0,0 +1,108 @@
+#!/bin/bash
+# =============================================================================
+# WaterFlow Docker Entrypoint
+# Unified CLI that dispatches commands to the appropriate scripts
+# =============================================================================
+
+set -e
+
+# Show help if no arguments
+show_help() {
+    cat << EOF
+WaterFlow Docker Container
+
+Usage: docker run waterflow <command> [options]
+
+Commands:
+  train           Run training pipeline
+  inference       Run inference on PDB files
+  generate-esm    Generate ESM3 embeddings for PDB files
+  python          Run arbitrary Python script or command
+  --help          Show this help message
+
+For interactive shell: docker run --entrypoint /bin/bash -it waterflow
+
+Examples:
+  # Training
+  docker run --gpus all -v /path/to/data:/data waterflow train \\
+      --train_list /data/splits/train.txt \\
+      --val_list /data/splits/val.txt \\
+      --encoder_type esm \\
+      --epochs 200
+
+  # Generate ESM embeddings
+  docker run --gpus all -v /path/to/data:/data waterflow generate-esm \\
+      --split_file /data/splits/train.txt
+
+  # Inference
+  docker run --gpus all -v /path/to/data:/data waterflow inference \\
+      --run_dir /data/checkpoints/run_name \\
+      --pdb_list /data/splits/test.txt \\
+      --output_dir /data/outputs
+
+Environment Variables:
+  WATERFLOW_PDB_DIR         Base directory for PDB files (default: /data/pdb)
+  WATERFLOW_CACHE_DIR       Cache directory for embeddings (default: /data/cache)
+  WATERFLOW_CHECKPOINT_DIR  Checkpoint directory (default: /data/checkpoints)
+  WATERFLOW_OUTPUT_DIR      Output directory (default: /data/outputs)
+  WATERFLOW_LOG_DIR         Log directory (default: /data/logs)
+  WATERFLOW_SPLITS_DIR      Split/list files directory (default: /data/splits)
+  WANDB_API_KEY             Weights & Biases API key (set via -e or docker-compose)
+  WANDB_PROJECT             W&B project name (set via -e or docker-compose)
+
+EOF
+}
+
+# Get the command
+COMMAND="${1:-}"
+
+case "$COMMAND" in
+    train)
+        shift
+        exec python /app/scripts/train.py \
+            --base_pdb_dir "${WATERFLOW_PDB_DIR}" \
+            --processed_dir "${WATERFLOW_CACHE_DIR}" \
+            --save_dir "${WATERFLOW_CHECKPOINT_DIR}" \
+            --wandb_dir "${WATERFLOW_LOG_DIR}" \
+            "$@"
+        ;;
+
+    inference)
+        shift
+        exec python /app/scripts/inference.py \
+            --base_pdb_dir "${WATERFLOW_PDB_DIR}" \
+            --processed_dir "${WATERFLOW_CACHE_DIR}" \
+            --output_dir "${WATERFLOW_OUTPUT_DIR}" \
+            "$@"
+        ;;
+
+    generate-esm)
+        shift
+        exec python /app/scripts/generate_esm_embeddings.py \
+            --base_pdb_dir "${WATERFLOW_PDB_DIR}" \
+            --cache_dir "${WATERFLOW_CACHE_DIR}" \
+            "$@"
+        ;;
+
+    python)
+        shift
+        exec python "$@"
+        ;;
+
+    --help|-h|help)
+        show_help
+        exit 0
+        ;;
+
+    "")
+        show_help
+        exit 0
+        ;;
+
+    *)
+        echo "Error: Unknown command '$COMMAND'"
+        echo ""
+        show_help
+        exit 1
+        ;;
+esac