Add model training with mini-batches

bin/emle-train

@@ -90,6 +90,23 @@ parser.add_argument(
 parser.add_argument(
     "--plot-data", type=str, metavar="name.mat", default=None, help="Data for plotting"
 )
+parser.add_argument(
+    "--use-minibatch",
+    action="store_true",
+    help="Use minibatch training",
+)
+parser.add_argument(
+    "--batch-size",
+    type=int,
+    metavar="",
+    default=1024,
+    help="Batch size for minibatch training",
+)
+parser.add_argument(
+    "--shuffle",
+    action="store_true",
+    help="Shuffle training data",
+)
 
 parser.add_argument("output", type=str, help="Output model file")
 args = parser.parse_args()
@@ -128,4 +145,7 @@ trainer.train(
     print_every=args.print_every,
     model_filename=args.output,
     plot_data_filename=args.plot_data,
+    use_minibatch=args.use_minibatch,
+    batch_size=args.batch_size,
+    shuffle=args.shuffle,
 )

emle/train/_ivm.py

@@ -24,6 +24,7 @@
 
 import numpy as _np
 import torch as _torch
+from loguru import logger as _logger
 
 from ._gpr import GPR as _GPR
 from ._utils import pad_to_max as _pad_to_max
@@ -143,3 +144,124 @@ def perform_ivm(aev_mols, z_mols, atom_ids, species, thr, sigma):
         aev_ivm_allz = [aev_z[ivm_idx_z] for aev_z, ivm_idx_z in zip(aev_allz, ivm_idx)]
 
         return ivm_mol_atom_ids_padded, aev_ivm_allz
+
+    @staticmethod
+    def perform_ivm_lazy(
+        aev_filenames,
+        z_batches,
+        atom_id_batches,
+        species,
+        sigma,
+        thr,
+        n_max=None,
+        device="cuda",
+    ):
+        """
+        Calculate representative feature vectors for each species from batches loaded from disk.
+
+        Parameters
+        ----------
+        aev_filenames: list of str
+            Paths to saved AEV batches.
+
+        z_batches: list of torch.Tensor(N_ATOMS_BATCH,)
+            Atomic numbers for each batch.
+
+        atom_id_batches: list of torch.Tensor(N_ATOMS_BATCH, 2)
+            (mol_idx, atom_idx) for each atom in each batch.
+
+        species: torch.Tensor(N_SPECIES,)
+            Unique species to select from.
+
+        sigma: float
+            Kernel width.
+
+        thr: float
+            Variance threshold.
+
+        n_max: int or None
+            Max number of reference AEVs to select per species.
+
+        device: torch.device
+            Device to use.
+
+        Returns
+        -------
+        ivm_mol_atom_ids_padded: torch.Tensor(N_SPECIES, MAX_N_REF, 2)
+            IVM selected (mol_idx, atom_idx) per species.
+
+        aev_ivm_allz: list of torch.Tensor(N_REF, AEV_DIM)
+            AEV features for reference atoms.
+        """
+        selected = {z.item(): [] for z in species}
+        selected_ids = {z.item(): [] for z in species}
+
+        for z in species:
+            z = z.item()
+            iter_count = 0
+            while True:
+                max_var = -float("inf")
+                best_aev = None
+                best_id = None
+
+                for fname, z_batch, id_batch in zip(
+                    aev_filenames, z_batches, atom_id_batches
+                ):
+                    aev_batch = _torch.load(fname, map_location=device).to(device)
+                    z_b = z_batch.to(device)
+                    id_b = id_batch.to(device)
+                    mask = (z_b == z).to(device)
+
+                    if not mask.any():
+                        continue
+
+                    aev_z = aev_batch[mask]
+                    ids_z = id_b[mask]
+
+                    if len(selected[z]) == 0:
+                        local_var = _torch.ones(len(aev_z), device=device)
+                    else:
+                        aev_sel = _torch.stack(selected[z]).to(device)
+                        k_sel = _GPR._aev_kernel(aev_sel, aev_sel)
+                        k_inv = _torch.linalg.inv(
+                            k_sel + _torch.eye(len(aev_sel), device=device) * sigma**2
+                        )
+                        k = _GPR._aev_kernel(aev_z, aev_sel)
+                        local_var = 1 - _torch.sum(k @ k_inv * k, dim=1)
+
+                    top_idx = _torch.argmax(local_var)
+                    if local_var[top_idx] > max_var:
+                        max_var = local_var[top_idx].item()
+                        best_aev = aev_z[top_idx].detach().cpu()
+                        best_id = ids_z[top_idx].detach().cpu()
+
+                if max_var < thr or (n_max is not None and len(selected[z]) >= n_max):
+                    break
+
+                selected[z].append(best_aev.to(device))
+                selected_ids[z].append(best_id)
+                iter_count += 1
+                _logger.info(
+                    f"IVM for species {z}: Iter {iter_count}: max var = {max_var:.5f}, n_selected = {len(selected[z])}"
+                )
+
+        ivm_mol_atom_ids_padded = _pad_to_max(
+            [
+                _torch.stack(v).to(device)
+                if len(v) > 0
+                else _torch.empty(0, 2, dtype=_torch.long, device=device)
+                for v in selected_ids.values()
+            ],
+            value=-1,
+        )
+
+        aev_ivm_allz = [
+            _torch.stack(v).to(device)
+            if len(v) > 0
+            else _torch.empty(
+                0, selected[next(iter(selected))][0].shape[-1], device=device
+            )
+            for v in selected.values()
+        ]
+
+        return ivm_mol_atom_ids_padded, aev_ivm_allz