Repository: Genesis-Embodied-AI/Genesis
File: genesis/utils/array_class.py (Python)

Code before:

def get_entities_state(solver):
    shape = solver._batch_shape(solver.n_entities_)
    kwargs = {
        "hibernated": V(dtype=gs.ti_int, shape=shape),
    }

    if use_ndarray:
        return StructEntitiesState(**kwargs)
    else:

        @ti.data_oriented
        class ClassEntitiesState:
            def __init__(self):
                for k, v in kwargs.items():
                    setattr(self, k, v)

        return ClassEntitiesState()


# =========================================== StaticRigidSimConfig ===========================================


@dataclasses.dataclass
class StaticRigidSimDummy:
    para_level: int = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    use_hibernation: bool = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    use_contact_island: bool = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    batch_links_info: bool = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    batch_dofs_info: bool = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    batch_joints_info: bool = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    enable_mujoco_compatibility: bool = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    enable_multi_contact: bool = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    enable_self_collision: bool = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    enable_adjacent_collision: bool = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    enable_collision: bool = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    box_box_detection: bool = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    integrator: int = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    sparse_solve: bool = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    solver_type: int = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})

    # dynamic properties
    substep_dt: float = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    iterations: int = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    tolerance: float = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    ls_iterations: int = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    ls_tolerance: float = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    n_equalities: int = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    n_equalities_candidate: int = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    hibernation_thresh_acc: float = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})
    hibernation_thresh_vel: float = dataclasses.field(metadata={FIELD_METADATA_CACHE_VALUE: True})

Review comment (performance):
I can see why you're going to ask me to make it possible to set `FIELD_METADATA_CACHE_VALUE` conceptually on an entire dataclass at some point :)

Code after:

def get_entities_state(solver):
    shape = solver._batch_shape(solver.n_entities_)
    kwargs = {
        "hibernated": V(dtype=gs.ti_int, shape=shape),
    }

    if use_ndarray:
        return StructEntitiesState(**kwargs)
    else:

        @ti.data_oriented
        class ClassEntitiesState:
            def __init__(self):
                for k, v in kwargs.items():
                    setattr(self, k, v)

        return ClassEntitiesState()


# =========================================== StaticRigidSimConfig ===========================================

cache_value = partial(dataclasses.field, metadata={FIELD_METADATA_CACHE_VALUE: True})


@dataclasses.dataclass
class StaticRigidSimCacheKey:
    para_level: int = cache_value()
    use_hibernation: bool = cache_value()
    batch_links_info: bool = cache_value()
    batch_dofs_info: bool = cache_value()
    batch_joints_info: bool = cache_value()
    enable_mujoco_compatibility: bool = cache_value()
    enable_multi_contact: bool = cache_value()
    enable_adjacent_collision: bool = cache_value()
    enable_collision: bool = cache_value()
    box_box_detection: bool = cache_value()
    integrator: int = cache_value()
    sparse_solve: bool = cache_value()
    solver_type: int = cache_value()


def get_static_rigid_sim_cache_key(solver):
    kwargs = {
        "para_level": solver.sim._para_level,
        "use_hibernation": getattr(solver, "_use_hibernation", False),
        "batch_links_info": getattr(solver._options, "batch_links_info", False),
        "batch_dofs_info": getattr(solver._options, "batch_dofs_info", False),
        "batch_joints_info": getattr(solver._options, "batch_joints_info", False),
        "enable_mujoco_compatibility": getattr(solver, "_enable_mujoco_compatibility", False),
        "enable_multi_contact": getattr(solver, "_enable_multi_contact", True),