1#![allow(clippy::multiple_crate_versions)] #[cfg(test)]
23use ::serde_json as _;
24
25use std::{borrow::Cow, num::NonZeroUsize};
26
27use ndarray::{Array, Array1, ArrayBase, ArrayViewMut, Axis, Data, DataMut, Dimension, IxDyn};
28use num_traits::Float;
29use numcodecs::{
30 AnyArray, AnyArrayDType, AnyArrayView, AnyArrayViewMut, AnyCowArray, Codec, StaticCodec,
31 StaticCodecConfig, StaticCodecVersion,
32};
33use schemars::{JsonSchema, Schema, SchemaGenerator, json_schema};
34use serde::{Deserialize, Deserializer, Serialize, Serializer};
35use thiserror::Error;
36
37type EbccCodecVersion = StaticCodecVersion<0, 1, 1>;
38
39#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
49#[schemars(deny_unknown_fields)]
50pub struct EbccCodec {
51 #[serde(flatten)]
53 pub residual: EbccResidualType,
54 #[serde(default = "default_base_cr")]
56 pub base_cr: Positive<f32>,
57 #[serde(default)]
59 pub chunk_shape: EbccChunkShape,
60 #[serde(default, rename = "_version")]
62 pub version: EbccCodecVersion,
63}
64
65#[derive(Debug, Copy, Clone, Default, Serialize, Deserialize, JsonSchema)]
66#[serde(deny_unknown_fields)]
67pub enum EbccChunkShape {
69 #[serde(rename = "auto")]
71 #[default]
72 Auto,
73 #[serde(untagged)]
75 Explicit([NonZeroUsize; 3]),
76}
77
78const fn default_base_cr() -> Positive<f32> {
79 Positive(100.0)
80}
81
82#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, JsonSchema)]
84#[serde(tag = "residual")]
85#[serde(deny_unknown_fields)]
86pub enum EbccResidualType {
87 #[serde(rename = "jpeg2000-only")]
88 Jpeg2000Only,
90 #[serde(rename = "absolute")]
91 AbsoluteError {
93 error: Positive<f32>,
95 },
96 #[serde(rename = "relative")]
97 RelativeError {
99 error: Positive<f32>,
101 },
102}
103
104impl Codec for EbccCodec {
105 type Error = EbccCodecError;
106
107 fn encode(&self, data: AnyCowArray) -> Result<AnyArray, Self::Error> {
108 match data {
109 AnyCowArray::F32(data) => Ok(AnyArray::U8(
110 Array1::from(compress(
111 data,
112 self.residual,
113 self.base_cr,
114 self.chunk_shape,
115 )?)
116 .into_dyn(),
117 )),
118 encoded => Err(EbccCodecError::UnsupportedDtype(encoded.dtype())),
119 }
120 }
121
122 fn decode(&self, encoded: AnyCowArray) -> Result<AnyArray, Self::Error> {
123 let AnyCowArray::U8(encoded) = encoded else {
124 return Err(EbccCodecError::EncodedDataNotBytes {
125 dtype: encoded.dtype(),
126 });
127 };
128
129 if !matches!(encoded.shape(), [_]) {
130 return Err(EbccCodecError::EncodedDataNotOneDimensional {
131 shape: encoded.shape().to_vec(),
132 });
133 }
134
135 decompress(&AnyCowArray::U8(encoded).as_bytes())
136 }
137
138 fn decode_into(
139 &self,
140 encoded: AnyArrayView,
141 decoded: AnyArrayViewMut,
142 ) -> Result<(), Self::Error> {
143 let AnyArrayView::U8(encoded) = encoded else {
144 return Err(EbccCodecError::EncodedDataNotBytes {
145 dtype: encoded.dtype(),
146 });
147 };
148
149 if !matches!(encoded.shape(), [_]) {
150 return Err(EbccCodecError::EncodedDataNotOneDimensional {
151 shape: encoded.shape().to_vec(),
152 });
153 }
154
155 match decoded {
156 AnyArrayViewMut::F32(decoded) => {
157 decompress_into(&AnyArrayView::U8(encoded).as_bytes(), decoded)
158 }
159 decoded => Err(EbccCodecError::UnsupportedDtype(decoded.dtype())),
160 }
161 }
162}
163
164impl StaticCodec for EbccCodec {
165 const CODEC_ID: &'static str = "ebcc.rs";
166
167 type Config<'de> = Self;
168
169 fn from_config(config: Self::Config<'_>) -> Self {
170 config
171 }
172
173 fn get_config(&self) -> StaticCodecConfig<'_, Self> {
174 StaticCodecConfig::from(self)
175 }
176}
177
178#[derive(Debug, thiserror::Error)]
180pub enum EbccCodecError {
181 #[error("Ebcc does not support the dtype {0}")]
183 UnsupportedDtype(AnyArrayDType),
184 #[error("Ebcc failed to encode the header")]
186 HeaderEncodeFailed {
187 source: EbccHeaderError,
189 },
190 #[error(
193 "Ebcc can only encode >2D data where the last two dimensions must be at least 32x32 but received an array of shape {shape:?}"
194 )]
195 InsufficientDimensions {
196 shape: Vec<usize>,
198 },
199 #[error("Ebcc failed to encode the data")]
201 EbccEncodeFailed {
202 source: EbccCodingError,
204 },
205 #[error("Ebcc failed to encode a 3D slice")]
207 SliceEncodeFailed {
208 source: EbccSliceError,
210 },
211 #[error(
214 "Ebcc can only decode one-dimensional byte arrays but received an array of dtype {dtype}"
215 )]
216 EncodedDataNotBytes {
217 dtype: AnyArrayDType,
219 },
220 #[error(
223 "Ebcc can only decode one-dimensional byte arrays but received a byte array of shape {shape:?}"
224 )]
225 EncodedDataNotOneDimensional {
226 shape: Vec<usize>,
228 },
229 #[error("Ebcc failed to decode the header")]
231 HeaderDecodeFailed {
232 source: EbccHeaderError,
234 },
235 #[error("Ebcc cannot decode an array of shape {decoded:?} into an array of shape {array:?}")]
237 DecodeIntoShapeMismatch {
238 decoded: Vec<usize>,
240 array: Vec<usize>,
242 },
243 #[error("Ebcc failed to decode a slice")]
245 SliceDecodeFailed {
246 source: EbccSliceError,
248 },
249 #[error("Ebcc failed to decode from an excessive number of slices")]
251 DecodeTooManySlices,
252 #[error("Ebcc failed to decode the data")]
254 EbccDecodeFailed {
255 source: EbccCodingError,
257 },
258}
259
260#[expect(clippy::derive_partial_eq_without_eq)] #[derive(Debug, Copy, Clone, PartialEq, PartialOrd, Hash)]
262pub struct Positive<T: Float>(T);
264
265impl<T: Float> PartialEq<T> for Positive<T> {
266 fn eq(&self, other: &T) -> bool {
267 self.0 == *other
268 }
269}
270
271impl Serialize for Positive<f32> {
272 fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
273 serializer.serialize_f32(self.0)
274 }
275}
276
277impl<'de> Deserialize<'de> for Positive<f32> {
278 fn deserialize<D: Deserializer<'de>>(deserializer: D) -> Result<Self, D::Error> {
279 let x = f32::deserialize(deserializer)?;
280
281 if x > 0.0 {
282 Ok(Self(x))
283 } else {
284 Err(serde::de::Error::invalid_value(
285 serde::de::Unexpected::Float(f64::from(x)),
286 &"a positive value",
287 ))
288 }
289 }
290}
291
292impl JsonSchema for Positive<f32> {
293 fn schema_name() -> Cow<'static, str> {
294 Cow::Borrowed("PositiveF32")
295 }
296
297 fn schema_id() -> Cow<'static, str> {
298 Cow::Borrowed(concat!(module_path!(), "::", "Positive<f32>"))
299 }
300
301 fn json_schema(_gen: &mut SchemaGenerator) -> Schema {
302 json_schema!({
303 "type": "number",
304 "exclusiveMinimum": 0.0
305 })
306 }
307}
308
309#[derive(Debug, Error)]
310#[error(transparent)]
311pub struct EbccHeaderError(postcard::Error);
313
314#[derive(Debug, Error)]
315#[error(transparent)]
316pub struct EbccSliceError(postcard::Error);
318
319#[derive(Debug, Error)]
320#[error(transparent)]
321pub struct EbccCodingError(ebcc::EBCCError);
323
324#[allow(clippy::missing_panics_doc)]
336pub fn compress<S: Data<Elem = f32>, D: Dimension>(
337 data: ArrayBase<S, D>,
338 residual: EbccResidualType,
339 base_cr: Positive<f32>,
340 chunk_shape: EbccChunkShape,
341) -> Result<Vec<u8>, EbccCodecError> {
342 let mut encoded = postcard::to_extend(
343 &CompressionHeader {
344 dtype: EbccDType::F32,
345 shape: Cow::Borrowed(data.shape()),
346 version: StaticCodecVersion,
347 },
348 Vec::new(),
349 )
350 .map_err(|err| EbccCodecError::HeaderEncodeFailed {
351 source: EbccHeaderError(err),
352 })?;
353
354 if data.is_empty() {
356 return Ok(encoded);
357 }
358
359 let mut chunk_size = Vec::from(data.shape());
360 let (width, height, depth) = match *chunk_size.as_mut_slice() {
361 [ref mut rest @ .., depth, height, width] => {
362 for r in rest {
363 *r = 1;
364 }
365 (width, height, depth)
366 }
367 [height, width] => (width, height, 1),
368 _ => {
369 return Err(EbccCodecError::InsufficientDimensions {
370 shape: Vec::from(data.shape()),
371 });
372 }
373 };
374
375 if (width < 32) || (height < 32) {
376 return Err(EbccCodecError::InsufficientDimensions {
377 shape: Vec::from(data.shape()),
378 });
379 }
380
381 for mut slice in data.into_dyn().exact_chunks(chunk_size.as_slice()) {
382 while slice.ndim() < 3 {
383 slice = slice.insert_axis(Axis(0));
384 }
385 #[expect(clippy::unwrap_used)]
386 let slice = slice.into_shape_with_order((depth, height, width)).unwrap();
389
390 let encoded_slice = ebcc::ebcc_encode_chunking_compat(
391 slice,
392 &ebcc::EBCCConfig {
393 base_cr: base_cr.0,
394 residual_compression_type: match residual {
395 EbccResidualType::Jpeg2000Only => ebcc::EBCCResidualType::Jpeg2000Only,
396 EbccResidualType::AbsoluteError { error } => {
397 ebcc::EBCCResidualType::AbsoluteError(error.0)
398 }
399 EbccResidualType::RelativeError { error } => {
400 ebcc::EBCCResidualType::RelativeError(error.0)
401 }
402 },
403 },
404 match chunk_shape {
405 EbccChunkShape::Auto => ebcc::EBCCCompatChunkShape::Auto,
406 EbccChunkShape::Explicit(chunk_shape) => {
407 ebcc::EBCCCompatChunkShape::Explicit(chunk_shape)
408 }
409 },
410 )
411 .map_err(|err| EbccCodecError::EbccEncodeFailed {
412 source: EbccCodingError(err),
413 })?;
414
415 encoded = postcard::to_extend(encoded_slice.as_slice(), encoded).map_err(|err| {
416 EbccCodecError::SliceEncodeFailed {
417 source: EbccSliceError(err),
418 }
419 })?;
420 }
421
422 Ok(encoded)
423}
424
425pub fn decompress(encoded: &[u8]) -> Result<AnyArray, EbccCodecError> {
436 fn decompress_typed(
437 encoded: &[u8],
438 shape: &[usize],
439 ) -> Result<Array<f32, IxDyn>, EbccCodecError> {
440 let mut decoded = Array::<f32, _>::zeros(shape);
441 decompress_into_typed(encoded, decoded.view_mut())?;
442 Ok(decoded)
443 }
444
445 let (header, encoded) =
446 postcard::take_from_bytes::<CompressionHeader>(encoded).map_err(|err| {
447 EbccCodecError::HeaderDecodeFailed {
448 source: EbccHeaderError(err),
449 }
450 })?;
451
452 if header.shape.iter().copied().any(|s| s == 0) {
454 return match header.dtype {
455 EbccDType::F32 => Ok(AnyArray::F32(Array::zeros(&*header.shape))),
456 };
457 }
458
459 match header.dtype {
460 EbccDType::F32 => Ok(AnyArray::F32(decompress_typed(encoded, &header.shape)?)),
461 }
462}
463
464pub fn decompress_into<S: DataMut<Elem = f32>, D: Dimension>(
477 encoded: &[u8],
478 decoded: ArrayBase<S, D>,
479) -> Result<(), EbccCodecError> {
480 let (header, encoded) =
481 postcard::take_from_bytes::<CompressionHeader>(encoded).map_err(|err| {
482 EbccCodecError::HeaderDecodeFailed {
483 source: EbccHeaderError(err),
484 }
485 })?;
486
487 if decoded.shape() != &*header.shape {
488 return Err(EbccCodecError::DecodeIntoShapeMismatch {
489 decoded: header.shape.into_owned(),
490 array: Vec::from(decoded.shape()),
491 });
492 }
493
494 if header.shape.iter().copied().any(|s| s == 0) {
496 return match header.dtype {
497 EbccDType::F32 => Ok(()),
498 };
499 }
500
501 match header.dtype {
502 EbccDType::F32 => decompress_into_typed(encoded, decoded.into_dyn().view_mut()),
503 }
504}
505
506fn decompress_into_typed(
507 mut encoded: &[u8],
508 mut decoded: ArrayViewMut<f32, IxDyn>,
509) -> Result<(), EbccCodecError> {
510 let mut chunk_size = Vec::from(decoded.shape());
511 let (width, height, depth) = match *chunk_size.as_mut_slice() {
512 [ref mut rest @ .., depth, height, width] => {
513 for r in rest {
514 *r = 1;
515 }
516 (width, height, depth)
517 }
518 [height, width] => (width, height, 1),
519 [width] => (width, 1, 1),
520 [] => (1, 1, 1),
521 };
522
523 for mut slice in decoded.exact_chunks_mut(chunk_size.as_slice()) {
524 let (encoded_slice, rest) =
525 postcard::take_from_bytes::<Cow<[u8]>>(encoded).map_err(|err| {
526 EbccCodecError::SliceDecodeFailed {
527 source: EbccSliceError(err),
528 }
529 })?;
530 encoded = rest;
531
532 while slice.ndim() < 3 {
533 slice = slice.insert_axis(Axis(0));
534 }
535 #[expect(clippy::unwrap_used)]
536 let slice = slice.into_shape_with_order((depth, height, width)).unwrap();
539
540 ebcc::ebcc_decode_chunking_into(&encoded_slice, slice).map_err(|err| {
541 EbccCodecError::EbccDecodeFailed {
542 source: EbccCodingError(err),
543 }
544 })?;
545 }
546
547 if !encoded.is_empty() {
548 return Err(EbccCodecError::DecodeTooManySlices);
549 }
550
551 Ok(())
552}
553
554#[derive(Serialize, Deserialize)]
555struct CompressionHeader<'a> {
556 dtype: EbccDType,
557 #[serde(borrow)]
558 shape: Cow<'a, [usize]>,
559 version: EbccCodecVersion,
560}
561
562#[derive(Copy, Clone, Debug, Serialize, Deserialize)]
564enum EbccDType {
565 #[serde(rename = "f32", alias = "float32")]
566 F32,
567}
568
569#[cfg(test)]
570mod tests {
571 use super::*;
572
573 #[test]
574 fn test_unsupported_dtype() {
575 let codec = EbccCodec {
576 residual: EbccResidualType::Jpeg2000Only,
577 base_cr: Positive(10.0),
578 version: StaticCodecVersion,
579 chunk_shape: EbccChunkShape::Auto,
580 };
581
582 let data = Array1::<i32>::zeros(100);
583 let result = codec.encode(AnyCowArray::I32(data.into_dyn().into()));
584
585 assert!(matches!(result, Err(EbccCodecError::UnsupportedDtype(_))));
586 }
587
588 #[test]
589 fn test_invalid_dimensions() {
590 let codec = EbccCodec {
591 residual: EbccResidualType::Jpeg2000Only,
592 base_cr: Positive(10.0),
593 version: StaticCodecVersion,
594 chunk_shape: EbccChunkShape::Auto,
595 };
596
597 let data = Array::zeros(32);
599 let result = codec.encode(AnyCowArray::F32(data.into_dyn().into()));
600 assert!(
601 matches!(result, Err(EbccCodecError::InsufficientDimensions { shape }) if shape == [32])
602 );
603
604 let data = Array::zeros((16, 16));
606 let result = codec.encode(AnyCowArray::F32(data.into_dyn().into()));
607 assert!(
608 matches!(result, Err(EbccCodecError::InsufficientDimensions { shape }) if shape == [16, 16])
609 );
610
611 let data = Array::zeros((1, 32, 16));
613 let result = codec.encode(AnyCowArray::F32(data.into_dyn().into()));
614 assert!(
615 matches!(result, Err(EbccCodecError::InsufficientDimensions { shape }) if shape == [1, 32, 16])
616 );
617
618 let data = Array::zeros((1, 32, 32));
620 let result = codec.encode(AnyCowArray::F32(data.into_dyn().into()));
621 assert!(result.is_ok());
622
623 let data = Array::zeros((2, 2, 2, 32, 32));
625 let result = codec.encode(AnyCowArray::F32(data.into_dyn().into()));
626 assert!(result.is_ok());
627 }
628
629 #[test]
630 fn test_large_array() -> Result<(), EbccCodecError> {
631 let height = 721; let width = 1440;
634 let frames = 1;
635
636 #[expect(clippy::suboptimal_flops, clippy::cast_precision_loss)]
637 let data = Array::from_shape_fn((frames, height, width), |(_k, i, j)| {
638 let lat = -90.0 + (i as f32 / height as f32) * 180.0;
639 let lon = -180.0 + (j as f32 / width as f32) * 360.0;
640 #[allow(clippy::let_and_return)]
641 let temp = 273.15 + 30.0 * (1.0 - lat.abs() / 90.0) + 5.0 * (lon / 180.0).sin();
642 temp
643 });
644
645 let codec_error = 0.1;
646 let codec = EbccCodec {
647 residual: EbccResidualType::AbsoluteError {
648 error: Positive(codec_error),
649 },
650 base_cr: Positive(20.0),
651 version: StaticCodecVersion,
652 chunk_shape: EbccChunkShape::Auto,
653 };
654
655 let encoded = codec.encode(AnyArray::F32(data.clone().into_dyn()).into_cow())?;
656 let decoded = codec.decode(encoded.cow())?;
657
658 let AnyArray::U8(encoded) = encoded else {
659 return Err(EbccCodecError::EncodedDataNotBytes {
660 dtype: encoded.dtype(),
661 });
662 };
663
664 let AnyArray::F32(decoded) = decoded else {
665 return Err(EbccCodecError::UnsupportedDtype(decoded.dtype()));
666 };
667
668 let original_size = data.len() * std::mem::size_of::<f32>();
670 #[allow(clippy::cast_precision_loss)]
671 let compression_ratio = original_size as f64 / encoded.len() as f64;
672
673 assert!(
674 compression_ratio > 5.0,
675 "Compression ratio {compression_ratio} should be at least 5:1",
676 );
677
678 let max_error = data
680 .iter()
681 .zip(decoded.iter())
682 .map(|(&orig, &decomp)| (orig - decomp).abs())
683 .fold(0.0f32, f32::max);
684
685 assert!(
686 max_error <= (codec_error + 1e-6),
687 "Max error {max_error} exceeds error bound {codec_error}",
688 );
689
690 Ok(())
691 }
692
693 #[test]
694 fn test_huge_array() -> Result<(), EbccCodecError> {
695 let height = 721 * 2; let width = 1440 * 2;
698 let frames = 2;
699
700 #[expect(clippy::suboptimal_flops, clippy::cast_precision_loss)]
701 let data = Array::from_shape_fn((frames, height, width), |(_k, i, j)| {
702 let lat = -90.0 + (i as f32 / height as f32) * 180.0;
703 let lon = -180.0 + (j as f32 / width as f32) * 360.0;
704 #[allow(clippy::let_and_return)]
705 let temp = 273.15 + 30.0 * (1.0 - lat.abs() / 90.0) + 5.0 * (lon / 180.0).sin();
706 temp
707 });
708
709 let codec_error = 0.1;
710 let codec = EbccCodec {
711 residual: EbccResidualType::AbsoluteError {
712 error: Positive(codec_error),
713 },
714 base_cr: Positive(20.0),
715 version: StaticCodecVersion,
716 chunk_shape: EbccChunkShape::Auto,
717 };
718
719 let encoded = codec.encode(AnyArray::F32(data.clone().into_dyn()).into_cow())?;
720 let decoded = codec.decode(encoded.cow())?;
721
722 let AnyArray::U8(encoded) = encoded else {
723 return Err(EbccCodecError::EncodedDataNotBytes {
724 dtype: encoded.dtype(),
725 });
726 };
727
728 let AnyArray::F32(decoded) = decoded else {
729 return Err(EbccCodecError::UnsupportedDtype(decoded.dtype()));
730 };
731
732 let original_size = data.len() * std::mem::size_of::<f32>();
734 #[allow(clippy::cast_precision_loss)]
735 let compression_ratio = original_size as f64 / encoded.len() as f64;
736
737 assert!(
738 compression_ratio > 5.0,
739 "Compression ratio {compression_ratio} should be at least 5:1",
740 );
741
742 let max_error = data
744 .iter()
745 .zip(decoded.iter())
746 .map(|(&orig, &decomp)| (orig - decomp).abs())
747 .fold(0.0f32, f32::max);
748
749 assert!(
750 max_error <= (codec_error + 1e-6),
751 "Max error {max_error} exceeds error bound {codec_error}",
752 );
753
754 Ok(())
755 }
756}