use ndarray::Array2;
use numpy;
use pyo3::prelude::*;
use std::collections::HashMap;
use crate::channel::*;
use crate::device::*;
use crate::instruction::*;
pub trait BaseExperiment {
fn devices(&self) -> &HashMap<String, Device>;
fn devices_(&mut self) -> &mut HashMap<String, Device>;
fn assert_has_device(&self, name: &str) {
assert!(
self.devices().contains_key(name),
"Physical device {} not found. Registered devices are {:?}",
name,
self.devices().keys().collect::<Vec<_>>()
);
}
fn assert_device_has_channel(&self, name: &str, chan_name: &str) {
self.assert_has_device(name);
let device = self.devices().get(name).unwrap();
assert!(
device.channels().contains_key(chan_name),
"Channel name {} not found in device {}. Registered channels are: {:?}",
chan_name,
name,
device.channels().keys().collect::<Vec<_>>()
);
}
fn add_device_base(&mut self, dev: Device) {
let name = dev.name();
assert!(
!self.devices().contains_key(name),
"Device {} already registered. Registered devices are {:?}",
name,
self.devices().keys().collect::<Vec<_>>()
);
self.devices_().insert(name.to_string(), dev);
}
fn add_ao_device(&mut self, name: &str, samp_rate: f64) {
self.add_device_base(Device::new(name, TaskType::AO, samp_rate));
}
fn add_do_device(&mut self, name: &str, samp_rate: f64) {
self.add_device_base(Device::new(name, TaskType::DO, samp_rate));
}
fn edit_stop_time(&self) -> f64 {
self.devices()
.values()
.map(|dev| dev.edit_stop_time())
.fold(0.0, f64::max)
}
fn compiled_stop_time(&self) -> f64 {
self.devices()
.values()
.map(|dev| dev.compiled_stop_time())
.fold(0.0, f64::max)
}
fn compile(&mut self) -> f64 {
let stop_time = self.edit_stop_time();
self.compile_with_stoptime(stop_time);
assert!(stop_time == self.compiled_stop_time());
stop_time
}
fn compile_with_stoptime(&mut self, stop_time: f64) {
assert!(
self.devices().values().all(|d| d.export_trig().is_none())
|| self
.devices()
.values()
.filter(|d| d.export_trig() == Some(true))
.count()
== 1,
"Cannot compile an experiment with devices expecting yet no device exporting trigger"
);
self.devices_()
.values_mut()
.for_each(|dev| dev.compile(((stop_time) * dev.samp_rate()) as usize));
}
fn compiled_devices(&self) -> Vec<&Device> {
self.devices()
.values()
.filter_map(|dev| if dev.is_compiled() { Some(&*dev) } else { None })
.collect()
}
fn is_edited(&self) -> bool {
self.devices().values().any(|dev| dev.is_edited())
}
fn is_compiled(&self) -> bool {
self.devices().values().any(|dev| dev.is_compiled())
}
fn is_fresh_compiled(&self) -> bool {
self.devices().values().all(|dev| dev.is_fresh_compiled())
}
fn clear_edit_cache(&mut self) {
self.devices_()
.values_mut()
.for_each(|dev| dev.clear_edit_cache());
}
fn clear_compile_cache(&mut self) {
self.devices_()
.values_mut()
.for_each(|dev| dev.clear_compile_cache());
}
fn typed_device_op<F, R>(&mut self, name: &str, task_type: TaskType, mut f: F) -> R
where
F: FnMut(&mut Device) -> R,
{
self.assert_has_device(name);
let dev = self.devices_().get_mut(name).unwrap();
assert!(
dev.task_type() == task_type,
"Device {} is incompatible with instruction",
name
);
f(dev)
}
fn device_op<F, R>(&mut self, name: &str, mut f: F) -> R
where
F: FnMut(&mut Device) -> R,
{
self.assert_has_device(name);
let dev = self.devices_().get_mut(name).unwrap();
f(dev)
}
fn typed_channel_op<F, R>(
&mut self,
name: &str,
chan_name: &str,
task_type: TaskType,
mut f: F,
) -> R
where
F: FnMut(&mut Channel) -> R,
{
self.assert_device_has_channel(name, chan_name);
let dev = self.devices_().get_mut(name).unwrap();
assert!(
dev.task_type() == task_type,
"Channel {}/{} is incompatible with instruction",
name,
chan_name
);
let chan = dev.channels_().get_mut(chan_name).unwrap();
f(chan)
}
fn channel_op<F, R>(&mut self, name: &str, chan_name: &str, mut f: F) -> R
where
F: FnMut(&mut Channel) -> R,
{
self.assert_device_has_channel(name, chan_name);
let chan = self
.devices_()
.get_mut(name)
.unwrap()
.channels_()
.get_mut(chan_name)
.unwrap();
f(chan)
}
fn add_ao_channel(&mut self, name: &str, channel_id: usize) {
self.typed_device_op(name, TaskType::AO, |dev| {
(*dev).add_channel(&format!("ao{}", channel_id))
});
}
fn add_do_channel(&mut self, name: &str, port_id: usize, line_id: usize) {
self.typed_device_op(name, TaskType::DO, |dev| {
(*dev).add_channel(&format!("port{}/line{}", port_id, line_id))
});
}
fn device_calc_signal_nsamps(
&mut self,
dev_name: &str,
start_pos: usize,
end_pos: usize,
nsamps: usize,
require_streamable: bool,
require_editable: bool,
) -> Array2<f64> {
self.device_op(dev_name, |dev| {
(*dev).calc_signal_nsamps(
start_pos,
end_pos,
nsamps,
require_streamable,
require_editable,
)
})
}
fn device_cfg_samp_clk_src(&mut self, name: &str, src: &str) {
self.device_op(name, |dev| (*dev).cfg_samp_clk_src(src))
}
fn device_cfg_trig(&mut self, name: &str, trig_line: &str, export_trig: bool) {
assert!(
!export_trig
|| (export_trig
&& self
.devices()
.values()
.all(|dev| dev.export_trig().is_none())),
"Device {} cannot export triggers since another device already exports triggers.",
name
);
self.device_op(name, |dev| (*dev).cfg_trig(trig_line, export_trig))
}
fn device_cfg_ref_clk(
&mut self,
name: &str,
ref_clk_line: &str,
ref_clk_rate: f64,
export_ref_clk: bool,
) {
self.device_op(name, |dev| {
(*dev).cfg_ref_clk(ref_clk_line, ref_clk_rate, export_ref_clk)
})
}
fn device_edit_stop_time(&mut self, name: &str) -> f64 {
self.device_op(name, |dev| (*dev).edit_stop_time())
}
fn device_compiled_stop_time(&mut self, name: &str) -> f64 {
self.device_op(name, |dev| (*dev).compiled_stop_time())
}
fn device_clear_compile_cache(&mut self, name: &str) {
self.device_op(name, |dev| (*dev).clear_compile_cache())
}
fn device_clear_edit_cache(&mut self, name: &str) {
self.device_op(name, |dev| (*dev).clear_edit_cache())
}
fn device_compiled_channel_names(
&mut self,
name: &str,
require_streamable: bool,
require_editable: bool,
) -> Vec<String> {
self.device_op(name, |dev| {
(*dev)
.compiled_channels(require_streamable, require_editable)
.iter()
.map(|chan| chan.name().to_string())
.collect()
})
}
fn constant(
&mut self,
dev_name: &str,
chan_name: &str,
t: f64,
duration: f64,
value: f64,
keep_val: bool,
) {
self.typed_channel_op(dev_name, chan_name, TaskType::AO, |chan| {
(*chan).constant(value, t, duration, keep_val);
});
}
fn sine(
&mut self,
dev_name: &str,
chan_name: &str,
t: f64,
duration: f64,
keep_val: bool,
freq: f64,
amplitude: Option<f64>,
phase: Option<f64>,
dc_offset: Option<f64>,
) {
self.typed_channel_op(dev_name, chan_name, TaskType::AO, |chan| {
let instr = Instruction::new_sine(freq, amplitude, phase, dc_offset);
(*chan).add_instr(instr, t, duration, keep_val)
});
}
fn high(&mut self, dev_name: &str, chan_name: &str, t: f64, duration: f64) {
self.typed_channel_op(dev_name, chan_name, TaskType::DO, |chan| {
(*chan).constant(1., t, duration, false);
});
}
fn low(&mut self, dev_name: &str, chan_name: &str, t: f64, duration: f64) {
self.typed_channel_op(dev_name, chan_name, TaskType::DO, |chan| {
(*chan).constant(0., t, duration, false);
});
}
fn go_high(&mut self, dev_name: &str, chan_name: &str, t: f64) {
self.typed_channel_op(dev_name, chan_name, TaskType::DO, |chan| {
(*chan).constant(1., t, 1. / (*chan).samp_rate(), true);
});
}
fn go_low(&mut self, dev_name: &str, chan_name: &str, t: f64) {
self.typed_channel_op(dev_name, chan_name, TaskType::DO, |chan| {
(*chan).constant(0., t, 1. / (*chan).samp_rate(), true);
});
}
fn channel_clear_edit_cache(&mut self, dev_name: &str, chan_name: &str) {
self.channel_op(dev_name, chan_name, |chan| (*chan).clear_edit_cache());
}
fn channel_calc_signal_nsamps(
&mut self,
dev_name: &str,
chan_name: &str,
start_time: f64,
end_time: f64,
num_samps: usize,
) -> Vec<f64> {
self.channel_op(dev_name, chan_name, |chan| {
(*chan).calc_signal_nsamps(start_time, end_time, num_samps)
})
}
fn channel_clear_compile_cache(&mut self, dev_name: &str, chan_name: &str) {
self.channel_op(dev_name, chan_name, |chan| (*chan).clear_compile_cache());
}
}
#[pyclass]
pub struct Experiment {
devices: HashMap<String, Device>,
}
#[macro_export]
macro_rules! impl_exp_boilerplate {
($exp_type: ty) => {
impl BaseExperiment for $exp_type {
fn devices(&self) -> &HashMap<String, Device> {
&self.devices
}
fn devices_(&mut self) -> &mut HashMap<String, Device> {
&mut self.devices
}
}
#[pymethods]
impl $exp_type {
fn add_ao_device(&mut self, name: &str, samp_rate: f64) {
BaseExperiment::add_ao_device(self, name, samp_rate);
}
fn add_do_device(&mut self, name: &str, samp_rate: f64) {
BaseExperiment::add_do_device(self, name, samp_rate);
}
pub fn edit_stop_time(&self) -> f64 {
BaseExperiment::edit_stop_time(self)
}
pub fn compiled_stop_time(&self) -> f64 {
BaseExperiment::compiled_stop_time(self)
}
pub fn compile(&mut self) -> f64 {
BaseExperiment::compile(self)
}
pub fn compile_with_stoptime(&mut self, stop_time: f64) {
BaseExperiment::compile_with_stoptime(self, stop_time);
}
pub fn is_edited(&self) -> bool {
BaseExperiment::is_edited(self)
}
pub fn is_compiled(&self) -> bool {
BaseExperiment::is_compiled(self)
}
pub fn is_fresh_compiled(&self) -> bool {
BaseExperiment::is_fresh_compiled(self)
}
pub fn clear_edit_cache(&mut self) {
BaseExperiment::clear_edit_cache(self);
}
pub fn clear_compile_cache(&mut self) {
BaseExperiment::clear_compile_cache(self);
}
pub fn add_ao_channel(&mut self, name: &str, channel_id: usize) {
BaseExperiment::add_ao_channel(self, name, channel_id);
}
pub fn add_do_channel(&mut self, name: &str, port_id: usize, line_id: usize) {
BaseExperiment::add_do_channel(self, name, port_id, line_id);
}
pub fn device_cfg_samp_clk_src(&mut self, name: &str, src: &str) {
BaseExperiment::device_cfg_samp_clk_src(self, name, src);
}
pub fn device_cfg_trig(&mut self, name: &str, trig_line: &str, export_trig: bool) {
BaseExperiment::device_cfg_trig(self, name, trig_line, export_trig);
}
pub fn device_cfg_ref_clk(
&mut self,
name: &str,
ref_clk_line: &str,
ref_clk_rate: f64,
export_ref_clk: bool,
) {
BaseExperiment::device_cfg_ref_clk(
self,
name,
ref_clk_line,
ref_clk_rate,
export_ref_clk,
);
}
pub fn device_compiled_channel_names(
&mut self,
name: &str,
require_streamable: bool,
require_editable: bool,
) -> Vec<String> {
BaseExperiment::device_compiled_channel_names(
self,
name,
require_streamable,
require_editable,
)
}
pub fn calc_signal(
&mut self,
name: &str,
t_start: f64,
t_end: f64,
nsamps: usize,
require_streamable: bool,
require_editable: bool,
py: Python,
) -> PyResult<PyObject> {
self.assert_has_device(name);
let samp_rate = self.devices().get(name).unwrap().samp_rate();
let arr = BaseExperiment::device_calc_signal_nsamps(
self,
name,
(t_start * samp_rate) as usize,
(t_end * samp_rate) as usize,
nsamps,
require_streamable,
require_editable,
);
Ok(numpy::PyArray::from_array(py, &arr).to_object(py))
}
pub fn device_edit_stop_time(&mut self, name: &str) -> f64 {
BaseExperiment::device_edit_stop_time(self, name)
}
pub fn device_compiled_stop_time(&mut self, name: &str) -> f64 {
BaseExperiment::device_compiled_stop_time(self, name)
}
pub fn device_clear_compile_cache(&mut self, name: &str) {
BaseExperiment::device_clear_compile_cache(self, name)
}
pub fn device_clear_edit_cache(&mut self, name: &str) {
BaseExperiment::device_clear_edit_cache(self, name)
}
pub fn constant(
&mut self,
dev_name: &str,
chan_name: &str,
t: f64,
duration: f64,
value: f64,
keep_val: bool,
) {
BaseExperiment::constant(self, dev_name, chan_name, t, duration, value, keep_val);
}
pub fn sine(
&mut self,
dev_name: &str,
chan_name: &str,
t: f64,
duration: f64,
keep_val: bool,
freq: f64,
amplitude: Option<f64>,
phase: Option<f64>,
dc_offset: Option<f64>,
) {
BaseExperiment::sine(
self, dev_name, chan_name, t, duration, keep_val, freq, amplitude, phase,
dc_offset,
);
}
pub fn high(&mut self, dev_name: &str, chan_name: &str, t: f64, duration: f64) {
BaseExperiment::high(self, dev_name, chan_name, t, duration);
}
pub fn low(&mut self, dev_name: &str, chan_name: &str, t: f64, duration: f64) {
BaseExperiment::low(self, dev_name, chan_name, t, duration);
}
pub fn go_high(&mut self, dev_name: &str, chan_name: &str, t: f64) {
BaseExperiment::go_high(self, dev_name, chan_name, t);
}
pub fn go_low(&mut self, dev_name: &str, chan_name: &str, t: f64) {
BaseExperiment::go_low(self, dev_name, chan_name, t);
}
pub fn channel_clear_compile_cache(&mut self, dev_name: &str, chan_name: &str) {
BaseExperiment::channel_clear_compile_cache(self, dev_name, chan_name);
}
pub fn channel_clear_edit_cache(&mut self, dev_name: &str, chan_name: &str) {
BaseExperiment::channel_clear_edit_cache(self, dev_name, chan_name);
}
pub fn channel_calc_signal_nsamps(
&mut self,
dev_name: &str,
chan_name: &str,
start_time: f64,
end_time: f64,
num_samps: usize,
) -> Vec<f64> {
BaseExperiment::channel_calc_signal_nsamps(
self, dev_name, chan_name, start_time, end_time, num_samps,
)
}
}
};
}
#[pymethods]
impl Experiment {
#[new]
pub fn new() -> Self {
Self {
devices: HashMap::new(),
}
}
}
impl_exp_boilerplate!(Experiment);