Single-file 2D Physics simulation: Multi-Body, Finite Element method, Co-simulation

sim.py

# context is available at https://wf34.ws/works/learnings-of-multibody-dynamics-workshop-2025/
# requirements: pip3 install typed-argument-parser numpy matplotlib
# 

#!/usr/bin/env python3
import bisect
import copy
import os
import typing
import numpy as np
from tap import Tap
from matplotlib import pyplot as plt
from matplotlib.patches import Circle, Rectangle

MB = 'mb'
FEA = 'fea'
CO = 'co'

class SimArgs(Tap):
    mode: typing.Literal[MB, FEA, CO]
    dest_dir: str

    def process_args(self):
        self.abs_dest_dir = os.path.abspath(self.dest_dir)
        self.gif_path = os.path.join(os.path.dirname(self.abs_dest_dir), f'_{self.mode}.gif')

        if os.path.exists(self.dest_dir):
            print(f'wont overwrite {self.dest_dir}')
            exit(1)

        if os.path.exists(self.gif_path):
            print(f'wont overwrite {self.gif_path}')
            exit(2)

        os.mkdir(self.dest_dir)


def render(mode, free_bodies, fixed_bodies, world_bounds, out_dir, index, timestamp):
    dest_file = os.path.join(out_dir, f'{index:06d}.png')
    fig, ax = plt.subplots(1, 1, figsize=(6, 6))
    ax.set_xlim(*world_bounds[0])
    ax.set_ylim(*world_bounds[1])
    ax.set_aspect('equal')
    ax.grid(True)

    for fb in fixed_bodies:
            x, y = fb.get_coords()
            ax.plot(x, y, color='k', linewidth=1)

    if mode in {MB, CO}:
        for bdy in free_bodies:
            theta = bdy.pos[0]
            R = bdy.radius
            circle = Circle(bdy.pos[1:], R, fill=False, edgecolor='k', linewidth=1)

            h_bar = Rectangle(bdy.pos[1:] - [R, 0.], 2 * R, 0., angle=np.degrees(theta), fill=False, edgecolor='k', linewidth=1, rotation_point='center')
            v_bar = Rectangle(bdy.pos[1:] - [0., R], 0, 2 * R, angle=np.degrees(theta), fill=False, edgecolor='k', linewidth=1, rotation_point='center')
            ax.add_patch(circle)
            ax.add_patch(h_bar)
            ax.add_patch(v_bar)



    ax.set_title(f't = {timestamp:.3f}\n')
    fig.savefig(dest_file, bbox_inches='tight')
    plt.close(fig)
    del fig


class body():
    def __init__(self, mass, pos, vel, inertia = None, radius = None):
        self.mass = mass
        self.pos = pos
        self.vel = vel

        assert inertia != radius

        if inertia is None:
            self.radius = radius
            self.inertia = .5 * mass * radius ** 2
        elif radius is None:
            self.inertia = inertia
        
        M = np.array([[self.inertia, 0, 0],
                      [0, self.mass, 0],
                      [0, 0, self.mass]])
        self.invM = np.linalg.inv(M)


class mesh_body():
    element_length = 1.
    P = 50 # points per element
    radius = 0.5

    def __init__(self, mass, pos, span):
        self.mass = mass
        self.pos = pos
        self.span = span

        L = self.element_length
        self.N = int((span[1] - span[0]) // L)  # elements count
        self.M = self.N+1  # nodes count
        self.rho = 2. * mass / ((span[1] - span[0]) * np.pi * self.radius ** 2)

        self.nodes_params = np.zeros((2*self.M, 1),)
        self.node_velocities = np.zeros((2*self.M, 1),)
        self.node_accelerations = np.zeros((2*self.M, 1),)
        self.K = None
        self.K_eff_red_inv = None
        self.M_global = None
        self.C = None

    
    def get_coords(self):
        assert self.nodes_params is not None
        X = []
        Y = []
        L = self.element_length 
        for element_index in range(self.N):
            x_offset = element_index * L
            for p in np.linspace(0, 1, self.P):
                n_1 = 1 - 3*p**2 + 2*p**3
                n_2 = L*(p - 2*p**2 + p**3)
                n_3 = 3*p**2 - 2*p**3
                n_4 = L*(-p**2 + p**3)
                w = np.dot([n_1, n_2, n_3, n_4], self.nodes_params[2*element_index:2*element_index+4, :].ravel()).item()
                X.append(x_offset + p * L)
                Y.append(self.pos+w)

        return X, Y


    def make_contact_force(self, mode, collision_index):
        if mode != CO:
            return None
        left_node_index = collision_index // self.P
        right_node_index = left_node_index + 1
        assert right_node_index < self.M
        a = (collision_index - left_node_index * self.P) / self.P
        b = 1. - a
        max_force_magnitude_newtons = -1e+5
        f = max_force_magnitude_newtons
        a2 = a ** 2
        b2 = b ** 2
        L = self.element_length
        L2 = L ** 2
        L3 = L ** 3
        # [v_i]
        # [theta_i]
        F = np.zeros((2*self.M, 1))
        F[2 * left_node_index + 0,  0] = b2 * (3*a+b) / L3 * f
        F[2 * left_node_index + 1,  0] = -a * b2 / L2 * f
        F[2 * right_node_index + 0, 0] = a2 * (3*b + a) / L3 * f
        F[2 * right_node_index + 1, 0] = a2 * b / L2 * f
        #print('coeffs ', b2 * (3*a+b) / L3, -a * b2 / L2, a2 * (3*b + a) / L3, a2 * b / L2)
        return F


    def solve_for_contact_normal(self, collision_index):
        half_step = self.P
        _, y = self.get_coords()
        left_b = max([0, collision_index-half_step])
        right_b =  min([len(y), collision_index+half_step])
        y_chunk = np.array(y[left_b:right_b])
        offset = self.P // 5
        y_chunk_fw = np.roll(y_chunk, offset)
        diff = y_chunk[offset:] - y_chunk_fw[offset:]
        slope_fit = np.mean(diff)
        slope_angle = np.arctan(slope_fit)
        return np.array([-np.sin(slope_angle), np.cos(slope_angle)]), slope_angle



def find_collision(fixed_body, free_body):
    EPS = 0.03
    ball_x, ball_y = free_body.pos[1:]
    x, y = fixed_body.get_coords()
    ind = bisect.bisect_left(x, ball_x)
    if ind == len(x):
        return False, None

    return y[ind] + free_body.radius + EPS > ball_y, ind


def propagate_fixed(mode, now, delta_t, fixed_bodies, Fs):
    assert len(Fs) <= 1
    delta_t_sq = delta_t ** 2
    E_steel = 2.e+6
    rez_bodies = []

    for fixed_body in fixed_bodies:
        rez_bodies.append(copy.deepcopy(fixed_body))
        if mode == MB:
            continue
        fb = rez_bodies[-1]

        N = fb.N
        M = fb.M
        newmark_beta = .25
        newmark_gamma = .5

        if fb.K_eff_red_inv is None:
            I = np.pi * fb.radius ** 4 / 4.

            L = fb.element_length
            L2 = L ** 2
            L3 = L ** 3
            K_e = np.array([[  12,    6*L,   -12,    6*L  ],
                            [  6*L,   4*L2,  -6*L,   2*L2 ],
                            [ -12,   -6*L,    12,   -6*L  ],
                            [  6*L,   2*L2,  -6*L,   4*L2 ]]) * E_steel * I / L3
            mu = fb.rho * np.pi * fb.radius ** 2 / 2
            M_e =  np.array([[  156,   22*L,   54,    -13*L ],
                             [ 22*L,   4*L2,   13*L,  -3*L2 ],
                             [   54,   13*L,   156,   -22*L ],
                             [-13*L,  -3*L2,  -22*L,   4*L2]]) * mu * L / 420.


            K = np.zeros((2*M, 2*M),)
            M_global = np.zeros((2*M, 2*M),)

            alpha_rayleigh = .9
            beta_rayleigh = .9

            connectivity = lambda e, k: 2*e + k
            for element_id in range(N):
                assert element_id+1 < M
                for i in range(4):
                    for j in range(4):
                        row = connectivity(element_id, i)
                        col = connectivity(element_id, j)
                        K[row, col] += K_e[i, j]
                        M_global += M_e[i, j]

            C = alpha_rayleigh * M_global + beta_rayleigh * K
            # C = np.zeros((2*M, 2*M),)
            fb.C = C
            fb.K = K
            fb.M_global = M_global 
            K_eff = K + (newmark_gamma / (newmark_beta * delta_t)) * C + 1. / (newmark_beta * delta_t_sq) * M_global

            print('K eff full', np.linalg.cond(K_eff))
            K_eff_red = K_eff[2:-2, 2:-2]
            print('K eff red', np.linalg.cond(K_eff_red))
            fb.K_eff_red_inv = np.linalg.inv(K_eff_red)

        
        u_prev = fb.nodes_params
        u_prev_dot = fb.node_velocities
        u_prev_ddot = fb.node_accelerations

        # 1)
        u_pred = u_prev + delta_t * u_prev_dot + (.5 - newmark_beta) * delta_t_sq * u_prev_ddot
        u_pred *= .99
        u_pred[:2, :] = 0.
        u_pred[-2:, :] = 0.
        u_pred_dot = u_prev_dot + (1. - newmark_gamma) * delta_t * u_prev_ddot
        u_pred_dot[:2, :] = 0.
        u_pred_dot[-2:, :] = 0.
        u_pred_dot *= .97
        u_pred_ddot = u_prev_ddot
        u_pred_ddot *= .9

        # 2) is above
        K_eff_red_inv = fb.K_eff_red_inv

        # 3)
        F = Fs[0] if len(Fs) == 1 else np.zeros((2*fb.M, 1))
        F_eff = F + fb.M_global @ (u_pred / (newmark_beta * delta_t_sq) + \
                                   u_pred_dot / (newmark_beta * delta_t) + \
                                   ((1. / (2. * newmark_beta)) - 1.) * u_pred_ddot) + \
                    fb.C @ (newmark_gamma * u_pred / (newmark_beta * delta_t) + u_pred_dot + \
                            delta_t * (newmark_gamma / (2. * newmark_beta) - 1.) * u_pred_ddot)
        # 4)
        F_eff = F_eff[2:-2,:]
        u_red = K_eff_red_inv @ F_eff
        u = np.zeros((2*M, 1),)
        u[2:-2, :] = u_red

        # 5)
        fb.nodes_params = u
        fb.node_accelerations = (u - u_pred)/(newmark_beta * delta_t_sq) - u_pred_dot / (newmark_beta * delta_t)
        fb.node_accelerations[:2, :] = 0.
        fb.node_accelerations[-2:, :] = 0.
        fb.node_velocities = u_pred_dot + newmark_gamma * delta_t * fb.node_accelerations
        fb.node_velocities[:2, :] = 0.
        fb.node_velocities[-2:, :] = 0.

        E_pot = (0.5 * u.T @ fb.K @ u).item()
        E_kin = (0.5 * fb.node_velocities.T @ fb.M_global @ fb.node_velocities).item()
        print(f'{now:.3f}, kin: {E_kin:.3f}, pot: {E_pot:.3f}')

    return rez_bodies


def make_fea_force(mode, t, force_time_end, beam: mesh_body):
    if FEA != mode or t >= force_time_end:
        return None

    mid_id = beam.M // 2
    max_force_magnitude_newtons = -1e+4
    a = t / force_time_end
    f = np.zeros((2*beam.M, 1))
    f[2*mid_id, 0] = np.sin(a * np.pi / 2) * max_force_magnitude_newtons
    return f


def propagate(mode, now, delta_t, free_bodies, fixed_bodies, forces):
    force_for_fea = list(filter(lambda f: f.shape[0] > 3, forces))
    pred_fixed_bodies = propagate_fixed(mode, now, delta_t, fixed_bodies, force_for_fea)

    prop_bodies = []
    contact_force = None
    for bdy in free_bodies:
        prop_bodies.append(copy.deepcopy(bdy))
        if mode == FEA:
            continue

        curr_b = prop_bodies[-1]
        a = (curr_b.invM @ forces[0][np.newaxis, :].T).T.ravel()
        curr_b.vel += a * delta_t
        curr_b.pos += curr_b.vel * delta_t

        for fb in pred_fixed_bodies:
            collides, collision_index = find_collision(fb, curr_b)
            if collides:
                contact_n, slope_angle = fb.solve_for_contact_normal(collision_index)
                print(f'collides! with slope angle {np.degrees(slope_angle)}, contact normal: {contact_n}')
                ra = contact_n * curr_b.radius
                perp_vec = lambda r: np.array([-r[1], r[0]])
                ra_perp = perp_vec(ra)
                v_rel = curr_b.vel[1:] + curr_b.vel[0] * ra_perp
                print(f'ra_perp: {ra_perp}, angular velocity: {curr_b.vel[0]}, spatial velocity: {curr_b.vel[1:]}')
                v_rel_n = np.dot(contact_n, v_rel)

                e = 0.85
                I_inv = 1. / curr_b.inertia
                m_inv = 1. / curr_b.mass
                beam_m_inv = 1. / fb.mass

                cross2d = lambda r, n: r[0] * n[1] - r[1] * n[0]
                denom = m_inv + beam_m_inv + I_inv * cross2d(ra, contact_n) ** 2
                j = -(1. + e) * v_rel_n / denom
                J = j * contact_n

                smth_omega = I_inv * cross2d(ra, J)
                smth_lin = m_inv * J
                print(f'at {now:.3f} v_rel {v_rel} || change: {smth_omega:.1f} {smth_lin}, impulse: {J}')

                curr_b.vel[0] += smth_omega
                curr_b.vel[1:] += smth_lin
                contact_force = fb.make_contact_force(mode, collision_index)
            else:
                contact_force = None

    return prop_bodies, pred_fixed_bodies, contact_force


def gif_maker(dest_dir, gif_path, mode):
    command = f'ffmpeg -framerate 30 -pattern_type glob -i "{dest_dir}/*.png" -vf "fps=30,scale=600:-1:flags=lanczos,split[s0][s1];[s0]palettegen[p];[s1][p]paletteuse" -loop 0 {gif_path}'
    os.system(command)


def do_sim(args: SimArgs):
    time_end = 11.
    delta_t = 3.e-2

    world_bounds = [[-.5, 20], [-1.5, 15]]

    ball = body(5., np.array([0., 0., 10.]), np.array([np.radians(3), 2., 3.]), None, .3)
    beam = mesh_body(80, 0., world_bounds[0])
    bodies = [ball]
    fixed_bodies = [beam]
    gravity = np.array([0, 0, -12.81])
    force_time_end = time_end / 4.
    dynamic_force = None
    last_contact_time = None
    nominal_contact_force = None

    for i, t in enumerate(np.arange(0, time_end, delta_t)):
        dynamic_force = make_fea_force(args.mode, t, force_time_end, beam)
        forces = [gravity]
        if dynamic_force is not None:
            forces.append(dynamic_force)
        elif last_contact_time is not None:
            if t > last_contact_time + force_time_end:
                last_contact_time = None
                nominal_contact_force = None
            else:
                assert nominal_contact_force is not None
                forces.append(copy.deepcopy(nominal_contact_force))
                forces[-1] *= np.cos(np.pi / 2 * (t-last_contact_time) / force_time_end)

        bodies, fixed_bodies, ret_contact_force = propagate(args.mode, t, delta_t, bodies, fixed_bodies, forces)
        if ret_contact_force is not None:
            nominal_contact_force = ret_contact_force
            last_contact_time = t
        if 0 == i % 3:
            render(args.mode, bodies, fixed_bodies, world_bounds, args.dest_dir, i, t)

    gif_maker(args.dest_dir, args.gif_path, args.mode)


if __name__ == '__main__':
    args = SimArgs().parse_args()
    do_sim(args)