### OMTRA Training Command Example

Source: https://github.com/gnina/omtra/blob/main/docs/training.md

Example of a full training command for OMTRA, specifying run name, task group, batch size, training duration, dataset parameters, worker count, and model-specific training configurations. This command also demonstrates multi-GPU training setup.

```console
python routines/train.py           
    name=chirality_ph5050 \             # Run name
    task_group=pharmit5050_cond_a \     # Task group
    edges_per_batch=230000 \            # Batch size
    trainer.val_check_interval=600 
    max_steps=600000 \                  # Total train time
    plinder_pskip_factor=1.0 \
    num_workers=6  \  
    trainer=distributed \               # Multi-GPU training
    trainer.devices=2 \                 # Multi-GPU training
    graph=more_prot \                
    model.train_t_dist=beta \
    model.t_alpha=1.8 \
    model.cat_loss_weight=0.8 \
    model.time_scaled_loss=true
    model/aux_losses=pairs \
```

--------------------------------

### Example Training Set Data Points

Source: https://github.com/gnina/omtra/blob/main/omtra_pipelines/crossdocked_dataset/README.md

Illustrates the format of training set samples, which are tuples of PDB and SDF file paths.

```python
('DYR_STAAU_2_158_0/4xe6_X_rec_3fqc_55v_lig_tt_docked_4_pocket10.pdb',
 'DYR_STAAU_2_158_0/4xe6_X_rec_3fqc_55v_lig_tt_docked_4.sdf')
```

```python
('TRY1_BOVIN_66_246_0/1k1j_A_rec_1yp9_uiz_lig_tt_docked_1_pocket10.pdb',
 'TRY1_BOVIN_66_246_0/1k1j_A_rec_1yp9_uiz_lig_tt_docked_1.sdf')
```

--------------------------------

### Start OMTRA Web Application Services

Source: https://github.com/gnina/omtra/blob/main/omtra_webapp/START.md

Navigate to the omtra_webapp directory and run this command to start all services in detached mode.

```bash
cd omtra_webapp
docker-compose up -d
```

--------------------------------

### OMTRA Sampling Example Usage

Source: https://github.com/gnina/omtra/blob/main/docs/reproducing_results.md

Example of how to use the `routines/sample.py` script for sampling OMTRA models. This command demonstrates setting the task, number of samples, replicates, and enabling visualization and ground truth atom usage.

```console
python routines/sample.py \
    <PATH>/<CHECKPOINT>.ckpt \
    --task=fixed_protein_ligand_denovo_condensed \
    --n_samples=10 \
    --n_replicates=10 \
    --use_gt_n_lig_atoms \
    --visualize \

```

--------------------------------

### Manual Installation of OMTRA Environment

Source: https://github.com/gnina/omtra/blob/main/readme.md

Installs OMTRA and its dependencies in a new conda/mamba environment. This is the recommended approach for most users.

```bash
mamba create -n omtra python=3.11
mamba activate omtra

git clone https://github.com/gnina/OMTRA.git
cd OMTRA

chmod +x build_env.sh
./build_env.sh
```

--------------------------------

### Run Crossdocked Processing with SLURM

Source: https://github.com/gnina/omtra/blob/main/omtra_pipelines/crossdocked_dataset/README.md

Example SLURM script to execute the crossdocked data processing pipeline. Ensure environment variables and paths are correctly set.

```bash
#!/bin/bash
#SBATCH -J crossdocked
#SBATCH --partition=dept_cpu
#SBATCH -o slurm_output/crossdocked_processing/%A_%a.out
#SBATCH -e slurm_output/crossdocked_processing/%A_%a.out
#SBATCH --cpus-per-task 18

hostname

source ~/.bashrc
mamba activate omtra

python omtra_pipelines/crossdocked_dataset/run_crossdocked_processing_external_splits.py \
  --cd_directory /net/galaxy/home/koes/paf46_shared/cd2020_v1.3/types \
  --pocket_cutoff 8.0 \
  --zarr_output_dir /net/galaxy/home/koes/jmgupta/omtra_2/data/crossdocked/external_split \
  --root_dir /net/galaxy/home/koes/paf46_shared/cd2020_v1.3 \
  --max_batches None \
  --batch_size 500 \
  --n_cpus 16 \
  --max_pending 32

```

--------------------------------

### Instantiate PharmitDataset

Source: https://github.com/gnina/omtra/blob/main/docs/pharmit_dataset.md

Instantiate the PharmitDataset class to load the dataset. Specify the data directory, split, and desired return type. This example returns RDKit molecules.

```python
dataset = PharmitDataset(data_dir='/net/galaxy/home/koes/icd3/moldiff/OMTRA/data/pharmit',
                         split='test',
                         return_type='rdkit')
mol = dataset[0]
```

--------------------------------

### Docker CLI Setup for OMTRA

Source: https://github.com/gnina/omtra/blob/main/readme.md

Sets up the OMTRA command-line interface using a pre-built Docker image. The script automatically pulls the latest image if not found locally.

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

--------------------------------

### OMTRA CLI Usage

Source: https://github.com/gnina/omtra/blob/main/readme.md

Basic command structure for using the OMTRA command-line interface after installation.

```bash
omtra --task <task> [options]
```

--------------------------------

### Example Pocket Embedding Output

Source: https://github.com/gnina/omtra/blob/main/omtra_pipelines/plinder_dataset/plinder_tutorial_embed.ipynb

This is an example output of a generated pocket embedding and its shape.

```text
[[[  57.46649     12.286621   -43.25647   ...  -64.02086     46.0403
     44.52671  ]
  [ 241.72176     20.567871    -6.1555786 ... -260.55695    201.69476
     93.41516  ]
  [ 149.90399   -185.30414     53.811676  ... -180.73975    104.03235
    132.5448   ]
  ...
  [  66.27484     65.163574    70.37598   ... -167.59225     51.625427
     72.36229  ]
  [  53.55011   -108.667175   107.80731   ... -116.176636    16.85482
    -58.165955 ]
  [ -37.105103     4.1192627   42.01111   ... -105.40344     36.250763
     68.942505 ]]]
(1, 50, 1536)
```

--------------------------------

### Load Protein and Ligand with PoseCheck

Source: https://github.com/gnina/omtra/blob/main/notebooks/docking_evaluation.ipynb

Loads protein structure from PDB/CIF and ligand from SDF, then calculates interactions using PoseCheck. Ensure RDKit, gemmi, and PoseCheck are installed.

```python
from posecheck import PoseCheck
from rdkit import Chem
import gemmi

prot_file = "/net/galaxy/home/koes/ltoft/OMTRA/pistack_example/protein_0_H.pdb"
st = gemmi.read_structure("/net/galaxy/home/koes/ltoft/OMTRA/pistack_example/protein.cif")
st.write_pdb(prot_file)


prot = Chem.MolFromPDBFile("/net/galaxy/home/koes/ltoft/OMTRA/pistack_example/protein_0.pdb", removeHs=False, sanitize=False)
prot = Chem.AddHs(prot, addCoords=True)
Chem.SanitizeMol(prot)
Chem.SetAromaticity(prot)
w = Chem.PDBWriter("/net/galaxy/home/koes/ltoft/OMTRA/pistack_example/protein_0_H.pdb")
w.write(prot)
w.close()


true_lig = Chem.SDMolSupplier('/net/galaxy/home/koes/ltoft/OMTRA/pistack_example/ligand.sdf', sanitize=False, removeHs=False)[0]
Chem.SanitizeMol(true_lig)
true_lig = Chem.AddHs(true_lig, addCoords=True)
Chem.SetAromaticity(true_lig)

pc = PoseCheck()
pc.load_protein_from_pdb(prot_file)
pc.load_ligands_from_mols([true_lig])
interactions = pc.calculate_interactions()
```

--------------------------------

### Set Working Directory and Environment Variable

Source: https://github.com/gnina/omtra/blob/main/omtra_pipelines/plinder_dataset/plinder_tutorial_embed.ipynb

Configures the environment by changing the current working directory and setting the PLINDER_MOUNT environment variable. These commands are crucial for proper setup before running pipeline scripts.

```python
#feel free to skip, needed these for proper environment setup
%cd /net/galaxy/home/koes/nodonkor/OMTRA
%env PLINDER_MOUNT=/net/galaxy/home/koes/tjkatz/.local/share/plinder/2024-06/v2
```

--------------------------------

### Ligand Information Structure

Source: https://github.com/gnina/omtra/blob/main/omtra_pipelines/plinder_dataset/plinder_tutorial.ipynb

This is an example of the structure of ligand information extracted from the system annotation. It includes details like SMILES, molecular weight, and atom counts.

```python
Result:
{'pdb_id': '7ueu',
 'biounit_id': '1',
 'ligands': [{'pdb_id': '7ueu',
   'biounit_id': '1',
   'asym_id': 'C',
   'instance': 2,
   'ccd_code': 'AN2',
   'plip_type': 'SMALLMOLECULE',
   'bird_id': '',
   'centroid': [-24.080585479736328, 10.235146522521973, -25.439708709716797],
   'smiles': 'Nc1ncnc2c1ncn2[C@@H]1O[C@H](CO[P@](=O)(O)O[P@@](N)(=O)O)[C@@H](O)[C@H]1O',
   'resolved_smiles': 'O[C@@H]1[C@@H](CO[P@@](=O)(O[P@@](=O)([O])N)O)O[C@H]([C@@H]1O)n1cnc2c1ncnc2N',
   'residue_numbers': [1],
   'rdkit_canonical_smiles': 'Nc1ncnc2c1ncn2[C@@H]1O[C@H](CO[P@](=O)(O)O[P@@](N)(=O)O)[C@@H](O)[C@H]1O',
   'molecular_weight': 426.0453993519999,
   'crippen_clogp': -1.7796000000000007,
   'num_rot_bonds': 6,
   'num_hbd': 6,
   'num_hba': 12,
   'num_rings': 3,
   'num_heavy_atoms': 27,
   'is_covalent': False,

```

--------------------------------

### Pharmacophore-Guided Ligand Design from SDF

Source: https://github.com/gnina/omtra/blob/main/readme.md

Generate ligands that match a pharmacophore extracted from a reference ligand's SDF file. This is a convenient way to guide design based on an existing molecule's features.

```bash
omtra --task denovo_ligand_from_pharmacophore_condensed \
  --pharmacophore_file reference_ligand.sdf \
  --n_samples 100 \
  --output_dir outputs/pharm_guided
```

--------------------------------

### Plinder Ligand Information Example

Source: https://github.com/gnina/omtra/blob/main/omtra_pipelines/plinder_dataset/plinder_tutorial.ipynb

This snippet displays a dictionary containing detailed information about a specific ligand, including its PDB ID, chemical properties, and interactions with neighboring residues and ligands. It also includes Posebusters validation results.

```python
{'pdb_id': '7ueu',
   'biounit_id': '1',
   'asym_id': 'D',
   'instance': 2,
   'ccd_code': 'MG',
   'plip_type': 'SMALLMOLECULE',
   'bird_id': '',
   'centroid': [-30.303308486938477, 14.382668495178223, -24.010414123535156],
   'smiles': '[Mg+2]',
   'resolved_smiles': '[Mg]',
   'residue_numbers': [1],
   'rdkit_canonical_smiles': '[Mg+2]',
   'molecular_weight': 23.98394454018,
   'crippen_clogp': -0.3808,
   'num_rot_bonds': 0,
   'num_hbd': 0,
   'num_hba': 0,
   'num_rings': 0,
   'num_heavy_atoms': 1,
   'is_covalent': False,
   'covalent_linkages': [],
   'neighboring_residues': {'2.A': [25, 27, 28, 32, 123, 146, 197, 199, 200]},
   'neighboring_ligands': ['2.C'],
   'interacting_residues': {},
   'interacting_ligands': [],
   'interactions': {},
   'neighboring_residue_threshold': 6.0,
   'neighboring_ligand_threshold': 4.0,
   'num_neighboring_ppi_atoms_within_4A_of_gap': 0,
   'num_neighboring_ppi_atoms_within_8A_of_gap': 0,
   'num_missing_ppi_interface_residues': 0,
   'num_pli_atoms_within_4A_of_gap': 0,
   'num_pli_atoms_within_8A_of_gap': 0,
   'num_missing_pli_interface_residues': 0,
   'num_resolved_heavy_atoms': 1,
   'num_unresolved_heavy_atoms': 0,
   'tpsa': 0.0,
   'qed': 0.3158896948418247,
   'is_ion': True,
   'is_lipinski': False,
   'is_fragment': False,
   'is_oligo': False,
   'is_cofactor': False,
   'in_artifact_list': False,
   'is_artifact': False,
   'is_other': True,
   'is_invalid': False,
   'posebusters_result': {
    'mol_pred_loaded': True,
    'mol_true_loaded': True,
    'mol_cond_loaded': True,
    'sanitization': True,
    'all_atoms_connected': True,
    'molecular_formula': True,
    'molecular_bonds': True,
    'double_bond_stereochemistry': True,
    'tetrahedral_chirality': True,
    'bond_lengths': True,
    'bond_angles': True,
    'internal_steric_clash': True,
    'aromatic_ring_flatness': True,
    'double_bond_flatness': True,
    'protein-ligand_maximum_distance': True,
    'minimum_distance_to_protein': True,
    'minimum_distance_to_organic_cofactors': True,
    'minimum_distance_to_inorganic_cofactors': True,
    'minimum_distance_to_waters': True,
    'volume_overlap_with_protein': True,
    'volume_overlap_with_organic_cofactors': True,
    'volume_overlap_with_inorganic_cofactors': True,
    'volume_overlap_with_waters': True,
    'rmsd_≤_2å': True,
    'passes_valence_checks': True,
    'passes_kekulization': True,
    'inchi_crystal_valid': True,
    'inchi_docked_valid': True,
    'inchi_crystal': 'InChI=1S/Mg/q+2',
    'inchi_docked': 'InChI=1S/Mg/q+2',
    'inchi_overall': True,
    'inchi_version': True,
    'stereochemistry_preserved': True,
    'hydrogens': True,
    'net_charge': True,
    'protons': True,
    'stereo_sp3': True,
    'stereo_sp3_inverted': True,
    'stereo_type': True,
    'number_bonds': None,
    'shortest_bond_relative_length': None,
    'longest_bond_relative_length': None,
    'number_short_outlier_bonds': None,
    'number_long_outlier_bonds': None,
    'number_angles': None,
    'most_extreme_relative_angle': None,
    'number_outlier_angles': None,
    'number_noncov_pairs': None,
    'shortest_noncovalent_relative_distance': None,
    'number_clashes': None,
    'aromatic_ring_maximum_distance_from_plane': None,
    'double_bond_maximum_distance_from_plane': None
   }
}
```

--------------------------------

### Load Configuration and Dataset

Source: https://github.com/gnina/omtra/blob/main/notebooks/ligand_viewer.ipynb

Loads a configuration using `quick_load.load_cfg` and instantiates a datamodule. It then loads the training dataset and extracts a specific graph object from the 'pharmit' subset.

```python
cfg = quick_load.load_cfg(overrides=['task_group=no_protein'], pharmit_path='/net/galaxy/home/koes/icd3/moldiff/OMTRA/data/pharmit') # you can provide overrides to the default config via command-line syntax here
datamodule = datamodule_from_config(cfg)
train_dataset = datamodule.load_dataset("train")
pharmit_train_dataset = train_dataset.datasets['pharmit']
g = pharmit_train_dataset[('denovo_ligand', 0)]
g.nodes["lig"].data.keys()
```

--------------------------------

### Run Docking Evaluation (With Samples Dir)

Source: https://github.com/gnina/omtra/blob/main/omtra_pipelines/docking_eval/readme.md

This command runs the docking evaluation by reading samples from a specified directory and system information file. Ensure the `--samples_dir` and `--sys_info_file` paths are correctly set.

```console
python docking_eval.py \
    --samples_dir OMTRA_ROOT/omtra_pipelines/docking_eval/outputs/rigid_docking_condensed/samples \
    --sys_info_file OMTRA_ROOT/omtra_pipelines/docking_eval/outputs/rigid_docking_condensed/rigid_docking_condensed_sys_info.csv \
    --task rigid_docking_condensed \
    --n_samples 10 \
    --n_replicates 10 \
    --max_batch_size 600 \
    --timeout 1200 
```

--------------------------------

### Instantiate Datamodule and Model

Source: https://github.com/gnina/omtra/blob/main/notebooks/docking_evaluation.ipynb

Loads a datamodule from configuration and an OMTRA model from a checkpoint. The model is moved to the specified device and set to evaluation mode.

```python
dm  = quick_load.datamodule_from_config(train_cfg)
multitask_dataset = dm.load_dataset('test')
model = quick_load.omtra_from_checkpoint(ckpt_path).to(device).eval()
```

--------------------------------

### Pharmacophore JSON Structure Example

Source: https://github.com/gnina/omtra/blob/main/readme.md

This is an example of the JSON structure for pharmacophore constraints. It includes points with names, coordinates, and an enabled flag.

```json
{
  "points": [
    {
      "name": "Aromatic",
      "x": 10.5,
      "y": 20.3,
      "z": 15.2,
      "enabled": true
    },
    {
      "name": "HydrogenAcceptor",
      "x": 8.2,
      "y": 18.7,
      "z": 14.1,
      "enabled": true
    }
  ]
}
```

--------------------------------

### Initialize PharmitDataset

Source: https://github.com/gnina/omtra/blob/main/notebooks/zarr_dev.ipynb

Initializes a PharmitDataset instance from a Zarr file. Replace 'test_ligand_dataset.zarr' with your Zarr file path.

```python
dataset = PharmitDataset('test_ligand_dataset.zarr')
```

--------------------------------

### Run Docking Evaluation (No Samples Dir)

Source: https://github.com/gnina/omtra/blob/main/omtra_pipelines/docking_eval/readme.md

This command initiates the docking evaluation task using a checkpoint path and specifying task-related parameters. It's used when the script generates its own samples.

```console
python docking_eval.py \
    --ckpt_path PATH/last.ckpt \
    --task rigid_docking_condensed \
    --n_samples 10 \
    --n_replicates 10 \
    --max_batch_size 600 \
    --timeout 1200 
```

--------------------------------

### Sample de novo ligands and visualize trajectory

Source: https://github.com/gnina/omtra/blob/main/notebooks/sampling_example.ipynb

Samples a de novo ligand and visualizes its generation trajectory, showing the progression from an initial state to the predicted final state. Requires the helper functions `mols_to_sdf` and `display_ligand_trajectory` to be defined.

```python
sys = model.sample(
    task_name='denovo_ligand',
    g_list=None,
    n_replicates=1,
    unconditional_n_atoms_dist='pharmit',
    n_timesteps=200,
    visualize=True
)[0]
display_ligand_trajectory(sys)
```

--------------------------------

### Import PharmitDataset

Source: https://github.com/gnina/omtra/blob/main/notebooks/zarr_dev.ipynb

Imports the necessary PharmitDataset class from the omtra library. Ensure omtra is installed.

```python
from omtra.dataset.pharmit import PharmitDataset
```

--------------------------------

### OMTRA CLI: All Features Disabled

Source: https://github.com/gnina/omtra/blob/main/readme.md

Create a pharmacophore JSON file where all features are disabled by default. This can be useful for starting with a blank hypothesis.

```bash
omtra mol2pharm ligand.sdf -o pharm.json --all-disabled
```

--------------------------------

### Import Libraries for Plinder Pipeline

Source: https://github.com/gnina/omtra/blob/main/omtra_pipelines/plinder_dataset/plinder_tutorial.ipynb

Imports necessary modules and classes for the Plinder data processing pipeline. Ensure these libraries are installed and accessible.

```python
from omtra_pipelines.plinder_dataset.plinder_pipeline import *
import plinder.core.utils.config
from plinder.core import PlinderSystem
from omtra_pipelines.plinder_dataset.utils import NPNDE_MAP, LIGAND_MAP
```

--------------------------------

### Make OMTRA Docker CLI Permanently Available

Source: https://github.com/gnina/omtra/blob/main/readme.md

Adds the OMTRA Docker CLI setup script to your shell configuration for persistent access.

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

--------------------------------

### Instantiate Sampled Systems

Source: https://github.com/gnina/omtra/blob/main/notebooks/docking_evaluation.ipynb

Instantiates SampledSystem objects for docking evaluation. Requires pre-loaded data and task configuration.

```python
from omtra.eval.system import SampledSystem
from omtra.data.condensed_atom_typing import CondensedAtomTyper
from omtra.tasks.tasks import Task
from omtra.tasks.register import task_name_to_class
from routines.sample import write_ground_truth


task_name = 'rigid_docking_condensed'
task: Task = task_name_to_class(task_name)

cond_a_typer = CondensedAtomTyper(fake_atoms=False, include_crossdocked=True)

ss = [] 

for i in range(len(g_list)):
    ss.append(SampledSystem(
                g=g_list[i],
                task=task,
                fake_atoms=False,
                cond_a_typer=cond_a_typer,
            ))
```

--------------------------------

### Get Plinder Configuration

Source: https://github.com/gnina/omtra/blob/main/omtra_pipelines/plinder_dataset/plinder_tutorial.ipynb

Retrieves the Plinder configuration, including local cache and remote data directories. This is useful for understanding where data is stored and accessed.

```python
cfg = plinder.core.get_config()
print(f"local cache directory: {cfg.data.plinder_dir}")
print(f"remote data directory: {cfg.data.plinder_remote}")
```

--------------------------------

### Load configuration and dataset for OMTRA

Source: https://github.com/gnina/omtra/blob/main/notebooks/sampling_example.ipynb

Loads the configuration, instantiates the datamodule, and prepares the training dataset, specifically the 'pharmit' subset. It also loads the OMTRA model from a checkpoint.

```python
cfg = quick_load.load_cfg(overrides=['task_group=pharmit5050_cond_a'], pharmit_path='/home/ian/projects/mol_diffusion/OMTRA/data/pharmit') # you can provide overrides to the default config via command-line syntax here
datamodule = datamodule_from_config(cfg)
train_dataset = datamodule.load_dataset("val")
pharmit_dataset = train_dataset.datasets['pharmit']
ckpt_path = '/home/ian/projects/mol_diffusion/OMTRA/local/runs_from_cluster/uncond_r179_iiftoej0/checkpoints/last.ckpt'
model = quick_load.omtra_from_checkpoint(ckpt_path).cuda().eval()
```

--------------------------------

### Load Model and Display Performance

Source: https://github.com/gnina/omtra/blob/main/notebooks/ligand_viewer.ipynb

Loads a Ligand VQ-VAE model from a checkpoint and displays its performance on a dataset, with an option to ignore bond orders.

```python
checkpoint = "/net/galaxy/home/koes/ltoft/OMTRA/outputs/2025-07-08/e_nonzero_loss_2025-07-08_10-21-078175/checkpoints/last.ckpt"

model = LigandVQVAE.load_from_checkpoint(checkpoint)

display_performance(model, pharmit_val_dataset, n_mols=9, use_bond_order=False)
```

--------------------------------

### Complete Pharmacophore JSON Example

Source: https://github.com/gnina/omtra/blob/main/readme.md

Defines a binding hypothesis with multiple pharmacophore features, including coordinates and enabled status. The 'PositiveIon' feature is shown as disabled.

```json
{
  "points": [
    {
      "name": "Aromatic",
      "x": 12.456,
      "y": 8.234,
      "z": 15.789,
      "enabled": true
    },
    {
      "name": "HydrogenDonor",
      "x": 10.123,
      "y": 11.456,
      "z": 14.234,
      "enabled": true
    },
    {
      "name": "HydrogenAcceptor",
      "x": 14.567,
      "y": 9.890,
      "z": 13.456,
      "enabled": true
    },
    {
      "name": "Hydrophobic",
      "x": 11.234,
      "y": 7.890,
      "z": 17.123,
      "enabled": true
    },
    {
      "name": "PositiveIon",
      "x": 13.890,
      "y": 12.345,
      "z": 16.789,
      "enabled": false
    }
  ]
}
```

--------------------------------

### Sample ligand conformers and visualize trajectory

Source: https://github.com/gnina/omtra/blob/main/notebooks/sampling_example.ipynb

Samples conformers for a given ligand and visualizes the trajectory of the conformer generation process. This requires the helper functions `mols_to_sdf` and `display_ligand_trajectory` to be defined.

```python
# get a ligand from the dataset to do conformer sampling with!
g = pharmit_dataset[('ligand_conformer', 0)].to('cuda')
g_list = [g]

sys = model.sample(
    task_name='ligand_conformer',
    g_list=g_list,
    device=g.device,
    n_replicates=1,
    unconditional_n_atoms_dist='pharmit',
    n_timesteps=200,
    visualize=True
)[0]
display_ligand_trajectory(sys)
```

--------------------------------

### Load Baseline Model and Display Embeddings

Source: https://github.com/gnina/omtra/blob/main/notebooks/ligand_viewer.ipynb

Loads a baseline LigandVQVAE model from a checkpoint and displays ligand embeddings. Ensure the checkpoint path is correct.

```python
# BASELINE:
checkpoint = "/net/galaxy/home/koes/ltoft/OMTRA/outputs/2025-06-30/EMA_baseline_2025-06-30_14-13-407448/checkpoints/last.ckpt"

model = LigandVQVAE.load_from_checkpoint(checkpoint)

display_ligand_embeddings(model, pharmit_val_dataset, index_to_color, indices=[4455608, 2443760, 3514035, 3498938, 2873876, 2638608, 2032191, 2359807, 2577854, 4243698, 3502419, 2784527])
```

--------------------------------

### Plot memory measurements

Source: https://github.com/gnina/omtra/blob/main/notebooks/zarr_memory_investigation.ipynb

Generates a plot of memory usage relative to the start over batch index. Vertical red dashed lines indicate chunk boundaries.

```python
plt.plot(mem_measurements)

# plot vertical lines every 1000 samples

start_idxs = np.arange(0, root['data'].shape[0] - batch_size, batch_size)
end_idxs = start_idxs + batch_size

chunk_boundaries = []
current_chunk = 0
for batch_idx, (start_idx, end_idx) in enumerate(zip(start_idxs, end_idxs)):
    start_chunk = start_idx // root['data'].chunks[0]
    end_chunk = end_idx // root['data'].chunks[0]
    if start_chunk != current_chunk or end_chunk != current_chunk:
        chunk_boundaries.append(batch_idx)
        current_chunk = end_chunk

for i , boundary_batch_idx in enumerate(chunk_boundaries):
    plt.axvline(x=boundary_batch_idx, color='r', linestyle='--', lw=1, label='Chunk Boundary' if i == 0 else None)
    

# plt.xlim(0, 10000)
legend = plt.legend(framealpha=1)
plt.xlabel('Batch Index')
plt.ylabel('Process Memory usage relative to start (MB)')
plt.show()
```

--------------------------------

### Initialize and Process System

Source: https://github.com/gnina/omtra/blob/main/omtra_pipelines/plinder_dataset/plinder_tutorial.ipynb

Initializes the SystemProcessor with custom parameters and processes a specified system. This is used to extract receptor, ligand, and pocket data.

```python
system_processor = SystemProcessor(
    ligand_atom_map=LIGAND_MAP,
    npnde_atom_map=NPNDE_MAP,
    pocket_cutoff=8.0,
    link_type="apo",
)
result = system_processor.process_system(
    system_id="7gcm__1__1.A__1.D", save_pockets=False
)
```

--------------------------------

### Customize OMTRA Docker Image

Source: https://github.com/gnina/omtra/blob/main/readme.md

Specifies a custom Docker image name or version for OMTRA CLI setup. Also shows how to disable GPU support.

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

export OMTRA_NO_GPU=1
```

--------------------------------

### Define current memory usage function

Source: https://github.com/gnina/omtra/blob/main/notebooks/zarr_memory_investigation.ipynb

A utility function to get the current Resident Set Size (RSS) memory usage of the current process in megabytes.

```python
process = psutil.Process(os.getpid())
def current_memory():
    # process = psutil.Process(os.getpid())
    memory_usage = process.memory_info().rss
    memory_usage_mb = memory_usage / (1024 ** 2)
    return memory_usage_mb
```

--------------------------------

### Conformer Generation

Source: https://github.com/gnina/omtra/blob/main/readme.md

Generate multiple conformers for a given molecule. Specify the input ligand file, number of samples, and output directory.

```bash
omtra --task ligand_conformer_condensed \
  --ligand_file molecule.sdf \
  --n_samples 20 \
  --output_dir outputs/conformers
```

--------------------------------

### Initialize Prolif Fingerprint

Source: https://github.com/gnina/omtra/blob/main/notebooks/docking_evaluation.ipynb

Initializes a Prolif fingerprint object to calculate PiStacking interactions. Requires Prolif and RDKit libraries.

```python
from prolif.fingerprint import Fingerprint
import prolif as plf
from prolif import Molecule as plfMol

#prot = plf.Molecule(prot)
plf_prot = plfMol.from_pdb("/net/galaxy/home/koes/ltoft/OMTRA/pistack_example/protein_0_H.pdb")

plf_prot = plf.Molecule.from_rdkit(prot)
plf_lig = plf.Molecule.from_rdkit(true_lig)

fingerprint = Fingerprint(interactions=['PiStacking'])
fingerprint.run_from_iterable(lig_iterable=[true_lig], prot_mol=prot)
```

--------------------------------

### Load and Compare De Novo Ligand Metrics

Source: https://github.com/gnina/omtra/blob/main/notebooks/stochastic_sampling_investigation.ipynb

Loads condensed metrics for stochastic and regular de novo ligand generation from pickle files and displays them as a pandas DataFrame. This is useful for comparing the quality and properties of newly generated ligands.

```python
metrics = {}

with open("/net/galaxy/home/koes/ltoft/OMTRA/samples/uncond_stoch_sampling/denovo_lig/denovo_ligand_condensed_metrics.pkl", "rb") as f:
    denovo_lig_stoch_sampling = pickle.load(f)
    metrics['stoch_sampling'] = denovo_lig_stoch_sampling

with open("/net/galaxy/home/koes/ltoft/OMTRA/samples/uncond/denovo_lig/denovo_ligand_condensed_metrics.pkl", "rb") as f:
     denovo_lig = pickle.load(f)
     metrics['regular'] = denovo_lig

metrics = pd.DataFrame(metrics)
metrics
```

--------------------------------

### Import necessary libraries for OMTRA

Source: https://github.com/gnina/omtra/blob/main/notebooks/ligand_viewer.ipynb

Imports essential libraries for OMTRA, including plotting, PyTorch, DGL, RDKit, and custom OMTRA modules. Ensure these libraries are installed.

```python
import matplotlib.pyplot as plt
import matplotlib
import numpy as np
from matplotlib.patches import Patch
import torch
import dgl
import random
import io


from rdkit.Chem import AllChem
from rdkit import Chem
from rdkit.Chem import Draw
from rdkit.Chem.Draw import rdMolDraw2D,IPythonConsole
from PIL import Image

from omtra.load.quick import datamodule_from_config
import omtra.load.quick as quick_load

from omtra.constants import charge_map, lig_atom_type_map

from omtra.models.ligand_encoder.vq import LigandVQVAE

from omtra.tasks.register import task_name_to_class
from omtra.eval.system import SampledSystem
from omtra.data.graph.utils import get_upper_edge_mask

from omtra.constants import *

```

--------------------------------

### Get Codebook Indices and Embeddings

Source: https://github.com/gnina/omtra/blob/main/notebooks/ligand_viewer.ipynb

Extracts codebook indices and latent space embeddings from a graph representation of a molecule using a VQ-VAE model's encoder and VQ-VAE layers.

```python
def get_codebook_indices(model, g):
    z_e, mask = model.encoder(g)  # Encoder output
    mask = mask.cpu().numpy().flatten().tolist()
    _, z_d, codebook_indices, _, _ = model.vq_vae(z_e)
    codebook_indices = codebook_indices.cpu().numpy().flatten().tolist()
    return codebook_indices, z_d, mask
```

--------------------------------

### Load and Compare Ligand Conformer Metrics

Source: https://github.com/gnina/omtra/blob/main/notebooks/stochastic_sampling_investigation.ipynb

Loads condensed metrics for stochastic and regular ligand conformer generation from pickle files and displays them as a pandas DataFrame. This is useful for comparing the stability and validity of generated conformers.

```python
metrics = {}

pd.options.display.float_format = "{:.3f}".format


with open("/net/galaxy/home/koes/ltoft/OMTRA/samples/uncond_stoch_sampling/lig_conformer/ligand_conformer_condensed_metrics.pkl", "rb") as f:
     lig_conf_stoch_sampling = pickle.load(f)
     metrics['stoch_sampling'] = lig_conf_stoch_sampling

with open("/net/galaxy/home/koes/ltoft/OMTRA/samples/uncond/lig_conformer/ligand_conformer_condensed_metrics.pkl", "rb") as f:
     lig_conf = pickle.load(f)
     metrics['regular'] = lig_conf

metrics = pd.DataFrame(metrics)
metrics
```

--------------------------------

### Plinder Dataset Structure Example

Source: https://github.com/gnina/omtra/blob/main/omtra_pipelines/plinder_dataset/plinder_tutorial_embed.ipynb

Illustrates the typical structure of a Plinder dataset, showing coordinate arrays, type arrays, and vector embeddings. This format is common for representing molecular or structural data.

```python
types=array([0, 0, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 1, 1, 1, 1, 5, 5]), vectors=array([[[-0.19175664, -0.98141224, -0.007707  ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[ 0.1284616 , -0.97854487,  0.16108247],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[ 0.13031483, -0.34575994, -0.92922985],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[-0.3236038 , -0.04536502,  0.94510454],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[ 0.78092973,  0.59654354, -0.18516092],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[ 0.33204687,  0.17050088, -0.92772535],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[ 0.74925527,  0.33311861,  0.57240591],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[ 0.04276047, -0.43860023,  0.8976644 ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[-0.56563373,  0.4859512 , -0.66626565],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[-0.33242906,  0.93829588, -0.09535075],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[-0.60015455, -0.52373369, -0.60458047],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[-0.80140211, -0.58458442, -0.12655323],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[ 0.58910435,  0.63314625,  0.50207757],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[ 0.33373653, -0.90478529, -0.26454395],
        [-0.65519135,  0.09169369,  0.74987769],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[ 0.8136989 , -0.29745378,  0.49941501],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[ 0.1868987 , -0.83008801, -0.52537869],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[ 0.23339875, -0.80215136,  0.54961642],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[ 0.10221835, -0.10091621,  0.98962989],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[-0.56004532,  0.81263243, -0.16117619],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]],

       [[ 0.1284616 , -0.97854487,  0.16108247],
        [ 0.        ,  0.        ,  0.        ],
        [ 0.        ,  0.        ,  0.        ]]])
```

--------------------------------

### Load and Compare Ligand Conformer Metrics

Source: https://github.com/gnina/omtra/blob/main/notebooks/stochastic_sampling_investigation.ipynb

Loads and compares condensed metrics for ligand conformer generation using stochastic sampling versus a regular method. This snippet is useful for evaluating the performance differences between the two sampling strategies.

```python
metrics = {}

pd.options.display.float_format = "{:.3f}".format

with open("/net/galaxy/home/koes/ltoft/OMTRA/samples/phcond_seeded_stoch_sampling/lig_conformer_from_pharm/ligand_conformer_from_pharmacophore_condensed_metrics.pkl", "rb") as f:
     lig_conf_stoch_sampling = pickle.load(f)
     metrics['stoch_sampling'] = lig_conf_stoch_sampling

with open("/net/galaxy/home/koes/ltoft/OMTRA/samples/phcond_seeded/lig_conformer_from_pharm/ligand_conformer_from_pharmacophore_condensed_metrics.pkl", "rb") as f:
      lig_conf = pickle.load(f)
      metrics['regular'] = lig_conf

metrics = pd.DataFrame(metrics)
metrics
```

--------------------------------

### Load OMTRA configuration and dataset

Source: https://github.com/gnina/omtra/blob/main/notebooks/ligand_viewer.ipynb

Loads the OMTRA configuration using a specified path and overrides, then loads the training dataset. This prepares the data for further processing.

```python
from omtra.data.condensed_atom_typing import CondensedAtomTyper
from omtra.eval.system import SampledSystem

cfg = quick_load.load_cfg(overrides=['task_group=pharmit5050_cond_a'], pharmit_path='/net/galaxy/home/koes/icd3/moldiff/OMTRA/data/pharmit') # you can provide overrides to the default config via command-line syntax here
datamodule = datamodule_from_config(cfg)
train_dataset = datamodule.load_dataset("train")
pharmit_train_dataset = train_dataset.datasets['pharmit']
g = pharmit_train_dataset[('denovo_ligand_condensed', 0)]

```

--------------------------------

### Pharmacophore-Guided Ligand Design from File

Source: https://github.com/gnina/omtra/blob/main/readme.md

Generate ligands that match a given pharmacophore. Use this command when you have a pharmacophore definition in a JSON or XYZ file.

```bash
omtra --task denovo_ligand_from_pharmacophore_condensed \
  --pharmacophore_file constraints.json \
  --n_samples 100 \
  --output_dir outputs/pharm_guided
```

--------------------------------

### Get Zarr Array Information

Source: https://github.com/gnina/omtra/blob/main/notebooks/zarr_dev.ipynb

Retrieves and displays detailed information about a specific Zarr array, including its data type, shape, chunk shape, and storage details. Useful for verifying array properties.

```python
root['node_data/a'].info
```

--------------------------------

### Load and Compare Sampling Metrics

Source: https://github.com/gnina/omtra/blob/main/notebooks/stochastic_sampling_investigation.ipynb

Loads metrics from both stochastic and regular de novo ligand generation processes and stores them in a pandas DataFrame for comparison. Ensure the pickle files exist at the specified paths.

```python
pd.options.display.float_format = "{:.3f}".format

metrics = {}

with open("/net/galaxy/home/koes/ltoft/OMTRA/samples/phcond_seeded_stoch_sampling/denovo_lig_from_pharm/denovo_ligand_from_pharmacophore_condensed_metrics.pkl", "rb") as f:
    denovo_lig_stoch_sampling = pickle.load(f)
    metrics['stoch_sampling'] = denovo_lig_stoch_sampling

with open("/net/galaxy/home/koes/ltoft/OMTRA/samples/phcond_seeded/denovo_lig_from_pharm/denovo_ligand_from_pharmacophore_condensed_metrics.pkl", "rb") as f:
     denovo_lig = pickle.load(f)
     metrics['regular'] = denovo_lig

metrics = pd.DataFrame(metrics)
metrics
```

--------------------------------

### Sample de novo ligands with OMTRA

Source: https://github.com/gnina/omtra/blob/main/notebooks/sampling_example.ipynb

Samples new ligand structures using the 'denovo_ligand_condensed' task. This snippet shows how to obtain the RDKit molecule object from the sampled system.

```python
sampled_systems = model.sample(
    task_name='denovo_ligand_condensed',
    g_list=None,
    n_replicates=1,
    unconditional_n_atoms_dist='pharmit',
    n_timesteps=200,
)

sampled_systems[0].get_rdkit_ligand()
```