DesignOptimization.py

"""
Cart–Pole rail travel + stepper running torque sizing helper
------------------------------------------------------------

What this script computes (from your robustness specs):
1) Capture displacement (how far the cart must move to "save" the pole)
2) Recommended half-rail and total rail length (with margin)
3) Required cart acceleration (assuming bang-bang accel/decel to cover the distance in a chosen time)
4) Required cart force
5) Required motor torque at speed through belt/pulley (and optional inertia torque)

IMPORTANT: Validate your chosen stepper against its *pull-out torque vs speed* curve
at the computed peak motor speed (NOT holding torque at 0 rpm).
"""

from dataclasses import dataclass
from math import sin, cos, sqrt, pi

# -------------------------
# Inputs you will tweak
# -------------------------

@dataclass
class CartPoleSpecs:
    # Geometry / physics
    l: float = 0.30            # [m] pivot -> pole COM distance
    g: float = 9.81            # [m/s^2]
    theta_max_deg: float = 10  # [deg] max recoverable angle
    theta_dot_max: float = 1.0 # [rad/s] max recoverable angular rate
    x_dot0: float = 0.0        # [m/s] initial cart speed (often 0)

@dataclass
class RailDesign:
    margin_frac: float = 0.75  # [-] add 75% margin on capture displacement
    extra_margin_m: float = 0.02  # [m] fixed extra margin for end-stops, etc.

@dataclass
class MotionPlan:
    catch_time: float = 0.30   # [s] time to cover the needed displacement (bang-bang accel)

@dataclass
class DriveTrain:
    pulley_radius: float = 0.02 # [m] effective drive pulley radius
    gear_ratio: float = 1.0     # [-] G = motor_speed / pulley_speed (1.0 if direct)
    efficiency: float = 0.90    # [-] belt+bearings+etc (0.85–0.95 typical)

@dataclass
class MassAndFriction:
    M_cart: float = 1.0        # [kg] cart + moving carriage mass
    m_pole: float = 0.2        # [kg] pole mass (include if it rides with cart dynamics)
    F_friction: float = 2.0    # [N] estimated Coulomb-ish friction (start conservative)
    F_disturbance: float = 0.0 # [N] optional extra lumped disturbance force

@dataclass
class MotorInertia:
    J_motor: float = 0.0       # [kg*m^2] rotor inertia (optional)
    J_extra: float = 0.0       # [kg*m^2] any extra rotating inertia on motor shaft (optional)

# -------------------------
# Helpers
# -------------------------

def deg2rad(deg: float) -> float:
    return deg * pi / 180.0

def radps_to_rpm(w: float) -> float:
    return w * 60.0 / (2.0 * pi)

# -------------------------
# Core calculations
# -------------------------

def compute_requirements(
    specs: CartPoleSpecs,
    rail: RailDesign,
    motion: MotionPlan,
    drive: DriveTrain,
    mass: MassAndFriction,
    inertia: MotorInertia,
    safety_factor: float = 2.5  # apply to motor torque for "robust"
) -> dict:
    # 1) Natural rate for linear inverted pendulum (LIPM-style)
    omega0 = sqrt(specs.g / specs.l)  # [1/s]

    # 2) Capture displacement (cart travel needed to "save" it)
    theta = deg2rad(specs.theta_max_deg)
    theta_dot = specs.theta_dot_max

    # Capture point shift relative to current cart position x:
    # Δx_cap = l*theta + (x_dot + l*theta_dot)/omega0
    dx_cap = specs.l * theta + (specs.x_dot0 + specs.l * theta_dot) / omega0  # [m]

    # 3) Rail length recommendation
    dx_design = dx_cap * (1.0 + rail.margin_frac) + rail.extra_margin_m  # [m]
    half_rail = dx_design
    rail_length = 2.0 * half_rail

    # 4) Motion: bang-bang accel to move dx_cap in catch_time
    # For bang-bang (accel then decel), displacement = a*T^2/4  => a = 4*dx/T^2
    T = motion.catch_time
    a_req = 4.0 * dx_cap / (T ** 2)  # [m/s^2]
    v_peak = a_req * (T / 2.0)       # [m/s]

    # 5) Force sizing (conservative)
    M_total = mass.M_cart + mass.m_pole
    F_req = M_total * a_req + mass.F_friction + mass.F_disturbance  # [N]

    # 6) Convert to motor torque through pulley + ratio
    r = drive.pulley_radius
    G = drive.gear_ratio
    eta = drive.efficiency

    # Load torque from cart force:
    tau_load = (F_req * r) / (G * eta)  # [N*m]

    # Motor speed requirement:
    # cart speed v -> pulley omega_p = v/r, motor omega_m = G*omega_p
    omega_m_peak = G * (v_peak / r)          # [rad/s]
    rpm_peak = radps_to_rpm(omega_m_peak)    # [rpm]

    # Motor angular acceleration:
    alpha_m = G * (a_req / r)  # [rad/s^2]

    # Inertia torque on motor shaft:
    J_total = inertia.J_motor + inertia.J_extra
    tau_inertia = J_total * alpha_m  # [N*m]

    # Peak torque magnitude estimate
    tau_peak = tau_load + tau_inertia
    tau_peak_sf = safety_factor * tau_peak

    return {
        "omega0_1ps": omega0,
        "capture_displacement_m": dx_cap,
        "half_rail_recommended_m": half_rail,
        "rail_length_recommended_m": rail_length,
        "catch_time_s": T,
        "required_cart_accel_mps2": a_req,
        "peak_cart_speed_mps": v_peak,
        "required_cart_force_N": F_req,
        "motor_speed_peak_radps": omega_m_peak,
        "motor_speed_peak_rpm": rpm_peak,
        "motor_alpha_peak_radps2": alpha_m,
        "motor_torque_load_Nm": tau_load,
        "motor_torque_inertia_Nm": tau_inertia,
        "motor_torque_peak_Nm": tau_peak,
        "motor_torque_peak_with_safety_Nm": tau_peak_sf,
        "notes": (
            "Compare motor_torque_peak_with_safety_Nm against stepper PULL-OUT torque "
            "at motor_speed_peak_rpm (not holding torque). Increase voltage/current, "
            "add gear reduction, or reduce catch_time if margin is insufficient."
        ),
    }

def pretty_print(results: dict) -> None:
    print("\n===== Cart-Pole Rail + Stepper Sizing Results =====")
    print(f"omega0                         = {results['omega0_1ps']:.3f} 1/s")
    print(f"capture displacement (dx_cap)  = {results['capture_displacement_m']:.4f} m")
    print(f"recommended half-rail travel   = {results['half_rail_recommended_m']:.4f} m")
    print(f"recommended total rail length  = {results['rail_length_recommended_m']:.4f} m")

    print("\n--- Motion plan---")
    print(f"catch time T                   = {results['catch_time_s']:.3f} s")
    print(f"required cart accel a_req      = {results['required_cart_accel_mps2']:.3f} m/s^2")
    print(f"peak cart speed v_peak         = {results['peak_cart_speed_mps']:.3f} m/s")

    print("\n--- Force / Motor ---")
    print(f"required cart force F_req      = {results['required_cart_force_N']:.3f} N")
    print(f"motor peak speed               = {results['motor_speed_peak_rpm']:.1f} rpm")
    print(f"motor load torque              = {results['motor_torque_load_Nm']:.4f} N·m")
    print(f"motor inertia torque           = {results['motor_torque_inertia_Nm']:.4f} N·m")
    print(f"motor peak torque (est)        = {results['motor_torque_peak_Nm']:.4f} N·m")
    print(f"motor peak torque w/ safety    = {results['motor_torque_peak_with_safety_Nm']:.4f} N·m")

    # print("\nNOTE:", results["notes"])
    print("===================================================\n")


# -------------------------
# Example usage
# -------------------------
if __name__ == "__main__":
    specs = CartPoleSpecs(
        l=0.30,
        theta_max_deg=10,
        theta_dot_max=1.0,
        x_dot0=0.0
    )

    rail = RailDesign(
        margin_frac=0.75,
        extra_margin_m=0.02
    )

    motion = MotionPlan(
        catch_time=0.30
    )

    drive = DriveTrain(
        pulley_radius=0.02,
        gear_ratio=1.0,
        efficiency=0.90
    )

    mass = MassAndFriction(
        M_cart=1.0,
        m_pole=0.2,
        F_friction=2.0,
        F_disturbance=0.0
    )

    inertia = MotorInertia(
        J_motor=0.0,
        J_extra=0.0
    )

    results = compute_requirements(
        specs, rail, motion, drive, mass, inertia,
        safety_factor=2.5
    )
    pretty_print(results)