Honor the CLI key specification in LMDBDataset

[lollcat]

Problem

LMDBDataset.__getitem__ hardcodes KeySpecification.from_defaults() when building each config:

config = config_from_atoms(
    atoms,
    key_specification=KeySpecification.from_defaults(),   # ignores CLI overrides
)

So CLI key overrides such as --total_charge_key / --total_spin_key are silently ignored for .aselmdb / .lmdb datasets — load_dataset_for_path never passes the parsed key specification to LMDBDataset, and __getitem__ would ignore it even if it did.

This bites OMol-style data. The defaults map total_charge <- "total_charge" and total_spin <- "total_spin", but OMol stores these under data["charge"] / data["spin"] (surfaced by get_atoms as atoms.info["charge"] / ["spin"]). The mapped keys are absent, so total_charge / total_spin fall back to 0.0 / 1.0 — every system is loaded as a neutral closed-shell singlet, regardless of the flags on the command line.

For example, mace-foundations/mace_omol/mace-omol.sh passes --total_charge_key="charge" --total_spin_key="spin", but these currently have no effect on the LMDB path.

Fix

Thread head_config.key_specification through load_dataset_for_path into the LMDBDataset constructors, and honor it in __getitem__ — falling back to KeySpecification.from_defaults() when none is supplied, so existing behavior is unchanged. LMDBDataset is the only on-the-fly reader that applied config_from_atoms with a hardcoded spec; HDF5 datasets bake the key mapping in at write time.

Tests

Two regression tests in tests/test_lmdb_database.py, each writing a one-molecule .aselmdb with data={"charge": -1, "spin": 3} (the OMol layout) and asserting the key spec built from --total_charge_key=charge --total_spin_key=spin yields total_charge == -1, total_spin == 3:

• test_lmdb_dataset_honors_key_specification — constructs LMDBDataset(..., key_specification=...) directly (also checks the no-spec default still gives 0 / 1). Guards the lmdb_dataset.py half; fails on pre-fix code (total_charge defaults to 0).
• test_load_dataset_for_path_forwards_key_specification — calls load_dataset_for_path with a minimal SimpleNamespace(head_name=..., key_specification=...). Guards the run_train_utils.py dispatch half; fails when only that half is reverted.

Reproduction

Self-contained (synthesizes a one-molecule .aselmdb, no external data):

import tempfile, os
from ase.build import molecule
from mace.tools import AtomicNumberTable
from mace.data.utils import KeySpecification, update_keyspec_from_kwargs
from mace.data.lmdb_dataset import LMDBDataset
from mace.tools.fairchem_dataset.lmdb_dataset_tools import LMDBDatabase

CHARGE, SPIN = -1, 3

# write a 1-molecule .aselmdb with charge/spin under the row `data` dict (as OMol does)
tmp = tempfile.mkdtemp()
path = os.path.join(tmp, "sample.aselmdb")
with LMDBDatabase(path, readonly=False) as db:
    db.write(molecule("H2O"), data={"charge": CHARGE, "spin": SPIN})

# key spec built by `--total_charge_key=charge --total_spin_key=spin`
key_spec = KeySpecification.from_defaults()
update_keyspec_from_kwargs(key_spec, {"total_charge_key": "charge", "total_spin_key": "spin"})

ds = LMDBDataset(path, r_max=5.0, z_table=AtomicNumberTable([1, 8]), key_specification=key_spec)
sample = ds[0]
print("total_charge =", float(sample.total_charge), " (expected", CHARGE, ")")
print("total_spin   =", float(sample.total_spin),   " (expected", SPIN, ")")

Before (develop):

total_charge = 0.0   (expected -1)
total_spin   = 1.0   (expected 3)

After (this branch):

total_charge = -1.0   (expected -1)
total_spin   = 3.0   (expected 3)

Checks

• black (26.5.1) and isort (8.0.1): clean on the changed source files.
• pylint with the repo's hook args: 10.00/10 on both changed files.
• pytest tests/test_lmdb_database.py: 3 passed (existing + 2 new). tests/test_data.py: 9 passed. Full suite collects (518 tests) with no import errors.
• tests/ is excluded from the formatting hooks in .pre-commit-config.yaml, so the added test is not reformatted.