Initial commit

crates/training/rtx-preprocessing/src/scalers/minmax_scaler.rs

@@ -0,0 +1,308 @@
+use crate::{InvertibleTransformer, PreprocessingError, Result, Transformer};
+use rtx_tensor::Tensor;
+use serde::{Deserialize, Serialize};
+
+/// MinMaxScaler: Transform features by scaling each feature to a given range.
+///
+/// This estimator scales and translates each feature individually such that it is in
+/// the given range on the training set, e.g. between zero and one.
+///
+/// The transformation is given by:
+/// X_std = (X - X.min(axis=0)) / (X.max(axis=0) - X.min(axis=0))
+/// X_scaled = X_std * (max - min) + min
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct MinMaxScaler {
+    /// The desired range of transformed data
+    feature_range: (f32, f32),
+    /// Minimum values for each feature (fitted)
+    data_min_: Option<Vec<f32>>,
+    /// Maximum values for each feature (fitted)
+    data_max_: Option<Vec<f32>>,
+    /// Range (max - min) for each feature (fitted)
+    data_range_: Option<Vec<f32>>,
+    /// Number of features seen during fit
+    n_features_: Option<usize>,
+}
+
+impl MinMaxScaler {
+    /// Create a new MinMaxScaler with default range [0, 1].
+    pub fn new() -> Self {
+        Self {
+            feature_range: (0.0, 1.0),
+            data_min_: None,
+            data_max_: None,
+            data_range_: None,
+            n_features_: None,
+        }
+    }
+
+    /// Create a MinMaxScaler with custom range.
+    ///
+    /// # Arguments
+    /// * `min` - Minimum value of the target range
+    /// * `max` - Maximum value of the target range
+    ///
+    /// # Panics
+    /// Panics if min >= max
+    pub fn with_range(min: f32, max: f32) -> Self {
+        assert!(min < max, "min must be less than max, got min={min}, max={max}");
+
+        Self {
+            feature_range: (min, max),
+            data_min_: None,
+            data_max_: None,
+            data_range_: None,
+            n_features_: None,
+        }
+    }
+
+    /// Get the target feature range.
+    pub fn feature_range(&self) -> (f32, f32) {
+        self.feature_range
+    }
+
+    /// Get the minimum values for each feature.
+    ///
+    /// # Panics
+    /// Panics if the scaler has not been fitted.
+    pub fn data_min(&self) -> &[f32] {
+        self.data_min_
+            .as_ref()
+            .expect("MinMaxScaler has not been fitted")
+    }
+
+    /// Get the maximum values for each feature.
+    ///
+    /// # Panics
+    /// Panics if the scaler has not been fitted.
+    pub fn data_max(&self) -> &[f32] {
+        self.data_max_
+            .as_ref()
+            .expect("MinMaxScaler has not been fitted")
+    }
+
+    /// Get the range (max - min) for each feature.
+    ///
+    /// # Panics
+    /// Panics if the scaler has not been fitted.
+    pub fn data_range(&self) -> &[f32] {
+        self.data_range_
+            .as_ref()
+            .expect("MinMaxScaler has not been fitted")
+    }
+
+    /// Compute min, max, and range for each feature.
+    fn compute_statistics(&self, data: &Tensor) -> Result<(Vec<f32>, Vec<f32>, Vec<f32>)> {
+        let shape = data.shape();
+        let _dims = shape.dims();
+        if shape.ndim() != 2 {
+            return Err(PreprocessingError::invalid_input(
+                "Expected 2D tensor (samples, features)",
+            ));
+        }
+
+        let dims = shape.dims();
+        let n_samples = dims[0];
+        let n_features = dims[1];
+
+        if n_samples == 0 {
+            return Err(PreprocessingError::EmptyDataset);
+        }
+
+        let cpu_data = data.to_cpu()?;
+        let values: Vec<f32> = cpu_data.clone();
+        let mut data_min = vec![f32::INFINITY; n_features];
+        let mut data_max = vec![f32::NEG_INFINITY; n_features];
+
+        // Compute min and max for each feature
+        for sample in 0..n_samples {
+            for feature in 0..n_features {
+                let idx = sample * n_features + feature;
+                let value = values[idx];
+
+                if !value.is_nan() {
+                    if value < data_min[feature] {
+                        data_min[feature] = value;
+                    }
+                    if value > data_max[feature] {
+                        data_max[feature] = value;
+                    }
+                }
+            }
+        }
+
+        // Handle edge cases and compute range
+        let mut data_range = vec![0.0; n_features];
+        for feature in 0..n_features {
+            if data_min[feature] == f32::INFINITY {
+                // All values were NaN
+                data_min[feature] = 0.0;
+                data_max[feature] = 0.0;
+                data_range[feature] = 0.0;
+            } else {
+                data_range[feature] = data_max[feature] - data_min[feature];
+
+                // Avoid division by zero for constant features
+                if data_range[feature] == 0.0 {
+                    data_range[feature] = 1.0;
+                }
+            }
+        }
+
+        Ok((data_min, data_max, data_range))
+    }
+}
+
+impl Default for MinMaxScaler {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
+impl Transformer for MinMaxScaler {
+    type Input = Tensor;
+    type Output = Tensor;
+
+    fn fit(&mut self, data: &Self::Input) -> Result<()> {
+        let (data_min, data_max, data_range) = self.compute_statistics(data)?;
+
+        self.data_min_ = Some(data_min);
+        self.data_max_ = Some(data_max);
+        self.data_range_ = Some(data_range);
+        self.n_features_ = Some(data.shape().dims()[1]);
+
+        Ok(())
+    }
+
+    fn transform(&self, data: &Self::Input) -> Result<Self::Output> {
+        if !self.is_fitted() {
+            return Err(PreprocessingError::NotFitted);
+        }
+
+        let shape = data.shape();
+        let _dims = shape.dims();
+        if shape.ndim() != 2 {
+            return Err(PreprocessingError::invalid_input(
+                "Expected 2D tensor (samples, features)",
+            ));
+        }
+
+        let dims = shape.dims();
+        let n_features = dims[1];
+        let expected_features = self.n_features_.unwrap();
+
+        if n_features != expected_features {
+            return Err(PreprocessingError::dimension_mismatch(
+                expected_features,
+                n_features,
+            ));
+        }
+
+        let cpu_data = data.to_cpu()?;
+        let values: Vec<f32> = cpu_data.clone();
+        let mut transformed_values = values.clone();
+
+        let data_min = self.data_min_.as_ref().unwrap();
+        let data_range = self.data_range_.as_ref().unwrap();
+        let (feature_min, feature_max) = self.feature_range;
+        let feature_range = feature_max - feature_min;
+
+        let n_samples = dims[0];
+
+        for sample in 0..n_samples {
+            for feature in 0..n_features {
+                let idx = sample * n_features + feature;
+                let value = transformed_values[idx];
+
+                // Scale to [0, 1]
+                let std_value = if data_range[feature] != 0.0 {
+                    (value - data_min[feature]) / data_range[feature]
+                } else {
+                    0.0 // Constant feature
+                };
+
+                // Scale to target range
+                let scaled_value = std_value * feature_range + feature_min;
+
+                transformed_values[idx] = scaled_value;
+            }
+        }
+
+        let transformed_f32: Vec<f32> = transformed_values.clone();
+        Ok(Tensor::from_slice(&transformed_f32, dims, data.device())?)
+    }
+
+    fn is_fitted(&self) -> bool {
+        self.data_min_.is_some()
+            && self.data_max_.is_some()
+            && self.data_range_.is_some()
+            && self.n_features_.is_some()
+    }
+
+    fn reset(&mut self) {
+        self.data_min_ = None;
+        self.data_max_ = None;
+        self.data_range_ = None;
+        self.n_features_ = None;
+    }
+}
+
+impl InvertibleTransformer for MinMaxScaler {
+    fn inverse_transform(&self, data: &Self::Output) -> Result<Self::Input> {
+        if !self.is_fitted() {
+            return Err(PreprocessingError::NotFitted);
+        }
+
+        let shape = data.shape();
+        let _dims = shape.dims();
+        if shape.ndim() != 2 {
+            return Err(PreprocessingError::invalid_input(
+                "Expected 2D tensor (samples, features)",
+            ));
+        }
+
+        let dims = shape.dims();
+        let n_features = dims[1];
+        let expected_features = self.n_features_.unwrap();
+
+        if n_features != expected_features {
+            return Err(PreprocessingError::dimension_mismatch(
+                expected_features,
+                n_features,
+            ));
+        }
+
+        let cpu_data = data.to_cpu()?;
+        let values: Vec<f32> = cpu_data.clone();
+        let mut inverse_values = values.clone();
+
+        let data_min = self.data_min_.as_ref().unwrap();
+        let data_range = self.data_range_.as_ref().unwrap();
+        let (feature_min, feature_max) = self.feature_range;
+        let feature_range = feature_max - feature_min;
+
+        let n_samples = dims[0];
+
+        for sample in 0..n_samples {
+            for feature in 0..n_features {
+                let idx = sample * n_features + feature;
+                let scaled_value = inverse_values[idx];
+
+                // Scale back from target range to [0, 1]
+                let std_value = if feature_range != 0.0 {
+                    (scaled_value - feature_min) / feature_range
+                } else {
+                    0.0
+                };
+
+                // Scale back to original range
+                let original_value = std_value * data_range[feature] + data_min[feature];
+
+                inverse_values[idx] = original_value;
+            }
+        }
+
+        let inverse_f32: Vec<f32> = inverse_values.clone();
+        Ok(Tensor::from_slice(&inverse_f32, dims, data.device())?)
+    }
+}