cuvs/
pca.rs

1use std::{ptr, ptr::NonNull};
2
3use cudf::Exec;
4
5use crate::{
6    device::{DeviceBuffer, DeviceTensorMut, MatrixLayout, OwnedMatrix},
7    dlpack::DlpackTensor,
8    error::{Result, check_cuvs, from_cudf, invalid_input, invalid_state, require_nonzero},
9    resources::CuvsResources,
10};
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum PcaSolver {
14    CovarianceDivideAndConquer,
15    CovarianceJacobi,
16}
17
18impl PcaSolver {
19    fn raw(self) -> cuvs_sys::cuvsPcaSolver {
20        match self {
21            Self::CovarianceDivideAndConquer => cuvs_sys::cuvsPcaSolver::CUVS_PCA_COV_EIG_DQ,
22            Self::CovarianceJacobi => cuvs_sys::cuvsPcaSolver::CUVS_PCA_COV_EIG_JACOBI,
23        }
24    }
25}
26
27#[derive(Debug, Clone)]
28pub struct PcaParams {
29    pub n_components: usize,
30    pub copy: bool,
31    pub whiten: bool,
32    pub solver: PcaSolver,
33    pub tolerance: f32,
34    pub iterations: i32,
35    pub flip_signs_based_on_u: bool,
36}
37
38impl PcaParams {
39    pub fn new(n_components: usize) -> Self {
40        Self {
41            n_components,
42            copy: true,
43            whiten: false,
44            solver: PcaSolver::CovarianceDivideAndConquer,
45            tolerance: 1e-7,
46            iterations: 15,
47            flip_signs_based_on_u: false,
48        }
49    }
50}
51
52/// Caller-owned device outputs for one PCA fit-transform operation.
53#[derive(Debug)]
54pub struct PcaOutputBuffers {
55    pub transformed: DeviceBuffer<f32>,
56    pub components: DeviceBuffer<f32>,
57    pub explained_variance: DeviceBuffer<f32>,
58    pub explained_variance_ratio: DeviceBuffer<f32>,
59    pub singular_values: DeviceBuffer<f32>,
60    pub mean: DeviceBuffer<f32>,
61    pub noise_variance: DeviceBuffer<f32>,
62}
63
64impl PcaOutputBuffers {
65    pub fn new(resources: &CuvsResources, rows: usize, columns: usize, n_components: usize) -> Result<Self> {
66        let operation = "PcaOutputBuffers::new";
67        require_nonzero(rows, operation, "rows")?;
68        require_nonzero(columns, operation, "columns")?;
69        require_nonzero(n_components, operation, "n_components")?;
70        let transformed_len = rows
71            .checked_mul(n_components)
72            .ok_or_else(|| invalid_input(operation, "transformed output length overflowed usize"))?;
73        let components_len = n_components
74            .checked_mul(columns)
75            .ok_or_else(|| invalid_input(operation, "components output length overflowed usize"))?;
76        Ok(Self {
77            transformed: DeviceBuffer::new(resources, transformed_len)?,
78            components: DeviceBuffer::new(resources, components_len)?,
79            explained_variance: DeviceBuffer::new(resources, n_components)?,
80            explained_variance_ratio: DeviceBuffer::new(resources, n_components)?,
81            singular_values: DeviceBuffer::new(resources, n_components)?,
82            mean: DeviceBuffer::new(resources, columns)?,
83            noise_variance: DeviceBuffer::new(resources, 1)?,
84        })
85    }
86
87    /// Consume the output owner and split the column-major transformed matrix
88    /// into independent cudf-owned device columns.
89    pub fn into_transformed_columns(self, fit: PcaFitTransform) -> Result<Vec<cudf::Series>> {
90        let operation = "PcaOutputBuffers::into_transformed_columns";
91        validate_output_lengths(&self, fit.rows, fit.columns, fit.n_components, operation)?;
92        if !self.transformed.native_input_validated() {
93            return Err(invalid_state(
94                operation,
95                "transformed PCA output has not passed post-native finite validation",
96            ));
97        }
98        let expected = fit
99            .rows
100            .checked_mul(fit.n_components)
101            .ok_or_else(|| invalid_input(operation, "transformed output length overflowed usize"))?;
102        if self.transformed.len() != expected {
103            return Err(invalid_input(
104                operation,
105                format!("transformed output contains {} elements, expected {expected}", self.transformed.len()),
106            ));
107        }
108
109        let flat =
110            cudf::Series::from_column(self.transformed.into_column()?).map_err(|error| from_cudf(operation, error))?;
111        (0..fit.n_components)
112            .map(|component| {
113                flat.on(Exec::RETAINED)
114                    .slice(component * fit.rows, fit.rows)
115                    .map_err(|error| from_cudf(operation, error))
116            })
117            .collect()
118    }
119
120    /// Consume the output owner and return the transformed PCA relation as a
121    /// cudf DataFrame with `component_0..N` columns.
122    pub fn into_transformed_dataframe(self, fit: PcaFitTransform) -> Result<cudf::DataFrame> {
123        let operation = "PcaOutputBuffers::into_transformed_dataframe";
124        let columns = self.into_transformed_columns(fit)?;
125        cudf::DataFrame::from_series(
126            columns.into_iter().enumerate().map(|(component, column)| (format!("component_{component}"), column)),
127        )
128        .map_err(|error| from_cudf(operation, error))
129    }
130
131    fn validate_native_outputs(&mut self, operation: &'static str) -> Result<()> {
132        self.transformed.validate_native_input(operation)?;
133        self.components.validate_native_input(operation)?;
134        self.explained_variance.validate_native_input(operation)?;
135        self.explained_variance_ratio.validate_native_input(operation)?;
136        self.singular_values.validate_native_input(operation)?;
137        self.mean.validate_native_input(operation)?;
138        self.noise_variance.validate_native_input(operation)?;
139        Ok(())
140    }
141}
142
143#[derive(Debug, Clone, Copy, PartialEq, Eq)]
144pub struct PcaFitTransform {
145    rows: usize,
146    columns: usize,
147    n_components: usize,
148}
149
150impl PcaFitTransform {
151    pub fn rows(&self) -> usize {
152        self.rows
153    }
154
155    pub fn columns(&self) -> usize {
156        self.columns
157    }
158
159    pub fn n_components(&self) -> usize {
160        self.n_components
161    }
162}
163
164struct ParamsHandle {
165    raw: Option<NonNull<cuvs_sys::cuvsPcaParams>>,
166}
167
168impl ParamsHandle {
169    fn new(params: &PcaParams, operation: &'static str) -> Result<Self> {
170        validate_params(params, operation)?;
171        let mut raw = ptr::null_mut();
172        unsafe { check_cuvs(cuvs_sys::cuvsPcaParamsCreate(&mut raw), operation)? };
173        let raw = NonNull::new(raw).ok_or_else(|| invalid_state(operation, "cuVS returned null PCA params"))?;
174        unsafe {
175            let native = raw.as_ptr();
176            (*native).n_components = params.n_components as i32;
177            (*native).copy = params.copy;
178            (*native).whiten = params.whiten;
179            (*native).algorithm = params.solver.raw();
180            (*native).tol = params.tolerance;
181            (*native).n_iterations = params.iterations;
182        }
183        Ok(Self { raw: Some(raw) })
184    }
185
186    fn as_ptr(&self, operation: &'static str) -> Result<cuvs_sys::cuvsPcaParams_t> {
187        self.raw.map(NonNull::as_ptr).ok_or_else(|| invalid_state(operation, "PCA params have been closed"))
188    }
189
190    fn close(&mut self, operation: &'static str) -> Result<()> {
191        let Some(raw) = self.raw.take() else {
192            return Ok(());
193        };
194        let result = unsafe { check_cuvs(cuvs_sys::cuvsPcaParamsDestroy(raw.as_ptr()), operation) };
195        if result.is_err() {
196            self.raw = Some(raw);
197        }
198        result
199    }
200}
201
202impl Drop for ParamsHandle {
203    fn drop(&mut self) {
204        if let Some(raw) = self.raw.take() {
205            unsafe {
206                let _ = cuvs_sys::cuvsPcaParamsDestroy(raw.as_ptr());
207            }
208        }
209    }
210}
211
212pub(crate) fn fit_transform_into(
213    resources: &CuvsResources,
214    input: &mut DeviceTensorMut<'_, f32>,
215    params: &PcaParams,
216    outputs: &mut PcaOutputBuffers,
217) -> Result<PcaFitTransform> {
218    let operation = "pca::fit_transform_into";
219    resources.validate_exact_binding(operation)?;
220    if input.binding_id() != resources.binding_id() {
221        return Err(invalid_input(operation, "PCA input and cuVS resources must use the same exact runtime binding")
222            .with_context("tensor_binding_id", input.binding_id())
223            .with_context("resource_binding_id", resources.binding_id()));
224    }
225    if input.device() != resources.device().ordinal() as i32 {
226        return Err(invalid_input(operation, "PCA input and cuVS resources must use the same CUDA device")
227            .with_context("tensor_device", input.device())
228            .with_context("resource_device", resources.device().ordinal()));
229    }
230    if matrix_layout_mut(input, operation)? != MatrixLayout::ColumnMajor {
231        return Err(invalid_input(operation, "PCA input must be a contiguous mutable column-major matrix"));
232    }
233    validate_params(params, operation)?;
234    let rows = input.rows()?;
235    let columns = input.columns()?;
236    validate_input_dimensions(rows, columns, params.n_components, operation)?;
237    validate_output_lengths(outputs, rows, columns, params.n_components, operation)?;
238
239    let transformed_strides = [1, rows];
240    let components_strides = [1, params.n_components];
241    let PcaOutputBuffers {
242        transformed,
243        components,
244        explained_variance,
245        explained_variance_ratio,
246        singular_values,
247        mean,
248        noise_variance,
249    } = outputs;
250    let mut transformed_view =
251        transformed.as_tensor_mut(resources, &[rows, params.n_components], Some(&transformed_strides), 0)?;
252    let mut components_view =
253        components.as_tensor_mut(resources, &[params.n_components, columns], Some(&components_strides), 0)?;
254    let mut explained_variance_view = explained_variance.as_tensor_mut(resources, &[params.n_components], None, 0)?;
255    let mut explained_variance_ratio_view =
256        explained_variance_ratio.as_tensor_mut(resources, &[params.n_components], None, 0)?;
257    let mut singular_values_view = singular_values.as_tensor_mut(resources, &[params.n_components], None, 0)?;
258    let mut mean_view = mean.as_tensor_mut(resources, &[columns], None, 0)?;
259    let mut noise_variance_view = noise_variance.as_tensor_mut(resources, &[1], None, 0)?;
260
261    resources.with_native(operation, |native| {
262        let mut params_handle = ParamsHandle::new(params, operation)?;
263        let mut input_tensor = DlpackTensor::from_mut(input);
264        let mut transformed_tensor = DlpackTensor::from_mut(&mut transformed_view);
265        let mut components_tensor = DlpackTensor::from_mut(&mut components_view);
266        let mut explained_variance_tensor = DlpackTensor::from_mut(&mut explained_variance_view);
267        let mut explained_variance_ratio_tensor = DlpackTensor::from_mut(&mut explained_variance_ratio_view);
268        let mut singular_values_tensor = DlpackTensor::from_mut(&mut singular_values_view);
269        let mut mean_tensor = DlpackTensor::from_mut(&mut mean_view);
270        let mut noise_variance_tensor = DlpackTensor::from_mut(&mut noise_variance_view);
271        let call_result = unsafe {
272            check_cuvs(
273                cuvs_sys::cuvsPcaFitTransform(
274                    native,
275                    params_handle.as_ptr(operation)?,
276                    input_tensor.as_mut_ptr(),
277                    transformed_tensor.as_mut_ptr(),
278                    components_tensor.as_mut_ptr(),
279                    explained_variance_tensor.as_mut_ptr(),
280                    explained_variance_ratio_tensor.as_mut_ptr(),
281                    singular_values_tensor.as_mut_ptr(),
282                    mean_tensor.as_mut_ptr(),
283                    noise_variance_tensor.as_mut_ptr(),
284                    params.flip_signs_based_on_u,
285                ),
286                operation,
287            )
288        };
289        match (call_result, params_handle.close(operation)) {
290            (Err(error), _) => Err(error),
291            (Ok(()), close) => close,
292        }
293    })?;
294    outputs.validate_native_outputs(operation)?;
295
296    Ok(PcaFitTransform { rows, columns, n_components: params.n_components })
297}
298
299fn validate_input_dimensions(rows: usize, columns: usize, n_components: usize, operation: &'static str) -> Result<()> {
300    if rows < 2 {
301        return Err(invalid_input(operation, format!("PCA input must contain at least two rows, found {rows}"))
302            .with_context("rows", rows)
303            .with_context("minimum_rows", 2));
304    }
305    if columns < 2 {
306        return Err(invalid_input(
307            operation,
308            format!("PCA input must contain at least two feature columns, found {columns}"),
309        )
310        .with_context("feature_columns", columns)
311        .with_context("minimum_feature_columns", 2));
312    }
313    if n_components > columns {
314        return Err(invalid_input(
315            operation,
316            format!("n_components ({n_components}) must be <= input columns ({columns})"),
317        ));
318    }
319    Ok(())
320}
321
322/// Run PCA against an exclusively owned packed matrix.
323///
324/// Accepts exclusively owned packed inputs for in-place native execution.
325pub fn fit_transform_owned(
326    resources: &CuvsResources,
327    input: &mut OwnedMatrix<f32>,
328    params: &PcaParams,
329    outputs: &mut PcaOutputBuffers,
330) -> Result<PcaFitTransform> {
331    let operation = "pca::fit_transform_owned";
332    if !input.is_input_validated() {
333        return Err(invalid_input(operation, "PCA input must carry finite-input validation"));
334    }
335    if input.layout() != MatrixLayout::ColumnMajor {
336        return Err(invalid_input(operation, "PCA owned input must be a contiguous column-major matrix"));
337    }
338    validate_params(params, operation)?;
339    validate_input_dimensions(input.rows(), input.columns(), params.n_components, operation)?;
340    validate_output_lengths(outputs, input.rows(), input.columns(), params.n_components, operation)?;
341    let mut tensor = input.as_tensor_mut(resources)?;
342    fit_transform_into(resources, &mut tensor, params, outputs).map_err(|error| error.with_operation(operation))
343}
344
345fn validate_params(params: &PcaParams, operation: &'static str) -> Result<()> {
346    require_nonzero(params.n_components, operation, "n_components")?;
347    if params.n_components > i32::MAX as usize {
348        return Err(invalid_input(operation, "n_components exceeds the cuVS i32 parameter range"));
349    }
350    if !params.tolerance.is_finite() || params.tolerance < 0.0 {
351        return Err(invalid_input(operation, "PCA tolerance must be finite and non-negative"));
352    }
353    if params.iterations <= 0 {
354        return Err(invalid_input(operation, "PCA iterations must be greater than zero"));
355    }
356    Ok(())
357}
358
359fn validate_output_lengths(
360    outputs: &PcaOutputBuffers,
361    rows: usize,
362    columns: usize,
363    n_components: usize,
364    operation: &'static str,
365) -> Result<()> {
366    let transformed = rows
367        .checked_mul(n_components)
368        .ok_or_else(|| invalid_input(operation, "transformed output length overflowed usize"))?;
369    let components = n_components
370        .checked_mul(columns)
371        .ok_or_else(|| invalid_input(operation, "components output length overflowed usize"))?;
372    let actual = [
373        ("transformed", outputs.transformed.len(), transformed),
374        ("components", outputs.components.len(), components),
375        ("explained_variance", outputs.explained_variance.len(), n_components),
376        ("explained_variance_ratio", outputs.explained_variance_ratio.len(), n_components),
377        ("singular_values", outputs.singular_values.len(), n_components),
378        ("mean", outputs.mean.len(), columns),
379        ("noise_variance", outputs.noise_variance.len(), 1),
380    ];
381    for (name, actual, expected) in actual {
382        if actual != expected {
383            return Err(invalid_input(operation, format!("{name} contains {actual} elements, expected {expected}")));
384        }
385    }
386    Ok(())
387}
388
389fn matrix_layout_mut(matrix: &DeviceTensorMut<'_, f32>, operation: &'static str) -> Result<MatrixLayout> {
390    let rows = matrix.rows()? as i64;
391    let columns = matrix.columns()? as i64;
392    match matrix.strides() {
393        None => Ok(MatrixLayout::RowMajor),
394        Some([row_stride, column_stride]) if *row_stride == columns && *column_stride == 1 => {
395            Ok(MatrixLayout::RowMajor)
396        },
397        Some([row_stride, column_stride]) if *row_stride == 1 && *column_stride == rows => {
398            Ok(MatrixLayout::ColumnMajor)
399        },
400        _ => Err(invalid_input(operation, "PCA matrix must be contiguous row-major or column-major")),
401    }
402}
403
404#[cfg(test)]
405mod tests {
406    use super::*;
407
408    #[test]
409    fn input_dimensions_reject_a_single_row_before_native_pca() {
410        let error = validate_input_dimensions(1, 2, 1, "test").expect_err("PCA requires at least two rows");
411
412        assert_eq!(error.kind(), crate::ErrorKind::InvalidInput);
413        assert!(error.message().contains("at least two rows"));
414        assert_eq!(error.context_value("rows"), Some("1"));
415    }
416
417    #[test]
418    fn input_dimensions_accept_the_native_boundary() {
419        validate_input_dimensions(2, 2, 2, "test").expect("two rows and two columns satisfy the native contract");
420    }
421}
cuvs/pca.rs

cuvs/
pca.rs
