用稠密行索引和滚动游标加速回测

This commit is contained in:
boris
2026-08-24 11:46:55 +08:00
parent 911074ae95
commit 0686532be0
2 changed files with 208 additions and 39 deletions
+159 -28
View File
@@ -78,6 +78,12 @@ pub enum PriceField {
Last,
}
#[derive(Debug, Clone, Copy)]
pub struct MarketRollingCursor {
current_end: usize,
decision_end: usize,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DailyMarketSnapshot {
#[serde(with = "date_format")]
@@ -548,14 +554,18 @@ impl AdjustedCloseSeries {
}
fn current_moving_average(&self, date: NaiveDate, lookback: usize) -> Option<f64> {
if lookback == 0 {
return None;
}
let end = match self.dates.binary_search(&date) {
Ok(index) => index + 1,
Err(0) => return None,
Err(index) => index,
};
self.current_moving_average_at_end(end, lookback)
}
fn current_moving_average_at_end(&self, end: usize, lookback: usize) -> Option<f64> {
if lookback == 0 || end > self.dates.len() {
return None;
}
if end < lookback {
return None;
}
@@ -577,14 +587,18 @@ impl AdjustedCloseSeries {
}
fn decision_moving_average(&self, date: NaiveDate, lookback: usize) -> Option<f64> {
if lookback == 0 {
return None;
}
let end = match self.dates.binary_search(&date) {
Ok(index) => index,
Err(0) => return None,
Err(index) => index,
};
self.decision_moving_average_at_end(end, lookback)
}
fn decision_moving_average_at_end(&self, end: usize, lookback: usize) -> Option<f64> {
if lookback == 0 || end > self.dates.len() {
return None;
}
if end < lookback {
return None;
}
@@ -854,6 +868,10 @@ impl SymbolPriceSeries {
fn decision_volume_moving_average(&self, date: NaiveDate, lookback: usize) -> Option<f64> {
let end = self.previous_completed_end_index(date)?;
self.decision_volume_moving_average_at_end(end, lookback)
}
fn decision_volume_moving_average_at_end(&self, end: usize, lookback: usize) -> Option<f64> {
self.valid_volume_window(end, lookback).map(|(start, end)| {
normalize_rolling_factor(
(self.valid_volume_sum_prefix[end] - self.valid_volume_sum_prefix[start])
@@ -865,6 +883,10 @@ impl SymbolPriceSeries {
fn current_volume_moving_average(&self, date: NaiveDate, lookback: usize) -> Option<f64> {
let end = self.end_index(date)?;
self.current_volume_moving_average_at_end(end, lookback)
}
fn current_volume_moving_average_at_end(&self, end: usize, lookback: usize) -> Option<f64> {
self.valid_volume_window(end, lookback).map(|(start, end)| {
normalize_rolling_factor(
(self.valid_volume_sum_prefix[end] - self.valid_volume_sum_prefix[start])
@@ -898,6 +920,13 @@ impl SymbolPriceSeries {
Some((start, end))
}
fn rolling_cursor(&self, date: NaiveDate) -> Option<MarketRollingCursor> {
Some(MarketRollingCursor {
current_end: self.end_index(date)?,
decision_end: self.previous_completed_end_index(date).unwrap_or(0),
})
}
fn valid_volume_values(&self, end: usize, lookback: usize) -> Option<Vec<f64>> {
let (start, end) = self.valid_volume_window(end, lookback)?;
let values = self.volumes[start..end]
@@ -1125,13 +1154,14 @@ pub struct DataSet {
instruments: HashMap<String, Instrument>,
calendar: TradingCalendar,
market_by_date: BTreeMap<NaiveDate, Vec<Arc<DailyMarketSnapshot>>>,
market_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
market_row_index_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
factor_by_date: BTreeMap<NaiveDate, Vec<Arc<DailyFactorSnapshot>>>,
factor_row_index_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
factor_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
factor_text_by_date: BTreeMap<NaiveDate, Vec<FactorTextValue>>,
factor_text_index: HashMap<(NaiveDate, String, String), FactorTextValue>,
candidate_by_date: BTreeMap<NaiveDate, Vec<Arc<CandidateEligibility>>>,
candidate_symbol_ids_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
candidate_row_index_by_date: Arc<BTreeMap<NaiveDate, Vec<u32>>>,
corporate_actions_by_date: BTreeMap<NaiveDate, Vec<CorporateAction>>,
execution_quotes_by_date: HashMap<NaiveDate, HashMap<String, Vec<IntradayExecutionQuote>>>,
order_book_depth_index: HashMap<(NaiveDate, String), Vec<IntradayOrderBookDepthLevel>>,
@@ -1351,16 +1381,20 @@ impl DataSet {
&factor_by_date,
&candidate_by_date,
);
let market_symbol_ids_by_date =
build_group_symbol_ids(&market_by_date, &symbol_id_by_code, |item| {
let market_row_index_by_date =
build_group_row_indices(&market_by_date, &symbol_id_by_code, |item| {
item.symbol.as_str()
});
let factor_row_index_by_date =
build_group_row_indices(&factor_by_date, &symbol_id_by_code, |item| {
item.symbol.as_str()
});
let factor_symbol_ids_by_date =
build_group_symbol_ids(&factor_by_date, &symbol_id_by_code, |item| {
item.symbol.as_str()
});
let candidate_symbol_ids_by_date =
build_group_symbol_ids(&candidate_by_date, &symbol_id_by_code, |item| {
let candidate_row_index_by_date =
build_group_row_indices(&candidate_by_date, &symbol_id_by_code, |item| {
item.symbol.as_str()
});
let mut market_series_by_symbol_id = vec![None; symbol_id_by_code.len()];
@@ -1391,13 +1425,14 @@ impl DataSet {
instruments,
calendar,
market_by_date,
market_symbol_ids_by_date: Arc::new(market_symbol_ids_by_date),
market_row_index_by_date: Arc::new(market_row_index_by_date),
factor_by_date,
factor_row_index_by_date: Arc::new(factor_row_index_by_date),
factor_symbol_ids_by_date: Arc::new(factor_symbol_ids_by_date),
factor_text_by_date,
factor_text_index,
candidate_by_date,
candidate_symbol_ids_by_date: Arc::new(candidate_symbol_ids_by_date),
candidate_row_index_by_date: Arc::new(candidate_row_index_by_date),
corporate_actions_by_date,
execution_quotes_by_date,
order_book_depth_index,
@@ -1465,9 +1500,9 @@ impl DataSet {
date: NaiveDate,
symbol_id: u32,
) -> Option<&DailyMarketSnapshot> {
find_arc_by_symbol_id(
find_arc_by_dense_row_index(
self.market_by_date.get(&date)?,
self.market_symbol_ids_by_date.get(&date)?,
self.market_row_index_by_date.get(&date)?,
symbol_id,
)
}
@@ -1482,6 +1517,15 @@ impl DataSet {
.as_deref()
}
pub fn market_rolling_cursor_by_symbol_id(
&self,
date: NaiveDate,
symbol_id: u32,
) -> Option<MarketRollingCursor> {
self.market_series_by_symbol_id(symbol_id)?
.rolling_cursor(date)
}
fn adjusted_close_series(&self, symbol: &str) -> Option<&AdjustedCloseSeries> {
self.adjusted_close_series_by_symbol
.get(symbol)
@@ -1504,9 +1548,9 @@ impl DataSet {
date: NaiveDate,
symbol_id: u32,
) -> Option<&DailyFactorSnapshot> {
find_arc_by_symbol_id(
find_arc_by_dense_row_index(
self.factor_by_date.get(&date)?,
self.factor_symbol_ids_by_date.get(&date)?,
self.factor_row_index_by_date.get(&date)?,
symbol_id,
)
}
@@ -1521,9 +1565,9 @@ impl DataSet {
date: NaiveDate,
symbol_id: u32,
) -> Option<&CandidateEligibility> {
find_arc_by_symbol_id(
find_arc_by_dense_row_index(
self.candidate_by_date.get(&date)?,
self.candidate_symbol_ids_by_date.get(&date)?,
self.candidate_row_index_by_date.get(&date)?,
symbol_id,
)
}
@@ -2617,6 +2661,38 @@ impl DataSet {
}
}
pub fn market_decision_numeric_moving_average_at_cursor(
&self,
date: NaiveDate,
symbol_id: u32,
symbol: &str,
field: &str,
lookback: usize,
cursor: MarketRollingCursor,
) -> Option<f64> {
let field = normalized_field(field);
match field.as_ref() {
"close" | "prev_close" | "stock_close" | "price" => self
.adjusted_close_series_by_symbol_id(symbol_id)
.and_then(|series| {
series.decision_moving_average_at_end(cursor.decision_end, lookback)
}),
"volume" | "stock_volume" => {
self.market_series_by_symbol_id(symbol_id)
.and_then(|series| {
series.decision_volume_moving_average_at_end(cursor.decision_end, lookback)
})
}
_ => self.market_decision_numeric_moving_average_by_symbol_id(
date,
symbol_id,
symbol,
field.as_ref(),
lookback,
),
}
}
pub fn market_current_numeric_moving_average(
&self,
date: NaiveDate,
@@ -2672,6 +2748,38 @@ impl DataSet {
}
}
pub fn market_current_numeric_moving_average_at_cursor(
&self,
date: NaiveDate,
symbol_id: u32,
symbol: &str,
field: &str,
lookback: usize,
cursor: MarketRollingCursor,
) -> Option<f64> {
let field = normalized_field(field);
match field.as_ref() {
"close" | "prev_close" | "stock_close" | "price" => self
.adjusted_close_series_by_symbol_id(symbol_id)
.and_then(|series| {
series.current_moving_average_at_end(cursor.current_end, lookback)
}),
"volume" | "stock_volume" => {
self.market_series_by_symbol_id(symbol_id)
.and_then(|series| {
series.current_volume_moving_average_at_end(cursor.current_end, lookback)
})
}
_ => self.market_current_numeric_moving_average_by_symbol_id(
date,
symbol_id,
symbol,
field.as_ref(),
lookback,
),
}
}
pub fn market_latest_back_adjusted_close(&self, date: NaiveDate, symbol: &str) -> Option<f64> {
self.adjusted_close_series(symbol)
.and_then(|series| series.latest_back_adjusted_close(date))
@@ -3360,17 +3468,40 @@ where
.collect()
}
fn find_arc_by_symbol_id<'a, T>(
fn build_group_row_indices<T, F>(
groups: &BTreeMap<NaiveDate, Vec<Arc<T>>>,
symbol_id_by_code: &HashMap<String, u32>,
symbol_of: F,
) -> BTreeMap<NaiveDate, Vec<u32>>
where
F: Fn(&T) -> &str + Copy,
{
groups
.iter()
.map(|(date, rows)| {
let mut row_indices = vec![u32::MAX; symbol_id_by_code.len()];
for (row_index, row) in rows.iter().enumerate() {
let symbol_id = *symbol_id_by_code
.get(symbol_of(row.as_ref()))
.expect("snapshot symbol missing from FIDC symbol index");
let slot = &mut row_indices[symbol_id as usize];
debug_assert_eq!(*slot, u32::MAX, "duplicate symbol in daily snapshot rows");
*slot = u32::try_from(row_index)
.expect("daily snapshot row index exceeds u32 capacity");
}
(*date, row_indices)
})
.collect()
}
fn find_arc_by_dense_row_index<'a, T>(
rows: &'a [Arc<T>],
symbol_ids: &[u32],
row_indices: &[u32],
symbol_id: u32,
) -> Option<&'a T> {
if rows.len() != symbol_ids.len() {
return None;
}
symbol_ids
.binary_search(&symbol_id)
.ok()
let row_index = *row_indices.get(symbol_id as usize)?;
(row_index != u32::MAX)
.then_some(row_index as usize)
.and_then(|index| rows.get(index))
.map(Arc::as_ref)
}
+49 -11
View File
@@ -8,8 +8,8 @@ use rhai::{AST, Dynamic, Engine, Map, Scope};
use crate::broker::{MatchingType, RebalanceCashMode, SlippageModel};
use crate::cost::ChinaAShareCostModel;
use crate::data::{
DailyMarketSnapshot, EligibleUniverseSnapshot, PriceField, decision_free_float_cap_bn,
decision_market_cap_bn,
DailyMarketSnapshot, EligibleUniverseSnapshot, MarketRollingCursor, PriceField,
decision_free_float_cap_bn, decision_market_cap_bn,
};
use crate::engine::BacktestError;
use crate::events::OrderSide;
@@ -616,6 +616,8 @@ struct DayExpressionState {
struct StockExpressionState {
symbol: String,
symbol_id: u32,
rolling_date: NaiveDate,
rolling_cursor: MarketRollingCursor,
market_cap: f64,
market_cap_bn: f64,
free_float_cap: f64,
@@ -3614,14 +3616,21 @@ impl PlatformExprStrategy {
ctx: &StrategyContext<'_>,
date: NaiveDate,
symbol_id: u32,
rolling_cursor: Option<MarketRollingCursor>,
symbol: &str,
field: &str,
lookback: usize,
) -> Option<f64> {
ctx.data
.market_decision_numeric_moving_average_by_symbol_id(
date, symbol_id, symbol, field, lookback,
)
match rolling_cursor {
Some(cursor) => ctx.data.market_decision_numeric_moving_average_at_cursor(
date, symbol_id, symbol, field, lookback, cursor,
),
None => ctx
.data
.market_decision_numeric_moving_average_by_symbol_id(
date, symbol_id, symbol, field, lookback,
),
}
}
fn stock_current_rolling_mean(
@@ -3629,13 +3638,19 @@ impl PlatformExprStrategy {
ctx: &StrategyContext<'_>,
date: NaiveDate,
symbol_id: u32,
rolling_cursor: Option<MarketRollingCursor>,
symbol: &str,
field: &str,
lookback: usize,
) -> Option<f64> {
ctx.data.market_current_numeric_moving_average_by_symbol_id(
date, symbol_id, symbol, field, lookback,
)
match rolling_cursor {
Some(cursor) => ctx.data.market_current_numeric_moving_average_at_cursor(
date, symbol_id, symbol, field, lookback, cursor,
),
None => ctx.data.market_current_numeric_moving_average_by_symbol_id(
date, symbol_id, symbol, field, lookback,
),
}
}
fn stock_state_at_time(
@@ -3716,6 +3731,16 @@ impl PlatformExprStrategy {
symbol: symbol.to_string(),
})
})?;
let rolling_cursor = ctx
.data
.market_rolling_cursor_by_symbol_id(date, symbol_id)
.ok_or_else(|| {
BacktestError::Data(crate::data::DataSetError::MissingSnapshot {
kind: "market rolling cursor",
date,
symbol: symbol.to_string(),
})
})?;
{
let mut cache_date = self.stock_state_cache_date.borrow_mut();
if *cache_date != Some(date) {
@@ -3760,8 +3785,16 @@ impl PlatformExprStrategy {
if !self.stock_rolling_requirements.requires(field, lookback) {
return f64::NAN;
}
self.stock_decision_rolling_mean(ctx, date, symbol_id, symbol, field, lookback)
.unwrap_or(f64::NAN)
self.stock_decision_rolling_mean(
ctx,
date,
symbol_id,
Some(rolling_cursor),
symbol,
field,
lookback,
)
.unwrap_or(f64::NAN)
};
let stock_ma_short = rolling("close", self.config.stock_short_ma_days);
let stock_ma_mid = rolling("close", self.config.stock_mid_ma_days);
@@ -3887,6 +3920,8 @@ impl PlatformExprStrategy {
let state = StockExpressionState {
symbol: symbol.to_string(),
symbol_id,
rolling_date: date,
rolling_cursor,
market_cap,
market_cap_bn,
free_float_cap,
@@ -5374,6 +5409,7 @@ impl PlatformExprStrategy {
ctx,
day.date,
stock.symbol_id,
(stock.rolling_date == day.date).then_some(stock.rolling_cursor),
&stock.symbol,
field,
lookback,
@@ -6038,6 +6074,7 @@ impl PlatformExprStrategy {
ctx,
day.date,
stock.symbol_id,
(stock.rolling_date == day.date).then_some(stock.rolling_cursor),
&stock.symbol,
other,
lookback,
@@ -6087,6 +6124,7 @@ impl PlatformExprStrategy {
ctx,
day.date,
stock.symbol_id,
(stock.rolling_date == day.date).then_some(stock.rolling_cursor),
&stock.symbol,
other,
lookback,