python数据可视化

Python 提供了许多流行的数据可视化库，其中最常用的包括 Matplotlib、Seaborn 和 Plotly。

下面给出一个例子：

# plot_utils.py
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import os

def plot_results(timestamps, positions, roll_angles, pitch_angles, yaw_angles, contact_forces, contact_torques):
    # Plot the end-effector position, orientation, and contact force curves
    plt.figure(figsize=(15, 10))

    # End-effector position
    plt.subplot(2, 3, 1)
    plt.plot(timestamps, positions[:, 0], label="x")
    plt.title("End-effector Position X")
    plt.xlabel("Time (s)")
    plt.ylabel("Position (m)")

    plt.subplot(2, 3, 2)
    plt.plot(timestamps, positions[:, 1], label="y")
    plt.title("End-effector Position Y")
    plt.xlabel("Time (s)")
    plt.ylabel("Position (m)")

    plt.subplot(2, 3, 3)
    plt.plot(timestamps, positions[:, 2], label="z")
    plt.title("End-effector Position Z")
    plt.xlabel("Time (s)")
    plt.ylabel("Position (m)")

    # Euler angles
    plt.subplot(2, 3, 4)
    plt.plot(timestamps, roll_angles, label="Roll")
    plt.title("End-effector Orientation (Roll)")
    plt.xlabel("Time (s)")
    plt.ylabel("Angle (rad)")

    plt.subplot(2, 3, 5)
    plt.plot(timestamps, pitch_angles, label="Pitch")
    plt.title("End-effector Orientation (Pitch)")
    plt.xlabel("Time (s)")
    plt.ylabel("Angle (rad)")

    plt.subplot(2, 3, 6)
    plt.plot(timestamps, yaw_angles, label="Yaw")
    plt.title("End-effector Orientation (Yaw)")
    plt.xlabel("Time (s)")
    plt.ylabel("Angle (rad)")

    plt.legend()
    plt.grid(True) 
    plt.tight_layout()
    plt.show()

    # Contact forces and torques (relative to end-effector frame)
    plt.figure(figsize=(15, 10))

    plt.subplot(2, 3, 1)
    plt.plot(timestamps, contact_forces[:, 0], label="Fx (EE)")
    plt.title("Contact Force X (End-effector Frame)")
    plt.xlabel("Time (s)")
    plt.ylabel("Force (N)")

    plt.subplot(2, 3, 2)
    plt.plot(timestamps, contact_forces[:, 1], label="Fy (EE)")
    plt.title("Contact Force Y (End-effector Frame)")
    plt.xlabel("Time (s)")
    plt.ylabel("Force (N)")

    plt.subplot(2, 3, 3)
    plt.plot(timestamps, contact_forces[:, 2], label="Fz (EE)")
    plt.title("Contact Force Z (End-effector Frame)")
    plt.xlabel("Time (s)")
    plt.ylabel("Force (N)")

    plt.subplot(2, 3, 4)
    plt.plot(timestamps, contact_torques[:, 0], label="Tx (EE)")
    plt.title("Contact Torque X (End-effector Frame)")
    plt.xlabel("Time (s)")
    plt.ylabel("Torque (Nm)")

    plt.subplot(2, 3, 5)
    plt.plot(timestamps, contact_torques[:, 1], label="Ty (EE)")
    plt.title("Contact Torque Y (End-effector Frame)")
    plt.xlabel("Time (s)")
    plt.ylabel("Torque (Nm)")

    plt.subplot(2, 3, 6)
    plt.plot(timestamps, contact_torques[:, 2], label="Tz (EE)")
    plt.title("Contact Torque Z (End-effector Frame)")
    plt.xlabel("Time (s)")
    plt.ylabel("Torque (Nm)")

    plt.tight_layout()
    plt.show()
    plt.grid(True) 

def save_force_torque_data(timestamps, contact_forces, contact_torques, folder_path="Data"):
    # Ensure the folder exists
    os.makedirs(folder_path, exist_ok=True)

    # Create a DataFrame for contact forces and torques
    data = {
        "Time (s)": timestamps,
        "Fx (N)": contact_forces[:, 0],
        "Fy (N)": contact_forces[:, 1],
        "Fz (N)": contact_forces[:, 2],
        "Tx (Nm)": contact_torques[:, 0],
        "Ty (Nm)": contact_torques[:, 1],
        "Tz (Nm)": contact_torques[:, 2]
    }

    df = pd.DataFrame(data)

    # Save to Excel file
    file_path = os.path.join(folder_path, "Fe_Data.xlsx")
    df.to_excel(file_path, index=False)

    print(f"Contact force and torque data saved to {file_path}")


def save_position_orientation_data(timestamps, positions, roll_angles, pitch_angles, yaw_angles, folder_path="Data"):
    # Ensure the folder exists
    os.makedirs(folder_path, exist_ok=True)

    # Create a DataFrame for position and orientation (roll, pitch, yaw) data
    data = {
        "Time (s)": timestamps,
        "Position X (m)": positions[:, 0],
        "Position Y (m)": positions[:, 1],
        "Position Z (m)": positions[:, 2],
        "Roll (rad)": roll_angles,
        "Pitch (rad)": pitch_angles,
        "Yaw (rad)": yaw_angles
    }

    df = pd.DataFrame(data)

    # Save to Excel file
    file_path = os.path.join(folder_path, "Position_Orientation_Data.xlsx")
    df.to_excel(file_path, index=False)

    print(f"Position and orientation data saved to {file_path}")