Skip to main content

fits_io/image/compression/
write.rs

1//! Writing an image as a tile-compressed table.
2//!
3//! This is the other half of the tiled image convention: an image is cut into
4//! tiles, each tile is compressed on its own, and the results become the rows of
5//! a binary table whose header describes, in keywords beginning `Z`, the image
6//! it stands for. A reader that knows nothing of compression sees a table; one
7//! that does sees the image.
8
9use crate::header::card_keys;
10use crate::header::{Bitpix, Header};
11use crate::image::compression::dither::{Dither, Quantization};
12use crate::image::compression::rice::BytesPerValue;
13use crate::image::compression::{dither, hcompress, plio, rice};
14use std::error::Error;
15
16/// Which algorithm the tiles are compressed with.
17///
18/// The convention defines several, and they suit different data. Rice is the
19/// usual choice for astronomical images: it is fast and, on data where each
20/// pixel is close to the one before it, it is the smallest of these. Gzip
21/// compresses anything at all, including floating point values that have not
22/// been quantised, which is what makes it the one lossless choice for a
23/// floating point image.
24#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
25pub enum Compression {
26    /// `RICE_1`: the differences between neighbouring pixels, coded so that
27    /// small ones are short. Integers only — a floating point image has to be
28    /// quantised first.
29    #[default]
30    Rice,
31    /// `GZIP_1`: deflate over the tile's bytes as they stand.
32    Gzip,
33    /// `GZIP_2`: deflate after gathering the first byte of every value, then
34    /// the second, and so on, which usually compresses better because the high
35    /// bytes of neighbouring values are alike.
36    ShuffledGzip,
37    /// `HCOMPRESS_1`: an image transform that gathers each two by two block
38    /// into a sum and three differences, coded as a quadtree. It compresses
39    /// smooth images better than Rice, and at a `scale` above one it does so by
40    /// throwing away the low bits of each coefficient — which is lossy, and
41    /// which [`ImageHDU::read_image`] can smooth back over.
42    ///
43    /// A scale of zero or one keeps every bit. Tiles are two-dimensional, so
44    /// this cannot be used with a tile that reaches along a third axis.
45    ///
46    /// [`ImageHDU::read_image`]: crate::hdu::ImageHDU::read_image
47    Hcompress {
48        /// What the transform's coefficients are divided by.
49        scale: i64,
50    },
51    /// `PLIO_1`: run-length coding for the integer masks IRAF writes, where a
52    /// tile is mostly long runs of the same small number. Values must be
53    /// non-negative and no wider than twenty-seven bits.
54    Plio,
55    /// `NOCOMPRESS`: the tiles are stored as they are. Useful for writing a
56    /// file in the tiled layout without paying to compress it.
57    None,
58}
59
60impl Compression {
61    /// The name a ZCMPTYPE card writes this algorithm under.
62    pub fn card_value(self) -> &'static str {
63        match self {
64            Compression::Rice => "RICE_1",
65            Compression::Gzip => "GZIP_1",
66            Compression::ShuffledGzip => "GZIP_2",
67            Compression::Hcompress { .. } => "HCOMPRESS_1",
68            Compression::Plio => "PLIO_1",
69            Compression::None => "NOCOMPRESS",
70        }
71    }
72
73    /// Whether this algorithm works on integers rather than on bytes.
74    fn needs_integers(self) -> bool {
75        matches!(
76            self,
77            Compression::Rice | Compression::Hcompress { .. } | Compression::Plio
78        )
79    }
80
81    /// How wide each element of the column holding a tile is.
82    ///
83    /// PLIO's instructions are sixteen bit words, and the column holds them as
84    /// such rather than as the bytes they are made of.
85    fn element_bytes(self) -> usize {
86        match self {
87            Compression::Plio => 2,
88            _ => 1,
89        }
90    }
91
92    /// The TFORMn of the column holding a tile.
93    fn column_format(self) -> &'static str {
94        match self {
95            Compression::Plio => "1PI",
96            _ => "1PB",
97        }
98    }
99}
100
101/// How a floating point image's values are turned into the integers a
102/// compressor can work on.
103///
104/// Quantising is what makes a floating point image compress at all well, and it
105/// is lossy: what comes back is within one step of what went in. The step is the
106/// choice being made here.
107#[derive(Debug, Clone, Copy, PartialEq, Default)]
108pub enum Quantize {
109    /// Do not quantise. The values are compressed as they stand, which only
110    /// [`Compression::Gzip`], [`Compression::ShuffledGzip`] and
111    /// [`Compression::None`] can do, and nothing is lost.
112    #[default]
113    Lossless,
114    /// Quantise in steps of this many units of the image's own values.
115    Step(f64),
116    /// Quantise in steps of the tile's own estimated noise divided by this.
117    ///
118    /// Four is the usual choice: it keeps the quantisation step well below the
119    /// noise already in the data, so nothing measurable is lost, while throwing
120    /// away the low bits that are noise anyway and would otherwise compress
121    /// terribly.
122    ///
123    /// The noise is estimated from the median absolute third-order difference
124    /// between neighbouring pixels, which is what the convention's own reference
125    /// implementation does. That implementation takes the smallest of three such
126    /// estimates rather than this one alone, so it may settle on a slightly
127    /// different step for the same tile; either is a fair reading of "the noise
128    /// in this tile".
129    NoiseLevel(f64),
130}
131
132/// Everything about how an image is to be compressed.
133///
134/// ```
135/// use fits_io::image::compression::{Compression, CompressionOptions, Quantize};
136///
137/// // Rice coding, in tiles of 100 by 100 pixels.
138/// let options = CompressionOptions::new(Compression::Rice).with_tile_size(&[100, 100]);
139///
140/// // A floating point image, quantised to a quarter of its own noise.
141/// let lossy = CompressionOptions::new(Compression::Rice)
142///     .with_quantization(Quantize::NoiseLevel(4.0));
143///
144/// // HCOMPRESS, throwing away the low bits of the transform's coefficients.
145/// let smooth = CompressionOptions::new(Compression::Hcompress { scale: 16 });
146/// ```
147#[derive(Debug, Clone, PartialEq)]
148pub struct CompressionOptions {
149    compression: Compression,
150    tile: Option<Vec<u32>>,
151    quantize: Quantize,
152    quantization: Quantization,
153    seed: i64,
154    block: usize,
155}
156
157impl Default for CompressionOptions {
158    fn default() -> Self {
159        Self::new(Compression::default())
160    }
161}
162
163impl CompressionOptions {
164    /// Compress with `compression`, in the convention's default tiles — one row
165    /// of the image each — and without quantising.
166    pub fn new(compression: Compression) -> Self {
167        Self {
168            compression,
169            tile: None,
170            quantize: Quantize::Lossless,
171            // Dithering costs nothing and is what keeps quantisation from
172            // laying a pattern over a smooth background, so it is on wherever
173            // quantisation is.
174            quantization: Quantization::SubtractiveDither1,
175            seed: 1,
176            block: 32,
177        }
178    }
179
180    /// Cut the image into tiles of this shape, fastest axis first.
181    ///
182    /// A shape shorter than the image's axes has the remaining axes tiled one
183    /// plane at a time. Bigger tiles compress a little better and cost more to
184    /// read a small part of the image from.
185    #[must_use]
186    pub fn with_tile_size(mut self, tile: &[u32]) -> Self {
187        self.tile = Some(tile.to_vec());
188        self
189    }
190
191    /// Quantise a floating point image before compressing it.
192    ///
193    /// This has no effect on an image that is already integers.
194    #[must_use]
195    pub fn with_quantization(mut self, quantize: Quantize) -> Self {
196        self.quantize = quantize;
197        self
198    }
199
200    /// Which dithering the quantisation uses.
201    ///
202    /// [`Quantization::SubtractiveDither1`] unless this says otherwise, which is
203    /// what a floating point image usually wants.
204    #[must_use]
205    pub fn with_dithering(mut self, quantization: Quantization) -> Self {
206        self.quantization = quantization;
207        self
208    }
209
210    /// Where in the dithering sequence the first tile starts, from 1 to 10000.
211    ///
212    /// Two files written from the same data with the same seed come out
213    /// identical, which is what makes a compressed file reproducible.
214    #[must_use]
215    pub fn with_dither_seed(mut self, seed: i64) -> Self {
216        self.seed = seed.rem_euclid(dither::SEQUENCE_LENGTH as i64).max(1);
217        self
218    }
219
220    /// How many values Rice coding fits its split point to at a time.
221    ///
222    /// Thirty-two unless this says otherwise, which is what the convention's
223    /// reference implementation uses.
224    #[must_use]
225    pub fn with_block_size(mut self, block: usize) -> Self {
226        self.block = block.max(1);
227        self
228    }
229
230    /// The algorithm the tiles are compressed with.
231    pub fn compression(&self) -> Compression {
232        self.compression
233    }
234}
235
236/// The value standing for a pixel the image does not define, and the ten values
237/// above it that the convention reserves alongside it.
238const NULL_VALUE: i64 = -2147483647;
239const RESERVED_VALUES: f64 = 10.0;
240
241/// The columns a compressed table is written with.
242const COMPRESSED_DATA: &str = "COMPRESSED_DATA";
243const SCALE: &str = "ZSCALE";
244const ZERO: &str = "ZZERO";
245
246/// Compresses an image into the header and data of the table that stands for it.
247///
248/// `header` describes the image as it is now, and `data` is its pixels as they
249/// sit in the file. What comes back is the header and data section of a binary
250/// table extension carrying the same image, compressed.
251pub(crate) fn compress(
252    header: &Header,
253    data: &[u8],
254    options: &CompressionOptions,
255) -> Result<(Header, Vec<u8>), Box<dyn Error + Send + Sync>> {
256    let bitpix = header
257        .bitpix()
258        .ok_or("An image needs a BITPIX card before it can be compressed")?;
259
260    let shape = shape_of(header);
261    if shape.is_empty() {
262        return Err("An image with no axes has nothing to compress".into());
263    }
264
265    let tile = tile_shape(options, &shape);
266    let quantizing = matches!(bitpix, Bitpix::F32 | Bitpix::F64)
267        && !matches!(options.quantize, Quantize::Lossless);
268
269    if bitpix.is_floating() && options.compression.needs_integers() && !quantizing {
270        return Err(format!(
271            "{} compresses integers, and this image holds floating point values. Either quantise \
272             it, which loses the low bits of every pixel, or compress it with GZIP_1, which does \
273             not.",
274            options.compression.card_value()
275        )
276        .into());
277    }
278
279    let pixels = read_pixels(data, bitpix, &shape)?;
280
281    // A quantised tile holds 32-bit integers whatever the image's own type is.
282    let stored = if quantizing { Bitpix::I32 } else { bitpix };
283
284    let tiles: Vec<usize> = shape
285        .iter()
286        .zip(&tile)
287        .map(|(length, tile)| length.div_ceil(*tile))
288        .collect();
289    let tile_count: usize = tiles.iter().product();
290
291    let mut rows = Vec::new();
292    let mut heap = Vec::new();
293    let mut any_blank = false;
294
295    let elements = options.compression.element_bytes();
296
297    for index in 0..tile_count {
298        let (values, extent) = gather(&pixels, &shape, &tile, &tiles, index);
299
300        let (integers, scale, zero) = if quantizing {
301            let (integers, scale, zero) = quantize_tile(&values, options, index);
302            any_blank |= values.iter().any(|value| !value.is_finite());
303            (integers, Some(scale), Some(zero))
304        } else {
305            (
306                values.iter().map(|value| *value as i64).collect(),
307                None,
308                None,
309            )
310        };
311
312        let compressed = encode(&integers, &values, stored, &extent, options)?;
313
314        // Each row carries a descriptor saying how long its array is and where
315        // in the heap it sits; the bytes themselves all go in the heap. The
316        // length is counted in the column's own elements, which are not always
317        // bytes.
318        rows.extend_from_slice(&((compressed.len() / elements) as u32).to_be_bytes());
319        rows.extend_from_slice(&(heap.len() as u32).to_be_bytes());
320        heap.extend_from_slice(&compressed);
321
322        if let (Some(scale), Some(zero)) = (scale, zero) {
323            rows.extend_from_slice(&scale.to_be_bytes());
324            rows.extend_from_slice(&zero.to_be_bytes());
325        }
326    }
327
328    let row_bytes = if quantizing { 8 + 16 } else { 8 };
329
330    let mut table = rows;
331    table.extend_from_slice(&heap);
332
333    let compressed_header = compressed_header(
334        header,
335        bitpix,
336        &shape,
337        &tile,
338        options,
339        quantizing,
340        any_blank,
341        tile_count,
342        row_bytes,
343        heap.len(),
344    )?;
345
346    Ok((compressed_header, table))
347}
348
349/// The pixels of an image, as the numbers they stand for.
350fn read_pixels(
351    data: &[u8],
352    bitpix: Bitpix,
353    shape: &[usize],
354) -> Result<Vec<f64>, Box<dyn Error + Send + Sync>> {
355    let count: usize = shape.iter().product();
356    let width = bitpix.byte_size();
357
358    if data.len() < count * width {
359        return Err(format!(
360            "This image says it holds {} pixels of {} bytes, and its data section is {} bytes",
361            count,
362            width,
363            data.len()
364        )
365        .into());
366    }
367
368    Ok(data[..count * width]
369        .chunks_exact(width)
370        .filter_map(|raw| bitpix.read_be(raw))
371        .collect())
372}
373
374/// The image's shape, fastest axis first.
375fn shape_of(header: &Header) -> Vec<usize> {
376    let axes = header.naxis().unwrap_or(0).max(0) as usize;
377
378    (0..axes)
379        .map(|axis| header.naxis_n(axis).unwrap_or(0).max(0) as usize)
380        .collect()
381}
382
383/// How far a tile reaches along each axis.
384fn tile_shape(options: &CompressionOptions, shape: &[usize]) -> Vec<usize> {
385    (0..shape.len())
386        .map(|axis| {
387            let asked = match &options.tile {
388                Some(tile) => tile.get(axis).map(|size| *size as usize),
389                // The convention's default is one row of the image.
390                None => Some(if axis == 0 { shape[0] } else { 1 }),
391            };
392
393            asked.unwrap_or(1).clamp(1, shape[axis].max(1))
394        })
395        .collect()
396}
397
398/// The values of one tile, in the order the tile stores them.
399fn gather(
400    pixels: &[f64],
401    shape: &[usize],
402    tile: &[usize],
403    tiles: &[usize],
404    index: usize,
405) -> (Vec<f64>, Vec<usize>) {
406    // Where this tile starts along each axis, and how far it reaches before the
407    // edge of the image cuts it short.
408    let mut origin = vec![0_usize; shape.len()];
409    let mut extent = vec![0_usize; shape.len()];
410
411    let mut rest = index;
412    for axis in 0..shape.len() {
413        origin[axis] = (rest % tiles[axis]) * tile[axis];
414        rest /= tiles[axis];
415        extent[axis] = tile[axis].min(shape[axis] - origin[axis]);
416    }
417
418    let run = extent[0];
419    let runs: usize = extent.iter().skip(1).product();
420    let mut values = Vec::with_capacity(run * runs);
421
422    let mut within = vec![0_usize; shape.len()];
423
424    for index in 0..runs {
425        let mut rest = index;
426        for axis in 1..shape.len() {
427            within[axis] = rest % extent[axis];
428            rest /= extent[axis];
429        }
430
431        let mut at = origin[0];
432        let mut stride = shape[0];
433        for axis in 1..shape.len() {
434            at += (origin[axis] + within[axis]) * stride;
435            stride *= shape[axis];
436        }
437
438        values.extend_from_slice(&pixels[at..at + run]);
439    }
440
441    (values, extent)
442}
443
444/// Turns one tile's values into the integers a compressor works on, and says
445/// what scales them back.
446fn quantize_tile(
447    values: &[f64],
448    options: &CompressionOptions,
449    tile: usize,
450) -> (Vec<i64>, f64, f64) {
451    let finite: Vec<f64> = values.iter().copied().filter(|v| v.is_finite()).collect();
452
453    let step = match options.quantize {
454        Quantize::Step(step) => step.abs(),
455        Quantize::NoiseLevel(level) => noise(&finite) / level.max(f64::MIN_POSITIVE),
456        // The caller has already established that this tile is being quantised.
457        Quantize::Lossless => 0.0,
458    };
459
460    // A tile with no spread at all, or one this crate can find no noise in,
461    // is stored in steps of one, which for such a tile loses nothing.
462    let step = if step > 0.0 && step.is_finite() {
463        step
464    } else {
465        1.0
466    };
467
468    let minimum = finite.iter().copied().fold(f64::INFINITY, f64::min);
469    let maximum = finite.iter().copied().fold(f64::NEG_INFINITY, f64::max);
470
471    let zero = if !minimum.is_finite() {
472        0.0
473    } else if finite.len() < values.len()
474        || options.quantization == Quantization::SubtractiveDither2
475    {
476        // Room has to be left below the smallest value for the values the
477        // convention reserves, of which the blank value is one.
478        minimum - step * (NULL_VALUE as f64 + RESERVED_VALUES)
479    } else {
480        // Otherwise the smallest value sits at the bottom of the range, on a
481        // whole number of steps — so that compressing a file that has already
482        // been through this comes out the same way twice.
483        let factor = (minimum / step + 0.5).floor();
484        factor * step
485    };
486
487    let blank = (finite.len() < values.len()).then_some(NULL_VALUE);
488    let integers = dither::quantize(
489        values,
490        step,
491        zero,
492        options.quantization,
493        blank,
494        Dither::for_tile(options.seed, tile),
495    );
496
497    let _ = maximum;
498
499    (integers, step, zero)
500}
501
502/// Estimates the noise in a tile from how much neighbouring pixels differ.
503///
504/// A third-order difference — twice a pixel less its two neighbours — cancels
505/// any smooth gradient the image has, so what is left of it is noise. Taking the
506/// median of those rather than their mean keeps a handful of stars from being
507/// mistaken for noise, and the constant turns that median into the standard
508/// deviation of a normal distribution that would produce it.
509fn noise(values: &[f64]) -> f64 {
510    /// Turns the median absolute third-order difference into a standard
511    /// deviation.
512    const MEDIAN_TO_SIGMA: f64 = 0.6052697;
513
514    if values.len() < 3 {
515        return 0.0;
516    }
517
518    let mut differences: Vec<f64> = values
519        .windows(3)
520        .map(|window| (2.0 * window[1] - window[0] - window[2]).abs())
521        .collect();
522
523    differences.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
524
525    let middle = differences.len() / 2;
526    let median = if differences.len().is_multiple_of(2) {
527        (differences[middle - 1] + differences[middle]) / 2.0
528    } else {
529        differences[middle]
530    };
531
532    MEDIAN_TO_SIGMA * median
533}
534
535/// Compresses one tile.
536///
537/// `integers` is the tile as whole numbers, for the algorithms that work on
538/// those, and `values` the same tile as it stands, for the ones that work on
539/// bytes.
540fn encode(
541    integers: &[i64],
542    values: &[f64],
543    stored: Bitpix,
544    extent: &[usize],
545    options: &CompressionOptions,
546) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
547    match options.compression {
548        Compression::Rice => {
549            let width = BytesPerValue::from_count(stored.byte_size() as i64)?;
550            Ok(rice::compress(integers, width, options.block))
551        }
552        Compression::Hcompress { scale } => {
553            // The transform is two-dimensional, and a tile that reaches along a
554            // third axis is not a plane for it to work on.
555            if extent.iter().skip(2).any(|length| *length > 1) {
556                return Err(format!(
557                    "HCOMPRESS compresses a plane at a time, and this tile is {:?}",
558                    extent
559                )
560                .into());
561            }
562
563            let columns = extent.first().copied().unwrap_or(0);
564            let rows = extent.get(1).copied().unwrap_or(1);
565
566            // The transform's own axis order is the other way round from the
567            // image's: its first dimension is the slow one.
568            hcompress::compress(integers, rows, columns, scale)
569        }
570        Compression::Plio => {
571            let words = plio::compress(integers)?;
572
573            Ok(words.iter().flat_map(|word| word.to_be_bytes()).collect())
574        }
575        Compression::None => Ok(to_be_bytes(integers, values, stored)),
576        Compression::Gzip => gzip(&to_be_bytes(integers, values, stored)),
577        Compression::ShuffledGzip => {
578            let bytes = to_be_bytes(integers, values, stored);
579            gzip(&shuffle(&bytes, stored.byte_size()))
580        }
581    }
582}
583
584/// A tile as the bytes it is stored in.
585fn to_be_bytes(integers: &[i64], values: &[f64], stored: Bitpix) -> Vec<u8> {
586    let mut bytes = Vec::with_capacity(integers.len() * stored.byte_size());
587
588    match stored {
589        Bitpix::U8 => bytes.extend(integers.iter().map(|value| *value as u8)),
590        Bitpix::I16 => {
591            for value in integers {
592                bytes.extend_from_slice(&(*value as i16).to_be_bytes());
593            }
594        }
595        Bitpix::I32 => {
596            for value in integers {
597                bytes.extend_from_slice(&(*value as i32).to_be_bytes());
598            }
599        }
600        // A floating point tile that was not quantised keeps its own values,
601        // NaNs and all, rather than the integers they do not have.
602        Bitpix::F32 => {
603            for value in values {
604                bytes.extend_from_slice(&(*value as f32).to_be_bytes());
605            }
606        }
607        Bitpix::F64 => {
608            for value in values {
609                bytes.extend_from_slice(&value.to_be_bytes());
610            }
611        }
612    }
613
614    bytes
615}
616
617/// Gathers the first byte of every value, then the second, and so on.
618fn shuffle(bytes: &[u8], width: usize) -> Vec<u8> {
619    if width <= 1 {
620        return bytes.to_vec();
621    }
622
623    let count = bytes.len() / width;
624    let mut out = vec![0_u8; count * width];
625
626    for byte in 0..width {
627        for value in 0..count {
628            out[byte * count + value] = bytes[value * width + byte];
629        }
630    }
631
632    out
633}
634
635#[cfg(feature = "gzip")]
636fn gzip(bytes: &[u8]) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
637    use std::io::Write;
638
639    let mut encoder = flate2::write::GzEncoder::new(Vec::new(), flate2::Compression::default());
640    encoder.write_all(bytes)?;
641
642    Ok(encoder.finish()?)
643}
644
645#[cfg(not(feature = "gzip"))]
646fn gzip(_bytes: &[u8]) -> Result<Vec<u8>, Box<dyn Error + Send + Sync>> {
647    Err("Compressing with gzip needs the `gzip` feature".into())
648}
649
650/// Builds the header of the table an image is compressed into.
651#[allow(clippy::too_many_arguments)]
652fn compressed_header(
653    header: &Header,
654    bitpix: Bitpix,
655    shape: &[usize],
656    tile: &[usize],
657    options: &CompressionOptions,
658    quantizing: bool,
659    any_blank: bool,
660    rows: usize,
661    row_bytes: usize,
662    heap: usize,
663) -> Result<Header, Box<dyn Error + Send + Sync>> {
664    let mut out = header.clone();
665
666    // The image's own structural cards are replaced by the table's, and its
667    // shape moves into the Z keywords.
668    out.remove_card(card_keys::NAXIS);
669    out.remove_prefixed(card_keys::PREFIX_NAXIS_N);
670
671    out.set_card(card_keys::BITPIX, 8_i64)?;
672    out.set_naxis_n(0, row_bytes as i64)?;
673    out.set_naxis_n(1, rows as i64)?;
674    out.set_card(card_keys::NAXIS, 2_i64)?;
675    out.set_card(card_keys::PCOUNT, heap as i64)?;
676    out.set_card(card_keys::GCOUNT, 1_i64)?;
677
678    let tiles = options.compression.column_format();
679
680    let columns: Vec<(&str, &str)> = if quantizing {
681        vec![(COMPRESSED_DATA, tiles), (SCALE, "1D"), (ZERO, "1D")]
682    } else {
683        vec![(COMPRESSED_DATA, tiles)]
684    };
685
686    out.set_card(card_keys::TFIELDS, columns.len() as i64)?;
687    for (index, (name, format)) in columns.iter().enumerate() {
688        out.set_card(
689            &format!("{}{}", card_keys::PREFIX_TTYPE_N, index + 1),
690            *name,
691        )?;
692        out.set_card(
693            &format!("{}{}", card_keys::PREFIX_TFORM_N, index + 1),
694            *format,
695        )?;
696    }
697
698    out.set_card(card_keys::ZIMAGE, true)?;
699    out.set_card(card_keys::ZBITPIX, i64::from(bitpix))?;
700    out.set_card(card_keys::ZNAXIS, shape.len() as i64)?;
701
702    for (axis, length) in shape.iter().enumerate() {
703        out.set_card(&format!("ZNAXIS{}", axis + 1), *length as i64)?;
704        out.set_card(&format!("ZTILE{}", axis + 1), tile[axis] as i64)?;
705    }
706
707    out.set_card(card_keys::ZCMPTYPE, options.compression.card_value())?;
708
709    // The algorithms take their settings as name and value pairs.
710    let mut parameters: Vec<(&str, i64)> = Vec::new();
711    match options.compression {
712        Compression::Rice => {
713            parameters.push(("BLOCKSIZE", options.block as i64));
714            parameters.push((
715                "BYTEPIX",
716                if quantizing {
717                    4
718                } else {
719                    bitpix.byte_size() as i64
720                },
721            ));
722        }
723        Compression::Hcompress { scale } => {
724            parameters.push(("SCALE", scale));
725            // Whether to smooth is the reader's choice, and a file that asks
726            // for it says so here.
727            parameters.push(("SMOOTH", 0));
728        }
729        _ => {}
730    }
731
732    for (index, (name, value)) in parameters.iter().enumerate() {
733        out.set_card(&format!("ZNAME{}", index + 1), *name)?;
734        out.set_card(&format!("ZVAL{}", index + 1), *value)?;
735    }
736
737    if quantizing {
738        out.set_card(card_keys::ZQUANTIZ, options.quantization.card_value())?;
739        out.set_card(card_keys::ZDITHER0, options.seed)?;
740
741        if any_blank {
742            out.set_card(card_keys::ZBLANK, NULL_VALUE)?;
743        }
744    }
745
746    Ok(out)
747}
748
749/// The BITPIX of a header, as the number the card holds.
750impl Bitpix {
751    /// Whether this type holds floating point values.
752    pub(crate) fn is_floating(self) -> bool {
753        matches!(self, Bitpix::F32 | Bitpix::F64)
754    }
755}
756
757#[cfg(test)]
758mod tests {
759    use super::{Compression, CompressionOptions, Quantize, compress, noise, shuffle};
760    use crate::header::{Bitpix, Header};
761
762    /// A header for a `width` by `height` image of `bitpix`.
763    fn header(bitpix: Bitpix, width: usize, height: usize) -> Header {
764        let mut header = Header::default();
765
766        header.set_card("BITPIX", i64::from(bitpix)).unwrap();
767        header.set_card("NAXIS", 2_i64).unwrap();
768        header.set_naxis_n(0, width as i64).unwrap();
769        header.set_naxis_n(1, height as i64).unwrap();
770
771        header
772    }
773
774    fn i16_data(values: &[i16]) -> Vec<u8> {
775        values.iter().flat_map(|v| v.to_be_bytes()).collect()
776    }
777
778    #[test]
779    fn a_compressed_header_describes_both_the_table_and_the_image() {
780        let values: Vec<i16> = (0..64).collect();
781        let (compressed, _) = compress(
782            &header(Bitpix::I16, 8, 8),
783            &i16_data(&values),
784            &CompressionOptions::new(Compression::Rice),
785        )
786        .expect("an image that can be compressed");
787
788        // The table it is stored as.
789        assert_eq!(compressed.bitpix(), Some(Bitpix::U8));
790        assert_eq!(compressed.naxis(), Some(2));
791        assert_eq!(compressed.table_fields(), Some(1));
792
793        // And the image it stands for.
794        assert!(compressed.is_compressed_image());
795        assert_eq!(compressed.compressed_bitpix(), Some(Bitpix::I16));
796        assert_eq!(compressed.compressed_naxis(), Some(2));
797        assert_eq!(compressed.compressed_naxis_n(0), Some(8));
798        assert_eq!(compressed.compressed_naxis_n(1), Some(8));
799        assert_eq!(compressed.compression_type(), Some("RICE_1"));
800        assert_eq!(compressed.compression_parameter("BYTEPIX"), Some(2));
801    }
802
803    #[test]
804    fn the_default_tile_is_one_row_of_the_image() {
805        let values: Vec<i16> = (0..64).collect();
806        let (compressed, _) = compress(
807            &header(Bitpix::I16, 8, 8),
808            &i16_data(&values),
809            &CompressionOptions::new(Compression::Rice),
810        )
811        .expect("an image that can be compressed");
812
813        assert_eq!(compressed.compressed_tile_size(0), 8);
814        assert_eq!(compressed.compressed_tile_size(1), 1);
815        // One row per tile means one table row per image row.
816        assert_eq!(compressed.naxis_n(1), Some(8));
817    }
818
819    #[test]
820    fn a_tile_size_larger_than_the_image_is_cut_down_to_it() {
821        let values: Vec<i16> = (0..64).collect();
822        let (compressed, _) = compress(
823            &header(Bitpix::I16, 8, 8),
824            &i16_data(&values),
825            &CompressionOptions::new(Compression::Rice).with_tile_size(&[1000, 1000]),
826        )
827        .expect("an image that can be compressed");
828
829        assert_eq!(compressed.compressed_tile_size(0), 8);
830        assert_eq!(compressed.compressed_tile_size(1), 8);
831        assert_eq!(compressed.naxis_n(1), Some(1));
832    }
833
834    #[test]
835    fn a_floating_point_image_cannot_be_rice_coded_without_being_quantised() {
836        let data: Vec<u8> = (0..64).flat_map(|i| (i as f32).to_be_bytes()).collect();
837
838        let error = compress(
839            &header(Bitpix::F32, 8, 8),
840            &data,
841            &CompressionOptions::new(Compression::Rice),
842        )
843        .expect_err("Rice coding works on integers");
844
845        assert!(error.to_string().contains("quantise"), "got: {error}");
846    }
847
848    #[test]
849    fn a_quantised_image_says_how_it_was_quantised() {
850        let data: Vec<u8> = (0..64)
851            .flat_map(|i| (i as f32 * 0.5).to_be_bytes())
852            .collect();
853
854        let (compressed, _) = compress(
855            &header(Bitpix::F32, 8, 8),
856            &data,
857            &CompressionOptions::new(Compression::Rice)
858                .with_quantization(Quantize::Step(0.01))
859                .with_dither_seed(42),
860        )
861        .expect("a quantised image compresses");
862
863        assert_eq!(
864            compressed.quantization_method(),
865            Some("SUBTRACTIVE_DITHER_1")
866        );
867        assert_eq!(compressed.dither_seed(), Some(42));
868        assert_eq!(compressed.table_fields(), Some(3));
869        assert_eq!(compressed.compression_parameter("BYTEPIX"), Some(4));
870
871        // The image is still a floating point image; only its tiles are
872        // integers.
873        assert_eq!(compressed.compressed_bitpix(), Some(Bitpix::F32));
874        assert_eq!(
875            compressed.card("TTYPE2").map(|v| v.value_to_string()),
876            Some("ZSCALE".to_string())
877        );
878    }
879
880    #[test]
881    fn the_noise_estimate_follows_the_noise() {
882        // Pixels drawn either side of a line: the estimate should follow how
883        // far they are drawn from it, and should not be fooled by the line.
884        let quiet: Vec<f64> = (0..200)
885            .map(|i| i as f64 + if i % 2 == 0 { 0.1 } else { -0.1 })
886            .collect();
887        let loud: Vec<f64> = (0..200)
888            .map(|i| i as f64 + if i % 2 == 0 { 5.0 } else { -5.0 })
889            .collect();
890
891        assert!(noise(&quiet) > 0.0);
892        assert!(
893            noise(&loud) > 10.0 * noise(&quiet),
894            "{} vs {}",
895            noise(&loud),
896            noise(&quiet)
897        );
898
899        // A perfectly smooth ramp has no noise in it at all.
900        let ramp: Vec<f64> = (0..200).map(|i| i as f64 * 3.0).collect();
901        assert_eq!(noise(&ramp), 0.0);
902    }
903
904    #[test]
905    fn shuffling_gathers_each_byte_of_every_value_together() {
906        // Two 16-bit values: the high bytes first, then the low ones.
907        assert_eq!(
908            shuffle(&[0x12, 0x34, 0x56, 0x78], 2),
909            vec![0x12, 0x56, 0x34, 0x78]
910        );
911
912        // Single bytes have nothing to gather.
913        assert_eq!(shuffle(&[1, 2, 3], 1), vec![1, 2, 3]);
914    }
915}