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