estimator

estimator

Aircraft fuel-flow estimation from trajectory data via an ONNX model.

Public API

FuelEstimator — DataFrame-in / DataFrame-out estimator. Any frame narwhals supports (pandas, polars, pyarrow, …) is accepted and the input type is preserved on output. Dispatches per aircraft typecode internally. AircraftFuelEstimator — typecode-bound, numpy-only estimator (faster per call, no per-call params lookup).

AircraftFuelEstimator

Typecode-bound fuel-flow estimator with numpy-only I/O.

Faster than :meth:FuelEstimator.estimate per call because the typecode params lookup, normalization arrays and engine-scale factor are precomputed at construction. Input is 1-D numpy arrays; output is a 1-D numpy array of fuel flow in kg/s (multiply by 3600 for kg/h).

Source code in src/acropole/estimator.py

class AircraftFuelEstimator:
    """Typecode-bound fuel-flow estimator with numpy-only I/O.

    Faster than :meth:`FuelEstimator.estimate` per call because the typecode
    params lookup, normalization arrays and engine-scale factor are precomputed
    at construction. Input is 1-D numpy arrays; output is a 1-D numpy array of
    fuel flow in **kg/s** (multiply by 3600 for kg/h).
    """

    DEFAULT_MASS = _DEFAULT_MASS

    def __init__(
        self,
        typecode: str,
        aircraft_params_path: Annotated[
            str | None, "CSV path; None -> package data"
        ] = None,
        model_path: Annotated[str | None, "ONNX path; None -> package data"] = None,
        dtype: Annotated[npt.DTypeLike, "intermediate numpy precision"] = np.float64,
    ) -> None:
        if aircraft_params_path is None:
            aircraft_params_path = str(
                files("acropole").joinpath("data/aircraft_params.csv")
            )
        row = _load_aircraft_params(aircraft_params_path).get(typecode)
        if row is None:
            raise ValueError(f"Aircraft type {typecode!r} not in aircraft_params")

        if model_path is None:
            model_path = str(
                files("acropole").joinpath("models/acropole_fuel_model.onnx")
            )
        session = ort.InferenceSession(
            str(model_path), providers=["CPUExecutionProvider"]
        )

        self._init_from(typecode, row, session, cast("FloatDType", np.dtype(dtype)))

    @classmethod
    def _from_shared(
        cls, estimator: FuelEstimator, typecode: str
    ) -> AircraftFuelEstimator:
        """Build bound to ``typecode``, reusing ``estimator``'s session/params."""
        params = estimator._params_by_type.get(typecode)
        if params is None:
            raise ValueError(f"Aircraft type {typecode!r} not in aircraft_params")
        obj = cls.__new__(cls)
        obj._init_from(typecode, params, estimator.session, estimator.dtype)
        return obj

    def _init_from(  # type: ignore[no-any-unimported]  # ort.InferenceSession: no stubs
        self,
        typecode: str,
        params: dict[str, float],
        session: ort.InferenceSession,
        dtype: FloatDType,
    ) -> None:
        self.typecode = typecode
        self.dtype = dtype
        self.session = session
        self._input_name = session.get_inputs()[0].name
        self._output_name = session.get_outputs()[0].name

        to_dtype = dtype.type
        self._engine_type = to_dtype(params["ENGINE_TYPE"])
        self._surface = to_dtype(params["SURFACE"])
        self._max_ope_alti = to_dtype(params["MAX_OPE_ALTI"])
        self._max_ope_speed = to_dtype(params["MAX_OPE_SPEED"])
        self._ope_empty_weight = to_dtype(params["OPE_EMPTY_WEIGHT"])
        self._mass_range = to_dtype(
            params["MAX_TO_WEIGHT"] - params["OPE_EMPTY_WEIGHT"]
        )
        self._fuel_scale = float(params["FUEL_FLOW_TO"] * params["ENGINE_NUM"])

        self._mins = np.array(_MINIMUMS, dtype=dtype)
        self._scale = np.array(_MAXIMUMS, dtype=dtype) - self._mins

    def estimate(
        self,
        groundspeed: FloatArray,
        altitude: FloatArray,
        vertical_rate: FloatArray,
        *,
        airspeed: FloatArray | None = None,
        mass: FloatArray | None = None,
        second: FloatArray | None = None,
        d_altitude: FloatArray | None = None,
        d_groundspeed: FloatArray | None = None,
        d_airspeed: FloatArray | None = None,
    ) -> FloatArray:
        """Predict per-sample fuel flow in **kg/s** (1-D array of length N)."""
        dtype = self.dtype
        gs = np.asarray(groundspeed, dtype=dtype)
        alt = np.asarray(altitude, dtype=dtype)
        vr = np.asarray(vertical_rate, dtype=dtype)
        n = gs.shape[0]
        air = np.asarray(airspeed, dtype=dtype) if airspeed is not None else gs

        mass_norm = self._mass_norm(mass, n, dtype)
        d_alt, d_gs, d_as = self._derivatives(
            n, dtype, alt, gs, air, vr, second, d_altitude, d_groundspeed, d_airspeed
        )

        inputs = np.empty((n, 12), dtype=dtype)
        inputs[:, 0] = self._engine_type
        inputs[:, 1] = d_alt
        inputs[:, 2] = d_gs
        inputs[:, 3] = d_as
        inputs[:, 4] = self._surface
        inputs[:, 5] = self._max_ope_alti
        inputs[:, 6] = self._max_ope_speed
        inputs[:, 7] = alt
        inputs[:, 8] = gs
        inputs[:, 9] = air
        inputs[:, 10] = vr
        inputs[:, 11] = mass_norm

        normalized = (inputs - self._mins) / self._scale
        if normalized.dtype != np.float32:
            normalized = normalized.astype(np.float32)

        values = self.session.run([self._output_name], {self._input_name: normalized})[
            0
        ]
        single = (
            values.squeeze(axis=-1)
            if values.ndim == _MODEL_OUTPUT_NDIM
            else values.ravel()
        )
        return np.asarray(single * self._fuel_scale)

    def _mass_norm(
        self, mass: FloatArray | None, n: int, dtype: FloatDType
    ) -> FloatArray:
        if mass is None:
            return np.full(n, self.DEFAULT_MASS, dtype=dtype)
        if self._mass_range == 0:
            warnings.warn(
                f"Aircraft {self.typecode!r} has MAX_TO_WEIGHT == OPE_EMPTY_WEIGHT; "
                "mass normalization is undefined (NaN).",
                stacklevel=2,
            )
            return np.full(n, np.nan, dtype=dtype)
        out = (
            np.asarray(mass, dtype=dtype) - self._ope_empty_weight
        ) / self._mass_range
        return np.asarray(out)

    def _derivatives(
        self,
        n: int,
        dtype: FloatDType,
        alt: FloatArray,
        gs: FloatArray,
        air: FloatArray,
        vr: FloatArray,
        second: FloatArray | None,
        d_altitude: FloatArray | None,
        d_groundspeed: FloatArray | None,
        d_airspeed: FloatArray | None,
    ) -> tuple[FloatArray, FloatArray, FloatArray]:
        if second is None:
            d_alt = (
                np.asarray(d_altitude, dtype=dtype)
                if d_altitude is not None
                else vr / dtype.type(60.0)
            )
            d_gs = (
                np.asarray(d_groundspeed, dtype=dtype)
                if d_groundspeed is not None
                else np.zeros(n, dtype=dtype)
            )
            d_as = (
                np.asarray(d_airspeed, dtype=dtype)
                if d_airspeed is not None
                else np.zeros(n, dtype=dtype)
            )
            return d_alt, d_gs, d_as

        dt = diff_bfill(np.asarray(second, dtype=dtype))
        d_alt = (
            np.asarray(d_altitude, dtype=dtype)
            if d_altitude is not None
            else safe_divide(diff_bfill(alt), dt)
        )
        d_gs = (
            np.asarray(d_groundspeed, dtype=dtype)
            if d_groundspeed is not None
            else safe_divide(diff_bfill(gs), dt)
        )
        d_as = (
            np.asarray(d_airspeed, dtype=dtype)
            if d_airspeed is not None
            else safe_divide(diff_bfill(air), dt)
        )
        return d_alt, d_gs, d_as

estimate(groundspeed, altitude, vertical_rate, *, airspeed=None, mass=None, second=None, d_altitude=None, d_groundspeed=None, d_airspeed=None)

Predict per-sample fuel flow in kg/s (1-D array of length N).

Source code in src/acropole/estimator.py

def estimate(
    self,
    groundspeed: FloatArray,
    altitude: FloatArray,
    vertical_rate: FloatArray,
    *,
    airspeed: FloatArray | None = None,
    mass: FloatArray | None = None,
    second: FloatArray | None = None,
    d_altitude: FloatArray | None = None,
    d_groundspeed: FloatArray | None = None,
    d_airspeed: FloatArray | None = None,
) -> FloatArray:
    """Predict per-sample fuel flow in **kg/s** (1-D array of length N)."""
    dtype = self.dtype
    gs = np.asarray(groundspeed, dtype=dtype)
    alt = np.asarray(altitude, dtype=dtype)
    vr = np.asarray(vertical_rate, dtype=dtype)
    n = gs.shape[0]
    air = np.asarray(airspeed, dtype=dtype) if airspeed is not None else gs

    mass_norm = self._mass_norm(mass, n, dtype)
    d_alt, d_gs, d_as = self._derivatives(
        n, dtype, alt, gs, air, vr, second, d_altitude, d_groundspeed, d_airspeed
    )

    inputs = np.empty((n, 12), dtype=dtype)
    inputs[:, 0] = self._engine_type
    inputs[:, 1] = d_alt
    inputs[:, 2] = d_gs
    inputs[:, 3] = d_as
    inputs[:, 4] = self._surface
    inputs[:, 5] = self._max_ope_alti
    inputs[:, 6] = self._max_ope_speed
    inputs[:, 7] = alt
    inputs[:, 8] = gs
    inputs[:, 9] = air
    inputs[:, 10] = vr
    inputs[:, 11] = mass_norm

    normalized = (inputs - self._mins) / self._scale
    if normalized.dtype != np.float32:
        normalized = normalized.astype(np.float32)

    values = self.session.run([self._output_name], {self._input_name: normalized})[
        0
    ]
    single = (
        values.squeeze(axis=-1)
        if values.ndim == _MODEL_OUTPUT_NDIM
        else values.ravel()
    )
    return np.asarray(single * self._fuel_scale)

FuelEstimator

Data pipeline for trajectory fuel-flow enhancement.

Accepts any eager DataFrame narwhals supports (pandas, polars, pyarrow, …) and returns the same type, adding fuel_flow (kg/s), fuel_flow_kgh (kg/h) and — when a second column is present — fuel_cumsum (kg). The second column (like the other optional features) is triggered by its presence in the frame, no keyword argument required. Frames mixing several aircraft typecodes are handled per typecode (each row scored with its own aircraft parameters).

Source code in src/acropole/estimator.py

class FuelEstimator:
    """Data pipeline for trajectory fuel-flow enhancement.

    Accepts any eager DataFrame narwhals supports (pandas, polars, pyarrow, …)
    and returns the same type, adding
    ``fuel_flow`` (kg/s), ``fuel_flow_kgh`` (kg/h) and — when a ``second``
    column is present — ``fuel_cumsum`` (kg). The ``second`` column (like the
    other optional features) is triggered by its presence in the frame, no
    keyword argument required. Frames mixing several aircraft typecodes are
    handled per typecode (each row scored with its own aircraft parameters).
    """

    DEFAULT_MASS = _DEFAULT_MASS

    def __init__(
        self,
        aircraft_params_path: Annotated[
            str | None, "CSV path; None -> package data"
        ] = None,
        model_path: Annotated[str | None, "ONNX path; None -> package data"] = None,
        dtype: Annotated[npt.DTypeLike, "intermediate numpy precision"] = np.float64,
    ) -> None:
        if aircraft_params_path is None:
            aircraft_params_path = str(
                files("acropole").joinpath("data/aircraft_params.csv")
            )
        self._params_by_type = _load_aircraft_params(aircraft_params_path)

        if model_path is None:
            model_path = str(
                files("acropole").joinpath("models/acropole_fuel_model.onnx")
            )
        self.session = ort.InferenceSession(
            str(model_path), providers=["CPUExecutionProvider"]
        )
        self.dtype: FloatDType = cast("FloatDType", np.dtype(dtype))

    def for_aircraft(self, typecode: str) -> AircraftFuelEstimator:
        """Return an :class:`AircraftFuelEstimator` bound to ``typecode``, reusing
        this estimator's already-loaded ONNX session and parameters (no reload)."""
        return AircraftFuelEstimator._from_shared(self, typecode)

    def estimate(self, flight: IntoDataFrameT, **kwargs: str) -> IntoDataFrameT:
        """Estimate fuel flow for ``flight``; see class docstring for columns.

        Optional features are triggered by the **presence of their column** in
        ``flight`` (matched by name): when a ``second`` column is present the
        real time-derivatives are used and a ``fuel_cumsum`` column is added; an
        absent ``second`` column falls back to quasi-steady derivatives and adds
        no ``fuel_cumsum`` (no kwarg required either way).

        Column-name overrides via kwargs map a non-standard column name onto a
        feature: ``typecode``, ``groundspeed``, ``altitude``, ``vertical_rate``,
        ``airspeed``, ``mass``, ``second``, ``d_altitude``, ``d_groundspeed``,
        ``d_airspeed``.

        Missing or unsupported typecodes warn and produce NaN fuel flow. With
        ``second``, derivatives are computed within each typecode group; a
        single-sample group has undefined derivatives unless supplied explicitly.
        ``fuel_cumsum`` integrates the entire frame in its existing row order.
        An unknown fuel contribution makes that and all subsequent totals NaN;
        missing consumption is never silently treated as zero.
        """
        df = nw.from_native(flight, eager_only=True)

        col: dict[str, str] = {
            name: kwargs.get(name, name)
            for name in (
                "typecode",
                "groundspeed",
                "altitude",
                "vertical_rate",
                "airspeed",
                "mass",
                "second",
                "d_altitude",
                "d_groundspeed",
                "d_airspeed",
            )
        }

        for required in ("typecode", "groundspeed", "altitude", "vertical_rate"):
            if col[required] not in df.columns:
                raise ValueError(f"Column {col[required]!r} not found")
        second_col = col["second"]
        if second_col in df.columns and not df[second_col].dtype.is_numeric():
            raise ValueError("column for second must be float or integer")

        fuel_flow = self._predict_grouped(df, col)

        backend = df.implementation
        out = df.with_columns(
            nw.new_series("fuel_flow", fuel_flow, backend=backend),
            nw.new_series("fuel_flow_kgh", fuel_flow * 3600.0, backend=backend),
        )
        if second_col in df.columns:
            sec = df[second_col].to_numpy().astype(self.dtype)
            out = out.with_columns(
                nw.new_series(
                    "fuel_cumsum",
                    np.cumsum(fuel_flow * diff_bfill(sec)),
                    backend=backend,
                )
            )
        return out.to_native()

    def _predict_grouped(
        self, df: nw.DataFrame[IntoDataFrameT], col: dict[str, str]
    ) -> FloatArray:
        """Run inference per typecode group, scattering results back to row order.

        Feature columns are pulled out of the frame once, up front; the per-group
        split is then a numpy mask. That keeps the backend-specific work to a
        single conversion per column instead of one frame filter per typecode.
        """
        codes = df[col["typecode"]]
        # Nullable pandas strings expose pd.NA, whose equality is not boolean.
        typecodes = codes.to_numpy().astype(object)
        typecodes[codes.is_null().to_numpy()] = None
        required, optional = self._extract(df, col)

        result = np.full(len(df), np.nan, dtype=self.dtype)
        for typecode in dict.fromkeys(typecodes.tolist()):  # unique, order-preserving
            if typecode not in self._params_by_type:
                warnings.warn(f"Aircraft type {typecode!r} not supported", stacklevel=3)
                continue  # leave NaN for unsupported rows
            mask = typecodes == typecode
            groundspeed, altitude, vertical_rate = (arr[mask] for arr in required)
            result[mask] = self.for_aircraft(typecode).estimate(
                groundspeed,
                altitude,
                vertical_rate,
                **{k: None if v is None else v[mask] for k, v in optional.items()},
            )
        return result

    def _extract(
        self, df: nw.DataFrame[IntoDataFrameT], col: dict[str, str]
    ) -> tuple[tuple[FloatArray, FloatArray, FloatArray], dict[str, FloatArray | None]]:
        """Split the frame into its required feature arrays and optional kwargs.

        The three required columns are guaranteed present by ``estimate``'s
        upfront validation, so they are returned non-optional; the rest become
        ``estimate``'s keyword-only arguments (``None`` when their column is
        absent).
        """
        dtype = self.dtype

        def arr(key: str) -> FloatArray | None:
            """Numpy array for column ``col[key]`` if present, else None."""
            name = col[key]
            if name not in df.columns:
                return None
            return df[name].to_numpy().astype(dtype)

        def required(key: str) -> FloatArray:
            """Array for a column ``estimate`` already validated as present."""
            return df[col[key]].to_numpy().astype(dtype)

        features = (
            required("groundspeed"),
            required("altitude"),
            required("vertical_rate"),
        )
        optional = {
            "airspeed": arr("airspeed"),
            "mass": arr("mass"),
            "second": arr("second"),
            "d_altitude": arr("d_altitude"),
            "d_groundspeed": arr("d_groundspeed"),
            "d_airspeed": arr("d_airspeed"),
        }
        return features, optional

estimate(flight, **kwargs)

Estimate fuel flow for flight; see class docstring for columns.

Optional features are triggered by the presence of their column in flight (matched by name): when a second column is present the real time-derivatives are used and a fuel_cumsum column is added; an absent second column falls back to quasi-steady derivatives and adds no fuel_cumsum (no kwarg required either way).

Column-name overrides via kwargs map a non-standard column name onto a feature: typecode, groundspeed, altitude, vertical_rate, airspeed, mass, second, d_altitude, d_groundspeed, d_airspeed.

Missing or unsupported typecodes warn and produce NaN fuel flow. With second, derivatives are computed within each typecode group; a single-sample group has undefined derivatives unless supplied explicitly. fuel_cumsum integrates the entire frame in its existing row order. An unknown fuel contribution makes that and all subsequent totals NaN; missing consumption is never silently treated as zero.

Source code in src/acropole/estimator.py

def estimate(self, flight: IntoDataFrameT, **kwargs: str) -> IntoDataFrameT:
    """Estimate fuel flow for ``flight``; see class docstring for columns.

    Optional features are triggered by the **presence of their column** in
    ``flight`` (matched by name): when a ``second`` column is present the
    real time-derivatives are used and a ``fuel_cumsum`` column is added; an
    absent ``second`` column falls back to quasi-steady derivatives and adds
    no ``fuel_cumsum`` (no kwarg required either way).

    Column-name overrides via kwargs map a non-standard column name onto a
    feature: ``typecode``, ``groundspeed``, ``altitude``, ``vertical_rate``,
    ``airspeed``, ``mass``, ``second``, ``d_altitude``, ``d_groundspeed``,
    ``d_airspeed``.

    Missing or unsupported typecodes warn and produce NaN fuel flow. With
    ``second``, derivatives are computed within each typecode group; a
    single-sample group has undefined derivatives unless supplied explicitly.
    ``fuel_cumsum`` integrates the entire frame in its existing row order.
    An unknown fuel contribution makes that and all subsequent totals NaN;
    missing consumption is never silently treated as zero.
    """
    df = nw.from_native(flight, eager_only=True)

    col: dict[str, str] = {
        name: kwargs.get(name, name)
        for name in (
            "typecode",
            "groundspeed",
            "altitude",
            "vertical_rate",
            "airspeed",
            "mass",
            "second",
            "d_altitude",
            "d_groundspeed",
            "d_airspeed",
        )
    }

    for required in ("typecode", "groundspeed", "altitude", "vertical_rate"):
        if col[required] not in df.columns:
            raise ValueError(f"Column {col[required]!r} not found")
    second_col = col["second"]
    if second_col in df.columns and not df[second_col].dtype.is_numeric():
        raise ValueError("column for second must be float or integer")

    fuel_flow = self._predict_grouped(df, col)

    backend = df.implementation
    out = df.with_columns(
        nw.new_series("fuel_flow", fuel_flow, backend=backend),
        nw.new_series("fuel_flow_kgh", fuel_flow * 3600.0, backend=backend),
    )
    if second_col in df.columns:
        sec = df[second_col].to_numpy().astype(self.dtype)
        out = out.with_columns(
            nw.new_series(
                "fuel_cumsum",
                np.cumsum(fuel_flow * diff_bfill(sec)),
                backend=backend,
            )
        )
    return out.to_native()

for_aircraft(typecode)

Return an :class:AircraftFuelEstimator bound to typecode, reusing this estimator's already-loaded ONNX session and parameters (no reload).

Source code in src/acropole/estimator.py

def for_aircraft(self, typecode: str) -> AircraftFuelEstimator:
    """Return an :class:`AircraftFuelEstimator` bound to ``typecode``, reusing
    this estimator's already-loaded ONNX session and parameters (no reload)."""
    return AircraftFuelEstimator._from_shared(self, typecode)