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

grib_reader/
lib.rs
