Per-element (per-atom) force loss weighting

[mikyneb]

Closes #1539.

Adds opt-in per-atom force loss weighting via a new xyz arrays column, REF_forces_atomic_weights (one scalar per atom). Per-element weighting is achieved by expanding an element→weight map into this column during preprocessing. Defaults to uniform, so existing training runs are unchanged.

Changes

As noted in the issue thread, loading already supports arbitrary per-atom arrays — the work is loss-side plus registering the key.

• default_keys.py / arg_parser.py / data/utils.py: register forces_atomic_weights as an arrays key, add --forces_atomic_weights_key to the train and preprocess parsers.
• data/atomic_data.py: forces_atomic_weights as a first-class node field, shape (n_atoms, 1), defaulting to ones in from_config (keeps batches homogeneous when only some configs carry the column, and keeps the feature strictly opt-in).
• modules/loss.py: a getattr-based helper multiplies the weights into every force path — mean_squared_error_forces, mean_normed_error_forces, conditional_mse_forces, conditional_huber_forces (via a new optional arg), and both branches of WeightedHuberEnergyForcesStressLoss. Covers the weighted, forces_only, stress, huber, universal, and l1l2energyforces losses.

Out of Scope

• Validation RMSE_F stays unweighted (filter_nonzero_weight still masks by config, not atom).
• Per-component (n, 3) weights — scalar-per-atom is the clean v1; mean_normed_error_forces can't use per-component weights anyway.

Usage

wmap = {29: 10.0, 6: 1.0}  # up-weight Cu vs C
atoms.arrays["REF_forces_atomic_weights"] = np.array(
    [wmap[z] for z in atoms.get_atomic_numbers()]
)

Tests

New tests/test_forces_atomic_weights.py (uniform-weights-match-default, linear scaling, zero-weight masking, per-element ratio, mixed/legacy batches, xyz round-trip, DDP-reduction consistency). Existing test_modules, test_data, test_run_train pass.