Alignment Throughput
====================
The following reports throughput for several alignment modes on a Linux HPC system.
They summarize how backend and batching choices affect wall-clock cost across workloads.
## Experiment
Experiments use 98 diverse drugs (50–800 g/mol), randomly paired under a fixed seed with shared embedded conformers.
Alignments are run through the `MoleculePair` and `MoleculePairBatch` APIs, which wrap the same optimization routines.
The figure below focuses on volumetric, volumetric + ESP (partial charge), and pharmacophore alignment for pair counts from 10 to 10,000 at each order of magnitude.
We also vary the number of SE(3) initializations (`num_repeats`) at 1, 5, and 50.
The default in `MoleculePair` and `MoleculePairBatch` is 50, which matters most for surface and surface + ESP similarity; for non-surface modes, five repeats are typically adequate.
The five-initialization setting corresponds to one identity alignment plus four principal-component–based starts.
The four principal-component starts match the intiialization count in RDKit/PubChem shape–color alignment.
Tabulated results below therefore use `num_repeats = 5`.
> Note:
For the sequential alignments, `optimize_{*}_overlay{*}` functions were directly used rather than through the API call in MoleculePair.
## Benchmark results
At small batch sizes, the PyTorch analytical-gradient path is generally fastest and remains the default in `MoleculePair` as it is easily parallelizable via multiprocessing across devices--unlike Jax.
For longer runs where the number of pairs grows, padded-batch Jax (`MoleculePairBatch` which unifies shapes under-the-hood) improves throughput once alignment count is on the order of 100 or higher, chiefly because a single JIT compilation amortizes over the batch.
Additional speedup is possible with multi-device parallelism via `shard_map`; the gain may be hardware- and environment-dependent.
Bucketing groups similarly sized molecules to reduce padding overhead but increases the number of distinct compiled kernels, so it trades wasted arithmetic against compilation cost; it is most valuable when molecular size varies widely across the set.
## Tabulated throughput for different modalities at `num_repeats = 5`
The tables list throughputs for `num_repeats = 5` by modality.
Configurations that exhausted memory—notably batched surface and surface+ESP runs—are omitted; batched surface alignment is not recommended in this setting.
Sequential alignment via `MoleculePair` for Jax and PyTorch modalities were only evaluated for <100 pairs because these do not scale.
Multiprocessed `MoleculePairBatch` was evaluated for >100 pairs.
### Volumetric alignment — alignments/s
| Approach | n = 10 | n = 100 | n = 1000 | n = 10000 |
|---|---:|---:|---:|---:|
| Jax Volume | 0.30 | 2.05 | - | - |
| PyTorch (autodiff) | 0.39 | 6.35 | - | - |
| PyTorch (analytical) | 1.00 | 12.54 | - | - |
| Jax Batch (1 bucket) | **15.54** | **59.57** | 103.69 | 110.59 |
| Jax Batch (4 buckets) | 5.63 | 42.63 | 101.42 | 103.18 |
| Jax Batch (8 buckets) | 3.08 | 25.01 | 93.00 | 120.69 |
| Jax Batch (4 cpus, 4 buckets) | - | 41.51 | **255.56** | 484.05 |
| Jax Batch (8 cpus, 8 buckets) | - | 31.99 | 225.60 | **506.64** |
### Volumetric + ESP alignment — alignments/s
| Approach | n = 10 | n = 100 | n = 1000 | n = 10000 |
|---|---:|---:|---:|---:|
| Jax Volume+ESP | 6.52 | 1.58 | - | - |
| PyTorch (autodiff) | 0.19 | 5.57 | - | - |
| PyTorch (analytical) | 0.29 | 9.84 | - | - |
| Jax Batch (1 bucket) | **14.35** | **72.49** | 81.25 | 123.70 |
| Jax Batch (4 buckets) | 4.54 | 37.55 | 81.90 | 121.04 |
| Jax Batch (8 buckets) | 3.02 | 23.59 | 115.76 | 143.04 |
| Jax Batch (4 cpus, 4 buckets) | - | 40.92 | **299.10** | 468.43 |
| Jax Batch (8 cpus, 8 buckets) | - | 26.68 | 212.16 | **506.33** |
### Pharmacophore alignment — alignments/s
| Approach | n = 10 | n = 100 | n = 1000 | n = 10000 |
|---|---:|---:|---:|---:|
| Jax Pharmacophore (base) | 2.01 | 0.91 | - | - |
| Jax Pharmacophore (vectorized) | 3.73 | 2.00 | - | - |
| PyTorch (autodiff) | 3.77 | 3.68 | - | - |
| PyTorch (analytical) | 11.90 | 9.45 | - | - |
| Jax Batch (1 bucket) | **14.15** | **47.61** | 68.13 | 67.26 |
| Jax Batch (4 buckets) | 5.17 | 44.84 | 89.49 | 106.80 |
| Jax Batch (8 buckets) | 4.25 | 2.29 | 95.04 | 164.53 |
| Jax Batch (4 cpus, 4 buckets) | - | 44.08 | **294.72** | 492.94 |
| Jax Batch (8 cpus, 8 buckets) | - | 26.46 | 216.67 | **532.53** |
### Surface alignment — alignments/s
| Approach | n = 10 | n = 100 | n = 1000 | n = 10000 |
|---|---:|---:|---:|---:|
| Jax Surface | 3.72 | 4.12 | - | - |
| PyTorch (autodiff) | 0.88 | 1.83 | - | - |
| PyTorch (analytical) | 1.78 | 3.70 | - | - |
| Jax Batch (1 bucket) | **4.76** | **4.88** | **6.63** | 4.34 |
| Jax Batch (4 cpus, 4 buckets) | - | - | - | **6.68** |
### Surface + ESP alignment — alignments/s
| Approach | n = 10 | n = 100 | n = 1000 | n = 10000 |
|---|---:|---:|---:|---:|
| Jax Surface+ESP | 2.65 | 3.80 | - | - |
| PyTorch (autodiff) | 0.65 | 2.76 | - | - |
| PyTorch (analytical) | **3.02** | **4.57** | - | - |
| Jax Batch (1 bucket) | 2.85 | 2.30 | **2.71** | - |
## Experiment results (semi-log plot)
