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

grib_reader/
lib.rs
