1#![allow(clippy::multiple_crate_versions)] use ::zstd_sys as _;
25
26#[cfg(test)]
27use ::serde_json as _;
28
29use std::{borrow::Cow, fmt, num::NonZeroUsize};
30
31use ndarray::{Array, Array1, ArrayBase, Axis, Data, Dimension, IxDyn, ShapeError};
32use num_traits::{Float, identities::Zero};
33use numcodecs::{
34 AnyArray, AnyArrayAssignError, AnyArrayDType, AnyArrayView, AnyArrayViewMut, AnyCowArray,
35 Codec, StaticCodec, StaticCodecConfig, StaticCodecVersion,
36};
37use schemars::{JsonSchema, Schema, SchemaGenerator, json_schema};
38use serde::{Deserialize, Deserializer, Serialize, Serializer};
39use thiserror::Error;
40
41type QpetSperrCodecVersion = StaticCodecVersion<0, 2, 0>;
42
43#[derive(Clone, Serialize, Deserialize, JsonSchema)]
44#[schemars(deny_unknown_fields)]
46pub struct QpetSperrCodec {
54 #[serde(flatten)]
56 pub mode: QpetSperrCompressionMode,
57 #[serde(default, rename = "_version")]
59 pub version: QpetSperrCodecVersion,
60}
61
62#[derive(Clone, Serialize, Deserialize, JsonSchema)]
63#[serde(tag = "mode")]
65pub enum QpetSperrCompressionMode {
66 #[serde(rename = "qoi-symbolic")]
68 SymbolicQuantityOfInterest {
69 qoi: String,
71 #[serde(default = "default_qoi_block_size")]
74 qoi_block_size: (NonZeroUsize, NonZeroUsize, NonZeroUsize),
75 qoi_pwe: Positive<f64>,
78 #[serde(default = "default_sperr_chunks")]
80 sperr_chunks: (NonZeroUsize, NonZeroUsize, NonZeroUsize),
81 #[serde(default)]
83 data_pwe: Option<Positive<f64>>,
84 #[serde(default = "default_qoi_k")]
86 qoi_k: Positive<f64>,
87 #[serde(default)]
89 high_prec: bool,
90 },
91}
92
93const fn default_qoi_block_size() -> (NonZeroUsize, NonZeroUsize, NonZeroUsize) {
94 const NON_ZERO_ONE: NonZeroUsize = NonZeroUsize::MIN;
95 (NON_ZERO_ONE, NON_ZERO_ONE, NON_ZERO_ONE)
97}
98
99const fn default_sperr_chunks() -> (NonZeroUsize, NonZeroUsize, NonZeroUsize) {
100 const NON_ZERO_256: NonZeroUsize = NonZeroUsize::MIN.saturating_add(255);
101 (NON_ZERO_256, NON_ZERO_256, NON_ZERO_256)
102}
103
104const fn default_qoi_k() -> Positive<f64> {
105 Positive(3.0)
107}
108
109impl Codec for QpetSperrCodec {
110 type Error = QpetSperrCodecError;
111
112 fn encode(&self, data: AnyCowArray) -> Result<AnyArray, Self::Error> {
113 match data {
114 AnyCowArray::F32(data) => Ok(AnyArray::U8(
115 Array1::from(compress(data, &self.mode)?).into_dyn(),
116 )),
117 AnyCowArray::F64(data) => Ok(AnyArray::U8(
118 Array1::from(compress(data, &self.mode)?).into_dyn(),
119 )),
120 encoded => Err(QpetSperrCodecError::UnsupportedDtype(encoded.dtype())),
121 }
122 }
123
124 fn decode(&self, encoded: AnyCowArray) -> Result<AnyArray, Self::Error> {
125 let AnyCowArray::U8(encoded) = encoded else {
126 return Err(QpetSperrCodecError::EncodedDataNotBytes {
127 dtype: encoded.dtype(),
128 });
129 };
130
131 if !matches!(encoded.shape(), [_]) {
132 return Err(QpetSperrCodecError::EncodedDataNotOneDimensional {
133 shape: encoded.shape().to_vec(),
134 });
135 }
136
137 decompress(&AnyCowArray::U8(encoded).as_bytes())
138 }
139
140 fn decode_into(
141 &self,
142 encoded: AnyArrayView,
143 mut decoded: AnyArrayViewMut,
144 ) -> Result<(), Self::Error> {
145 let decoded_in = self.decode(encoded.cow())?;
146
147 Ok(decoded.assign(&decoded_in)?)
148 }
149}
150
151impl StaticCodec for QpetSperrCodec {
152 const CODEC_ID: &'static str = "qpet-sperr.rs";
153
154 type Config<'de> = Self;
155
156 fn from_config(config: Self::Config<'_>) -> Self {
157 config
158 }
159
160 fn get_config(&self) -> StaticCodecConfig<'_, Self> {
161 StaticCodecConfig::from(self)
162 }
163}
164
165#[derive(Debug, Error)]
166pub enum QpetSperrCodecError {
168 #[error("QpetSperr does not support the dtype {0}")]
170 UnsupportedDtype(AnyArrayDType),
171 #[error("QpetSperr failed to encode the header")]
173 HeaderEncodeFailed {
174 source: QpetSperrHeaderError,
176 },
177 #[error("QpetSperr failed to encode the data")]
179 QpetSperrEncodeFailed {
180 source: QpetSperrCodingError,
182 },
183 #[error("QpetSperr failed to encode a slice")]
185 SliceEncodeFailed {
186 source: QpetSperrSliceError,
188 },
189 #[error(
192 "QpetSperr can only decode one-dimensional byte arrays but received an array of dtype {dtype}"
193 )]
194 EncodedDataNotBytes {
195 dtype: AnyArrayDType,
197 },
198 #[error(
201 "QpetSperr can only decode one-dimensional byte arrays but received a byte array of shape {shape:?}"
202 )]
203 EncodedDataNotOneDimensional {
204 shape: Vec<usize>,
206 },
207 #[error("QpetSperr failed to decode the header")]
209 HeaderDecodeFailed {
210 source: QpetSperrHeaderError,
212 },
213 #[error("QpetSperr failed to decode a slice")]
215 SliceDecodeFailed {
216 source: QpetSperrSliceError,
218 },
219 #[error("QpetSperr failed to decode from an excessive number of slices")]
221 DecodeTooManySlices,
222 #[error("QpetSperr failed to decode the data")]
224 SperrDecodeFailed {
225 source: QpetSperrCodingError,
227 },
228 #[error("QpetSperr decoded into an invalid shape not matching the data size")]
230 DecodeInvalidShape {
231 source: ShapeError,
233 },
234 #[error("QpetSperr cannot decode into the provided array")]
236 MismatchedDecodeIntoArray {
237 #[from]
239 source: AnyArrayAssignError,
240 },
241}
242
243#[derive(Debug, Error)]
244#[error(transparent)]
245pub struct QpetSperrHeaderError(postcard::Error);
247
248#[derive(Debug, Error)]
249#[error(transparent)]
250pub struct QpetSperrSliceError(postcard::Error);
252
253#[derive(Debug, Error)]
254#[error(transparent)]
255pub struct QpetSperrCodingError(qpet_sperr::Error);
257
258#[allow(clippy::missing_panics_doc)]
270pub fn compress<T: QpetSperrElement, S: Data<Elem = T>, D: Dimension>(
271 data: ArrayBase<S, D>,
272 mode: &QpetSperrCompressionMode,
273) -> Result<Vec<u8>, QpetSperrCodecError> {
274 let mut encoded = postcard::to_extend(
275 &CompressionHeader {
276 dtype: T::DTYPE,
277 shape: Cow::Borrowed(data.shape()),
278 version: StaticCodecVersion,
279 },
280 Vec::new(),
281 )
282 .map_err(|err| QpetSperrCodecError::HeaderEncodeFailed {
283 source: QpetSperrHeaderError(err),
284 })?;
285
286 if data.is_empty() {
288 return Ok(encoded);
289 }
290
291 let mut chunk_size = Vec::from(data.shape());
292 let (width, height, depth) = match *chunk_size.as_mut_slice() {
293 [ref mut rest @ .., depth, height, width] => {
294 for r in rest {
295 *r = 1;
296 }
297 (width, height, depth)
298 }
299 [height, width] => (width, height, 1),
300 [width] => (width, 1, 1),
301 [] => (1, 1, 1),
302 };
303
304 for mut slice in data.into_dyn().exact_chunks(chunk_size.as_slice()) {
305 while slice.ndim() < 3 {
306 slice = slice.insert_axis(Axis(0));
307 }
308 #[allow(clippy::unwrap_used)]
309 let slice = slice.into_shape_with_order((depth, height, width)).unwrap();
312
313 let QpetSperrCompressionMode::SymbolicQuantityOfInterest {
314 qoi,
315 qoi_block_size,
316 qoi_pwe,
317 sperr_chunks,
318 data_pwe,
319 qoi_k,
320 high_prec,
321 } = mode;
322
323 let encoded_slice = qpet_sperr::compress_3d(
324 slice,
325 qpet_sperr::CompressionMode::SymbolicQuantityOfInterest {
326 qoi: qoi.as_str(),
327 qoi_block_size: *qoi_block_size,
328 qoi_pwe: qoi_pwe.0,
329 data_pwe: data_pwe.map(|data_pwe| data_pwe.0),
330 qoi_k: qoi_k.0,
331 high_prec: *high_prec,
332 },
333 (
334 sperr_chunks.0.get(),
335 sperr_chunks.1.get(),
336 sperr_chunks.2.get(),
337 ),
338 )
339 .map_err(|err| QpetSperrCodecError::QpetSperrEncodeFailed {
340 source: QpetSperrCodingError(err),
341 })?;
342
343 encoded = postcard::to_extend(encoded_slice.as_slice(), encoded).map_err(|err| {
344 QpetSperrCodecError::SliceEncodeFailed {
345 source: QpetSperrSliceError(err),
346 }
347 })?;
348 }
349
350 Ok(encoded)
351}
352
353pub fn decompress(encoded: &[u8]) -> Result<AnyArray, QpetSperrCodecError> {
366 fn decompress_typed<T: QpetSperrElement>(
367 mut encoded: &[u8],
368 shape: &[usize],
369 ) -> Result<Array<T, IxDyn>, QpetSperrCodecError> {
370 let mut decoded = Array::<T, _>::zeros(shape);
371
372 let mut chunk_size = Vec::from(shape);
373 let (width, height, depth) = match *chunk_size.as_mut_slice() {
374 [ref mut rest @ .., depth, height, width] => {
375 for r in rest {
376 *r = 1;
377 }
378 (width, height, depth)
379 }
380 [height, width] => (width, height, 1),
381 [width] => (width, 1, 1),
382 [] => (1, 1, 1),
383 };
384
385 for mut slice in decoded.exact_chunks_mut(chunk_size.as_slice()) {
386 let (encoded_slice, rest) =
387 postcard::take_from_bytes::<Cow<[u8]>>(encoded).map_err(|err| {
388 QpetSperrCodecError::SliceDecodeFailed {
389 source: QpetSperrSliceError(err),
390 }
391 })?;
392 encoded = rest;
393
394 while slice.ndim() < 3 {
395 slice = slice.insert_axis(Axis(0));
396 }
397 #[allow(clippy::unwrap_used)]
398 let slice = slice.into_shape_with_order((depth, height, width)).unwrap();
401
402 qpet_sperr::decompress_into_3d(&encoded_slice, slice).map_err(|err| {
403 QpetSperrCodecError::SperrDecodeFailed {
404 source: QpetSperrCodingError(err),
405 }
406 })?;
407 }
408
409 if !encoded.is_empty() {
410 return Err(QpetSperrCodecError::DecodeTooManySlices);
411 }
412
413 Ok(decoded)
414 }
415
416 let (header, encoded) =
417 postcard::take_from_bytes::<CompressionHeader>(encoded).map_err(|err| {
418 QpetSperrCodecError::HeaderDecodeFailed {
419 source: QpetSperrHeaderError(err),
420 }
421 })?;
422
423 if header.shape.iter().copied().product::<usize>() == 0 {
425 return match header.dtype {
426 QpetSperrDType::F32 => Ok(AnyArray::F32(Array::zeros(&*header.shape))),
427 QpetSperrDType::F64 => Ok(AnyArray::F64(Array::zeros(&*header.shape))),
428 };
429 }
430
431 match header.dtype {
432 QpetSperrDType::F32 => Ok(AnyArray::F32(decompress_typed(encoded, &header.shape)?)),
433 QpetSperrDType::F64 => Ok(AnyArray::F64(decompress_typed(encoded, &header.shape)?)),
434 }
435}
436
437pub trait QpetSperrElement: qpet_sperr::Element + Zero {
439 const DTYPE: QpetSperrDType;
441}
442
443impl QpetSperrElement for f32 {
444 const DTYPE: QpetSperrDType = QpetSperrDType::F32;
445}
446impl QpetSperrElement for f64 {
447 const DTYPE: QpetSperrDType = QpetSperrDType::F64;
448}
449
450#[expect(clippy::derive_partial_eq_without_eq)] #[derive(Copy, Clone, PartialEq, PartialOrd, Hash)]
452pub struct Positive<T: Float>(T);
454
455impl<T: Float> Positive<T> {
456 #[must_use]
457 pub const fn get(self) -> T {
459 self.0
460 }
461}
462
463impl Serialize for Positive<f64> {
464 fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
465 serializer.serialize_f64(self.0)
466 }
467}
468
469impl<'de> Deserialize<'de> for Positive<f64> {
470 fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
471 let x = f64::deserialize(deserializer)?;
472
473 if x > 0.0 {
474 Ok(Self(x))
475 } else {
476 Err(serde::de::Error::invalid_value(
477 serde::de::Unexpected::Float(x),
478 &"a positive value",
479 ))
480 }
481 }
482}
483
484impl JsonSchema for Positive<f64> {
485 fn schema_name() -> Cow<'static, str> {
486 Cow::Borrowed("PositiveF64")
487 }
488
489 fn schema_id() -> Cow<'static, str> {
490 Cow::Borrowed(concat!(module_path!(), "::", "Positive<f64>"))
491 }
492
493 fn json_schema(_gen: &mut SchemaGenerator) -> Schema {
494 json_schema!({
495 "type": "number",
496 "exclusiveMinimum": 0.0
497 })
498 }
499}
500
501#[derive(Serialize, Deserialize)]
502struct CompressionHeader<'a> {
503 dtype: QpetSperrDType,
504 #[serde(borrow)]
505 shape: Cow<'a, [usize]>,
506 version: QpetSperrCodecVersion,
507}
508
509#[derive(Copy, Clone, Debug, Serialize, Deserialize)]
511#[expect(missing_docs)]
512pub enum QpetSperrDType {
513 #[serde(rename = "f32", alias = "float32")]
514 F32,
515 #[serde(rename = "f64", alias = "float64")]
516 F64,
517}
518
519impl fmt::Display for QpetSperrDType {
520 fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
521 fmt.write_str(match self {
522 Self::F32 => "f32",
523 Self::F64 => "f64",
524 })
525 }
526}
527
528#[cfg(test)]
529#[allow(clippy::unwrap_used)]
530mod tests {
531 use std::f64;
532
533 use ndarray::{Ix0, Ix1, Ix2, Ix3, Ix4};
534
535 use super::*;
536
537 #[test]
538 fn zero_length() {
539 let encoded = compress(
540 Array::<f32, _>::from_shape_vec([3, 0], vec![]).unwrap(),
541 &QpetSperrCompressionMode::SymbolicQuantityOfInterest {
542 qoi: String::from("x"),
543 qoi_block_size: default_qoi_block_size(),
544 qoi_pwe: Positive(42.0),
545 sperr_chunks: default_sperr_chunks(),
546 data_pwe: None,
547 qoi_k: default_qoi_k(),
548 high_prec: false,
549 },
550 )
551 .unwrap();
552 let decoded = decompress(&encoded).unwrap();
553
554 assert_eq!(decoded.dtype(), AnyArrayDType::F32);
555 assert!(decoded.is_empty());
556 assert_eq!(decoded.shape(), &[3, 0]);
557 }
558
559 #[test]
560 fn small_2d() {
561 let encoded = compress(
562 Array::<f32, _>::from_shape_vec([1, 1], vec![42.0]).unwrap(),
563 &QpetSperrCompressionMode::SymbolicQuantityOfInterest {
564 qoi: String::from("x"),
565 qoi_block_size: default_qoi_block_size(),
566 qoi_pwe: Positive(42.0),
567 sperr_chunks: default_sperr_chunks(),
568 data_pwe: None,
569 qoi_k: default_qoi_k(),
570 high_prec: false,
571 },
572 )
573 .unwrap();
574 let decoded = decompress(&encoded).unwrap();
575
576 assert_eq!(decoded.dtype(), AnyArrayDType::F32);
577 assert_eq!(decoded.len(), 1);
578 assert_eq!(decoded.shape(), &[1, 1]);
579 }
580
581 #[test]
582 fn large_3d() {
583 let encoded = compress(
584 Array::<f64, _>::zeros((64, 64, 64)),
585 &QpetSperrCompressionMode::SymbolicQuantityOfInterest {
586 qoi: String::from("x"),
587 qoi_block_size: default_qoi_block_size(),
588 qoi_pwe: Positive(42.0),
589 sperr_chunks: default_sperr_chunks(),
590 data_pwe: None,
591 qoi_k: default_qoi_k(),
592 high_prec: false,
593 },
594 )
595 .unwrap();
596 let decoded = decompress(&encoded).unwrap();
597
598 assert_eq!(decoded.dtype(), AnyArrayDType::F64);
599 assert_eq!(decoded.len(), 64 * 64 * 64);
600 assert_eq!(decoded.shape(), &[64, 64, 64]);
601 }
602
603 #[test]
604 fn all_modes() {
605 for mode in [QpetSperrCompressionMode::SymbolicQuantityOfInterest {
606 qoi: String::from("x^2"),
607 qoi_block_size: default_qoi_block_size(),
608 qoi_pwe: Positive(0.1),
609 sperr_chunks: default_sperr_chunks(),
610 data_pwe: None,
611 qoi_k: default_qoi_k(),
612 high_prec: false,
613 }] {
614 let encoded = compress(Array::<f64, _>::zeros((64, 64, 64)), &mode).unwrap();
615 let decoded = decompress(&encoded).unwrap();
616
617 assert_eq!(decoded.dtype(), AnyArrayDType::F64);
618 assert_eq!(decoded.len(), 64 * 64 * 64);
619 assert_eq!(decoded.shape(), &[64, 64, 64]);
620 }
621 }
622
623 #[test]
624 fn many_dimensions() {
625 for data in [
626 Array::<f32, Ix0>::from_shape_vec([], vec![42.0])
627 .unwrap()
628 .into_dyn(),
629 Array::<f32, Ix1>::from_shape_vec([2], vec![1.0, 2.0])
630 .unwrap()
631 .into_dyn(),
632 Array::<f32, Ix2>::from_shape_vec([2, 2], vec![1.0, 2.0, 3.0, 4.0])
633 .unwrap()
634 .into_dyn(),
635 Array::<f32, Ix3>::from_shape_vec(
636 [2, 2, 2],
637 vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0],
638 )
639 .unwrap()
640 .into_dyn(),
641 Array::<f32, Ix4>::from_shape_vec(
642 [2, 2, 2, 2],
643 vec![
644 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0,
645 15.0, 16.0,
646 ],
647 )
648 .unwrap()
649 .into_dyn(),
650 ] {
651 let encoded = compress(
652 data.view(),
653 &QpetSperrCompressionMode::SymbolicQuantityOfInterest {
654 qoi: String::from("x"),
655 qoi_block_size: default_qoi_block_size(),
656 qoi_pwe: Positive(f64::EPSILON),
657 sperr_chunks: default_sperr_chunks(),
658 data_pwe: None,
659 qoi_k: default_qoi_k(),
660 high_prec: false,
661 },
662 )
663 .unwrap();
664 let decoded = decompress(&encoded).unwrap();
665
666 assert_eq!(decoded, AnyArray::F32(data));
667 }
668 }
669
670 #[test]
671 fn zero_square_qoi() {
672 let encoded = compress(
673 Array::<f64, _>::zeros((64, 64, 1)),
674 &QpetSperrCompressionMode::SymbolicQuantityOfInterest {
675 qoi: String::from("x^2"),
676 qoi_block_size: default_qoi_block_size(),
677 qoi_pwe: Positive(0.1),
678 sperr_chunks: default_sperr_chunks(),
679 data_pwe: None,
680 qoi_k: default_qoi_k(),
681 high_prec: false,
682 },
683 )
684 .unwrap();
685 let decoded = decompress(&encoded).unwrap();
686
687 assert_eq!(decoded.dtype(), AnyArrayDType::F64);
688 assert_eq!(decoded.len(), 64 * 64 * 1);
689 assert_eq!(decoded.shape(), &[64, 64, 1]);
690 }
691}