planners/trajectory_generator.py

"""
In this file, you should implement your trajectory generation class or function.
Your method must generate a smooth 3-axis trajectory (x(t), y(t), z(t)) that 
passes through all the previously computed path points. A positional deviation 
up to 0.1 m from each path point is allowed.

You should output the generated trajectory and visualize it. The figure must
contain three subplots showing x, y, and z, respectively, with time t (in seconds)
as the horizontal axis. Additionally, you must plot the original discrete path 
points on the same figure for comparison.

You are expected to write the implementation yourself. Do NOT copy or reuse any 
existing trajectory generation code from others. Avoid using external packages 
beyond general scientific libraries such as numpy, math, or scipy. If you decide 
to use additional packages, you must clearly explain the reason in your report.
"""

import numpy as np
import matplotlib.pyplot as plt

# --------------------------------------------------------------------------------------------------- #
# Trajectory planning: Natural cubic spline through the discrete path points
# The trajectory is continuous and smooth and passes through all path points.
# We parameterize the spline by cumulative chord length (time ~ distance / speed).
# Then we plot x(t), y(t), z(t) with the discrete path points overlayed.
# --------------------------------------------------------------------------------------------------- #


def natural_cubic_spline_coeffs(t, y):
    """
    Compute natural cubic spline second derivatives (m values) for points (t, y).
    Returns array m of length n (m[0] = m[-1] = 0).
    """
    n = len(t)
    if n < 2:
        raise ValueError("Need at least two points for spline.")
    h = np.diff(t)
    # handle degenerate h
    if np.any(h <= 0):
        raise ValueError("t must be strictly increasing.", t)
    # set up tridiagonal system for m[1..n-2]
    A = np.zeros((n, n))
    rhs = np.zeros(n)
    A[0, 0] = 1.0
    A[-1, -1] = 1.0
    for i in range(1, n - 1):
        A[i, i - 1] = h[i - 1]
        A[i, i] = 2 * (h[i - 1] + h[i])
        A[i, i + 1] = h[i]
        rhs[i] = 6 * ((y[i + 1] - y[i]) / h[i] - (y[i] - y[i - 1]) / h[i - 1])
    m = np.linalg.solve(A, rhs)
    return m


def build_spline_segments(t, y):
    """
    Given t and y (arrays of length n), compute spline coefficients for each segment.
    Returns list of dicts with keys a,b,c,d and interval [t_i, t_{i+1}).
    Spline form on segment i: S_i(s) = a + b*s + c*s^2 + d*s^3, where s = (t - t_i)
    """
    n = len(t)
    h = np.diff(t)
    m = natural_cubic_spline_coeffs(t, y)
    segments = []
    for i in range(n - 1):
        a = y[i]
        c = m[i] / 2.0
        d = (m[i + 1] - m[i]) / (6.0 * h[i])
        b = (y[i + 1] - y[i]) / h[i] - (h[i] * (2 * m[i] + m[i + 1])) / 6.0
        segments.append({'t0': t[i], 't1': t[i + 1], 'a': a, 'b': b, 'c': c, 'd': d, 'h': h[i]})
    return segments


def evaluate_spline(segments, t_query):
    """
    Evaluate spline (list of segments) at times t_query (array).
    """
    t_query = np.array(t_query)
    y_out = np.zeros_like(t_query, dtype=float)
    for i, ti in enumerate(t_query):
        # find segment
        # segments are ordered, so we can do a simple search
        if ti <= segments[0]['t0']:
            seg = segments[0]
        elif ti >= segments[-1]['t1']:
            seg = segments[-1]
        else:
            # binary search could be used; linear search is fine for modest sizes
            for seg in segments:
                if seg['t0'] <= ti <= seg['t1']:
                    break
        s = ti - seg['t0']
        y_out[i] = seg['a'] + seg['b'] * s + seg['c'] * s**2 + seg['d'] * s**3
    return y_out


def plan_trajectory_through_path(path, desired_speed=1.0, samples_per_meter=10):
    """
    Given path (N x 3), compute a smooth trajectory (x(t), y(t), z(t)).
    desired_speed sets time scaling (seconds per meter ~ 1/desired_speed).
    samples_per_meter controls temporal resolution.
    Returns time array t_out and trajectory array traj_out (len(t_out) x 3)
    """
    path = np.array(path, dtype=float)
    n = path.shape[0]
    # parameterize by cumulative chord length
    dists = np.linalg.norm(np.diff(path, axis=0), axis=1)
    chord = np.insert(np.cumsum(dists), 0, 0.0)
    total_length = chord[-1]
    if total_length == 0:
        t = np.array([0.0])
        traj = path.copy()
        return t, traj

    # time vector proportional to distance, with constant desired_speed
    T_total = total_length / desired_speed
    t_pts = chord / desired_speed  # times at which we must pass through path points

    # build spline for each axis
    seg_x = build_spline_segments(t_pts, path[:, 0])
    seg_y = build_spline_segments(t_pts, path[:, 1])
    seg_z = build_spline_segments(t_pts, path[:, 2])

    # sample times
    n_samples = max(200, int(np.ceil(total_length * samples_per_meter)))
    t_out = np.linspace(0.0, t_pts[-1], n_samples)

    x_out = evaluate_spline(seg_x, t_out)
    y_out = evaluate_spline(seg_y, t_out)
    z_out = evaluate_spline(seg_z, t_out)

    traj_out = np.vstack((x_out, y_out, z_out)).T
    return t_out, traj_out

def plot_trajectory(t_traj, traj, path):
    # Plot x(t), y(t), z(t) in separate subplots and overlay discrete path points at their times
    fig = plt.figure(figsize=(8, 8))
    axs = []
    axs.append(plt.subplot2grid((7, 1), (0, 0), rowspan=3, colspan=1))
    axs.append(plt.subplot2grid((7, 1), (3, 0), rowspan=3, colspan=1))
    axs.append(plt.subplot2grid((7, 1), (6, 0), rowspan=1, colspan=1))
    axs[0].plot(t_traj, traj[:, 0], label='x(t)')
    axs[1].plot(t_traj, traj[:, 1], label='y(t)', color='orange')
    axs[2].plot(t_traj, traj[:, 2], label='z(t)', color='green')

    # overlay discrete path points
    # compute discrete path times (same parameterization used for spline)
    path = np.array(path)
    dists = np.linalg.norm(np.diff(path, axis=0), axis=1)
    chord = np.insert(np.cumsum(dists), 0, 0.0)
    path_times = chord / 1.0  # desired_speed = 1.0 used above

    axs[0].scatter(path_times, path[:, 0], color='k', zorder=10)
    axs[1].scatter(path_times, path[:, 1], color='k', zorder=10)
    axs[2].scatter(path_times, path[:, 2], color='k', zorder=10)

    axs[2].set_xlabel('time (s)')
    axs[0].set_ylabel('x (m)')
    axs[1].set_ylabel('y (m)')
    axs[2].set_ylabel('z (m)')

    axs[0].grid(True)
    axs[1].grid(True)
    axs[2].grid(True)

    axs[0].set_title('Trajectory')

    plt.tight_layout()
    plt.show()