copc_rs/
writer.rs

1//! COPC file writer.
2
3use crate::compressor::CopcCompressor;
4use crate::copc::{CopcInfo, Entry, HierarchyPage, OctreeNode, VoxelKey};
5
6use las::{Builder, Header};
7
8use std::collections::HashMap;
9use std::fs::File;
10use std::io::{BufWriter, Cursor, Seek, SeekFrom, Write};
11use std::path::Path;
12
13// enum for point data record format upgrades
14enum UpgradePdrf {
15    From1to6,  // upgrades (1=>6)
16    From3to7,  // upgrades (3=>7)
17    NoUpgrade, // 6, 7 and 8
18}
19
20impl UpgradePdrf {
21    fn log_string(&self) -> &str {
22        match self {
23            UpgradePdrf::From1to6 => "Upgrading LAS PDRF from 1 to 6",
24            UpgradePdrf::From3to7 => "Upgrading LAS PDRF from 3 to 7",
25            UpgradePdrf::NoUpgrade => "COPC supports the given PDRF",
26        }
27    }
28}
29
30/// COPC file writer
31pub struct CopcWriter<'a, W: 'a + Write + Seek> {
32    is_closed: bool,
33    start: u64,
34    // point writer
35    compressor: CopcCompressor<'a, W>,
36    header: Header,
37    // a page of the written entries
38    hierarchy: HierarchyPage,
39    min_node_size: i32,
40    max_node_size: i32,
41    copc_info: CopcInfo,
42    // root node in octree, access point for the tree
43    root_node: OctreeNode,
44    // a hashmap to store chunks that are not full yet
45    open_chunks: HashMap<VoxelKey, Cursor<Vec<u8>>>,
46}
47
48impl CopcWriter<'_, BufWriter<File>> {
49    /// Creates a new COPC-writer for a path,
50    /// creates a file at that path and wraps it in a BufWrite for you
51    /// and passes it along to [new]
52    ///
53    /// see [new] for usage
54    ///
55    /// [new]: Self::new
56    pub fn from_path<P: AsRef<Path>>(
57        path: P,
58        header: Header,
59        min_size: i32,
60        max_size: i32,
61    ) -> crate::Result<Self> {
62        let copc_ext = Path::new(match path.as_ref().file_stem() {
63            Some(copc) => copc,
64            None => return Err(crate::Error::WrongCopcExtension),
65        })
66        .extension();
67
68        match (copc_ext, path.as_ref().extension()) {
69            (Some(copc), Some(laz)) => match (&copc.to_str(), &laz.to_str()) {
70                (Some(copc_str), Some(laz_str)) => {
71                    if &copc_str.to_lowercase() != "copc" || &laz_str.to_lowercase() != "laz" {
72                        return Err(crate::Error::WrongCopcExtension);
73                    }
74                }
75                _ => return Err(crate::Error::WrongCopcExtension),
76            },
77            _ => return Err(crate::Error::WrongCopcExtension),
78        }
79
80        File::create(path)
81            .map_err(crate::Error::from)
82            .and_then(|file| CopcWriter::new(BufWriter::new(file), header, min_size, max_size))
83    }
84}
85
86/// public API
87impl<W: Write + Seek> CopcWriter<'_, W> {
88    /// Create a COPC file writer for the write- and seekable `write`
89    /// configured with the provided [las::Header]
90    /// recommended to use [from_path] for writing to file
91    ///
92    /// The `bounds` field in the `header` is used as the bounds for the octree
93    /// the bounds are checked for being normal
94    ///
95    /// `max_size` is the maximal number of [las::Point]s an octree node can hold
96    /// any max_size < 1 sets the max_size to [crate::MAX_NODE_SIZE_DEFAULT]
97    /// this is a soft limit
98    ///
99    /// `min_size` is the minimal number of [las::Point]s an octree node can hold
100    /// any min_size < 1 sets the min_size to [crate::MIN_NODE_SIZE_DEFAULT]
101    /// this is a hard limit
102    ///
103    /// `min_size` greater or equal to `max_size` after checking values < 1
104    /// results in a [crate::Error::InvalidNodeSize] error
105    ///
106    ///
107    /// This writer is strictly following the LAS 1.4 spec and the COPC spec
108    /// which means that any provided header not compatible with those will lead
109    /// to an Err
110    /// That being said, LAS 1.2 headers and PDRFs 1 and 3 are accepted and upgraded to
111    /// their matching LAS 1.4 versions
112    /// GeoTiff CRS VLR's are parsed and written to WKT CRS VLR's
113    /// A CRS VLR is __MANDATORY__ and without one
114    ///
115    /// [from_path]: Self::from_path
116    pub fn new(mut write: W, header: Header, min_size: i32, max_size: i32) -> crate::Result<Self> {
117        let start = write.stream_position()?;
118
119        let min_node_size = if min_size < 1 {
120            crate::MIN_NODE_SIZE_DEFAULT
121        } else {
122            min_size
123        };
124
125        let max_node_size = if max_size < 1 {
126            crate::MAX_NODE_SIZE_DEFAULT
127        } else {
128            max_size
129        };
130
131        if min_node_size >= max_node_size {
132            return Err(crate::Error::InvalidNodeSize);
133        }
134
135        if header.version() != las::Version::new(1, 4) {
136            log::log!(log::Level::Info, "Old Las version. Upgrading");
137        }
138
139        let mut has_wkt_vlr = false;
140
141        // store the vlrs contained in the header for forwarding
142        let mut forward_vlrs = Vec::with_capacity(header.vlrs().len());
143        for vlr in header.vlrs() {
144            match (vlr.user_id.to_lowercase().as_str(), vlr.record_id) {
145                ("lasf_projection", 2112) => {
146                    has_wkt_vlr = true;
147                    forward_vlrs.push(vlr.clone());
148                }
149                // not forwarding these vlrs
150                ("lasf_projection", 34735..=34737) => (), // geo-tiff crs
151                ("copc", 1 | 1000) => (),
152                ("laszip encoded", 22204) => (),
153                ("lasf_spec", 100..355 | 65535) => (), // wave form packet descriptors
154                // forwarding all other vlrs
155                _ => forward_vlrs.push(vlr.clone()),
156            }
157        }
158
159        // store the evlrs contained in the header for forwarding
160        let mut forward_evlrs = Vec::with_capacity(header.evlrs().len());
161        for evlr in header.evlrs() {
162            match (evlr.user_id.to_lowercase().as_str(), evlr.record_id) {
163                ("lasf_projection", 2112) => {
164                    has_wkt_vlr = true;
165                    forward_evlrs.push(evlr.clone());
166                }
167                // not forwarding these vlrs
168                ("lasf_projection", 34735..=34737) => (), // geo-tiff crs
169                ("copc", 1 | 1000) => (),                 // 1 should never be a evlr
170                ("laszip encoded", 22204) => (),          // should never be a evlr
171                ("lasf_spec", 100..355 | 65535) => (),    // waveform data packets
172                // forwarding all other evlrs
173                _ => forward_evlrs.push(evlr.clone()),
174            }
175        }
176
177        // las version 1.4 says pdrf 6-10 must have a wkt crs
178        // copc is only valid for las 1.4 and says only pdrf 6-8 is supported
179        //
180        // which means that any geotiff crs must be converted to a wkt crs
181        //
182        // could just use header.has_wkt_vlr(), but so many las files are wrongly written
183        // so I don't trust it
184        //
185        // ignores any vertical crs that might stored in geotiff
186        if !has_wkt_vlr {
187            let epsg = las_crs::parse_las_crs(&header)?;
188            let wkt_data = match crs_definitions::from_code(epsg.horizontal) {
189                Some(wkt) => wkt,
190                None => return Err(crate::Error::InvalidEPSGCode(epsg.horizontal)),
191            }
192            .wkt
193            .as_bytes()
194            .to_owned();
195
196            let mut user_id = [0; 16];
197            for (i, c) in "LASF_Projection".as_bytes().iter().enumerate() {
198                user_id[i] = *c;
199            }
200
201            let crs_vlr = las::raw::Vlr {
202                reserved: 0,
203                user_id,
204                record_id: 2112,
205                record_length_after_header: las::raw::vlr::RecordLength::Vlr(wkt_data.len() as u16),
206                description: [0; 32],
207                data: wkt_data,
208            };
209
210            forward_vlrs.push(las::Vlr::new(crs_vlr));
211        }
212
213        // check bounds are normal
214        let bounds = header.bounds();
215        if !(bounds.max.x - bounds.min.x).is_normal()
216            || !(bounds.max.y - bounds.min.y).is_normal()
217            || !(bounds.max.z - bounds.min.z).is_normal()
218        {
219            return Err(crate::Error::InvalidBounds(bounds));
220        }
221
222        let mut raw_head = header.into_raw()?;
223
224        // mask off the two leftmost bits corresponding to compression of pdrf
225        let pdrf = raw_head.point_data_record_format & 0b00111111;
226        let upgrade_pdrf = match pdrf {
227            1 => {
228                let upgrade = UpgradePdrf::From1to6;
229
230                log::log!(log::Level::Info, "{}", upgrade.log_string());
231                upgrade
232            }
233            3 => {
234                let upgrade = UpgradePdrf::From3to7;
235
236                log::log!(log::Level::Info, "{}", upgrade.log_string());
237                upgrade
238            }
239            0 | 2 => {
240                return Err(las::Error::InvalidPointFormat(las::point::Format::new(
241                    raw_head.point_data_record_format,
242                )?))?;
243            }
244            4..=5 | 9.. => {
245                return Err(las::Error::InvalidPointFormat(las::point::Format::new(
246                    raw_head.point_data_record_format,
247                )?))?;
248            }
249            6..=8 => UpgradePdrf::NoUpgrade,
250        };
251
252        // adjust and clear some fields
253        raw_head.version = las::Version::new(1, 4);
254        raw_head.point_data_record_format += match upgrade_pdrf {
255            UpgradePdrf::NoUpgrade => 0,
256            UpgradePdrf::From1to6 => 5,
257            UpgradePdrf::From3to7 => 4,
258        };
259        raw_head.point_data_record_format |= 0b11000000; // make sure the compress bits are set
260        raw_head.point_data_record_length += match upgrade_pdrf {
261            UpgradePdrf::NoUpgrade => 0,
262            _ => 2,
263        };
264        raw_head.global_encoding |= 0b10000; // make sure wkt crs bit is set
265        raw_head.number_of_point_records = 0;
266        raw_head.number_of_points_by_return = [0; 5];
267        raw_head.large_file = None;
268        raw_head.evlr = None;
269        raw_head.padding = vec![];
270
271        let mut software_buffer = [0_u8; 32];
272        for (i, byte) in format!("COPC-rs v{}", crate::VERSION).bytes().enumerate() {
273            software_buffer[i] = byte;
274        }
275        raw_head.generating_software = software_buffer;
276
277        // start building a real header from the raw header
278        let mut builder = Builder::new(raw_head)?;
279        // add a blank COPC-vlr as the first vlr
280        builder.vlrs.push(CopcInfo::default().into_vlr()?);
281
282        // create the laz vlr
283        let point_format = builder.point_format;
284        let mut laz_items = laz::laszip::LazItemRecordBuilder::new();
285        laz_items.add_item(laz::LazItemType::Point14);
286        if point_format.has_color {
287            if point_format.has_nir {
288                laz_items.add_item(laz::LazItemType::RGBNIR14);
289            } else {
290                laz_items.add_item(laz::LazItemType::RGB14);
291            }
292        }
293        if point_format.extra_bytes > 0 {
294            laz_items.add_item(laz::LazItemType::Byte14(point_format.extra_bytes));
295        }
296
297        let laz_vlr = laz::LazVlrBuilder::new(laz_items.build())
298            .with_variable_chunk_size()
299            .build();
300        let mut cursor = Cursor::new(Vec::<u8>::new());
301        laz_vlr.write_to(&mut cursor)?;
302        let laz_vlr = las::Vlr {
303            user_id: laz::LazVlr::USER_ID.to_owned(),
304            record_id: laz::LazVlr::RECORD_ID,
305            description: laz::LazVlr::DESCRIPTION.to_owned(),
306            data: cursor.into_inner(),
307        };
308        builder.vlrs.push(laz_vlr);
309
310        // add the forwarded vlrs
311        builder.vlrs.extend(forward_vlrs);
312        builder.evlrs.extend(forward_evlrs);
313        // the EPT-hierarchy evlr is not yet added
314
315        let header = builder.into_header()?;
316
317        // write the header and vlrs
318        // this is just to reserve the space
319        header.write_to(&mut write)?;
320
321        let center_point = las::Vector {
322            x: (bounds.min.x + bounds.max.x) / 2.,
323            y: (bounds.min.y + bounds.max.y) / 2.,
324            z: (bounds.min.z + bounds.max.z) / 2.,
325        };
326        let halfsize = (center_point.x - bounds.min.x)
327            .max((center_point.y - bounds.min.y).max(center_point.z - bounds.min.z));
328
329        let mut root_node = OctreeNode::new();
330
331        root_node.bounds = las::Bounds {
332            min: las::Vector {
333                x: center_point.x - halfsize,
334                y: center_point.y - halfsize,
335                z: center_point.z - halfsize,
336            },
337            max: las::Vector {
338                x: center_point.x + halfsize,
339                y: center_point.y + halfsize,
340                z: center_point.z + halfsize,
341            },
342        };
343        root_node.entry.key.level = 0;
344        root_node.entry.offset = write.stream_position()?;
345
346        let copc_info = CopcInfo {
347            center: center_point,
348            halfsize,
349            spacing: 0.,
350            root_hier_offset: 0,
351            root_hier_size: 0,
352            gpstime_minimum: f64::MAX,
353            gpstime_maximum: f64::MIN,
354        };
355
356        Ok(CopcWriter {
357            is_closed: false,
358            start,
359            compressor: CopcCompressor::new(write, header.laz_vlr()?)?,
360            header,
361            hierarchy: HierarchyPage { entries: vec![] },
362            min_node_size,
363            max_node_size,
364            copc_info,
365            root_node,
366            open_chunks: HashMap::default(),
367        })
368    }
369
370    /// Write anything that implements [IntoIterator]
371    /// over [las::Point] to the COPC [Write]
372    /// Only one iterator can be written so a call to [Self::write] closes the writer.
373    ///
374    /// `num_points` is the number of points in the iterator
375    /// the number of points is used for stochastically filling the nodes
376    /// if `num_points` is < 1 a greedy filling strategy is used
377    /// this should only be used if the passed iterator is randomly ordered
378    /// which most of the time not is the case
379    /// if `num_points` is not equal to the actual number of points in the
380    /// iterator all points will still be written but the point distribution
381    /// in a node will not represent of the entire distribution over that node
382    /// i.e. only full resolution queries will look right which means the point cloud
383    /// will look wierd in any viewer which utilizes the COPC information
384    ///
385    /// returns an `Err`([crate::Error::ClosedWriter]) if the writer has already been closed.
386    ///
387    /// If a point is outside the copc `bounds` or not matching the
388    /// [las::point::Format] of the writer's header `Err` is returned
389    /// [crate::PointAddError::PointAttributesDoNotMatch] take precedence over
390    /// [crate::PointAddError::PointNotInBounds]
391    /// All the points inside the bounds and matching the point format are written regardless
392    ///
393    /// All points which both match the point format and are inside the bounds are added
394    ///
395    /// Lastly [Self::close] is called. If closing fails an [crate::Error] is returned and
396    /// the state of the [Write] is undefined
397    ///
398    /// If all points match the format, are inside the bounds and [Self::close] is successfull `Ok(())` is returned
399    pub fn write<D: IntoIterator<Item = las::Point>>(
400        &mut self,
401        data: D,
402        num_points: i32,
403    ) -> crate::Result<()> {
404        if self.is_closed {
405            return Err(crate::Error::ClosedWriter);
406        }
407
408        let result = if num_points < self.max_node_size + self.min_node_size {
409            // greedy filling strategy
410            self.write_greedy(data)
411        } else {
412            // stochastic filling strategy
413            self.write_stochastic(data, num_points as usize)
414        };
415
416        self.close()?;
417        result
418    }
419
420    /// Whether this writer is closed or not
421    pub fn is_closed(&self) -> bool {
422        self.is_closed
423    }
424
425    /// number of points in the largest node
426    pub fn max_node_size(&self) -> i32 {
427        self.max_node_size
428    }
429
430    /// number of points in the smallest node
431    pub fn min_node_size(&self) -> i32 {
432        self.min_node_size
433    }
434
435    /// This writer's header, some fields are updated on closing of the writer
436    pub fn header(&self) -> &Header {
437        &self.header
438    }
439
440    /// This writer's EPT Hierarchy
441    pub fn hierarchy_entries(&self) -> &HierarchyPage {
442        &self.hierarchy
443    }
444
445    /// This writer's COPC info
446    pub fn copc_info(&self) -> &CopcInfo {
447        &self.copc_info
448    }
449}
450
451/// private functions
452impl<W: Write + Seek> CopcWriter<'_, W> {
453    /// Greedy strategy for writing points
454    fn write_greedy<D: IntoIterator<Item = las::Point>>(&mut self, data: D) -> crate::Result<()> {
455        let mut invalid_points = Ok(());
456
457        for p in data.into_iter() {
458            if !p.matches(self.header.point_format()) {
459                invalid_points = Err(crate::Error::InvalidPoint(
460                    crate::PointAddError::PointAttributesDoNotMatch(*self.header.point_format()),
461                ));
462                continue;
463            }
464            if !bounds_contains_point(&self.root_node.bounds, &p) {
465                if invalid_points.is_ok() {
466                    invalid_points = Err(crate::Error::InvalidPoint(
467                        crate::PointAddError::PointNotInBounds,
468                    ));
469                }
470                continue;
471            }
472
473            self.add_point_greedy(p)?;
474        }
475        invalid_points
476    }
477
478    /// Stochastic strategy for writing points
479    fn write_stochastic<D: IntoIterator<Item = las::Point>>(
480        &mut self,
481        data: D,
482        num_points: usize,
483    ) -> crate::Result<()> {
484        let mut invalid_points = Ok(());
485
486        // the number of expected levels in the copc hierarchy
487        // assuming that the lidar scans cover a way bigger horizontal span than vertical
488        // effectivly dividing every level into 4 instead of 8
489        // (removing this assumption would lead to a division by 3 instead of 2, and thus fewer expected levels)
490        //
491        // each level then holds 4^i * max_points_per_node points
492        //
493        // solve for l:
494        // num_points / sum_i=0^l 4^i = max_points_per_node
495        //
496        // sum_i=0^l 4^i = 1/3 (4^(l+1) - 1)
497        // =>
498        // l = (log_2( 3*num_points/max_points_per_node + 1) - 2)/2
499        let expected_levels =
500            ((((3 * num_points) as f64 / self.max_node_size as f64 + 1.).log2() - 2.) / 2.).ceil()
501                as usize;
502
503        for (i, p) in data.into_iter().enumerate() {
504            if !p.matches(self.header.point_format()) {
505                invalid_points = Err(crate::Error::InvalidPoint(
506                    crate::PointAddError::PointAttributesDoNotMatch(*self.header.point_format()),
507                ));
508                continue;
509            }
510            if !bounds_contains_point(&self.root_node.bounds, &p) {
511                if invalid_points.is_ok() {
512                    invalid_points = Err(crate::Error::InvalidPoint(
513                        crate::PointAddError::PointNotInBounds,
514                    ));
515                }
516                continue;
517            }
518
519            // if the given num_points was smaller than the actual number of points
520            // and we have passed that number revert to the greedy strategy
521            if num_points <= i {
522                self.add_point_greedy(p)?;
523            } else {
524                self.add_point_stochastic(p, expected_levels)?;
525            }
526        }
527        invalid_points
528    }
529
530    /// Close is called after the last point is written
531    fn close(&mut self) -> crate::Result<()> {
532        if self.is_closed {
533            return Err(crate::Error::ClosedWriter);
534        }
535        if self.header.number_of_points() < 1 {
536            return Err(crate::Error::EmptyCopcFile);
537        }
538
539        // write the unclosed chunks, order does not matter
540        for (key, chunk) in self.open_chunks.drain() {
541            let inner = chunk.into_inner();
542            if inner.is_empty() {
543                continue;
544            }
545            let (chunk_table_entry, chunk_offset) = self.compressor.compress_chunk(inner)?;
546            self.hierarchy.entries.push(Entry {
547                key,
548                offset: chunk_offset,
549                byte_size: chunk_table_entry.byte_count as i32,
550                point_count: chunk_table_entry.point_count as i32,
551            })
552        }
553
554        self.compressor.done()?;
555
556        let start_of_first_evlr = self.compressor.get_mut().stream_position()?;
557
558        let raw_evlrs: Vec<las::Result<las::raw::Vlr>> = self
559            .header
560            .evlrs()
561            .iter()
562            .map(|evlr| evlr.clone().into_raw(true))
563            .collect();
564
565        // write copc-evlr
566        self.hierarchy
567            .clone()
568            .into_evlr()?
569            .into_raw(true)?
570            .write_to(self.compressor.get_mut())?;
571        // write the rest of the evlrs
572        for raw_evlr in raw_evlrs {
573            raw_evlr?.write_to(self.compressor.get_mut())?;
574        }
575
576        self.compressor
577            .get_mut()
578            .seek(SeekFrom::Start(self.start))?;
579        self.header.clone().into_raw().and_then(|mut raw_header| {
580            if let Some(mut e) = raw_header.evlr {
581                e.start_of_first_evlr = start_of_first_evlr;
582                e.number_of_evlrs += 1;
583            } else {
584                raw_header.evlr = Some(las::raw::header::Evlr {
585                    start_of_first_evlr,
586                    number_of_evlrs: 1,
587                });
588            }
589            raw_header.write_to(self.compressor.get_mut())
590        })?;
591
592        // update the copc info vlr and write it
593        self.copc_info.spacing =
594            2. * self.copc_info.halfsize / (self.root_node.entry.point_count as f64);
595        self.copc_info.root_hier_offset = start_of_first_evlr + 60; // the header is 60bytes
596        self.copc_info.root_hier_size = self.hierarchy.byte_size();
597
598        self.copc_info
599            .clone()
600            .into_vlr()?
601            .into_raw(false)?
602            .write_to(self.compressor.get_mut())?;
603
604        self.compressor
605            .get_mut()
606            .seek(SeekFrom::Start(self.start))?;
607
608        self.is_closed = true;
609        Ok(())
610    }
611
612    // find the first non-full octree-node that contains the point
613    // and add it to the node, if the node now is full
614    // add the node to the hierarchy page and write to file
615    fn add_point_greedy(&mut self, point: las::Point) -> crate::Result<()> {
616        self.header.add_point(&point);
617
618        if point.gps_time.unwrap() < self.copc_info.gpstime_minimum {
619            self.copc_info.gpstime_minimum = point.gps_time.unwrap();
620        } else if point.gps_time.unwrap() > self.copc_info.gpstime_maximum {
621            self.copc_info.gpstime_maximum = point.gps_time.unwrap();
622        }
623
624        let mut node_key = None;
625        let mut write_chunk = false;
626
627        let root_bounds = self.root_node.bounds;
628
629        // starting from the root walk thorugh the octree
630        // and find the correct node to add the point to
631        let mut nodes_to_check = vec![&mut self.root_node];
632        while let Some(node) = nodes_to_check.pop() {
633            if !bounds_contains_point(&node.bounds, &point) {
634                // the point does not belong to this subtree
635                continue;
636            }
637            if node.is_full(self.max_node_size) {
638                // the point belongs to the subtree, but this node is full
639                // need to push the node's children to the nodes_to_check stack
640                if node.children.is_empty() {
641                    // the node does not have any children
642                    // so lets add children to the node
643                    let child_keys = node.entry.key.children();
644                    for key in child_keys {
645                        let child_bounds = key.bounds(&root_bounds);
646                        node.children.push(OctreeNode {
647                            entry: Entry {
648                                key,
649                                offset: 0,
650                                byte_size: 0,
651                                point_count: 0,
652                            },
653                            bounds: child_bounds,
654                            children: Vec::with_capacity(8),
655                        })
656                    }
657                }
658                // push the children to the stack
659                for child in node.children.iter_mut() {
660                    nodes_to_check.push(child);
661                }
662            } else {
663                // we've found the first non-full node that contains the point
664                node_key = Some(node.entry.key.clone());
665                node.entry.point_count += 1;
666
667                // check if the node now is full
668                write_chunk = node.is_full(self.max_node_size);
669                break;
670            }
671        }
672        let Some(node_key) = node_key else {
673            return Err(crate::Error::PointNotAddedToAnyNode);
674        };
675
676        let raw_point = point.into_raw(self.header.transforms())?;
677
678        if !self.open_chunks.contains_key(&node_key) {
679            let mut val = Cursor::new(vec![]);
680            raw_point.write_to(&mut val, self.header.point_format())?;
681
682            self.open_chunks.insert(node_key.clone(), val);
683        } else {
684            let buffer = self.open_chunks.get_mut(&node_key).unwrap();
685            raw_point.write_to(buffer, self.header.point_format())?;
686        }
687
688        if write_chunk {
689            let chunk = self.open_chunks.remove(&node_key).unwrap();
690            let (chunk_table_entry, chunk_offset) =
691                self.compressor.compress_chunk(chunk.into_inner())?;
692            self.hierarchy.entries.push(Entry {
693                key: node_key,
694                offset: chunk_offset,
695                byte_size: chunk_table_entry.byte_count as i32,
696                point_count: chunk_table_entry.point_count as i32,
697            });
698        }
699        Ok(())
700    }
701
702    fn add_point_stochastic(
703        &mut self,
704        point: las::Point,
705        expected_levels: usize,
706    ) -> crate::Result<()> {
707        // strategy: find the deepest node that contains this point
708        // choose at (weighted) random this node or one of its parents
709        // add point to that node
710        // write full nodes to file
711
712        let root_bounds = self.root_node.bounds;
713
714        let mut node_candidates = vec![];
715
716        // starting from the root walk thorugh the octree
717        let mut nodes_to_check = vec![&mut self.root_node];
718        while let Some(node) = nodes_to_check.pop() {
719            if !bounds_contains_point(&node.bounds, &point) {
720                // the point does not belong to this subtree
721                continue;
722            }
723
724            if node.children.is_empty() && node.entry.key.level < expected_levels as i32 {
725                let child_keys = node.entry.key.children();
726                for key in child_keys {
727                    let child_bounds = key.bounds(&root_bounds);
728                    node.children.push(OctreeNode {
729                        entry: Entry {
730                            key,
731                            offset: 0,
732                            byte_size: 0,
733                            point_count: 0,
734                        },
735                        bounds: child_bounds,
736                        children: Vec::with_capacity(8),
737                    })
738                }
739            }
740            if !node.is_full(self.max_node_size) {
741                node_candidates.push(&mut node.entry);
742            }
743            // push the children to the stack
744            for child in node.children.iter_mut() {
745                nodes_to_check.push(child);
746            }
747        }
748
749        if node_candidates.is_empty() {
750            // we need to add a new level, revert to greedy approach
751            return self.add_point_greedy(point);
752        }
753
754        // weighted by the inverse of the area (should volume be used?) the nodes cover
755        let chosen_index = get_random_weighted_index(&node_candidates);
756
757        let chosen_entry = &mut node_candidates[chosen_index];
758
759        chosen_entry.point_count += 1;
760
761        let write_chunk = chosen_entry.point_count > self.max_node_size;
762
763        let node_key = chosen_entry.key.clone();
764
765        self.header.add_point(&point);
766
767        if point.gps_time.unwrap() < self.copc_info.gpstime_minimum {
768            self.copc_info.gpstime_minimum = point.gps_time.unwrap();
769        } else if point.gps_time.unwrap() > self.copc_info.gpstime_maximum {
770            self.copc_info.gpstime_maximum = point.gps_time.unwrap();
771        }
772
773        let raw_point = point.into_raw(self.header.transforms())?;
774
775        if !self.open_chunks.contains_key(&node_key) {
776            let mut val = Cursor::new(vec![]);
777            raw_point.write_to(&mut val, self.header.point_format())?;
778
779            self.open_chunks.insert(node_key.clone(), val);
780        } else {
781            let buffer = self.open_chunks.get_mut(&node_key).unwrap();
782            raw_point.write_to(buffer, self.header.point_format())?;
783        }
784
785        if write_chunk {
786            let chunk = self.open_chunks.remove(&node_key).unwrap();
787            let (chunk_table_entry, chunk_offset) =
788                self.compressor.compress_chunk(chunk.into_inner())?;
789            self.hierarchy.entries.push(Entry {
790                key: node_key,
791                offset: chunk_offset,
792                byte_size: chunk_table_entry.byte_count as i32,
793                point_count: chunk_table_entry.point_count as i32,
794            });
795        }
796        Ok(())
797    }
798}
799
800fn get_random_weighted_index(entries: &Vec<&mut Entry>) -> usize {
801    // calculate weights
802    let levels: Vec<i32> = entries.iter().map(|e| e.key.level).collect();
803    let zero_level = levels[0];
804
805    // for each level down the side lengths are halved i.e area is a quarter
806    let areas: Vec<f64> = levels
807        .iter()
808        .map(|l| (0.25_f64).powi(l - zero_level))
809        .collect();
810    // total inv area
811    let inv_sum = areas.iter().fold(0., |acc, a| acc + 1. / a);
812
813    let weights: Vec<f64> = areas.iter().map(|a| (1. / a) / inv_sum).collect();
814
815    // get random index
816    let random = fastrand::f64();
817    let mut chosen_index = weights.len() - 1;
818
819    for i in 0..weights.len() - 1 {
820        if (weights[i]..=weights[i + 1]).contains(&random) {
821            chosen_index = i;
822            break;
823        }
824    }
825    chosen_index
826}
827
828impl<W: Write + Seek> Drop for CopcWriter<'_, W> {
829    fn drop(&mut self) {
830        if !self.is_closed {
831            // can only happen if the writer is created but no points is written
832            // or something goes wrong while writing
833            self.close()
834                .expect("Error when dropping the writer. No points written.");
835        }
836    }
837}
838
839#[inline]
840fn bounds_contains_point(b: &las::Bounds, p: &las::Point) -> bool {
841    !(b.max.x < p.x
842        || b.max.y < p.y
843        || b.max.z < p.z
844        || b.min.x > p.x
845        || b.min.y > p.y
846        || b.min.z > p.z)
847}