SPB Git

spb/wp40_uqo

Public MIT
7.8 KB · 103 lines cpp
1// nanobind bindings for RAID. Arrays are passed as NumPy views (no copy for contiguous float64/int64), the GIL is2// released around every heavy call, and results come back as dicts (JSON-like) plus NumPy arrays for vectors.3#include <nanobind/nanobind.h>4#include <nanobind/ndarray.h>5#include <nanobind/stl/string.h>6#include <nanobind/stl/vector.h>7#include <nanobind/stl/optional.h>8#include <nlohmann/json.hpp>9#include "raid/core/types.hpp"10#include "raid/core/error.hpp"11#include "raid/backend/backend.hpp"12#include "raid/estimators/ols.hpp"13#include "raid/estimators/realized.hpp"14#include "raid/models/garch.hpp"15#include "raid/data/bar_series.hpp"16#include "raid/data/calendar.hpp"17#include "raid/data/returns.hpp"1819namespace nb = nanobind;20using namespace raid;21using json = nlohmann::json;22using F64 = nb::ndarray<const double, nb::c_contig, nb::device::cpu>;23using I64 = nb::ndarray<const int64_t, nb::c_contig, nb::device::cpu>;2425namespace {26nb::object json_to_py(const json& j) {27  switch (j.type()) {28    case json::value_t::null: return nb::none();29    case json::value_t::boolean: return nb::bool_(j.get<bool>());30    case json::value_t::number_integer: return nb::int_(j.get<long long>());31    case json::value_t::number_unsigned: return nb::int_(j.get<unsigned long long>());32    case json::value_t::number_float: return nb::float_(j.get<double>());33    case json::value_t::string: return nb::str(j.get<std::string>().c_str());34    case json::value_t::array: { nb::list l; for (auto& v : j) l.append(json_to_py(v)); return l; }35    case json::value_t::object: { nb::dict d; for (auto& [k, v] : j.items()) d[k.c_str()] = json_to_py(v); return d; }36    default: return nb::none();37  }38}39nb::ndarray<nb::numpy, double> to_np(std::vector<double> v) {40  auto* data = new std::vector<double>(std::move(v));41  nb::capsule owner(data, [](void* p) noexcept { delete static_cast<std::vector<double>*>(p); });42  return nb::ndarray<nb::numpy, double>(data->data(), {data->size()}, owner);43}44nb::ndarray<nb::numpy, int64_t> to_np_i64(std::vector<int64_t> v) {45  auto* data = new std::vector<int64_t>(std::move(v));46  nb::capsule owner(data, [](void* p) noexcept { delete static_cast<std::vector<int64_t>*>(p); });47  return nb::ndarray<nb::numpy, int64_t>(data->data(), {data->size()}, owner);48}49}  // namespace5051NB_MODULE(_core, m) {52  m.doc() = "RAID compiled core (C++20, CPU + optional Metal)";53  m.attr("__version__") = RAID_VERSION;54  nb::exception<RaidError>(m, "RaidError");5556  m.def("devices", []() { json out = json::array(); for (auto& d : backend::list_devices()) out.push_back(d.to_json()); return json_to_py(out); });5758  m.def("ols", [](F64 X, F64 y, const std::string& cov, int hac_lags, std::optional<I64> clusters, double ridge, std::vector<usize> unpen, std::optional<F64> weights, std::vector<std::string> names) {59    RAID_REQUIRE(X.ndim() == 2 && y.ndim() == 1 && X.shape(0) == y.shape(0), invalid_argument, "X (n,k), y (n,)");60    const Eigen::Index n = static_cast<Eigen::Index>(X.shape(0)), k = static_cast<Eigen::Index>(X.shape(1));61    Eigen::Map<const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> Xm(X.data(), n, k);62    Eigen::Map<const Eigen::VectorXd> ym(y.data(), n);63    est::OlsOptions o; o.cov.type = stats::covtype_from_string(cov); o.cov.hac_lags = hac_lags; o.ridge_lambda = ridge; o.unpenalized = unpen; o.names = names;64    std::vector<i64> cl; if (clusters) { cl.assign(clusters->data(), clusters->data() + clusters->shape(0)); o.cov.clusters = cl; }65    std::vector<f64> w; if (weights) { w.assign(weights->data(), weights->data() + weights->shape(0)); o.weights = w; }66    Eigen::MatrixXd Xd = Xm; json out; { nb::gil_scoped_release rel; auto r = est::ols(Xd, ym, o); out = r.estimate.to_json(); }67    return json_to_py(out);68  }, nb::arg("X"), nb::arg("y"), nb::arg("cov") = "classical", nb::arg("hac_lags") = -1, nb::arg("clusters") = nb::none(), nb::arg("ridge_lambda") = 0.0, nb::arg("unpenalized") = std::vector<usize>{}, nb::arg("weights") = nb::none(), nb::arg("names") = std::vector<std::string>{});6970  m.def("realized_measures", [](F64 r, I64 off, const std::string& backend_s, bool bpv_fs, bool deterministic) {71    est::RealizedOptions o; o.plan.backend = backend_from_string(backend_s); o.bpv_finite_sample = bpv_fs; o.plan.deterministic = deterministic;72    backend::Batch b{std::span<const f64>(r.data(), r.shape(0)), std::span<const i64>(off.data(), off.shape(0))};73    est::RealizedMeasures mres; { nb::gil_scoped_release rel; mres = est::realized_measures_batch(b, o); }74    nb::dict d; d["rv"] = to_np(mres.rv); d["bpv"] = to_np(mres.bpv); d["rsv_neg"] = to_np(mres.rsv_neg); d["rsv_pos"] = to_np(mres.rsv_pos); d["rq"] = to_np(mres.rq); d["jump"] = to_np(mres.jump); d["tripv"] = to_np(mres.tripv); d["ret"] = to_np(mres.daily_return);75    d["backend"] = std::string(to_string(mres.stats.backend)); d["precision"] = std::string(to_string(mres.stats.precision)); d["wall_seconds"] = mres.stats.wall_seconds; d["device_seconds"] = mres.stats.device_seconds; d["note"] = mres.stats.note; d["warnings"] = mres.warnings;76    return d;77  }, nb::arg("returns"), nb::arg("offsets"), nb::arg("backend") = "auto", nb::arg("bpv_finite_sample") = true, nb::arg("deterministic") = true);7879  m.def("garch_fit", [](F64 r, const std::string& model, const std::string& dist, int horizon, const std::string& cov, int max_iter, double tol) {80    models::GarchFitOptions o; o.model = model == "gjr" ? models::GarchModel::gjr : (model == "egarch" ? models::GarchModel::egarch : models::GarchModel::garch);81    o.dist = dist == "student_t" || dist == "t" ? models::GarchDist::student_t : models::GarchDist::gaussian; o.forecast_horizon = horizon; o.cov = cov; o.max_iter = max_iter; o.tol = tol;82    std::span<const f64> rs(r.data(), r.shape(0)); json out; { nb::gil_scoped_release rel; out = models::garch_fit(rs, o).to_json(); }83    return json_to_py(out);84  }, nb::arg("returns"), nb::arg("model") = "garch", nb::arg("dist") = "gaussian", nb::arg("horizon") = 0, nb::arg("cov") = "qmle", nb::arg("max_iter") = 500, nb::arg("tol") = 1e-8);8586  m.def("validate_bars", [](I64 t, F64 o, F64 h, F64 l, F64 c, F64 v, const std::string& symbol, const std::string& interval, const std::string& session) {87    BarSeries b; b.symbol = symbol; b.interval = interval_from_string(interval); const usize n = t.shape(0); b.reserve(n);88    for (usize i = 0; i < n; ++i) b.push_back(Timestamp(t.data()[i]), o.data()[i], h.data()[i], l.data()[i], c.data()[i], v.data()[i]);89    json out; { nb::gil_scoped_release rel; auto st = b.normalize(); MarketCalendar cal; auto q = b.validate(cal, session == "extended" ? SessionPolicy::extended : SessionPolicy::regular); q.duplicates = st.duplicates_removed; q.unsorted = st.reordered; out = q.to_json(); }90    return json_to_py(out);91  });9293  m.def("calendar_sessions", [](const std::string& a, const std::string& b) { MarketCalendar cal; std::vector<std::string> out; for (auto d : cal.sessions(date_from_iso(a), date_from_iso(b))) out.push_back(d.iso()); return out; });9495  m.def("intraday_returns", [](I64 t, F64 c, const std::string& kind) {96    BarSeries b; b.interval = Interval::m1; const usize n = t.shape(0); b.reserve(n);97    for (usize i = 0; i < n; ++i) b.push_back(Timestamp(t.data()[i]), c.data()[i], c.data()[i], c.data()[i], c.data()[i], 0.0);98    data::IntradayReturns ir; { nb::gil_scoped_release rel; b.normalize(); MarketCalendar cal; auto reg = b.filter_session(cal, SessionPolicy::regular); ir = data::intraday_returns(reg, kind == "arithmetic" ? data::ReturnKind::arithmetic : data::ReturnKind::log); }99    nb::dict d; d["r"] = to_np(ir.r); d["offsets"] = to_np_i64(ir.session_offsets); std::vector<std::string> s; for (auto dd : ir.session_dates) s.push_back(dd.iso()); d["sessions"] = s; d["overnight"] = to_np(ir.overnight); d["close_to_close"] = to_np(ir.close_to_close); d["open_to_close"] = to_np(ir.open_to_close);100    return d;101  });102}103
Files · 114c1f28e7a32db3097a169b3ffc45c8b166be67
Menu

Theme