Implement versioned serialization for Primitive hierarchies

[joelaforet]

Is your feature request related to a problem? Please describe.

MuPT still does not provide a native way to save and reload a Primitive representation. This is a real workflow bottleneck for systems that take meaningful effort to build, especially SAAMR-compliant systems used in mupt-examples.

A common workflow is:
build locally -> transfer to HPC -> run simulation -> retrieve trajectory -> analyze later

Today, if the original Python session is gone, the MuPT representation is gone as well. That means users lose the tree-like relational hierarchy of Primitive objects and their children, including the information needed to recover MuPT-native analysis pathways such as primitive_to_mdanalysis().

This issue is intended as a concrete implementation-oriented continuation of #11, which discussed canonical forms and serialization at a broader level. The goal here is to define and implement a practical, versioned serialization/deserialization path for Primitive-based representations.

Describe the solution you'd like

Add a versioned serialization/deserialization protocol for Primitive hierarchies.

The main requirement is faithful preservation of the relational structure of a MuPT representation, including:

• parent/child hierarchy
• child handles / labels
• connector information
• internal and external connector mappings
• topology / graph connectivity
• geometry / shape information
• element assignments
• relevant metadata

Initial support can focus on SAAMR-compliant systems, but the design should ideally extend to more general Primitive trees.

A simple public API would be useful, for example:

• save(primitive, path)
• load(path) -> Primitive

The on-disk format should be versioned so future schema evolution is explicit and old files can be handled gracefully.

A schema-driven implementation is likely the safest path. In particular, a Pydantic-based schema layer is worth serious consideration, not as a replacement for MuPT core runtime classes, but as a typed serialization boundary:

• MuPT runtime objects (Primitive, Connector, TopologicalStructure, shapes, registries) remain unchanged
• dedicated serialization DTOs define the persisted structure
• MuPT objects are converted to DTOs on save, and DTOs are converted back to MuPT objects on load

This approach is attractive because MuPT needs a rigorously typed, recursive, versionable schema for tree-structured molecular representations. Pydantic is well suited to nested validation, recursive models, explicit serialization rules, and JSON Schema generation. JSON Schema would also give maintainers a concrete contract for .mupt contents.

YAML may also be worth considering as the human-readable wire format if inspectability is a priority, but the schema should remain independent of the encoding choice. JSON would also be acceptable. The key requirement is preserving the full Primitive tree and its relational cross-references in a stable, versioned form.

Describe alternatives you've considered

• pickle / dill: easy to prototype, but fragile, unsafe for shared files, and not appropriate for a stable interchange format
• direct ad hoc JSON/YAML dumps of runtime objects: possible, but likely brittle without an explicit schema layer
• making Primitive and related runtime classes themselves Pydantic models: likely too invasive, since MuPT core objects carry runtime behavior (NodeMixin, UniqueRegistry, networkx graphs, mutable geometry) that should not be conflated with the persisted representation
• relying only on downstream exported files: preserves simulation artifacts, but not the MuPT-native hierarchy and relationships
• folding this entirely into canonicalization work from Canonical forms and serialization for Primitives #11: likely too broad; serialization is valuable on its own and can move first

Additional context

• Primitive is centered in mupt/mupr/primitives.py
• primitive_to_mdanalysis() already exists in mupt/interfaces/mdanalysis/exporters.py, so reloadable MuPT representations would immediately enable downstream analysis workflows
• mupt-examples is the main user-facing surface where this pain is visible
• Canonical forms and serialization for Primitives #11 should be treated as prior discussion/context; this issue is meant to request an actionable implementation path
• preserving the tree-like hierarchy of primitives and children is the core design constraint
• a likely implementation detail is that some currently freeform fields, especially metadata: dict[Hashable, Any], may need tighter serialization rules for v1

Acceptance criteria

• Add a versioned save/load path for Primitive-based representations
• Preserve hierarchy, handles/labels, topology, connectors, and relevant metadata through round-trip serialization
• Support SAAMR-compliant systems as the initial target
• Use an explicit schema layer rather than relying on direct pickling of runtime objects
• Fail clearly on unsupported objects or schema versions
• Add tests covering build -> save -> load -> analyze workflows
• Reference Canonical forms and serialization for Primitives #11 as prior context in implementation/PR discussion

[timbernat]

Thanks for the detailed write-up! I agree this is a high-value feature; a large part of why I've held off on implementing something like this is because of the number of objects (which you've listed) with very disparate internal data all of which need to be preserved in a uniform and lossless way. I'll direct some more development attention to this!

freeform fields, especially metadata: dict[Hashable, Any], may need tighter serialization rules for v1

This is another big challenge. With the disparate objects in the repr we at least know what they should contain, but with metadata all bets are off; as the adage goes in software engineering, "all input is evil until proven otherwise". I've put together frameworks before for doing this sort of thing in general for JSON: the main idea is you register converters for each type of data serialized and preserve info about type when serializing to ensure the process is invertible (i.e. that the write function is injective with respect to type).

pickle / dill: easy to prototype, but fragile, unsafe for shared files

Agreed, as any change to the underlying class between write and subsequent read would render even completely-valid serialized hierarchies unloadable.

JSON/YAML dumps of runtime objects

Widely-used file formats are definitely a big plus. As human-readable as these particular formats are, I don't suspect inspecting molecular systems from file alone will be a particularly-common task, especially for large, complex polymer systems which can't be read off "at a glance" like small molecules (this is the sort of automation we are trying to provide!). More likely one would inspect/edit a molecular system after loading the file into an interpreter session or notebook, so compact (e.g. binary) formats are likely offer a greater benefit than human-readable ones.

I've been considering HDF5 as an appealing storage format, since it permits nested fields, high-density data, and metadata storage in a relatively compact binary format which is fairly widely recognized. Another (substantially worse) alternative I've considered is to serialize each hierarchy component (e.g. Connector, Primitive, etc) one-per-file in a serialized directory. The latter has a number of obvious problems, and the only really useful nugget there is to handle serialization per-component, implementing a generic interface via Protocol templates) and having each type dictate how read/write is actually handled.

#11 : likely too broad; serialization is valuable on its own and can move first

In its entirety absolutely, too broad for here, though there are elements of canonicalization that will be necessary for serializing to be useful. The main one that comes to mind is determining consistent rules for ordering Primitives even across levels, which will be essential for keeping "nearby" neighbors near to one another in transit, and for preserving invariants for downstream formats which are more index-based, such as MDAnalysis.

[joelaforet]

I don't have strong preferences for the file type we settle on, only that one exists and that it accomplishes the main goal of being a lossless "round trip" from Primitive -> (On Disk Somehow) -> Identical Primitive in memory.

I can prototype using HDF5 and json (using patterns from the Polymerist repo) if there are no blockers + other suggestions.