6DoF Pdg Stc Ifthen

6-DoF Powered Descent Guidance with compound STCs via ox.stl.IfThen

Same physics, boundaries, and trigger thresholds as 6DoF_pdg_stc.py, but the state-triggered rules are written as STL implications instead of hand-rolled ReLU products.

How the STCs are encoded

Each rule is an implication IfThen(trigger, consequent) with lite=True so the GMSR residual integrates cleanly under CTCS. Compound triggers use ox.stl.And / ox.stl.Or:

IfThen(h < 100 m, |δ_e| ≤ 1°)
IfThen(And(h < 100 m, h > 2 m), tight glideslope)
IfThen(And(||v|| < 35, tilt < 60°), single-engine thrust)
IfThen(Or(||v|| > 35, tilt > 60°), three-engine thrust)

Compound state-triggered constraints

h < 100 m                       → tight gimbal, tilt, glideslope
h < 110 m  (alpha-tightened)    → tight angular-rate and speed
h < 200 m  (220 m alpha) & h>2  → LOS boresight cone toward the pad
||v|| < 35 m/s  AND  tilt < 60° → single-engine thrust limits
||v|| > 35 m/s  OR   tilt > 60° → three-engine thrust limits

Sibling of 6DoF_pdg_stc.py (custom product encoding). Reference: CT-cSTC/CT-cSTC.ipynb

When run as a script, launches three viser windows after solving: 1. Animated trajectory – thrust plume, attitude frame, velocity-colored trail 2. SCP convergence – node positions across iterations 3. Snapshot grid – evenly-spaced body poses along the final path

Viser scene uses the same ENU frame as the model (x, y horizontal; z = altitude up).

File: examples/rocket/6DoF_pdg_stc_ifthen.py

import os
import sys

import numpy as np

current_dir = os.path.dirname(os.path.abspath(__file__))
grandparent_dir = os.path.dirname(os.path.dirname(current_dir))
sys.path.append(grandparent_dir)

import openscvx as ox
from examples.plotting_viser import (
    AnimatedServerHandle,
    create_animated_plotting_server,
    create_scp_animated_plotting_server,
    create_snapshot_plotting_server,
)
from examples.rocket._cstc_plotting import (
    CstcLimits,
    CstcScales,
    add_altitude_trigger_discs,
    add_cstc_phase_markers,
    add_los_viewcone,
    add_site_markers,
    apply_panel_dark_theme,
    node_trigger_indices,
    plot_cstc_panel,
    plot_cstc_trajectory_3d,
    prepare_for_viser,
    save_plotly_figure,
)
from openscvx import Problem
from openscvx.plotting.viser import add_animation_controls, add_glideslope_cone

# ── Physical parameters (notebook cell 16 / cell 33) ──────────────────────────
G0 = 9.806  # m/s²
ISP = 330.0  # s
M_WET = 100_000.0  # kg
M_DRY = 85_000.0  # kg

# Initial conditions
R_I_INIT = np.array([200.0, 200.0, 500.0])  # m
V_I_INIT = np.array([0.0, 0.0, -50.0])  # m/s
# 90° tilt about x-axis: euler_to_quat([90,0,0]) → [w, x, y, z] = [√2/2, √2/2, 0, 0]
# OpenSCvx quaternion convention [x, y, z, w]: Q_INIT = [√2/2, 0, 0, √2/2]
Q_INIT = np.array([np.sin(np.pi / 4), 0.0, 0.0, np.cos(np.pi / 4)])
W_INIT = np.zeros(3)  # rad/s

# Terminal conditions
R_I_FINAL = np.array([0.0, 0.0, 0.001])  # m — touchdown target (slightly above pad)
GS_LOS_APEX_OFFSET_M = 0.001  # m — GS / LoS cone apex sits below touchdown
R_I_APEX = R_I_FINAL - np.array([0.0, 0.0, GS_LOS_APEX_OFFSET_M])  # cone vertex (pad)
V_I_FINAL = np.array([0.0, 0.0, -5.0])  # m/s (gentle touchdown)
Q_FINAL = np.array([0.0, 0.0, 0.0, 1.0])  # upright
W_FINAL = np.zeros(3)

# Thrust limits (N)
T_MAX = 2_200_000.0 * 3.0  # 3-engine
T_MIN = 2_200_000.0 * 0.4 * 3.0
T_MAX_AFT = 2_200_000.0  # 1-engine aft phase
T_MIN_AFT = 2_200_000.0 * 0.4

# Hard control angle limits (deg)
DELTA_ENGINE_MAX_DEG = 10.0
DELTA_BORESIGHT_MAX_DEG = 20.0

# State constraint limits
THETA_MAX_DEG = 90.0  # max tilt from vertical
W_B_MAX_RAD_S = np.pi / 2  # max angular rate
GS_MAX_DEG = 90.0 - 35.0  # glideslope cone half-angle from horizontal = 55°

# cSTC trigger thresholds (notebook cell 16)
ALT_TRIGGER_H1_M = 100.0  # h < 100 m  → tight gimbal / tilt / glideslope
ALT_TRIGGER_H2_M = 200.0  # h < 200 m  → LOS boresight cone
ALT_DONE_M = 2.0  # h > 2 m    → finalizer (disable cone near touchdown)
SPD_STC_TRIG = 35.0  # m/s        → speed trigger for single/multi-engine
THETA_STC_TRIG = 60.0  # deg        → tilt trigger for single/multi-engine

# Trigger alpha-tightening factors (notebook cell 16/33)
_ALPHA_ALT = 1.1  # altitude trigger tightening
_ALPHA_ALT_LOS = 1.1
_ALPHA_T_MIN = 1.02

# cSTC constraint alpha-scaling factors (notebook cell 16)
_ALPHA_SPD = 0.95
_ALPHA_OMEGA = 0.99
_ALPHA_THETA = 0.92
_ALPHA_GS = 0.92
_ALPHA_LOS = 0.98

# cSTC tight constraint values (alpha applied)
DELTA_STC_DEG = 1.0
THETA_STC_DEG = 5.0 * _ALPHA_THETA  # ≈ 4.6°
OMEGA_STC_RAD_S = 2.5 * np.pi / 180 * _ALPHA_OMEGA
V_STC_CONS = 20.0 * _ALPHA_SPD  # ≈ 19 m/s
GS_STC_DEG = (90.0 - 5.0) * _ALPHA_GS  # ≈ 78.2°
LOS_STC_DEG = 5.0 * _ALPHA_LOS  # ≈ 4.9°

# Inertia (diagonal, kg·m²)
J_B_DIAG = np.array([60.0, 60.0, 1.5])
J_B_INV_DIAG = np.array([1 / 60.0, 1 / 60.0, 1 / 1.5])

# Gravity (m/s²)
G_I = np.array([0.0, 0.0, -G0])

# Fuel consumption  ṁ = −α_m · T
ALPHA_M = 1.0 / (ISP * G0)

# Aerodynamics
RHO_AIR = 1.225  # kg/m³
S_AREA = 545.0  # m²
C_AERO_DIAG = np.array([0.4068, 0.4068, 0.0522])

# Lever arms (m) — engine to CM, aero-centre to engine
R_CM = np.array([0.0, 0.0, -14.0])
R_CP = np.array([0.0, 0.0, 3.0])

# ── Scaling (notebook cell 18 / cell 37) ──────────────────────────────────────
R_SCALE = np.linalg.norm(R_I_INIT)  # ≈ 574.46 m
M_SCALE = M_WET

_alpha_m_s = ALPHA_M * R_SCALE
_g_I_s = G_I / R_SCALE
_r_cm_s = R_CM / R_SCALE
_r_cp_s = R_CP / R_SCALE
_J_B_s = J_B_DIAG / R_SCALE**2
_J_B_inv_s = J_B_INV_DIAG * R_SCALE**2
_rho_s = RHO_AIR / (M_SCALE * R_SCALE) * R_SCALE**2  # = RHO_AIR * R_SCALE / M_SCALE

_T_max_s = T_MAX / (M_SCALE * R_SCALE)
_T_min_s = T_MIN / (M_SCALE * R_SCALE)
_T_max_aft_s = T_MAX_AFT / (M_SCALE * R_SCALE)
_T_min_aft_s = T_MIN_AFT / (M_SCALE * R_SCALE)

_r_init_s = R_I_INIT / R_SCALE
_r_final_s = R_I_FINAL / R_SCALE
_r_apex_s = R_I_APEX / R_SCALE
_v_init_s = V_I_INIT / R_SCALE
_v_final_s = V_I_FINAL / R_SCALE
_m_wet_s = M_WET / M_SCALE  # 1.0
_m_dry_s = M_DRY / M_SCALE  # 0.85

_v_stc_s = V_STC_CONS / R_SCALE


# ── Precomputed constraint thresholds ─────────────────────────────────────────
def _tilt_sq_bound(theta_deg: float) -> float:
    """Squared quaternion tilt bound: ||[q_x, q_y]||² ≤ (1-cos θ)/2."""
    return (1.0 - np.cos(np.pi / 180 * theta_deg)) / 2.0


_tilt_sq_max = _tilt_sq_bound(THETA_MAX_DEG)  # 0.5
_tilt_sq_stc = _tilt_sq_bound(THETA_STC_DEG)  # tight tilt (≈4.6°)
_tilt_sq_trig = _tilt_sq_bound(THETA_STC_TRIG)  # tilt trigger (60° → 0.25)
_tan_gs_max = np.tan(np.pi / 180 * GS_MAX_DEG)
_tan_gs_stc = np.tan(np.pi / 180 * GS_STC_DEG)
_cos_psi_stc = np.cos(np.pi / 180 * LOS_STC_DEG)
_delta_stc_rad = np.pi / 180 * DELTA_STC_DEG
_omega_sq_stc = OMEGA_STC_RAD_S**2

# Scaled altitude / speed trigger thresholds
_alt_h1_s = ALT_TRIGGER_H1_M / R_SCALE
_alt_h1a_s = _ALPHA_ALT * ALT_TRIGGER_H1_M / R_SCALE
_alt_h2a_s = _ALPHA_ALT_LOS * ALT_TRIGGER_H2_M / R_SCALE
_alt_done_s = ALT_DONE_M / R_SCALE
_spd_trig_s = SPD_STC_TRIG / R_SCALE

# ── Discretization ────────────────────────────────────────────────────────────
N = 15

# ── States ────────────────────────────────────────────────────────────────────
mass = ox.State("mass", shape=(1,))
mass.max = [_m_wet_s]
mass.min = [_m_dry_s]
mass.initial = [_m_wet_s]
mass.final = [ox.Maximize(_m_dry_s)]  # fuel-optimal objective

position = ox.State("position", shape=(3,))
position.max = [1.5, 1.5, 1.5]
position.min = [-1.5, -1.5, 0.0]  # z ≥ 0 (above ground)
position.initial = [
    ox.Free(float(_r_init_s[0])),
    ox.Free(float(_r_init_s[1])),
    ox.Free(float(_r_init_s[2])),
]
position.final = [float(_r_final_s[0]), float(_r_final_s[1]), float(_r_final_s[2])]

velocity = ox.State("velocity", shape=(3,))
_v_box = 150.0 / R_SCALE
velocity.max = [_v_box, _v_box, _v_box]
velocity.min = [-_v_box, -_v_box, -_v_box]
velocity.initial = [float(_v_init_s[0]), float(_v_init_s[1]), float(_v_init_s[2])]
velocity.final = [float(_v_final_s[0]), float(_v_final_s[1]), float(_v_final_s[2])]

attitude = ox.State("attitude", shape=(4,))
attitude.max = [1.0, 1.0, 1.0, 1.0]
attitude.min = [-1.0, -1.0, -1.0, -1.0]
attitude.initial = [
    ox.Free(float(Q_INIT[0])),
    ox.Free(float(Q_INIT[1])),
    ox.Free(float(Q_INIT[2])),
    ox.Free(float(Q_INIT[3])),
]
attitude.final = list(Q_FINAL)

angular_velocity = ox.State("angular_velocity", shape=(3,))
angular_velocity.max = [W_B_MAX_RAD_S, W_B_MAX_RAD_S, W_B_MAX_RAD_S]
angular_velocity.min = [-W_B_MAX_RAD_S, -W_B_MAX_RAD_S, -W_B_MAX_RAD_S]
angular_velocity.initial = [1e-8, 0.0, 0.0]
angular_velocity.final = [0.0, 0.0, 0.0]

# ── Controls: [T, δ_e, φ_e, δ_b, φ_b] ────────────────────────────────────────
thrust_mag = ox.Control("thrust_mag", shape=(1,))
thrust_mag.max = [_T_max_s]
thrust_mag.min = [_T_min_aft_s]  # always at least single-engine minimum
thrust_mag.guess = np.full((N, 1), float((_T_max_s + _T_min_aft_s) / 2))

gimbal_elev = ox.Control("gimbal_elev", shape=(1,))
gimbal_elev.max = [np.radians(DELTA_ENGINE_MAX_DEG)]
gimbal_elev.min = [-np.radians(DELTA_ENGINE_MAX_DEG)]
gimbal_elev.guess = np.zeros((N, 1))

gimbal_az = ox.Control("gimbal_az", shape=(1,))
gimbal_az.max = [np.pi]
gimbal_az.min = [-np.pi]
gimbal_az.guess = np.zeros((N, 1))

los_elev = ox.Control("los_elev", shape=(1,))
los_elev.max = [np.radians(DELTA_BORESIGHT_MAX_DEG)]
los_elev.min = [-np.radians(DELTA_BORESIGHT_MAX_DEG)]
los_elev.guess = np.full((N, 1), np.radians(DELTA_BORESIGHT_MAX_DEG) * 0.5)

los_az = ox.Control("los_az", shape=(1,))
los_az.max = [np.radians(DELTA_BORESIGHT_MAX_DEG)]
los_az.min = [-np.radians(DELTA_BORESIGHT_MAX_DEG)]
los_az.guess = np.zeros((N, 1))

# ── State & Control ───────────────────────────────────────────────────────────
states = [mass, position, velocity, attitude, angular_velocity]
controls = [thrust_mag, gimbal_elev, gimbal_az, los_elev, los_az]

# ── Quaternion kinematics ─────────────────────────────────────────────────────
# OpenSCvx convention: attitude = [q_x, q_y, q_z, q_w]
q1, q2, q3, q4 = attitude[0], attitude[1], attitude[2], attitude[3]

# CBI: inertial→body DCM  (CBI.T is body→inertial)
CBI = ox.Block(
    [
        [q4**2 + q1**2 - q2**2 - q3**2, 2 * (q1 * q2 - q4 * q3), 2 * (q4 * q2 + q1 * q3)],
        [2 * (q4 * q3 + q1 * q2), q4**2 - q1**2 + q2**2 - q3**2, 2 * (q2 * q3 - q4 * q1)],
        [2 * (q1 * q3 - q4 * q2), 2 * (q4 * q1 + q2 * q3), q4**2 - q1**2 - q2**2 + q3**2],
    ]
).T

w1, w2, w3 = angular_velocity[0], angular_velocity[1], angular_velocity[2]
q1_dot = 0.5 * (w1 * q4 - w2 * q3 + w3 * q2)
q2_dot = 0.5 * (w1 * q3 - w3 * q1 + w2 * q4)
q3_dot = 0.5 * (w2 * q1 - w1 * q2 + w3 * q4)
q4_dot = -0.5 * (w1 * q1 + w2 * q2 + w3 * q3)
attitude_dot = ox.Concat(q1_dot, q2_dot, q3_dot, q4_dot)


def cross(a, b):
    """Symbolic 3-D cross product."""
    return ox.Concat(
        a[1] * b[2] - a[2] * b[1],
        a[2] * b[0] - a[0] * b[2],
        a[0] * b[1] - a[1] * b[0],
    )


# ── Thrust vector in body frame (polar parameterisation) ──────────────────────
T = thrust_mag[0]
de = gimbal_elev[0]
pe = gimbal_az[0]
T_B = ox.Concat(
    T * ox.Sin(de) * ox.Cos(pe),
    T * ox.Sin(de) * ox.Sin(pe),
    T * ox.Cos(de),
)

# ── Aerodynamic drag in body frame ────────────────────────────────────────────
C_aero_ox = ox.Diag(ox.Concat(float(C_AERO_DIAG[0]), float(C_AERO_DIAG[1]), float(C_AERO_DIAG[2])))
A_B = -0.5 * _rho_s * ox.linalg.Norm(velocity) * S_AREA * C_aero_ox @ CBI @ velocity

# ── Inertia tensors (diagonal, scaled) ────────────────────────────────────────
# Mass-varying model: J_eff = J_s × m_s  →  ω̇ = J_inv_s @ (τ/m_s − ω × J_s @ ω)
J_B_ox = ox.Diag(ox.Concat(float(_J_B_s[0]), float(_J_B_s[1]), float(_J_B_s[2])))
J_B_inv_ox = ox.Diag(ox.Concat(float(_J_B_inv_s[0]), float(_J_B_inv_s[1]), float(_J_B_inv_s[2])))
r_cm_ox = ox.Concat(float(_r_cm_s[0]), float(_r_cm_s[1]), float(_r_cm_s[2]))
r_cp_ox = ox.Concat(float(_r_cp_s[0]), float(_r_cp_s[1]), float(_r_cp_s[2]))
g_I_ox = ox.Concat(float(_g_I_s[0]), float(_g_I_s[1]), float(_g_I_s[2]))

torque = cross(r_cm_ox, T_B) + cross(r_cp_ox, A_B)
gyro = cross(angular_velocity, J_B_ox @ angular_velocity)

# ── Dynamics (all in scaled units) ────────────────────────────────────────────
dynamics = {
    "mass": -_alpha_m_s * T,
    "position": velocity,
    "velocity": CBI.T @ (T_B + A_B) / mass[0] + g_I_ox,
    "attitude": attitude_dot,
    "angular_velocity": J_B_inv_ox @ (torque / mass[0] - gyro),
}

# ── Shared sub-expressions for constraints ────────────────────────────────────
# Tilt: ||[q_x, q_y]||²  (matches notebook's ||[x[8], x[9]]||² in [w,x,y,z] ordering)
tilt_sq = attitude[0] ** 2 + attitude[1] ** 2
speed = ox.linalg.Norm(velocity)
omega_sq = ox.linalg.Norm(angular_velocity) ** 2
z_alt = position[2]  # height above pad (z = 0); triggers use absolute altitude

# Position relative to GS / LoS cone apex (below touchdown, not at touchdown)
pos_rel_x = position[0] - float(_r_apex_s[0])
pos_rel_y = position[1] - float(_r_apex_s[1])
pos_rel_z = position[2] - float(_r_apex_s[2])
r_xy_norm = ox.linalg.Norm(ox.Concat(pos_rel_x, pos_rel_y))
r_los_norm = ox.linalg.Norm(ox.Concat(pos_rel_x, pos_rel_y, pos_rel_z))

# LOS boresight in inertial frame via body→inertial rotation
db = los_elev[0]
pb = los_az[0]
los_B = ox.Concat(
    ox.Sin(db) * ox.Cos(pb),
    ox.Sin(db) * ox.Sin(pb),
    ox.Cos(db),
)
los_I = CBI.T @ los_B  # unit-norm boresight in inertial frame
r_dot_los = pos_rel_x * los_I[0] + pos_rel_y * los_I[1] + pos_rel_z * los_I[2]


# ── Trigger / consequent predicates (Constraint objects for IfThen) ───────────
# Triggers (antecedents)
alt_lt_100 = z_alt <= _alt_h1_s  # h < 100 m
alt_lt_110 = z_alt <= _alt_h1a_s  # h < 110 m (alpha-tightened)
alt_lt_220 = z_alt <= _alt_h2a_s  # h < 220 m (LOS, alpha)
alt_gt_2 = z_alt >= _alt_done_s  # h >   2 m (finalizer)
v_lt_35 = speed <= _spd_trig_s  # ||v|| < 35 m/s
v_gt_35 = speed >= _spd_trig_s  # ||v|| > 35 m/s
tilt_lt_60 = tilt_sq <= _tilt_sq_trig  # tilt < 60°
tilt_gt_60 = tilt_sq >= _tilt_sq_trig  # tilt > 60°

# GMSR smoothing shared by all STC predicates (lite for CTCS integration)
_STC_C = 1e-12
_STC_KW = dict(c=_STC_C, lite=True)
_STC_OVER = dict(penalty="huber")

# Consequents
gimbal_tight = ox.stl.And(
    de <= _delta_stc_rad,
    -de <= _delta_stc_rad,
    **_STC_KW,
)  # |δ_e| ≤ 1°
gs_tight = r_xy_norm * _tan_gs_stc - pos_rel_z <= 0  # tight glideslope
omega_tight = omega_sq <= _omega_sq_stc  # tight angular rate
tilt_tight = tilt_sq <= _tilt_sq_stc  # tight tilt
spd_tight = speed <= _v_stc_s  # tight speed
los_tight = r_los_norm * _cos_psi_stc - r_dot_los <= 0  # LOS cone
thrust_aft = ox.stl.And(
    T <= _T_max_aft_s,
    T >= _T_min_aft_s,
    **_STC_KW,
)  # single-engine band
thrust_multi = ox.stl.And(
    T <= _T_max_s,
    T >= _ALPHA_T_MIN * _T_min_s,
    **_STC_KW,
)  # three-engine band

# Compound triggers
alt100_and_above_pad = ox.stl.And(alt_lt_100, alt_gt_2, **_STC_KW)
alt220_and_above_pad = ox.stl.And(alt_lt_220, alt_gt_2, **_STC_KW)
aft_phase = ox.stl.And(v_lt_35, tilt_lt_60, **_STC_KW)
multi_phase = ox.stl.Or(v_gt_35, tilt_gt_60, **_STC_KW)

# ── Constraints ───────────────────────────────────────────────────────────────
constraints = []

# Boundary conditions (convex equality constraints)
constraints.append((position == _r_init_s).convex().at([0]))
constraints.append((attitude == Q_INIT).convex().at([0]))
constraints.append((velocity == _v_init_s).convex().at([0]))
constraints.append((angular_velocity == W_INIT).convex().at([0]))
constraints.append((position == _r_final_s).convex().at([N - 1]))
constraints.append((velocity == _v_final_s).convex().at([N - 1]))
constraints.append((attitude == Q_FINAL).convex().at([N - 1]))
constraints.append((angular_velocity == W_FINAL).convex().at([N - 1]))

# ── Always-on CTCS (entire trajectory) ────────────────────────────────────────
constraints.append(ox.ctcs(tilt_sq - _tilt_sq_max <= 0))  # tilt ≤ 90°
constraints.append(ox.ctcs(omega_sq - W_B_MAX_RAD_S**2 <= 0, penalty="huber"))  # angular rate
# glideslope 55°
constraints.append(ox.ctcs(r_xy_norm * _tan_gs_max - pos_rel_z <= 0, penalty="huber"))
constraints.append(ox.ctcs(_m_dry_s - mass[0] <= 0, penalty="huber"))  # dry-mass floor

# ── Compound state-triggered constraints via IfThen ───────────────────────────
# lite=True: residual is non-negative and integrates under CTCS like the
# hand-rolled ReLU products in 6DoF_pdg_stc.py. Small c → tight Boolean approx.

# h < 100 m → tight gimbal / tilt
constraints.append(ox.stl.IfThen(alt_lt_100, gimbal_tight, **_STC_KW).over((0, N - 1), **_STC_OVER))
constraints.append(ox.stl.IfThen(alt_lt_100, tilt_tight, **_STC_KW).over((0, N - 1), **_STC_OVER))
# h < 100 m AND h > 2 m → tight glideslope cone
constraints.append(
    ox.stl.IfThen(alt100_and_above_pad, gs_tight, **_STC_KW).over((0, N - 1), **_STC_OVER)
)
# h < 110 m → tight angular rate and tight speed
constraints.append(ox.stl.IfThen(alt_lt_110, omega_tight, **_STC_KW).over((0, N - 1), **_STC_OVER))
constraints.append(ox.stl.IfThen(alt_lt_110, spd_tight, **_STC_KW).over((0, N - 1), **_STC_OVER))
# h < 220 m AND h > 2 m → LOS boresight cone toward R_I_APEX
constraints.append(
    ox.stl.IfThen(alt220_and_above_pad, los_tight, **_STC_KW).over((0, N - 1), **_STC_OVER)
)

# ||v|| < 35 m/s AND tilt < 60° → single-engine thrust limits
constraints.append(ox.stl.IfThen(aft_phase, thrust_aft, **_STC_KW).over((0, N - 1), **_STC_OVER))
# ||v|| > 35 m/s OR tilt > 60° → three-engine thrust limits
constraints.append(
    ox.stl.IfThen(multi_phase, thrust_multi, **_STC_KW).over((0, N - 1), **_STC_OVER)
)

# ── Time (free final time with per-segment dilation) ─────────────────────────
_t_f_guess = 21.0

time = ox.Time(
    initial=0.0, final=ox.Free(_t_f_guess), min=0.0, max=_t_f_guess * 1.5, uniform_time_grid=True
)

# ── Problem Assembly ──────────────────────────────────────────────────────────
problem = Problem(
    N=N,
    states=states,
    controls=controls,
    dynamics=dynamics,
    constraints=constraints,
    time=time,
    float_dtype="float64",
    licq_max=1e-6,
    algorithm={
        # A high *constant* constraint-violation weight (lam_vc) plays the role
        # of the notebook's w_con_dyn, driving the integrated STC penalty to
        # zero.  PTR (the default acceptance-ratio loop) keeps the SCP stable;
        # the adaptive AugmentedLagrangian was found to overshoot and diverge.
        "lam_vc": 4e0,
        "lam_cost": 1e-2,
        "k_max": 800,
        "autotuner": ox.ConstantProximalWeight(),
        # "autotuner": ox.AugmentedLagrangian(ep=1E-1, eta_lambda=1e3),
    },
    discretizer={
        "diffrax_kwargs": {"atol": 1e-6, "rtol": 1e-6},
    },
)

# Prop tolerances
problem.settings.prp.atol = 1e-12
problem.settings.prp.rtol = 1e-12

# ── Viser display parameters ──────────────────────────────────────────────────
# cSTC uses ENU inertial frame (x, y horizontal; z = altitude). Viser is Z-up,
# so positions map directly without any axis permutation.
SCENE_SCALE = 10.0  # 1 viser unit = 10 m
PLUME_SCALE = 8.0
ATTITUDE_AXES_LENGTH = 2.0
VIEWCONE_SCALE = 4.0
_VISER_UP_AXIS = (0.0, 0.0, 1.0)
GS_HALFANGLE_DEG = 90.0 - GS_MAX_DEG
GS_STC_HALFANGLE_DEG = 90.0 - GS_STC_DEG

SCALES = CstcScales(r_scale=R_SCALE, m_scale=M_SCALE)
LIMITS = CstcLimits(
    gs_apex_m=R_I_APEX,
    los_apex_m=R_I_APEX,
    pad_z_m=0.0,
    alt_trigger_h1_m=ALT_TRIGGER_H1_M,
    alt_trigger_h2_m=ALT_TRIGGER_H2_M,
    spd_stc_trig=SPD_STC_TRIG,
    theta_stc_trig_deg=THETA_STC_TRIG,
    t_max=T_MAX,
    t_min=T_MIN,
    t_max_aft=T_MAX_AFT,
    t_min_aft=T_MIN_AFT,
    v_stc_cons=V_STC_CONS,
    theta_max_deg=THETA_MAX_DEG,
    theta_stc_deg=THETA_STC_DEG,
    delta_engine_max_deg=DELTA_ENGINE_MAX_DEG,
    delta_stc_deg=DELTA_STC_DEG,
    w_b_max_rad_s=W_B_MAX_RAD_S,
    omega_stc_rad_s=OMEGA_STC_RAD_S,
    gs_max_deg=GS_MAX_DEG,
    gs_stc_deg=GS_STC_DEG,
    delta_boresight_max_deg=DELTA_BORESIGHT_MAX_DEG,
    los_stc_deg=LOS_STC_DEG,
    m_dry=M_DRY,
)

SITE_MARKERS = [
    ("/cone_apex", R_I_APEX, (255, 220, 80), 0.08),
    ("/landing_pad", R_I_FINAL, (50, 255, 80), 0.12),
]
ALT_TRIGGERS = [
    (ALT_TRIGGER_H2_M, (80, 160, 255), "h < 200 m → LOS boresight cone"),
    (ALT_TRIGGER_H1_M, (255, 80, 80), "h < 100 m → tight gimbal/tilt/ω/speed/GS"),
]


def launch_viser_servers(result) -> None:
    """Create trajectory, SCP convergence, and snapshot viser servers."""
    pos = np.asarray(result.trajectory["position"])
    node_indices = node_trigger_indices(result.nodes, LIMITS, scales=SCALES)

    handle = create_animated_plotting_server(
        result,
        thrust_key="thrust_body",
        thrust_style="plume",
        thrust_scale=PLUME_SCALE,
        thrust_plume_half_angle_deg=13.0,
        thrust_plume_color=(255, 140, 40),
        thrust_plume_opacity=0.55,
        thrust_remap_world_to_viser=False,
        attitude_key="attitude",
        attitude_axes_length=ATTITUDE_AXES_LENGTH,
        show_viewcone=False,
        trail_point_size=0.08,
        show_grid=True,
        dark_mode=True,
        scene_scale=1.0,
        controls="manual",
    )
    assert isinstance(handle, AnimatedServerHandle)
    traj_server = handle.server
    _, update_viewcone = add_los_viewcone(
        traj_server, result, half_angle_deg=LOS_STC_DEG, scale=VIEWCONE_SCALE
    )
    add_animation_controls(
        traj_server, handle.traj_time, handle.update_callbacks + [update_viewcone], loop=True
    )

    # Two cones share the scene, so each needs its own scene path: the wide
    # always-on glideslope and the tight one the h < 100 m cSTC switches to.
    apex_vis = tuple((R_I_APEX / SCENE_SCALE).tolist())
    add_glideslope_cone(
        traj_server,
        apex=apex_vis,
        height=float(np.max(pos[:, 2])) * 1.15,
        glideslope_angle_deg=GS_HALFANGLE_DEG,
        axis=_VISER_UP_AXIS,
        color=(80, 200, 80),
        opacity=0.12,
        name="/constraints/glideslope_wide",
    )
    add_glideslope_cone(
        traj_server,
        apex=apex_vis,
        height=ALT_TRIGGER_H1_M / SCENE_SCALE,
        glideslope_angle_deg=GS_STC_HALFANGLE_DEG,
        axis=_VISER_UP_AXIS,
        color=(255, 180, 40),
        opacity=0.15,
        name="/constraints/glideslope_tight",
    )

    add_altitude_trigger_discs(
        traj_server,
        pos,
        center_xy=(0.0, 0.0),
        base_z_m=0.0,
        scene_scale=SCENE_SCALE,
        triggers=ALT_TRIGGERS,
    )
    add_site_markers(traj_server, SITE_MARKERS, scene_scale=SCENE_SCALE)
    add_cstc_phase_markers(traj_server, pos, node_indices, n_nodes=N, limits=LIMITS)

    scp_server = create_scp_animated_plotting_server(
        result,
        position_slice=slice(1, 4),
        attitude_slice=slice(7, 11),
        show_attitudes=True,
        attitude_stride=3,
        attitude_axes_length=ATTITUDE_AXES_LENGTH,
        frame_duration_ms=80,
        scene_scale=1.0,
    )
    snap_server = create_snapshot_plotting_server(
        result,
        attitude_axes_length=ATTITUDE_AXES_LENGTH,
        show_body_frame=True,
        initial_n_snapshots=6,
        show_grid=True,
        background_color=(240, 240, 245),
    )
    for server in (scp_server, snap_server):
        add_altitude_trigger_discs(
            server,
            pos,
            center_xy=(0.0, 0.0),
            base_z_m=0.0,
            scene_scale=SCENE_SCALE,
            triggers=ALT_TRIGGERS,
        )

    traj_server.sleep_forever()


def plot_cstc_results(
    result,
    *,
    show: bool = True,
    save_prefix: str = "cstc_ifthen",
    dark_mode: bool = False,
    save_dark: bool = True,
) -> tuple:
    """Nine-panel states/controls figure and the speed-coloured 3-D path.

    The panel is authored for print; ``save_dark`` additionally writes a
    dark-themed copy as ``{save_prefix}_states_controls_dark`` for slides, and
    ``dark_mode`` makes that copy the one returned and displayed.
    """
    from openscvx.plotting.publication import show_plotly_with_latin_modern

    fig_panel = plot_cstc_panel(result, LIMITS, scales=SCALES)
    fig_3d = plot_cstc_trajectory_3d(
        result,
        scales=SCALES,
        markers=[("Landing pad", R_I_FINAL, "lime")],
        title="6-DoF PDG Trajectory (IfThen cSTC)",
    )

    if save_prefix and not dark_mode:
        save_plotly_figure(fig_panel, f"{save_prefix}_states_controls")
    if dark_mode or (save_prefix and save_dark):
        import plotly.graph_objects as go

        fig_dark = go.Figure(fig_panel)
        apply_panel_dark_theme(fig_dark)
        if save_prefix and save_dark:
            save_plotly_figure(fig_dark, f"{save_prefix}_states_controls_dark")
        if dark_mode:
            fig_panel = fig_dark
    if save_prefix:
        save_plotly_figure(fig_3d, f"{save_prefix}_trajectory_3d")

    if show:
        show_plotly_with_latin_modern(fig_panel)
        show_plotly_with_latin_modern(fig_3d)

    return fig_panel, fig_3d


if __name__ == "__main__":
    problem.settings.dev.debug = True  # bypass JAX export (see repo issue)
    problem.initialize()
    result = problem.solve()
    result = problem.post_process()

    traj = result.trajectory
    pos = np.asarray(traj["position"]) * R_SCALE
    vel = np.asarray(traj["velocity"]) * R_SCALE
    m = np.asarray(traj["mass"]) * M_SCALE

    print("\n── Solution summary ─────────────────────────────────────────────")
    print(f"  Final position (m):   {pos[-1]}")
    print(f"  Final velocity (m/s): {vel[-1]}")
    print(f"  Final mass (kg):      {m[-1, 0]:.1f}  (dry mass: {M_DRY:.0f} kg)")
    print(f"  Fuel used (kg):       {M_WET - m[-1, 0]:.1f}")

    print("\n── Generating plots ─────────────────────────────────────────────")
    plot_cstc_results(result, show=True)

    prepare_for_viser(result, scales=SCALES, scene_scale=SCENE_SCALE)
    launch_viser_servers(result)