grib_reader/
lib.rs

1//! Pure-Rust GRIB file reader.
2//!
3//! The current implementation supports the production-critical baseline for both
4//! GRIB1 and GRIB2: regular latitude/longitude grids, simple packing, and
5//! GRIB2 complex packing with general group splitting.
6//!
7//! # Example
8//!
9//! ```no_run
10//! use grib_reader::GribFile;
11//!
12//! let file = GribFile::open("gfs.grib2")?;
13//! println!("messages: {}", file.message_count());
14//!
15//! for msg in file.messages() {
16//!     println!(
17//!         "  {} {:?} {:04}-{:02}-{:02}T{:02}:{:02}:{:02}Z",
18//!         msg.parameter_name(),
19//!         msg.grid_shape(),
20//!         msg.reference_time().year,
21//!         msg.reference_time().month,
22//!         msg.reference_time().day,
23//!         msg.reference_time().hour,
24//!         msg.reference_time().minute,
25//!         msg.reference_time().second,
26//!     );
27//! }
28//!
29//! let data = file.message(0)?.read_data_as_f64()?;
30//! println!("shape: {:?}", data.shape());
31//! # Ok::<(), grib_reader::Error>(())
32//! ```
33
34pub mod data;
35pub mod error;
36pub mod grib1;
37pub mod grid;
38pub mod indicator;
39pub mod metadata;
40pub mod parameter;
41pub mod product;
42pub mod sections;
43
44pub use data::{
45    ComplexPackingParams, DataRepresentation, DecodeSample, SimplePackingParams,
46    SpatialDifferencingParams,
47};
48pub use error::{Error, Result};
49pub use grib1::{BinaryDataSection, GridDescription, ProductDefinition as Grib1ProductDefinition};
50pub use grid::{GridDefinition, LatLonGrid};
51pub use metadata::{ForecastTimeUnit, Parameter, ReferenceTime};
52pub use product::{
53    AnalysisOrForecastTemplate, FixedSurface, Identification, ProductDefinition,
54    ProductDefinitionTemplate,
55};
56
57use std::path::Path;
58
59use memmap2::Mmap;
60use ndarray::{ArrayD, IxDyn};
61
62use crate::data::{
63    bitmap_payload as grib2_bitmap_payload, count_bitmap_present_points, decode_field_into,
64    decode_payload_into,
65};
66use crate::indicator::Indicator;
67use crate::sections::{index_fields, FieldSections, SectionRef};
68
69#[cfg(feature = "rayon")]
70use rayon::prelude::*;
71
72const GRIB_MAGIC: &[u8; 4] = b"GRIB";
73
74/// Configuration for opening GRIB data.
75#[derive(Debug, Clone, Copy)]
76pub struct OpenOptions {
77    /// When `true`, the first malformed GRIB candidate aborts opening.
78    ///
79    /// When `false`, candidate offsets with invalid framing are skipped and
80    /// scanning continues. Once a candidate has a valid indicator, message
81    /// length, and end marker, any indexing or decoding error is returned.
82    pub strict: bool,
83}
84
85impl Default for OpenOptions {
86    fn default() -> Self {
87        Self { strict: true }
88    }
89}
90
91/// A parsed GRIB field.
92#[derive(Debug, Clone)]
93pub struct MessageMetadata {
94    pub edition: u8,
95    pub center_id: u16,
96    pub subcenter_id: u16,
97    pub discipline: Option<u8>,
98    pub reference_time: ReferenceTime,
99    pub parameter: Parameter,
100    pub grid: GridDefinition,
101    pub data_representation: DataRepresentation,
102    pub forecast_time_unit: Option<u8>,
103    pub forecast_time: Option<u32>,
104    pub message_offset: u64,
105    pub message_length: u64,
106    pub field_index_in_message: usize,
107    grib1_product: Option<grib1::ProductDefinition>,
108    grib2_identification: Option<Identification>,
109    grib2_product: Option<ProductDefinition>,
110}
111
112/// A GRIB file containing one or more logical fields.
113pub struct GribFile {
114    data: GribData,
115    messages: Vec<MessageIndex>,
116}
117
118#[derive(Clone)]
119struct MessageIndex {
120    offset: usize,
121    length: usize,
122    metadata: MessageMetadata,
123    decode_plan: DecodePlan,
124}
125
126#[derive(Clone, Copy)]
127enum DecodePlan {
128    Grib1 {
129        bitmap: Option<SectionRef>,
130        data: SectionRef,
131    },
132    Grib2(FieldSections),
133}
134
135enum GribData {
136    Mmap(Mmap),
137    Bytes(Vec<u8>),
138}
139
140impl GribData {
141    fn as_bytes(&self) -> &[u8] {
142        match self {
143            GribData::Mmap(m) => m,
144            GribData::Bytes(b) => b,
145        }
146    }
147}
148
149impl GribFile {
150    /// Open a GRIB file from disk using memory-mapped I/O.
151    pub fn open<P: AsRef<Path>>(path: P) -> Result<Self> {
152        Self::open_with_options(path, OpenOptions::default())
153    }
154
155    /// Open a GRIB file from disk using explicit decoder options.
156    pub fn open_with_options<P: AsRef<Path>>(path: P, options: OpenOptions) -> Result<Self> {
157        let file = std::fs::File::open(path.as_ref())
158            .map_err(|e| Error::Io(e, path.as_ref().display().to_string()))?;
159        let mmap = unsafe { Mmap::map(&file) }
160            .map_err(|e| Error::Io(e, path.as_ref().display().to_string()))?;
161        Self::from_data(GribData::Mmap(mmap), options)
162    }
163
164    /// Open a GRIB file from an owned byte buffer.
165    pub fn from_bytes(data: Vec<u8>) -> Result<Self> {
166        Self::from_bytes_with_options(data, OpenOptions::default())
167    }
168
169    /// Open a GRIB file from an owned byte buffer using explicit decoder options.
170    pub fn from_bytes_with_options(data: Vec<u8>, options: OpenOptions) -> Result<Self> {
171        Self::from_data(GribData::Bytes(data), options)
172    }
173
174    fn from_data(data: GribData, options: OpenOptions) -> Result<Self> {
175        let messages = scan_messages(data.as_bytes(), options)?;
176        if messages.is_empty() {
177            return Err(Error::NoMessages);
178        }
179        Ok(Self { data, messages })
180    }
181
182    /// Returns the GRIB edition of the first field.
183    pub fn edition(&self) -> u8 {
184        self.messages
185            .first()
186            .map(|message| message.metadata.edition)
187            .unwrap_or(0)
188    }
189
190    /// Returns the number of logical fields in the file.
191    pub fn message_count(&self) -> usize {
192        self.messages.len()
193    }
194
195    /// Access a field by index.
196    pub fn message(&self, index: usize) -> Result<Message<'_>> {
197        let record = self
198            .messages
199            .get(index)
200            .ok_or(Error::MessageNotFound(index))?;
201        let bytes = &self.data.as_bytes()[record.offset..record.offset + record.length];
202        Ok(Message {
203            index,
204            bytes,
205            metadata: &record.metadata,
206            decode_plan: record.decode_plan,
207        })
208    }
209
210    /// Iterate over all fields.
211    pub fn messages(&self) -> MessageIter<'_> {
212        MessageIter {
213            file: self,
214            index: 0,
215        }
216    }
217
218    /// Decode every field in the file.
219    pub fn read_all_data_as_f64(&self) -> Result<Vec<ArrayD<f64>>> {
220        #[cfg(feature = "rayon")]
221        {
222            (0..self.message_count())
223                .into_par_iter()
224                .map(|index| self.message(index)?.read_data_as_f64())
225                .collect()
226        }
227        #[cfg(not(feature = "rayon"))]
228        {
229            (0..self.message_count())
230                .map(|index| self.message(index)?.read_data_as_f64())
231                .collect()
232        }
233    }
234
235    /// Decode every field in the file as `f32`.
236    pub fn read_all_data_as_f32(&self) -> Result<Vec<ArrayD<f32>>> {
237        #[cfg(feature = "rayon")]
238        {
239            (0..self.message_count())
240                .into_par_iter()
241                .map(|index| self.message(index)?.read_data_as_f32())
242                .collect()
243        }
244        #[cfg(not(feature = "rayon"))]
245        {
246            (0..self.message_count())
247                .map(|index| self.message(index)?.read_data_as_f32())
248                .collect()
249        }
250    }
251}
252
253/// A single logical GRIB field.
254pub struct Message<'a> {
255    bytes: &'a [u8],
256    metadata: &'a MessageMetadata,
257    decode_plan: DecodePlan,
258    index: usize,
259}
260
261impl<'a> Message<'a> {
262    pub fn edition(&self) -> u8 {
263        self.metadata.edition
264    }
265
266    pub fn index(&self) -> usize {
267        self.index
268    }
269
270    pub fn metadata(&self) -> &MessageMetadata {
271        self.metadata
272    }
273
274    pub fn reference_time(&self) -> &ReferenceTime {
275        &self.metadata.reference_time
276    }
277
278    pub fn parameter(&self) -> &Parameter {
279        &self.metadata.parameter
280    }
281
282    pub fn center_id(&self) -> u16 {
283        self.metadata.center_id
284    }
285
286    pub fn subcenter_id(&self) -> u16 {
287        self.metadata.subcenter_id
288    }
289
290    pub fn identification(&self) -> Option<&Identification> {
291        self.metadata.grib2_identification.as_ref()
292    }
293
294    pub fn product_definition(&self) -> Option<&ProductDefinition> {
295        self.metadata.grib2_product.as_ref()
296    }
297
298    pub fn grib1_product_definition(&self) -> Option<&grib1::ProductDefinition> {
299        self.metadata.grib1_product.as_ref()
300    }
301
302    pub fn grid_definition(&self) -> &GridDefinition {
303        &self.metadata.grid
304    }
305
306    pub fn parameter_name(&self) -> &'static str {
307        self.metadata.parameter.short_name
308    }
309
310    pub fn parameter_description(&self) -> &'static str {
311        self.metadata.parameter.description
312    }
313
314    pub fn forecast_time_unit(&self) -> Option<u8> {
315        self.metadata.forecast_time_unit
316    }
317
318    pub fn forecast_time_unit_kind(&self) -> Option<ForecastTimeUnit> {
319        let unit = self.metadata.forecast_time_unit?;
320        ForecastTimeUnit::from_edition_and_code(self.metadata.edition, unit)
321    }
322
323    pub fn forecast_time(&self) -> Option<u32> {
324        self.metadata.forecast_time
325    }
326
327    pub fn valid_time(&self) -> Option<ReferenceTime> {
328        let unit = self.metadata.forecast_time_unit?;
329        let lead = self.metadata.forecast_time?;
330        self.metadata
331            .reference_time
332            .checked_add_forecast_time_by_edition(self.metadata.edition, unit, lead)
333    }
334
335    pub fn grid_shape(&self) -> (usize, usize) {
336        self.metadata.grid.shape()
337    }
338
339    pub fn latitudes(&self) -> Option<Vec<f64>> {
340        match &self.metadata.grid {
341            GridDefinition::LatLon(grid) => Some(grid.latitudes()),
342            GridDefinition::Unsupported(_) => None,
343        }
344    }
345
346    pub fn longitudes(&self) -> Option<Vec<f64>> {
347        match &self.metadata.grid {
348            GridDefinition::LatLon(grid) => Some(grid.longitudes()),
349            GridDefinition::Unsupported(_) => None,
350        }
351    }
352
353    pub fn decode_into<T: DecodeSample>(&self, out: &mut [T]) -> Result<()> {
354        let grid = match &self.metadata.grid {
355            GridDefinition::LatLon(grid) => grid,
356            GridDefinition::Unsupported(template) => {
357                return Err(Error::UnsupportedGridTemplate(*template));
358            }
359        };
360
361        match self.decode_plan {
362            DecodePlan::Grib2(field) => {
363                let data_section = section_bytes(self.bytes, field.data);
364                let bitmap_section = match field.bitmap {
365                    Some(section) => grib2_bitmap_payload(section_bytes(self.bytes, section))?,
366                    None => None,
367                };
368                decode_field_into(
369                    data_section,
370                    &self.metadata.data_representation,
371                    bitmap_section,
372                    self.metadata.grid.num_points(),
373                    out,
374                )?
375            }
376            DecodePlan::Grib1 { bitmap, data } => {
377                let bitmap_section = match bitmap {
378                    Some(section) => grib1::bitmap_payload(section_bytes(self.bytes, section))?,
379                    None => None,
380                };
381                let data_section = section_bytes(self.bytes, data);
382                if data_section.len() < 11 {
383                    return Err(Error::InvalidSection {
384                        section: 4,
385                        reason: format!("expected at least 11 bytes, got {}", data_section.len()),
386                    });
387                }
388                decode_payload_into(
389                    &data_section[11..],
390                    &self.metadata.data_representation,
391                    bitmap_section,
392                    self.metadata.grid.num_points(),
393                    out,
394                )?
395            }
396        }
397
398        grid.reorder_for_ndarray_in_place(out)
399    }
400
401    pub fn read_flat_data_as_f64(&self) -> Result<Vec<f64>> {
402        let mut decoded = vec![0.0; self.metadata.grid.num_points()];
403        self.decode_into(&mut decoded)?;
404        Ok(decoded)
405    }
406
407    pub fn read_flat_data_as_f32(&self) -> Result<Vec<f32>> {
408        let mut decoded = vec![0.0_f32; self.metadata.grid.num_points()];
409        self.decode_into(&mut decoded)?;
410        Ok(decoded)
411    }
412
413    pub fn read_data_as_f64(&self) -> Result<ArrayD<f64>> {
414        let ordered = self.read_flat_data_as_f64()?;
415        ArrayD::from_shape_vec(IxDyn(&self.metadata.grid.ndarray_shape()), ordered)
416            .map_err(|e| Error::Other(format!("failed to build ndarray from decoded field: {e}")))
417    }
418
419    pub fn read_data_as_f32(&self) -> Result<ArrayD<f32>> {
420        let ordered = self.read_flat_data_as_f32()?;
421        ArrayD::from_shape_vec(IxDyn(&self.metadata.grid.ndarray_shape()), ordered)
422            .map_err(|e| Error::Other(format!("failed to build ndarray from decoded field: {e}")))
423    }
424
425    pub fn raw_bytes(&self) -> &[u8] {
426        self.bytes
427    }
428}
429
430/// Iterator over fields in a GRIB file.
431pub struct MessageIter<'a> {
432    file: &'a GribFile,
433    index: usize,
434}
435
436impl<'a> Iterator for MessageIter<'a> {
437    type Item = Message<'a>;
438
439    fn next(&mut self) -> Option<Self::Item> {
440        if self.index >= self.file.message_count() {
441            return None;
442        }
443        let message = self.file.message(self.index).ok()?;
444        self.index += 1;
445        Some(message)
446    }
447}
448
449fn section_bytes(msg_bytes: &[u8], section: SectionRef) -> &[u8] {
450    &msg_bytes[section.offset..section.offset + section.length]
451}
452
453fn scan_messages(data: &[u8], options: OpenOptions) -> Result<Vec<MessageIndex>> {
454    let mut messages = Vec::new();
455    let mut pos = 0usize;
456
457    while pos + 8 <= data.len() {
458        if &data[pos..pos + 4] != GRIB_MAGIC {
459            pos += 1;
460            continue;
461        }
462
463        let (indicator, next_pos) = match locate_message(data, pos) {
464            Ok(located) => located,
465            Err(err) if !options.strict && is_recoverable_candidate_error(&err) => {
466                pos += 4;
467                continue;
468            }
469            Err(err) => return Err(err),
470        };
471
472        let message_bytes = &data[pos..next_pos];
473        let indexed = match indicator.edition {
474            1 => index_grib1_message(message_bytes, pos)?,
475            2 => index_grib2_message(message_bytes, pos, &indicator)?,
476            other => return Err(Error::UnsupportedEdition(other)),
477        };
478
479        messages.extend(indexed);
480        pos = next_pos;
481    }
482
483    Ok(messages)
484}
485
486fn is_recoverable_candidate_error(err: &Error) -> bool {
487    matches!(err, Error::InvalidMessage(_) | Error::Truncated { .. })
488}
489
490fn locate_message(data: &[u8], pos: usize) -> Result<(Indicator, usize)> {
491    let indicator = Indicator::parse(&data[pos..]).ok_or_else(|| {
492        Error::InvalidMessage(format!("failed to parse indicator at byte offset {pos}"))
493    })?;
494    let length = indicator.total_length as usize;
495    if length < 12 {
496        return Err(Error::InvalidMessage(format!(
497            "message at byte offset {pos} reports impossible length {length}"
498        )));
499    }
500    let end = pos
501        .checked_add(length)
502        .ok_or_else(|| Error::InvalidMessage("message length overflow".into()))?;
503    if end > data.len() {
504        return Err(Error::Truncated { offset: end as u64 });
505    }
506    if &data[end - 4..end] != b"7777" {
507        return Err(Error::InvalidMessage(format!(
508            "message at byte offset {pos} does not end with 7777"
509        )));
510    }
511
512    Ok((indicator, end))
513}
514
515fn index_grib1_message(message_bytes: &[u8], offset: usize) -> Result<Vec<MessageIndex>> {
516    let sections = grib1::parse_message_sections(message_bytes)?;
517    let grid_ref = sections.grid.ok_or_else(|| {
518        Error::InvalidSectionOrder(
519            "GRIB1 decoding requires an explicit grid definition section".into(),
520        )
521    })?;
522    let grid_description = GridDescription::parse(section_bytes(message_bytes, grid_ref))?;
523    let grid = grid_description.grid;
524
525    let bitmap_present_count = match sections.bitmap {
526        Some(bitmap) => {
527            let bitmap_payload = grib1::bitmap_payload(section_bytes(message_bytes, bitmap))?
528                .ok_or(Error::MissingBitmap)?;
529            count_bitmap_present_points(bitmap_payload, grid.num_points())?
530        }
531        None => grid.num_points(),
532    };
533
534    let (_bds, data_representation) = BinaryDataSection::parse(
535        section_bytes(message_bytes, sections.data),
536        sections.product.decimal_scale,
537        bitmap_present_count,
538    )?;
539    let parameter = sections.product.parameter();
540
541    Ok(vec![MessageIndex {
542        offset,
543        length: message_bytes.len(),
544        decode_plan: DecodePlan::Grib1 {
545            bitmap: sections.bitmap,
546            data: sections.data,
547        },
548        metadata: MessageMetadata {
549            edition: 1,
550            center_id: sections.product.center_id as u16,
551            subcenter_id: sections.product.subcenter_id as u16,
552            discipline: None,
553            reference_time: sections.product.reference_time,
554            parameter,
555            grid,
556            data_representation,
557            forecast_time_unit: Some(sections.product.forecast_time_unit),
558            forecast_time: sections.product.forecast_time(),
559            message_offset: offset as u64,
560            message_length: message_bytes.len() as u64,
561            field_index_in_message: 0,
562            grib1_product: Some(sections.product),
563            grib2_identification: None,
564            grib2_product: None,
565        },
566    }])
567}
568
569fn index_grib2_message(
570    message_bytes: &[u8],
571    offset: usize,
572    indicator: &Indicator,
573) -> Result<Vec<MessageIndex>> {
574    let fields = index_fields(message_bytes)?;
575    let mut messages = Vec::with_capacity(fields.len());
576
577    for (field_index_in_message, field_sections) in fields.into_iter().enumerate() {
578        let identification =
579            Identification::parse(section_bytes(message_bytes, field_sections.identification))?;
580        let grid = GridDefinition::parse(section_bytes(message_bytes, field_sections.grid))?;
581        let product =
582            ProductDefinition::parse(section_bytes(message_bytes, field_sections.product))?;
583        let data_representation = DataRepresentation::parse(section_bytes(
584            message_bytes,
585            field_sections.data_representation,
586        ))?;
587        let parameter = Parameter::new_grib2(
588            indicator.discipline,
589            product.parameter_category,
590            product.parameter_number,
591            product.parameter_name(indicator.discipline),
592            product.parameter_description(indicator.discipline),
593        );
594
595        messages.push(MessageIndex {
596            offset,
597            length: message_bytes.len(),
598            decode_plan: DecodePlan::Grib2(field_sections),
599            metadata: MessageMetadata {
600                edition: indicator.edition,
601                center_id: identification.center_id,
602                subcenter_id: identification.subcenter_id,
603                discipline: Some(indicator.discipline),
604                reference_time: ReferenceTime {
605                    year: identification.reference_year,
606                    month: identification.reference_month,
607                    day: identification.reference_day,
608                    hour: identification.reference_hour,
609                    minute: identification.reference_minute,
610                    second: identification.reference_second,
611                },
612                parameter,
613                grid,
614                data_representation,
615                forecast_time_unit: product.forecast_time_unit(),
616                forecast_time: product.forecast_time(),
617                message_offset: offset as u64,
618                message_length: message_bytes.len() as u64,
619                field_index_in_message,
620                grib1_product: None,
621                grib2_identification: Some(identification),
622                grib2_product: Some(product),
623            },
624        });
625    }
626
627    Ok(messages)
628}
629
630#[cfg(test)]
631mod tests {
632    use super::*;
633
634    fn grib_i32_bytes(value: i32) -> [u8; 4] {
635        if value >= 0 {
636            (value as u32).to_be_bytes()
637        } else {
638            ((-value) as u32 | 0x8000_0000).to_be_bytes()
639        }
640    }
641
642    fn build_indicator(total_len: usize, discipline: u8) -> Vec<u8> {
643        let mut indicator = Vec::with_capacity(16);
644        indicator.extend_from_slice(b"GRIB");
645        indicator.extend_from_slice(&[0, 0]);
646        indicator.push(discipline);
647        indicator.push(2);
648        indicator.extend_from_slice(&(total_len as u64).to_be_bytes());
649        indicator
650    }
651
652    fn build_identification() -> Vec<u8> {
653        let mut section = vec![0u8; 21];
654        section[..4].copy_from_slice(&(21u32).to_be_bytes());
655        section[4] = 1;
656        section[5..7].copy_from_slice(&7u16.to_be_bytes());
657        section[7..9].copy_from_slice(&0u16.to_be_bytes());
658        section[9] = 35;
659        section[10] = 1;
660        section[11] = 1;
661        section[12..14].copy_from_slice(&2026u16.to_be_bytes());
662        section[14] = 3;
663        section[15] = 20;
664        section[16] = 12;
665        section[17] = 0;
666        section[18] = 0;
667        section[19] = 0;
668        section[20] = 1;
669        section
670    }
671
672    fn build_grid(ni: u32, nj: u32, scanning_mode: u8) -> Vec<u8> {
673        let mut section = vec![0u8; 72];
674        section[..4].copy_from_slice(&(72u32).to_be_bytes());
675        section[4] = 3;
676        section[6..10].copy_from_slice(&(ni * nj).to_be_bytes());
677        section[12..14].copy_from_slice(&0u16.to_be_bytes());
678        section[30..34].copy_from_slice(&ni.to_be_bytes());
679        section[34..38].copy_from_slice(&nj.to_be_bytes());
680        section[46..50].copy_from_slice(&grib_i32_bytes(50_000_000));
681        section[50..54].copy_from_slice(&grib_i32_bytes(-120_000_000));
682        section[55..59].copy_from_slice(&grib_i32_bytes(49_000_000));
683        section[59..63].copy_from_slice(&grib_i32_bytes(-119_000_000));
684        section[63..67].copy_from_slice(&1_000_000u32.to_be_bytes());
685        section[67..71].copy_from_slice(&1_000_000u32.to_be_bytes());
686        section[71] = scanning_mode;
687        section
688    }
689
690    fn build_product(parameter_category: u8, parameter_number: u8) -> Vec<u8> {
691        let mut section = vec![0u8; 34];
692        section[..4].copy_from_slice(&(34u32).to_be_bytes());
693        section[4] = 4;
694        section[7..9].copy_from_slice(&0u16.to_be_bytes());
695        section[9] = parameter_category;
696        section[10] = parameter_number;
697        section[11] = 2;
698        section[17] = 1;
699        section[18..22].copy_from_slice(&0u32.to_be_bytes());
700        section[22] = 103;
701        section[23] = 0;
702        section[24..28].copy_from_slice(&850u32.to_be_bytes());
703        section[28] = 255;
704        section
705    }
706
707    fn build_simple_representation(encoded_values: usize, bits_per_value: u8) -> Vec<u8> {
708        let mut section = vec![0u8; 21];
709        section[..4].copy_from_slice(&(21u32).to_be_bytes());
710        section[4] = 5;
711        section[5..9].copy_from_slice(&(encoded_values as u32).to_be_bytes());
712        section[9..11].copy_from_slice(&0u16.to_be_bytes());
713        section[11..15].copy_from_slice(&0f32.to_be_bytes());
714        section[19] = bits_per_value;
715        section[20] = 0;
716        section
717    }
718
719    fn pack_u8_values(values: &[u8]) -> Vec<u8> {
720        values.to_vec()
721    }
722
723    fn build_bitmap(bits: &[bool]) -> Vec<u8> {
724        let payload_len = bits.len().div_ceil(8) + 1;
725        let mut section = vec![0u8; payload_len + 5];
726        section[..4].copy_from_slice(&((payload_len + 5) as u32).to_be_bytes());
727        section[4] = 6;
728        section[5] = 0;
729        for (index, bit) in bits.iter().copied().enumerate() {
730            if bit {
731                section[6 + index / 8] |= 1 << (7 - (index % 8));
732            }
733        }
734        section
735    }
736
737    fn build_data(payload: &[u8]) -> Vec<u8> {
738        let mut section = vec![0u8; payload.len() + 5];
739        section[..4].copy_from_slice(&((payload.len() + 5) as u32).to_be_bytes());
740        section[4] = 7;
741        section[5..].copy_from_slice(payload);
742        section
743    }
744
745    fn assemble_grib2_message(sections: &[Vec<u8>]) -> Vec<u8> {
746        let total_len = 16 + sections.iter().map(|section| section.len()).sum::<usize>() + 4;
747        let mut message = build_indicator(total_len, 0);
748        for section in sections {
749            message.extend_from_slice(section);
750        }
751        message.extend_from_slice(b"7777");
752        message
753    }
754
755    fn build_grib1_message(values: &[u8]) -> Vec<u8> {
756        let mut pds = vec![0u8; 28];
757        pds[..3].copy_from_slice(&[0, 0, 28]);
758        pds[3] = 2;
759        pds[4] = 7;
760        pds[5] = 255;
761        pds[6] = 0;
762        pds[7] = 0b1000_0000;
763        pds[8] = 11;
764        pds[9] = 100;
765        pds[10..12].copy_from_slice(&850u16.to_be_bytes());
766        pds[12] = 26;
767        pds[13] = 3;
768        pds[14] = 20;
769        pds[15] = 12;
770        pds[16] = 0;
771        pds[17] = 1;
772        pds[18] = 0;
773        pds[19] = 0;
774        pds[20] = 0;
775        pds[24] = 21;
776        pds[25] = 0;
777
778        let mut gds = vec![0u8; 32];
779        gds[..3].copy_from_slice(&[0, 0, 32]);
780        gds[5] = 0;
781        gds[6..8].copy_from_slice(&2u16.to_be_bytes());
782        gds[8..10].copy_from_slice(&2u16.to_be_bytes());
783        gds[10..13].copy_from_slice(&[0x01, 0x4d, 0x50]);
784        gds[13..16].copy_from_slice(&[0x81, 0xd4, 0xc0]);
785        gds[16] = 0x80;
786        gds[17..20].copy_from_slice(&[0x01, 0x49, 0x68]);
787        gds[20..23].copy_from_slice(&[0x81, 0xd0, 0xd8]);
788        gds[23..25].copy_from_slice(&1000u16.to_be_bytes());
789        gds[25..27].copy_from_slice(&1000u16.to_be_bytes());
790
791        let mut bds = vec![0u8; 11 + values.len()];
792        let len = bds.len() as u32;
793        bds[..3].copy_from_slice(&[(len >> 16) as u8, (len >> 8) as u8, len as u8]);
794        bds[3] = 0;
795        bds[10] = 8;
796        bds[11..].copy_from_slice(values);
797
798        let total_len = 8 + pds.len() + gds.len() + bds.len() + 4;
799        let mut message = Vec::new();
800        message.extend_from_slice(b"GRIB");
801        message.extend_from_slice(&[
802            ((total_len >> 16) & 0xff) as u8,
803            ((total_len >> 8) & 0xff) as u8,
804            (total_len & 0xff) as u8,
805            1,
806        ]);
807        message.extend_from_slice(&pds);
808        message.extend_from_slice(&gds);
809        message.extend_from_slice(&bds);
810        message.extend_from_slice(b"7777");
811        message
812    }
813
814    #[test]
815    fn scans_single_grib2_message() {
816        let message = assemble_grib2_message(&[
817            build_identification(),
818            build_grid(2, 2, 0),
819            build_product(0, 0),
820            build_simple_representation(4, 8),
821            build_data(&pack_u8_values(&[1, 2, 3, 4])),
822        ]);
823        let messages = scan_messages(&message, OpenOptions::default()).unwrap();
824        assert_eq!(messages.len(), 1);
825        assert_eq!(messages[0].metadata.parameter.short_name, "TMP");
826    }
827
828    #[test]
829    fn decodes_simple_grib2_message_to_ndarray() {
830        let message = assemble_grib2_message(&[
831            build_identification(),
832            build_grid(2, 2, 0),
833            build_product(0, 0),
834            build_simple_representation(4, 8),
835            build_data(&pack_u8_values(&[1, 2, 3, 4])),
836        ]);
837        let file = GribFile::from_bytes(message).unwrap();
838        let array = file.message(0).unwrap().read_data_as_f64().unwrap();
839        assert_eq!(array.shape(), &[2, 2]);
840        assert_eq!(
841            array.iter().copied().collect::<Vec<_>>(),
842            vec![1.0, 2.0, 3.0, 4.0]
843        );
844    }
845
846    #[test]
847    fn applies_bitmap_to_missing_values() {
848        let message = assemble_grib2_message(&[
849            build_identification(),
850            build_grid(2, 2, 0),
851            build_product(0, 1),
852            build_simple_representation(3, 8),
853            build_bitmap(&[true, false, true, true]),
854            build_data(&pack_u8_values(&[10, 20, 30])),
855        ]);
856        let file = GribFile::from_bytes(message).unwrap();
857        let array = file.message(0).unwrap().read_data_as_f64().unwrap();
858        let values = array.iter().copied().collect::<Vec<_>>();
859        assert_eq!(values[0], 10.0);
860        assert!(values[1].is_nan());
861        assert_eq!(values[2], 20.0);
862        assert_eq!(values[3], 30.0);
863    }
864
865    #[test]
866    fn indexes_multiple_fields_in_one_grib2_message() {
867        let message = assemble_grib2_message(&[
868            build_identification(),
869            build_grid(2, 2, 0),
870            build_product(0, 0),
871            build_simple_representation(4, 8),
872            build_data(&pack_u8_values(&[1, 2, 3, 4])),
873            build_product(0, 2),
874            build_simple_representation(4, 8),
875            build_data(&pack_u8_values(&[5, 6, 7, 8])),
876        ]);
877        let file = GribFile::from_bytes(message).unwrap();
878        assert_eq!(file.message_count(), 2);
879        assert_eq!(file.message(0).unwrap().parameter_name(), "TMP");
880        assert_eq!(file.message(1).unwrap().parameter_name(), "POT");
881    }
882
883    #[test]
884    fn decodes_simple_grib1_message_to_ndarray() {
885        let message = build_grib1_message(&[1, 2, 3, 4]);
886        let file = GribFile::from_bytes(message).unwrap();
887        assert_eq!(file.edition(), 1);
888        let field = file.message(0).unwrap();
889        assert_eq!(field.parameter_name(), "TMP");
890        assert!(field.identification().is_none());
891        assert!(field.grib1_product_definition().is_some());
892        let array = field.read_data_as_f64().unwrap();
893        assert_eq!(array.shape(), &[2, 2]);
894        assert_eq!(
895            array.iter().copied().collect::<Vec<_>>(),
896            vec![1.0, 2.0, 3.0, 4.0]
897        );
898    }
899}
grib_reader/lib.rs

grib_reader/
lib.rs
