Skip to main content

legume_numeric/matrix/
traits.rs

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