SPB Git

spb/wp40_uqo

Public MIT
19.7 KB · 184 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/estimators/har.hpp"16#include "raid/models/var.hpp"17#include "raid/statistics/bootstrap.hpp"18#include "raid/statistics/reductions.hpp"19#include "raid/data/bar_series.hpp"20#include "raid/data/calendar.hpp"21#include "raid/data/returns.hpp"22#include "raid/data/findata.hpp"23#include "raid/data/io.hpp"2425namespace nb = nanobind;26using namespace raid;27using json = nlohmann::json;28using F64 = nb::ndarray<const double, nb::c_contig, nb::device::cpu>;29using I64 = nb::ndarray<const int64_t, nb::c_contig, nb::device::cpu>;3031namespace {32nb::object json_to_py(const json& j) {33  switch (j.type()) {34    case json::value_t::null: return nb::none();35    case json::value_t::boolean: return nb::bool_(j.get<bool>());36    case json::value_t::number_integer: return nb::int_(j.get<long long>());37    case json::value_t::number_unsigned: return nb::int_(j.get<unsigned long long>());38    case json::value_t::number_float: return nb::float_(j.get<double>());39    case json::value_t::string: return nb::str(j.get<std::string>().c_str());40    case json::value_t::array: { nb::list l; for (auto& v : j) l.append(json_to_py(v)); return l; }41    case json::value_t::object: { nb::dict d; for (auto& [k, v] : j.items()) d[k.c_str()] = json_to_py(v); return d; }42    default: return nb::none();43  }44}45nb::ndarray<nb::numpy, double> to_np(std::vector<double> v) {46  auto* data = new std::vector<double>(std::move(v));47  nb::capsule owner(data, [](void* p) noexcept { delete static_cast<std::vector<double>*>(p); });48  return nb::ndarray<nb::numpy, double>(data->data(), {data->size()}, owner);49}50nb::ndarray<nb::numpy, int64_t> to_np_i64(std::vector<int64_t> v) {51  auto* data = new std::vector<int64_t>(std::move(v));52  nb::capsule owner(data, [](void* p) noexcept { delete static_cast<std::vector<int64_t>*>(p); });53  return nb::ndarray<nb::numpy, int64_t>(data->data(), {data->size()}, owner);54}55}  // namespace5657NB_MODULE(_core, m) {58  m.doc() = "RAID compiled core (C++20, CPU + optional Metal)";59  m.attr("__version__") = RAID_VERSION;60  nb::exception<RaidError>(m, "RaidError");6162  m.def("devices", []() { json out = json::array(); for (auto& d : backend::list_devices()) out.push_back(d.to_json()); return json_to_py(out); });6364  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) {65    RAID_REQUIRE(X.ndim() == 2 && y.ndim() == 1 && X.shape(0) == y.shape(0), invalid_argument, "X (n,k), y (n,)");66    const Eigen::Index n = static_cast<Eigen::Index>(X.shape(0)), k = static_cast<Eigen::Index>(X.shape(1));67    Eigen::Map<const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> Xm(X.data(), n, k);68    Eigen::Map<const Eigen::VectorXd> ym(y.data(), n);69    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;70    std::vector<i64> cl; if (clusters) { cl.assign(clusters->data(), clusters->data() + clusters->shape(0)); o.cov.clusters = cl; }71    std::vector<f64> w; if (weights) { w.assign(weights->data(), weights->data() + weights->shape(0)); o.weights = w; }72    Eigen::MatrixXd Xd = Xm; json out; { nb::gil_scoped_release rel; auto r = est::ols(Xd, ym, o); out = r.estimate.to_json(); }73    return json_to_py(out);74  }, 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>{});7576  m.def("realized_measures", [](F64 r, I64 off, const std::string& backend_s, bool bpv_fs, bool deterministic) {77    est::RealizedOptions o; o.plan.backend = backend_from_string(backend_s); o.bpv_finite_sample = bpv_fs; o.plan.deterministic = deterministic;78    backend::Batch b{std::span<const f64>(r.data(), r.shape(0)), std::span<const i64>(off.data(), off.shape(0))};79    est::RealizedMeasures mres; { nb::gil_scoped_release rel; mres = est::realized_measures_batch(b, o); }80    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);81    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;82    return d;83  }, nb::arg("returns"), nb::arg("offsets"), nb::arg("backend") = "auto", nb::arg("bpv_finite_sample") = true, nb::arg("deterministic") = true);8485  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) {86    models::GarchFitOptions o; o.model = model == "gjr" ? models::GarchModel::gjr : (model == "egarch" ? models::GarchModel::egarch : models::GarchModel::garch);87    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;88    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(); }89    return json_to_py(out);90  }, 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);9192  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) {93    BarSeries b; b.symbol = symbol; b.interval = interval_from_string(interval); const usize n = t.shape(0); b.reserve(n);94    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]);95    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(); }96    return json_to_py(out);97  });9899  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; });100101  // ---- HAR family ----102  m.def("har_fit", [](F64 rv, std::vector<usize> lags, int horizon, bool log_model, bool harq, std::optional<F64> rq, const std::string& cov, int hac_lags, const std::string& target) {103    est::HarSpec s; s.lags = lags; s.horizon = horizon; s.log = log_model; s.harq = harq; s.cov.type = stats::covtype_from_string(cov); s.cov.hac_lags = hac_lags; s.target = target == "single" ? est::HarTarget::single : est::HarTarget::average;104    std::span<const f64> rvs(rv.data(), rv.shape(0)); std::span<const f64> rqs; if (rq) rqs = std::span<const f64>(rq->data(), rq->shape(0));105    json out; { nb::gil_scoped_release rel; auto e = est::har_fit(rvs, s, rqs); out = e.to_json(); out["forecast_next"] = est::har_forecast(e, rvs, s, rqs); }106    return json_to_py(out);107  }, nb::arg("rv"), nb::arg("lags") = std::vector<usize>{1, 5, 22}, nb::arg("horizon") = 1, nb::arg("log") = false, nb::arg("harq") = false, nb::arg("rq") = nb::none(), nb::arg("cov") = "hac", nb::arg("hac_lags") = -1, nb::arg("target") = "average");108  m.def("har_rolling", [](F64 rv, std::vector<usize> lags, int horizon, bool log_model, usize window, usize step, bool harq, std::optional<F64> rq) {109    est::HarSpec s; s.lags = lags; s.horizon = horizon; s.log = log_model; s.harq = harq; std::span<const f64> rvs(rv.data(), rv.shape(0)); std::span<const f64> rqs; if (rq) rqs = std::span<const f64>(rq->data(), rq->shape(0));110    est::HarRolling r; { nb::gil_scoped_release rel; r = est::har_rolling(rvs, s, window, step, rqs); }111    nb::dict d; std::vector<int64_t> o(r.origin.begin(), r.origin.end()); d["origin"] = to_np_i64(std::move(o)); d["forecast"] = to_np(r.forecast); d["actual"] = to_np(r.actual); d["n_fits"] = r.n_fits; d["seconds"] = r.seconds;112    auto l = est::forecast_loss(r.actual, r.forecast); d["mse"] = l.mse; d["qlike"] = l.qlike; d["n_loss"] = l.n; return d;113  }, nb::arg("rv"), nb::arg("lags") = std::vector<usize>{1, 5, 22}, nb::arg("horizon") = 1, nb::arg("log") = false, nb::arg("window") = 1000, nb::arg("step") = 1, nb::arg("harq") = false, nb::arg("rq") = nb::none());114  m.def("har_batch", [](F64 rv, I64 off, std::vector<usize> lags, int horizon, bool log_model, const std::string& backend_s) {115    est::HarSpec s; s.lags = lags; s.horizon = horizon; s.log = log_model; ExecutionPlan p; p.backend = backend_from_string(backend_s);116    est::HarBatchResult r; { nb::gil_scoped_release rel; r = est::har_batch(std::span<const f64>(rv.data(), rv.shape(0)), std::span<const i64>(off.data(), off.shape(0)), s, p); }117    nb::list fits; for (auto& e : r.fits) fits.append(json_to_py(e.to_json())); nb::dict d; d["fits"] = fits; d["backend"] = std::string(to_string(r.stats.backend)); d["precision"] = std::string(to_string(r.stats.precision)); d["wall_seconds"] = r.stats.wall_seconds; d["device_seconds"] = r.stats.device_seconds; return d;118  }, nb::arg("rv"), nb::arg("offsets"), nb::arg("lags") = std::vector<usize>{1, 5, 22}, nb::arg("horizon") = 1, nb::arg("log") = false, nb::arg("backend") = "auto");119  // ---- VAR / connectedness ----120  m.def("var_connectedness", [](F64 Y, int p, int H, usize window, usize step) {121    RAID_REQUIRE(Y.ndim() == 2, invalid_argument, "Y must be (T, N)"); const Eigen::Index T = static_cast<Eigen::Index>(Y.shape(0)), N = static_cast<Eigen::Index>(Y.shape(1));122    Eigen::MatrixXd Ym = Eigen::Map<const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(Y.data(), T, N);123    nb::dict d; { nb::gil_scoped_release rel; auto f = models::var_fit(Ym, p); auto c = models::connectedness(f, H);124      std::vector<f64> tbl(static_cast<usize>(N * N)); for (Eigen::Index i = 0; i < N; ++i) for (Eigen::Index j = 0; j < N; ++j) tbl[static_cast<usize>(i * N + j)] = c.table(i, j);125      std::vector<f64> fr(c.from.data(), c.from.data() + N), to(c.to.data(), c.to.data() + N), net(c.net.data(), c.net.data() + N);126      nb::gil_scoped_acquire acq; d["total"] = c.total; d["table_flat"] = to_np(tbl); d["from"] = to_np(fr); d["to"] = to_np(to); d["net"] = to_np(net); d["stable"] = f.stable; d["max_root"] = f.max_root; d["loglik"] = f.loglik; d["n_obs"] = f.n_obs;127      if (window > 0) { auto r = models::rolling_connectedness(Ym, p, H, window, step); std::vector<int64_t> e(r.end_index.begin(), r.end_index.end()); d["rolling_end"] = to_np_i64(std::move(e)); d["rolling_total"] = to_np(r.total); d["rolling_failures"] = r.failures; d["rolling_seconds"] = r.seconds; } }128    return d;129  }, nb::arg("Y"), nb::arg("p") = 1, nb::arg("H") = 10, nb::arg("window") = 0, nb::arg("step") = 1);130  m.def("ledoit_wolf", [](F64 X) { const Eigen::Index n = static_cast<Eigen::Index>(X.shape(0)), p = static_cast<Eigen::Index>(X.shape(1)); Eigen::MatrixXd Xm = Eigen::Map<const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>>(X.data(), n, p); auto r = models::ledoit_wolf(Xm); std::vector<f64> c(static_cast<usize>(p * p)); for (Eigen::Index i = 0; i < p; ++i) for (Eigen::Index j = 0; j < p; ++j) c[static_cast<usize>(i * p + j)] = r.cov(i, j); nb::dict d; d["cov_flat"] = to_np(c); d["shrinkage"] = r.shrinkage; d["mu"] = r.mu; return d; });131132  // ---- bootstrap ----133  m.def("bootstrap", [](F64 x, const std::string& stat, const std::string& scheme, double mean_block, usize replicas, uint64_t seed, double level) {134    stats::BootstrapOptions o; o.scheme = scheme == "circular" ? stats::BootstrapScheme::circular_block : (scheme == "iid" ? stats::BootstrapScheme::iid : stats::BootstrapScheme::stationary); o.mean_block = mean_block; o.replicas = replicas; o.seed = seed; o.ci_level = level;135    std::function<f64(std::span<const f64>)> f; if (stat == "mean") f = [](std::span<const f64> v) { return stats::mean<f64>(v); }; else if (stat == "rv" || stat == "sum_squares") f = [](std::span<const f64> v) { return stats::sum_squares<f64>(v); }; else if (stat == "variance") f = [](std::span<const f64> v) { return stats::variance<f64>(v); }; else throw RaidError(ErrorCode::invalid_argument, "stat must be mean|rv|variance");136    stats::BootstrapResult r; { nb::gil_scoped_release rel; r = stats::bootstrap(std::span<const f64>(x.data(), x.shape(0)), f, o); }137    nb::dict d; d["point"] = r.point; d["se"] = r.se; d["ci_lo"] = r.ci_lo; d["ci_hi"] = r.ci_hi; d["bias"] = r.bias; d["failures"] = r.failures; d["block_length"] = r.block_length; d["statistics"] = to_np(r.statistics); d["seconds"] = r.seconds; return d;138  }, nb::arg("x"), nb::arg("stat") = "mean", nb::arg("scheme") = "stationary", nb::arg("mean_block") = 0.0, nb::arg("replicas") = 999, nb::arg("seed") = 20261102, nb::arg("level") = 0.95);139140  // ---- FinData client (keys from FINDATA_API_KEY / FINDATA_API_KEY_FILE; never passed through Python) ----141  m.def("findata_bars", [](const std::string& symbol, const std::string& interval, const std::string& start, const std::string& end, const std::string& adjustment, const std::string& cache_dir, bool offline) {142    findata::Config cfg = findata::Config::from_env(); if (!cache_dir.empty()) cfg.cache_dir = cache_dir; cfg.offline = offline; findata::Client cl(cfg);143    findata::BarsRequest r; r.symbol = symbol; r.interval = interval_from_string(interval); r.start = date_from_iso(start); r.end = date_from_iso(end); r.adjustment = adjustment;144    findata::BarsResult res; { nb::gil_scoped_release rel; res = cl.bars(r); }145    nb::dict d; std::vector<int64_t> t(res.bars.t_open.size()); for (usize i = 0; i < t.size(); ++i) t[i] = res.bars.t_open[i].us;146    d["timestamp_us"] = to_np_i64(std::move(t)); d["open"] = to_np(res.bars.open); d["high"] = to_np(res.bars.high); d["low"] = to_np(res.bars.low); d["close"] = to_np(res.bars.close); d["volume"] = to_np(res.bars.volume);147    std::vector<int64_t> so(res.bars.session_offsets.begin(), res.bars.session_offsets.end()); d["session_offsets"] = to_np_i64(std::move(so)); std::vector<std::string> sd; for (auto x : res.bars.session_dates) sd.push_back(x.iso()); d["sessions"] = sd;148    d["symbol"] = res.bars.symbol; d["instrument_id"] = res.bars.instrument_id; d["dataset_id"] = res.dataset_id; d["dataset_version"] = res.dataset_version; d["from_cache"] = res.from_cache; d["cache_path"] = res.cache_path; d["duplicates_dropped"] = res.duplicates_dropped;149    std::vector<std::string> rid; for (auto& e : res.exchanges) rid.push_back(e.request_id); d["request_ids"] = rid; d["key_fingerprint"] = cl.key_fingerprint();150    return d;151  }, nb::arg("symbol"), nb::arg("interval") = "1m", nb::arg("start"), nb::arg("end"), nb::arg("adjustment") = "raw", nb::arg("cache_dir") = "", nb::arg("offline") = false);152  m.def("findata_export", [](std::vector<std::string> symbols, const std::string& dataset_id, const std::string& start, const std::string& end, const std::string& adjustment, const std::string& out_path, const std::string& cache_dir) {153    findata::Config cfg = findata::Config::from_env(); if (!cache_dir.empty()) cfg.cache_dir = cache_dir; findata::Client cl(cfg);154    findata::ExportRequest r; r.symbols = symbols; r.dataset_id = dataset_id; r.start = date_from_iso(start); r.end = date_from_iso(end); r.adjustment = adjustment;155    findata::ExportResult res; { nb::gil_scoped_release rel; res = cl.export_bars(r, out_path); }156    nb::dict d; d["export_id"] = res.export_id; d["path"] = res.path; d["sha256"] = res.sha256; d["dataset_version"] = res.dataset_version; d["rows"] = res.rows; d["bytes"] = res.bytes; d["from_cache"] = res.from_cache; d["seconds"] = res.seconds_total; d["key_fingerprint"] = cl.key_fingerprint();157    return d;158  }, nb::arg("symbols"), nb::arg("dataset_id") = "market_bars.stock.1m.raw", nb::arg("start"), nb::arg("end"), nb::arg("adjustment") = "raw", nb::arg("out_path") = "", nb::arg("cache_dir") = "");159  m.def("findata_json", [](const std::string& what, const std::string& a, const std::string& b, const std::string& c, int n) {160    findata::Client cl; json out; { nb::gil_scoped_release rel;161      if (what == "health") out = cl.health(); else if (what == "catalog") out = cl.catalog_datasets(); else if (what == "dataset") out = cl.dataset(a);162      else if (what == "sessions") out = cl.sessions(a, interval_from_string(b.empty() ? "1m" : b), date_from_iso(c.substr(0, 10)), date_from_iso(c.substr(11, 10)));163      else if (what == "coverage") out = cl.coverage_symbols_all(a.empty() ? "market_bars.stock.1m.raw" : a, n, 100000);164      else if (what == "fundamentals") { std::vector<std::string> m; if (!c.empty()) { usize p0 = 0; for (usize i = 0; i <= c.size(); ++i) if (i == c.size() || c[i] == ',') { m.push_back(c.substr(p0, i - p0)); p0 = i + 1; } } out = cl.fundamentals(a, b, m); }165      else if (what == "query") out = cl.query(a, static_cast<usize>(n)); else throw RaidError(ErrorCode::invalid_argument, "unknown findata_json kind " + what); }166    return json_to_py(out);167  }, nb::arg("what"), nb::arg("a") = "", nb::arg("b") = "", nb::arg("c") = "", nb::arg("n") = 1000);168#ifdef RAID_WITH_ARROW169  m.def("read_bars_parquet", [](const std::string& path, std::vector<std::string> symbols) {170    data::ParquetBars pb; { nb::gil_scoped_release rel; pb = data::read_bars_parquet(path, symbols); }171    nb::dict out; for (auto& [sym, b] : pb.series) { nb::dict d; std::vector<int64_t> t(b.t_open.size()); for (usize i = 0; i < t.size(); ++i) t[i] = b.t_open[i].us; d["timestamp_us"] = to_np_i64(std::move(t)); d["open"] = to_np(b.open); d["high"] = to_np(b.high); d["low"] = to_np(b.low); d["close"] = to_np(b.close); d["volume"] = to_np(b.volume); std::vector<int64_t> so(b.session_offsets.begin(), b.session_offsets.end()); d["session_offsets"] = to_np_i64(std::move(so)); std::vector<std::string> sd; for (auto x : b.session_dates) sd.push_back(x.iso()); d["sessions"] = sd; d["instrument_id"] = b.instrument_id; out[sym.c_str()] = d; }172    nb::dict meta; for (auto& [k, v] : pb.metadata) meta[k.c_str()] = v; out["__meta__"] = meta; out["__rows__"] = pb.rows_read; out["__duplicates_dropped__"] = pb.duplicates_dropped; return out;173  }, nb::arg("path"), nb::arg("symbols") = std::vector<std::string>{});174#endif175176  m.def("intraday_returns", [](I64 t, F64 c, const std::string& kind) {177    BarSeries b; b.interval = Interval::m1; const usize n = t.shape(0); b.reserve(n);178    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);179    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); }180    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);181    return d;182  });183}184
Files · d628c807b304dba507dfb8130278ae5bb0f92411
Menu

Theme