conspire/io/netcdf/variable/
mod.rs1use crate::io::netcdf::{
2 DefineVariable, GetVariable, NcType, NetCDF, PutVariable, State, VarBuild,
3 format::{self, Storage, decode_be, decode_le, hdf5},
4 nc_lock, reject_nul,
5};
6use std::{
7 ffi::NulError,
8 io::{Seek, SeekFrom, Write},
9};
10
11impl DefineVariable for NetCDF {
12 fn define_variable<T: NcType>(
13 &mut self,
14 name: &str,
15 ndims: usize,
16 dim_names: &[&str],
17 ) -> Result<(), NulError> {
18 reject_nul(name)?;
19 for dim_name in dim_names {
20 reject_nul(dim_name)?;
21 }
22 assert_eq!(ndims, dim_names.len(), "ndims must equal dim_names.len()");
23 let _guard = nc_lock();
24 self.writer_defining().variables.push(VarBuild {
25 name: name.to_string(),
26 xtype: T::XTYPE,
27 dim_names: dim_names.iter().map(|dim| dim.to_string()).collect(),
28 attributes: Vec::new(),
29 });
30 Ok(())
31 }
32}
33
34impl PutVariable for NetCDF {
35 fn put_variable<T: NcType>(&mut self, name: &str, data: &[T]) -> Result<(), NulError> {
36 reject_nul(name)?;
37 let _guard = nc_lock();
38 let writer = match &mut self.state {
39 State::Write(writer) => writer,
40 State::Read(_) => panic!("put_variable on a NetCDF opened for reading"),
41 };
42 let netcdf4 = writer.netcdf4.is_some();
43 let dims = std::mem::take(&mut writer.dims);
44 let output = writer
45 .output
46 .as_mut()
47 .unwrap_or_else(|| panic!("put_variable before end_definition"));
48 let index = output
49 .variables
50 .iter()
51 .position(|spec| spec.name == name)
52 .unwrap_or_else(|| panic!("no variable named {name}"));
53 let spec = &output.variables[index];
54 assert_eq!(
55 spec.xtype,
56 T::XTYPE,
57 "type mismatch writing variable {name}"
58 );
59 if netcdf4 {
60 assert_eq!(
61 data.len() as u64,
62 spec.elements(&dims),
63 "wrong element count for variable {name}"
64 );
65 let mut buffer = Vec::with_capacity(data.len() * T::SIZE);
66 format::encode_le(data, &mut buffer);
67 output.data[index] = buffer;
68 } else {
69 let (begin, vsize) = (spec.begin, spec.vsize as usize);
70 assert_eq!(
71 data.len(),
72 vsize / T::SIZE,
73 "wrong element count for variable {name}"
74 );
75 let mut buffer = Vec::with_capacity(vsize);
76 format::encode_be(data, &mut buffer);
77 buffer.resize(vsize, 0);
78 output
79 .file
80 .seek(SeekFrom::Start(begin))
81 .expect("seek failed");
82 output
83 .file
84 .write_all(&buffer)
85 .expect("variable write failed");
86 }
87 writer.dims = dims;
88 Ok(())
89 }
90}
91
92impl NetCDF {
93 fn read_variable<T: NcType>(&self, name: &str, len: usize) -> Option<Vec<T>> {
94 let reader = match &self.state {
95 State::Read(reader) => reader,
96 State::Write(_) => panic!("get_variable on a NetCDF opened for writing"),
97 };
98 let spec = reader.parsed.vars.iter().find(|spec| spec.name == name)?;
99 assert_eq!(
100 spec.xtype,
101 T::XTYPE,
102 "type mismatch reading variable {name}"
103 );
104 match &spec.storage {
105 Storage::Classic => {
106 let start = spec.begin as usize;
107 let end = start + len * T::SIZE;
108 assert!(
109 end <= reader.bytes.len(),
110 "variable {name} data runs past end of file"
111 );
112 Some(decode_be(&reader.bytes[start..end]))
113 }
114 Storage::Hdf5 {
115 little_endian,
116 shape,
117 layout,
118 filters,
119 fill,
120 } => {
121 assert_eq!(
122 len as u64,
123 shape.iter().product::<u64>().max(1),
124 "wrong element count for variable {name}"
125 );
126 let decode = |bytes: &[u8]| {
127 if *little_endian {
128 decode_le(bytes)
129 } else {
130 decode_be(bytes)
131 }
132 };
133 if let (hdf5::Layout::Contiguous { addr, size }, true) =
134 (layout, filters.is_empty())
135 {
136 return Some(decode(&reader.bytes[*addr..*addr + *size]));
137 }
138 let start = vec![0usize; shape.len()];
139 let count: Vec<usize> = shape.iter().map(|&d| d as usize).collect();
140 let raw = hdf5::read_data(
141 &reader.bytes,
142 layout,
143 filters,
144 shape,
145 &start,
146 &count,
147 T::SIZE,
148 fill.as_deref().unwrap_or(&[]),
149 );
150 Some(decode(&raw))
151 }
152 }
153 }
154
155 fn read_variable_slice<T: NcType>(
156 &self,
157 name: &str,
158 start: &[usize],
159 count: &[usize],
160 ) -> Option<Vec<T>> {
161 let reader = match &self.state {
162 State::Read(reader) => reader,
163 State::Write(_) => panic!("get_variable_slice on a NetCDF opened for writing"),
164 };
165 let spec = reader.parsed.vars.iter().find(|spec| spec.name == name)?;
166 assert_eq!(
167 spec.xtype,
168 T::XTYPE,
169 "type mismatch reading variable {name}"
170 );
171 match &spec.storage {
172 Storage::Hdf5 {
173 little_endian,
174 shape,
175 layout,
176 filters,
177 fill,
178 } => {
179 check_slice(name, shape, start, count);
180 let raw = hdf5::read_data(
181 &reader.bytes,
182 layout,
183 filters,
184 shape,
185 start,
186 count,
187 T::SIZE,
188 fill.as_deref().unwrap_or(&[]),
189 );
190 Some(if *little_endian {
191 decode_le(&raw)
192 } else {
193 decode_be(&raw)
194 })
195 }
196 Storage::Classic => {
197 let shape = spec.shape(&reader.parsed.dims);
198 check_slice(name, &shape, start, count);
199 let bytes = shape.iter().product::<u64>().max(1) as usize * T::SIZE;
200 let layout = hdf5::Layout::Contiguous {
201 addr: spec.begin as usize,
202 size: bytes,
203 };
204 let raw = hdf5::read_data(
205 &reader.bytes,
206 &layout,
207 &[],
208 &shape,
209 start,
210 count,
211 T::SIZE,
212 &[],
213 );
214 Some(decode_be(&raw))
215 }
216 }
217 }
218}
219
220fn check_slice(name: &str, shape: &[u64], start: &[usize], count: &[usize]) {
221 assert_eq!(start.len(), shape.len(), "start has wrong rank for {name}");
222 assert_eq!(count.len(), shape.len(), "count has wrong rank for {name}");
223 for (i, (&s, &c)) in start.iter().zip(count).enumerate() {
224 assert!(
225 s + c <= shape[i] as usize,
226 "slice out of bounds on axis {i} of {name}"
227 );
228 }
229}
230
231impl GetVariable for NetCDF {
232 fn get_variable<T: NcType>(&self, name: &str, len: usize) -> Result<Vec<T>, NulError> {
233 reject_nul(name)?;
234 let _guard = nc_lock();
235 Ok(self
236 .read_variable(name, len)
237 .unwrap_or_else(|| panic!("no variable named {name}")))
238 }
239 fn try_get_variable<T: NcType>(
240 &self,
241 name: &str,
242 len: usize,
243 ) -> Result<Option<Vec<T>>, NulError> {
244 reject_nul(name)?;
245 let _guard = nc_lock();
246 Ok(self.read_variable(name, len))
247 }
248 fn get_variable_slice<T: NcType>(
249 &self,
250 name: &str,
251 start: &[usize],
252 count: &[usize],
253 ) -> Result<Vec<T>, NulError> {
254 reject_nul(name)?;
255 let _guard = nc_lock();
256 Ok(self
257 .read_variable_slice(name, start, count)
258 .unwrap_or_else(|| panic!("no variable named {name}")))
259 }
260}