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
12pub(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
54pub(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
96pub(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
114pub(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#[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#[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}