use crate::core::native_engine::{DailysSoA, PairsSoA};
use anyhow::Context;
use errors::WbtError;
use polars::prelude::*;
use std::path::Path;
mod backtest;
pub mod backtest_result_wire;
pub mod cal_yearly_days;
pub mod daily_performance;
pub mod errors;
mod evaluate_pairs;
pub mod is_good_strategy;
pub mod key_trades;
pub mod native_engine;
pub mod period_win_rates;
pub mod position_risk;
mod report;
pub mod rolling_daily_performance;
pub mod top_drawdowns;
pub mod trade_dir;
pub mod utils;
pub mod yearly_return;
pub use evaluate_pairs::EvaluatePairs;
pub use report::{Report, StatsReport, SymbolsReport};
pub use utils::WeightType;
pub struct WeightBacktest {
dfw: DataFrame,
digits: i64,
fee_rate: f64,
symbols: Vec<Arc<str>>,
dailys_soa: Option<DailysSoA>,
pairs_soa: Option<PairsSoA>,
daily_return_cache: Option<DataFrame>,
dailys_cache: Option<DataFrame>,
pairs_cache: Option<DataFrame>,
agg_pairs_cache: Option<Vec<key_trades::AggRow>>,
weight_type: Option<WeightType>,
report: Option<Report>,
yearly_days: usize,
}
impl WeightBacktest {
pub fn dfw(&self) -> &DataFrame {
&self.dfw
}
pub fn digits(&self) -> i64 {
self.digits
}
pub fn fee_rate(&self) -> f64 {
self.fee_rate
}
pub fn symbols(&self) -> &[Arc<str>] {
&self.symbols
}
pub fn report(&self) -> Option<&Report> {
self.report.as_ref()
}
pub fn yearly_days(&self) -> usize {
self.yearly_days
}
pub fn weight_type(&self) -> Option<WeightType> {
self.weight_type
}
pub fn new(dfw: DataFrame, digits: i64, fee_rate: Option<f64>) -> Result<Self, WbtError> {
if !(0..=10).contains(&digits) {
return Err(WbtError::InvalidInput(format!(
"digits must be in 0..=10, got {digits}"
)));
}
let dfw = Self::validate_input(dfw)?;
let mut dfw = Self::convert_datetime(dfw)
.map_err(|e| WbtError::InvalidInput(format!("column 'dt': {e}")))?;
Self::round_weight(&mut dfw, digits).context("Failed to round weight")?;
let symbols = Self::unique_symbols(&dfw).context("Failed to unique_symbols")?;
let dfw = {
let n_rows = dfw.height();
let n_syms = symbols.len();
let mut order_map: hashbrown::HashMap<&str, u32> =
hashbrown::HashMap::with_capacity(n_syms);
for (idx, sym) in symbols.iter().enumerate() {
order_map.insert(sym.as_ref(), idx as u32);
}
let sym_ca = dfw.column("symbol")?.as_materialized_series().str()?;
let sym_ids: Vec<u32> = sym_ca
.into_iter()
.map(|opt_s| {
opt_s
.and_then(|s| order_map.get(s).copied())
.ok_or_else(|| {
WbtError::InvalidInput("column 'symbol': missing symbol mapping".into())
})
})
.collect::<Result<_, _>>()?;
drop(order_map);
let mut bucket_counts = vec![0u32; n_syms];
for &sid in &sym_ids {
bucket_counts[sid as usize] += 1;
}
let mut write_pos = vec![0u32; n_syms];
let mut acc = 0u32;
for i in 0..n_syms {
write_pos[i] = acc;
acc += bucket_counts[i];
}
let mut perm = vec![0u32; n_rows];
for (i, &sid_val) in sym_ids.iter().enumerate().take(n_rows) {
let sid = sid_val as usize;
perm[write_pos[sid] as usize] = i as u32;
write_pos[sid] += 1;
}
let perm_idx = IdxCa::new(PlSmallStr::from("idx"), &perm);
let sym_id_vals: Vec<u32> = perm.iter().map(|&i| sym_ids[i as usize]).collect();
DataFrame::new_infer_height(vec![
Column::new("sym_id".into(), sym_id_vals),
dfw.column("dt")?
.as_materialized_series()
.take(&perm_idx)?
.into_column(),
dfw.column("weight")?
.as_materialized_series()
.take(&perm_idx)?
.into_column(),
dfw.column("price")?
.as_materialized_series()
.take(&perm_idx)?
.into_column(),
dfw.column("symbol")?
.as_materialized_series()
.take(&perm_idx)?
.into_column(),
])?
};
let wb = Self {
dfw,
digits,
symbols,
fee_rate: fee_rate.unwrap_or(0.0002),
dailys_soa: None,
pairs_soa: None,
daily_return_cache: None,
dailys_cache: None,
pairs_cache: None,
agg_pairs_cache: None,
weight_type: None,
report: None,
yearly_days: 252,
};
Ok(wb)
}
pub fn from_file(path: &str, digits: i64, fee_rate: Option<f64>) -> Result<Self, WbtError> {
let p = Path::new(path);
let ext = p
.extension()
.and_then(|e| e.to_str())
.unwrap_or("")
.to_lowercase();
let df = match ext.as_str() {
"csv" => {
CsvReader::new(std::fs::File::open(p).map_err(|e| WbtError::Io(e.to_string()))?)
.finish()
.map_err(WbtError::Polars)?
}
"parquet" => {
let file = std::fs::File::open(p).map_err(|e| WbtError::Io(e.to_string()))?;
ParquetReader::new(file)
.finish()
.map_err(WbtError::Polars)?
}
"feather" | "arrow" => {
let file = std::fs::File::open(p).map_err(|e| WbtError::Io(e.to_string()))?;
IpcReader::new(file).finish().map_err(WbtError::Polars)?
}
_ => {
return Err(WbtError::Io(format!(
"Unsupported file format: '{}'. Supported: csv, parquet, feather, arrow",
ext
)));
}
};
Self::new(df, digits, fee_rate)
}
pub fn backtest(
&mut self,
n_jobs: Option<usize>,
weight_type: WeightType,
yearly_days: usize,
) -> Result<(), WbtError> {
let n_jobs = n_jobs.unwrap_or(4);
let pool = rayon::ThreadPoolBuilder::new()
.stack_size(64 * 1024 * 1024)
.num_threads(n_jobs)
.build()
.context("Failed to create thread pool")?;
pool.install(|| self.do_backtest(weight_type, yearly_days))
}
pub fn daily_return_df(&mut self) -> Result<&DataFrame, WbtError> {
if self.daily_return_cache.is_none() {
let dailys_soa = self
.dailys_soa
.as_ref()
.ok_or_else(|| WbtError::NoneValue("dailys_soa not computed yet".into()))?;
let report = self
.report
.as_ref()
.ok_or_else(|| WbtError::NoneValue("report not computed yet".into()))?;
let weight_type = self
.weight_type
.ok_or_else(|| WbtError::NoneValue("weight_type not computed yet".into()))?;
let df = Self::build_daily_return_df(dailys_soa, &report.daily_totals, weight_type)?;
self.daily_return_cache = Some(df);
}
Ok(self.daily_return_cache.as_ref().unwrap())
}
pub fn dailys_df(&mut self) -> Result<&DataFrame, WbtError> {
if self.dailys_cache.is_none() {
let df = self
.dailys_soa
.as_ref()
.ok_or_else(|| WbtError::NoneValue("dailys_soa not computed yet".into()))?
.to_dataframe()?;
self.dailys_cache = Some(df);
}
Ok(self.dailys_cache.as_ref().unwrap())
}
pub fn pairs_df(&mut self) -> Result<Option<&DataFrame>, WbtError> {
if self.pairs_soa.is_none() {
return Ok(None);
}
if self.pairs_cache.is_none() {
let df = self.pairs_soa.as_ref().unwrap().to_dataframe()?;
self.pairs_cache = Some(df);
}
Ok(self.pairs_cache.as_ref())
}
fn ensure_agg_pairs(&mut self) -> Result<bool, WbtError> {
if self.pairs_soa.is_none() {
return Ok(false);
}
if self.agg_pairs_cache.is_none() {
let agg = key_trades::aggregate_pairs(self.pairs_soa.as_ref().unwrap());
self.agg_pairs_cache = Some(agg);
}
Ok(true)
}
pub fn aggregated_pairs_df(&mut self) -> Result<Option<DataFrame>, WbtError> {
if !self.ensure_agg_pairs()? {
return Ok(None);
}
let soa = self.pairs_soa.as_ref().unwrap();
let agg = self.agg_pairs_cache.as_ref().unwrap();
Ok(Some(key_trades::agg_to_df(soa, agg)?))
}
pub fn key_trades_df(&mut self, top: usize) -> Result<Option<DataFrame>, WbtError> {
if !self.ensure_agg_pairs()? {
return Ok(None);
}
let soa = self.pairs_soa.as_ref().unwrap();
let agg = self.agg_pairs_cache.as_ref().unwrap();
Ok(Some(key_trades::key_trades_to_df(soa, agg, top)?))
}
pub fn yearly_return_df(&mut self, min_days: usize) -> Result<DataFrame, WbtError> {
let wide = self.daily_return_df()?;
yearly_return::compute_yearly_returns(wide, min_days)
}
pub fn alpha_df(&self) -> Result<DataFrame, WbtError> {
let report = self
.report
.as_ref()
.ok_or_else(|| WbtError::NoneValue("report not computed yet".into()))?;
let dt = &report.daily_totals;
let n = dt.totals.len();
let epoch = chrono::NaiveDate::from_ymd_opt(1970, 1, 1).unwrap();
let dr_dates: Vec<i32> = dt
.date_keys
.iter()
.map(|dk| {
let nd = utils::date_key_to_naive_date(*dk);
(nd - epoch).num_days() as i32
})
.collect();
let excess: Vec<f64> = (0..n)
.map(|i| dt.totals[i] - dt.benchmark_means[i])
.collect();
DataFrame::new_infer_height(vec![
Series::new("date".into(), dr_dates)
.cast(&DataType::Date)
.map_err(WbtError::Polars)?
.into_column(),
Series::new("超额".into(), excess).into_column(),
Series::new("策略".into(), &dt.totals).into_column(),
Series::new("基准".into(), &dt.benchmark_means).into_column(),
])
.map_err(WbtError::Polars)
}
}
impl WeightBacktest {
fn validate_input(mut df: DataFrame) -> Result<DataFrame, WbtError> {
for name in ["dt", "symbol", "weight", "price"] {
let column = df
.column(name)
.map_err(|_| WbtError::InvalidInput(format!("Missing required column '{name}'")))?;
if column.null_count() > 0 {
return Err(WbtError::InvalidInput(format!(
"column '{name}' must not contain null values"
)));
}
}
let symbols = df.column("symbol")?.as_materialized_series();
let symbols = symbols
.str()
.map_err(|_| WbtError::InvalidInput("column 'symbol' must contain strings".into()))?;
if symbols.into_no_null_iter().any(|s| s.trim().is_empty()) {
return Err(WbtError::InvalidInput(
"column 'symbol' must not contain empty strings".into(),
));
}
for name in ["weight", "price"] {
let column = df.column(name)?.as_materialized_series();
if !column.dtype().is_primitive_numeric() {
return Err(WbtError::InvalidInput(format!(
"column '{name}' must be numeric, got {:?}",
column.dtype()
)));
}
let values = column.strict_cast(&DataType::Float64).map_err(|e| {
WbtError::InvalidInput(format!("column '{name}' cannot convert to Float64: {e}"))
})?;
if values.f64()?.into_no_null_iter().any(|v| !v.is_finite()) {
return Err(WbtError::InvalidInput(format!(
"column '{name}' must contain only finite values (no NaN or infinity)"
)));
}
df.replace(name, values.into())?;
}
Ok(df)
}
pub(crate) fn unique_symbols(df: &DataFrame) -> Result<Vec<Arc<str>>, WbtError> {
let symbols_series = df.column("symbol")?.as_materialized_series().str()?;
let mut unique_symbols_set = hashbrown::HashSet::new();
for symbol in symbols_series.into_iter().flatten() {
unique_symbols_set.insert(symbol);
}
let mut unique_symbols: Vec<Arc<str>> =
unique_symbols_set.into_iter().map(Arc::from).collect();
unique_symbols.sort_unstable();
Ok(unique_symbols)
}
fn sort_by_dt(df: DataFrame) -> Result<DataFrame, WbtError> {
df.lazy()
.sort(
["dt"],
SortMultipleOptions::default().with_order_descending(false),
)
.collect()
.map_err(|e| anyhow::anyhow!("Failed to sort by dt: {e}").into())
}
pub(crate) fn convert_datetime(mut df: DataFrame) -> Result<DataFrame, WbtError> {
let dt_col = df.column("dt")?.as_materialized_series().clone();
let dt_type = dt_col.dtype().clone();
match &dt_type {
DataType::Datetime(TimeUnit::Nanoseconds, _) => Ok(Self::sort_by_dt(df)?),
DataType::Datetime(TimeUnit::Milliseconds, _) => {
let dt_cast = dt_col.cast(&DataType::Datetime(TimeUnit::Milliseconds, None))?;
let _ = df.replace("dt", dt_cast.into())?;
Ok(Self::sort_by_dt(df)?)
}
DataType::Int64 => {
let parsed_col = dt_col
.i64()?
.into_iter()
.map(|opt_ts| opt_ts.map(|ts| ts * 1000));
let dt_s = Series::from_iter(parsed_col)
.cast(&DataType::Datetime(TimeUnit::Milliseconds, None))?;
let _ = df.replace("dt", dt_s.into())?;
Ok(Self::sort_by_dt(df)?)
}
DataType::String => {
let df = df
.lazy()
.with_column(col("dt").str().to_datetime(
Some(TimeUnit::Milliseconds),
None,
StrptimeOptions {
format: Some("%Y-%m-%d %H:%M:%S".into()),
strict: true,
exact: false,
cache: true,
},
lit("raise"),
))
.sort(
["dt"],
SortMultipleOptions::default().with_order_descending(false),
)
.collect()
.context("Failed to convert datetime")?;
Ok(df)
}
_ => Err(anyhow::anyhow!("Unsupported datetime type: {:?}", dt_type).into()),
}
}
pub(crate) fn round_weight(df: &mut DataFrame, digits: i64) -> Result<(), WbtError> {
let scale = 10_f64.powi(digits as i32);
let weight_s = df.column("weight")?.as_materialized_series().clone();
let rounded = weight_s
.f64()?
.into_iter()
.map(|opt| opt.map(|val| (val * scale).round() / scale))
.collect::<Float64Chunked>();
let _ = df.replace("weight", rounded.into_series().into())?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
fn raw_example_data() -> DataFrame {
df! {
"dt" => &[
"2019-01-02 09:01:00",
"2019-01-03 09:02:00",
"2019-01-04 09:03:00",
"2019-01-05 09:04:00",
"2019-01-06 09:05:00"
],
"symbol" => &["DLi9001"; 5],
"weight" => &[
0.511,
0.000,
-0.250,
0.000,
0.000
],
"price" => &[
961.695,
960.720,
962.669,
960.720,
961.695
]
}
.unwrap()
}
#[test]
fn test_round_weight() {
let mut df = raw_example_data();
WeightBacktest::round_weight(&mut df, 4).unwrap();
let weights: Vec<f64> = df
.column("weight")
.unwrap()
.as_materialized_series()
.f64()
.unwrap()
.into_no_null_iter()
.collect();
assert_eq!(weights, vec![0.511, 0.0, -0.25, 0.0, 0.0]);
}
#[test]
fn test_convert_datetime() {
let df = raw_example_data();
let df = WeightBacktest::convert_datetime(df).unwrap();
assert!(matches!(
df.column("dt").unwrap().dtype(),
DataType::Datetime(_, _)
));
assert_eq!(df.height(), 5);
}
#[test]
fn test_unique_symbols() {
let df = raw_example_data();
let symbols = WeightBacktest::unique_symbols(&df).unwrap();
assert_eq!(symbols, vec![Arc::from("DLi9001")]);
}
#[test]
fn new_valid_dataframe() {
let df = raw_example_data();
let wb = WeightBacktest::new(df, 2, None).unwrap();
assert_eq!(wb.fee_rate, 0.0002);
assert_eq!(wb.digits, 2);
assert!(!wb.symbols.is_empty());
}
#[test]
fn new_rejects_out_of_range_digits() {
for digits in [-1, -50, 11, 400] {
let df = raw_example_data();
match WeightBacktest::new(df, digits, None) {
Ok(_) => panic!("digits={digits} must be rejected"),
Err(err) => assert!(
err.to_string().contains("digits must be in 0..=10"),
"unexpected error for digits={digits}: {err}"
),
}
}
}
#[test]
fn new_normalizes_integer_weights_and_prices() {
let mut df = raw_example_data();
df.replace(
"weight",
Series::new("weight".into(), [1_i64, 0, -1, 1, 0]).into(),
)
.unwrap();
df.replace(
"price",
Series::new("price".into(), [100_i32, 102, 99, 103, 101]).into(),
)
.unwrap();
let mut reference = df.clone();
for name in ["weight", "price"] {
let values = reference
.column(name)
.unwrap()
.as_materialized_series()
.cast(&DataType::Float64)
.unwrap();
reference.replace(name, values.into()).unwrap();
}
let mut wb = WeightBacktest::new(df, 2, None).unwrap();
let mut expected = WeightBacktest::new(reference, 2, None).unwrap();
wb.backtest(Some(1), WeightType::TS, 252).unwrap();
expected.backtest(Some(1), WeightType::TS, 252).unwrap();
assert!(
wb.daily_return_df()
.unwrap()
.equals_missing(expected.daily_return_df().unwrap())
);
for name in ["weight", "price"] {
assert_eq!(wb.dfw.column(name).unwrap().dtype(), &DataType::Float64);
}
}
#[test]
fn new_rejects_missing_and_null_required_columns() {
for name in ["dt", "symbol", "weight", "price"] {
let df = raw_example_data();
let mut null_df = df.clone();
null_df
.replace(
name,
Series::full_null(name.into(), df.height(), df.column(name).unwrap().dtype())
.into(),
)
.unwrap();
for invalid in [df.drop(name).unwrap(), null_df] {
let err = WeightBacktest::new(invalid, 2, None)
.err()
.expect("must reject invalid input");
assert!(matches!(err, WbtError::InvalidInput(_)));
assert!(err.to_string().contains(name));
}
}
}
#[test]
fn new_rejects_partial_null_symbol_instead_of_mapping_to_first_symbol() {
let mut df = raw_example_data();
df.replace(
"symbol",
Series::new(
"symbol".into(),
[Some("A"), None, Some("B"), Some("B"), Some("A")],
)
.into(),
)
.unwrap();
assert!(matches!(
WeightBacktest::new(df, 2, None),
Err(WbtError::InvalidInput(_))
));
}
#[test]
fn new_rejects_non_numeric_weights_and_prices() {
for name in ["weight", "price"] {
let mut df = raw_example_data();
df.replace(name, Series::new(name.into(), ["1"; 5]).into())
.unwrap();
let err = WeightBacktest::new(df, 2, None)
.err()
.expect("must reject strings");
assert!(matches!(err, WbtError::InvalidInput(_)));
assert!(err.to_string().contains(name));
}
}
#[test]
fn new_rejects_non_finite_weights_and_prices() {
for name in ["weight", "price"] {
for value in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY] {
for dtype in [DataType::Float32, DataType::Float64] {
let mut df = raw_example_data();
let values = Series::new(name.into(), [1.0, value, 1.0, 1.0, 1.0])
.cast(&dtype)
.unwrap();
df.replace(name, values.into()).unwrap();
let err = WeightBacktest::new(df, 2, None)
.err()
.expect("must reject non-finite input");
assert!(matches!(err, WbtError::InvalidInput(_)));
assert!(err.to_string().contains(name));
assert!(err.to_string().contains("finite"));
}
}
}
}
#[test]
fn weight_type_ts_cs_aggregate_returns() {
for short in [false, true] {
for (mode, divisor) in [("ts", 2.0), ("cs", 1.0)] {
let b_weight = if short { -0.25 } else { 0.5 };
let df = df! {
"dt" => &[
"2024-01-01 09:00:00", "2024-01-02 09:00:00",
"2024-01-03 09:00:00", "2024-01-04 09:00:00",
"2024-01-01 09:00:00", "2024-01-02 09:00:00",
"2024-01-03 09:00:00", "2024-01-04 09:00:00",
],
"symbol" => &["A", "A", "A", "A", "B", "B", "B", "B"],
"weight" => &[0.5, 0.5, 0.5, 0.5, b_weight, b_weight, b_weight, b_weight],
"price" => &[100.0, 110.0, 99.0, 108.9, 200.0, 240.0, 216.0, 237.6],
}
.unwrap();
let mut wb = WeightBacktest::new(df, 2, Some(0.0)).unwrap();
wb.backtest(Some(1), mode.parse().unwrap(), 252).unwrap();
let totals = wb
.daily_return_df()
.unwrap()
.column("total")
.unwrap()
.f64()
.unwrap();
let expected = if short {
[0.0, -0.025, 0.025]
} else {
[0.15, -0.10, 0.10]
};
assert_eq!(totals.len(), expected.len());
for (actual, expected) in totals.into_no_null_iter().zip(expected) {
assert!(
(actual - expected / divisor).abs() < 1e-12,
"{mode}: {actual}"
);
}
}
}
}
#[test]
fn new_custom_fee_rate() {
let df = raw_example_data();
let wb = WeightBacktest::new(df, 2, Some(0.001)).unwrap();
assert_eq!(wb.fee_rate, 0.001);
}
#[test]
fn yearly_return_df_end_to_end_minimal() {
let df = raw_example_data();
let mut wb = WeightBacktest::new(df, 2, None).unwrap();
wb.backtest(Some(1), WeightType::TS, 252).unwrap();
let y = wb.yearly_return_df(1).unwrap();
assert_eq!(y.height(), 2);
let years: Vec<i32> = y
.column("year")
.unwrap()
.as_materialized_series()
.i32()
.unwrap()
.into_no_null_iter()
.collect();
assert_eq!(years, vec![2019, 2019]);
let syms: Vec<String> = y
.column("symbol")
.unwrap()
.as_materialized_series()
.str()
.unwrap()
.into_no_null_iter()
.map(|s: &str| s.to_string())
.collect();
assert_eq!(syms, vec!["DLi9001".to_string(), "total".to_string()]);
}
#[test]
fn yearly_return_df_filters_when_below_min_days() {
let df = raw_example_data();
let mut wb = WeightBacktest::new(df, 2, None).unwrap();
wb.backtest(Some(1), WeightType::TS, 252).unwrap();
let y = wb.yearly_return_df(120).unwrap();
assert_eq!(y.height(), 0);
}
#[test]
fn explicit_recompute_invalidates_all_materialized_caches() {
let mut wb = WeightBacktest::new(raw_example_data(), 2, None).unwrap();
wb.backtest(Some(1), WeightType::TS, 252).unwrap();
let detached = wb.daily_return_df().unwrap().clone().clear();
assert_eq!(detached.height(), 0);
assert!(wb.daily_return_df().unwrap().height() > 0);
wb.dailys_df().unwrap();
wb.pairs_df().unwrap();
wb.aggregated_pairs_df().unwrap();
wb.backtest(Some(1), WeightType::CS, 365).unwrap();
assert!(wb.daily_return_cache.is_none());
assert!(wb.dailys_cache.is_none());
assert!(wb.pairs_cache.is_none());
assert!(wb.agg_pairs_cache.is_none());
assert_eq!(wb.yearly_days(), 365);
let mut fresh = WeightBacktest::new(raw_example_data(), 2, None).unwrap();
fresh.backtest(Some(1), WeightType::CS, 365).unwrap();
assert!(
wb.daily_return_df()
.unwrap()
.equals_missing(fresh.daily_return_df().unwrap())
);
assert_eq!(
wb.report().unwrap().stats.daily_performance,
fresh.report().unwrap().stats.daily_performance
);
}
#[test]
fn daily_return_cache_is_lazy_and_reused() {
let df = raw_example_data();
let mut wb = WeightBacktest::new(df, 2, None).unwrap();
wb.backtest(Some(1), WeightType::TS, 252).unwrap();
assert!(wb.daily_return_cache.is_none());
let first_ptr = {
let df = wb.daily_return_df().unwrap();
df as *const DataFrame
};
assert!(wb.daily_return_cache.is_some());
let second_ptr = {
let df = wb.daily_return_df().unwrap();
df as *const DataFrame
};
assert_eq!(first_ptr, second_ptr);
}
#[test]
fn new_missing_column() {
let df = df! {
"dt" => &["2019-01-02 09:01:00"],
"symbol" => &["A"],
"weight" => &[0.5_f64]
}
.unwrap();
assert!(WeightBacktest::new(df, 2, None).is_err());
}
#[test]
fn convert_datetime_int64() {
let df = df! {
"dt" => &[1546398060_i64, 1546484520_i64],
"symbol" => &["A", "A"],
"weight" => &[0.5_f64, -0.5],
"price" => &[100.0, 101.0]
}
.unwrap();
let result = WeightBacktest::convert_datetime(df);
assert!(result.is_ok());
let df = result.unwrap();
assert!(matches!(
df.column("dt").unwrap().dtype(),
DataType::Datetime(_, _)
));
}
#[test]
fn round_weight_precision() {
let mut df = df! {
"dt" => &["2019-01-02 09:01:00"],
"symbol" => &["A"],
"weight" => &[0.12345678_f64],
"price" => &[100.0]
}
.unwrap();
WeightBacktest::round_weight(&mut df, 4).unwrap();
let w = df
.column("weight")
.unwrap()
.as_materialized_series()
.f64()
.unwrap()
.get(0)
.unwrap();
assert_eq!(w, 0.1235);
}
#[test]
fn round_weight_zero() {
let mut df = df! {
"dt" => &["2019-01-02 09:01:00"],
"symbol" => &["A"],
"weight" => &[0.0_f64],
"price" => &[100.0]
}
.unwrap();
WeightBacktest::round_weight(&mut df, 4).unwrap();
let w = df
.column("weight")
.unwrap()
.as_materialized_series()
.f64()
.unwrap()
.get(0)
.unwrap();
assert_eq!(w, 0.0);
}
#[test]
fn new_normalizes_weight_with_digits() {
let df = df! {
"dt" => &["2019-01-02 09:01:00"],
"symbol" => &["A"],
"weight" => &[0.125_f64],
"price" => &[100.0_f64]
}
.unwrap();
let wb = WeightBacktest::new(df, 2, None).unwrap();
let weight = wb
.dfw
.column("weight")
.unwrap()
.as_materialized_series()
.f64()
.unwrap()
.get(0)
.unwrap();
assert_eq!(weight, 0.13);
}
#[test]
fn from_file_csv() {
let dir = std::env::temp_dir().join("wbt_test_from_file");
std::fs::create_dir_all(&dir).unwrap();
let csv_path = dir.join("test.csv");
let csv_content = "dt,symbol,weight,price\n\
2024-01-01 09:30:00,SYM_A,0.5,100.0\n\
2024-01-02 09:30:00,SYM_A,-0.3,101.0\n\
2024-01-01 09:30:00,SYM_B,0.2,50.0\n\
2024-01-02 09:30:00,SYM_B,0.0,51.0\n";
std::fs::write(&csv_path, csv_content).unwrap();
let wb = WeightBacktest::from_file(csv_path.to_str().unwrap(), 2, None).unwrap();
assert_eq!(wb.symbols.len(), 2);
assert!(wb.symbols.contains(&std::sync::Arc::from("SYM_A")));
assert!(wb.symbols.contains(&std::sync::Arc::from("SYM_B")));
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn from_file_missing_column() {
let dir = std::env::temp_dir().join("wbt_test_missing_col");
std::fs::create_dir_all(&dir).unwrap();
let csv_path = dir.join("bad.csv");
std::fs::write(&csv_path, "dt,symbol,weight\n2024-01-01,A,0.5\n").unwrap();
let result = WeightBacktest::from_file(csv_path.to_str().unwrap(), 2, None);
assert!(result.is_err());
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn from_file_unsupported_ext() {
let result = WeightBacktest::from_file("/tmp/test.xlsx", 2, None);
assert!(result.is_err());
let err_msg = match result {
Err(e) => e.to_string(),
Ok(_) => unreachable!(),
};
assert!(err_msg.contains("Unsupported"));
}
#[test]
fn from_file_parquet() {
let dir = std::env::temp_dir().join("wbt_test_from_file_parquet");
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("test.parquet");
let df = df! {
"dt" => &[
"2024-01-01 09:30:00",
"2024-01-02 09:30:00",
"2024-01-01 09:30:00",
"2024-01-02 09:30:00",
],
"symbol" => &["SYM_A", "SYM_A", "SYM_B", "SYM_B"],
"weight" => &[0.5_f64, -0.3, 0.2, 0.0],
"price" => &[100.0_f64, 101.0, 50.0, 51.0]
}
.unwrap();
let file = std::fs::File::create(&path).unwrap();
ParquetWriter::new(file).finish(&mut df.clone()).unwrap();
let wb = WeightBacktest::from_file(path.to_str().unwrap(), 2, None).unwrap();
assert_eq!(wb.symbols.len(), 2);
assert!(wb.symbols.contains(&std::sync::Arc::from("SYM_A")));
assert!(wb.symbols.contains(&std::sync::Arc::from("SYM_B")));
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn from_file_feather() {
let dir = std::env::temp_dir().join("wbt_test_from_file_feather");
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("test.feather");
let df = df! {
"dt" => &[
"2024-01-01 09:30:00",
"2024-01-02 09:30:00",
"2024-01-01 09:30:00",
"2024-01-02 09:30:00",
],
"symbol" => &["SYM_A", "SYM_A", "SYM_B", "SYM_B"],
"weight" => &[0.5_f64, -0.3, 0.2, 0.0],
"price" => &[100.0_f64, 101.0, 50.0, 51.0]
}
.unwrap();
let file = std::fs::File::create(&path).unwrap();
IpcWriter::new(file).finish(&mut df.clone()).unwrap();
let wb = WeightBacktest::from_file(path.to_str().unwrap(), 2, None).unwrap();
assert_eq!(wb.symbols.len(), 2);
assert!(wb.symbols.contains(&std::sync::Arc::from("SYM_A")));
assert!(wb.symbols.contains(&std::sync::Arc::from("SYM_B")));
std::fs::remove_dir_all(&dir).ok();
}
#[test]
fn unique_symbols_sorted_order() {
let df = df! {
"dt" => &["2019-01-02", "2019-01-02", "2019-01-02"],
"symbol" => &["C", "A", "B"],
"weight" => &[0.1, 0.2, 0.3],
"price" => &[1.0, 2.0, 3.0]
}
.unwrap();
let syms = WeightBacktest::unique_symbols(&df).unwrap();
assert_eq!(syms, vec![Arc::from("A"), Arc::from("B"), Arc::from("C")]);
}
#[test]
fn convert_datetime_nanoseconds_passthrough() {
let dates: Vec<i64> = vec![
1_704_067_200_000_000_000, 1_704_153_600_000_000_000, ];
let dt_series = Series::new("dt".into(), dates)
.cast(&DataType::Datetime(TimeUnit::Nanoseconds, None))
.unwrap();
let df = DataFrame::new_infer_height(vec![
dt_series.into_column(),
Series::new("symbol".into(), &["A", "A"]).into_column(),
Series::new("weight".into(), &[0.5_f64, 0.0]).into_column(),
Series::new("price".into(), &[100.0_f64, 101.0]).into_column(),
])
.unwrap();
let result = WeightBacktest::convert_datetime(df);
assert!(result.is_ok());
let df = result.unwrap();
assert!(matches!(
df.column("dt").unwrap().dtype(),
DataType::Datetime(TimeUnit::Nanoseconds, _)
));
assert_eq!(df.height(), 2);
}
}