bio_tools 0.1.1

Install, run, and inspect computational biology and chemistry tools, e.g. AlphaFold, Boltz, RFdiffusion, and ProteinMPNN
Documentation
use std::{path::Path, process::Command};

use super::{
    InstallError, Installer,
    common::{PipOptions, TorchBackend},
};
use crate::tool_definitions::Tool;

const LIGAND_WEIGHTS_ROOT: &str = "https://files.ipd.uw.edu/pub/ligandmpnn";
const ABMPNN_WEIGHTS: &str = "https://zenodo.org/records/8164693/files/abmpnn.pt?download=1";

pub(super) fn install_ligand(installer: &mut Installer) -> Result<(), InstallError> {
    const SLUG: &str = Tool::LigandMpnn.slug();
    install_runtime(installer, SLUG)?;
    let target = installer.tools_root().join("LigandMPNN");
    installer.clone_or_update("https://github.com/dauparas/LigandMPNN", &target)?;

    for filename in [
        "ligandmpnn_v_32_010_25.pt",
        "proteinmpnn_v_48_020.pt",
        "solublempnn_v_48_020.pt",
    ] {
        installer.download(
            &format!("{LIGAND_WEIGHTS_ROOT}/{filename}"),
            &target.join("model_params").join(filename),
        )?;
    }
    require_nonempty(
        &target.join("model_params/ligandmpnn_v_32_010_25.pt"),
        "the LigandMPNN weights did not download",
    )?;
    installer.note(format!("LigandMPNN installed at {}", target.display()));
    Ok(())
}

pub(super) fn install_protein(installer: &mut Installer) -> Result<(), InstallError> {
    const SLUG: &str = Tool::ProteinMpnn.slug();
    install_runtime(installer, SLUG)?;
    let target = installer.tools_root().join("ProteinMPNN");
    installer.clone_or_update("https://github.com/dauparas/ProteinMPNN", &target)?;
    require_nonempty(
        &target.join("vanilla_model_weights/v_48_020.pt"),
        "the ProteinMPNN checkout does not contain its vanilla model weights",
    )?;
    installer.download(ABMPNN_WEIGHTS, &target.join("abmpnn_weights/v_48_020.pt"))?;
    convert_weights(installer, &target);
    installer.note(format!("ProteinMPNN installed at {}", target.display()));
    Ok(())
}

fn install_runtime(installer: &mut Installer, slug: &str) -> Result<(), InstallError> {
    let backend = installer.select_torch_backend()?;
    installer.create_venv(slug, "3.12")?;
    installer.install_torch(slug, &["torch==2.7.1"], backend)?;
    installer.pip_install(slug, &["numpy<2"], PipOptions::default())?;
    if backend == TorchBackend::Cuda126 && !installer.torch_cuda_works(slug) {
        installer.note("Warning: Torch cannot reach the GPU; MPNN will run on CPU");
    }
    Ok(())
}

fn convert_weights(installer: &Installer, target: &Path) {
    let converter = [
        "scripts/convert_mpnn_weights.py",
        "install_scripts/convert_mpnn_weights.py",
        "convert_mpnn_weights.py",
    ]
    .into_iter()
    .find_map(|candidate| installer.support_file(candidate));
    let Some(converter) = converter else {
        installer.note(
            "convert_mpnn_weights.py was not supplied; skipping native ddG weight conversion",
        );
        return;
    };

    installer.step("Converting ProteinMPNN weights for the native ddG scanner");
    let mut command = Command::new(installer.venv_python(Tool::ProteinMpnn.slug()));
    command
        .arg(converter)
        .arg("--checkpoint")
        .arg(target.join("vanilla_model_weights/v_48_020.pt"))
        .arg("--output")
        .arg(target.join("converted/v_48_020.mcnn"))
        .arg("--repo")
        .arg(target);
    if installer.succeeds(&mut command) {
        installer.note("Native ddG scanning weights are available");
    } else {
        installer.note("Weight conversion failed; Python MPNN remains installed and usable");
    }
}

fn require_nonempty(path: &Path, message: &str) -> Result<(), InstallError> {
    if path.metadata().is_ok_and(|metadata| metadata.len() > 0) {
        Ok(())
    } else {
        Err(InstallError::InvalidConfiguration(format!(
            "{message}: {}",
            path.display()
        )))
    }
}