1use 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
25pub enum Compression {
26 #[default]
30 Rice,
31 Gzip,
33 ShuffledGzip,
37 Hcompress {
48 scale: i64,
50 },
51 Plio,
55 None,
58}
59
60impl Compression {
61 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 fn needs_integers(self) -> bool {
75 matches!(
76 self,
77 Compression::Rice | Compression::Hcompress { .. } | Compression::Plio
78 )
79 }
80
81 fn element_bytes(self) -> usize {
86 match self {
87 Compression::Plio => 2,
88 _ => 1,
89 }
90 }
91
92 fn column_format(self) -> &'static str {
94 match self {
95 Compression::Plio => "1PI",
96 _ => "1PB",
97 }
98 }
99}
100
101#[derive(Debug, Clone, Copy, PartialEq, Default)]
108pub enum Quantize {
109 #[default]
113 Lossless,
114 Step(f64),
116 NoiseLevel(f64),
130}
131
132#[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 pub fn new(compression: Compression) -> Self {
167 Self {
168 compression,
169 tile: None,
170 quantize: Quantize::Lossless,
171 quantization: Quantization::SubtractiveDither1,
175 seed: 1,
176 block: 32,
177 }
178 }
179
180 #[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 #[must_use]
195 pub fn with_quantization(mut self, quantize: Quantize) -> Self {
196 self.quantize = quantize;
197 self
198 }
199
200 #[must_use]
205 pub fn with_dithering(mut self, quantization: Quantization) -> Self {
206 self.quantization = quantization;
207 self
208 }
209
210 #[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 #[must_use]
225 pub fn with_block_size(mut self, block: usize) -> Self {
226 self.block = block.max(1);
227 self
228 }
229
230 pub fn compression(&self) -> Compression {
232 self.compression
233 }
234}
235
236const NULL_VALUE: i64 = -2147483647;
239const RESERVED_VALUES: f64 = 10.0;
240
241const COMPRESSED_DATA: &str = "COMPRESSED_DATA";
243const SCALE: &str = "ZSCALE";
244const ZERO: &str = "ZZERO";
245
246pub(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 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 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
349fn 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
374fn 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
383fn 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 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
398fn gather(
400 pixels: &[f64],
401 shape: &[usize],
402 tile: &[usize],
403 tiles: &[usize],
404 index: usize,
405) -> (Vec<f64>, Vec<usize>) {
406 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
444fn 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 Quantize::Lossless => 0.0,
458 };
459
460 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 minimum - step * (NULL_VALUE as f64 + RESERVED_VALUES)
479 } else {
480 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
502fn noise(values: &[f64]) -> f64 {
510 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
535fn 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 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 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
584fn 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 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
617fn 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#[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 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 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 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
749impl Bitpix {
751 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 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 assert_eq!(compressed.bitpix(), Some(Bitpix::U8));
790 assert_eq!(compressed.naxis(), Some(2));
791 assert_eq!(compressed.table_fields(), Some(1));
792
793 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 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 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 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 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 assert_eq!(
908 shuffle(&[0x12, 0x34, 0x56, 0x78], 2),
909 vec![0x12, 0x56, 0x34, 0x78]
910 );
911
912 assert_eq!(shuffle(&[1, 2, 3], 1), vec![1, 2, 3]);
914 }
915}