rlx-umap 0.2.4

Parametric UMAP on RLX — full fit/transform API and k-NN building blocks.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, version 3.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program. If not, see <https://www.gnu.org/licenses/>.

//! High-level `Umap::fit` API (parametric UMAP on RLX).

use rlx_driver::Device;
use rlx_runtime::device_ext;

use crate::config::UmapConfig;
use crate::fitted::FittedUmap;
use crate::train::EpochProgress;
use crate::training::{FitOptions, fit_with_progress};

/// Parametric UMAP trainer (mirrors fast-umap `Umap<B>`).
pub struct Umap {
    config: UmapConfig,
    device: Device,
}

impl Umap {
    pub fn new(config: UmapConfig) -> Self {
        Self {
            config,
            device: Device::Cpu,
        }
    }

    pub fn with_device(config: UmapConfig, device: Device) -> Self {
        assert!(
            device_ext::is_available(device),
            "device {device:?} is not available"
        );
        Self { config, device }
    }

    pub fn config(&self) -> &UmapConfig {
        &self.config
    }

    /// Fit parametric UMAP on `data` (`n_samples × n_features`).
    pub fn fit(self, data: Vec<Vec<f64>>) -> FittedUmap {
        self.fit_with_signal(data, crate::interrupt::install_ctrlc_handler())
    }

    /// Fit with a progress callback (invoked when loss is read back).
    pub fn fit_with_progress(
        self,
        data: Vec<Vec<f64>>,
        on_progress: impl Fn(EpochProgress) + Send + 'static,
    ) -> FittedUmap {
        self.fit_with_signal_and_progress(
            data,
            crate::interrupt::install_ctrlc_handler(),
            on_progress,
        )
    }

    /// Fit with an external cancellation channel (e.g. Ctrl-C).
    pub fn fit_with_signal(
        self,
        data: Vec<Vec<f64>>,
        exit_rx: crossbeam_channel::Receiver<()>,
    ) -> FittedUmap {
        self.fit_with_signal_and_progress(data, exit_rx, |_| {})
    }

    fn fit_with_signal_and_progress(
        self,
        data: Vec<Vec<f64>>,
        exit_rx: crossbeam_channel::Receiver<()>,
        on_progress: impl Fn(EpochProgress) + Send + 'static,
    ) -> FittedUmap {
        let options = FitOptions {
            device: self.device,
            exit_rx: Some(exit_rx),
        };
        fit_with_progress(self.config, data, options, on_progress)
    }
}