Skip to main content

fpzip_rs/
encoder.rs

1use crate::codec::range_encoder::RangeEncoder;
2use crate::codec::rc_qs_model::RCQsModel;
3use crate::core::front::Front;
4use crate::core::pc_encoder::{symbol_count, PCEncoderDouble, PCEncoderFloat};
5use crate::core::pc_map;
6
7#[cfg(feature = "rayon")]
8extern crate alloc;
9#[cfg(feature = "rayon")]
10use alloc::vec::Vec;
11
12/// Compresses a 3D float array at the given bit precision.
13pub(crate) fn compress_3d_float(
14    encoder: &mut RangeEncoder,
15    data: &[f32],
16    nx: i32,
17    ny: i32,
18    nz: i32,
19    bits: u32,
20) {
21    let symbols = symbol_count(bits);
22    let mut model = RCQsModel::with_defaults(true, symbols);
23    let zero_mapped = pc_map::forward_f32(0.0f32, bits);
24    let mut front = Front::<u32>::new(nx, ny, zero_mapped);
25
26    let mut data_index = 0;
27
28    front.advance(0, 0, 1);
29    for _ in 0..nz {
30        front.advance(0, 1, 0);
31        for _ in 0..ny {
32            front.advance(1, 0, 0);
33            for _ in 0..nx {
34                let p = front
35                    .get(1, 0, 0)
36                    .wrapping_sub(front.get(0, 1, 1))
37                    .wrapping_add(front.get(0, 1, 0))
38                    .wrapping_sub(front.get(1, 0, 1))
39                    .wrapping_add(front.get(0, 0, 1))
40                    .wrapping_sub(front.get(1, 1, 0))
41                    .wrapping_add(front.get(1, 1, 1));
42
43                let a = pc_map::forward_f32(data[data_index], bits);
44                data_index += 1;
45
46                let mut pc = PCEncoderFloat::new(encoder, &mut model, bits);
47                let a = pc.encode(a, p);
48                front.push(a);
49            }
50        }
51    }
52}
53
54/// Compresses a 3D double array at the given bit precision.
55pub(crate) fn compress_3d_double(
56    encoder: &mut RangeEncoder,
57    data: &[f64],
58    nx: i32,
59    ny: i32,
60    nz: i32,
61    bits: u32,
62) {
63    let symbols = symbol_count(bits);
64    let mut model = RCQsModel::with_defaults(true, symbols);
65    let zero_mapped = pc_map::forward_f64(0.0f64, bits);
66    let mut front = Front::<u64>::new(nx, ny, zero_mapped);
67
68    let mut data_index = 0;
69
70    front.advance(0, 0, 1);
71    for _ in 0..nz {
72        front.advance(0, 1, 0);
73        for _ in 0..ny {
74            front.advance(1, 0, 0);
75            for _ in 0..nx {
76                let p = front
77                    .get(1, 0, 0)
78                    .wrapping_sub(front.get(0, 1, 1))
79                    .wrapping_add(front.get(0, 1, 0))
80                    .wrapping_sub(front.get(1, 0, 1))
81                    .wrapping_add(front.get(0, 0, 1))
82                    .wrapping_sub(front.get(1, 1, 0))
83                    .wrapping_add(front.get(1, 1, 1));
84
85                let a = pc_map::forward_f64(data[data_index], bits);
86                data_index += 1;
87
88                let mut pc = PCEncoderDouble::new(encoder, &mut model, bits);
89                let a = pc.encode(a, p);
90                front.push(a);
91            }
92        }
93    }
94}
95
96/// Compresses a 4D float array (multiple fields).
97pub(crate) fn compress_4d_float(
98    encoder: &mut RangeEncoder,
99    data: &[f32],
100    nx: i32,
101    ny: i32,
102    nz: i32,
103    nf: i32,
104    bits: u32,
105) {
106    let field_size = (nx as usize) * (ny as usize) * (nz as usize);
107    for f in 0..nf as usize {
108        let start = f * field_size;
109        let end = start + field_size;
110        compress_3d_float(encoder, &data[start..end], nx, ny, nz, bits);
111    }
112}
113
114/// Compresses a 4D double array (multiple fields).
115pub(crate) fn compress_4d_double(
116    encoder: &mut RangeEncoder,
117    data: &[f64],
118    nx: i32,
119    ny: i32,
120    nz: i32,
121    nf: i32,
122    bits: u32,
123) {
124    let field_size = (nx as usize) * (ny as usize) * (nz as usize);
125    for f in 0..nf as usize {
126        let start = f * field_size;
127        let end = start + field_size;
128        compress_3d_double(encoder, &data[start..end], nx, ny, nz, bits);
129    }
130}
131
132#[allow(dead_code)]
133/// Parallel compression of 4D float data (each field compressed independently).
134///
135/// **Note**: Produces different byte output than sequential mode because each field
136/// starts with a fresh encoder state. The output is only decompressible by
137/// `decompress_4d_float_parallel`.
138#[cfg(feature = "rayon")]
139pub(crate) fn compress_4d_float_parallel(
140    data: &[f32],
141    nx: i32,
142    ny: i32,
143    nz: i32,
144    nf: i32,
145    bits: u32,
146) -> Vec<Vec<u8>> {
147    use rayon::prelude::*;
148
149    let field_size = (nx as usize) * (ny as usize) * (nz as usize);
150    (0..nf as usize)
151        .into_par_iter()
152        .map(|f| {
153            let start = f * field_size;
154            let end = start + field_size;
155            let mut encoder = RangeEncoder::new();
156            compress_3d_float(&mut encoder, &data[start..end], nx, ny, nz, bits);
157            encoder.finish()
158        })
159        .collect()
160}
161
162#[allow(dead_code)]
163/// Parallel compression of 4D double data.
164#[cfg(feature = "rayon")]
165pub(crate) fn compress_4d_double_parallel(
166    data: &[f64],
167    nx: i32,
168    ny: i32,
169    nz: i32,
170    nf: i32,
171    bits: u32,
172) -> Vec<Vec<u8>> {
173    use rayon::prelude::*;
174
175    let field_size = (nx as usize) * (ny as usize) * (nz as usize);
176    (0..nf as usize)
177        .into_par_iter()
178        .map(|f| {
179            let start = f * field_size;
180            let end = start + field_size;
181            let mut encoder = RangeEncoder::new();
182            compress_3d_double(&mut encoder, &data[start..end], nx, ny, nz, bits);
183            encoder.finish()
184        })
185        .collect()
186}