archimedes.tree.struct¶

archimedes.tree.struct(

cls: T | None = None,

**kwargs,

) → T | Callable¶

Decorator to convert a class into a tree-compatible frozen dataclass.

This decorator creates a structured data class that can be seamlessly used with Archimedes’ tree functions. The class will be registered with the tree system, allowing its instances to be flattened, mapped over, and transformed while preserving its structure.

Parameters:

cls (type, optional) – The class to convert into a tree-compatible dataclass.
**kwargs (dict) – Additional keyword arguments passed to dataclasses.dataclass(). By default, frozen=True is set unless explicitly overridden.

Returns:

decorated_class – The decorated class, now a frozen dataclass registered as tree-compatible.

Return type:

type

Notes

The “frozen” attribute makes the class immutable, meaning that once an instance is created, its fields cannot be modified. This is useful for ensuring that the state of the object remains consistent during operations. The replace() method allows you to create modified copies of the object with new values for specific fields.

Fields are automatically classified as either “data” (dynamic values that change during operations) or “static” (configuration parameters). By default, all fields are treated as data unless marked with field(static=True).

The decorated class:

Is frozen (immutable) by default
Has a replace() method for creating modified copies
Will be properly handled by tree.flatten(), tree.map(), etc.
Can be nested within other tree nodes (structs, dicts, tuples, etc.)

Examples

>>> import archimedes as arc
>>> import numpy as np
>>>
>>> @arc.struct
>>> class Vehicle:
...     # Dynamic state variables (included in transformations)
...     position: np.ndarray
...     velocity: np.ndarray
...
...     # Static configuration parameters (preserved during transformations)
...     mass: float = arc.field(static=True, default=1000.0)
...     drag_coef: float = arc.field(static=True, default=0.3)
...
...     def kinetic_energy(self):
...         return 0.5 * self.mass * np.sum(self.velocity**2)
>>>
>>> # Create an instance
>>> car = Vehicle(
...     position=np.array([0.0, 0.0]),
...     velocity=np.array([10.0, 0.0]),
... )
>>>
>>> # Create a modified copy
>>> car2 = car.replace(position=np.array([5.0, 0.0]))
>>>
>>> # Apply a transformation (only to dynamic fields)
>>> scaled = arc.tree.map(lambda x: x * 2, car)
>>> print(scaled.position)    # [0. 0.] -> [0. 0.]
>>> print(scaled.velocity)    # [10. 0.] -> [20. 0.]
>>> print(scaled.mass)        # 1000.0 (unchanged)
>>>
>>> # Nested structs
>>> @arc.struct
>>> class System:
...     vehicle1: Vehicle
...     vehicle2: Vehicle
...
...     def total_energy(self):
...         return self.vehicle1.kinetic_energy() + self.vehicle2.kinetic_energy()
>>>
>>> system = System(car, car2)
>>> # This transformation applies to all dynamic fields in the entire hierarchy
>>> scaled_system = arc.tree.map(lambda x: x * 0.5, system)