sdr_heatmap/
lib.rs

1use csv::StringRecord;
2use flate2::read::GzDecoder;
3use image::png::PNGEncoder;
4use log::*;
5use rayon::prelude::*;
6use std::f32;
7use std::io::prelude::*;
8use std::path::Path;
9use std::{cmp::Ordering, fs::File};
10pub mod colors;
11use colors::scale_tocolor;
12use colors::Palettes;
13
14#[derive(Debug)]
15struct Measurement {
16    date: String,
17    time: String,
18    freq_low: u32,
19    freq_high: u32,
20    freq_step: f64,
21    samples: u32,
22    values: Vec<f32>,
23}
24
25impl Measurement {
26    fn get_values(&self) -> Vec<(f64, f32)> {
27        self.values
28            .iter()
29            .zip(0..)
30            .map(|(value, i)| {
31                (
32                    ((i) as f64) * self.freq_step + (self.freq_low as f64),
33                    *value,
34                )
35            })
36            .collect()
37    }
38    fn new(record: StringRecord) -> Measurement {
39        let mut values: Vec<_> = record.iter().skip(6).map(|s| s.parse().unwrap()).collect();
40        values.truncate(record.len() - 7);
41        Measurement {
42            date: record.get(0).unwrap().to_string(),
43            time: record.get(1).unwrap().to_string(),
44            freq_low: parse(record.get(2).unwrap()).unwrap(),
45            freq_high: parse(record.get(3).unwrap()).unwrap(),
46            freq_step: parse(record.get(4).unwrap()).unwrap(),
47            samples: parse(record.get(5).unwrap()).unwrap(),
48            values,
49        }
50    }
51}
52
53#[derive(PartialEq, Debug)]
54pub struct Summary {
55    pub min: f32,
56    pub max: f32,
57}
58
59impl Summary {
60    fn empty() -> Self {
61        Self {
62            min: f32::INFINITY,
63            max: f32::NEG_INFINITY,
64        }
65    }
66
67    fn combine_f32(s: Self, i: f32) -> Self {
68        Summary::combine(s, Summary { min: i, max: i })
69    }
70
71    fn combine(a: Self, b: Self) -> Self {
72        Self {
73            min: {
74                let a = a.min;
75                let b = b.min;
76                if a.is_finite() {
77                    a.min(b)
78                } else {
79                    b
80                }
81            },
82            max: {
83                let a = a.max;
84                let b = b.max;
85                if a.is_finite() {
86                    a.max(b)
87                } else {
88                    b
89                }
90            },
91        }
92    }
93}
94
95fn parse<T: std::str::FromStr>(
96    string: &str,
97) -> std::result::Result<T, <T as std::str::FromStr>::Err> {
98    let parsed = string.parse::<T>();
99    debug_assert!(parsed.is_ok(), "Could not parse {}", string);
100    parsed
101}
102
103pub fn open_file(path: &Path) -> Box<dyn std::io::Read> {
104    let file = File::open(path).unwrap();
105    if path.extension().unwrap() == "gz" {
106        Box::new(GzDecoder::new(file))
107    } else {
108        Box::new(file)
109    }
110}
111
112fn read_file<T: std::io::Read>(file: T) -> csv::Reader<T> {
113    csv::ReaderBuilder::new()
114        .has_headers(false)
115        .from_reader(file)
116}
117
118pub fn main(path: &Path) {
119    info!("Loading: {}", path.display());
120    //Preprocess
121    let file = open_file(path);
122    let summary = preprocess_iter(file);
123    info!("Color values {} to {}", summary.min, summary.max);
124    //Process
125    let file = open_file(path);
126    let reader = read_file(file);
127    let (datawidth, dataheight, img) = process(reader, summary.min, summary.max);
128    //Draw
129    let (height, imgdata) = create_image(datawidth, dataheight, img);
130    let dest = path.with_extension("png");
131    save_image(datawidth, height, imgdata, dest.to_str().unwrap()).unwrap();
132}
133
134pub fn preprocess(file: Box<dyn Read>) -> Summary {
135    let reader = read_file(file);
136    let mut min = f32::INFINITY;
137    let mut max = f32::NEG_INFINITY;
138    for result in reader.into_records() {
139        let record = {
140            let mut x = result.unwrap();
141            x.trim();
142            x
143        };
144        let values: Vec<f32> = record
145            .iter()
146            .skip(6)
147            .map(|s| {
148                if s == "-nan" {
149                    f32::NAN
150                } else {
151                    s.parse::<f32>()
152                        .unwrap_or_else(|e| panic!("{} should be a valid float: {:?}", s, e))
153                }
154            })
155            .collect();
156        for value in values {
157            if value != f32::INFINITY && value != f32::NEG_INFINITY {
158                if value > max {
159                    max = value
160                }
161                if value < min {
162                    min = value
163                }
164            }
165        }
166    }
167    Summary { min, max }
168}
169
170pub fn preprocess_iter(file: Box<dyn Read>) -> Summary {
171    read_file(file)
172        .into_records()
173        .map(|x| x.unwrap())
174        .flat_map(|line| {
175            line.into_iter()
176                .skip(6)
177                .map(|s| {
178                    if s == "-nan" {
179                        f32::NAN
180                    } else {
181                        s.trim().parse::<f32>().unwrap_or_else(|e| {
182                            panic!("'{}' should be a valid float: '{:?}'", s, e)
183                        })
184                    }
185                })
186                .collect::<Vec<f32>>()
187        })
188        .fold(Summary::empty(), Summary::combine_f32)
189}
190
191pub fn preprocess_par_iter(file: Box<dyn Read>) -> Summary {
192    read_file(file)
193        .into_records()
194        .collect::<Vec<_>>()
195        .into_iter()
196        .par_bridge()
197        .map(|x| x.unwrap())
198        .flat_map(|line| {
199            line.into_iter()
200                .skip(6)
201                .map(|s| s.to_owned())
202                .collect::<Vec<String>>()
203        })
204        .map(|s| {
205            if s == "-nan" {
206                f32::NAN
207            } else {
208                s.trim()
209                    .parse::<f32>()
210                    .unwrap_or_else(|e| panic!("'{}' should be a valid float: '{:?}'", s, e))
211            }
212        })
213        .fold(Summary::empty, Summary::combine_f32)
214        .reduce(Summary::empty, Summary::combine)
215}
216
217fn process(
218    reader: csv::Reader<Box<dyn std::io::Read>>,
219    min: f32,
220    max: f32,
221) -> (usize, usize, std::vec::Vec<u8>) {
222    let mut date: String = "".to_string();
223    let mut time: String = "".to_string();
224    let mut batch = 0;
225    let mut datawidth = 0;
226    let mut img = Vec::new();
227    for result in reader.into_records() {
228        let mut record = result.unwrap();
229        record.trim();
230        assert!(record.len() > 7);
231        let m = Measurement::new(record);
232        let vals = m.get_values();
233        if date == m.date && time == m.time {
234        } else {
235            if datawidth == 0 {
236                datawidth = batch;
237            }
238            debug_assert_eq! {datawidth,batch}
239            batch = 0;
240            date = m.date;
241            time = m.time;
242        }
243        for (_, v) in vals {
244            let pixel = scale_tocolor(Palettes::Default, v, min, max);
245            img.extend(pixel.iter());
246            batch += 1;
247        }
248    }
249    if datawidth == 0 {
250        datawidth = batch;
251    }
252    info!("Img data {}x{}", datawidth, batch);
253    (datawidth, img.len() / 3 / datawidth, img)
254}
255
256fn tape_measure(width: usize, imgdata: &mut Vec<u8>) {
257    let length = width * 26 * 3;
258    imgdata.append(&mut vec![0; length]);
259}
260
261fn create_image(width: usize, height: usize, mut img: Vec<u8>) -> (usize, std::vec::Vec<u8>) {
262    info!("Raw {}x{}", width, height);
263    let mut imgdata: Vec<u8> = Vec::new();
264    tape_measure(width, &mut imgdata);
265    imgdata.append(&mut img);
266    let height = height + 26;
267    let expected_length = width * height * 3;
268    match expected_length.cmp(&imgdata.len()) {
269        Ordering::Greater => {
270            warn!("Image is missing some values, was the file cut early? Filling black.",);
271            imgdata.append(&mut vec![0; expected_length - imgdata.len()]);
272        }
273        Ordering::Less => {
274            warn!("Image has too many values, was the file cut early? Trimming.",);
275            imgdata.truncate(expected_length);
276        }
277        Ordering::Equal => {}
278    }
279    (height, imgdata)
280}
281
282fn save_image(
283    width: usize,
284    height: usize,
285    imgdata: Vec<u8>,
286    dest: &str,
287) -> std::result::Result<(), image::error::ImageError> {
288    info!("Saving {} {}x{}", dest, width, height);
289    let f = std::fs::File::create(dest).unwrap();
290    PNGEncoder::new(f).encode(
291        &imgdata,
292        width as u32,
293        height as u32,
294        image::ColorType::Rgb8,
295    )
296}
297
298#[cfg(test)]
299mod tests {
300    use crate::*;
301    use pretty_assertions::{assert_eq, assert_ne};
302
303    #[test]
304    fn preprocess_result() {
305        let res = preprocess(open_file(Path::new("samples/sample1.csv.gz")));
306        assert_eq!(
307            res,
308            Summary {
309                min: -29.4,
310                max: 21.35
311            }
312        );
313    }
314
315    #[test]
316    fn preprocess_iter_result() {
317        let res = preprocess_iter(open_file(Path::new("samples/sample1.csv.gz")));
318        assert_eq!(
319            res,
320            Summary {
321                min: -29.4,
322                max: 21.35
323            }
324        );
325    }
326
327    #[test]
328    fn preprocess_par_iter_result() {
329        let res = preprocess_par_iter(open_file(Path::new("samples/sample1.csv.gz")));
330        assert_eq!(
331            res,
332            Summary {
333                min: -29.4,
334                max: 21.35
335            }
336        );
337    }
338
339    #[test]
340    fn complete() {
341        main(Path::new("samples/sample1.csv.gz"))
342    }
343}
sdr_heatmap/lib.rs

sdr_heatmap/
lib.rs
