use crate::common::{validate_inputs, validate_options};
pub use crate::indicator_types::TIndicatorState;
use crate::types::{DisplayGroup, DisplayType, IndicatorError, IndicatorType, Info};
use serde::{Deserialize, Serialize};
use std::f64::consts::PI;
use std::simd::{num::SimdFloat, Simd, StdFloat};
include!(concat!(env!("OUT_DIR"), "/msw_twiddles.rs"));
#[inline(always)]
pub fn twiddles_for_period(period: usize) -> (&'static [f64], &'static [f64]) {
debug_assert!((6..=50).contains(&period), "period {period} outside [6, 50]");
let idx = period.saturating_sub(6).min(44);
(&MSW_COS_TWIDDLES[idx][..period], &MSW_SIN_TWIDDLES[idx][..period])
}
pub const INPUTS_WIDTH: usize = 1;
pub const OPTIONS_WIDTH: usize = 1;
#[cfg(feature = "simd_assets")]
pub use crate::indicators::simd_indicators::msw_simd::indicator_by_assets;
#[cfg(feature = "simd_options")]
pub use crate::indicators::simd_indicators::msw_simd::indicator_by_options;
#[cfg(feature = "simd_assets")]
pub mod by_assets {
pub use crate::indicators::simd_indicators::msw_simd::indicator_by_assets as indicator;
}
#[cfg(feature = "simd_options")]
pub mod by_options {
pub use crate::indicators::simd_indicators::msw_simd::indicator_by_options as indicator;
}
const TPI: f64 = PI * 2.0;
const HPI: f64 = PI * 0.5;
const INV_SQRT2: f64 = std::f64::consts::FRAC_1_SQRT_2;
pub const INFO: Info = Info {
name: "msw",
full_name: "Mesa Sine Wave",
indicator_type: IndicatorType::Cycle,
inputs: &["real"],
options: &["period"],
outputs: &["msw_sine", "msw_lead"],
optional_outputs: &[],
display_groups: &[DisplayGroup {
offset: None,
id: "msw",
label: "MSW",
display_type: DisplayType::Indicator,
outputs: &["msw_sine", "msw_lead"],
}],
};
#[derive(Serialize, Deserialize)]
pub struct State {
pub rp: f64,
pub ip: f64,
pub wr: f64,
pub wi: f64,
}
impl State {
pub fn new(multiplier: f64) -> Self {
let angle = TPI * multiplier;
let (wi, wr) = angle.sin_cos();
Self {
rp: 0.0,
ip: 0.0,
wr,
wi,
}
}
}
pub struct MSWConstants<const N: usize>;
impl<const N: usize> MSWConstants<N> {
pub const TPI: Simd<f64, N> = Simd::splat(TPI);
}
#[derive(Serialize, Deserialize)]
pub struct IndicatorState {
state: State,
real: Vec<f64>,
period: usize,
cos_twiddles: Vec<f64>,
sin_twiddles: Vec<f64>,
}
impl IndicatorState {
pub fn new(real: &[f64], period: usize, multiplier: f64) -> Self {
let (cos_twiddles, sin_twiddles) = precompute_twiddles(period, multiplier);
let (rp, ip) = match period {
0..=7 => {
dot_product_simd::<4>(&real[real.len() - period..], &cos_twiddles, &sin_twiddles)
}
_ => dot_product_simd::<8>(&real[real.len() - period..], &cos_twiddles, &sin_twiddles),
};
Self {
state: State {
rp,
ip,
..State::new(multiplier)
},
real: real[real.len() - period..].to_vec(),
period,
cos_twiddles,
sin_twiddles,
}
}
}
impl TIndicatorState<1> for IndicatorState {
fn batch_indicator(
&mut self,
inputs: &[&[f64]; INPUTS_WIDTH],
_optional_outputs: Option<&[bool]>,
) -> Result<Vec<Vec<f64>>, IndicatorError> {
validate_inputs(inputs, 1)?;
let new_data = inputs[0];
self.real.extend_from_slice(new_data);
let n = new_data.len();
let (mut sine_line, mut lead_line) =
(crate::uninit_vec!(f64, n), crate::uninit_vec!(f64, n));
cycle_sdft(
&self.real,
self.period,
&mut self.state,
&mut sine_line,
&mut lead_line,
);
self.real.drain(..self.real.len() - self.period);
Ok(vec![sine_line, lead_line])
}
}
pub fn min_data(options: &[f64]) -> usize {
options[0] as usize + 1
}
pub fn output_length(data_len: usize, options: &[f64]) -> usize {
data_len - min_data(options) + 1
}
pub fn multiplier(period: usize) -> f64 {
1.0 / period as f64
}
pub fn indicator(
inputs: &[&[f64]; INPUTS_WIDTH],
options: &[f64; OPTIONS_WIDTH],
_optional_outputs: Option<&[bool]>,
) -> Result<(Vec<Vec<f64>>, IndicatorState), IndicatorError> {
validate_options(options)?;
let period = options[0] as usize;
let mult = multiplier(period);
validate_inputs(inputs, min_data(options))?;
let real = inputs[0];
let capacity = output_length(real.len(), options);
let mut sine_line = crate::uninit_vec!(f64, capacity);
let mut lead_line = crate::uninit_vec!(f64, capacity);
let (cos_twiddles, sin_twiddles) = precompute_twiddles(period, mult);
let state = match period {
0..=7 => cycle_msw_sdft::<4>(
real,
period,
mult,
&cos_twiddles,
&sin_twiddles,
&mut sine_line,
&mut lead_line,
),
_ => cycle_msw_sdft::<8>(
real,
period,
mult,
&cos_twiddles,
&sin_twiddles,
&mut sine_line,
&mut lead_line,
),
};
Ok((
vec![sine_line, lead_line],
IndicatorState {
state,
real: real[real.len() - period..].to_vec(),
period,
cos_twiddles,
sin_twiddles,
},
))
}
#[inline(always)]
pub fn calc(state: &mut State, new_sample: f64, old_sample: f64) -> (f64, f64) {
let rp = state.wr.mul_add(state.rp, -(state.wi * state.ip)) + (new_sample - old_sample);
let ip = state.wr.mul_add(state.ip, state.wi * state.rp);
state.rp = rp;
state.ip = ip;
phase_from_rp_ip(rp, ip)
}
#[inline(always)]
pub fn calc_rp_ip<const N: usize>(window: &[f64], multiplier: f64) -> (f64, f64) {
let (cos_twiddles, sin_twiddles) = precompute_twiddles(window.len(), multiplier);
dot_product_simd::<N>(window, &cos_twiddles, &sin_twiddles)
}
#[inline(always)]
pub fn calc_full<const N: usize>(window: &[f64], multiplier: f64) -> (f64, f64) {
let (rp, ip) = calc_rp_ip::<N>(window, multiplier);
phase_from_rp_ip(rp, ip)
}
#[inline(always)]
fn cycle_sdft(
real: &[f64],
period: usize,
state: &mut State,
sine_line: &mut [f64],
lead_line: &mut [f64],
) {
for i in period..real.len() {
let j = i - period;
let (sine, lead) = calc(state, real[i], real[i - period]);
unsafe {
*sine_line.get_unchecked_mut(j) = sine;
*lead_line.get_unchecked_mut(j) = lead;
}
}
}
fn cycle_msw_sdft<const N: usize>(
real: &[f64],
period: usize,
multiplier: f64,
cos_twiddles: &[f64],
sin_twiddles: &[f64],
sine_line: &mut [f64],
lead_line: &mut [f64],
) -> State {
let (rp, ip) = dot_product_simd::<N>(&real[..period], cos_twiddles, sin_twiddles);
let mut state = State {
rp,
ip,
..State::new(multiplier)
};
cycle_sdft(real, period, &mut state, sine_line, lead_line);
state
}
#[inline(always)]
pub fn phase_from_rp_ip(rp: f64, ip: f64) -> (f64, f64) {
let mut p = if rp.abs() > 0.001 {
(ip / rp).atan()
} else {
PI * if ip < 0.0 { -1.0 } else { 1.0 }
};
if rp < 0.0 {
p += PI;
}
p += HPI;
if p < 0.0 {
p += TPI;
}
if p > TPI {
p -= TPI;
}
let (sp, cp) = p.sin_cos();
let lead = INV_SQRT2.mul_add(sp, INV_SQRT2 * cp);
(sp, lead)
}
pub fn precompute_twiddles(period: usize, multiplier: f64) -> (Vec<f64>, Vec<f64>) {
let mut cos_twiddles = Vec::with_capacity(period);
let mut sin_twiddles = Vec::with_capacity(period);
for k in 0..period {
let angle = TPI * (period - 1 - k) as f64 * multiplier;
let (s, c) = angle.sin_cos();
cos_twiddles.push(c);
sin_twiddles.push(s);
}
(cos_twiddles, sin_twiddles)
}
#[inline(always)]
pub(crate) fn dot_product_simd<const N: usize>(
price_window: &[f64],
cos_twiddles: &[f64],
sin_twiddles: &[f64],
) -> (f64, f64) {
let mut rp_acc = Simd::<f64, N>::splat(0.0);
let mut ip_acc = Simd::<f64, N>::splat(0.0);
let full_chunks = price_window.len() / N;
for i in 0..full_chunks {
let offset = i * N;
let w = Simd::<f64, N>::from_slice(&price_window[offset..]);
let c = Simd::<f64, N>::from_slice(&cos_twiddles[offset..]);
let s = Simd::<f64, N>::from_slice(&sin_twiddles[offset..]);
rp_acc = c.mul_add(w, rp_acc);
ip_acc = s.mul_add(w, ip_acc);
}
let mut rp = rp_acc.reduce_sum();
let mut ip = ip_acc.reduce_sum();
for i in (full_chunks * N)..price_window.len() {
rp = cos_twiddles[i].mul_add(price_window[i], rp);
ip = sin_twiddles[i].mul_add(price_window[i], ip);
}
(rp, ip)
}