pmcore/estimation/nonparametric/
posterior.rs

1pub use anyhow::{bail, Result};
2use faer::Mat;
3use serde::{Deserialize, Serialize};
4
5use crate::estimation::nonparametric::{psi::Psi, weights::Weights};
6
7#[derive(Debug, Clone)]
8pub struct Posterior {
9    mat: Mat<f64>,
10}
11
12impl Posterior {
13    fn new(mat: Mat<f64>) -> Self {
14        Posterior { mat }
15    }
16
17    pub fn calculate(psi: &Psi, w: &Weights) -> Result<Self> {
18        if psi.matrix().ncols() != w.weights().nrows() {
19            bail!(
20                "Number of rows in psi ({}) and number of weights ({}) do not match.",
21                psi.matrix().nrows(),
22                w.weights().nrows()
23            );
24        }
25
26        let psi_matrix = psi.matrix();
27        let py = psi_matrix * w.weights();
28
29        let posterior = Mat::from_fn(psi_matrix.nrows(), psi_matrix.ncols(), |i, j| {
30            psi_matrix.get(i, j) * w.weights().get(j) / py.get(i)
31        });
32
33        Ok(posterior.into())
34    }
35
36    pub fn matrix(&self) -> &Mat<f64> {
37        &self.mat
38    }
39
40    pub fn to_csv<W: std::io::Write>(&self, writer: W) -> Result<()> {
41        let mut csv_writer = csv::Writer::from_writer(writer);
42
43        for i in 0..self.mat.nrows() {
44            let row: Vec<f64> = (0..self.mat.ncols()).map(|j| *self.mat.get(i, j)).collect();
45            csv_writer.serialize(row)?;
46        }
47
48        csv_writer.flush()?;
49        Ok(())
50    }
51
52    pub fn from_csv<R: std::io::Read>(reader: R) -> Result<Self> {
53        let mut csv_reader = csv::Reader::from_reader(reader);
54        let mut rows: Vec<Vec<f64>> = Vec::new();
55
56        for result in csv_reader.deserialize() {
57            let row: Vec<f64> = result?;
58            rows.push(row);
59        }
60
61        if rows.is_empty() {
62            bail!("CSV file is empty");
63        }
64
65        let nrows = rows.len();
66        let ncols = rows[0].len();
67
68        for (i, row) in rows.iter().enumerate() {
69            if row.len() != ncols {
70                bail!("Row {} has {} columns, expected {}", i, row.len(), ncols);
71            }
72        }
73
74        let mat = Mat::from_fn(nrows, ncols, |i, j| rows[i][j]);
75
76        Ok(Posterior::new(mat))
77    }
78}
79
80impl From<Mat<f64>> for Posterior {
81    fn from(mat: Mat<f64>) -> Self {
82        Posterior::new(mat)
83    }
84}
85
86impl Serialize for Posterior {
87    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
88    where
89        S: serde::Serializer,
90    {
91        use serde::ser::SerializeSeq;
92
93        let mut seq = serializer.serialize_seq(Some(self.mat.nrows()))?;
94
95        for i in 0..self.mat.nrows() {
96            let row: Vec<f64> = (0..self.mat.ncols()).map(|j| *self.mat.get(i, j)).collect();
97            seq.serialize_element(&row)?;
98        }
99
100        seq.end()
101    }
102}
103
104impl<'de> Deserialize<'de> for Posterior {
105    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
106    where
107        D: serde::Deserializer<'de>,
108    {
109        use serde::de::{SeqAccess, Visitor};
110        use std::fmt;
111
112        struct PosteriorVisitor;
113
114        impl<'de> Visitor<'de> for PosteriorVisitor {
115            type Value = Posterior;
116
117            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
118                formatter.write_str("a sequence of rows (vectors of f64)")
119            }
120
121            fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
122            where
123                A: SeqAccess<'de>,
124            {
125                let mut rows: Vec<Vec<f64>> = Vec::new();
126
127                while let Some(row) = seq.next_element::<Vec<f64>>()? {
128                    rows.push(row);
129                }
130
131                if rows.is_empty() {
132                    return Err(serde::de::Error::custom("Empty matrix not allowed"));
133                }
134
135                let nrows = rows.len();
136                let ncols = rows[0].len();
137
138                for (i, row) in rows.iter().enumerate() {
139                    if row.len() != ncols {
140                        return Err(serde::de::Error::custom(format!(
141                            "Row {} has {} columns, expected {}",
142                            i,
143                            row.len(),
144                            ncols
145                        )));
146                    }
147                }
148
149                let mat = Mat::from_fn(nrows, ncols, |i, j| rows[i][j]);
150
151                Ok(Posterior::new(mat))
152            }
153        }
154
155        deserializer.deserialize_seq(PosteriorVisitor)
156    }
157}
158
159pub fn posterior(psi: &Psi, w: &Weights) -> Result<Posterior> {
160    Posterior::calculate(psi, w)
161}
pmcore/estimation/nonparametric/posterior.rs

pmcore/estimation/nonparametric/
posterior.rs
