Batch the torch Cell List

[CompRhys]

Somewhat pointless optimization as alchemiops is way better choice in most cases. Serves some pedagocial value but code hard to read batching into a 4th dimension. Left serial code around as a fall back for single systems as batched was like 30% slower.

============================================================
CORRECTNESS CHECKS
============================================================
  PASS [single_Ar_fcc]: 428 pairs match exactly
  PASS [batch_Ar+Si]: 2148 pairs match exactly
  PASS [batch_periodic+nonperiodic]: 440 pairs match exactly
  PASS [mixed_pbc_TFT]: 140 pairs match exactly
  PASS [size_disparity_8vs32]: 5144 pairs match exactly
  PASS [large_108atoms]: 11556 pairs match exactly
  PASS [batch_3x_Ar]: 1284 pairs match exactly
  PASS [self_interaction]: 86 pairs match exactly

ALL CORRECTNESS CHECKS PASSED

============================================================
TIMING COMPARISON
============================================================
Config                                 Orig ms   Batch ms    Speedup   Orig MiB  Batch MiB
-------------------------------------------------------------------------------------------
1x Ar FCC (4 atoms)                      0.523      0.451      1.16x      442.3      437.5
4x Ar FCC batch (4 atoms each)           1.905      0.913      2.09x      443.1      436.8
1x Ar 3x3x3 (108 atoms)                  1.430      1.427      1.00x      442.4      437.5
8x Ar FCC batch                          3.506      1.147      3.06x      438.9      437.0
16x Ar FCC batch                         6.983      1.435      4.87x      438.5      440.2
32x Ar FCC batch                        12.860      1.639      7.84x      439.4      441.1

[CompRhys]

Used [2,2,2] supercells to make it clearer that the batching helps.

(base) ➜  torch-sim git:(batched-cell-nl-torch) ✗     uv run --with-editable . examples/benchmarking/neighborlists.py \
        --source wbm --n-structures 100 --device cpu
Reading ASE Atoms from zip_filename='/Users/radical-rhys/.cache/matbench-discovery/wbm/2024-08-04-wbm-initial-atoms.extxyz.zip': 100%|█| 256963
{
  "source": "wbm",
  "n_structures": 100,
  "n_atoms": 6560,
  "cutoff_angstrom": 5.0,
  "seed": 0,
  "device": "cpu",
  "dtype": "float64",
  "n_repeats": 7,
  "load_s": 21.585,
  "backends": [
      {
      "nl_backend": "torch_linked_cell",  # serial current
      "n_pairs": 178482,
      "median_nl_s": 0.172468,
      "timings_s": [
        0.182636,
        0.173212,
        0.171743,
        0.171029,
        0.172265,
        0.172468,
        0.176357
      ],
      "atoms_per_s": 38035.9
    },
    {
      "nl_backend": "torch_linked_cell_batched",
      "n_pairs": 178482,
      "median_nl_s": 0.054274,
      "timings_s": [
        0.055468,
        0.054274,
        0.055742,
        0.054746,
        0.053414,
        0.052164,
        0.05415
      ],
      "atoms_per_s": 120869.3
    },
    {
      "nl_backend": "torch_n2",  # this is batched and padded.
      "n_pairs": 178482,
      "median_nl_s": 0.141705,
      "timings_s": [
        0.139523,
        0.140596,
        0.142936,
        0.140069,
        0.145619,
        0.145857,
        0.141705
      ],
      "atoms_per_s": 46293.2
    },
    {
      "nl_backend": "vesin",
      "n_pairs": 178482,
      "median_nl_s": 0.022575,
      "timings_s": [
        0.024077,
        0.022916,
        0.022501,
        0.022493,
        0.022575,
        0.023135,
        0.022549
      ],
      "atoms_per_s": 290585.3
    },
    {
      "nl_backend": "alchemi_n2",
      "n_pairs": 178482,
      "median_nl_s": 0.014629,
      "timings_s": [
        0.015335,
        0.01494,
        0.014477,
        0.013703,
        0.014629,
        0.014699,
        0.014465
      ],
      "atoms_per_s": 448414.2
    },
    {
      "nl_backend": "alchemi_cell",
      "n_pairs": 178482,
      "median_nl_s": 0.02513,
      "timings_s": [
        0.02513,
        0.0249,
        0.024896,
        0.025311,
        0.025492,
        0.024845,
        0.025243
      ],
      "atoms_per_s": 261039.5
    }
  ]
}