OMTRA: A Multi-Task Generative model for Structure-Based Drug Design

<p align="center"> <img src="assets/omtra_bot.png" alt="OMTRA Banner" width="800"> </p>

OMTRA is a flow-matching based generative model for small-molecule + protein systems. It supports a variety of tasks relevant to structure-based drug design, including:

• Unconditional 3D de novo molecule generation
• Unconditional ligand conformer generation
• Protein Pocket-conditioned de novo molecule design
• Protein-ligand docking (rigid and, flexible coming soon)
• Pharmacophore-conditioned molecule generation
• Pharmacophore-conditioned conformer generation
• Protein AND pharmacophore-conditioned molecule design
• Protein AND pharmacophore-conditioned docking

OMTRA is described in our preprint: https://arxiv.org/abs/2512.05080 and will be presented at MLSB 2025.

<p align="center"> <img src="assets/omtra_fig.png" alt="OMTRA Overview" width="75%"> </p>

---

Table of Contents

• Installation
  • Manual Installation (Recommended)
  • Docker Installation
• Model Weights
• Sampling
  • CLI Reference
  • Available Tasks
  • CLI Examples
  • Pharmacophore File Formats
  • Web Application
• Training
• Additional Documentation

---

Installation

There are two ways to set up OMTRA:

1. Manual Installation — Build the environment manually in a conda/mamba environment (recommended for most users)
2. Docker Installation — Use a Docker container for isolated, reproducible environments

System Requirements

• Linux System
• NVIDIA GPU with CUDA support (CUDA 12.1 recommended)
• Python 3.11

Manual Installation (Recommended)

This approach gives you direct control over the environment and is recommended for development and most use cases.

# Create and activate conda/mamba environment
mamba create -n omtra python=3.11
mamba activate omtra

# Clone the repository
git clone https://github.com/gnina/OMTRA.git
cd OMTRA

# Run the build script
chmod +x build_env.sh
./build_env.sh

The build script installs:

• CUDA-enabled versions of PyTorch, DGL, and PyG
• OMTRA package and all dependencies

After installation, the omtra command will be available:

omtra --task <task> [options]

Docker Installation

Docker provides an isolated environment and is particularly useful for deployment or if you want to use the web application interface.

Prerequisites

• Docker and Docker Compose installed
• NVIDIA Container Toolkit installed (installation guide)
• Model weights downloaded to omtra/trained_models/ directory (see Model Weights)

Using the Pre-built Image

The CLI image is available on Docker Hub and will be automatically pulled when you first use it:

cd OMTRA
source docker-cli-setup.sh
omtra --task <task> [options]

The setup script will automatically pull gnina/omtra:latest from Docker Hub if it's not already available locally.

Building the Docker Image Locally (Optional)

If you prefer to build the image yourself:

cd OMTRA
docker build -t gnina/omtra:latest .

Then set up the CLI wrapper:

source docker-cli-setup.sh

Making the CLI Available Permanently

Add the following to your shell configuration (~/.bashrc or ~/.zshrc):

source /path/to/OMTRA/docker-cli-setup.sh

Customizing the Docker Image

You can specify a custom image name or version by setting the OMTRA_CLI_IMAGE environment variable before sourcing the setup script:

export OMTRA_CLI_IMAGE="gnina/omtra:v1.0.0"
source docker-cli-setup.sh

To disable GPU support (for testing on CPU-only machines):

export OMTRA_NO_GPU=1

---

Model Weights

Download the pre-trained model weights using wget:

wget -r -np -nH --cut-dirs=3 -R "index.html*" -P omtra/trained_models https://bits.csb.pitt.edu/files/OMTRA/omtra_v0_weights/

This will create the omtra/trained_models/ directory with the checkpoint files. The CLI automatically selects the appropriate checkpoint based on the task. You can also specify a checkpoint explicitly with the --checkpoint flag.

---

Sampling

There are two ways to sample from a trained OMTRA model:

1. Command-Line Interface (CLI) — For scripting and batch processing
2. Web Application — For interactive exploration

CLI Reference

Basic Usage

omtra --task <task> [options]

The omtra command is available after either installation method. With manual installation, it's installed via pip install -e .. With Docker, the docker-cli-setup.sh script creates a shell function that wraps the Docker container.

Core Arguments

| Argument | Type | Default | Description | |----------|------|---------|-------------| | --task | string | required | The sampling task to perform (see Available Tasks) | | --checkpoint | path | auto | Path to model checkpoint (auto-detected from task if not provided) | | --n_samples | int | 100 | Number of samples to generate | | --n_timesteps | int | 250 | Number of integration steps during sampling | | --output_dir | path | None | Directory to save output files | | --metrics | flag | False | Compute evaluation metrics on generated samples |

Input File Arguments

For conditional generation tasks, you can provide input structures directly:

| Argument | Type | Description | |----------|------|-------------| | --protein_file | path | Protein structure file (PDB or CIF format) | | --ligand_file | path | Ligand structure file (SDF format) | | --pharmacophore_file | path | Pharmacophore file (JSON from Pharmit, XYZ, or SDF format) | | | | Pocket definition (choose one): | | --pocket_ligand | path | Path to reference ligand file (SDF) to define pocket around ligand atoms | | --pocket_center | string | Pocket center coordinates as 'x,y,z' | | --pocket_residues | string | Pocket residues as 'CHAIN:RESID,CHAIN:START-END' (e.g., 'A:123-125,B:200') | | --bbox_length | float | Bounding box length (Angstroms) when using --pocket_center (default: 23.0) |

When input files are provided, --n_samples specifies how many samples to generate for that single input system.

Advanced Sampling Options

| Argument | Type | Default | Description | |----------|------|---------|-------------| | --stochastic_sampling | flag | False | Enable stochastic (vs deterministic) sampling | | --noise_scaler | float | 1.0 | Scaling factor for noise in stochastic sampling | | --eps | float | 0.01 | Small epsilon value for numerical stability | | --visualize | flag | False | Generate visualization of sampling trajectory |

Ligand Size Control

| Argument | Type | Default | Description | |----------|------|---------|-------------| | --use_gt_n_lig_atoms | flag | False | Match ground truth ligand atom count | | --n_lig_atom_margin | float | 0.15 | Margin (±%) around ground truth atom count | | --n_lig_atoms_mean | float | None | Mean for normal distribution of atom counts | | --n_lig_atoms_std | float | None | Std dev for normal distribution of atom counts |

Available Tasks

OMTRA supports multiple drug design tasks. Use the --task argument to select one:

Note: Tasks marked with ⚠️ do not have pre-trained checkpoints available yet.

Unconditional Generation

| Task | Description | |------|-------------| | denovo_ligand_condensed | Generate novel drug-like molecules from scratch |

Protein-Conditioned Generation

| Task | Description | |------|-------------| | fixed_protein_ligand_denovo_condensed | Design ligands for a fixed protein binding site | | protein_ligand_denovo_condensed ⚠️ | Joint generation of ligand with flexible protein | | exp_apo_conditioned_denovo_ligand_condensed ⚠️ | De novo ligand generation starting from experimental apo structure | | pred_apo_conditioned_denovo_ligand_condensed ⚠️ | De novo ligand generation starting from predicted apo structure |

Docking Tasks

| Task | Description | |------|-------------| | rigid_docking_condensed | Dock a known ligand into a fixed protein structure | | flexible_docking_condensed ⚠️ | Dock with protein flexibility | | expapo_conditioned_ligand_docking_condensed ⚠️ | Docking starting from experimental apo structure | | predapo_conditioned_ligand_docking_condensed ⚠️ | Docking starting from predicted apo structure |

Conformer Generation

| Task | Description | |------|-------------| | ligand_conformer_condensed | Generate 3D conformations for a given ligand |

Pharmacophore-Conditioned Tasks

| Task | Description | |------|-------------| | denovo_ligand_pharmacophore_condensed ⚠️ | Generate ligand and pharmacophore jointly | | denovo_ligand_from_pharmacophore_condensed | Design ligand matching a given pharmacophore | | ligand_conformer_from_pharmacophore_condensed | Generate conformer satisfying pharmacophore | | fixed_protein_pharmacophore_ligand_denovo_condensed | Design ligand for protein with pharmacophore constraints | | rigid_docking_pharmacophore_condensed | Dock ligand with pharmacophore constraints | | protein_ligand_pharmacophore_denovo_condensed ⚠️ | Joint generation of ligand, protein, and pharmacophore |

CLI Examples