eng-lab/hwe/pipeline.py

"""Run every unlocked analysis on the design and collect the results.

The EL01 pipeline, in order (each stage reads the design and earlier stages only):
  1. cad        build the arm with YOUR hwe.core.arm, measure its volume and mass,
                and cross-check the volume against YOUR arm_volume_by_hand
  2. mass       the mass budget from [[component]] entries, with the arm's CAD mass
                in place of its estimate (YOUR hwe.core.massbudget)
  3. hover      T/W, hover thrust fraction, maximum climb acceleration
                (YOUR hwe.core.hover), from the thrust table's maximum thrust
  4. energy     battery energy, hover power from the table's efficiency at the
                hover thrust, and flight time (YOUR hwe.core.energy)
  5. sweep      the mass spiral: flight time as the battery grows (section 2.10)
  6. requirements  every requirement in design/requirements.toml, checked
  7. phone      the EL01 accelerometer test, if its data file exists
                (YOUR hwe.core.kinematics)

Results are stored flat under dotted names ("hover.thrust_to_weight"), always
as Quantities. A stage whose core code is not written yet says so in
`stages`; it is never skipped silently. Nothing here strips units except the
CAD boundary and the plotting code.
"""
from __future__ import annotations

import tomllib
import traceback
from dataclasses import dataclass, field
from pathlib import Path

from hwe import catalog, core, requirements as reqs
from hwe.design import ROOT, Design, Item, load
from hwe.units import Q_, Quantity

PROGRESS_FILE = ROOT / "design" / "progress.toml"

# Which module unlocks each stage. The report lists locked stages as the road ahead.
STAGES = [
    ("cad", "EL01", "the arm: CAD volume and mass, cross-checked by hand"),
    ("mass", "EL01", "mass budget with growth allowances"),
    ("hover", "EL01", "thrust-to-weight, hover thrust fraction, climb acceleration"),
    ("energy", "EL01", "battery energy, hover power, flight time"),
    ("sweep", "EL01", "the mass spiral: flight time against battery size"),
    ("requirements", "EL01", "requirement compliance"),
    ("phone", "EL01", "accelerometer test: predicted versus measured"),
    ("balance", "EL02", "centre of mass and each motor's share of the thrust"),
    ("momentum", "EL03", "momentum theory: hover power from first principles"),
    ("electrical", "EL04", "battery sag, motor operating point, currents"),
    ("trade", "EL05", "trade studies across the design space"),
    ("stress", "EL07", "arm stress and deflection"),
    ("fea", "EL08", "finite-element model of the arm"),
    ("modes", "EL09", "vibration modes against motor speeds"),
    ("clearance", "EL10", "clearances and tolerance stacks"),
]


@dataclass
class Run:
    design: Design
    results: dict[str, Quantity] = field(default_factory=dict)
    stages: dict[str, str] = field(default_factory=dict)       # stage -> "ok" | "not written: ..." | "error: ..." | "locked" | "no data"
    cross_checks: list[dict] = field(default_factory=list)
    budget: list[Item] = field(default_factory=list)
    by_subsystem: dict[str, Quantity] = field(default_factory=dict)
    sweep: dict | None = None
    requirement_status: list = field(default_factory=list)
    phone: dict | None = None
    part: object = None
    unlocked: list[str] = field(default_factory=list)


def unlocked_modules() -> list[str]:
    if not PROGRESS_FILE.exists():
        return ["EL01"]
    with open(PROGRESS_FILE, "rb") as f:
        return list(tomllib.load(f).get("unlocked", ["EL01"]))


def _stage(run: Run, name: str, func) -> None:
    module = next(m for s, m, _ in STAGES if s == name)
    if module not in run.unlocked:
        run.stages[name] = "locked"
        return
    try:
        outcome = func(run)
        run.stages[name] = outcome or "ok"
    except NotImplementedError as e:
        run.stages[name] = f"not written: {e}"
    except Exception as e:                                      # report it, keep going
        run.stages[name] = f"error: {type(e).__name__}: {e}"
        run.results.setdefault("_tracebacks", []).append(traceback.format_exc())


# ----------------------------------------------------------------------------
def stage_cad(run: Run):
    arm_mod = core.load("arm")
    p = {k: v for k, v in run.design.section("arm").items() if k not in ("count", "material")}
    part = arm_mod.arm(**p)
    run.part = part
    from hwe.cad import boundary
    material = catalog.materials()[run.design.get("arm.material")]
    volume = boundary.volume(part)
    run.results["arm.volume"] = volume
    run.results["arm.mass"] = boundary.mass(part, material["density"])
    by_hand = arm_mod.arm_volume_by_hand(**p)
    run.results["arm.volume_by_hand"] = by_hand.to("mm^3")
    rel = abs((volume - by_hand) / by_hand).to("dimensionless").m
    run.cross_checks.append(dict(name="arm volume: CAD versus hand geometry", a=volume, b=by_hand.to("mm^3"),
                                 rel=rel, tol=1e-6, ok=rel <= 1e-6, kind="independent"))


def battery_mass(run: Run, capacity: Quantity | None = None) -> Quantity:
    """The pack's mass follows its energy: m = E / specific energy, so a bigger capacity weighs more.

    Specific energy (W*h per kg) comes from a real pack in the catalog; holding it
    fixed while the capacity changes is the modelling assumption of the sweep.
    """
    en = core.load("energy")
    bat = run.design.section("battery")
    energy = en.battery_energy(capacity if capacity is not None else bat["capacity"], bat["cells"], bat["cell_voltage"])
    return (energy / bat["specific_energy"]).to("g")


def _budget_items(run: Run) -> list[Item]:
    items = []
    bat = run.design.section("battery")
    items.append(Item("battery", "power", battery_mass(run), bat.get("maturity", "catalog"),
                      bat.get("source", "")))
    payload = run.design.section("payload")
    if payload and payload["mass"].m > 0:
        items.append(Item("payload", "payload", payload["mass"], payload.get("maturity", "estimated"),
                          payload.get("source", "")))
    for spec in run.design.items:
        if spec.mass is not None:
            items.append(spec)
            continue
        cad_mass = run.results.get("arm.mass")          # mass = "cad": taken from the CAD stage
        if cad_mass is not None:
            items.append(Item(spec.name, spec.subsystem, cad_mass * spec.count, "cad", "CAD: hwe.core.arm"))
        elif spec.fallback is not None:
            items.append(Item(spec.name, spec.subsystem, spec.fallback, "estimated", "estimate (CAD not run)"))
        else:
            raise ValueError(f"{spec.name}: mass comes from CAD, which has not run, and has no fallback")
    return items


def stage_mass(run: Run):
    mb = core.load("massbudget")
    allowances = run.design.section("growth_allowance")
    items = _budget_items(run)
    run.budget = items
    run.results["mass.basic"] = mb.basic_mass(items).to("g")
    run.results["mass.predicted"] = mb.predicted_mass(items, allowances).to("g")
    run.by_subsystem = {k: v.to("g") for k, v in mb.mass_by_subsystem(items, allowances).items()}
    total = sum(run.by_subsystem.values(), Q_(0, "g"))
    rel = abs(((total - run.results["mass.predicted"]) / run.results["mass.predicted"]).to("").m)
    run.cross_checks.append(dict(name="mass by subsystem adds up to the predicted total",
                                 a=total, b=run.results["mass.predicted"], rel=rel, tol=1e-9, ok=rel <= 1e-9,
                                 kind="identity"))
    table = _thrust_table(run)
    min_ttw = run.design.get("requirements.min_thrust_to_weight")
    n = run.design.get("propulsion.motor_count")
    g = run.design.get("environment.g")
    allow = mb.allowable_mass(table.max_thrust(), n, min_ttw, g).to("g")
    run.results["mass.allowable"] = allow
    run.results["mass.margin"] = mb.mass_margin(run.results["mass.predicted"], allow)


def _thrust_table(run: Run) -> catalog.ThrustTable:
    prop = run.design.section("propulsion")
    return catalog.thrust_table(prop["thrust_table"], prop.get("source", ""))


def _check_table_voltage(run: Run, table: catalog.ThrustTable) -> None:
    """A thrust table is valid only at the voltage it was measured at; say so if the pack differs."""
    bat = run.design.section("battery")
    pack = (bat["cells"] * bat["cell_voltage"]).to("V")
    try:
        tested = table.table.find("voltage").mean().to("V")
    except KeyError:
        return
    rel = abs(((pack - tested) / tested).to("").m)
    run.cross_checks.append(dict(name="battery voltage matches the thrust table's test voltage",
                                 a=pack, b=tested, rel=rel, tol=0.10, ok=rel <= 0.10, kind="assumption"))


def stage_hover(run: Run):
    hv = core.load("hover")
    table = _thrust_table(run)
    _check_table_voltage(run, table)
    n = run.design.get("propulsion.motor_count")
    g = run.design.get("environment.g")
    m = run.results["mass.predicted"]
    t_max = table.max_thrust()
    run.results["hover.max_thrust_per_motor"] = t_max.to("gf")
    run.results["hover.weight"] = hv.weight(m, g).to("N")
    run.results["hover.thrust_per_motor"] = hv.hover_thrust_per_motor(m, n, g).to("gf")
    ttw = hv.thrust_to_weight(t_max, n, m, g)
    run.results["hover.thrust_to_weight"] = ttw
    frac = hv.hover_thrust_fraction(t_max, n, m, g)
    run.results["hover.hover_fraction"] = frac
    run.results["hover.max_climb_accel"] = hv.max_climb_acceleration(t_max, n, m, g).to("m/s^2")
    # Recomputed: the same formula and inputs as plain gram-force arithmetic, no core code. It catches
    # slips in your code, not a wrong model or a wrong table.
    by_hand = (n * t_max.to("gf").m) / (m.to("g").m * (g / Q_(9.80665, "m/s^2")).to("").m)
    rel = abs(ttw.m - by_hand) / by_hand
    run.cross_checks.append(dict(name="T/W: your code versus plain gram-force arithmetic",
                                 a=ttw, b=Q_(by_hand, ""), rel=rel, tol=1e-9, ok=rel <= 1e-9, kind="recomputed"))
    rel2 = abs(frac.m * ttw.m - 1.0)
    run.cross_checks.append(dict(name="hover fraction x T/W = 1", a=Q_(frac.m * ttw.m, ""), b=Q_(1.0, ""),
                                 rel=rel2, tol=1e-9, ok=rel2 <= 1e-9, kind="identity"))


def _hover_power(run: Run, mass: Quantity) -> Quantity:
    """Motors plus electronics, at hover, for a given total mass."""
    hv, en = core.load("hover"), core.load("energy")
    table = _thrust_table(run)
    n = run.design.get("propulsion.motor_count")
    g = run.design.get("environment.g")
    per_motor = hv.hover_thrust_per_motor(mass, n, g)
    eff = table.efficiency_at(per_motor)
    return en.propulsion_power(per_motor, eff, n) + run.design.get("electronics.power"), eff


def stage_energy(run: Run):
    en = core.load("energy")
    bat = run.design.section("battery")
    energy = en.battery_energy(bat["capacity"], bat["cells"], bat["cell_voltage"])
    run.results["energy.battery_energy"] = energy.to("W*h")
    power, eff = _hover_power(run, run.results["mass.predicted"])
    run.results["energy.hover_efficiency"] = eff
    run.results["energy.hover_power"] = power.to("W")
    t = en.flight_time(energy, power, bat["usable_fraction"])
    run.results["energy.flight_time"] = t.to("min")
    pack_voltage = (bat["cells"] * bat["cell_voltage"]).to("V")
    run.results["energy.hover_current"] = (power / pack_voltage).to("A")
    # Identity: charge over a current computed from the same power. I = P / V makes Q / I the same
    # formula as E / P, so it must agree exactly; it catches arithmetic slips only.
    q_usable = bat["capacity"] * bat["usable_fraction"]
    t2 = (q_usable / run.results["energy.hover_current"]).to("min")
    rel = abs(((t - t2) / t2).to("").m)
    run.cross_checks.append(dict(name="flight time: energy/power versus charge over P/V",
                                 a=t.to("min"), b=t2, rel=rel, tol=1e-9, ok=rel <= 1e-9, kind="identity"))
    # Independent: the thrust table's MEASURED current column at the hover thrust (the main result used
    # its efficiency column). The table was measured at its own test voltage, not the pack's, so expect
    # a gap about as large as the difference between those voltages.
    hv = core.load("hover")
    per_motor = hv.hover_thrust_per_motor(run.results["mass.predicted"], run.design.get("propulsion.motor_count"),
                                          run.design.get("environment.g"))
    table_current = (run.design.get("propulsion.motor_count") * _thrust_table(run).current_at(per_motor)
                     + run.design.get("electronics.power") / pack_voltage)
    run.results["energy.hover_current_table"] = table_current.to("A")
    t3 = (q_usable / table_current).to("min")
    rel3 = abs(((t - t3) / t3).to("").m)
    run.cross_checks.append(dict(name="flight time: energy/power versus the table's measured current",
                                 a=t.to("min"), b=t3, rel=rel3, tol=0.10, ok=rel3 <= 0.10, kind="independent"))


def stage_sweep(run: Run):
    """Flight time against battery size, holding its specific energy and cell count fixed."""
    if "energy.flight_time" not in run.results:
        return "skipped: needs the energy stage's flight time"
    en = core.load("energy")
    bat = run.design.section("battery")
    allowance = run.design.section("growth_allowance")[bat.get("maturity", "catalog")]
    specific = bat["specific_energy"].to("W*h/kg")
    others = run.results["mass.predicted"] - battery_mass(run) * (1 + allowance)
    hv = core.load("hover")
    table = _thrust_table(run)
    n = run.design.get("propulsion.motor_count")
    g = run.design.get("environment.g")
    min_ttw = run.design.get("requirements.min_thrust_to_weight")
    capacities, masses, times, ttws = [], [], [], []
    base = bat["capacity"].to("mA*h").m
    for k in range(2, 65):                       # 0.25x to 8x the design's capacity, in eighths
        cap = Q_(base * k / 8.0, "mA*h")
        energy = en.battery_energy(cap, bat["cells"], bat["cell_voltage"])
        mb = (energy / specific).to("g")
        total = others + mb * (1 + allowance)
        try:
            power, _ = _hover_power(run, total)
        except ValueError:
            continue                   # hover thrust outside the tested table: no data, no point
        capacities.append(cap)
        masses.append(mb)
        times.append(en.flight_time(energy, power, bat["usable_fraction"]).to("min"))
        ttws.append(hv.thrust_to_weight(table.max_thrust(), n, total, g))
    run.sweep = dict(capacity=capacities, battery_mass=masses, flight_time=times, thrust_to_weight=ttws,
                     min_thrust_to_weight=min_ttw, specific_energy=specific,
                     current_capacity=bat["capacity"], current_time=run.results["energy.flight_time"])
    ok = [i for i, r in enumerate(ttws) if r.m >= min_ttw]
    if ok:
        run.results["sweep.largest_capacity_meeting_ttw"] = capacities[ok[-1]]
    if times:
        best = max(range(len(times)), key=lambda i: times[i].m)
        run.results["sweep.best_capacity"] = capacities[best]
        run.results["sweep.best_flight_time"] = times[best]


def stage_requirements(run: Run):
    path = ROOT / run.design.raw.get("requirements", {}).get("file", "design/requirements.toml")
    records = reqs.load(path)
    problems = reqs.validate(records, run.results)
    run.results["_requirement_problems"] = problems
    run.requirement_status = reqs.check(records, run.results)
    # Assumption: the T/W minimum that sets the allowable mass (drone.toml) must be the one the
    # requirements promise. A number kept in two files drifts; this catches it.
    min_ttw = run.design.get("requirements.min_thrust_to_weight")
    for r in records:
        if r.result == "hover.thrust_to_weight" and r.op == ">=":
            promised = r.threshold.to("").m
            rel = abs(min_ttw - promised) / promised
            run.cross_checks.append(dict(name=f"allowable mass uses {r.id}'s T/W minimum",
                                         a=Q_(min_ttw, ""), b=Q_(promised, ""), rel=rel, tol=1e-9,
                                         ok=rel <= 1e-9, kind="assumption"))
    return "ok" if not problems else f"file problems: {len(problems)}"


def stage_phone(run: Run):
    from hwe.testdata import read_csv
    cfg = run.design.raw.get("tests", {}).get("phone", {})
    data_path = ROOT / cfg.get("data", "data/tests/el01_phone/lift.csv")
    if not data_path.exists():
        return "no data"
    km = core.load("kinematics")
    table = read_csv(data_path)
    t = table.find("time")
    a = table.find("acceleration", "z") * float(cfg.get("sign", 1))   # sign = -1 if your app reads -g at rest
    still = Q_(cfg.get("still_until", "1.0 s"))
    v, h = km.velocity_and_height(t, a, still)
    run.phone = dict(t=t, a=a, v=v, h=h, path=data_path, still_until=still)
    run.results["phone.final_height"] = h[-1].to("cm")
    run.results["phone.peak_accel_reading"] = a.max().to("m/s^2")
    run.results["phone.rest_reading"] = a[t < still].mean().to("m/s^2")


def run(overrides: dict[str, str] | None = None, design_path: Path | None = None) -> Run:
    design = load(design_path or ROOT / "design" / "drone.toml", overrides)
    r = Run(design=design, unlocked=unlocked_modules())
    for name, func in [("cad", stage_cad), ("mass", stage_mass), ("hover", stage_hover),
                       ("energy", stage_energy), ("sweep", stage_sweep),
                       ("requirements", stage_requirements), ("phone", stage_phone)]:
        _stage(r, name, func)
    for name, module, _ in STAGES:
        r.stages.setdefault(name, "locked" if module not in r.unlocked else "not built yet")
    return r