diff --git a/hexit/src/combine.rs b/hexit/src/combine.rs index aacfeca..6e4846e 100644 --- a/hexit/src/combine.rs +++ b/hexit/src/combine.rs @@ -1,9 +1,7 @@ -use crate::{mask, options::Combine, progress}; +use crate::{options::Combine, progress}; use anyhow::Result; use byteorder::{LittleEndian as LE, ReadBytesExt, WriteBytesExt}; use flate2::bufread::GzDecoder; -use geo::{coord, GeometryCollection, Intersects}; -use h3o::{CellIndex, LatLng}; use hextree::{compaction::EqCompactor, disktree::DiskTree, Cell, HexTreeMap}; use indicatif::MultiProgress; use std::{ffi::OsStr, fs::File, io::BufReader, path::Path}; @@ -13,9 +11,8 @@ impl Combine { assert!(!self.input.is_empty()); let mut hextree: HexTreeMap = HexTreeMap::with_compactor(EqCompactor); let progress_group = MultiProgress::new(); - let mask = mask::open(self.mask.as_deref())?; for tess_file_path in &self.input { - Self::read_tessellation(tess_file_path, &progress_group, mask.as_ref(), &mut hextree)?; + Self::read_tessellation(tess_file_path, &progress_group, &mut hextree)?; } self.write_disktree(&hextree, &progress_group)?; self.verify_disktree(&hextree, &progress_group)?; @@ -25,7 +22,6 @@ impl Combine { fn read_tessellation( tess_file_path: &Path, progress_group: &MultiProgress, - mask: Option<&GeometryCollection>, hextree: &mut HexTreeMap, ) -> Result<()> { let tess_file = File::open(tess_file_path)?; @@ -39,21 +35,10 @@ impl Combine { let n_samples = rdr.read_u64::()?; let pb = progress_group.add(progress::bar(tess_file_name.to_string(), n_samples)); for _sample_n in 0..n_samples { + let raw_cell = rdr.read_u64::()?; + let cell = hextree::Cell::from_raw(raw_cell)?; let elevation = rdr.read_i16::()?; - let n_cells = rdr.read_u16::()?; - for _cell_n in 0..n_cells { - let raw_cell = rdr.read_u64::()?; - let cell = hextree::Cell::from_raw(raw_cell)?; - let mask_contains_cell = mask.as_ref().map_or(true, |mask| { - let cell_index = CellIndex::try_from(cell.into_raw()).unwrap(); - let cell_center = LatLng::from(cell_index); - let coord = coord!(x: cell_center.lng(), y: cell_center.lat()); - mask.intersects(&coord) - }); - if mask_contains_cell { - hextree.insert(cell, elevation); - } - } + hextree.insert(cell, elevation); pb.inc(1); } assert!( @@ -77,7 +62,7 @@ impl Combine { .expect("already opened, therefore path must be a file"); let disktree_len = hextree.len(); let pb = progress_group.add(progress::bar( - disktree_file_name.to_string(), + format!("Writing {disktree_file_name}"), disktree_len as u64, )); hextree.to_disktree(disktree_file, |wtr, val| { @@ -98,8 +83,13 @@ impl Combine { } let mut disktree = DiskTree::open(&self.out)?; + let disktree_file_name = self + .out + .file_name() + .and_then(OsStr::to_str) + .expect("already opened, therefore path must be a file"); let pb = progress_group.add(progress::bar( - "Validating disktree".to_string(), + format!("Validating {disktree_file_name}"), hextree.len() as u64, )); let mut count = 0; diff --git a/hexit/src/lookup.rs b/hexit/src/lookup.rs index 91f6a73..3cb3f72 100644 --- a/hexit/src/lookup.rs +++ b/hexit/src/lookup.rs @@ -2,6 +2,7 @@ use crate::options::Lookup; use anyhow::Result; use byteorder::{LittleEndian as LE, ReadBytesExt}; use hextree::{disktree::DiskTree, Cell}; +use std::{fs::File, io::Write}; impl Lookup { pub fn run(&self) -> Result<()> { @@ -11,17 +12,45 @@ impl Lookup { .or_else(|_| u64::from_str_radix(&self.cell, 16))?; let cell = Cell::try_from(raw_cell)?; let mut disktree = DiskTree::open(&self.disktree)?; + + if self.iter { + Self::by_iter(cell, &mut disktree) + } else { + Self::by_get(cell, &mut disktree) + } + } + + fn by_get(cell: Cell, disktree: &mut DiskTree) -> Result<()> { let t0 = std::time::Instant::now(); match disktree.seek_to_cell(cell)? { None => (), Some((cell, rdr)) => { let t_seek = t0.elapsed(); - let val = rdr.read_i16::()?; + let elev = rdr.read_i16::()?; let t_tot = t0.elapsed(); - println!("{cell}: {val}"); + println!("{cell}: {elev}"); println!("{t_seek:?} {t_tot:?}"); } } Ok(()) } + + fn by_iter(_target_cell: Cell, disktree: &mut DiskTree) -> Result<()> { + fn read_elev(res: hextree::Result<(Cell, &mut File)>) -> Result> { + let (cell, rdr) = res?; + let mask = Cell::try_from(0x8126bffffffffff)?; + if cell.is_related_to(&mask) { + Ok(Some((cell, rdr.read_i16::()?))) + } else { + Ok(None) + } + } + let mut stderr = std::io::stderr().lock(); + for res in disktree.iter(read_elev)? { + if let Some((cell, elev)) = res? { + writeln!(&mut stderr, "{cell}: {elev}")?; + } + } + Ok(()) + } } diff --git a/hexit/src/options.rs b/hexit/src/options.rs index d87ec9a..4c94352 100644 --- a/hexit/src/options.rs +++ b/hexit/src/options.rs @@ -65,6 +65,9 @@ pub struct Combine { /// Lookup value for H3 cell in a disktree. #[derive(Debug, Clone, Args)] pub struct Lookup { + /// Iterate through the disktree instead of `get`ting the value. + #[arg(short, long)] + pub iter: bool, pub disktree: PathBuf, pub cell: String, } diff --git a/hexit/src/tesselate.rs b/hexit/src/tesselate.rs index 10496f1..037c897 100644 --- a/hexit/src/tesselate.rs +++ b/hexit/src/tesselate.rs @@ -5,8 +5,9 @@ use flate2::{write::GzEncoder, Compression}; use geo::{GeometryCollection, Intersects}; use h3o::{ geom::{PolyfillConfig, Polygon, ToCells}, - CellIndex, Resolution, + Resolution, }; +use hextree::{compaction::EqCompactor, Cell, HexTreeMap}; use indicatif::{MultiProgress, ProgressBar}; use nasadem::{Sample, Tile}; use rayon::prelude::*; @@ -33,17 +34,19 @@ impl Tesselate { mask: Option<&GeometryCollection>, progress_group: &MultiProgress, ) -> Result<()> { - let out_file_name = { + let (in_file_name, out_file_name) = { let file_name = height_file_path .file_name() .and_then(OsStr::to_str) .expect("already opened, therefore path must be a file"); - format!("{file_name}.res{}.h3tez", self.resolution) + ( + file_name, + format!("{file_name}.res{}.h3tez", self.resolution), + ) }; let out_file_path = self.out_dir.clone().join(&out_file_name); if out_file_path.exists() && !self.overwrite { - // Exit early if we've already processed this input. return Ok(()); } @@ -54,46 +57,69 @@ impl Tesselate { }; let tile = Tile::memmap(height_file_path)?; - let intersects = mask - .as_ref() - .map_or(true, |mask| mask.intersects(&tile.polygon())); + let intersects = mask.as_ref().map_or(true, |mask| { + let polygon = tile.polygon(); + mask.intersects(&polygon) + }); if intersects { - let pb = progress_group.add(progress::bar(out_file_name, tile.len() as u64)); + let pb = progress_group.add(progress::bar( + format!("Tesselate {in_file_name}"), + tile.len() as u64, + )); + let hextree = self.tesselate_tile(&tile, &pb)?; let tmp_out_file = File::create(&out_file_tmp_path)?; let tmp_out_wtr = GzEncoder::new(tmp_out_file, Compression::new(self.compression)); - self.polyfill_tile(&tile, &pb, BufWriter::new(tmp_out_wtr))?; + let wtr = BufWriter::new(tmp_out_wtr); + let pb = progress_group.add(progress::bar( + format!("Write {out_file_name}"), + tile.len() as u64, + )); + Self::write_to_disk(&hextree, &pb, wtr)?; fs::rename(out_file_tmp_path, out_file_path)?; } Ok(()) } - fn polyfill_tile( + fn tesselate_tile( &self, tile: &Tile, progress_bar: &ProgressBar, - mut out: impl Write, - ) -> Result<()> { - out.write_u64::(tile.len() as u64)?; + ) -> Result> { + let mut hextree: HexTreeMap = HexTreeMap::with_compactor(EqCompactor); for sample in tile.iter() { - let (elev, hexes) = polyfill_sample(&sample, self.resolution)?; - out.write_i16::(elev)?; - out.write_u16::(u16::try_from(hexes.len())?)?; + let (elev, hexes) = Self::tesselate_sample(&sample, self.resolution)?; for hex in hexes { - out.write_u64::(hex)?; + hextree.insert(Cell::try_from(hex)?, elev); } progress_bar.inc(1); } - Ok(()) + Ok(hextree) } -} -fn polyfill_sample(sample: &Sample, resolution: Resolution) -> Result<(i16, Vec)> { - let elevation = sample.elevation(); - let polygon = Polygon::from_degrees(sample.polygon())?; - let cell_iter = polygon.to_cells(PolyfillConfig::new(resolution)); - let mut cells: Vec = CellIndex::compact(cell_iter)?.map(u64::from).collect(); - cells.sort_unstable(); - cells.dedup(); - Ok((elevation, cells)) + fn tesselate_sample(sample: &Sample, resolution: Resolution) -> Result<(i16, Vec)> { + let elevation = sample.elevation(); + let polygon = Polygon::from_degrees(sample.polygon())?; + let mut cells: Vec = polygon + .to_cells(PolyfillConfig::new(resolution)) + .map(u64::from) + .collect(); + cells.sort_unstable(); + cells.dedup(); + Ok((elevation, cells)) + } + + fn write_to_disk( + hextree: &HexTreeMap, + progress_bar: &ProgressBar, + mut out: impl Write, + ) -> Result<()> { + out.write_u64::(hextree.len() as u64)?; + for (cell, elev) in hextree.iter() { + out.write_u64::(cell.into_raw())?; + out.write_i16::(*elev)?; + progress_bar.inc(1); + } + Ok(()) + } }