Support for model architecture manipulation

Summary

This MR adds functionality for modifying the architecture of a trained NEP model so that it can serve as the starting point for further training, avoiding a full restart from random parameters. Five new Model methods cover the main use cases:

Method	What it does
augment()	Expand neuron count, enable higher-body descriptor terms, or add a charge output head
prune()	Reduce neuron count, disable descriptor terms, or remove the charge output head
remove_species()	Remove one or more species (refactored; also fixes bugs)
keep_species()	Retain a subset of species, dropping all others
add_species()	Extend an existing model with new chemical species

All methods return a new Model object, leaving the source unchanged. They share a common scheme for the SNES restart file: existing parameters receive adaptive sigma (max(sigma_floor, sigma_factor × |μ|)), preserving trained values and keeping near-zero parameters dormant; new parameters are initialized to match the fresh-model behavior of the nep executable of the GPUMD package (uniform μ ∈ [-1, 1], σ = sigma_new).

A training_parameters property derives the correct nep.in architecture fields from any modified model. This addresses a critical issue: GPUMD reads nep.in (not nep.txt) to determine the model size when loading nep.restart for training. Using a stale nep.in after augmentation causes a stride mismatch that scrambles all parameter values, producing training-from-scratch-level initial loss even with sigma --> 0.

---

Changes

Model.augment() — new method

Applies one or more structural changes atomically and returns a new Model:

• Neuron expansion (n_neuron): pads w0, b0, w1 with zeros; new rows/columns receive sigma_new in the restart.
• Descriptor expansion (l_max_4b, l_max_5b, has_q_112/123/233/134): adds columns to w0 and appends entries to q_scaler. Increasing an already-enabled l_max value only updates the header — no new descriptor dimensions are created.
• Charge head (charge_head=True): promotes potential --> potential_with_charges; adds w1_charge per species (initialized to zero) and initializes sqrt_epsilon_infinity = 1 (ε∞ = 1, the vacuum value).
• Requires restart; raises ValueError otherwise.

Model.prune() — new method

Applies one or more structure-reducing changes atomically and returns a new Model:

• Neuron reduction (n_neuron): ranks neurons by importance score mean_s(||w0_s[n,:]||₂ × |w1_s[n]|) (charge models add |w1_charge_s[n]|), keeps the top-scoring neurons, and discards the rest.
• Descriptor reduction (l_max_4b=0, has_q_112=False, etc.): removes the corresponding descriptor columns from w0 and q_scaler. Setting a non-zero l_max to a smaller non-zero value only updates the header.
• Charge head removal (charge_head=True): demotes potential_with_charges --> potential; removes w1_charge per species and sqrt_epsilon_infinity from the restart.
• Surviving parameters receive adaptive sigma. Raises ValueError if the target would expand the model (use augment() instead) or if restart_parameters is not loaded.

Model.remove_species() — refactored

• Breaking change (intentional): now returns a new Model instead of modifying in-place, consistent with the other methods.
• Bug fix: n_ann_output_weights and n_bias were computed incorrectly for potential_with_charges models (missing ×2 multiplier and wrong bias count), causing write-->read roundtrips to fail.
• Bug fix: typewise radial_cutoff/angular_cutoff lists were not pruned when a species was removed, so surviving species inherited the deleted species' cutoff values.
• Bug fix: sqrt_epsilon_infinity was dropped from restart parameters when pruning a charge model.
• Applies adaptive sigma to surviving parameters when restart is loaded.

Model.keep_species() — new method

Convenience complement to remove_species(). Accepts the species to retain and derives the removal list automatically. Thin wrapper around remove_species().

Model.add_species() — new method

• Adds a per-species ANN sub-network (NEP4 only; raises for NEP3) and all new (s1, s2) descriptor weight pairs for each new element.
• New ANN weights and descriptor coefficients are initialized with μ drawn uniformly from [-1, 1], matching GPUMD's fresh-model initialization. Charge-specific parameters (w1_charge) are kept at μ = 0.
• sigma_new defaults to 0.1, matching GPUMD's sigma0 default. Accepts a seed parameter for reproducible initialization.
• Handles typewise cutoffs: radial_cutoff/angular_cutoff arguments are required when the model uses per-type cutoffs.

Model.training_parameters — new property

Returns the model's architecture fields (version, type, cutoff, n_max, basis_size, l_max, neuron, and optionally zbl) in the format accepted by write_nepfile():

params = read_nepfile('nep.in')
params.update(aug.training_parameters)
write_nepfile(params, '.')

Bug fix — get_parity_data transpose in calorine/nep/io.py

The transpose of force/BEC/atomic-virial arrays (extracted_property.T) was applied inside the for select in selection: loop, flipping the orientation on every pass. For multi-component selections (e.g. ['x', 'y']) this silently returned wrong data on even-length selections. The transpose is now applied once, before the loop.

Internal refactoring — calorine/nep/model.py

• _adaptive_sigma() extracted as a module-level function (was duplicated inside each method).
• _apply_adaptive_sigma_to_restart() extracted as a module-level helper, replacing ~25-line blocks repeated across augment, prune, add_species, and remove_species. The unified helper also corrects a pre-existing inconsistency: remove_species previously used the raw max() formula for b1/sqrt_epsilon_infinity scalars while all other methods used _adaptive_sigma.
• _recalculate_parameter_counts() extracted as a module-level helper, replacing duplicate parameter-count blocks across the same four methods. The unified helper also adds the previously missing polarizability doubling in add_species, augment, and prune.

Display fixes (Model.__str__, Model._repr_html_)

restart_parameters previously printed its full nested dict; it now shows "available" or "not available".

Tutorial notebook (tutorials/modifying_models.ipynb)

New notebook added to the Training NEP models section of the documentation. Covers all five methods and the training_parameters property with concrete before/after output. Includes:

• Bar charts of parameter-count change under augmentation and pruning
• Sigma-distribution histograms: new vs. existing species (after add_species()), and before vs. after pruning (adaptive \sigma \propto |\mu| for surviving parameters)
• Chained-operation examples (remove_species --> add_species, augment --> add_species, prune --> add_species)
• training_parameters / write_nepfile integration example

---

Tests

130 tests pass; calorine/nep/model.py reaches 100% coverage.