Skip to main content

causal_hub/datasets/table/gaussian/
complete.rs

1use std::{
2    borrow::Cow,
3    io::{Read, Write},
4    sync::Arc,
5};
6
7use csv::{ReaderBuilder, WriterBuilder};
8use ndarray::prelude::*;
9use serde::{Deserialize, Serialize};
10
11use crate::{
12    datasets::{Dataset, GaussEv, GaussEvT},
13    io::CsvIO,
14    models::{GaussSupport, HasLabels},
15    types::{Error, Labels, Result, Set},
16};
17
18/// A type alias for a gaussian variable.
19pub type GaussType = f64;
20/// A type alias for a gaussian sample.
21pub type GaussSample = Array1<GaussType>;
22
23/// A struct representing a gaussian dataset.
24#[derive(Clone, Debug, Serialize, Deserialize)]
25pub struct GaussTable {
26    labels: Labels,
27    values: Array2<GaussType>,
28}
29
30/// Concrete iterator over Gaussian table evidences.
31pub struct GaussTableEvidenceIter<'a> {
32    rows: ndarray::iter::LanesIter<'a, GaussType, Ix1>,
33    labels: &'a Labels,
34}
35
36impl<'a> Iterator for GaussTableEvidenceIter<'a> {
37    type Item = Result<GaussEv>;
38
39    fn next(&mut self) -> Option<Self::Item> {
40        let row = self.rows.next()?;
41
42        let evidences = row
43            .iter()
44            .enumerate()
45            .map(|(event, &value)| GaussEvT::CertainPositive { event, value });
46
47        Some(GaussEv::new(self.labels.clone(), evidences))
48    }
49}
50
51impl HasLabels for GaussTable {
52    #[inline]
53    fn labels(&self) -> &Labels {
54        &self.labels
55    }
56}
57
58impl GaussTable {
59    /// Creates a new gaussian dataset.
60    ///
61    /// # Arguments
62    ///
63    /// * `labels` - The labels of the variables.
64    /// * `values` - The values of the variables.
65    ///
66    /// # Panics
67    ///
68    /// * Panics if the number of columns in `values` does not match the number of `labels`.
69    ///
70    /// # Results
71    ///
72    /// A new gaussian dataset instance.
73    ///
74    pub fn new(mut labels: Labels, mut values: Array2<GaussType>) -> Result<Self> {
75        // Check if the number of labels matches the number of columns in values.
76        if labels.len() != values.ncols() {
77            return Err(Error::IncompatibleShape(
78                &labels.len().to_string(),
79                &values.ncols().to_string(),
80            ));
81        }
82
83        // Sort labels and values accordingly.
84        if !labels.is_sorted() {
85            // Allocate indices to sort labels.
86            let mut indices: Vec<usize> = (0..labels.len()).collect();
87            // Sort the indices by labels.
88            indices.sort_by_key(|&i| &labels[i]);
89            // Sort the labels.
90            labels.sort();
91            // Allocate new values.
92            let mut new_values = values.clone();
93            // Sort the new values according to the sorted indices.
94            indices.into_iter().enumerate().for_each(|(i, j)| {
95                new_values.column_mut(i).assign(&values.column(j));
96            });
97            // Update values.
98            values = new_values;
99        }
100        // Check values are finite.
101        if !values.iter().all(|&x| x.is_finite()) {
102            return Err(Error::InvalidParameter("values", "must be finite"));
103        }
104
105        Ok(Self { labels, values })
106    }
107}
108
109impl Dataset for GaussTable {
110    type Values = Array2<GaussType>;
111    type Support = GaussSupport;
112    type Evidence = GaussEv;
113    type EvidenceIter<'a> = GaussTableEvidenceIter<'a>;
114
115    #[inline]
116    fn values(&self) -> &Self::Values {
117        &self.values
118    }
119
120    fn support(&self) -> Cow<'_, Self::Support> {
121        Cow::Owned(
122            self.labels
123                .iter()
124                .map(|l| (l.clone(), (f64::NEG_INFINITY, f64::INFINITY)))
125                .collect(),
126        )
127    }
128
129    fn evidence_iter(&self) -> Self::EvidenceIter<'_> {
130        GaussTableEvidenceIter {
131            rows: self.values.rows().into_iter(),
132            labels: &self.labels,
133        }
134    }
135
136    #[inline]
137    fn sample_size(&self) -> f64 {
138        self.values.nrows() as f64
139    }
140
141    fn select(&self, x: &Set<usize>) -> Result<Self> {
142        // Check that the indices are valid.
143        if let Some(&i) = x.iter().find(|&&i| i >= self.values.ncols()) {
144            return Err(Error::IndexOutOfBounds(i));
145        }
146
147        // Select the labels.
148        let labels: Labels = x
149            .iter()
150            .map(|&i| {
151                self.labels
152                    .get_index(i)
153                    .cloned()
154                    .ok_or_else(|| Error::IndexOutOfBounds(i))
155            })
156            .collect::<Result<_>>()?;
157
158        // Select the values.
159        let mut new_values = Array2::zeros((self.values.nrows(), x.len()));
160        // Copy the selected columns.
161        x.iter().enumerate().for_each(|(j, &i)| {
162            new_values.column_mut(j).assign(&self.values.column(i));
163        });
164        // Update the values.
165        let values = new_values;
166
167        // Return the new dataset.
168        Self::new(labels, values)
169    }
170}
171
172impl CsvIO for GaussTable {
173    fn from_csv_reader<R: Read>(reader: R) -> Result<Self> {
174        // Create a CSV reader from the string.
175        let mut reader = ReaderBuilder::new().has_headers(true).from_reader(reader);
176
177        // Check if the reader has headers.
178        if !reader.has_headers() {
179            return Err(Error::MissingHeader());
180        }
181
182        // Read the headers.
183        let labels: Labels = reader
184            .headers()?
185            .into_iter()
186            .map(|x| x.to_owned())
187            .collect();
188
189        // Read the records.
190        let values: Vec<GaussType> = reader
191            .into_records()
192            .enumerate()
193            .map(|(i, row)| {
194                // Get the record row.
195                let row = row.map_err(|evidence| Error::Csv(Arc::new(evidence)))?;
196                // Get the record values and convert to indices.
197                row.into_iter()
198                    .enumerate()
199                    .map(|(j, x)| {
200                        // Check for missing values.
201                        if x.is_empty() {
202                            return Err(Error::MissingValue(i + 1, j + 1));
203                        }
204                        // Cast the value.
205                        Ok(x.parse::<GaussType>()?)
206                    })
207                    .collect::<Result<Vec<_>>>()
208            })
209            .collect::<Result<Vec<_>>>()?
210            .into_iter()
211            .flatten()
212            .collect();
213
214        // Convert the values to an array.
215        let values = Array1::from_vec(values);
216
217        // Get the number of rows and columns.
218        let ncols = labels.len();
219        let nrows = values.len() / ncols;
220        // Reshape the values to the correct shape.
221        let values = values.into_shape_with_order((nrows, ncols))?;
222
223        // Construct the dataset.
224        Self::new(labels, values)
225    }
226
227    fn to_csv_writer<W: Write>(&self, writer: W) -> Result<()> {
228        // Create the CSV writer.
229        let mut writer = WriterBuilder::new().has_headers(true).from_writer(writer);
230
231        // Write the headers.
232        writer.write_record(self.labels.iter())?;
233
234        // Write the records.
235        self.values
236            .rows()
237            .into_iter()
238            .try_for_each(|row| -> Result<_> {
239                // Map the row values to strings.
240                let record = row.iter().map(|x| x.to_string());
241                // Write the record.
242                writer.write_record(record)?;
243
244                Ok(())
245            })?;
246
247        Ok(())
248    }
249}