Add support for weighted-sum heads on multi-head checkpoints

[zyxwwxyz]

Summary

Adds a way to attach fixed-coefficient linear-combination heads to a trained multi-head checkpoint - e.g. an energy/mix head equal to 0.25 * energy/pbe + 0.75 * energy/pbesol without any retraining. Linear-combination coefficients are hard-coded in the head, so the head contains no free parameters, and needs no retraining.

• New upet._weighted_sum module exposing create_weighted_sum_checkpoint / extract_wrapped_checkpoint (also re-exported from upet), plus a CLI (python -c "from upet._weighted_sum import _main; _main()") driven by a YAML spec of sources: {head: coefficient} per new head.
• Combination is computed on the wrapped model's own physical predictions, so forces/stresses for a combined energy head fall out of a single backward pass.
• Coefficients are used as given; it is possible to normalize them so that their sum is one by setting normalize_coefficients: true (raises an error if the provided coefficients sum to zero).
• specs (for create_weighted_sum_checkpoint/WeightedSumModel) takes WeightedSumHead instances (a small dataclass: sources, normalize_coefficients, description) rather than an untyped dict, so a misspelled or unsupported field fails immediately with a TypeError.
• Each combination head gets a description (auto-generated from sources/coefficients, or overridable) stored on the exported model.
• src/upet/_models.py: _load_model_with_custom_heads teaches upet's loading path to recognize these wrapped checkpoints (via architecture_name) and reconstruct them through WeightedSumModel.load_checkpoint, since they aren't a registered metatrain architecture and can't go through mtt export/mtt eval or plain load_model directly.
• Adds pyyaml as a new dependency (for the YAML spec format).
• Docs: new docs/src/weighted-sum-heads.rst page with CLI and Python API usage, linked from index.rst/fine-tuning.rst; the two top-level functions are also documented in api.rst alongside get_upet/save_upet.
• Compatibility with metatrain 2026.3.1 is verified.

Example

# wsum_heads.yaml
heads:
  energy/mix:
    sources:
      energy/pbe: 0.25
      energy/pbesol: 0.75
    description: "25/75 PBE/PBEsol mix"
  energy/diff:
    sources:
      energy/pbe: 1.0
      energy/pbesol: -1.0
    normalize_coefficients: false
  energy/calibrated-mix:
    sources:
      energy/pbe: 1
      energy/pbesol: 3
    normalize_coefficients: true

python -c "from upet._weighted_sum import _main; _main()" \
    model.ckpt wsum_heads.yaml model-wsum.ckpt

model-wsum.ckpt still works with get_upet/save_upet via checkpoint_path=...; use extract_wrapped_checkpoint to pull the original checkpoint back out for further fine-tuning.

Notes for reviewers

• A weighted-sum checkpoint embeds the original checkpoint unmodified, so it still works with get_upet/save_upet via checkpoint_path=.... To fine-tune further, pull the original back out first with extract_wrapped_checkpoint.
• Tests in tests/upet/test_weighted_sum.py cover spec parsing (YAML + inline --head), normalize_coefficients rescaling (including the zero-sum-coefficients error), WeightedSumHead's own field validation, description defaults/overrides, checkpoint round-tripping, the _main CLI end to end, and forces/stresses correctness against the wrapped model.

Test plan

• pytest tests/upet/test_weighted_sum.py
• Full test suite (a0828f9 added more rigorous coverage)
• Docs build (docs/src/weighted-sum-heads.rst renders correctly, sphinx refs resolve)

[Luthaf]

Not sure if this is worth having as a new public function. Could it go through the normal get_upet & co?

[zyxwwxyz]

I am not finding a nice way to do it. The two functions are a distinct checkpoint-authoring step that is not directly compatible with standard get_upet & co. Maybe I can tighten the namespace/docs to make that separation clearer.

[Luthaf]

ah, I misunderstood what extract_wrapped_checkpoint was doing. I'm not fully convinced it is required (the unwrapped checkpoint should exist elsewhere anyway), but it is fine to keep it!

[zyxwwxyz]

I would honestly keep it as well. If I want to publish a model on huggingface, I think it is easier to publish the .ckpt file with weighted heads only, and leave users (and the future me) the opportunity of stripping these heads and do further fine-tuning, rather than uploading both the models with and without weighted sum heads.

[Luthaf]

I don't love this API, because one does not know what needs to go in the specs without reading the docs. I would rather have separate keyword arguments with descriptive names directly

[zyxwwxyz]

I tried to solve it by replacing WeightedSumModel's untyped dict specs with a typed WeightedSumHead dataclass.

[abmazitov]

I don't understand the need for such a functionality, especially within a CLI command. I we need a weighted sum - we can get one by running run_model method of the underlying MetatomicCalculator.
Current implementation just goes orthogonally to a scope and design of a upet repo

[zyxwwxyz]

I don't understand the need for such a functionality, especially within a CLI command. I we need a weighted sum - we can get one by running run_model method of the underlying MetatomicCalculator. Current implementation just goes orthogonally to a scope and design of a upet repo

The reason of this PR is to implement the pet-exp protocol (https://arxiv.org/html/2604.24607v1), in order to do things like MD/phonon calculations with lammps/i-pi/ase using a weighted sum of various heads of a model as a force field.
My original idea was to do one single edit in one code to allow for this, and upet seemed to me the best place. If I am not wrong, changing MetatomicCalculator would only work in ase, and I would still need to define a custom i-pi driver and edit lammps pair_metatomic.cpp to reproduce the same behaviour in these codes. We can discuss about moving this to metatomic, but I have the feeling it would be a bit orthogonal there as well (as I think this weighted sum modification is PET-specific, while metatomic is a broader-purpose library)

About the CLI command, I think this is a minor side issue that we can easily remove.