rustytorch/crates/specialized/rtx-neuro-forward/src/gain.rs

//! Gain matrix (lead field) utilities.
//!
//! The gain matrix G relates source activity to sensor measurements:
//! M = G * S
//! where M is [n_sensors x n_times], G is [n_sensors x n_sources], and S is [n_sources x n_times].

use crate::ForwardResult;

/// Gain matrix wrapper with metadata
#[derive(Debug, Clone)]
pub struct GainMatrix {
    /// The gain matrix data [n_sensors x n_sources]
    data: Vec<Vec<f64>>,
    /// Number of sensors
    n_sensors: usize,
    /// Number of sources (or 3*n_sources for free orientation)
    n_source_columns: usize,
    /// Whether sources have free orientation (3 DOF each)
    free_orientation: bool,
    /// Sensor names
    sensor_names: Vec<String>,
}

impl GainMatrix {
    /// Create a new gain matrix
    pub fn new(
        data: Vec<Vec<f64>>,
        free_orientation: bool,
        sensor_names: Vec<String>,
    ) -> ForwardResult<Self> {
        let n_sensors = data.len();
        let n_source_columns = data.first().map_or(0, std::vec::Vec::len);

        Ok(Self {
            data,
            n_sensors,
            n_source_columns,
            free_orientation,
            sensor_names,
        })
    }

    /// Get the gain matrix data
    pub fn data(&self) -> &Vec<Vec<f64>> {
        &self.data
    }

    /// Get number of sensors
    pub fn n_sensors(&self) -> usize {
        self.n_sensors
    }

    /// Get number of source columns
    pub fn n_source_columns(&self) -> usize {
        self.n_source_columns
    }

    /// Get number of sources (accounting for orientation)
    pub fn n_sources(&self) -> usize {
        if self.free_orientation {
            self.n_source_columns / 3
        } else {
            self.n_source_columns
        }
    }

    /// Check if sources have free orientation
    pub fn is_free_orientation(&self) -> bool {
        self.free_orientation
    }

    /// Get sensor names
    pub fn sensor_names(&self) -> &[String] {
        &self.sensor_names
    }

    /// Get a row (all sources for one sensor)
    pub fn get_row(&self, sensor_idx: usize) -> Option<&[f64]> {
        self.data.get(sensor_idx).map(std::vec::Vec::as_slice)
    }

    /// Get a column (all sensors for one source/orientation)
    pub fn get_column(&self, source_idx: usize) -> Option<Vec<f64>> {
        if source_idx >= self.n_source_columns {
            return None;
        }
        Some(self.data.iter().map(|row| row[source_idx]).collect())
    }

    /// Compute the norm of each source's lead field
    ///
    /// For fixed orientation, this is the L2 norm of the column.
    /// For free orientation, this combines all 3 orientations.
    pub fn source_norms(&self) -> Vec<f64> {
        let n_sources = self.n_sources();
        let mut norms = Vec::with_capacity(n_sources);

        for src in 0..n_sources {
            if self.free_orientation {
                // Combine x, y, z columns
                let mut sum_sq = 0.0;
                for ori in 0..3 {
                    let col_idx = 3 * src + ori;
                    for row in &self.data {
                        sum_sq += row[col_idx] * row[col_idx];
                    }
                }
                norms.push(sum_sq.sqrt());
            } else {
                // Single column
                let sum_sq: f64 = self.data.iter().map(|row| row[src] * row[src]).sum();
                norms.push(sum_sq.sqrt());
            }
        }

        norms
    }

    /// Apply the gain matrix to source activity
    ///
    /// M = G * S
    ///
    /// # Arguments
    /// * `sources` - Source activity [n_source_columns x n_times]
    ///
    /// # Returns
    /// Sensor data [n_sensors x n_times]
    pub fn apply(&self, sources: &[Vec<f64>]) -> ForwardResult<Vec<Vec<f64>>> {
        if sources.len() != self.n_source_columns {
            return Err(crate::ForwardError::DimensionMismatch(format!(
                "Expected {} source columns, got {}",
                self.n_source_columns,
                sources.len()
            )));
        }

        let n_times = sources.first().map_or(0, std::vec::Vec::len);
        let mut result = vec![vec![0.0; n_times]; self.n_sensors];

        for (sens_idx, row) in self.data.iter().enumerate() {
            for (src_idx, &gain) in row.iter().enumerate() {
                for (t, &val) in sources[src_idx].iter().enumerate() {
                    result[sens_idx][t] += gain * val;
                }
            }
        }

        Ok(result)
    }

    /// Compute G^T * G (source covariance induced by sensor covariance)
    pub fn gram_matrix(&self) -> Vec<Vec<f64>> {
        let n = self.n_source_columns;
        let mut gram = vec![vec![0.0; n]; n];

        for i in 0..n {
            for j in i..n {
                let dot: f64 = self.data.iter().map(|row| row[i] * row[j]).sum();
                gram[i][j] = dot;
                gram[j][i] = dot;
            }
        }

        gram
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_gain_matrix() {
        let data = vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]];
        let names = vec!["S1".to_string(), "S2".to_string()];

        let gain = GainMatrix::new(data, false, names).unwrap();
        assert_eq!(gain.n_sensors(), 2);
        assert_eq!(gain.n_sources(), 3);
    }

    #[test]
    fn test_apply() {
        let data = vec![vec![1.0, 0.0], vec![0.0, 1.0]];
        let names = vec!["S1".to_string(), "S2".to_string()];
        let gain = GainMatrix::new(data, false, names).unwrap();

        let sources = vec![
            vec![1.0, 2.0], // Source 1
            vec![3.0, 4.0], // Source 2
        ];

        let result = gain.apply(&sources).unwrap();
        assert_eq!(result[0][0], 1.0);
        assert_eq!(result[1][0], 3.0);
    }

    #[test]
    fn test_source_norms() {
        let data = vec![vec![3.0, 0.0], vec![4.0, 1.0]];
        let names = vec!["S1".to_string(), "S2".to_string()];
        let gain = GainMatrix::new(data, false, names).unwrap();

        let norms = gain.source_norms();
        assert!((norms[0] - 5.0).abs() < 1e-10); // sqrt(9+16) = 5
        assert!((norms[1] - 1.0).abs() < 1e-10);
    }
}