feat: add yahoo module

stochastic-rs-quant/Cargo.toml

@@ -5,20 +5,25 @@ edition = "2021"
 
 [dependencies]
 chrono = "0.4.38"
-derive_builder = "0.20.1"
+dotenvy = "0.15.7"
 levenberg-marquardt = "0.14.0"
 mimalloc = { version = "0.1.43", optional = true }
 nalgebra = "0.33.0"
 num-complex = "0.4.6"
+polars = { version = "0.43.1", features = ["lazy"] }
 quadrature = "0.1.2"
 rand = "0.8.5"
 rand_distr = "0.4.3"
 stochastic-rs = { path = "../stochastic-rs-core" }
 tikv-jemallocator = { version = "0.6.0", optional = true }
+time = { version = "0.3.36", features = ["formatting", "parsing"] }
+tokio-test = "0.4.4"
+yahoo_finance_api = "2.2.1"
 
 [features]
 mimalloc = ["dep:mimalloc"]
 jemalloc = ["dep:tikv-jemallocator"]
+yahoo = []
 default = []
 
 [lib]

stochastic-rs-quant/src/lib.rs

@@ -5,6 +5,8 @@ use std::mem::ManuallyDrop;
 pub mod calibrator;
 pub mod heston;
 pub mod pricer;
+#[cfg(feature = "yahoo")]
+pub mod yahoo;
 
 #[cfg(feature = "mimalloc")]
 #[global_allocator]

stochastic-rs-quant/src/yahoo.rs

@@ -0,0 +1,237 @@
+use std::{borrow::Cow, fmt::Display};
+
+use polars::prelude::*;
+use time::OffsetDateTime;
+use tokio_test;
+use yahoo_finance_api::YahooConnector;
+
+/// Yahoo struct
+pub struct Yahoo<'a> {
+  /// YahooConnector
+  pub(crate) provider: YahooConnector,
+  /// Symbol
+  pub(crate) symbol: Option<Cow<'a, str>>,
+  /// Start date
+  pub(crate) start_date: Option<OffsetDateTime>,
+  /// End date
+  pub(crate) end_date: Option<OffsetDateTime>,
+  /// Options
+  pub options: Option<DataFrame>,
+  /// Price history
+  pub price_history: Option<DataFrame>,
+  /// Returns
+  pub returns: Option<DataFrame>,
+}
+
+pub enum ReturnType {
+  Arithmetic,
+  Logarithmic,
+  Absolute,
+}
+
+impl Display for ReturnType {
+  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+    match self {
+      ReturnType::Arithmetic => write!(f, "arithmetic"),
+      ReturnType::Logarithmic => write!(f, "logarithmic"),
+      ReturnType::Absolute => write!(f, "absolute"),
+    }
+  }
+}
+
+impl<'a> Default for Yahoo<'a> {
+  #[must_use]
+  fn default() -> Self {
+    Self {
+      provider: YahooConnector::new().unwrap(),
+      symbol: None,
+      start_date: Some(OffsetDateTime::UNIX_EPOCH),
+      end_date: Some(OffsetDateTime::now_utc()),
+      options: None,
+      price_history: None,
+      returns: None,
+    }
+  }
+}
+
+impl<'a> Yahoo<'a> {
+  /// Set symbol
+  pub fn set_symbol(&mut self, symbol: &'a str) {
+    self.symbol = Some(Cow::Borrowed(&symbol));
+  }
+
+  /// Set start date
+  pub fn set_start_date(&mut self, start_date: OffsetDateTime) {
+    self.start_date = Some(start_date);
+  }
+
+  /// Set end date
+  pub fn set_end_date(&mut self, end_date: OffsetDateTime) {
+    self.end_date = Some(end_date);
+  }
+
+  /// Get price history for symbol
+  pub fn get_price_history(&mut self) {
+    let res = tokio_test::block_on(self.provider.get_quote_history(
+      self.symbol.as_deref().unwrap(),
+      self.start_date.unwrap(),
+      self.end_date.unwrap(),
+    ))
+    .unwrap();
+
+    let history = res.quotes().unwrap();
+    let df = df!(
+        "timestamp" => Series::new("timestamp".into(), &history.iter().map(|r| r.timestamp / 86_400).collect::<Vec<_>>()).cast(&DataType::Date).unwrap(),
+        "volume" => &history.iter().map(|r| r.volume).collect::<Vec<_>>(),
+        "open" => &history.iter().map(|r| r.open).collect::<Vec<_>>(),
+        "high" => &history.iter().map(|r| r.high).collect::<Vec<_>>(),
+        "low" => &history.iter().map(|r| r.low).collect::<Vec<_>>(),
+        "close" => &history.iter().map(|r| r.close).collect::<Vec<_>>(),
+        "adjclose" => &history.iter().map(|r| r.adjclose).collect::<Vec<_>>(),
+    )
+    .unwrap();
+
+    self.price_history = Some(df);
+  }
+
+  /// Get options for symbol
+  pub fn get_options_chain(&mut self) {
+    let res = tokio_test::block_on(
+      self
+        .provider
+        .search_options(self.symbol.as_deref().unwrap()),
+    )
+    .unwrap();
+
+    let df = df!(
+        "name" => &res.options.iter().map(|o| o.name.clone()).collect::<Vec<_>>(),
+        "strike" => &res.options.iter().map(|o| o.strike).collect::<Vec<_>>(),
+        "last_trade_date" => &res.options.iter().map(|o| o.last_trade_date.clone()).collect::<Vec<_>>(),
+        "last_price" => &res.options.iter().map(|o| o.last_price).collect::<Vec<_>>(),
+        "bid" => &res.options.iter().map(|o| o.bid).collect::<Vec<_>>(),
+        "ask" => &res.options.iter().map(|o| o.ask).collect::<Vec<_>>(),
+        "change" => &res.options.iter().map(|o| o.change).collect::<Vec<_>>(),
+        "change_pct" => &res.options.iter().map(|o| o.change_pct).collect::<Vec<_>>(),
+        "volume" => &res.options.iter().map(|o| o.volume as u64).collect::<Vec<_>>(),
+        "open_interest" => &res.options.iter().map(|o| o.open_interest as u64).collect::<Vec<_>>(),
+        "impl_volatility" => &res.options.iter().map(|o| o.impl_volatility).collect::<Vec<_>>(),
+    ).unwrap();
+
+    self.options = Some(df);
+  }
+
+  /// Get returns for symbol
+  pub fn get_returns(&mut self, r#type: ReturnType) {
+    if self.price_history.is_none() {
+      self.get_price_history();
+    }
+
+    let cols = || col("*").exclude(["timestamp", "volume"]);
+    let df = match r#type {
+      ReturnType::Arithmetic => self
+        .price_history
+        .as_ref()
+        .unwrap()
+        .clone()
+        .lazy()
+        .select(&[
+          col("timestamp"),
+          col("volume"),
+          (cols() / cols().shift(lit(1)) - lit(1))
+            .name()
+            .suffix(&format!("_{}", &r#type)),
+        ])
+        .collect()
+        .unwrap(),
+      ReturnType::Absolute => self
+        .price_history
+        .as_ref()
+        .unwrap()
+        .clone()
+        .lazy()
+        .select(&[
+          col("timestamp"),
+          col("volume"),
+          (cols() / cols().shift(lit(1)))
+            .name()
+            .suffix(&format!("_{}", &r#type)),
+        ])
+        .collect()
+        .unwrap(),
+      ReturnType::Logarithmic => {
+        let ln = |col: &Series| -> Series {
+          col
+            .f64()
+            .unwrap()
+            .apply(|v| Some(v.unwrap().ln()))
+            .into_series()
+        };
+
+        let mut price_history = self.price_history.as_ref().unwrap().clone();
+        price_history.apply("open", ln).unwrap();
+        price_history.apply("high", ln).unwrap();
+        price_history.apply("low", ln).unwrap();
+        price_history.apply("close", ln).unwrap();
+        price_history.apply("adjclose", ln).unwrap();
+
+        price_history
+          .lazy()
+          .select(&[
+            col("timestamp"),
+            col("volume"),
+            (cols() - cols().shift(lit(1)))
+              .name()
+              .suffix(&format!("_{}", &r#type)),
+          ])
+          .collect()
+          .unwrap()
+      }
+    };
+
+    self.returns = Some(df);
+  }
+}
+
+#[cfg(test)]
+mod tests {
+  use super::*;
+
+  #[test]
+  fn test_yahoo_get_price_history() {
+    let mut yahoo = Yahoo::default();
+    yahoo.set_symbol("AAPL");
+    yahoo.get_price_history();
+    println!("{:?}", yahoo.price_history);
+    assert!(yahoo.price_history.is_some());
+  }
+
+  #[test]
+  fn test_yahoo_get_options_chain() {
+    let mut yahoo = Yahoo::default();
+    yahoo.set_symbol("AAPL");
+    yahoo.get_options_chain();
+    println!("{:?}", yahoo.options);
+    assert!(yahoo.options.is_some());
+  }
+
+  #[test]
+  fn test_yahoo_get_returns() {
+    let mut yahoo = Yahoo::default();
+    yahoo.set_symbol("AAPL");
+    yahoo.get_returns(ReturnType::Arithmetic);
+    println!("{:?}", yahoo.returns);
+    assert!(yahoo.returns.is_some());
+
+    let mut yahoo = Yahoo::default();
+    yahoo.set_symbol("AAPL");
+    yahoo.get_returns(ReturnType::Logarithmic);
+    println!("{:?}", yahoo.returns);
+    assert!(yahoo.returns.is_some());
+
+    let mut yahoo = Yahoo::default();
+    yahoo.set_symbol("AAPL");
+    yahoo.get_returns(ReturnType::Absolute);
+    println!("{:?}", yahoo.returns);
+    assert!(yahoo.returns.is_some());
+  }
+}