Skip to main content

legume_numeric/matrix/
traits.rs

1use crate::matrix::common_io::{Delimiter, ReadLinesOut};
2use candle_core::{Device, Tensor};
3use num_traits::Float;
4
5/// Trait for running statistics operations
6///
7/// Provides a common interface for both dense (ndarray-based) and
8/// sparse running statistics implementations.
9pub trait RunningStatOps<T>
10where
11    T: Float,
12{
13    type Output;
14
15    fn clear(&mut self);
16    fn count_positives(&self) -> Self::Output;
17    fn sum(&self) -> Self::Output;
18    fn mean(&self) -> Self::Output;
19    fn variance(&self) -> Self::Output;
20    fn std(&self) -> Self::Output;
21}
22
23/// some linear algebra routines
24pub trait RandomizedAlgs {
25    type InMat;
26    type OutMat;
27    type DVec;
28    type Scalar;
29
30    /// randomized singular value decomposition
31    /// # input
32    /// * `X`: `n x d` matrix
33    /// # output
34    /// * `U`: `n x k`
35    /// * `D`: `k x 1`
36    /// * `V`: `d x k`
37    fn rsvd(&self, max_rank: usize) -> anyhow::Result<(Self::OutMat, Self::DVec, Self::OutMat)>;
38}
39
40/// Convert to and from the vector of triplets
41pub trait MatTriplets {
42    type Mat;
43    type Scalar;
44
45    fn from_nonzero_triplets<I>(
46        nrow: usize,
47        ncol: usize,
48        triplets: &[(I, I, Self::Scalar)],
49    ) -> anyhow::Result<Self::Mat>
50    where
51        I: TryInto<usize> + Copy,
52        <I as TryInto<usize>>::Error: std::fmt::Debug;
53
54    fn to_nonzero_triplets(&self) -> anyhow::Result<NRowNColTriplets<Self::Scalar>>;
55}
56
57pub struct NRowNColTriplets<Scalar> {
58    pub nrow: usize,
59    pub ncol: usize,
60    pub triplets: Vec<(usize, usize, Scalar)>,
61}
62
63/// Reading off from `Tensor`
64pub trait ConvertMatOps {
65    type Mat;
66    type Scalar;
67
68    fn from_tensor(_: &Tensor) -> anyhow::Result<Self::Mat>;
69    fn to_tensor(&self, dev: &Device) -> anyhow::Result<Tensor>;
70}
71
72/// normalize, sum_to_one, scale, and centre columns
73pub trait MatOps {
74    type Mat;
75    type Scalar;
76
77    /// make each column sum to 1
78    fn sum_to_one_columns_inplace(&mut self);
79    /// make each column sum to 1
80    fn sum_to_one_columns(&self) -> Self::Mat;
81
82    /// make each row sum to 1
83    fn sum_to_one_rows_inplace(&mut self);
84    /// make each row sum to 1
85    fn sum_to_one_rows(&self) -> Self::Mat;
86
87    /// normalize logits after taking exp `(log-sum-exp)`
88    fn normalize_exp_logits_columns_inplace(&mut self);
89    /// normalize logits after taking exp `(log-sum-exp)`
90    fn normalize_exp_logits_columns(&self) -> Self::Mat;
91
92    /// column-wise log-softmax: subtract each column's log-sum-exp so the
93    /// `exp` of each column sums to 1. Returns log-probabilities (unlike
94    /// [`Self::normalize_exp_logits_columns`], which returns probabilities).
95    fn log_softmax_columns_inplace(&mut self);
96    /// column-wise log-softmax (see [`Self::log_softmax_columns_inplace`])
97    fn log_softmax_columns(&self) -> Self::Mat;
98
99    /// vector norm for each column
100    fn normalize_columns_inplace(&mut self);
101    /// vector norm for each column
102    fn normalize_columns(&self) -> Self::Mat;
103
104    /// standardization for each column
105    fn scale_columns_inplace(&mut self);
106    /// standardization for each column
107    fn scale_columns(&self) -> Self::Mat;
108
109    /// standardization for each row
110    fn scale_rows_inplace(&mut self);
111    /// standardization for each row
112    fn scale_rows(&self) -> Self::Mat;
113
114    /// centering for each column
115    fn centre_columns_inplace(&mut self);
116    /// centering for each column
117    fn centre_columns(&self) -> Self::Mat;
118}
119
120pub trait AdjustByDivisionOp<Other, Scalar> {
121    /// Adjust each column with the column of the matching batch index
122    ///
123    /// Assume: `Y[g] ~ Poisson(X[g] * λ)`
124    /// (1) Estimate the λ parameter by taking overall ratio, namely,
125    /// `λ = Σ Y[g] / Σ X[g]`
126    ///
127    /// (2) Take the residual (in the log space)
128    /// `ln Y[g] - ln (λ X[g])` or `Y[g]/λX[g]` if `X[g] > 0`
129    /// otherwise, do nothing
130    fn adjust_by_division_of_selected_inplace(&mut self, denom_db: &Other, batches: &[usize]);
131
132    /// adjust each column with the corresponding column of the denom
133    ///
134    /// Assume: `Y[g] ~ Poisson(X[g] * λ)`
135    /// (1) Estimate the λ parameter by taking overall ratio, namely,
136    /// `λ = Σ Y[g] / Σ X[g]`
137    ///
138    /// (2) Take the residual (in the log space)
139    /// `ln Y[g] - ln (λ X[g])` or `Y[g]/λX[g]` if `X[g] > 0`
140    /// otherwise, do nothing
141    fn adjust_by_division_inplace(&mut self, denom: &Other);
142}
143
144pub trait MatElemOps {
145    type Mat;
146    type Scalar;
147    fn log1p_inplace(&mut self);
148    fn log1p(&self) -> Self::Mat;
149}
150
151/// Elementwise chains fused into ONE pass, because candle's CPU backend runs them
152/// one core at a time.
153///
154/// Only matmul reaches `gemm`, which candle drives with `Parallelism::Rayon`;
155/// `unary_map` and `binary_map` are plain serial iterators, and the vectorized
156/// `f32_vec` path is `#[cfg(feature = "mkl" / "accelerate")]` — SIMD, still one
157/// core. So a loop whose matmuls scale across every core stalls on the
158/// elementwise ops between them, and any chain over a large tensor is worth
159/// collapsing into a single rayon pass.
160///
161/// Implemented for `Tensor` on CPU only. Off CPU the device's own kernels are
162/// already parallel, and each method falls back to the op chain it stands in for
163/// — same numbers either way, which the tests assert bitwise.
164pub trait FusedTensorOps: Sized {
165    /// `exp(min(self + offset, ceiling))`, i.e. the Poisson rate from a linear
166    /// predictor with the overflow guard `exp` needs (f32 overflows at 88).
167    ///
168    /// Replaces `self.broadcast_add(offset)?.minimum(ceiling)?.exp()`. `self` is
169    /// `[N, F]`; `offset` is any shape that chain broadcasts against it — `[N, F]`,
170    /// `[1, F]` or `[N, 1]`.
171    ///
172    /// # The receiver must be unaliased
173    ///
174    /// On CPU this overwrites `self`'s storage and hands it back, so the whole
175    /// chain costs one buffer instead of three. `Tensor` is an `Arc`, so taking
176    /// `self` by value does **not** prove exclusivity — `x.reshape(..)` yields a
177    /// contiguous tensor sharing `x`'s storage and would pass every guard here,
178    /// silently overwriting `x`. Pass a freshly computed tensor (a `matmul`
179    /// result), never a `clone`, `narrow` or `reshape` of one still in use.
180    ///
181    /// Back-prop is unsupported by construction (candle's in-place custom ops
182    /// carry no backward), which is why the callers take their gradients in
183    /// closed form.
184    ///
185    /// Deliberately single-offset. A chain carrying **two** offsets — a `[N, 1]`
186    /// column and a `[1, F]` row, as the joint velocity solver in
187    /// `graph-embedding-util` does — needs its own method rather than a caller
188    /// pre-broadcasting one of them, which would cost the full-size op this
189    /// exists to remove.
190    fn clamped_exp_add_inplace(self, offset: &Tensor, ceiling: f64) -> anyhow::Result<Self>;
191}
192
193/// TF-IDF (Term Frequency–Inverse Document Frequency) transformation
194///
195/// A numerical statistic reflecting how important a word (term) is to a document
196/// in a collection or corpus. (Wikipedia)
197///
198/// Treats the matrix as a term-document matrix where:
199/// - Rows are "terms" (e.g., genes, words)
200/// - Columns are "documents" (e.g., cell types, text documents)
201///
202/// **TF-IDF(t, d) = TF(t, d) × IDF(t)**
203///
204/// where:
205/// - TF(t, d) = term frequency of term t in document d (matrix values)
206/// - IDF(t) = log(N / df(t)) = inverse document frequency
207/// - N = total number of documents (columns)
208/// - df(t) = document frequency = number of documents containing term t
209///
210/// Terms appearing in many documents get lower weight; terms specific to few
211/// documents get higher weight.
212pub trait TfIdfOps {
213    type Mat;
214
215    /// Apply TF-IDF transformation
216    ///
217    /// IDF(t) = log(N / (df(t) + 1)) where df(t) = number of non-zero entries in row t
218    fn tfidf(&self) -> Self::Mat;
219
220    /// Apply TF-IDF followed by L2 column normalization
221    ///
222    /// Useful for cosine similarity comparisons between documents (columns)
223    fn tfidf_normalize_columns(&self) -> Self::Mat;
224}
225
226/// Operations to sample random matrices, only works for
227/// `nalgebra::DMatrix` and `ndarray::Array2`
228pub trait SampleOps {
229    type Mat;
230    type Scalar;
231
232    /// Sample a matrix from a uniform distribution `U(0,1)`.
233    ///
234    /// Unseeded: draws fresh entropy each call. For reproducible output use
235    /// [`SampleOps::runif_seeded`].
236    fn runif(dd: usize, nn: usize) -> Self::Mat;
237
238    /// Sample a matrix from a normal distribution `N(0,1)`.
239    ///
240    /// Unseeded: draws fresh entropy each call. For reproducible output use
241    /// [`SampleOps::rnorm_seeded`].
242    fn rnorm(dd: usize, nn: usize) -> Self::Mat;
243
244    /// Sample a matrix from a gamma distribution with `param` is
245    /// `(shape α, scale θ)`
246    ///
247    /// $$f(x|\alpha,\theta) = \frac{\theta^{-\alpha}}{\Gamma(\alpha)} x^{\alpha - 1} e^{-x/\theta}$$
248    ///
249    /// Note: `rate = 1/scale` or $\beta = 1/\theta$
250    ///
251    /// Unseeded: draws fresh entropy each call. For reproducible output use
252    /// [`SampleOps::rgamma_seeded`].
253    fn rgamma(dd: usize, nn: usize, param: (f32, f32)) -> Self::Mat;
254
255    /// Seeded, thread-order-independent `U(0,1)` sample. Byte-identical across
256    /// runs, thread counts, and machines for a fixed `seed`. See
257    /// [`crate::matrix::rand_util`].
258    fn runif_seeded(dd: usize, nn: usize, seed: u64) -> Self::Mat;
259
260    /// Seeded, thread-order-independent `N(0,1)` sample. Byte-identical across
261    /// runs, thread counts, and machines for a fixed `seed`. See
262    /// [`crate::matrix::rand_util`].
263    fn rnorm_seeded(dd: usize, nn: usize, seed: u64) -> Self::Mat;
264
265    /// Seeded, thread-order-independent gamma sample (`param = (shape α, scale θ)`).
266    /// Byte-identical across runs, thread counts, and machines for a fixed
267    /// `seed`. See [`crate::matrix::rand_util`].
268    fn rgamma_seeded(dd: usize, nn: usize, param: (f32, f32), seed: u64) -> Self::Mat;
269}
270
271pub trait DistanceOps {
272    type Scalar;
273    type Other;
274
275    /// A vector of Euclidean distances between sources and targets `other`
276    ///
277    /// * `other`: other data matrix
278    fn euclidean_distance(
279        &self,
280        other: &Self::Other,
281    ) -> anyhow::Result<Vec<(usize, usize, Self::Scalar)>>;
282
283    /// A vector of Euclidean distances between sources and targets `other`
284    ///
285    /// * `other`: other data matrix
286    /// * `select_columns_in_other`: specific columns
287    fn euclidean_distance_on_select_columns(
288        &self,
289        other: &Self::Other,
290        select_columns_in_other: &[usize],
291    ) -> anyhow::Result<Vec<(usize, usize, Self::Scalar)>>;
292}
293
294pub trait EncodingOps
295where
296    Self: Sized,
297{
298    type Mat;
299    type Scalar;
300
301    /// Sinusoidal Positional Encoding
302    /// * `emb_dim` - embedding dimension, say `d`
303    /// * returns each column's embedding results (row x 2d)
304    ///
305    /// for each element r of each column c:
306    ///  ret[r, 2i] = sin(x[r,c]/10000^(2i/d))
307    ///  ret[r, 2i + 1] = cos(x[r,c]/10000^(2i/d))
308    /// where i in [0, d/2-1]
309    fn positional_embedding_columns(&self, emb_dim: usize) -> anyhow::Result<Self::Mat>;
310}
311
312/// Operations that involves multiple types
313pub trait CompositeOps {
314    type Scalar;
315    type Mat;
316    type Other;
317
318    /// `self[:,col] += other[:,col]`
319    /// * `other`: `CscMatrix`
320    /// * `col`: column index
321    fn add_assign_column(&mut self, other: &Self::Other, col: usize);
322
323    /// `self += other`
324    /// * `other`: `CscMatrix`
325    fn add_assign(&mut self, other: &Self::Other);
326}
327
328/// Read and write matrices from and to files
329pub trait IoOps {
330    type Scalar;
331    type Mat;
332
333    fn read_file_delim(
334        file_path: &str,
335        delim: impl Into<Delimiter>,
336        skip: Option<usize>,
337    ) -> anyhow::Result<Self::Mat>;
338
339    /// Read the data matrix with row and column names
340    ///
341    /// * `file_path`: data file name
342    /// * `delim`: delimiter (`char` vector or string)
343    /// * `header_row`: header line (0-based); `None` will find no header
344    /// * `row_name_column_index`: column index (0-based) corresponds to row name
345    /// * `select_column_indices`: column indices (0-based) to include
346    /// * `select_column_names`: column names to include
347    ///
348    fn read_data(
349        file_path: &str,
350        delim: impl Into<Delimiter>,
351        header_row: Option<usize>,
352        row_name_column_index: Option<usize>,
353        select_column_indices: Option<&[usize]>,
354        select_column_names: Option<&[Box<str>]>,
355    ) -> anyhow::Result<MatWithNames<Self::Mat>>;
356
357    /// Read the data matrix with row and column names
358    ///
359    /// * `file_path`: data file name
360    /// * `delim`: delimiter (`char` vector or string)
361    /// * `header_row`: header line (0-based); `None` will find no header
362    /// * `header_column`: column index (0-based) corresponds to row name
363    ///
364    fn read_data_with_names(
365        file_path: &str,
366        delim: impl Into<Delimiter>,
367        header_row: Option<usize>,
368        header_column: Option<usize>,
369    ) -> anyhow::Result<MatWithNames<Self::Mat>> {
370        Self::read_data(file_path, delim, header_row, header_column, None, None)
371    }
372
373    #[allow(clippy::type_complexity)]
374    fn read_data_vec_with_indices_names(
375        file_path: &str,
376        delim: impl Into<Delimiter>,
377        header_line: Option<usize>,
378        row_name_index: Option<usize>,
379        column_indices: Option<&[usize]>,
380        column_names: Option<&[Box<str>]>,
381    ) -> anyhow::Result<(Vec<Box<str>>, Vec<Box<str>>, Vec<Self::Scalar>)>
382    where
383        Self::Scalar: std::str::FromStr,
384        <Self::Scalar as std::str::FromStr>::Err: std::fmt::Debug,
385    {
386        let hdr_line = match header_line {
387            Some(skip) => skip as i64,
388            None => -1, // no skipping
389        };
390
391        let ReadLinesOut { mut lines, header } =
392            crate::matrix::common_io::read_lines_of_words_delim(file_path, delim, hdr_line)?;
393
394        // A blank line tokenizes to one empty field, not to zero fields, so it
395        // must be dropped here or it caps the width check below at 1 and then
396        // breaks the value loop. Dropping it up front fixes both at once.
397        lines.retain(|w| !(w.len() == 1 && w[0].is_empty()));
398
399        let data_width = lines.iter().map(|w| w.len()).min().unwrap_or(header.len());
400        // R's write.table omits a name for the row-label column, so the header
401        // is one field short of the data rows and every header position names
402        // the data column one to its RIGHT. Detect that shape once; both the
403        // name matching and the naming lookup below shift through it.
404        let header_offset = usize::from(
405            !header.is_empty() && header.len() + 1 == data_width && row_name_index == Some(0),
406        );
407
408        let mut relevant_indices: Vec<usize> = vec![];
409
410        let indices_given = column_indices.is_some_and(|ix| !ix.is_empty());
411        if let Some(indices) = column_indices {
412            relevant_indices.extend(indices.iter().copied());
413        }
414
415        // Explicit indices OVERRIDE names, as the callers' help documents; a
416        // union would quietly widen the selection with every default name that
417        // happens to be present in the header.
418        if !indices_given {
419            if let Some(names) = column_names {
420                // The tokenizer has already unquoted both sides.
421                let name_indices: Vec<usize> = header
422                    .iter()
423                    .enumerate()
424                    .filter_map(|(i, name)| {
425                        if names.iter().any(|n| n == name) {
426                            Some(i + header_offset)
427                        } else {
428                            None
429                        }
430                    })
431                    .collect();
432                relevant_indices.extend(name_indices);
433            }
434        }
435
436        // Neither selector given: take EVERY column except the row-name one.
437        // Without this the selection stays empty and the reader silently returns
438        // a 0-column matrix, so `read_data(.., None, None)` — the form used by
439        // `senna`'s `read_mat` and `data-beans-sim`'s topic-file loader — could
440        // never read a delimited file at all.
441        if column_indices.is_none() && column_names.is_none() {
442            let n_col = header
443                .len()
444                .max(lines.first().map_or(0, |words| words.len()));
445            relevant_indices.extend((0..n_col).filter(|j| Some(*j) != row_name_index));
446        }
447
448        relevant_indices.sort_unstable();
449        relevant_indices.dedup();
450
451        // Every subscript below is checked here first: the row-name column, the
452        // selected data columns, and the header lookup that names them. Width is
453        // the NARROWEST data row, ignoring blank ones, because a ragged file
454        // otherwise passes this and panics later in the value loop. The header
455        // is checked separately, since it can be one field short of the data
456        // rows when a writer omits a name for the row-label column.
457        let n_columns = data_width;
458        let mut to_check: Vec<usize> = relevant_indices.clone();
459        to_check.extend(row_name_index);
460        if let Some(&bad) = to_check.iter().find(|&&j| j >= n_columns) {
461            return Err(anyhow::anyhow!(
462                "column index {bad} is out of range: the file has {n_columns} column(s). \
463                 Name the columns to read, or pass indices within range."
464            ));
465        }
466
467        let row_names: Vec<Box<str>> = match row_name_index {
468            // Unquoted, for the same reason the header is: a fully-quoted csv
469            // would otherwise yield row names carrying their quotes, which then
470            // match nothing when joined against a matrix's own names.
471            Some(row_name_index) => lines
472                .iter()
473                .map(|words| words[row_name_index].clone())
474                .collect(),
475            _ => (0..lines.len())
476                .map(|x| x.to_string().into_boxed_str())
477                .collect(),
478        };
479
480        // Indices can come from a caller's fallback rather than from a name
481        // match, so they are not guaranteed to exist in this file. Say which
482        // column was asked for and how many the file has, instead of panicking
483        // on the subscript several lines later.
484        let column_names: Vec<Box<str>> = if header.is_empty() {
485            relevant_indices
486                .iter()
487                .map(|x| x.to_string().into_boxed_str())
488                .collect()
489        } else {
490            relevant_indices
491                .iter()
492                // A header can be narrower than the data rows; fall back to the
493                // position rather than panicking on a name that was never written.
494                .map(|&j| {
495                    j.checked_sub(header_offset)
496                        .and_then(|k| header.get(k))
497                        .cloned()
498                        .unwrap_or_else(|| j.to_string().into_boxed_str())
499                })
500                .collect()
501        };
502
503        let data: Vec<Vec<Self::Scalar>> = lines
504            .iter()
505            .map(|words| {
506                relevant_indices
507                    .iter()
508                    .map(|&i| words[i].parse::<Self::Scalar>().expect("failed to parse"))
509                    .collect()
510            })
511            .collect();
512
513        let data = data.into_iter().flatten().collect::<Vec<_>>();
514
515        Ok((row_names, column_names, data))
516    }
517
518    /// Read a `tsv` file while skipping until the header row
519    fn from_tsv(tsv_file: &str, skip: Option<usize>) -> anyhow::Result<Self::Mat> {
520        Self::read_file_delim(tsv_file, "\t", skip)
521    }
522
523    /// Read a `csv` file while skipping until the header row
524    fn from_csv(csv_file: &str, skip: Option<usize>) -> anyhow::Result<Self::Mat> {
525        Self::read_file_delim(csv_file, ",", skip)
526    }
527
528    /// write the matrix down to a file with delimiter
529    /// * `file_path`: output file path
530    /// * `delim`: separation character or string
531    fn write_file_delim(&self, file: &str, delim: &str) -> anyhow::Result<()>;
532
533    /// write the matrix down to a tsv file
534    /// * `file_path`: output file path
535    fn to_tsv(&self, tsv_file: &str) -> anyhow::Result<()> {
536        self.write_file_delim(tsv_file, "\t")
537    }
538
539    /// write the matrix down to a csv file
540    /// * `file_path`: output file path
541    fn to_csv(&self, csv_file: &str) -> anyhow::Result<()> {
542        self.write_file_delim(csv_file, ",")
543    }
544
545    /// write the matrix down to parquet with full control over naming
546    /// * `file_path`: output file path
547    /// * `row_names`: Tuple of (optional row_names, optional row_column_name)
548    ///   - `(None, None)`: use numeric row names `[0, n)` with "row" column name
549    ///   - `(None, Some("cell_pair"))`: use numeric row names with "cell_pair" column name
550    ///   - `(Some(names), None)`: use provided names with "row" column name
551    ///   - `(Some(names), Some("gene"))`: use provided names with "gene" column name
552    /// * `column_names`: if `None`, just add `[0, n)` numbers.
553    fn to_parquet_with_names(
554        &self,
555        file_path: &str,
556        row_names: (Option<&[Box<str>]>, Option<&str>),
557        column_names: Option<&[Box<str>]>,
558    ) -> anyhow::Result<()>;
559
560    /// write the matrix down to parquet with default names
561    /// * `file_path`: output file path
562    ///   Uses numeric row/column names and default "row" column name
563    fn to_parquet(&self, file_path: &str) -> anyhow::Result<()> {
564        self.to_parquet_with_names(file_path, (None, None), None)
565    }
566
567    /// Read a real-valued numeric matrix with the default row
568    /// index(0) and all the other available columns.
569    /// Assumes column 0 contains row names.
570    ///
571    fn from_parquet(file_path: &str) -> anyhow::Result<MatWithNames<Self::Mat>> {
572        Self::from_parquet_with_indices(file_path, Some(0), None)
573    }
574
575    /// Read a real-valued numeric matrix treating all columns as data.
576    /// Row names will be generated as "0", "1", "2", ...
577    ///
578    fn from_parquet_no_row_names(file_path: &str) -> anyhow::Result<MatWithNames<Self::Mat>> {
579        Self::from_parquet_with_indices(file_path, None, None)
580    }
581
582    /// Read a real-valued numeric matrix from the parquet file. We
583    /// can specify row name index. We can specify the row name column
584    /// index and desired column indices.
585    /// * `row_name_index`: column index (0-based) corresponds to row name
586    fn from_parquet_with_row_names(
587        file_path: &str,
588        row_name_index: Option<usize>,
589    ) -> anyhow::Result<MatWithNames<Self::Mat>> {
590        Self::from_parquet_with_indices_names(file_path, row_name_index, None, None)
591    }
592
593    /// Read a real-valued numeric matrix from the parquet file. We
594    /// can specify row name index. We can specify the row name column
595    /// index and desired column indices.
596    /// * `row_name_index`: column index (0-based) corresponds to row name
597    /// * `column_indices`: column indices (0-based) to include
598    fn from_parquet_with_indices(
599        file_path: &str,
600        row_name_index: Option<usize>,
601        column_indices: Option<&[usize]>,
602    ) -> anyhow::Result<MatWithNames<Self::Mat>> {
603        Self::from_parquet_with_indices_names(file_path, row_name_index, column_indices, None)
604    }
605
606    /// Read a real-valued numeric matrix from the parquet file.  We
607    /// can specify row name index.  We can specify the row name
608    /// column index and desired column names.
609    /// * `row_name_index`: column index (0-based) corresponds to row name
610    /// * `column_names`: column names to include
611    fn from_parquet_with_names(
612        file_path: &str,
613        row_name_index: Option<usize>,
614        column_names: Option<&[Box<str>]>,
615    ) -> anyhow::Result<MatWithNames<Self::Mat>> {
616        Self::from_parquet_with_indices_names(file_path, row_name_index, None, column_names)
617    }
618
619    /// Read a real-valued numeric matrix from the parquet file.  We
620    /// can specify row name index.  We can specify the row name
621    /// column index and desired column indices and names.
622    /// * `row_name_index`: column index (0-based) corresponds to row name
623    /// * `column_indices`: column indices (0-based) to include
624    /// * `column_names`: column names to include
625    fn from_parquet_with_indices_names(
626        file_path: &str,
627        row_name_index: Option<usize>,
628        column_indices: Option<&[usize]>,
629        column_names: Option<&[Box<str>]>,
630    ) -> anyhow::Result<MatWithNames<Self::Mat>>;
631}
632
633/// intput data matrix `mat` with `rows` and `cols`
634pub struct MatWithNames<M> {
635    pub rows: Vec<Box<str>>,
636    pub cols: Vec<Box<str>>,
637    pub mat: M,
638}
639
640pub trait MeltOps {
641    type Scalar;
642    type Mat;
643    /// melt a matrix with indices
644    fn melt_with_indexes(&self) -> (Vec<Self::Scalar>, Vec<Vec<usize>>);
645    /// melt a matrix
646    fn melt(&self) -> Vec<Self::Scalar>;
647    /// Melt multiple matrices/tensors together in a single traversal for cache efficiency.
648    /// All inputs must have the same dimensions.
649    /// Returns (values for each input, indices for each dimension).
650    fn melt_many_with_indexes(&self, others: &[&Self])
651        -> (Vec<Vec<Self::Scalar>>, Vec<Vec<usize>>);
652}
653
654pub trait CandleDataLoaderOps {
655    type Scalar;
656    type Mat;
657    // /// unify transpose
658    // fn transpose(&self) -> Self::Mat;
659    /// take each row vector as a sample
660    fn rows_to_tensor_vec(&self) -> Vec<Tensor>;
661
662    /// Return (nrows, ncols) dimensions
663    fn data_shape(&self) -> (usize, usize);
664
665    /// Extract row i as Vec<f32>.
666    ///
667    /// WARNING: default creates ALL row tensors then picks one — O(N*D) for O(D) work.
668    /// Implementors should override this.
669    fn row_to_f32_vec(&self, i: usize) -> Vec<f32> {
670        let t = &self.rows_to_tensor_vec()[i];
671        t.flatten_all().unwrap().to_vec1::<f32>().unwrap()
672    }
673}