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;
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct GenerationDiagnostics {
within_condition: f32,
between_condition: f32,
condition_count: usize,
sample_count: usize,
}
impl GenerationDiagnostics {
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,
})
}
pub fn within_condition(&self) -> f32 {
self.within_condition
}
pub fn between_condition(&self) -> f32 {
self.between_condition
}
pub fn condition_count(&self) -> usize {
self.condition_count
}
pub fn sample_count(&self) -> usize {
self.sample_count
}
}
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)
}