wxbox/crates/ar2/src/lib.rs

use nexrad_data::volume::File;
use nexrad_decode::messages::MessageContents;
use nexrad_decode::messages::digital_radar_data::{GenericDataBlock, RadialStatus};
use nexrad_decode::result::Error;
use std::fmt::Debug;
use rayon::prelude::*;

pub mod sites;


pub struct Scan {
    pub coverage_pattern_number: u16,
    pub sweeps: Vec<Sweep>,
}

pub struct Sweep {
    pub elevation_number: u8,
    pub radials: Vec<Radial>,
}

pub struct Radial {
    pub collection_timestamp: i64,
    pub azimuth_number: u16,
    pub azimuth_angle_degrees: f32,
    pub azimuth_spacing_degrees: f32,
    pub radial_status: RadialStatus,
    pub elevation_number: u8,
    pub elevation_number_degrees: f32,
    pub reflectivity: Option<MomentData>,
    pub velocity: Option<MomentData>,
    pub spectrum_width: Option<MomentData>,
    pub differential_reflectivity: Option<MomentData>,
    pub differential_phase: Option<MomentData>,
    pub correlation_coefficient: Option<MomentData>,
    pub specific_differential_phase: Option<MomentData>,
}

#[derive(Debug)]
pub struct MomentData {
    pub scale: f32,
    pub offset: f32,
    pub values: Vec<u8>,
    pub start_range: u16,
    pub sample_interval: u16,
}

impl MomentData {
    /// Values from this data moment corresponding to gates in the radial.

    pub fn values(&self) -> Vec<MomentValue> {
        let copied_values = self.values.iter().copied();

        if self.scale == 0.0 {
            return copied_values
                .map(|raw_value| MomentValue::Value(raw_value as f32))
                .collect();
        }

        copied_values
            .map(|raw_value| match raw_value {
                0 => MomentValue::BelowThreshold,
                1 => MomentValue::RangeFolded,
                _ => MomentValue::Value((raw_value as f32 - self.offset) / self.scale),
            })
            .collect()
    }
}

/// The data moment value for a product in a radial's gate. The value may be a floating-point number

/// or a special case such as "below threshold" or "range folded".

#[derive(Debug, Clone, Copy, PartialEq)]

pub enum MomentValue {
    /// The data moment value for a gate.
    Value(f32),

    /// The value for this gate was below the signal threshold.
    BelowThreshold,

    /// The value for this gate exceeded the maximum unambiguous range.
    RangeFolded,
}

pub fn parse(input: Vec<u8>) -> nexrad_data::result::Result<Scan> {
    let file = File::new(input);

    let mut vcp: Option<u16> = None;
    let mut radials = vec![];

    let mut radial_chunks: Vec<nexrad_data::result::Result<(Vec<Radial>, Option<u16>)>> = file.records()
        .par_iter_mut()
        .map(|record| {
            let mut vcp = None;
            let mut radials = vec![];
            if record.compressed() {
                *record = record.decompress()?; // this specifically seems to be most problematic. bzip is SLOW
            }
            let messages = record.messages()?;

            for message in messages {
                let contents = message.contents();
                if let MessageContents::DigitalRadarData(radar_data_message) = contents {
                    if vcp.is_none() {
                        if let Some(volume_block) = &radar_data_message.volume_data_block {
                            vcp = Some(volume_block.volume_coverage_pattern_number);
                        }
                    }
                    radials.push(into_radial(*radar_data_message.clone())?)
                }
            }

            Ok((radials, vcp))
        })
        .collect();

    for radial_chunk in radial_chunks {
        let (mut radial_chunk, this_vcp) = radial_chunk?;
        if vcp.is_none() && this_vcp.is_some() {
            vcp = this_vcp;
        }
        radials.append(&mut radial_chunk);
    }

    Ok(Scan {
        coverage_pattern_number: vcp.ok_or(Error::DecodingError("no vcp".to_string()))?,
        sweeps: Sweep::from_radials(radials),
    })
}

#[inline]
fn into_radial(
    message: nexrad_decode::messages::digital_radar_data::Message,
) -> nexrad_data::result::Result<Radial> {
    Ok(Radial {
        collection_timestamp: message
            .header
            .date_time()
            .ok_or(Error::MessageMissingDateError)?
            .timestamp_millis(),
        azimuth_number: message.header.azimuth_number,
        azimuth_angle_degrees: message.header.azimuth_angle,
        azimuth_spacing_degrees: message.header.azimuth_resolution_spacing as f32 * 0.5,
        radial_status: message.header.radial_status(),
        elevation_number: message.header.elevation_number,
        elevation_number_degrees: message.header.elevation_angle,
        reflectivity: message.reflectivity_data_block.map(|u| into_moment_data(u)),
        velocity: message.velocity_data_block.map(|u| into_moment_data(u)),
        spectrum_width: message
            .spectrum_width_data_block
            .map(|u| into_moment_data(u)),
        differential_reflectivity: message
            .differential_reflectivity_data_block
            .map(|u| into_moment_data(u)),
        differential_phase: message
            .differential_phase_data_block
            .map(|u| into_moment_data(u)),
        correlation_coefficient: message
            .correlation_coefficient_data_block
            .map(|u| into_moment_data(u)),
        specific_differential_phase: message
            .specific_diff_phase_data_block
            .map(|u| into_moment_data(u)),
    })
}

#[inline(always)]
fn into_moment_data(block: GenericDataBlock) -> MomentData {
    MomentData {
        scale: block.header.scale,
        offset: block.header.offset,
        values: block.encoded_data,
        start_range: block.header.data_moment_range,
        sample_interval: block.header.data_moment_range_sample_interval,
    }
}

impl Sweep {
    pub fn new(elevation_number: u8, radials: Vec<Radial>) -> Self {
        Self {
            elevation_number,

            radials,
        }
    }
    pub fn from_radials(radials: Vec<Radial>) -> Vec<Self> {
        let mut sweeps = Vec::new();

        let mut sweep_elevation_number = None;

        let mut sweep_radials = Vec::new();

        for radial in radials {
            if let Some(elevation_number) = sweep_elevation_number {
                if elevation_number != radial.elevation_number {
                    sweeps.push(Sweep::new(elevation_number, sweep_radials));

                    sweep_radials = Vec::new();
                }
            }

            sweep_elevation_number = Some(radial.elevation_number);

            sweep_radials.push(radial);
        }

        sweeps
    }
}