# Adding Model Terms This page is a checklist for adding a new `UFPTerm`. It is written for both human contributors and coding agents: follow the order here, keep behavior compatible with existing terms, and avoid broad hot-path refactors unless a benchmark or speed gate motivates them. ## Start With The Contract Every term should make these decisions explicit before implementation: - Which base class fits the term: `OneBodyTerm`, `PairTerm`, `ThreeBodyTerm`, or plain `UFPTerm`. - Whether the term needs a neighbor list, a full directed neighbor list, or charge/spin state from `UFPInputState`. - Whether the term provides analytic forces. If not, model APIs can derive forces from differentiable energy when called with `derive_forces=True`. - Whether any exposed parameters are linear coefficients. Only linear coefficients should be exposed through `ParameterBlock` for least-squares fitting. - Whether the term must round-trip through workflow checkpoints. If yes, add a schema codec. - Whether training should use a custom optimizer group such as `embedding`, `state`, or `charge_spin`. Use extension-facing imports from `ufp.terms.contracts`: ```python from ufp.terms.contracts import ( ParameterBlock, ParameterBlockCacheChannel, ParameterBlockCacheDescriptor, TermInputRequirements, UFPTerm, copy_parameter_data, ) ``` Built-in term modules may also use shared internal helpers such as `ufp.terms.categories`, `ufp.terms._constraints`, and `ufp.terms._edge`. Keep external extension code on `ufp.terms.contracts` unless a helper is promoted to the public API. ## Input Requirements Declare input requirements on the term and call `self.validate_inputs(inputs)` at the start of `forward()` and any assembly path that needs the same inputs. This keeps missing-neighbor-list and missing-state errors consistent across terms. ```python class ChargeAwareTerm(UFPTerm): @property def input_requirements(self) -> TermInputRequirements: return TermInputRequirements( full_neighbor_list=True, state_fields=("atomic_charges", "system_charges"), ) def forward(self, inputs): self.validate_inputs(inputs) charges = inputs.atomic_charges assert charges is not None ... ``` Do not store charges, spin moments, or related state only in `metadata`. Use the first-class state fields on `UFPInput`: - `atomic_charges` - `atomic_spin_moments` - `system_charges` - `system_spin_moments` ASE extraction already looks for common arrays and info keys when inputs are prepared. Add new state fields to `UFPInputState` before depending on them in a term. ## Category Layout Reuse the shared category helpers for all pair and source-distinguished triplet layouts. This preserves selector ordering, cache-channel keys, projection compatibility, and checkpoint behavior. ```python from ufp.terms.categories import ( active_pair_mask, active_triplet_mask, pair_categories, triplet_categories, ) pairs = pair_categories(self.atomic_types, symmetric=True) active = active_pair_mask(pairs, active_pairs=active_pairs, symmetric=True) triplets = triplet_categories(self.atomic_types) active_triplets = active_triplet_mask(triplets, active_triplets=active_triplets) ``` For hot paths, keep tensor loops local when that avoids extra materialization. The category helpers are for layout and validation, not a mandate to unify every runtime kernel. ## Parameter Blocks Expose a `ParameterBlock` only when the block has stable semantics and is linear in the values the solver will write. The block should provide: - a stable `name`, `kind`, and human-readable `label`; - `read` and `write` callbacks that do not replace parameter objects; - `regularization_group` when ridge or shape penalties should target it; - `assembler` only if the term has least-squares support; - `fittable=False` for nonlinear, generated, or diagnostic parameters. ```python def parameter_blocks(self) -> tuple[ParameterBlock, ...]: return ( ParameterBlock( name="coeffs", kind="myterm", shape=tuple(int(dim) for dim in self.coeffs.shape), read=lambda: self.coeffs, write=lambda values: copy_parameter_data(self.coeffs, values), label=f"myterm[{self.atomic_types}]", regularization_group="myterm", fittable=self.fittable, frozen=self.frozen, assembler=self.assemble_linear_block, ), ) ``` Do not expose nonlinear constrained parameters as ordinary least-squares coefficients. Use PyTorch training for those terms, or expose a generated linear block only when writes can be mapped back unambiguously. ## Cache Descriptors Attach a `ParameterBlockCacheDescriptor` when a full parameter block can be stored in a reusable semantic assembled cache. The descriptor must describe the meaning of the block, not its current layout index. ```python ParameterBlock( ..., cache_descriptor=ParameterBlockCacheDescriptor( family={ "kind": "myterm_spline", "atomic_types": [int(value) for value in self.atomic_types], "spline": str(self.spline), "first_knot": float(self.first_knot), "knot_spacing": float(self.knot_spacing), "coeff_shape": [int(dim) for dim in self.coeffs.shape[1:]], }, channels=tuple( ParameterBlockCacheChannel( kind="triplet", values=self.triplet_categories[index], start=int(index) * channel_size, stop=(int(index) + 1) * channel_size, ) for index in self._active_triplet_indices ), ), ) ``` Best practices for cache descriptors: - Include all hyperparameters that change design-matrix columns in `family`. - Use stable channel values such as `(Z1, Z2)` or `(source, Zj, Zk)`. - Mark `reusable=False` or omit the descriptor for partial, generated, nonlinear, or layout-dependent blocks. - Keep the old hard-coded descriptor path working until the term has tests proving descriptor and fallback metadata match. Least-squares cache projection only uses reusable descriptors for full selected blocks. Partial selectors fall back to exact-layout caches. ## Schema Codecs Workflow checkpoints reconstruct model architecture from registered schema codecs before loading the state dict. Add a codec for every term that should be checkpoint-reconstructable. ```python from ufp.workflows import register_term_schema_codec def _encode_myterm(term: object, provider_index: int | None) -> dict[str, object]: typed = cast(MyTerm, term) return { "coeff_shape": list(typed.coeffs.shape), "active_channels": [list(channel) for channel in typed.active_channels], "trainable": bool(typed.coeffs.requires_grad), "fittable": bool(typed.fittable), "frozen": bool(typed.frozen), } def _decode_myterm(entry, context) -> MyTerm: return MyTerm( cutoff=float(entry["cutoff"]), atomic_types=context.atomic_types, coeffs=torch.zeros(tuple(entry["coeff_shape"]), dtype=context.dtype), active_channels=[tuple(item) for item in entry.get("active_channels", ())], trainable=bool(entry.get("trainable", True)), fittable=bool(entry.get("fittable", True)), frozen=bool(entry.get("frozen", False)), ) register_term_schema_codec(MyTerm, _encode_myterm, _decode_myterm) ``` Keep codecs declarative: - Store shapes and constructor metadata, not large parameter values. The state dict carries values. - Preserve constructor defaults for old checkpoints by using `entry.get(...)`. - Include coefficient provider indices only for terms that already support `AlchemicalCoefficients`. - Avoid changing class names in schemas. If a rename is unavoidable, add an explicit compatibility migration in `ufp.workflows.models` and test an old payload. ## Optimizer Groups Terms can request custom workflow optimizer grouping: ```python class EmbeddingTerm(UFPTerm): @property def optimizer_group(self) -> str | None: return "embedding" ``` Current workflow defaults understand `embedding`, `state`, and `charge_spin`, and callers can pass additional `term_weight_decays`. Existing one-body, pair, and three-body weight-decay arguments still apply to terms in those module groups unless a custom group consumes the parameters first. ## Tests To Add For a new term, add focused tests before broad integration tests: - Forward pass shape, dtype, device, and force behavior. - Missing neighbor-list and missing state errors through `input_requirements`. - Category ordering and active-category canonicalization if the term has channels. - Checkpoint schema round trip through `save_checkpoint()` and `load_model_from_checkpoint()`. - `ParameterLayout` and coefficient selector behavior for every exposed block. - Least-squares assembly agreement with direct evaluation when the term is linear. - Cache descriptor metadata equivalence to any legacy fallback, plus projection tests when channel reuse matters. - Optimizer grouping if the term defines `optimizer_group`. - Relevant speed gates when touching `UFPInput`, pair hot paths, three-body bucketing, or least-squares assembly. Keep tests small enough to diagnose the contract that failed. Do not rely on a single end-to-end workflow test to cover schema, cache, and input-state behavior. ## Implementation Checklist 1. Add the term class with explicit `cutoff`, `atomic_types`, `provides_forces`, `input_requirements`, and dtype/device handling. 2. Add or reuse category helpers and state fields before implementing the runtime path. 3. Add `parameter_blocks()` only for truly linear exposed coefficients. 4. Add cache descriptors only for reusable full-block semantics. 5. Register a schema codec if checkpoint reconstruction should work. 6. Add optimizer grouping only when default onebody/pair/threebody/default groups are not enough. 7. Add tests for the specific contracts above. 8. Run targeted tests, `tox -e type`, and the relevant speed gates before changing hot-path kernels.