Hybrid orbit: analytic physics + a learned correction

Every trajectory is a planar orbit around a point mass. The textbook two-body formula misses drag, so a small tanh network is trained in torch to predict exactly the missing piece of the acceleration. hawk traces ONE kernel that performs a whole RK4 step with that network evaluated inline, compiles it, and eagle runs the compiled kernel over the whole batch.

This notebook runs the --host arm (CPU threads); the same kernel also runs as a replayed CUDA graph with --device on a machine with a GPU. Equivalent command line:

python examples/two_body_hybrid.py --host --steps 50 --batch 256

Needs eagle installed next to hawk, plus torch.

Colab or Kaggle: run this first cell (it installs the packages and downloads the helper files this notebook uses). Elsewhere it only checks the imports.

Show code cell source

Hide code cell source

# Colab setup, generated by docs/tools/nb_colab_setup.py (do not edit).
import os

try:
    import google.colab  # noqa: F401
    HOSTED = True  # a Colab runtime
except ImportError:
    HOSTED = "KAGGLE_KERNEL_RUN_TYPE" in os.environ  # a Kaggle kernel

if HOSTED:
    import importlib.metadata, pathlib, re, shutil, subprocess, sys, urllib.request

    gpu = bool(shutil.which("nvidia-smi")) and subprocess.run(["nvidia-smi", "-L"], capture_output=True).returncode == 0
    extra = ""
    if gpu:
        # the CUDA major the runtime already has: a CuPy that is installed, else torch's, else the driver's
        major = None
        for m in ("13", "12"):
            try:
                importlib.metadata.version(f"cupy-cuda{m}x")
                major = m
                break
            except importlib.metadata.PackageNotFoundError:
                pass
        if major is None:
            try:
                import torch
                major = "13" if str(torch.version.cuda).startswith("13") else None
            except ImportError:
                pass
        if major is None:
            smi = re.search(r"CUDA Version: (\d+)", subprocess.run(["nvidia-smi"], capture_output=True, text=True).stdout)
            major = "13" if smi and int(smi.group(1)) >= 13 else "12"
        extra = f"[cuda{major}]"
        print(f"GPU found: installing the {extra} extras to match this runtime.")
    pre = ["--pre"] if os.environ.get("RAPTOR_PRE") else []  # optional: test a release candidate
    ref = os.environ.get("RAPTOR_REF", "main")  # optional: fetch the helper files from this git ref
    if pre or "RAPTOR_REF" in os.environ:
        print(f"Pre-release test: installing with --pre, helper files from {ref}.")
    # our packages with --upgrade, so a runtime that kept an older release gets the current one
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--upgrade", *pre, "raptor-core", f"raptor-eagle{extra}", f"raptor-hawk{extra}"])
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "torch"])
    base = os.environ.get("RAPTOR_RAW_BASE", "https://raw.githubusercontent.com")  # optional: a local mirror
    top = os.environ.get("RAPTOR_WORKDIR") or ("/kaggle/working" if "KAGGLE_KERNEL_RUN_TYPE" in os.environ else "/content")
    root = pathlib.Path(top) / "hawk"
    for rel in (
        "examples/two_body_hybrid.py",
    ):
        (root / rel).parent.mkdir(parents=True, exist_ok=True)
        urllib.request.urlretrieve(f"{base}/amasat01/hawk/{ref}/{rel}", root / rel)
    (root / "docs/content/examples").mkdir(parents=True, exist_ok=True)
    os.chdir(root / "docs/content/examples")
    if not gpu:
        print("Running on the CPU (no GPU in this runtime). For a GPU: " + ("Settings → Accelerator → GPU T4 x2." if "KAGGLE_KERNEL_RUN_TYPE" in os.environ else "Runtime → Change runtime type → T4 GPU."))

try:
    import eagle, hawk, raptor  # noqa: F401
except ImportError as err:
    raise ImportError(
        "eagle, hawk, raptor not found. Install with `pip install raptor-core raptor-eagle raptor-hawk torch`. For a GPU, add [cuda12] (or [cuda13]) to raptor-eagle and raptor-hawk."
    ) from err

Show code cell source

Hide code cell source

import sys, pathlib, warnings

# This build's torch pulls in a deprecated pynvml that warns on the
# first torch.cuda touch; the warning's own traceback cites this
# machine's install path, which the docs' path-leak gate refuses --
# silence this one, known-irrelevant warning rather than the path leak.
warnings.filterwarnings(
    "ignore", message="The pynvml package is deprecated", category=FutureWarning
)

# These examples live in the repository's own examples/ directory, not in
# the installed package — add it to sys.path so the module below
# imports straight from the checkout.
sys.path.insert(0, str(pathlib.Path("../../../examples").resolve()))
import two_body_hybrid as demo
print(demo.__name__, "from examples/" + pathlib.Path(demo.__file__).name)

two_body_hybrid from examples/two_body_hybrid.py

The kernel is traced once, as ordinary Python — the MLP correction is evaluated inline, four times per RK4 step, with its weights read from shared Tables:

import inspect
print(inspect.getsource(demo.step_kernel))

def step_kernel(hidden: int = HIDDEN):
    """Trace the RK4 step. The MLP is evaluated inline, four times per step.

    The weights are four ``Table`` arguments: one copy shared by every
    trajectory, not one per sample. The kernel stays a pure function of
    ``(state, weights, dt)`` and carries no state of its own.
    """
    import hawk
    import hawk.math as hm

    def hybrid_rhs(s, w1, b1, w2, b2):
        """Gravity plus the network's correction, as one traced expression."""
        r2 = s[0] * s[0] + s[1] * s[1]
        inv_r3 = hm.rsqrt(r2 * r2 * r2)
        act = []
        for h in range(hidden):
            pre = b1.at(h=h)
            for i in range(4):
                pre = pre + w1.at(h=h, i=i) * s[i]
            act.append(hm.tanh(pre))
        corr = []
        for o in range(2):
            out = b2.at(o=o)
            for h in range(hidden):
                out = out + w2.at(o=o, h=h) * act[h]
            corr.append(out)
        return hm.vec(
            s[2],
            s[3],
            -MU * s[0] * inv_r3 + corr[0],
            -MU * s[1] * inv_r3 + corr[1],
        )

    @hawk.kernel
    def two_body_hybrid_step(
        x: hawk.Vector[4],
        w1: hawk.Table["h":hidden, "i":4],
        b1: hawk.Table["h":hidden],
        w2: hawk.Table["o":2, "h":hidden],
        b2: hawk.Table["o":2],
        dt: hawk.Param,
        x_next: hawk.Mutable[hawk.Vector[4]],
    ):
        k1 = hybrid_rhs(x, w1, b1, w2, b2)
        k2 = hybrid_rhs(x + (0.5 * dt) * k1, w1, b1, w2, b2)
        k3 = hybrid_rhs(x + (0.5 * dt) * k2, w1, b1, w2, b2)
        k4 = hybrid_rhs(x + dt * k3, w1, b1, w2, b2)
        x_next = x + (dt / 6.0) * (k1 + 2.0 * k2 + 2.0 * k3 + k4)

    return two_body_hybrid_step

By hand: train the correction, then let hawk publish the compiled kernel into a manifest raptor validates.

import pathlib, tempfile
import numpy as np
import hawk.artifact
from raptor.schema import validate_manifest
import eagle.exec as eexec
from eagle import plan as eplan

state = demo.seed_states(batch=256, seed=0)
weights = demo.train_correction(seed=0)

bundle = hawk.build(
    [demo.step_kernel()], pathlib.Path(tempfile.mkdtemp()),
    mode="float32", targets=("host",))
validate_manifest(bundle.manifest)
print("exec_targets:", bundle.manifest["exec_targets"])

exec_targets: ['host']

Then hand it to eagle and run it — this is the one call that actually launches the kernel. The weights are four Tables: one buffer that every trajectory reads, however large the batch.

plugin = next(iter(hawk.artifact.plugins(bundle).values()))
plan = eplan.plan(plugin, structure=eexec.HostTeam)
tables = demo.weight_tables(weights, np)   # one shared copy, not one per trajectory
x_next = np.zeros_like(state, dtype=np.float32)
bound = plan.bind(x=state.astype(np.float32), x_next=x_next, dt=0.01, **tables)
bound.launch()
print("x_next[:, 0]:", x_next[:, 0])

x_next[:, 0]: [-0.64382726 -0.7641566   0.77130795 -0.65005666]

The pieces above, assembled into one call — training, building, running and checking it, for reference:

results = demo.run_demo(steps=50, batch=256, target="host")

target            : host (manifest exec_targets ['host'])
batch x steps     : 256 x 50 at dt = 0.01
training loss     : 1.230e-06
(a) max |compiled - torch eager| : 2.384e-07 (tolerance 1e-04)
(b) final-state error vs truth   : corrected 6.695e-04 vs two-body only 2.405e-02
(c) wall per step (demo measurement, not a benchmark): compiled 0.741 ms, torch eager 0.621 ms

match is acceptance (a): the compiled rollout tracks a torch-eager rollout of the same weights to a tight tolerance. hybrid_error vs. two_body_error is acceptance (b): the learned correction should land closer to the true (gravity + drag) trajectory than gravity alone.

Line (c) runs on CPU threads at a small batch, where torch eager is faster: the step is a small dense network, which torch runs as matrix products, and each compiled step pays a fixed start-up cost. The compiled step pays off on a GPU (target="device") and at larger batches. The wall times are this run’s own measurement, not a committed benchmark.