Primitive blocks¶
This notebook demomstrates how to consturct composable Diagrams in Collimator. Let's implement a damped pendulum:
The pendulum dynamics are:
\begin{equation} \begin{bmatrix} \dot{\theta} \\ \dot{\omega} \end{bmatrix} = \begin{bmatrix} \omega \\ -\dfrac{g}{L} \sin(\theta) - \dfrac{b}{m L^2} \omega + \dfrac{u}{m L^2} \end{bmatrix}. \end{equation}
It has a continuous vector state of size 2 (angle $\theta$ and angular velocity $\omega = \dot{\theta}$) and a torque input $u=\tau$. The pendulum parameters are mass $m$, length $L$, damping coefficients $b$, and acceleration due to gravity $g$. The output of the system is the full state $[\theta, \omega]$.
Using primitive blocks, one way to implement a block diagram implementing the above equations is shown below.
We can create the above diagram in code using the Collimator's primitive blocks of Integrator
, Gain
, Adder
, Demultiplexer
, and Trigonometric
blocks. This is shown below.
import jax.numpy as jnp
import matplotlib.pyplot as plt
import collimator
from collimator import DiagramBuilder
from collimator.library import *
from collimator.simulation import ODESolverOptions
def make_pendulum(
x0=[1.0, 0.0],
m=1.0,
g=9.81,
L=1.0,
b=0.0,
name="pendulum",
):
"""
Model of a pendulum's dynamics with damping and external torque.
Equations of motion:
𝜃̇ = ω
ω̇ = -g/L * sin(𝜃) - b/(mL²) * ω + 1/(mL²) * τ
where:
- 𝜃 (theta) is the angular displacement,
- ω (omega) is the angular velocity,
- g is the acceleration due to gravity,
- L is the length of the pendulum,
- b is the damping coefficient,
- m is the mass of the pendulum,
- τ (tau) is the external torque.
The output is the full state [𝜃, ω] .
"""
builder = DiagramBuilder()
Integrator_0 = builder.add(Integrator(x0, name="Integrator_0"))
Demux_0 = builder.add(Demultiplexer(2, name="Demux_0"))
builder.connect(Integrator_0.output_ports[0], Demux_0.input_ports[0])
Sine_0 = builder.add(Trigonometric("sin", name="Sine_0"))
builder.connect(Demux_0.output_ports[0], Sine_0.input_ports[0])
Gain_0 = builder.add(Gain(-g / L, name="Gain_0"))
builder.connect(Sine_0.output_ports[0], Gain_0.input_ports[0])
# Damping
Gain_1 = builder.add(Gain(-b / m / L / L, name="Gain_1"))
builder.connect(Demux_0.output_ports[1], Gain_1.input_ports[0])
Adder_0 = builder.add(Adder(2, name="Adder_0"))
builder.connect(Gain_0.output_ports[0], Adder_0.input_ports[0])
builder.connect(Gain_1.output_ports[0], Adder_0.input_ports[1])
# Add an optional torque input
Gain_2 = builder.add(Gain(1.0 / m / L / L, name="Gain_2"))
Adder_1 = builder.add(Adder(2, name="Adder_1"))
builder.connect(Adder_0.output_ports[0], Adder_1.input_ports[0])
builder.connect(Gain_2.output_ports[0], Adder_1.input_ports[1])
Mux_0 = builder.add(Multiplexer(2, name="Mux_0"))
builder.connect(Demux_0.output_ports[1], Mux_0.input_ports[0])
builder.connect(Adder_1.output_ports[0], Mux_0.input_ports[1])
builder.connect(Mux_0.output_ports[0], Integrator_0.input_ports[0])
# Diagram-level inport
builder.export_input(Gain_2.input_ports[0])
# Full state-output
builder.export_output(Integrator_0.output_ports[0])
return builder.build(name=name)
pendulum = make_pendulum()
pendulum.pprint()
Initialized callback Integrator_0:in_0 with prereqs [] Initialized callback Integrator_0:Integrator_0_ode with prereqs [12] Initialized callback Integrator_0:out_0 with prereqs [2] Initialized callback Demux_0:u_0 with prereqs [] Initialized callback Demux_0:y_0 with prereqs [15] Initialized callback Demux_0:y_1 with prereqs [15] Initialized callback Sine_0:u_0 with prereqs [] Initialized callback Sine_0:y_0 with prereqs [18] Initialized callback Gain_0:u_0 with prereqs [] Initialized callback Gain_0:y_0 with prereqs [20] Initialized callback Gain_1:u_0 with prereqs [] Initialized callback Gain_1:y_0 with prereqs [22] Initialized callback Adder_0:u_0 with prereqs [] Initialized callback Adder_0:u_1 with prereqs [] Initialized callback Adder_0:y_0 with prereqs [24, 25] Initialized callback Gain_2:u_0 with prereqs [] Initialized callback Gain_2:y_0 with prereqs [27] Initialized callback Adder_1:u_0 with prereqs [] Initialized callback Adder_1:u_1 with prereqs [] Initialized callback Adder_1:y_0 with prereqs [29, 30] Initialized callback Mux_0:u_0 with prereqs [] Initialized callback Mux_0:u_1 with prereqs [] Initialized callback Mux_0:y_0 with prereqs [32, 33] Initialized callback pendulum:Gain_2_u_0 with prereqs [] Initialized callback pendulum:Integrator_0_out_0 with prereqs [14] |-- pendulum |-- Integrator_0(id=1) |-- Demux_0(id=2) |-- Sine_0(id=3) |-- Gain_0(id=4) |-- Gain_1(id=5) |-- Adder_0(id=6) |-- Gain_2(id=7) |-- Adder_1(id=8) |-- Mux_0(id=9)
Simulating the system¶
The above pendulum
plant can be simulated on its own. We just need to provide an input. Let's simulate the pendulum by fixing the input torque to zero.
pendulum.input_ports[0].fix_value(0.0)
context = pendulum.create_context()
t = jnp.linspace(0.0, 10.0, 100)
sol = collimator.odeint(
pendulum,
context,
(t[0], t[-1]),
t_eval=t,
ode_options=ODESolverOptions(atol=1e-8, rtol=1e-6),
)
plt.figure(figsize=(7, 3))
plt.plot(t, sol[0], label=[r"$\theta$", r"$\omega$"])
plt.legend()
plt.title("Pendulum plant")
plt.grid()
plt.show()
pendulum = make_pendulum()
builder = DiagramBuilder()
builder.add(pendulum)
sinusoidal_torque = builder.add(
Sine(amplitude=5.0, frequency=5.0, name="sinusoidal_torque")
)
builder.connect(sinusoidal_torque.output_ports[0], pendulum.input_ports[0])
diagram = builder.build()
diagram.pprint()
context = diagram.create_context()
t = jnp.linspace(0.0, 10.0, 100)
sol = collimator.odeint(
diagram,
context,
(t[0], t[-1]),
t_eval=t,
ode_options=ODESolverOptions(max_step_size=0.01, atol=1e-8, rtol=1e-6),
)
plt.figure(figsize=(7, 3))
plt.plot(t, sol[0], label=[r"$\theta$", r"$\omega$"])
plt.legend()
plt.grid()
plt.title("Sinusoidal torque input to the pendulum")
plt.show()
Initialized callback Integrator_0:in_0 with prereqs [] Initialized callback Integrator_0:Integrator_0_ode with prereqs [37] Initialized callback Integrator_0:out_0 with prereqs [2] Initialized callback Demux_0:u_0 with prereqs [] Initialized callback Demux_0:y_0 with prereqs [40] Initialized callback Demux_0:y_1 with prereqs [40] Initialized callback Sine_0:u_0 with prereqs [] Initialized callback Sine_0:y_0 with prereqs [43] Initialized callback Gain_0:u_0 with prereqs [] Initialized callback Gain_0:y_0 with prereqs [45] Initialized callback Gain_1:u_0 with prereqs [] Initialized callback Gain_1:y_0 with prereqs [47] Initialized callback Adder_0:u_0 with prereqs [] Initialized callback Adder_0:u_1 with prereqs [] Initialized callback Adder_0:y_0 with prereqs [49, 50] Initialized callback Gain_2:u_0 with prereqs [] Initialized callback Gain_2:y_0 with prereqs [52] Initialized callback Adder_1:u_0 with prereqs [] Initialized callback Adder_1:u_1 with prereqs [] Initialized callback Adder_1:y_0 with prereqs [54, 55] Initialized callback Mux_0:u_0 with prereqs [] Initialized callback Mux_0:u_1 with prereqs [] Initialized callback Mux_0:y_0 with prereqs [57, 58] Initialized callback pendulum:Gain_2_u_0 with prereqs [] Initialized callback pendulum:Integrator_0_out_0 with prereqs [39] Initialized callback sinusoidal_torque:out_0 with prereqs [1] |-- root |-- pendulum |-- Integrator_0(id=11) |-- Demux_0(id=12) |-- Sine_0(id=13) |-- Gain_0(id=14) |-- Gain_1(id=15) |-- Adder_0(id=16) |-- Gain_2(id=17) |-- Adder_1(id=18) |-- Mux_0(id=19) |-- sinusoidal_torque(id=21)
PD controller for the pendulum¶
As another example of composability, we demonstrate how the pendulum plant can be controlled by a PD control to drive its state to a zero vector. A PD controller block diagram is shown in below.
This can be implemented programatticaly as shown below.
pendulum = make_pendulum()
builder = DiagramBuilder()
builder.add(pendulum)
# PD feedback controller:
kp = 1.0
kd = 1.0
Demux_1 = builder.add(Demultiplexer(2, name="Demux_1"))
Gain_1 = builder.add(Gain(-kp, name="Gain_1"))
Gain_2 = builder.add(Gain(-kd, name="Gain_2"))
Adder_1 = builder.add(Adder(2, name="Adder_1"))
builder.connect(pendulum.output_ports[0], Demux_1.input_ports[0])
builder.connect(Demux_1.output_ports[0], Gain_1.input_ports[0])
builder.connect(Demux_1.output_ports[1], Gain_2.input_ports[0])
builder.connect(Gain_1.output_ports[0], Adder_1.input_ports[0])
builder.connect(Gain_2.output_ports[0], Adder_1.input_ports[1])
builder.connect(Adder_1.output_ports[0], pendulum.input_ports[0])
diagram = builder.build()
diagram.pprint()
Initialized callback Integrator_0:in_0 with prereqs [] Initialized callback Integrator_0:Integrator_0_ode with prereqs [63] Initialized callback Integrator_0:out_0 with prereqs [2] Initialized callback Demux_0:u_0 with prereqs [] Initialized callback Demux_0:y_0 with prereqs [66] Initialized callback Demux_0:y_1 with prereqs [66] Initialized callback Sine_0:u_0 with prereqs [] Initialized callback Sine_0:y_0 with prereqs [69] Initialized callback Gain_0:u_0 with prereqs [] Initialized callback Gain_0:y_0 with prereqs [71] Initialized callback Gain_1:u_0 with prereqs [] Initialized callback Gain_1:y_0 with prereqs [73] Initialized callback Adder_0:u_0 with prereqs [] Initialized callback Adder_0:u_1 with prereqs [] Initialized callback Adder_0:y_0 with prereqs [75, 76] Initialized callback Gain_2:u_0 with prereqs [] Initialized callback Gain_2:y_0 with prereqs [78] Initialized callback Adder_1:u_0 with prereqs [] Initialized callback Adder_1:u_1 with prereqs [] Initialized callback Adder_1:y_0 with prereqs [80, 81] Initialized callback Mux_0:u_0 with prereqs [] Initialized callback Mux_0:u_1 with prereqs [] Initialized callback Mux_0:y_0 with prereqs [83, 84] Initialized callback pendulum:Gain_2_u_0 with prereqs [] Initialized callback pendulum:Integrator_0_out_0 with prereqs [65] Initialized callback Demux_1:u_0 with prereqs [] Initialized callback Demux_1:y_0 with prereqs [88] Initialized callback Demux_1:y_1 with prereqs [88] Initialized callback Gain_1:u_0 with prereqs [] Initialized callback Gain_1:y_0 with prereqs [91] Initialized callback Gain_2:u_0 with prereqs [] Initialized callback Gain_2:y_0 with prereqs [93] Initialized callback Adder_1:u_0 with prereqs [] Initialized callback Adder_1:u_1 with prereqs [] Initialized callback Adder_1:y_0 with prereqs [95, 96] |-- root |-- pendulum |-- Integrator_0(id=23) |-- Demux_0(id=24) |-- Sine_0(id=25) |-- Gain_0(id=26) |-- Gain_1(id=27) |-- Adder_0(id=28) |-- Gain_2(id=29) |-- Adder_1(id=30) |-- Mux_0(id=31) |-- Demux_1(id=33) |-- Gain_1(id=34) |-- Gain_2(id=35) |-- Adder_1(id=36)
context = diagram.create_context()
t = jnp.linspace(0.0, 10.0, 100)
sol = collimator.odeint(
diagram,
context,
(t[0], t[-1]),
t_eval=t,
ode_options=ODESolverOptions(atol=1e-8, rtol=1e-6),
)
plt.figure(figsize=(7, 3))
plt.plot(t, sol[0], label=[r"$\theta$", r"$\omega$"])
plt.legend()
plt.grid()
plt.title("PD controlled pendulum")
plt.show()
We see that the PD controller does indeed drives the Pendulum state to a zero vector.