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:

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.

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.

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.

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.

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.

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:

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.