forked from pz4kybsvg/Conception
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
100 lines
3.7 KiB
100 lines
3.7 KiB
import unittest
|
|
|
|
from pydrake.examples import (
|
|
PendulumGeometry,
|
|
PendulumInput,
|
|
PendulumParams,
|
|
PendulumPlant,
|
|
PendulumState,
|
|
)
|
|
from pydrake.geometry import SceneGraph
|
|
from pydrake.systems.analysis import Simulator
|
|
from pydrake.systems.framework import DiagramBuilder
|
|
|
|
|
|
class TestPendulum(unittest.TestCase):
|
|
def test_plant(self):
|
|
pendulum = PendulumPlant()
|
|
self.assertEqual(pendulum.get_input_port(), pendulum.get_input_port(0))
|
|
self.assertEqual(pendulum.get_state_output_port(),
|
|
pendulum.get_output_port(0))
|
|
context = pendulum.CreateDefaultContext()
|
|
self.assertIsInstance(pendulum.get_state(context=context),
|
|
PendulumState)
|
|
self.assertIsInstance(pendulum.get_mutable_state(context=context),
|
|
PendulumState)
|
|
self.assertIsInstance(pendulum.get_parameters(context=context),
|
|
PendulumParams)
|
|
self.assertIsInstance(pendulum.get_mutable_parameters(context=context),
|
|
PendulumParams)
|
|
|
|
def test_input(self):
|
|
input = PendulumInput()
|
|
input.set_tau(tau=1.)
|
|
self.assertEqual(input.tau(), 1.)
|
|
input_tau = input.with_tau(tau=2.)
|
|
self.assertEqual(input_tau.tau(), 2.)
|
|
|
|
def test_params(self):
|
|
params = PendulumParams()
|
|
params.set_mass(mass=1.)
|
|
params.set_length(length=2.)
|
|
params.set_damping(damping=3.)
|
|
params.set_gravity(gravity=4.)
|
|
self.assertEqual(params.mass(), 1.)
|
|
self.assertEqual(params.length(), 2.)
|
|
self.assertEqual(params.damping(), 3.)
|
|
self.assertEqual(params.gravity(), 4.)
|
|
params_mass = params.with_mass(mass=5.)
|
|
params_length = params.with_length(length=6.)
|
|
params_damping = params.with_damping(damping=7.)
|
|
params_gravity = params.with_gravity(gravity=8.)
|
|
self.assertEqual(params_mass.mass(), 5.)
|
|
self.assertEqual(params_length.length(), 6.)
|
|
self.assertEqual(params_damping.damping(), 7.)
|
|
self.assertEqual(params_gravity.gravity(), 8.)
|
|
|
|
def test_state(self):
|
|
state = PendulumState()
|
|
state.set_theta(theta=1.)
|
|
state.set_thetadot(thetadot=2.)
|
|
self.assertEqual(state.theta(), 1.)
|
|
self.assertEqual(state.thetadot(), 2.)
|
|
state_theta = state.with_theta(theta=3.)
|
|
state_thetadot = state.with_thetadot(thetadot=4.)
|
|
self.assertEqual(state_theta.theta(), 3.)
|
|
self.assertEqual(state_thetadot.thetadot(), 4.)
|
|
|
|
def test_geometry(self):
|
|
builder = DiagramBuilder()
|
|
plant = builder.AddSystem(PendulumPlant())
|
|
scene_graph = builder.AddSystem(SceneGraph())
|
|
geom = PendulumGeometry.AddToBuilder(
|
|
builder=builder, pendulum_state_port=plant.get_output_port(0),
|
|
scene_graph=scene_graph)
|
|
builder.Build()
|
|
self.assertIsInstance(geom, PendulumGeometry)
|
|
|
|
def test_simulation(self):
|
|
# Basic constant-torque pendulum simulation.
|
|
pendulum = PendulumPlant()
|
|
|
|
# Create the simulator.
|
|
simulator = Simulator(pendulum)
|
|
context = simulator.get_mutable_context()
|
|
|
|
# Set an input torque.
|
|
input = PendulumInput()
|
|
input.set_tau(1.)
|
|
pendulum.get_input_port().FixValue(context, input)
|
|
|
|
# Set the initial state.
|
|
state = context.get_mutable_continuous_state_vector()
|
|
state.set_theta(1.)
|
|
state.set_thetadot(0.)
|
|
|
|
# Simulate (and make sure the state actually changes).
|
|
initial_state = state.CopyToVector()
|
|
simulator.AdvanceTo(1.0)
|
|
self.assertFalse((state.CopyToVector() == initial_state).any())
|