archimedes.tree.struct.register_dataclass¶

archimedes.tree.struct.register_dataclass(nodetype: Typ, data_fields: Sequence[str] | None = None, meta_fields: Sequence[str] | None = None, drop_fields: Sequence[str] = ()) → Typ¶

This function registers a dataclass type as a pytree node, with control over which fields are treated as leaf data versus metadata. Fields marked as metadata are excluded from transformations and treated as static configuration.

Parameters:

nodetype (type) – The dataclass type to register as a pytree node.
data_fields (sequence of str, optional) – Names of fields that should be treated as data (leaf values). If None and the type is a dataclass, fields are inferred based on metadata.
meta_fields (sequence of str, optional) – Names of fields that should be treated as metadata. If None and the type is a dataclass, fields are inferred based on metadata.
drop_fields (sequence of str, optional) – Names of fields to exclude from both data and metadata categories.

Returns:

nodetype – The input type, now registered as a pytree node.

Return type:

type

Notes

When to use: - For fine-grained control over how dataclass fields are handled in pytree ops - When you need some fields to be treated as static configuration rather than data - For advanced customization of pytree behavior for complex data models

Usually, instead of using this function directly, you’ll want to use the struct.pytree_node decorator which handles registration automatically and allows field classification via the struct.field(static=True) parameter. This function is mainly used internally to register the decorated classes.

Data fields are included when flattening a pytree and are considered leaf values that can be transformed. Meta fields are static configuration not included in transformations but preserved during reconstruction.

Raises:

TypeError – If data_fields and meta_fields aren’t both specified when either is specified, or if they are both None and the type is not a dataclass.
ValueError – If the specified fields don’t match the actual dataclass fields with init=True.

Examples

>>> import archimedes as arc
>>> from dataclasses import dataclass
>>> import numpy as np
>>>
>>> # Define a dataclass
>>> @dataclass
>>> class Vehicle:
...     position: np.ndarray  # Data field - changes during simulation
...     velocity: np.ndarray  # Data field - changes during simulation
...     mass: float           # Metadata - configuration parameter
...     name: str             # Metadata - configuration parameter
>>>
>>> # Register with explicit field classification
>>> arc.tree.register_dataclass(
...     Vehicle,
...     data_fields=["position", "velocity"],
...     meta_fields=["mass", "name"]
... )
>>>
>>> # Create an instance
>>> car = Vehicle(
...     position=np.array([0.0, 0.0]),
...     velocity=np.array([10.0, 0.0]),
...     mass=1500.0,
...     name="sedan"
... )
>>>
>>> # Flatten - only data fields are included in leaves
>>> leaves, treedef = arc.tree.flatten(car)
>>> print(leaves)
[array([0., 0.]), array([10., 0.])]
>>>
>>> # When unflattening, metadata is preserved
>>> doubled_leaves = [leaf * 2 for leaf in leaves]
>>> new_car = arc.tree.unflatten(treedef, doubled_leaves)
>>> print(new_car.velocity)  # Data field multiplied by 2
[20., 0.]
>>> print(new_car.mass)      # Metadata preserved
1500.0