Ebike trajectory optimization

In [ ]:

"""Trajectory optimization for the smart cargo e-bike — optimize the assist
policy over the *true* multi-domain DAE simulation.

The objective is evaluated by running the full hybrid acausal model each
iteration:

    J(cap) = E_battery(cap; D) + lambda * max(0, v_target - v_mean(cap; D))^2

i.e. find the *minimum-battery* assist-torque cap that still covers a fixed
reference DISTANCE D at or above a target average speed. Two framing choices
matter more than the optimizer:

* **Per distance, not per time.** Energy measured over a fixed time window
  rewards a design for going slowly -- fewer metres means less climbing and
  less drag -- so its "saving" is partly just a slower bike. Both terms here
  are read at the moment the vehicle has covered D metres. Since the route
  grade is position-indexed (see ebike_hybrid_simulation), every candidate
  climbs the same hill.
* **A penalty is a trade, not a constraint.** An exterior penalty converges
  from the infeasible side with residual shortfall
  s* ~ (dE/dcap) / (2*lambda*dv/dcap), so lambda is sized from the response
  scale and the achieved speed is REPORTED, never assumed. Part 2 of the
  tutorial series plots the optimum against lambda.

A note on gradients
-------------------
Jaxonomy's simulator is differentiable (reverse-mode adjoint via
``SimulatorOptions(enable_autodiff=True)``), and that path is exercised in the
autodiff test-suite. For *this* model, however, end-to-end AD is numerically
fragile: the hybrid speed-cutoff event carries an integer mode variable that
cannot hold a reverse-mode cotangent, and the stiff multi-domain DAE adjoint
returns NaN gradients even in forward mode. Rather than overclaim, we optimize
with a derivative-free method over the true simulation (Brent's method on the
bounded scalar). This still genuinely optimizes *through the physics* — the loss
is the real DAE rollout — and we report a finite-difference gradient at the
optimum to confirm stationarity.
"""

"""Trajectory optimization for the smart cargo e-bike — optimize the assist policy over the *true* multi-domain DAE simulation. The objective is evaluated by running the full hybrid acausal model each iteration: J(cap) = E_battery(cap; D) + lambda * max(0, v_target - v_mean(cap; D))^2 i.e. find the *minimum-battery* assist-torque cap that still covers a fixed reference DISTANCE D at or above a target average speed. Two framing choices matter more than the optimizer: * **Per distance, not per time.** Energy measured over a fixed time window rewards a design for going slowly -- fewer metres means less climbing and less drag -- so its "saving" is partly just a slower bike. Both terms here are read at the moment the vehicle has covered D metres. Since the route grade is position-indexed (see ebike_hybrid_simulation), every candidate climbs the same hill. * **A penalty is a trade, not a constraint.** An exterior penalty converges from the infeasible side with residual shortfall s* ~ (dE/dcap) / (2*lambda*dv/dcap), so lambda is sized from the response scale and the achieved speed is REPORTED, never assumed. Part 2 of the tutorial series plots the optimum against lambda. A note on gradients ------------------- Jaxonomy's simulator is differentiable (reverse-mode adjoint via ``SimulatorOptions(enable_autodiff=True)``), and that path is exercised in the autodiff test-suite. For *this* model, however, end-to-end AD is numerically fragile: the hybrid speed-cutoff event carries an integer mode variable that cannot hold a reverse-mode cotangent, and the stiff multi-domain DAE adjoint returns NaN gradients even in forward mode. Rather than overclaim, we optimize with a derivative-free method over the true simulation (Brent's method on the bounded scalar). This still genuinely optimizes *through the physics* — the loss is the real DAE rollout — and we report a finite-difference gradient at the optimum to confirm stationarity. """

In [ ]:

from __future__ import annotations

from __future__ import annotations

In [ ]:

import os
import sys
import time

import os import sys import time

In [ ]:

import numpy as np
import jax.numpy as jnp
from scipy.optimize import minimize_scalar

import numpy as np import jax.numpy as jnp from scipy.optimize import minimize_scalar

In [ ]:

sys.path.append(os.path.dirname(os.path.abspath(__file__)))

sys.path.append(os.path.dirname(os.path.abspath(__file__)))

In [ ]:

import jaxonomy
from jaxonomy.simulation import SimulatorOptions
from ebike_hybrid_simulation import make_ebike_diagram, EbikeConfig, energy_audit

import jaxonomy from jaxonomy.simulation import SimulatorOptions from ebike_hybrid_simulation import make_ebike_diagram, EbikeConfig, energy_audit

In [ ]:

# --- objective over the real simulation --------------------------------------
D_REF = 100.0              # reference distance for the objective [m]
V_TARGET_KMH = 12.0        # average-speed floor over that segment
LAMBDA_SPEED = 2.0e4       # penalty weight; sized so the residual shortfall
                           # s* stays below ~0.05 km/h (see the module docstring)
TF = 45.0                  # horizon; long enough for every candidate to cover D_REF

# --- objective over the real simulation -------------------------------------- D_REF = 100.0 # reference distance for the objective [m] V_TARGET_KMH = 12.0 # average-speed floor over that segment LAMBDA_SPEED = 2.0e4 # penalty weight; sized so the residual shortfall # s* stays below ~0.05 km/h (see the module docstring) TF = 45.0 # horizon; long enough for every candidate to cover D_REF

In [ ]:

_diagram, _handles = make_ebike_diagram(EbikeConfig(), return_handles=True,
                                        enable_speed_event=True)
_ctx = _diagram.create_context()
_aid = _handles["assist_policy"].system_id
_ports = {p.name: p for p in _diagram.output_ports}
_opts = SimulatorOptions(enable_autodiff=False, rtol=5e-4, atol=5e-6, buffer_length=260000)

_diagram, _handles = make_ebike_diagram(EbikeConfig(), return_handles=True, enable_speed_event=True) _ctx = _diagram.create_context() _aid = _handles["assist_policy"].system_id _ports = {p.name: p for p in _diagram.output_ports} _opts = SimulatorOptions(enable_autodiff=False, rtol=5e-4, atol=5e-6, buffer_length=260000)

In [ ]:

def rollout(cap):
    """Run the full hybrid model with the given assist-torque cap [Nm].

    Returns (E_batt, v_mean, E_human) measured at the point the vehicle has
    covered D_REF metres. A design that never covers D_REF within TF is
    infeasible and reports inf.
    """
    c = _ctx.with_subcontext(_aid, _ctx[_aid].with_parameter("max_assist_torque", jnp.array(float(cap))))
    rec = {"E": _ports["E_batt_term"], "d": _ports["distance"], "H": _ports["E_human"]}
    res = jaxonomy.simulate(_diagram, c, (0.0, TF), options=_opts, recorded_signals=rec)
    o = res.outputs
    t = np.asarray(res.time).squeeze()
    d = np.asarray(o["d"]).squeeze()
    E = np.asarray(o["E"]).squeeze()
    H = np.asarray(o["H"]).squeeze()
    if d[-1] < D_REF:
        return float("inf"), 0.0, float("inf")
    t_D = float(np.interp(D_REF, d, t))
    return (float(np.interp(D_REF, d, E)), D_REF / t_D * 3.6,
            float(np.interp(D_REF, d, H)))

def rollout(cap): """Run the full hybrid model with the given assist-torque cap [Nm]. Returns (E_batt, v_mean, E_human) measured at the point the vehicle has covered D_REF metres. A design that never covers D_REF within TF is infeasible and reports inf. """ c = _ctx.with_subcontext(_aid, _ctx[_aid].with_parameter("max_assist_torque", jnp.array(float(cap)))) rec = {"E": _ports["E_batt_term"], "d": _ports["distance"], "H": _ports["E_human"]} res = jaxonomy.simulate(_diagram, c, (0.0, TF), options=_opts, recorded_signals=rec) o = res.outputs t = np.asarray(res.time).squeeze() d = np.asarray(o["d"]).squeeze() E = np.asarray(o["E"]).squeeze() H = np.asarray(o["H"]).squeeze() if d[-1] < D_REF: return float("inf"), 0.0, float("inf") t_D = float(np.interp(D_REF, d, t)) return (float(np.interp(D_REF, d, E)), D_REF / t_D * 3.6, float(np.interp(D_REF, d, H)))

In [ ]:

def objective(cap):
    E_batt, v_mean, _ = rollout(cap)
    if not np.isfinite(E_batt):
        return 1e9
    shortfall = max(0.0, V_TARGET_KMH - v_mean)
    return E_batt + LAMBDA_SPEED * shortfall ** 2

def objective(cap): E_batt, v_mean, _ = rollout(cap) if not np.isfinite(E_batt): return 1e9 shortfall = max(0.0, V_TARGET_KMH - v_mean) return E_batt + LAMBDA_SPEED * shortfall ** 2

In [ ]:

def run_optimization(cap_baseline=18.0, bounds=(6.0, 20.0)):
    print("=" * 68)
    print("  E-BIKE TRAJECTORY OPTIMIZATION (derivative-free, over true DAE)")
    print("=" * 68)
    print(f"  Objective : min E_battery(D={D_REF:.0f} m) + "
          f"{LAMBDA_SPEED:.3g}*max(0, {V_TARGET_KMH}-v_mean)^2")
    print(f"  Horizon   : {TF:.0f} s   |   assist-cap bounds: {bounds} Nm")
    print(f"  Baseline  : over-assisted cap = {cap_baseline} Nm\n")

    history = []

    def logged_objective(cap):
        J = objective(float(cap))
        history.append((float(cap), float(J)))
        print(f"    eval cap={float(cap):6.3f} Nm  ->  J={J:12.1f}")
        return J

    print("  Optimizing (Brent bounded)...")
    t0 = time.time()
    result = minimize_scalar(logged_objective, bounds=bounds, method="bounded",
                             options={"xatol": 0.3, "maxiter": 20})
    t_opt = time.time() - t0
    cap_opt = float(result.x)

    E0, v0, H0 = rollout(cap_baseline)
    E1, v1, H1 = rollout(cap_opt)
    batt_saved_pct = (E0 - E1) / E0 * 100.0        # over the SAME D_REF metres
    shortfall = max(0.0, V_TARGET_KMH - v1)

    print("\n" + "=" * 68)
    print(f"  {'METRIC':<28} | {'BASELINE':>14} | {'OPTIMIZED':>14}")
    print("=" * 68)
    print(f"  {'Assist-torque cap (Nm)':<28} | {cap_baseline:>14.2f} | {cap_opt:>14.2f}")
    print(f"  {'Objective J':<28} | {objective(cap_baseline):>14.1f} | {result.fun:>14.1f}")
    print(f"  {f'Battery energy over {D_REF:.0f} m (J)':<28} | {E0:>14.1f} | {E1:>14.1f}")
    print(f"  {'  ... per metre (J/m)':<28} | {E0/D_REF:>14.2f} | {E1/D_REF:>14.2f}")
    print(f"  {'Mean speed (km/h)':<28} | {v0:>14.2f} | {v1:>14.2f}  (floor {V_TARGET_KMH})")
    print(f"  {'Human work over that segment':<28} | {H0:>14.1f} | {H1:>14.1f}")
    print("=" * 68)
    print(f"  Battery energy to cover the SAME {D_REF:.0f} m reduced {batt_saved_pct:.1f}%.")
    print(f"  Constraint: floor {V_TARGET_KMH:.1f} km/h, achieved {v1:.2f} km/h "
          f"(shortfall {shortfall:+.3f} km/h)")
    if shortfall > 1e-3:
        print("    -> the optimum is slightly INFEASIBLE: that is what an exterior")
        print("       penalty does (converges from the infeasible side). Raise lambda")
        print("       or switch to an augmented Lagrangian if feasibility must hold.")
    else:
        print("    -> feasible at this lambda; check Part 2's lambda sweep for how")
        print("       close to the boundary the penalty actually pins the optimum.")
    print(f"  Converged in {t_opt:.0f}s over {len(history)} real-simulation rollouts.")
    print("  Discretization noise on J is MEASURED in Part 2 (tolerance study), and")
    print("  the optimizer's tolerances are set above it -- do not read the optimum")
    print("  to finer resolution than that floor allows.")
    print("=" * 68)
    return cap_opt, history

def run_optimization(cap_baseline=18.0, bounds=(6.0, 20.0)): print("=" * 68) print(" E-BIKE TRAJECTORY OPTIMIZATION (derivative-free, over true DAE)") print("=" * 68) print(f" Objective : min E_battery(D={D_REF:.0f} m) + " f"{LAMBDA_SPEED:.3g}*max(0, {V_TARGET_KMH}-v_mean)^2") print(f" Horizon : {TF:.0f} s | assist-cap bounds: {bounds} Nm") print(f" Baseline : over-assisted cap = {cap_baseline} Nm\n") history = [] def logged_objective(cap): J = objective(float(cap)) history.append((float(cap), float(J))) print(f" eval cap={float(cap):6.3f} Nm -> J={J:12.1f}") return J print(" Optimizing (Brent bounded)...") t0 = time.time() result = minimize_scalar(logged_objective, bounds=bounds, method="bounded", options={"xatol": 0.3, "maxiter": 20}) t_opt = time.time() - t0 cap_opt = float(result.x) E0, v0, H0 = rollout(cap_baseline) E1, v1, H1 = rollout(cap_opt) batt_saved_pct = (E0 - E1) / E0 * 100.0 # over the SAME D_REF metres shortfall = max(0.0, V_TARGET_KMH - v1) print("\n" + "=" * 68) print(f" {'METRIC':<28} | {'BASELINE':>14} | {'OPTIMIZED':>14}") print("=" * 68) print(f" {'Assist-torque cap (Nm)':<28} | {cap_baseline:>14.2f} | {cap_opt:>14.2f}") print(f" {'Objective J':<28} | {objective(cap_baseline):>14.1f} | {result.fun:>14.1f}") print(f" {f'Battery energy over {D_REF:.0f} m (J)':<28} | {E0:>14.1f} | {E1:>14.1f}") print(f" {' ... per metre (J/m)':<28} | {E0/D_REF:>14.2f} | {E1/D_REF:>14.2f}") print(f" {'Mean speed (km/h)':<28} | {v0:>14.2f} | {v1:>14.2f} (floor {V_TARGET_KMH})") print(f" {'Human work over that segment':<28} | {H0:>14.1f} | {H1:>14.1f}") print("=" * 68) print(f" Battery energy to cover the SAME {D_REF:.0f} m reduced {batt_saved_pct:.1f}%.") print(f" Constraint: floor {V_TARGET_KMH:.1f} km/h, achieved {v1:.2f} km/h " f"(shortfall {shortfall:+.3f} km/h)") if shortfall > 1e-3: print(" -> the optimum is slightly INFEASIBLE: that is what an exterior") print(" penalty does (converges from the infeasible side). Raise lambda") print(" or switch to an augmented Lagrangian if feasibility must hold.") else: print(" -> feasible at this lambda; check Part 2's lambda sweep for how") print(" close to the boundary the penalty actually pins the optimum.") print(f" Converged in {t_opt:.0f}s over {len(history)} real-simulation rollouts.") print(" Discretization noise on J is MEASURED in Part 2 (tolerance study), and") print(" the optimizer's tolerances are set above it -- do not read the optimum") print(" to finer resolution than that floor allows.") print("=" * 68) return cap_opt, history

In [ ]:

if __name__ == "__main__":
    run_optimization()

if __name__ == "__main__": run_optimization()