use crate::indicators::simd_indicators::road_train::{Asset, Driver, PrimeMover};
pub use crate::indicators::simd_indicators::tsf_simd::{calc_simd, SimdState};
use crate::indicators::tsf::{
min_data, output_length, IndicatorState, State, INPUTS_WIDTH, OPTIONS_WIDTH,
};
use crate::types::IndicatorError;
use crate::{common::validate_options, common_simd::assets::validate_inputs};
use std::simd::Simd;
struct TsfDriver {
want_optional_outputs: (bool, bool, bool, bool),
period: usize,
}
impl Driver<State> for TsfDriver {
fn next_run<const N: usize>(
&mut self,
inputs: Vec<Vec<&[f64]>>,
mut outputs: Vec<Vec<&mut [f64]>>,
mut states: Vec<&mut State>,
_options: Vec<Option<&()>>,
) {
let mut state = SimdState::<N>::new_mut_ref(&states);
let len = inputs[0][0].len();
let simd_period = Simd::splat(self.period as f64);
let (has_optional, want_linreg, want_slope, want_intercept) = self.want_optional_outputs;
let (tsf_line_ptr, linreg_line_ptr, slope_line_ptr, intercept_line_ptr) = crate::extract_output_ptrs!(
outputs,
N,
tsf_line_ptr,
linreg_line_ptr,
slope_line_ptr,
intercept_line_ptr
);
let real_ptrs = crate::extract_input_ptrs!(inputs, N, real_ptrs);
for (j, i) in (self.period..len).enumerate() {
let (real, prev_real) = crate::extract_simd_at_indices!(N, real_ptrs,
real @ i,
prev_real @ i + 1 - self.period
);
let (tsf, linreg, slope, intercept) =
calc_simd(&mut state, prev_real, real, simd_period);
crate::write_simd_at_indices!(N, j,
tsf_line_ptr => tsf
);
if has_optional {
crate::store_simd_optional_outputs!(j, N,
want_linreg, linreg_line_ptr => linreg,
want_slope, slope_line_ptr => slope,
want_intercept, intercept_line_ptr => intercept
);
}
}
state.write_states(&mut states);
}
}
pub fn indicator_by_assets<const N: usize>(
inputs: &[&[&[f64]; INPUTS_WIDTH]; N], options: &[f64; OPTIONS_WIDTH],
optional_outputs: Option<&[bool]>,
) -> Result<(Vec<Vec<Vec<f64>>>, Vec<IndicatorState>), IndicatorError> {
validate_inputs::<INPUTS_WIDTH>(inputs, min_data(options))?;
validate_options(options)?;
let period = options[0] as usize;
let mut road_train = PrimeMover::<N, State>::new();
let mut want_optional_outputs = (false, false, false, false);
let mut output_buffers = Vec::with_capacity(N);
for i in 0..N {
let asset_inputs = vec![
inputs[i][0], ];
let (tsf_line, linreg_line, slope_line, intercept_line);
{
let capacity = output_length(inputs[i][0].len(), options);
(linreg_line, slope_line, intercept_line) = crate::init_optional_outputs_eff!(
optional_outputs, &[false, false, false],
linreg_line: capacity,
slope_line: capacity,
intercept_line: capacity
);
tsf_line = crate::uninit_vec!(f64, capacity);
}
let state = State::init_state(&inputs[i][0][1..period], period);
if i == 0 {
want_optional_outputs =
crate::calc_want_flags!(linreg_line, slope_line, intercept_line);
}
let mut output_buffer = vec![tsf_line, linreg_line, slope_line, intercept_line];
let mut asset_outputs = Vec::with_capacity(output_buffer.len());
for j in 0..output_buffer.len() {
unsafe {
let output_buffer = &mut output_buffer[j];
asset_outputs.push(std::slice::from_raw_parts_mut(
output_buffer.as_mut_ptr(), output_buffer.len(), ));
}
}
road_train.add_asset(Asset::new(
asset_inputs,
asset_outputs,
i,
period,
period,
state,
None,
));
output_buffers.push(output_buffer);
}
let mut driver = TsfDriver {
period,
want_optional_outputs,
};
let states_vec = road_train.drive(&mut driver);
let mut states = Vec::with_capacity(N);
for (i, state) in states_vec.into_iter().enumerate() {
states.push(IndicatorState::new(
state,
unsafe { inputs.get_unchecked(i).get_unchecked(0) },
period,
));
}
Ok((output_buffers, states))
}