Source code for brom_drake.utils.initial_condition_manager.manager

from dataclasses import dataclass, field
import numpy as np
from pydrake.all import (
    Context,
    Frame,
    ModelInstanceIndex,
    MultibodyPlant,
    RigidTransform,
)
from typing import List

# Internal Imports
from brom_drake.utils.initial_condition_manager.initial_condition import (
    InitialCondition,
)




@dataclass
class InitialConditionManager:
    _ic_tuples: List[InitialCondition] = field(default_factory=list)

    def add_initial_condition(self, ic_tuple: InitialCondition):
        self._ic_tuples.append(ic_tuple)

    def add_initial_configuration(
        self, model_instance_index: ModelInstanceIndex, configuration: np.ndarray
    ):
        ic_tuple = InitialCondition(
            model_instance_index=model_instance_index, configuration=configuration
        )
        self._ic_tuples.append(ic_tuple)

    def add_initial_pose(
        self, model_instance_index: ModelInstanceIndex, pose_wrt_parent: RigidTransform
    ):
        ic_tuple = InitialCondition(
            model_instance_index=model_instance_index, pose_wrt_parent=pose_wrt_parent
        )
        self._ic_tuples.append(ic_tuple)



    def set_all_initial_conditions(
        self, plant: MultibodyPlant, diagram_context: Context = None
    ):
        """
        *Description*

        Set all initial conditions (both the pose and configuration)
        for the given plant.
        """
        self.set_all_initial_configurations(plant, diagram_context=diagram_context)
        self.set_all_initial_poses(plant, diagram_context=diagram_context)




    def set_all_initial_configurations(
        self, plant: MultibodyPlant, diagram_context: Context = None
    ):
        """Set all known initial configurations for the given plant."""
        for ic_tuple in self._ic_tuples:
            ic_tuple.set_initial_configuration(plant, diagram_context=diagram_context)




    def set_all_initial_poses(
        self, plant: MultibodyPlant, diagram_context: Context = None
    ):
        """Set all known initial poses for the given plant."""
        for ic_tuple in self._ic_tuples:
            ic_tuple.set_initial_pose(plant, diagram_context=diagram_context)