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.

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.

Affected training recipes (mace-foundations)

• mace_omol/mace-omol.sh already passes --total_charge_key="charge" --total_spin_key="spin" and --energy_key=REF_energy --forces_key=REF_forces. Those charge/spin flags are currently dead and start working with this change; its energy/forces keys are unaffected.

• mace_polar_1/mace-polar-1-{medium,large}.sh do not pass charge/spin keys, and set --forces_key='forces'. Companion PR: mace_polar_1: pass charge/spin keys and use REF_energy/REF_forces mace-foundations#79.

  ⚠️ Coupling worth noting for reviewers: on current develop, --forces_key='forces' is silently ignored and forces fall back to the REF_forces default, so they load correctly by accident. Once overrides are honored (this PR), --forces_key='forces' actively selects the empty atoms.arrays["forces"] and forces load as zero. The companion PR updates the polar scripts to REF_forces (and adds the charge/spin keys), matching mace-omol.sh. Verified empirically both ways.

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 a 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.
• 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.

Both fail on the corresponding pre-fix code.

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
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 = 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: total_charge = 0.0, total_spin = 1.0 — After: total_charge = -1.0, total_spin = 3.0.

[aacostadiaz]

Thanks for the fix and the very thorough write-up!

LGTM! I'll go ahead and merge this now.