fits-io 0.2.0

//! The dithering a floating point image is quantised with.
//!
//! Compressing a floating point image means turning it into integers first, and
//! rounding every pixel to the nearest step of the same grid lays a visible
//! pattern over a smooth background — contours where the sky crosses from one
//! step to the next. The convention's answer is to add a known pseudo-random
//! number to each value before rounding it and to take the same number off again
//! on the way back, which spreads the rounding error out into noise the eye does
//! not organise into shapes.
//!
//! "Known" is what makes it work: the sequence comes from a generator the
//! convention fixes, so a reader reproduces exactly the numbers the writer used.
//! A reader that ignores ZQUANTIZ gets an image that is wrong by up to half a
//! quantisation step, in a pattern rather than at random.

use std::error::Error;

/// How many numbers the dithering sequence holds before it repeats.
pub(crate) const SEQUENCE_LENGTH: usize = 10000;

/// The quantised value that a `SUBTRACTIVE_DITHER_2` tile uses for an exact
/// zero, which it stores rather than dithers so that zero stays zero.
///
/// It sits just above [`NULL_VALUE`], at the bottom of the range the convention
/// reserves.
const ZERO_VALUE: i64 = -2147483646;

/// How a floating point image was turned into integers.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum Quantization {
    /// The values were rounded to the nearest step, with nothing added.
    #[default]
    NoDither,
    /// Every value was dithered, zeros included.
    SubtractiveDither1,
    /// As `SUBTRACTIVE_DITHER_1`, but a pixel that was exactly zero is stored
    /// as such and comes back as exactly zero.
    SubtractiveDither2,
}

impl Quantization {
    /// Reads the method a ZQUANTIZ card names.
    ///
    /// # Errors
    ///
    /// Returns an error for a method this crate does not implement: undithering
    /// with the wrong sequence is worse than saying so, because the image it
    /// produces looks right.
    pub(crate) fn from_card(value: Option<&str>) -> Result<Self, Box<dyn Error + Send + Sync>> {
        match value.map(str::trim) {
            None | Some("") | Some("NONE") | Some("NO_DITHER") => Ok(Quantization::NoDither),
            Some("SUBTRACTIVE_DITHER_1") => Ok(Quantization::SubtractiveDither1),
            Some("SUBTRACTIVE_DITHER_2") => Ok(Quantization::SubtractiveDither2),
            Some(other) => Err(format!(
                "ZQUANTIZ {:?} is not a quantisation method this crate implements; it reads \
                 NO_DITHER, SUBTRACTIVE_DITHER_1 and SUBTRACTIVE_DITHER_2",
                other
            )
            .into()),
        }
    }

    /// The name a ZQUANTIZ card writes this method under.
    pub(crate) fn card_value(self) -> &'static str {
        match self {
            Quantization::NoDither => "NO_DITHER",
            Quantization::SubtractiveDither1 => "SUBTRACTIVE_DITHER_1",
            Quantization::SubtractiveDither2 => "SUBTRACTIVE_DITHER_2",
        }
    }

    /// Whether this method adds anything to a value before rounding it.
    pub(crate) fn dithers(self) -> bool {
        !matches!(self, Quantization::NoDither)
    }
}

/// The sequence of numbers the convention dithers with.
///
/// It is generated by the Park-Miller "minimal standard" generator from a seed
/// of one, exactly as the reference implementation does, so that the numbers are
/// the same ones the writer of a file used. Doing it any other way — a better
/// generator, or the same one in a different arithmetic — undithers an image
/// with numbers nobody added to it.
pub(crate) fn sequence() -> &'static [f32; SEQUENCE_LENGTH] {
    use std::sync::OnceLock;

    static SEQUENCE: OnceLock<[f32; SEQUENCE_LENGTH]> = OnceLock::new();

    SEQUENCE.get_or_init(|| {
        let mut values = [0.0_f32; SEQUENCE_LENGTH];

        for (value, seed) in values.iter_mut().zip(seeds()) {
            // Single precision is what the reference implementation keeps the
            // result in, and the numbers have to be the same ones to the last
            // bit.
            *value = (seed / MODULUS) as f32;
        }

        values
    })
}

/// The seeds the generator passes through, in order.
///
/// Every product here is under 2^53, so the arithmetic is exact and the
/// sequence is the same on every machine.
fn seeds() -> impl Iterator<Item = f64> {
    let mut seed = 1.0_f64;

    std::iter::repeat_with(move || {
        let temp = MULTIPLIER * seed;
        seed = temp - MODULUS * (temp / MODULUS).floor();
        seed
    })
}

/// The multiplier and modulus of the generator the convention fixes.
const MULTIPLIER: f64 = 16807.0;
const MODULUS: f64 = 2147483647.0;

/// Walks the dithering sequence for one tile.
///
/// Each tile starts at its own place in the sequence, worked out from the file's
/// ZDITHER0 and the tile's number, so that tiles do not all dither alike.
#[derive(Debug, Clone, Copy)]
pub(crate) struct Dither {
    /// Where in the sequence this tile's starting point was drawn from.
    start: usize,
    /// The index of the next number to use.
    next: usize,
}

impl Dither {
    /// The dithering for tile `tile` of a file seeded with `seed`, both counting
    /// from one and zero respectively.
    pub(crate) fn for_tile(seed: i64, tile: usize) -> Self {
        // The seed and the tile number both step through the sequence, so that
        // two files of the same shape do not dither identically and neither do
        // two tiles of one file.
        let start = (seed - 1).rem_euclid(SEQUENCE_LENGTH as i64) as usize;
        let start = (start + tile) % SEQUENCE_LENGTH;

        Self {
            start,
            next: Self::first(start),
        }
    }

    /// Where in the sequence a tile starting at `start` takes its first number
    /// from.
    fn first(start: usize) -> usize {
        (sequence()[start] * 500.0) as usize % SEQUENCE_LENGTH
    }

    /// The next number of the sequence, as the reference implementation's
    /// single precision value widened.
    fn next_value(&mut self) -> f64 {
        self.take() as f64
    }

    /// The next number of the sequence.
    fn take(&mut self) -> f32 {
        let value = sequence()[self.next];

        self.next += 1;
        if self.next == SEQUENCE_LENGTH {
            // The sequence has run out, so the tile draws a fresh starting point
            // from the next entry of its own.
            self.start = (self.start + 1) % SEQUENCE_LENGTH;
            self.next = Self::first(self.start);
        }

        value
    }
}

/// Turns a tile's quantised integers back into the values they stood for.
///
/// `blank` is the integer standing for a pixel the image does not define, which
/// comes back as `NaN` rather than as whatever that integer scales to.
pub(crate) fn unquantize(
    values: &[f64],
    scale: f64,
    zero: f64,
    method: Quantization,
    blank: Option<f64>,
    mut dither: Dither,
) -> Vec<f64> {
    values
        .iter()
        .map(|value| {
            // A blank pixel still draws its number: the sequence has to stay in
            // step with the one the writer used, whatever this pixel holds.
            let random = if method.dithers() {
                dither.next_value()
            } else {
                0.5
            };

            if Some(*value) == blank {
                return f64::NAN;
            }

            if method == Quantization::SubtractiveDither2 && *value == ZERO_VALUE as f64 {
                return 0.0;
            }

            zero + scale * (*value - random + 0.5)
        })
        .collect()
}

/// Turns a tile's values into the integers a compressor can work on, the inverse
/// of [`unquantize`].
pub(crate) fn quantize(
    values: &[f64],
    scale: f64,
    zero: f64,
    method: Quantization,
    blank: Option<i64>,
    mut dither: Dither,
) -> Vec<i64> {
    values
        .iter()
        .map(|value| {
            let random = if method.dithers() {
                dither.next_value()
            } else {
                0.5
            };

            if !value.is_finite() {
                // An undefined pixel is stored as the blank value, and there is
                // nowhere to put one if the caller reserved no such value.
                return blank.unwrap_or(0);
            }

            if method == Quantization::SubtractiveDither2 && *value == 0.0 {
                return ZERO_VALUE;
            }

            // Rounded the way the reference implementation rounds, so that a
            // value on the boundary between two steps goes the same way here as
            // it would there.
            ((value - zero) / scale + random - 0.5).round() as i64
        })
        .collect()
}

#[cfg(test)]
mod tests {
    use super::{Dither, Quantization, SEQUENCE_LENGTH, quantize, seeds, sequence, unquantize};

    /// The integers a compressor would hold, as the decompressor hands them
    /// back: whatever the coding was, they arrive as numbers.
    fn as_values(quantised: &[i64]) -> Vec<f64> {
        quantised.iter().map(|value| *value as f64).collect()
    }

    #[test]
    fn the_sequence_is_the_one_the_convention_fixes() {
        let sequence = sequence();

        // The first values of the Park-Miller generator seeded with one, as
        // 16807/2147483647, 16807^2 mod m / m, and so on.
        assert!((sequence[0] as f64 - 16807.0 / 2147483647.0).abs() < 1e-7);
        assert!((sequence[1] as f64 - 282475249.0 / 2147483647.0).abs() < 1e-7);
        assert!((sequence[2] as f64 - 1622650073.0 / 2147483647.0).abs() < 1e-7);

        // The reference implementation checks itself against this: the seed
        // behind the last number of the sequence is 1043618065, and an
        // implementation that produces anything else is not producing the
        // sequence the convention fixed.
        assert_eq!(seeds().nth(SEQUENCE_LENGTH - 1), Some(1043618065.0));

        // It never leaves the unit interval.
        assert!(sequence.iter().all(|value| (0.0..1.0).contains(value)));
    }

    #[test]
    fn a_quantisation_method_is_read_from_its_card() {
        assert_eq!(
            Quantization::from_card(None).unwrap(),
            Quantization::NoDither
        );
        assert_eq!(
            Quantization::from_card(Some("SUBTRACTIVE_DITHER_1")).unwrap(),
            Quantization::SubtractiveDither1
        );
        assert!(Quantization::from_card(Some("SOMETHING_ELSE")).is_err());
    }

    #[test]
    fn two_tiles_do_not_dither_alike() {
        let values = [100.0_f64; 8];

        let first = unquantize(
            &values,
            0.5,
            0.0,
            Quantization::SubtractiveDither1,
            None,
            Dither::for_tile(1, 0),
        );
        let second = unquantize(
            &values,
            0.5,
            0.0,
            Quantization::SubtractiveDither1,
            None,
            Dither::for_tile(1, 1),
        );

        assert_ne!(first, second);
    }

    #[test]
    fn quantising_and_undoing_it_lands_within_one_step() {
        let values: Vec<f64> = (0..64).map(|index| 10.0 + index as f64 * 0.017).collect();
        let scale = 0.01;

        for method in [
            Quantization::NoDither,
            Quantization::SubtractiveDither1,
            Quantization::SubtractiveDither2,
        ] {
            let quantised = as_values(&quantize(
                &values,
                scale,
                10.0,
                method,
                None,
                Dither::for_tile(7, 3),
            ));
            let back = unquantize(
                &quantised,
                scale,
                10.0,
                method,
                None,
                Dither::for_tile(7, 3),
            );

            for (original, returned) in values.iter().zip(&back) {
                assert!(
                    (original - returned).abs() <= scale,
                    "{original} came back as {returned}, further than one step of {scale}"
                );
            }
        }
    }

    #[test]
    fn dithering_keeps_the_average_of_a_flat_patch_where_it_was() {
        // The point of dithering. A patch of sky sitting four tenths of a step
        // above a quantisation level rounds, plainly, to that level in every
        // pixel: the patch comes back four tenths of a step too dark, with no
        // trace left that it was ever anywhere else. Dithered, the pixels fall
        // on either side in the right proportion, and the patch keeps its
        // brightness even though no single pixel does.
        let scale = 0.01;
        let value = 10.0 + 0.4 * scale;
        let values = vec![value; 2000];

        let mean = |method| {
            let quantised = as_values(&quantize(
                &values,
                scale,
                0.0,
                method,
                None,
                Dither::for_tile(1, 0),
            ));
            let back = unquantize(&quantised, scale, 0.0, method, None, Dither::for_tile(1, 0));

            back.iter().sum::<f64>() / back.len() as f64
        };

        let plain = mean(Quantization::NoDither);
        let dithered = mean(Quantization::SubtractiveDither1);

        assert!(
            (plain - value).abs() > 0.3 * scale,
            "plain rounding should lose the offset, got {plain}"
        );
        assert!(
            (dithered - value).abs() < 0.05 * scale,
            "dithering should keep the average at {value}, got {dithered}"
        );
    }

    #[test]
    fn dither_two_keeps_zero_exactly_zero() {
        let quantised = as_values(&quantize(
            &[0.0, 1.0],
            0.5,
            0.0,
            Quantization::SubtractiveDither2,
            None,
            Dither::for_tile(1, 0),
        ));
        let back = unquantize(
            &quantised,
            0.5,
            0.0,
            Quantization::SubtractiveDither2,
            None,
            Dither::for_tile(1, 0),
        );

        assert_eq!(back[0], 0.0);
    }

    #[test]
    fn a_blank_value_comes_back_undefined() {
        let back = unquantize(
            &[5.0, -32768.0, 7.0],
            1.0,
            0.0,
            Quantization::NoDither,
            Some(-32768.0),
            Dither::for_tile(1, 0),
        );

        assert!(back[1].is_nan(), "got {back:?}");
        assert!(back[0].is_finite() && back[2].is_finite(), "got {back:?}");
    }
}