hanging_chain_dynamics.py

Go to the documentation of this file.

from typing import Union
import casadi as cs
import numpy as np
from casadi import vertcat as vc
 
 
class HangingChain:
    def __init__(self, N: int, dim: int):
        self.NN = N
        self.dimdim = dim
 
        self.y1y1 = cs.SX.sym("y1", dim * N, 1)  # state: balls 1→N positions
        self.y2y2 = cs.SX.sym("y2", dim * N, 1)  # state: balls 1→N velocities
        self.y3y3 = cs.SX.sym("y3", dim, 1)  # state: ball  1+N position
        self.uu = cs.SX.sym("u", dim, 1)  # input: ball  1+N velocity
        self.yy = vc(self.y1y1, self.y2y2, self.y3y3)  # full state vector
 
        self.mm = cs.SX.sym("m")  # mass
        self.DD = cs.SX.sym("D")  # spring constant
        self.LL = cs.SX.sym("L")  # spring length
 
        self.paramsparams = vc(self.mm, self.DD, self.LL)
 
        self.gg = np.array([0, 0, -9.81] if dim == 3 else [0, -9.81])  # gravity
        self.x0x0 = np.zeros((dim, ))  # ball 0 position
        self.x_endx_end = np.eye(1, dim, 0).ravel()  # ball N+1 reference position
 
    def dynamics(self, Ts=0.05):
        y, y1, y2, y3, u = self.yy, self.y1y1, self.y2y2, self.y3y3, self.uu
        dist = lambda xa, xb: cs.norm_2(xa - xb)
        N, d = self.NN, self.dimdim
        p = self.paramsparams
 
        # Continuous-time dynamics y' = f(y, u; p)
 
        f1 = [y2]
        f2 = []
        for i in range(N):
            xi = y1[d * i:d * i + d]
            xip1 = y1[d * i + d:d * i + d * 2] if i < N - 1 else y3
            Fiip1 = self.DD * (1 - self.LL / dist(xip1, xi)) * (xip1 - xi)
            xim1 = y1[d * i - d:d * i] if i > 0 else self.x0x0
            Fim1i = self.DD * (1 - self.LL / dist(xi, xim1)) * (xi - xim1)
            fi = (Fiip1 - Fim1i) / self.mm + self.gg
            f2 += [fi]
        f3 = [u]
 
        f_expr = vc(*f1, *f2, *f3)
        self.ff = cs.Function("f", [y, u, p], [f_expr], ["y", "u", "p"], ["y'"])
 
        # Discretize dynamics y[k+1] = f_d(y[k], u[k]; p)
 
        opt = {"tf": Ts, "simplify": True, "number_of_finite_elements": 4}
        intg = cs.integrator("intg", "rk", {
            "x": y,
            "p": vc(u, p),
            "ode": f_expr
        }, opt)
 
        f_d_expr = intg(x0=y, p=vc(u, p))["xf"]
        self.f_df_d = cs.Function("f_d", [y, u, p], [f_d_expr], ["y", "u", "p"],
                               ["y+"])
 
        return self.f_df_d
 
    def state_to_pos(self, y):
        N, d = self.NN, self.dimdim
        rav = lambda x: np.array(x).ravel()
        xdim = lambda y, i: np.concatenate(
            ([0], rav(y[i:d * N:d]), rav(y[-d + i])))
        if d == 3:
            return (xdim(y, 0), xdim(y, 1), xdim(y, 2))
        else:
            return (xdim(y, 0), xdim(y, 1), np.zeros((N + 1, )))
 
    def input_to_matrix(self, u):
        """
        Reshape the input signal from a vector into a dim × N_horiz matrix (note
        that CasADi matrices are stored column-wise and NumPy arrays row-wise)
        """
        if isinstance(u, np.ndarray):
            return u.reshape((self.dimdim, u.shape[0] // self.dimdim), order='F')
        else:
            return u.reshape((self.dimdim, u.shape[0] // self.dimdim))
 
    def simulate(self, N_sim: int, y_0: np.ndarray, u: Union[np.ndarray, list,
                                                             cs.SX.sym],
                 p: Union[np.ndarray, list, cs.SX.sym]):
        if isinstance(u, list):
            u = np.array(u)
        if isinstance(u, np.ndarray):
            if u.ndim == 1 or (u.ndim == 2 and u.shape[1] == 1):
                if u.shape[0] == self.dimdim:
                    u = np.tile(u, (N_sim, 1)).T
        return self.f_df_d.mapaccum(N_sim)(y_0, u, p)
 
    def initial_state(self):
        N, d = self.NN, self.dimdim
        y1_0 = np.zeros((d * N))
        y1_0[0::d] = np.arange(1, N + 1) / (N + 1)
        y2_0 = np.zeros((d * N))
        y3_0 = np.zeros((d, ))
        y3_0[0] = 1
 
        y_null = np.concatenate((y1_0, y2_0, y3_0))
        u_null = np.zeros((d, ))
 
        return y_null, u_null
 
    def generate_cost_fun(self, α=25, β=1, γ=0.01):
        N, d = self.NN, self.dimdim
        y1t = cs.SX.sym("y1t", d * N, 1)
        y2t = cs.SX.sym("y2t", d * N, 1)
        y3t = cs.SX.sym("y3t", d, 1)
        ut = cs.SX.sym("ut", d, 1)
        yt = cs.vertcat(y1t, y2t, y3t)
 
        L_cost = α * cs.sumsqr(y3t - self.x_endx_end) + γ * cs.sumsqr(ut)
        for i in range(N):
            xdi = y2t[d * i:d * i + d]
            L_cost += β * cs.sumsqr(xdi)
        return cs.Function("L_cost", [yt, ut], [L_cost])