1use crate::{BacktestConfig, BacktestEngine, BacktestError, BacktestReport};
7use polars::prelude::{RankMethod, RankOptions, *};
8
9#[derive(Debug, Clone, PartialEq)]
11pub struct CrossSectionalConfig {
12 pub factor_col: String,
13 pub top_frac: f64,
15 pub bottom_frac: f64,
17}
18
19impl CrossSectionalConfig {
20 pub fn long_short(factor_col: impl Into<String>, top_frac: f64, bottom_frac: f64) -> Self {
21 Self {
22 factor_col: factor_col.into(),
23 top_frac,
24 bottom_frac,
25 }
26 }
27}
28
29pub fn neutralize_factor(lf: LazyFrame, factor_col: &str, group_col: &str) -> LazyFrame {
31 lf.with_column(
32 (col(factor_col) - col(factor_col).mean().over([col(group_col)])).alias(factor_col),
33 )
34}
35
36pub fn zscore_factor(lf: LazyFrame, factor_col: &str, timestamp_col: &str) -> LazyFrame {
38 let mean = col(factor_col).mean().over([col(timestamp_col)]);
39 let std = col(factor_col).std(1).over([col(timestamp_col)]);
40 lf.with_column(((col(factor_col) - mean) / std).alias(factor_col))
41}
42
43pub fn winsorize_factor(
45 lf: LazyFrame,
46 factor_col: &str,
47 timestamp_col: &str,
48 lower_pct: f64,
49 upper_pct: f64,
50) -> LazyFrame {
51 let lower = col(factor_col)
52 .quantile(lit(lower_pct), QuantileMethod::Nearest)
53 .over([col(timestamp_col)]);
54 let upper = col(factor_col)
55 .quantile(lit(upper_pct), QuantileMethod::Nearest)
56 .over([col(timestamp_col)]);
57
58 lf.with_column(
59 when(col(factor_col).lt(lower.clone()))
60 .then(lower)
61 .when(col(factor_col).gt(upper.clone()))
62 .then(upper)
63 .otherwise(col(factor_col))
64 .alias(factor_col),
65 )
66}
67
68pub fn assign_long_short_exposure(
72 lf: LazyFrame,
73 timestamp_col: &str,
74 _symbol_col: &str,
75 cs: &CrossSectionalConfig,
76 exposure_col: &str,
77) -> Result<LazyFrame, BacktestError> {
78 if cs.top_frac <= 0.0 || cs.bottom_frac <= 0.0 {
79 return Err(BacktestError::InvalidInput(
80 "top_frac and bottom_frac must be > 0".into(),
81 ));
82 }
83 if cs.top_frac + cs.bottom_frac > 1.0 {
84 return Err(BacktestError::InvalidInput(
85 "top_frac + bottom_frac must be <= 1".into(),
86 ));
87 }
88
89 let n_per_ts = col(timestamp_col)
90 .count()
91 .over([col(timestamp_col)])
92 .cast(DataType::Float64);
93
94 let rank_best = col(&cs.factor_col)
96 .rank(
97 RankOptions {
98 method: RankMethod::Min,
99 descending: true,
100 },
101 None,
102 )
103 .over([col(timestamp_col)])
104 .cast(DataType::Float64);
105
106 let top_slots = when(
107 (n_per_ts.clone() * lit(cs.top_frac))
108 .cast(DataType::Int64)
109 .cast(DataType::Float64)
110 .lt(lit(1.0)),
111 )
112 .then(lit(1.0))
113 .otherwise(
114 (n_per_ts.clone() * lit(cs.top_frac))
115 .cast(DataType::Int64)
116 .cast(DataType::Float64),
117 );
118 let bottom_slots = when(
119 (n_per_ts.clone() * lit(cs.bottom_frac))
120 .cast(DataType::Int64)
121 .cast(DataType::Float64)
122 .lt(lit(1.0)),
123 )
124 .then(lit(1.0))
125 .otherwise(
126 (n_per_ts.clone() * lit(cs.bottom_frac))
127 .cast(DataType::Int64)
128 .cast(DataType::Float64),
129 );
130 let short_cut = n_per_ts - bottom_slots.clone() + lit(1.0);
131
132 let exposure = when(rank_best.clone().lt_eq(top_slots.clone()))
133 .then(lit(1.0) / top_slots.clone())
134 .when(rank_best.gt_eq(short_cut))
135 .then(lit(-1.0) / bottom_slots)
136 .otherwise(lit(0.0))
137 .alias(exposure_col);
138
139 Ok(lf.with_column(exposure))
140}
141
142pub fn run_cross_sectional_backtest(
144 lf: LazyFrame,
145 cs: &CrossSectionalConfig,
146 mut base_config: BacktestConfig,
147) -> Result<BacktestReport, BacktestError> {
148 let symbol_col = base_config.symbol_col.clone().ok_or_else(|| {
149 BacktestError::InvalidInput("symbol_col required for cross-sectional".into())
150 })?;
151
152 const EXPOSURE: &str = "cs_exposure";
153 let with_exp =
154 assign_long_short_exposure(lf, &base_config.timestamp_col, &symbol_col, cs, EXPOSURE)?;
155 base_config.signal_col = EXPOSURE.to_string();
156 BacktestEngine::new(base_config).backtest_with_report(with_exp)
157}
158
159#[cfg(test)]
160#[allow(clippy::panic)]
161mod tests {
162 use super::*;
163 use approx::assert_relative_eq;
164
165 fn panel_df() -> DataFrame {
166 let timestamps = vec![1i64, 1, 1, 1, 2, 2, 2, 2];
168 let symbols = vec!["A", "B", "C", "D", "A", "B", "C", "D"];
169 let closes = vec![10.0, 10.0, 10.0, 10.0, 11.0, 11.0, 11.0, 11.0];
170 let factor = vec![4.0, 3.0, 2.0, 1.0, 4.0, 3.0, 2.0, 1.0];
171 DataFrame::new(vec![
172 Column::new("timestamp".into(), timestamps),
173 Column::new("symbol".into(), symbols),
174 Column::new("close".into(), closes),
175 Column::new("score".into(), factor),
176 ])
177 .unwrap()
178 }
179
180 #[test]
181 fn test_factor_neutralize_demean_within_sector() {
182 let df = DataFrame::new(vec![
183 Column::new("sector".into(), vec!["Tech", "Tech", "Fin", "Fin"]),
184 Column::new("score".into(), vec![10.0, 20.0, 100.0, 200.0]),
185 ])
186 .unwrap();
187
188 let out = neutralize_factor(df.lazy(), "score", "sector")
189 .collect()
190 .unwrap();
191 let scores = out.column("score").unwrap().f64().unwrap();
192
193 let ts: Vec<f64> = scores.into_iter().map(|v| v.unwrap()).collect();
194 assert_relative_eq!(ts[0], -5.0, epsilon = 1e-9);
195 assert_relative_eq!(ts[1], 5.0, epsilon = 1e-9);
196 assert_relative_eq!(ts[2], -50.0, epsilon = 1e-9);
197 assert_relative_eq!(ts[3], 50.0, epsilon = 1e-9);
198 }
199
200 #[test]
201 fn test_factor_zscore_zero_mean_smoke() {
202 let df = panel_df();
203 let out = zscore_factor(df.lazy(), "score", "timestamp")
204 .collect()
205 .unwrap();
206 let scores = out.column("score").unwrap().f64().unwrap();
207 let ts: Vec<f64> = scores.into_iter().map(|v| v.unwrap()).collect();
208
209 let mean_1 = (ts[0] + ts[1] + ts[2] + ts[3]) / 4.0;
212 assert_relative_eq!(mean_1, 0.0, epsilon = 1e-9);
213 }
214
215 #[test]
216 fn test_factor_winsorize_clips_extremes() {
217 let df = DataFrame::new(vec![
218 Column::new("timestamp".into(), vec![1i64, 1, 1, 1, 1]),
219 Column::new("score".into(), vec![0.0, 10.0, 20.0, 30.0, 100.0]),
220 ])
221 .unwrap();
222
223 let out = winsorize_factor(df.lazy(), "score", "timestamp", 0.2, 0.8)
224 .collect()
225 .unwrap();
226 let scores = out.column("score").unwrap().f64().unwrap();
227 let ts: Vec<f64> = scores.into_iter().map(|v| v.unwrap()).collect();
228
229 assert_relative_eq!(ts[0], 10.0, epsilon = 1e-9); assert_relative_eq!(ts[1], 10.0, epsilon = 1e-9);
234 assert_relative_eq!(ts[4], 30.0, epsilon = 1e-9); }
236
237 #[test]
238 fn test_assign_long_short_exposure_top_bottom() {
239 let cs = CrossSectionalConfig::long_short("score", 0.25, 0.25);
240 let out =
241 assign_long_short_exposure(panel_df().lazy(), "timestamp", "symbol", &cs, "exposure")
242 .unwrap()
243 .collect()
244 .unwrap();
245
246 let exposure = out.column("exposure").unwrap().f64().unwrap();
247 let ts1: Vec<f64> = exposure.into_iter().take(4).map(|v| v.unwrap()).collect();
249 assert_eq!(ts1.iter().filter(|&&x| x > 0.0).count(), 1);
250 assert_eq!(ts1.iter().filter(|&&x| x < 0.0).count(), 1);
251 assert_relative_eq!(
252 ts1.iter().map(|x| x.abs()).sum::<f64>(),
253 2.0,
254 epsilon = 1e-9
255 );
256 }
257
258 #[test]
259 fn test_cross_sectional_backtest_smoke() {
260 let cs = CrossSectionalConfig::long_short("score", 0.25, 0.25);
261 let cfg = BacktestConfig {
262 cost_model: crate::CostModel {
263 commission_bps: 0.0,
264 slippage_bps: 0.0,
265 initial_cash: 100_000.0,
266 },
267 symbol_col: Some("symbol".into()),
268 ..Default::default()
269 };
270 let mut df = panel_df();
272 df = df
273 .lazy()
274 .with_column(lit(1.0).alias("score"))
275 .collect()
276 .unwrap();
277
278 let report = run_cross_sectional_backtest(df.lazy(), &cs, cfg.clone()).unwrap();
279 assert!(report.metrics.final_equity.is_finite());
280 }
281
282 #[test]
283 fn test_cross_sectional_invalid_fracs_error() {
284 let cs = CrossSectionalConfig::long_short("score", 0.6, 0.6);
285 let result =
286 assign_long_short_exposure(panel_df().lazy(), "timestamp", "symbol", &cs, "exposure");
287 let Err(err) = result else {
288 panic!("expected invalid frac error");
289 };
290 assert!(
291 err.to_string().contains("top_frac"),
292 "unexpected error message"
293 );
294 }
295}