1use rand::seq::SliceRandom;
  2use rand_core::RngCore;
  3use rand_distr::{Distribution, weighted::WeightedIndex};
  4
  5use crate::{
  6    alloc_aligned,
  7    api::{DataView, DataViewMut, FillUniform, Reset, ZnxInfos, ZnxSliceSize, ZnxView, ZnxViewMut, ZnxZero},
  8    layouts::{Data, DataMut, DataRef, ReaderFrom, ToOwnedDeep, VecZnx, WriterTo},
  9    source::Source,
 10};
 11
 12#[derive(PartialEq, Eq, Debug, Clone)]
 13pub struct ScalarZnx<D: Data> {
 14    pub data: D,
 15    pub n: usize,
 16    pub cols: usize,
 17}
 18
 19impl<D: DataRef> ToOwnedDeep for ScalarZnx<D> {
 20    type Owned = ScalarZnx<Vec<u8>>;
 21    fn to_owned_deep(&self) -> Self::Owned {
 22        ScalarZnx {
 23            data: self.data.as_ref().to_vec(),
 24            n: self.n,
 25            cols: self.cols,
 26        }
 27    }
 28}
 29
 30impl<D: Data> ZnxInfos for ScalarZnx<D> {
 31    fn cols(&self) -> usize {
 32        self.cols
 33    }
 34
 35    fn rows(&self) -> usize {
 36        1
 37    }
 38
 39    fn n(&self) -> usize {
 40        self.n
 41    }
 42
 43    fn size(&self) -> usize {
 44        1
 45    }
 46}
 47
 48impl<D: Data> ZnxSliceSize for ScalarZnx<D> {
 49    fn sl(&self) -> usize {
 50        self.n()
 51    }
 52}
 53
 54impl<D: Data> DataView for ScalarZnx<D> {
 55    type D = D;
 56    fn data(&self) -> &Self::D {
 57        &self.data
 58    }
 59}
 60
 61impl<D: Data> DataViewMut for ScalarZnx<D> {
 62    fn data_mut(&mut self) -> &mut Self::D {
 63        &mut self.data
 64    }
 65}
 66
 67impl<D: DataRef> ZnxView for ScalarZnx<D> {
 68    type Scalar = i64;
 69}
 70
 71impl<D: DataMut> ScalarZnx<D> {
 72    pub fn fill_ternary_prob(&mut self, col: usize, prob: f64, source: &mut Source) {
 73        let choices: [i64; 3] = [-1, 0, 1];
 74        let weights: [f64; 3] = [prob / 2.0, 1.0 - prob, prob / 2.0];
 75        let dist: WeightedIndex<f64> = WeightedIndex::new(weights).unwrap();
 76        self.at_mut(col, 0)
 77            .iter_mut()
 78            .for_each(|x: &mut i64| *x = choices[dist.sample(source)]);
 79    }
 80
 81    pub fn fill_ternary_hw(&mut self, col: usize, hw: usize, source: &mut Source) {
 82        assert!(hw <= self.n());
 83        self.at_mut(col, 0)[..hw]
 84            .iter_mut()
 85            .for_each(|x: &mut i64| *x = (((source.next_u32() & 1) as i64) << 1) - 1);
 86        self.at_mut(col, 0).shuffle(source);
 87    }
 88
 89    pub fn fill_binary_prob(&mut self, col: usize, prob: f64, source: &mut Source) {
 90        let choices: [i64; 2] = [0, 1];
 91        let weights: [f64; 2] = [1.0 - prob, prob];
 92        let dist: WeightedIndex<f64> = WeightedIndex::new(weights).unwrap();
 93        self.at_mut(col, 0)
 94            .iter_mut()
 95            .for_each(|x: &mut i64| *x = choices[dist.sample(source)]);
 96    }
 97
 98    pub fn fill_binary_hw(&mut self, col: usize, hw: usize, source: &mut Source) {
 99        assert!(hw <= self.n());
100        self.at_mut(col, 0)[..hw]
101            .iter_mut()
102            .for_each(|x: &mut i64| *x = (source.next_u32() & 1) as i64);
103        self.at_mut(col, 0).shuffle(source);
104    }
105
106    pub fn fill_binary_block(&mut self, col: usize, block_size: usize, source: &mut Source) {
107        assert!(self.n().is_multiple_of(block_size));
108        let max_idx: u64 = (block_size + 1) as u64;
109        let mask_idx: u64 = (1 << ((u64::BITS - max_idx.leading_zeros()) as u64)) - 1;
110        for block in self.at_mut(col, 0).chunks_mut(block_size) {
111            let idx: usize = source.next_u64n(max_idx, mask_idx) as usize;
112            if idx != block_size {
113                block[idx] = 1;
114            }
115        }
116    }
117}
118
119impl ScalarZnx<Vec<u8>> {
120    pub fn alloc_bytes(n: usize, cols: usize) -> usize {
121        n * cols * size_of::<i64>()
122    }
123
124    pub fn alloc(n: usize, cols: usize) -> Self {
125        let data: Vec<u8> = alloc_aligned::<u8>(Self::alloc_bytes(n, cols));
126        Self { data, n, cols }
127    }
128
129    pub fn from_bytes(n: usize, cols: usize, bytes: impl Into<Vec<u8>>) -> Self {
130        let data: Vec<u8> = bytes.into();
131        assert!(data.len() == Self::alloc_bytes(n, cols));
132        Self { data, n, cols }
133    }
134}
135
136impl<D: DataMut> ZnxZero for ScalarZnx<D> {
137    fn zero(&mut self) {
138        self.raw_mut().fill(0)
139    }
140    fn zero_at(&mut self, i: usize, j: usize) {
141        self.at_mut(i, j).fill(0);
142    }
143}
144
145impl<D: DataMut> FillUniform for ScalarZnx<D> {
146    fn fill_uniform(&mut self, source: &mut Source) {
147        source.fill_bytes(self.data.as_mut());
148    }
149}
150
151impl<D: DataMut> Reset for ScalarZnx<D> {
152    fn reset(&mut self) {
153        self.zero();
154        self.n = 0;
155        self.cols = 0;
156    }
157}
158
159pub type ScalarZnxOwned = ScalarZnx<Vec<u8>>;
160
161impl<D: Data> ScalarZnx<D> {
162    pub fn from_data(data: D, n: usize, cols: usize) -> Self {
163        Self { data, n, cols }
164    }
165}
166
167pub trait ScalarZnxToRef {
168    fn to_ref(&self) -> ScalarZnx<&[u8]>;
169}
170
171impl<D: DataRef> ScalarZnxToRef for ScalarZnx<D> {
172    fn to_ref(&self) -> ScalarZnx<&[u8]> {
173        ScalarZnx {
174            data: self.data.as_ref(),
175            n: self.n,
176            cols: self.cols,
177        }
178    }
179}
180
181pub trait ScalarZnxToMut {
182    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]>;
183}
184
185impl<D: DataMut> ScalarZnxToMut for ScalarZnx<D> {
186    fn to_mut(&mut self) -> ScalarZnx<&mut [u8]> {
187        ScalarZnx {
188            data: self.data.as_mut(),
189            n: self.n,
190            cols: self.cols,
191        }
192    }
193}
194
195impl<D: DataRef> ScalarZnx<D> {
196    pub fn as_vec_znx(&self) -> VecZnx<&[u8]> {
197        VecZnx {
198            data: self.data.as_ref(),
199            n: self.n,
200            cols: self.cols,
201            size: 1,
202            max_size: 1,
203        }
204    }
205}
206
207impl<D: DataMut> ScalarZnx<D> {
208    pub fn as_vec_znx_mut(&mut self) -> VecZnx<&mut [u8]> {
209        VecZnx {
210            data: self.data.as_mut(),
211            n: self.n,
212            cols: self.cols,
213            size: 1,
214            max_size: 1,
215        }
216    }
217}
218
219use byteorder::{LittleEndian, ReadBytesExt, WriteBytesExt};
220
221impl<D: DataMut> ReaderFrom for ScalarZnx<D> {
222    fn read_from<R: std::io::Read>(&mut self, reader: &mut R) -> std::io::Result<()> {
223        self.n = reader.read_u64::<LittleEndian>()? as usize;
224        self.cols = reader.read_u64::<LittleEndian>()? as usize;
225        let len: usize = reader.read_u64::<LittleEndian>()? as usize;
226        let buf: &mut [u8] = self.data.as_mut();
227        if buf.len() != len {
228            return Err(std::io::Error::new(
229                std::io::ErrorKind::UnexpectedEof,
230                format!("self.data.len()={} != read len={}", buf.len(), len),
231            ));
232        }
233        reader.read_exact(&mut buf[..len])?;
234        Ok(())
235    }
236}
237
238impl<D: DataRef> WriterTo for ScalarZnx<D> {
239    fn write_to<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
240        writer.write_u64::<LittleEndian>(self.n as u64)?;
241        writer.write_u64::<LittleEndian>(self.cols as u64)?;
242        let buf: &[u8] = self.data.as_ref();
243        writer.write_u64::<LittleEndian>(buf.len() as u64)?;
244        writer.write_all(buf)?;
245        Ok(())
246    }
247}