diff --git a/.gitignore b/.gitignore index c1abc0d..3226e41 100644 --- a/.gitignore +++ b/.gitignore @@ -151,3 +151,6 @@ cython_debug/ # and can be added to the global gitignore or merged into this file. For a more nuclear # option (not recommended) you can uncomment the following to ignore the entire idea folder. #.idea/ + +# Ignore apptainers +*.sif diff --git a/README.md b/README.md index d996570..6ba254b 100644 --- a/README.md +++ b/README.md @@ -20,8 +20,9 @@ You can find the environment file under `env` folder. ### Apptainer +Run this from this root directory: ``` -apptainer build -s --nv ${CONTAINER_PATH} ./env/Apptainer +apptainer build --nv SpatialPPIv2.sif ./env/Apptainer ``` ### Docker @@ -33,25 +34,12 @@ docker build -t ${IMAGE_NAME} ./env ### Conda ``` -conda create -n pytorch221 python=3.11 -conda activate pytorch221 -conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia -pip install numpy pandas seaborn tensorboard -pip install torch_geometric==2.5.1 -pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.2.0+cu118.html -conda install lightning -c conda-forge -pip install lightning[extra] -pip install matplotlib==3.8.3 -pip install biopandas==0.5.1 -pip install biopython==1.83 -pip install transformers==4.40.2 -pip install sentencepiece==0.2.0 -pip install torchsummary==1.5.1 -pip install scipy==1.12.0 -pip install torch_cluster -f https://data.pyg.org/whl/torch-2.2.0+cu118.html -pip install git+https://github.com/yusuf1759/prodigy-cryst.git -pip install pinder[all] -pip install fair-esm +conda env create -f env/environment.yaml +``` + +or if you have mamba (should be faster) +``` +mamba env create -f env/environment.yaml ``` ## Inference protein interactions @@ -79,6 +67,31 @@ When there is no suitable protein structure file, `fasta` files can be used to m python inference.py --A demo/D3INY1.fasta --B demo/P62593.fasta ``` +### Batch inference: one query protein vs a folder + +You can screen one query protein against all supported files in a folder and save interaction probabilities to a CSV file: + +``` +python inference_batch.py --query demo/P33895.pdb --targets demo --output results.csv +``` + +Options: + +- `--query`: Query protein file (`.pdb`/`.cif` or `.fasta`/`.fa`/`.faa`/`.fna`) +- `--targets`: Folder containing target protein files +- `--output`: Output CSV path. Default is `inference_batch.csv` +- `--chain_query`: Chain ID for query structure files. Default is `first` +- `--chain_targets`: Chain ID for target structure files. Default is `first` +- `--model`: Model for structure inputs (`ProtT5` or `ESM-2+ac`). For FASTA query, `ESM-2+ac` is forced +- `--device`: Either `cuda` or `cpu` + +Notes: + +- Query and target files must be the same modality: + - FASTA query requires all FASTA targets + - Structure query (`pdb`/`cif`) requires all structure targets +- The output CSV contains columns: `query`, `target`, `probability`, `status`, `error` + ### Visualization of inference process Please check the [notebook](https://github.com/ohuelab/SpatialPPIv2/blob/main/demo/example_visualize.ipynb). diff --git a/env/Apptainer b/env/Apptainer index 738b517..e3e25a4 100644 --- a/env/Apptainer +++ b/env/Apptainer @@ -1,23 +1,41 @@ Bootstrap: docker -From: pytorch/pytorch:2.2.1-cuda11.8-cudnn8-devel +From: nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04 + +%files + env/environment.yaml environment.yaml %post - apt-get update && apt-get install -y git - - pip install -U pip - pip install numpy pandas seaborn tensorboard - pip install torch_geometric==2.5.1 - pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.2.0+cu118.html - conda install lightning -c conda-forge - pip install lightning[extra] - pip install matplotlib==3.8.3 - pip install biopandas==0.5.1 - pip install biopython==1.83 - pip install transformers==4.40.2 - pip install sentencepiece==0.2.0 - pip install torchsummary==1.5.1 - pip install scipy==1.12.0 - pip install torch_cluster -f https://data.pyg.org/whl/torch-2.2.0+cu118.html - pip install git+https://github.com/yusuf1759/prodigy-cryst.git - pip install pinder[all] - pip install fair-esm + set -e + export DEBIAN_FRONTEND=noninteractive + + apt-get update + apt-get install -y --no-install-recommends \ + bzip2 \ + ca-certificates \ + curl \ + git + rm -rf /var/lib/apt/lists/* + + curl -Ls https://micro.mamba.pm/api/micromamba/linux-64/latest \ + | tar -xvj -C /usr/local/bin --strip-components=1 bin/micromamba + + export MAMBA_ROOT_PREFIX=/opt/conda + mkdir -p "${MAMBA_ROOT_PREFIX}" + + micromamba create -y -n SpatialPPIv2 -f environment.yaml + micromamba clean -a -y + +%environment + export MAMBA_ROOT_PREFIX=/opt/conda + export MAMBA_EXE=/usr/local/bin/micromamba + export CONDA_DEFAULT_ENV=SpatialPPIv2 + export PATH=/opt/conda/envs/SpatialPPIv2/bin:/opt/conda/bin:$PATH + export CUDA_HOME=/usr/local/cuda + export LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH} + +%runscript + if [ "$#" -eq 0 ]; then + exec /bin/bash + else + exec micromamba run -n SpatialPPIv2 "$@" + fi diff --git a/env/Dockerfile b/env/Dockerfile index ed41ad1..440523c 100644 --- a/env/Dockerfile +++ b/env/Dockerfile @@ -1,20 +1,41 @@ -FROM pytorch/pytorch:2.2.1-cuda11.8-cudnn8-devel -RUN apt-get update && apt-get install -y git -# Install dependency. -RUN pip install -U pip -RUN pip install numpy pandas seaborn tensorboard -RUN pip install torch_geometric==2.5.1 -RUN pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.2.0+cu118.html -RUN conda install lightning -c conda-forge -RUN pip install lightning[extra] -RUN pip install matplotlib==3.8.3 -RUN pip install biopandas==0.5.1 -RUN pip install biopython==1.83 -RUN pip install transformers==4.40.2 -RUN pip install sentencepiece==0.2.0 -RUN pip install torchsummary==1.5.1 -RUN pip install scipy==1.12.0 -RUN pip install torch_cluster -f https://data.pyg.org/whl/torch-2.2.0+cu118.html -RUN pip install git+https://github.com/yusuf1759/prodigy-cryst.git -RUN pip install pinder[all] -RUN pip install fair-esm +FROM nvidia/cuda:12.1.1-cudnn8-devel-ubuntu22.04 + +ENV DEBIAN_FRONTEND=noninteractive +ENV MAMBA_ROOT_PREFIX=/opt/conda +ENV MAMBA_EXE=/usr/local/bin/micromamba +ENV CONDA_DEFAULT_ENV=SpatialPPIv2 +ENV CUDA_HOME=/usr/local/cuda +ENV PATH=/opt/conda/envs/SpatialPPIv2/bin:/opt/conda/bin:${PATH} +ENV LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH} + +WORKDIR /workspace + +RUN apt-get update && apt-get install -y --no-install-recommends \ + bzip2 \ + ca-certificates \ + curl \ + git \ + && rm -rf /var/lib/apt/lists/* + +RUN curl -Ls https://micro.mamba.pm/api/micromamba/linux-64/latest \ + | tar -xvj -C /usr/local/bin --strip-components=1 bin/micromamba + +COPY env/environment.yaml /tmp/environment.yaml + +RUN mkdir -p "${MAMBA_ROOT_PREFIX}" \ + && micromamba create -y -n SpatialPPIv2 -f /tmp/environment.yaml \ + && micromamba clean -a -y \ + && rm -f /tmp/environment.yaml + +RUN printf '%s\n' \ + '#!/usr/bin/env bash' \ + 'set -e' \ + 'if [ "$#" -eq 0 ]; then' \ + ' exec /bin/bash' \ + 'else' \ + ' exec micromamba run -n SpatialPPIv2 "$@"' \ + 'fi' \ + > /usr/local/bin/docker-entrypoint.sh \ + && chmod +x /usr/local/bin/docker-entrypoint.sh + +ENTRYPOINT ["/usr/local/bin/docker-entrypoint.sh"] diff --git a/env/environment.yaml b/env/environment.yaml new file mode 100644 index 0000000..5776003 --- /dev/null +++ b/env/environment.yaml @@ -0,0 +1,287 @@ +name: SpatialPPIv2 +channels: + - conda-forge + - nvidia + - pkgs/main + - pytorch +dependencies: + - _libgcc_mutex=0.1=main + - _openmp_mutex=5.1=1_gnu + - alsa-lib=1.2.15.3=hb03c661_0 + - aom=3.12.1=h7934f7d_0 + - attr=2.5.2=h39aace5_0 + - binutils_impl_linux-64=2.40=h5293946_0 + - binutils_linux-64=2.40.0=hc2dff05_2 + - blas=1.0=mkl + - brotlicffi=1.2.0.0=py311h7354ed3_0 + - bzip2=1.0.8=h5eee18b_6 + - ca-certificates=2026.1.4=hbd8a1cb_0 + - cairo=1.18.4=h44eff21_0 + - certifi=2026.01.04=py311h06a4308_0 + - cffi=2.0.0=py311h4eded50_1 + - charset-normalizer=3.4.4=py311h06a4308_0 + - cuda-cudart=12.1.105=0 + - cuda-cupti=12.1.105=0 + - cuda-libraries=12.1.0=0 + - cuda-nvrtc=12.1.105=0 + - cuda-nvtx=12.1.105=0 + - cuda-opencl=13.1.115=h4f1e1d6_0 + - cuda-runtime=12.1.0=0 + - cuda-version=13.1=hd92462c_3 + - dav1d=1.2.1=h5eee18b_0 + - dbus=1.16.2=h24cb091_1 + - expat=2.7.4=h7354ed3_0 + - ffmpeg=4.2.2=h20bf706_0 + - fftw=3.3.10=nompi_h3b011a4_111 + - filelock=3.20.3=py311h06a4308_0 + - font-ttf-dejavu-sans-mono=2.37=hab24e00_0 + - font-ttf-inconsolata=3.000=h77eed37_0 + - font-ttf-source-code-pro=2.038=h77eed37_0 + - font-ttf-ubuntu=0.83=h77eed37_3 + - fontconfig=2.15.0=h2c49b7f_0 + - fonts-conda-ecosystem=1=0 + - fonts-conda-forge=1=hc364b38_1 + - freetype=2.14.1=hf5b9546_0 + - fribidi=1.0.16=h9fb5f84_0 + - gcc_impl_linux-64=11.2.0=h1234567_1 + - gcc_linux-64=11.2.0=h5c386dc_2 + - gdk-pixbuf=2.44.4=hbd8e730_0 + - giflib=5.2.2=h5eee18b_0 + - glib=2.86.3=h617169b_0 + - glib-tools=2.86.3=h9fad118_0 + - glslang=16.2.0=h96af755_1 + - gmp=6.3.0=h6a678d5_0 + - gmpy2=2.2.2=py311ha78e65c_0 + - gnutls=3.6.15=he1e5248_0 + - graphite2=1.3.14=h295c915_1 + - gstreamer-orc=0.4.42=h82d0256_0 + - gxx_impl_linux-64=11.2.0=h1234567_1 + - gxx_linux-64=11.2.0=hc2dff05_2 + - harfbuzz=10.4.0=h76408a6_0 + - icu=75.1=he02047a_0 + - idna=3.11=py311h06a4308_0 + - intel-gmmlib=22.9.0=hb700be7_0 + - intel-media-driver=25.3.4=hecca717_0 + - intel-openmp=2023.0.0=h9e868ea_25371 + - jack=1.9.22=hf4617a5_3 + - jinja2=3.1.6=py311h06a4308_0 + - jpeg=9f=h5ce9db8_0 + - kernel-headers_linux-64=4.18.0=h3108a97_1 + - lame=3.100=hbd0596d_1 + - lcms2=2.17=heab6991_0 + - ld_impl_linux-64=2.40=h12ee557_0 + - leptonica=1.82.0=hfdeec58_3 + - lerc=4.0.0=h6a678d5_0 + - libarchive=3.8.2=h3ec8f01_0 + - libass=0.17.4=h6f62b05_1 + - libavif=1.3.0=h3539ee5_0 + - libcap=2.77=h3ff7636_0 + - libcublas=12.1.0.26=0 + - libcufft=11.0.2.4=0 + - libcufile=1.16.1.26=h3b4bcfc_0 + - libcurand=10.4.1.81=h1b6c897_0 + - libcusolver=11.4.4.55=0 + - libcusparse=12.0.2.55=0 + - libdeflate=1.22=h5eee18b_0 + - libdrm=2.4.125=hb03c661_1 + - libegl=1.7.0=ha4b6fd6_2 + - libexpat=2.7.4=h7354ed3_0 + - libffi=3.4.4=h6a678d5_1 + - libflac=1.5.0=he200343_1 + - libgcc=15.2.0=h69a1729_7 + - libgcc-devel_linux-64=11.2.0=h1234567_1 + - libgcc-ng=15.2.0=h166f726_7 + - libgfortran=15.2.0=h166f726_7 + - libgfortran-ng=15.2.0=h166f726_7 + - libgfortran5=15.2.0=hc633d37_7 + - libgl=1.7.0=ha4b6fd6_2 + - libglib=2.86.3=h8b17d9a_0 + - libglvnd=1.7.0=ha4b6fd6_2 + - libglx=1.7.0=ha4b6fd6_2 + - libgomp=15.2.0=h4751f2c_7 + - libiconv=1.18=h3b78370_2 + - libidn2=2.3.4=h5eee18b_0 + - libjpeg-turbo=2.0.0=h9bf148f_0 + - libltdl=2.4.3a=h5888daf_0 + - liblzma=5.8.2=hb03c661_0 + - liblzma-devel=5.8.2=hb03c661_0 + - libnpp=12.0.2.50=0 + - libnsl=2.0.0=h5eee18b_0 + - libnvjitlink=12.1.105=0 + - libnvjpeg=12.1.1.14=0 + - libogg=1.3.5=hd0c01bc_1 + - libopenblas=0.3.30=h46f56fc_2 + - libopenjpeg=2.5.4=hee96239_1 + - libopus=1.6.1=h9f10d21_0 + - libpciaccess=0.18=hb9d3cd8_0 + - libpng=1.6.54=hee55ce4_0 + - librsvg=2.58.4=h49af25d_2 + - libsndfile=1.2.2=hc7d488a_2 + - libstdcxx=15.2.0=h39759b7_7 + - libstdcxx-devel_linux-64=11.2.0=h1234567_1 + - libstdcxx-ng=15.2.0=hc03a8fd_7 + - libsystemd0=259.1=h6569c3e_0 + - libtasn1=4.21.0=h27ab2c4_0 + - libtheora=1.2.0=h32ad74f_1 + - libtiff=4.7.1=h029b1ac_0 + - libtool=2.5.4=h5888daf_0 + - libudev1=259.1=h6569c3e_0 + - libunistring=0.9.10=h27cfd23_0 + - libuuid=1.41.5=h5eee18b_0 + - libva=2.23.0=he1eb515_0 + - libvorbis=1.3.7=h54a6638_2 + - libvpl=2.15.0=h54a6638_1 + - libvpx=1.7.0=h439df22_0 + - libwebp=1.6.0=h089d785_0 + - libwebp-base=1.6.0=hb7bb969_0 + - libxcb=1.17.0=h9b100fa_0 + - libxml2=2.13.9=h04c0eec_0 + - libzlib=1.3.1=hb25bd0a_0 + - llvm-openmp=14.0.6=h9e868ea_0 + - lz4-c=1.9.4=h6a678d5_1 + - markupsafe=3.0.2=py311h5eee18b_0 + - mkl=2023.1.0=h213fc3f_46344 + - mkl-service=2.4.0=py311h5eee18b_2 + - mkl_fft=1.3.11=py311h5eee18b_0 + - mkl_random=1.2.8=py311ha02d727_0 + - mpc=1.3.1=h5eee18b_0 + - mpfr=4.2.1=h5eee18b_0 + - mpg123=1.32.9=hc50e24c_0 + - mpmath=1.3.0=py311h06a4308_0 + - ncurses=6.5=h7934f7d_0 + - nettle=3.7.3=hbbd107a_1 + - networkx=3.6.1=py311h06a4308_0 + - numpy=2.0.1=py311h08b1b3b_1 + - numpy-base=2.0.1=py311hf175353_1 + - ocl-icd=2.3.3=h47b2149_0 + - opencl-headers=2025.07.22=hfb20e49_0 + - openfst=1.8.4=h84d6215_1 + - openh264=2.1.1=h4ff587b_0 + - openjpeg=2.5.4=h55fea9a_0 + - openssl=3.6.1=h35e630c_1 + - optree=0.18.0=py311hdf67eae_0 + - packaging=25.0=py311h06a4308_1 + - pango=1.56.3=h861ebed_0 + - pcre2=10.46=hf426167_0 + - pillow=12.0.0=py311h3b88751_1 + - pip=26.0.1=pyhc872135_0 + - pixman=0.46.4=h7934f7d_0 + - pthread-stubs=0.3=h0ce48e5_1 + - pulseaudio=17.0=haebf07f_3 + - pulseaudio-client=17.0=h9a6aba3_3 + - pulseaudio-daemon=17.0=h33dcb6b_3 + - pybind11=3.0.1=pyh7a1b43c_0 + - pybind11-global=3.0.1=pyhc7ab6ef_0 + - pycparser=2.23=py311h06a4308_0 + - pysocks=1.7.1=py311h06a4308_1 + - python=3.11.14=h6fa692b_0 + - python_abi=3.11=2_cp311 + - pytorch=2.5.1=py3.11_cuda12.1_cudnn9.1.0_0 + - pytorch-cuda=12.1=ha16c6d3_6 + - pytorch-mutex=1.0=cuda + - pyyaml=6.0.3=py311h591646f_0 + - readline=8.3=hc2a1206_0 + - requests=2.32.5=py311h06a4308_1 + - setuptools=80.10.2=py311h06a4308_0 + - shaderc=2025.5=h718be3e_1 + - soxr=0.1.3=h0b41bf4_3 + - spirv-tools=2026.1=hb700be7_0 + - sqlite=3.51.1=he0a8d7e_0 + - svt-av1=3.1.2=hecca717_0 + - sympy=1.14.0=py311h06a4308_1 + - sysroot_linux-64=2.28=h3108a97_1 + - tbb=2021.8.0=hdb19cb5_0 + - tesseract=5.2.0=hb0d2e87_3 + - tk=8.6.15=h54e0aa7_0 + - torchaudio=2.5.1=py311_cu121 + - torchtriton=3.1.0=py311 + - torchvision=0.20.1=py311_cu121 + - typing-extensions=4.15.0=py311h06a4308_0 + - typing_extensions=4.15.0=py311h06a4308_0 + - tzdata=2025c=he532380_0 + - urllib3=2.6.3=py311h06a4308_0 + - wayland=1.24.0=hdac8c69_0 + - wayland-protocols=1.47=hd8ed1ab_0 + - wheel=0.46.3=py311h06a4308_0 + - x264=1!157.20191217=h7b6447c_0 + - xorg-libice=1.1.2=hb9d3cd8_0 + - xorg-libsm=1.2.6=h9b100fa_0 + - xorg-libx11=1.8.12=h9b100fa_1 + - xorg-libxau=1.0.12=h9b100fa_0 + - xorg-libxdmcp=1.1.5=h9b100fa_0 + - xorg-libxext=1.3.6=h9b100fa_0 + - xorg-libxfixes=6.0.2=hb03c661_0 + - xorg-libxi=1.8.2=hb9d3cd8_0 + - xorg-libxrender=0.9.12=h9b100fa_0 + - xorg-libxtst=1.2.5=hb9d3cd8_3 + - xorg-xorgproto=2024.1=h5eee18b_1 + - xz=5.8.2=ha02ee65_0 + - xz-gpl-tools=5.8.2=ha02ee65_0 + - xz-tools=5.8.2=hb03c661_0 + - yaml=0.2.5=h7b6447c_0 + - zlib=1.3.1=hb25bd0a_0 + - zstd=1.5.7=h11fc155_0 + - git + - pip: + - --find-links https://data.pyg.org/whl/torch-2.5.1+cu121.html + - Cython==3.0.2 + - Markdown==3.10.2 + - Werkzeug==3.1.5 + - absl-py==2.4.0 + - aiohappyeyeballs==2.6.1 + - aiohttp==3.13.3 + - aiosignal==1.4.0 + - attrs==25.4.0 + - biopandas==0.5.1 + - biopython==1.83 + - contourpy==1.3.3 + - cycler==0.12.1 + - fair-esm==2.0.0 + - fonttools==4.61.1 + - freesasa==2.2.0.post3 + - frozenlist==1.8.0 + - fsspec==2026.2.0 + - grpcio==1.78.0 + - hf-xet==1.2.0 + - huggingface_hub==0.36.2 + - joblib==1.5.3 + - kiwisolver==1.4.9 + - lightning==2.6.1 + - lightning-utilities==0.15.2 + - looseversion==1.1.2 + - matplotlib==3.8.3 + - mmtf-python==1.1.3 + - msgpack==1.1.2 + - multidict==6.7.1 + - numpy==1.26.4 + - pandas==3.0.0 + - git+https://github.com/yusuf1759/prodigy-cryst.git + - propcache==0.4.1 + - protobuf==6.33.5 + - psutil==7.2.2 + - pyg-lib==0.4.0+pt25cu121 + - pyparsing==3.3.2 + - python-dateutil==2.9.0.post0 + - pytorch-lightning==2.6.1 + - regex==2026.1.15 + - safetensors==0.7.0 + - scikit-learn==1.2.2 + - scipy==1.12.0 + - seaborn==0.13.2 + - sentencepiece==0.2.0 + - six==1.17.0 + - sympy==1.13.1 + - tensorboard==2.20.0 + - tensorboard-data-server==0.7.2 + - threadpoolctl==3.6.0 + - tokenizers==0.19.1 + - torch_cluster==1.6.3 + - torch_geometric==2.5.1 + - torch_scatter==2.1.2 + - torch_sparse==0.6.18 + - torch_spline_conv==1.2.2 + - torchmetrics==1.8.2 + - torchsummary==1.5.1 + - tqdm==4.67.3 + - transformers==4.40.2 + - yarl==1.22.0 diff --git a/inference_batch.py b/inference_batch.py new file mode 100644 index 0000000..0eea1a3 --- /dev/null +++ b/inference_batch.py @@ -0,0 +1,200 @@ +import argparse +import csv +from pathlib import Path + +import torch + +from utils.dataset import build_data, build_data_from_adj +from utils.model import getModel +from utils.tool import Embed, extractPDB, getConfig, read_fasta + + +FASTA_EXTENSIONS = {".fasta", ".fa", ".faa", ".fna"} +STRUCTURE_EXTENSIONS = {".pdb", ".cif"} + + +def is_fasta_file(path: Path) -> bool: + return path.suffix.lower() in FASTA_EXTENSIONS + + +def is_structure_file(path: Path) -> bool: + return path.suffix.lower() in STRUCTURE_EXTENSIONS + + +def is_supported_file(path: Path) -> bool: + return is_fasta_file(path) or is_structure_file(path) + + +def load_query_features(query_path: Path, chain: str, model_type: str, embedder: Embed): + if model_type == "ESM-2+ac": + seq = read_fasta(str(query_path)) + feature, adj = embedder.encode(seq, attention_contact=True) + return {"feature": feature, "adj": adj} + + seq, coords = extractPDB(str(query_path), chain) + feature = embedder.encode(seq) + return {"feature": feature, "coords": coords} + + +def infer_pair( + model, + device, + model_type: str, + query_cached: dict, + target_path: Path, + query_name: str, + chain_target: str, + embedder: Embed, +): + if model_type == "ESM-2+ac": + seq_b = read_fasta(str(target_path)) + feat_b, adj_b = embedder.encode(seq_b, attention_contact=True) + input_data = build_data_from_adj( + features=[query_cached["feature"], feat_b], + adjs=[query_cached["adj"], adj_b], + ).to(device) + else: + seq_b, coord_b = extractPDB(str(target_path), chain_target) + feat_b = embedder.encode(seq_b) + input_data = build_data( + node_feature=torch.cat([query_cached["feature"], feat_b]), + coords=[query_cached["coords"], coord_b], + ).to(device) + + with torch.no_grad(): + probability = model(input_data).cpu().tolist()[0] + + return { + "query": query_name, + "target": target_path.name, + "probability": probability, + "status": "ok", + "error": "", + } + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--query", + required=True, + help="Path to query protein file (.fasta/.fa/.faa/.fna or .pdb/.cif).", + ) + parser.add_argument( + "--targets", + required=True, + help="Path to folder containing target protein files.", + ) + parser.add_argument( + "--output", + default="inference_batch.csv", + help="Output CSV path. Default: inference_batch.csv", + ) + parser.add_argument( + "--chain_query", + default="first", + help="Chain ID for query structure file. Ignored for FASTA. Default: first", + ) + parser.add_argument( + "--chain_targets", + default="first", + help="Chain ID for target structure files. Ignored for FASTA. Default: first", + ) + parser.add_argument( + "--model", + default="ProtT5", + choices=["ProtT5", "ESM-2+ac"], + help="Model for structure inputs. FASTA inputs always force ESM-2+ac.", + ) + parser.add_argument( + "--device", + default="cuda", + help="Device to use (e.g., cuda or cpu). Default: cuda", + ) + args = parser.parse_args() + + query_path = Path(args.query) + target_dir = Path(args.targets) + output_path = Path(args.output) + + if not query_path.is_file(): + raise FileNotFoundError(f"Query file not found: {query_path}") + if not is_supported_file(query_path): + raise ValueError(f"Unsupported query file type: {query_path.suffix}") + if not target_dir.is_dir(): + raise NotADirectoryError(f"Target folder not found: {target_dir}") + + target_files = sorted([p for p in target_dir.iterdir() if p.is_file() and is_supported_file(p)]) + if not target_files: + raise ValueError(f"No supported files found in target folder: {target_dir}") + + device = torch.device(args.device if torch.cuda.is_available() else "cpu") + print("Using device:", device) + + query_is_fasta = is_fasta_file(query_path) + targets_are_fasta = [is_fasta_file(p) for p in target_files] + + if query_is_fasta: + if not all(targets_are_fasta): + raise ValueError("Query is FASTA, so all target files must also be FASTA.") + model_type = "ESM-2+ac" + else: + if any(targets_are_fasta): + raise ValueError("Query is structure file, so all target files must be structure files (.pdb/.cif).") + model_type = args.model + + print("Selected model:", model_type) + print("Loading embedder. Language model weights may download on first run.") + if model_type == "ESM-2+ac": + ckpt = "checkpoint/SpatialPPIv2_ESM.ckpt" + embedder = Embed("esm2_t33_650M_UR50D", device) + else: + ckpt = "checkpoint/SpatialPPIv2_ProtT5.ckpt" + embedder = Embed("Rostlab/prot_t5_xl_uniref50", device) + + cfg = getConfig("config/default.yaml") + cfg["basic"]["num_features"] = embedder.featureLen + + model = getModel(cfg, ckpt=ckpt).to(device) + model.eval() + + print("Encoding query protein once for batch screening...") + query_cached = load_query_features(query_path, args.chain_query, model_type, embedder) + + results = [] + for target in target_files: + try: + rec = infer_pair( + model=model, + device=device, + model_type=model_type, + query_cached=query_cached, + target_path=target, + query_name=query_path.name, + chain_target=args.chain_targets, + embedder=embedder, + ) + print(f"{target.name}: {rec['probability']:.6f}") + results.append(rec) + except Exception as exc: + print(f"{target.name}: failed ({exc})") + results.append( + { + "query": query_path.name, + "target": target.name, + "probability": "", + "status": "error", + "error": str(exc), + } + ) + + output_path.parent.mkdir(parents=True, exist_ok=True) + with open(output_path, "w", newline="", encoding="utf-8") as f: + writer = csv.DictWriter( + f, + fieldnames=["query", "target", "probability", "status", "error"], + ) + writer.writeheader() + writer.writerows(results) + + print(f"Saved {len(results)} results to {output_path}")