Как визуализировать коллективное поведение самодвижущегося стержня в двумерном пространстве?

Несколько самоподвижных стержней (моделируемых с использованием нечётного числа соединённых твердых сфер) с фиксированной скоростью перемещения движутся в среде (2D) с тремя различными диффузионными константами для вращательной и трансляционной диффузий (x и y), каждый стержень имеет разную ориентацию. Когда две сферы от двух различных стержней попадают в определённый радиус взаимодействия, на них начинает действовать сила взаимодействия, которая возникает из потенциала Леннард-Джонса, который является функцией расстояния между двумя взаимодействующими сферами. Компоненты силы взаимодействия складываются со скоростью трансляции, а крутящий момент, который создаётся, медленно вращает оба стержня, и они начинают выстраиваться. Уравнение движения для стержней и форма потенциала приведены в приложениивведите описание изображения здесьвведите описание изображения здесь

Это код, который я написал,

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import matplotlib.patches as patches

# Константы и параметры
WIDTH, HEIGHT = 100, 100
NUM_SPHERES_PER_ROD = 9  # Нечетное число сфер на стержень
SPHERE_RADIUS = 0.5
NUM_RODS = 10
VELOCITY = 2
FRAMES = 10
NEIGHBOR_DISTANCE_THRESHOLD = 2 * SPHERE_RADIUS  # Диаметр сферы
EPSILON = 0.000001  # Глубина потенциальной ямы Леннард-Джонса
SIGMA = 2 * SPHERE_RADIUS  # Расстояние, на котором потенциал равен нулю

def initialize_rods(num_rods):
    positions = np.random.rand(num_rods, 2) * [WIDTH, HEIGHT]
    angles = np.random.uniform(0, 2 * np.pi, num_rods)
    return positions, angles

def update_positions(positions, angles, velocity):
    theta = angles
    dx = velocity * np.cos(theta)
    dy = velocity * np.sin(theta)
    positions[:, 0] += dx
    positions[:, 1] += dy

    # Периодическая граница
    positions[:, 0] = positions[:, 0] % WIDTH
    positions[:, 1] = positions[:, 1] % HEIGHT
    return positions

def get_sphere_positions(start_pos, angle, num_spheres):
    sphere_positions = []
    for i in range(int((num_spheres + 1) / 2)):
        x = start_pos[0] - 2 * i * SPHERE_RADIUS * np.cos(angle)
        y = start_pos[1] - 2 * i * SPHERE_RADIUS * np.sin(angle)
        x1 = start_pos[0] + 2 * i * SPHERE_RADIUS * np.cos(angle)
        y1 = start_pos[1] + 2 * i * SPHERE_RADIUS * np.sin(angle)
        sphere_positions.append((x, y))
        sphere_positions.append((x1, y1))
    return sphere_positions

def lennard_jones_force(r):
    if r == 0:
        return 0  # Избегаем деления на ноль
    r6 = (SIGMA / r) ** 6
    r12 = r6 * r6
    force_magnitude = 24 * EPSILON * (2 * r12 - r6) / r
    return force_magnitude

def compute_forces_and_neighbors(positions, angles):
    num_rods = len(positions)
    neighbor_sphere_pairs = []
    forces = np.zeros_like(positions)  # Силы, которые будут применены к каждому стержню

    for i in range(num_rods):
        sphere_positions_i = get_sphere_positions(positions[i], angles[i], NUM_SPHERES_PER_ROD)
        for j in range(i + 1, num_rods):
            sphere_positions_j = get_sphere_positions(positions[j], angles[j], NUM_SPHERES_PER_ROD)
            for sphere_pos_i in sphere_positions_i:
                for sphere_pos_j in sphere_positions_j:
                    dist_vector = np.array(sphere_pos_j) - np.array(sphere_pos_i)
                    dist = np.linalg.norm(dist_vector)
                    if dist < NEIGHBOR_DISTANCE_THRESHOLD:
                        # Записываем пару соседних сфер
                        neighbor_sphere_pairs.append((sphere_pos_i, sphere_pos_j))
                        # Вычисляем силу Леннард-Джонса
                        force_mag = lennard_jones_force(dist)
                        force_vec = force_mag * (dist_vector / dist)  # Нормализуем для получения направления
                        # Применяем силы к стержням
                        forces[i] += force_vec
                        forces[j] -= force_vec  # Равные и противоположные силы

    return forces, neighbor_sphere_pairs

def update_angles(angles, forces):
    for i in range(len(angles)):
        # Вычисляем результирующую силу
        resultant_force = forces[i]
        if np.linalg.norm(resultant_force) > 0:
            # Корректируем угол в зависимости от результирующего направления
            angles[i] = np.arctan2(resultant_force[1], resultant_force[0])
    return angles

def draw_sphere(ax, center, radius, COLOUR):
    circle = patches.Circle(center, radius, color=COLOUR, fill=True)
    ax.add_patch(circle)

def draw_rod(ax, start_pos, angle, num_spheres):
    sphere_positions = get_sphere_positions(start_pos, angle, num_spheres)
    for idx, (x, y) in enumerate(sphere_positions):
        COLOUR = "blue" if idx % 2 == 0 else "red"
        draw_sphere(ax, (x, y), SPHERE_RADIUS, COLOUR)

def animate(frame):
    global positions, angles
    positions = update_positions(positions, angles, VELOCITY)
    forces, neighbor_sphere_pairs = compute_forces_and_neighbors(positions, angles)  # Находим соседей и вычисляем силы
    angles = update_angles(angles, forces)  # Обновляем углы стержней в зависимости от сил

    ax.clear()
    for idx, (pos, angle) in enumerate(zip(positions, angles)):
        draw_rod(ax, pos, angle, NUM_SPHERES_PER_ROD)

    ax.set_xlim(0, WIDTH)
    ax.set_ylim(0, HEIGHT)
    ax.set_aspect('equal')

# Инициализация
positions, angles = initialize_rods(NUM_RODS)

# Настройка фигуры и оси
fig, ax = plt.subplots(figsize=(10, 8))

# Создание анимации
anim = animation.FuncAnimation(fig, animate, frames=FRAMES, interval=10, repeat=True)
plt.show()

Но вращение стержней после взаимодействия не точно. Может быть, у вас есть предложения?