Skip to content

Models and force fields

fm.load is fastmace.model.loading.load_pristine.

load_pristine 

load_pristine(
    source,
    head=None,
    *,
    sha256=None,
    trusted=False,
    download=False,
    charge=None,
    spin_multiplicity=None,
    checkpoint=None,
    environ=None,
)

fm.load (SPEC §7.1): the checkpoint as a pristine model, e3nn in float64 on the CPU.

source is a label ("mh1-r2scan"), an alias with head= ("mh1"), a path or a URL. A multi-head checkpoint needs head= or a label that names one; otherwise ValueError lists its heads. A path or URL is unpickled only with a matching sha256= or trusted=True; a label's file must match the table (SPEC §7.6). download=True fetches a missing file into the models cache and checks it before it is kept. The trust decision and the hash are in provenance(). The plan converts the model to its backend and precision; the pristine model itself is never changed (fm.ForceField). Before MACE is imported, the OEq runtime is prepared where the hybrid is possible (prepare_hybrid_if_possible; the outcome is oeq_runtime). The expandable-segments allocator is configured first (configure_allocator, a no-op if PYTORCH_CUDA_ALLOC_CONF is set): preparing the hybrid can touch CUDA, and the allocator reads the variable only once (5-T07).

ForceField

ForceField(model, atoms=None, *, precision="float32", rewrites="auto", species=None).

  • model: an fm.load result (a LoadedModel, normally the pristine e3nn float64 CPU model) or a label, which is loaded with fm.load. A path needs fm.load first, because the trust decision needs one (§7.6);
  • atoms: optional; its species are the default superset (§7.1 passes it). Without it the superset is the species of the first structure the force field is used with;
  • precision: the force precision, "float32" or "float64" (§7.3);
  • rewrites: "auto" or a RewriteConfig; resolved per plan;
  • species: the declared superset (§8.5), symbols or atomic numbers, compared as a set; a species without an embedding raises ModelElementError. It may be larger than the structure's species (phase 3, rule 8): every species-indexed table is built for it, so an atom can become any superset species in place (3-T02);
  • dispersion: None, "d3bj" or an fm.D3 (SPEC §7.7, §11.2). "d3bj" expands from the model's training functional (fastmace.longrange.d3.expand: ValueError when it is unknown or already contains dispersion); an explicit record whose functional differs from the training one is kept with a warning in dispersion_warnings. The complete record is dispersion and part of the identity. Refused (PlanRefused): the ATM term (d3-three-body) and PolarMACE (d3-polar, decision 41). Plans carry the record and put the D3 term into their force field (fastmace.model.dispersed, 3b-T04).

traits property 

traits

The model traits (capabilities.ModelTraits) of the pristine model.

identity property 

identity

The state.ForceFieldIdentity (SPEC §10.6); None until the superset is known.

Its force precision is the effective one (effective_precision), so a resume on the same MACE re-plans the same way and the identity matches. Known limitation (2-T04): it assumes float64_energy=True. A plan with PrecisionConfig(float64_energy=False), which is for A/B against the research engine only, on a MACE whose float64 energy forward does not hash runs at float32 but records a float64 identity.

superset_for 

superset_for(atoms) -> tuple[int, ...]

The superset, fixed from atoms at first use if none was declared (atomic numbers, sorted).

effective_precision 

effective_precision(float64_energy=True)

(precision, reason): the force precision plans run at, and why it differs from precision.

On a MACE whose float64 energy forward does not hash, a float32 force field runs at float64 (fastmace.rewrites.versions.effective_precision; plan task 2-T04); reason is then a Reason, else None.

execution_model 

execution_model(plan)

The model converted to plan's backend, precision and device (a LoadedModel).

Called by the planner only, once per plan. The checkpoint is loaded again, its hash checked against the pristine model's, so the pristine model is never changed and no two plans share a model (the per-model state that SPEC §8.7's instance isolation is about stays per plan). On CUDA with cuEq or hybrid, a later plan of the same conversion gets a private copy of the first plan's converted model instead (fastmace.model.templates; opt-1 OPT-T06), unless EngineConfig.model_template_cache is False.

template_count 

template_count() -> int

How many converted models this force field keeps as templates (opt-1 OPT-T06).

release_templates 

release_templates() -> int

Drop the converted models kept as templates; returns the bytes their tensors held.

The next plan of each conversion converts the checkpoint again. Use it to give the memory back, for example before an out-of-memory retry (opt-1 OPT-T06).

evaluation_plan 

evaluation_plan(atoms)

The one-off evaluation plan evaluate uses for structures like atoms (planned once).

evaluate 

evaluate(
    atoms: Any, stress: bool = False
) -> tuple[float, Any, Any]

(energy, forces, stress) of atoms as host copies, through a one-off minimal plan; stress is None unless asked for. The energy is accumulated in float64 (§7.3).

Workload dataclass

What the run will do (SPEC §9.2). Only max_temperature_K is required.

  • max_temperature_K: the highest thermostat target of the run, K; sizes the skin and the bucket headroom of MC and species-change workloads (rule 11; plain MD keeps the 0.3 Å skin). Simulation.run refuses a schedule above it (§7.6);
  • species_changes: swaps, transmutations or species events (§8.5);
  • atom_count_changes: reserved; the planner refuses it in 1.0 (§9.3 rule 7, §8.6);
  • min_mc_temperature_K: the lowest MC target temperature, K; the ΔE gate threshold 0.01 k_B T_min (§12.3). Required when mc_moves is non-empty, and positive;
  • host_callbacks_every: user Python every n steps (rule 9); None = no callbacks;
  • needs_per_atom: per-atom outputs the run needs ("energies", "charges", ...);
  • mc_moves: declared move types ("swap", "transmute", "displace");
  • reproducible: bitwise within the scope of §9.3; the planner refuses what it cannot guarantee. Qualified (2-T16, planner.reproducible): plain cuEq with the unfused convolution, the MH-1 architecture and the pinned cuequivariance_ops version;
  • mc_energy_precision: §10.4; "auto" = the force precision if the ΔE gate passes;
  • max_displacement_A: the largest displacement δ, Å, any Displace move of the run may use (adaptive δ included). Required when "displace" is in mc_moves; rule 11 adds 2δ to the skin (§10.4), and attaching a Displace with a larger delta_A is refused.