helper.py

import numpy as np 
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from config import TIME_DATA, N, SAVE_ANIMATION


# This script contains utility functions for plotting and animations.


# Define a set of 11 different colors for plotting and animations
colors = [
    "#1f77b4", "#ff7f0e", "#2ca02c", "#d62728", "#9467bd",
     "#8c564b", "#e377c2", "#7f7f7f", "#bcbd22", "#17becf", "#92A8D1"   
]

# Animation generator function. 
def animate(time_data, pos_data, num_agent, L_obstacles, save_animation=True):
    fig, ax = plt.subplots(figsize=(6, 6))
    ax.set_title("Agent 2D Position (X-Y)")
    ax.set_xlabel("X position (m)")
    ax.set_ylabel("Y position (m)")
    ax.grid(True)

    # Axis limits: include both trajectory and circle extents
    all_xy = np.vstack([pd[:, :2] for pd in pos_data])  # (N*T, 2)
    x_min, x_max = all_xy[:, 0].min(), all_xy[:, 0].max()
    y_min, y_max = all_xy[:, 1].min(), all_xy[:, 1].max()
    x_min = min(x_min, L_obstacles[0].center[0] - 4.0) - 0.5
    x_max = max(x_max, L_obstacles[1].center[0] + 4.0) + 0.5
    y_min = min(y_min, L_obstacles[1].center[1] - 4.0) - 0.5
    y_max = max(y_max, L_obstacles[1].center[1] + 4.0) + 0.5
    ax.set_xlim(x_min, x_max)
    ax.set_ylim(y_min, y_max)

    # Draw the L regions
    for L_obstacle in L_obstacles: L_obstacle.plot(ax, color='lightgrey')

    # Plot the points representing agents
    scatters = [ax.plot([], [], 'o',color = colors[i], label=f"Agent {i}")[0] for i in range(num_agent)]

    def update(frame):
        for i, sc in enumerate(scatters):
            sc.set_data([pos_data[i][frame, 0]], [pos_data[i][frame, 1]])
        return scatters

    ani = animation.FuncAnimation(fig, update, frames=len(time_data), interval=10, blit=True)
    plt.tight_layout()
    plt.show()
    if save_animation:
        print('Saving animation as video. Please wait.')
        writer = animation.PillowWriter(fps=100)  # adjust fps if you like
        ani.save("agent_sim1.gif", writer=writer)
        print('Done.')


# Plot the evolutions of the agent positions (up to 3-dimensional positon state)
def plot_position(time_data, pos_data,dimension):
    dim_labels = ['x', 'y', 'z']
    fig, axes = plt.subplots(dimension, 1, sharex=True)

    # Plot the agent's position in each dimension
    for dim in range(dimension):
        ax = axes if dim == 0 else axes[dim]
        for i, agent_pos_data in enumerate(pos_data):
            ax.plot(time_data[:, 0], agent_pos_data[:, dim],color = colors[i], label=f'Drone {i+1}')

        ax.set_ylabel(rf'Position ${dim_labels[dim]}$ (m)')

        ax.grid(True)
        if dim == 0:
            ax.set_title('Drone Position Over Time')
        if dim == 2:
            ax.set_xlabel('Time (s)')
    
    plt.tight_layout()


# Plot the agents' control inputs (velocities)
def plot_input(time_data, vel_data,dimension):
    dim_labels = ['x', 'y', 'z']
    fig, axes = plt.subplots(dimension, 1, sharex=True)

    # Plot the agent's control input in each dimension.
    for dim in range(dimension):
        ax = axes if dim == 0 else axes[dim]
        for i, agent_pos_data in enumerate(vel_data):
            ax.plot(time_data[:, 0], agent_pos_data[:, dim],color = colors[i], label=f'Drone {i+1}')
        ax.set_ylabel(rf'Control Input ${dim_labels[dim]}$ (m/s)')
        ax.grid(True)
        if dim == 0:
            ax.set_title('Control Inputs Over Time')
        if dim == 2:
            ax.set_xlabel('Time (s)')
    

    plt.tight_layout()

# Plot the \max(h_{L_1}(x_j), h_{L_2}(x_j) for each agent j 
def plot_barrier_functions(time_data,eig_data,num_agent):
    fig, ax = plt.subplots(1, 1, figsize=(11.5, 2.5))
    for i in range(num_agent):
        plt.plot(time_data, np.array(eig_data[i]),color = colors[i], label=rf'$Agent {i+1}$')
    
    plt.xlabel(r"Time [ticks]")
    plt.ylabel(r"$\max(h_{L_1}(x_j), h_{L_2}(x_j)$")
    plt.title(r"CBF Values")
    plt.grid(True)
    # plt.legend()  
    plt.tight_layout()


# Visualzie the results with the given data
def visualize_results(pos_data, vel_data, h_data, L1, L2):

    # Plot the max(h_{L_1}(x_j), h_{L_2}(x_j))
    plot_barrier_functions(TIME_DATA[:-1],h_data,N)

    # Plot the positions of the agents
    plot_position(TIME_DATA, pos_data,dimension=1)
    
    # Plot the control inputs of agents
    plot_input(TIME_DATA,vel_data,dimension=1)

    # Animate the simulation
    animate(TIME_DATA, pos_data, N, [L1, L2], SAVE_ANIMATION)
    plt.show()