fmi_sim/sim/
io.rs

1use std::{fmt::Display, sync::Arc};
2
3use anyhow::Context;
4use arrow::{
5    array::{ArrayBuilder, ArrayRef, Float64Array, Float64Builder, downcast_array, make_builder},
6    datatypes::{DataType, Field, Schema},
7    downcast_primitive_array,
8    record_batch::RecordBatch,
9};
10use fmi::traits::FmiInstance;
11
12use crate::Error;
13
14use super::{
15    interpolation::{Interpolate, PreLookup, find_index},
16    params::SimParams,
17    traits::{ImportSchemaBuilder, InstSetValues},
18    util::project_input_data,
19};
20
21/// Container for holding initial values for the FMU.
22pub struct StartValues<VR> {
23    pub structural_parameters: Vec<(VR, ArrayRef)>,
24    pub variables: Vec<(VR, ArrayRef)>,
25}
26
27pub struct InputState<Inst: FmiInstance> {
28    pub(crate) input_data: Option<RecordBatch>,
29    // Map schema column index to ValueReference
30    pub(crate) continuous_inputs: Vec<(Field, Inst::ValueRef)>,
31    // Map schema column index to ValueReference
32    pub(crate) discrete_inputs: Vec<(Field, Inst::ValueRef)>,
33}
34
35impl<Inst> Display for InputState<Inst>
36where
37    Inst: FmiInstance,
38{
39    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        let continuous_inputs = self
41            .continuous_inputs
42            .iter()
43            .map(|(field, _)| field.name())
44            .collect::<Vec<_>>();
45
46        let discrete_inputs = self
47            .discrete_inputs
48            .iter()
49            .map(|(field, _)| field.name())
50            .collect::<Vec<_>>();
51
52        f.write_str("InputState {\n")?;
53        if let Some(input_data) = &self.input_data {
54            writeln!(
55                f,
56                "input_data:\n{}",
57                arrow::util::pretty::pretty_format_batches(std::slice::from_ref(input_data))
58                    .unwrap()
59            )?;
60        } else {
61            writeln!(f, "input_data: None")?;
62        }
63        writeln!(f, "continuous_inputs: {continuous_inputs:?}")?;
64        writeln!(f, "discrete_inputs: {discrete_inputs:?}")?;
65        write!(f, "}}")
66    }
67}
68
69impl<Inst> InputState<Inst>
70where
71    Inst: FmiInstance,
72{
73    pub fn new<Import: ImportSchemaBuilder<ValueRef = Inst::ValueRef>>(
74        import: &Import,
75        input_data: Option<RecordBatch>,
76    ) -> anyhow::Result<Self> {
77        let model_input_schema = Arc::new(import.inputs_schema());
78        let continuous_inputs = import.continuous_inputs().collect();
79        let discrete_inputs = import.discrete_inputs().collect();
80
81        let input_data = input_data
82            .map(|input_data| project_input_data(&input_data, model_input_schema.clone()))
83            .transpose()?;
84
85        Ok(Self {
86            input_data,
87            continuous_inputs,
88            discrete_inputs,
89        })
90    }
91}
92
93impl<Inst> InputState<Inst>
94where
95    Inst: InstSetValues,
96{
97    pub fn apply_input<I: Interpolate>(
98        &mut self,
99        time: f64,
100        inst: &mut Inst,
101        discrete: bool,
102        continuous: bool,
103        after_event: bool,
104    ) -> Result<(), Error> {
105        if let Some(input_data) = &self.input_data {
106            let time_array: Float64Array = downcast_array(
107                input_data
108                    .column_by_name("time")
109                    .context("Input data must have a column named 'time' with the time values")?,
110            );
111
112            if continuous {
113                let pl = PreLookup::new(&time_array, time, after_event);
114
115                for (field, vr) in &self.continuous_inputs {
116                    if let Some(input_col) = input_data.column_by_name(field.name()) {
117                        assert_eq!(input_col.data_type(), field.data_type());
118
119                        //let ary = arrow::compute::cast(input_col, field.data_type()).map_err(|_| anyhow::anyhow!("Error casting type"))?;
120
121                        //log::trace!( "Applying continuous input {}={input_col:?} at time {time}", field.name());
122
123                        inst.set_interpolated::<I>(*vr, &pl, input_col)?;
124                    }
125                }
126            }
127
128            if discrete {
129                // TODO: Refactor the interpolation code to separate index lookup from interpolation
130                let input_idx = find_index(&time_array, time, after_event);
131
132                for (field, vr) in &self.discrete_inputs {
133                    if let Some(input_col) = input_data.column_by_name(field.name()) {
134                        let ary = arrow::compute::cast(input_col, field.data_type())
135                            .map_err(|_| anyhow::anyhow!("Error casting type"))?;
136
137                        let values = &ary.slice(input_idx, 1);
138
139                        //log::trace!( "Applying discrete input {}={values:#?} at time {time:.2}", field.name());
140
141                        inst.set_array(&[*vr], values);
142                    }
143                }
144            }
145        }
146
147        Ok(())
148    }
149
150    /// Get the time of the next input event after the given time.
151    /// If no such event exists, returns `f64::INFINITY`.
152    pub fn next_input_event(&self, time: f64) -> f64 {
153        if let Some(input_data) = &self.input_data {
154            let time_array: Float64Array =
155                downcast_array(input_data.column_by_name("time").unwrap());
156
157            for i in 0..(time_array.len() - 1) {
158                let t0 = time_array.value(i);
159                let t1 = time_array.value(i + 1);
160
161                if time >= t1 {
162                    continue;
163                }
164
165                if t0 == t1 {
166                    return t0; // discrete change of a continuous variable
167                }
168
169                // TODO: This could be computed once and cached
170
171                // skip continuous variables
172                for (field, _vr) in &self.discrete_inputs {
173                    if let Some(input_col) = input_data.column_by_name(field.name())
174                        && downcast_primitive_array!(
175                            input_col => input_col.value(i) != input_col.value(i + 1),
176                            t => panic!("Unsupported datatype {}", t)
177                        )
178                    {
179                        return t1;
180                    }
181                }
182            }
183        }
184        f64::INFINITY
185    }
186}
187
188pub struct Recorder<Inst: FmiInstance> {
189    pub(crate) field: Field,
190    pub(crate) value_reference: Inst::ValueRef,
191    pub(crate) builder: Box<dyn ArrayBuilder>,
192    pub(crate) binary_max_size: Option<usize>,
193}
194
195pub struct RecorderState<Inst: FmiInstance> {
196    pub(crate) time: Float64Builder,
197    pub(crate) recorders: Vec<Recorder<Inst>>,
198}
199
200impl<Inst> RecorderState<Inst>
201where
202    Inst: FmiInstance,
203{
204    pub fn new<Import: ImportSchemaBuilder<ValueRef = Inst::ValueRef>>(
205        import: &Import,
206        sim_params: &SimParams,
207    ) -> Self {
208        let num_points = ((sim_params.stop_time - sim_params.start_time)
209            / sim_params.output_interval)
210            .ceil() as usize;
211
212        let time = Float64Builder::with_capacity(num_points);
213
214        let recorders = import
215            .outputs()
216            .map(|(field, vr)| {
217                let builder = make_builder(field.data_type(), num_points);
218                let binary_max_size = if field.data_type() == &DataType::Binary {
219                    import.binary_max_size(vr)
220                } else {
221                    None
222                };
223                Recorder {
224                    field,
225                    value_reference: vr,
226                    builder,
227                    binary_max_size,
228                }
229            })
230            .collect();
231
232        Self { time, recorders }
233    }
234
235    /// Finish the output state and return the RecordBatch.
236    pub fn finish(self) -> RecordBatch {
237        let Self {
238            mut time,
239            recorders,
240        } = self;
241
242        let recorders = recorders.into_iter().map(
243            |Recorder {
244                 field,
245                 value_reference: _,
246                 mut builder,
247                 ..
248             }| { (field, builder.finish()) },
249        );
250
251        let time = std::iter::once((
252            Field::new("time", DataType::Float64, false),
253            Arc::new(time.finish()) as _,
254        ));
255
256        let (fields, columns): (Vec<_>, Vec<_>) = time.chain(recorders).unzip();
257        let schema = Arc::new(Schema::new(fields));
258        RecordBatch::try_new(schema, columns).unwrap()
259    }
260}
fmi_sim/sim/io.rs

fmi_sim/sim/
io.rs
