helena 0.1.0

use serde::{Deserialize, Serialize};

use crate::{Error, Result};

/// A logit-normal sampling density over schedule time `t` in `[0, 1]`.
///
/// A standard-normal draw `n` maps to `t = logistic(location + scale * n)`, so the
/// median lands at `logistic(location)` and `scale` sets the spread. Concentrating
/// the draw on the mid band biases training mass toward the noise levels a few-step
/// sampler traverses coarsely, instead of the trivial high-SNR tail a uniform-in-`t`
/// draw over-samples (the scarce-data objective brief, lever C; Esser et al. 2024).
///
/// The density only places mass; it does not alter the schedule. It composes the
/// existing [`CosineSchedule`](super::CosineSchedule), which still maps the drawn
/// `t` to a [`NoiseLevel`](super::NoiseLevel). Total and deterministic: `sample` is
/// a pure function of the draw, with no Burn and no RNG of its own (the trainer owns
/// the RNG so replay stays seed-determined).
#[derive(Clone, Copy, Debug, PartialEq, Serialize, Deserialize)]
#[serde(try_from = "LogitNormalRepr")]
pub struct LogitNormalDensity {
    location: f32,
    scale: f32,
}

/// Deserialization mirror routed through [`LogitNormalDensity::new`].
///
/// `deny_unknown_fields` makes a typo in the nested config object (e.g. `scle`)
/// fail closed, matching the resolver's fail-closed handling of unknown top-level
/// `extra` keys rather than silently dropping the misspelled field.
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct LogitNormalRepr {
    location: f32,
    scale: f32,
}

impl TryFrom<LogitNormalRepr> for LogitNormalDensity {
    type Error = Error;

    fn try_from(repr: LogitNormalRepr) -> Result<Self> {
        Self::new(repr.location, repr.scale)
    }
}

impl LogitNormalDensity {
    /// Centered location: the median time is `logistic(0) = 0.5`, the mid band of
    /// the cosine schedule.
    pub const DEFAULT_LOCATION: f32 = 0.0;

    /// Unit spread, the SD3 logit-normal starting point.
    pub const DEFAULT_SCALE: f32 = 1.0;

    /// Build a density, requiring `location` finite and `scale` finite and
    /// strictly positive.
    ///
    /// A non-positive scale would collapse the draw to a point (`scale = 0`) or
    /// invert its order (`scale < 0`), so it is rejected at construction and the
    /// field invariants hold for every value of the type.
    pub fn new(location: f32, scale: f32) -> Result<Self> {
        if !location.is_finite() {
            return Err(Error::validation(format!(
                "logit-normal location must be finite, got {location}"
            )));
        }
        if !scale.is_finite() || scale <= 0.0 {
            return Err(Error::validation(format!(
                "logit-normal scale must be finite and positive, got {scale}"
            )));
        }
        Ok(Self { location, scale })
    }

    /// Map a standard-normal draw `n` to a schedule time `logistic(location + scale * n)`.
    ///
    /// Total on every finite `n`: the logistic map saturates into the closed
    /// `[0, 1]` range [`CosineSchedule`](super::CosineSchedule) accepts, rather than
    /// diverging, so a drawn time is always a valid schedule input.
    pub fn sample(&self, standard_normal: f32) -> f32 {
        logistic(self.location + self.scale * standard_normal)
    }

    /// The location: the underlying normal's mean; the median time is `logistic(location)`.
    pub fn location(&self) -> f32 {
        self.location
    }

    /// The scale: the underlying normal's standard deviation.
    pub fn scale(&self) -> f32 {
        self.scale
    }
}

/// The standard logistic `1 / (1 + e^-x)`, total and in `[0, 1]` on every finite
/// input (saturating to the closed endpoints rather than diverging).
fn logistic(x: f32) -> f32 {
    1.0 / (1.0 + (-x).exp())
}