helena 0.1.0

Core types and component interfaces for helena, a latent data-to-waveform generation platform.
Documentation
use crate::latent::{Continuous, FrameSeq};
use crate::{Conditioning, Error, GenerationParams, Latent, LatentGenerator, Result, Seed};

use super::super::common::{continuous_set, elementwise_mean, mean_elementwise_variance, shape};
use super::sampling::generate_grid;

/// A condition-by-seed grid decomposed into diversity and controllability axes.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct GenerationDiagnostics {
    within_condition: f32,
    between_condition: f32,
    condition_count: usize,
    sample_count: usize,
}

impl GenerationDiagnostics {
    /// Measure diagnostics over `grid`, where `grid[c]` is one condition's samples.
    pub fn measure(grid: &[Vec<Latent<Continuous>>]) -> Result<Self> {
        if grid.is_empty() {
            return Err(Error::validation("generation grid has no conditions"));
        }

        let seeds = grid[0].len();
        let mut within_acc = 0.0f64;
        let mut condition_means: Vec<FrameSeq> = Vec::with_capacity(grid.len());
        let mut sample_count = 0usize;
        let mut expected_shape: Option<[usize; 2]> = None;
        for samples in grid {
            if samples.is_empty() {
                return Err(Error::validation("a condition has no samples"));
            }
            if samples.len() != seeds {
                return Err(Error::validation(format!(
                    "every condition needs the same seed budget; expected {seeds}, got {}",
                    samples.len()
                )));
            }
            let seqs = continuous_set(samples)?;
            match &expected_shape {
                None => expected_shape = Some(shape(seqs[0])),
                Some(expected) if *expected != shape(seqs[0]) => {
                    return Err(Error::validation(format!(
                        "all conditions must share a shape; expected {:?}, got {:?}",
                        expected,
                        shape(seqs[0])
                    )));
                }
                _ => {}
            }
            within_acc += mean_elementwise_variance(&seqs);
            condition_means.push(elementwise_mean(&seqs));
            sample_count += samples.len();
        }

        let within_condition = (within_acc / grid.len() as f64) as f32;
        let mean_refs: Vec<&FrameSeq> = condition_means.iter().collect();
        let between_condition = mean_elementwise_variance(&mean_refs) as f32;

        Ok(Self {
            within_condition,
            between_condition,
            condition_count: grid.len(),
            sample_count,
        })
    }

    /// Mean spread of samples sharing a conditioning.
    pub fn within_condition(&self) -> f32 {
        self.within_condition
    }

    /// Spread of per-condition sample means.
    pub fn between_condition(&self) -> f32 {
        self.between_condition
    }

    /// Number of distinct conditionings measured.
    pub fn condition_count(&self) -> usize {
        self.condition_count
    }

    /// Total generations across all conditions.
    pub fn sample_count(&self) -> usize {
        self.sample_count
    }
}

/// Sample a condition-by-seed grid from `generator` and score it.
pub fn sample_grid<G: LatentGenerator<Continuous>>(
    generator: &G,
    conditionings: &[Conditioning],
    seeds: &[Seed],
    params: &GenerationParams,
) -> Result<GenerationDiagnostics> {
    let grid = generate_grid(generator, conditionings, seeds, params)?;
    GenerationDiagnostics::measure(&grid)
}