Skip to main content

gwseq_io/bbi/
reader.rs

1//! The bigWig / bigBed reader and its request builders.
2//!
3//! The plumbing only; the extraction kernels are in [`super::extract`].
4
5use std::fmt::Write as _;
6use std::io::Write as _;
7use std::sync::Arc;
8
9use ndarray::{Array1, Array2};
10
11use crate::bbi::extract::Extraction;
12use crate::bbi::header::{BbiHeader, BbiKind, TotalSummary, ZoomHeader};
13use crate::error::{Error, Result};
14use crate::genomic::{BinMode, ChrMap, IndexedLocs, LocBatch, Locs, Reduce};
15use crate::parallel::Executor;
16use crate::progress::{ProgressFn, ProgressTracker};
17use crate::source::ByteSource;
18
19/// Which zoom level a read uses.
20#[derive(Debug, Clone, Copy, Default)]
21pub enum Zoom {
22    /// Full data. The default, and the only option for a bigBed — which carries
23    /// no zoom data, so naming a level there is refused rather than ignored.
24    #[default]
25    Full,
26    /// The coarsest level whose bin size is under `bin_size * zoom_correction`,
27    /// which may be the full data.
28    Auto,
29    Level(usize),
30}
31
32/// Everything a reader holds that `close()` gives back.
33///
34/// Split out so that `close()` is `self.inner.take()`: dropping it drops the
35/// executor, which joins its threads, and the source, which releases the
36/// handle. The headers live outside it and stay readable after close, as the
37/// API documents.
38#[derive(Debug)]
39struct Inner {
40    source: Arc<dyn ByteSource>,
41    executor: Executor,
42}
43
44#[derive(Debug)]
45pub struct BbiReader {
46    inner: Option<Inner>,
47    path: String,
48    zoom_correction: f64,
49
50    // Read once at open, and readable after close.
51    pub(crate) header: BbiHeader,
52    pub(crate) zoom_headers: Vec<ZoomHeader>,
53    pub(crate) total_summary: TotalSummary,
54    pub(crate) chr_map: ChrMap,
55    /// Chromosome names indexed by the file's own chromosome index.
56    pub(crate) chr_names: Vec<String>,
57    /// bigBed only: the entry columns the file declares.
58    pub(crate) auto_sql: indexmap::IndexMap<String, String>,
59}
60
61/// The three resolutions a read pulls apart, which are not always the request's
62/// own `bin_size`.
63///
64/// One struct rather than three positional arguments because two of them are an
65/// `Option<f64>` and swapping those would still compile.
66#[derive(Default, Clone, Copy)]
67struct Grid {
68    /// Output bins per locus. `None` lets the locus span and the bin size
69    /// decide between them.
70    bin_count: Option<usize>,
71    /// What the loci snap to, in base pairs. `None` is the request's own
72    /// `bin_size`. The entry paths pass 1.0: an entry is not a bin, so snapping
73    /// a whole-chromosome locus down to a bin boundary would only lose the
74    /// entries past the last whole one.
75    snap: Option<f64>,
76    /// What a zoom level has to be finer than, when the output bins do not say.
77    /// `quantify` pins `bin_count` to 1 and still reads at `bin_size`, so its
78    /// single output bin says nothing about the resolution asked for.
79    zoom: Option<f64>,
80}
81
82impl Grid {
83    /// What `read_entries`, `read_all_entries` and `to_bed` read on: one bin
84    /// per locus, and the loci on the base grid rather than the request's.
85    fn entries() -> Self {
86        Self {
87            bin_count: Some(1),
88            snap: Some(1.0),
89            zoom: None,
90        }
91    }
92}
93
94impl BbiReader {
95    pub fn open(
96        path: &str,
97        parallel: i64,
98        zoom_correction: f64,
99        block_size: Option<u64>,
100        max_blocks: Option<usize>,
101    ) -> Result<Self> {
102        let source = crate::source::open(path, block_size, max_blocks)?;
103        Self::from_source(source, path, parallel, zoom_correction)
104    }
105
106    /// Open over an already-built source. What `gwseq_io::open` calls once it
107    /// has sniffed the magic, so the file is not opened twice.
108    pub(crate) fn from_source(
109        source: Arc<dyn ByteSource>,
110        path: &str,
111        parallel: i64,
112        zoom_correction: f64,
113    ) -> Result<Self> {
114        let header = super::header::read_header(source.as_ref())?;
115        let zoom_headers = super::header::read_zoom_headers(source.as_ref(), header.zoom_levels)?;
116        let total_summary =
117            super::header::read_total_summary(source.as_ref(), header.total_summary_offset)?;
118        let (chr_map, _tree) = super::chr_tree::read(source.as_ref(), header.chr_tree_offset)?;
119        // Only a bigBed has columns; a bigWig's field_count is 0 and its
120        // autoSql offset points at nothing.
121        let auto_sql = if header.kind.is_bigbed() {
122            super::header::read_auto_sql(
123                source.as_ref(),
124                header.auto_sql_offset,
125                header.field_count,
126            )?
127        } else {
128            indexmap::IndexMap::new()
129        };
130        // Chromosome names by the index the file gives them, which is what data
131        // records and R-tree items carry and what an entry has to be named from.
132        let mut chr_names =
133            vec![String::new(); chr_map.iter().map(|e| e.index + 1).max().unwrap_or(0)];
134        for entry in chr_map.iter() {
135            chr_names[entry.index] = entry.id.clone();
136        }
137        let executor = Executor::new(parallel)?;
138        Ok(Self {
139            inner: Some(Inner { source, executor }),
140            path: path.to_string(),
141            zoom_correction,
142            header,
143            zoom_headers,
144            total_summary,
145            chr_map,
146            chr_names,
147            auto_sql,
148        })
149    }
150
151    pub fn kind(&self) -> BbiKind {
152        self.header.kind
153    }
154    pub fn path(&self) -> &str {
155        &self.path
156    }
157    pub fn chr_sizes(&self) -> &ChrMap {
158        &self.chr_map
159    }
160    pub fn header(&self) -> &BbiHeader {
161        &self.header
162    }
163    pub fn zoom_headers(&self) -> &[ZoomHeader] {
164        &self.zoom_headers
165    }
166    pub fn total_summary(&self) -> &TotalSummary {
167        &self.total_summary
168    }
169    /// bigBed only: the entry columns the file declares, in file order, read at
170    /// open from the autoSql block. Empty for a bigWig, and empty for a bigBed
171    /// whose header names no autoSql offset.
172    pub fn auto_sql(&self) -> &indexmap::IndexMap<String, String> {
173        &self.auto_sql
174    }
175    pub fn is_closed(&self) -> bool {
176        self.inner.is_none()
177    }
178    pub fn parallel(&self) -> usize {
179        self.inner.as_ref().map_or(0, |i| i.executor.parallel())
180    }
181
182    /// Give back the threads and the file handle. Idempotent.
183    pub fn close(&mut self) {
184        if let Some(inner) = self.inner.take() {
185            inner.source.close();
186        }
187    }
188
189    /// At the top of every call that reads, not deep inside one, so the error
190    /// names the file rather than a handle.
191    fn inner(&self) -> Result<&Inner> {
192        self.inner.as_ref().ok_or_else(|| Error::Closed {
193            path: self.path.clone(),
194        })
195    }
196
197    /// Everything the read paths do before they diverge: resolve the loci,
198    /// split them into batches, pick a zoom level, and find the root of the
199    /// index that level's data hangs off.
200    #[allow(clippy::type_complexity)]
201    fn prepare<'a>(
202        &'a self,
203        inner: &'a Inner,
204        locs: &Locs,
205        common: &ReadCommon,
206        grid: Grid,
207    ) -> Result<(
208        IndexedLocs,
209        Vec<LocBatch>,
210        Option<usize>,
211        u64,
212        ProgressTracker,
213    )> {
214        // bigBed files carry no zoom data: the entries are piled up from the
215        // full index whatever is asked for, so an explicit request is refused
216        // rather than silently ignored.
217        if self.header.kind.is_bigbed() && !matches!(common.zoom, Zoom::Full) {
218            return Err(Error::invalid("zoom is only supported for bigwig files"));
219        }
220        let indexed = IndexedLocs::build(
221            &self.chr_map,
222            locs,
223            grid.snap.unwrap_or(common.bin_size),
224            grid.bin_count,
225            common.full_bin,
226        )?;
227        let (batches, coverage) = indexed.batches(inner.executor.parallel());
228        let level = self.select_zoom(
229            grid.zoom
230                .unwrap_or_else(|| indexed.effective_bin_size(common.bin_size)),
231            common.zoom,
232        )?;
233        let index_offset = match level {
234            Some(i) => self.zoom_headers[i].index_offset,
235            None => self.header.full_index_offset,
236        };
237        super::header::check_data_tree_magic(inner.source.as_ref(), index_offset)?;
238        let tree_root = index_offset + super::header::DATA_TREE_HEADER_SIZE;
239        Ok((
240            indexed,
241            batches,
242            level,
243            tree_root,
244            ProgressTracker::with_callback(coverage, common.progress.clone()),
245        ))
246    }
247
248    fn extraction<'a>(
249        &'a self,
250        inner: &'a Inner,
251        indexed: &'a IndexedLocs,
252        batches: &'a [LocBatch],
253        level: Option<usize>,
254        tree_root: u64,
255        tracker: &'a ProgressTracker,
256    ) -> Extraction<'a> {
257        Extraction {
258            source: inner.source.as_ref(),
259            locs: indexed,
260            batches,
261            tree_root,
262            zoom: level.is_some(),
263            uncompress_buffer_size: self.header.uncompress_buffer_size,
264            tracker,
265        }
266    }
267
268    /// (loci, bins) of `f32`.
269    pub fn read_values(&self, req: &ValuesRequest) -> Result<Array2<f32>> {
270        let inner = self.inner()?;
271        let (indexed, batches, level, root, tracker) = self.prepare(
272            inner,
273            &req.common.locs,
274            &req.common,
275            Grid {
276                bin_count: req.common.bin_count,
277                ..Grid::default()
278            },
279        )?;
280        let ex = self.extraction(inner, &indexed, &batches, level, root, &tracker);
281        // For a bigBed the "values" are the depth of coverage its entries make,
282        // which has no bin_mode to choose between — the three coincide.
283        let flat = if self.header.kind.is_bigbed() {
284            super::extract::entries_pileup(
285                &ex,
286                &inner.executor,
287                &self.auto_sql,
288                req.common.def_value,
289            )?
290        } else {
291            super::extract::values(&ex, &inner.executor, req.bin_mode, req.common.def_value)?
292        };
293        tracker.done_report();
294        let rows = indexed.locs.len();
295        let cols = indexed.bin_count;
296        Array2::from_shape_vec((rows, cols), flat)
297            .map_err(|e| Error::invalid(format!("output shape {rows}x{cols}: {e}")))
298    }
299
300    /// One value per locus.
301    pub fn quantify(&self, req: &QuantifyRequest) -> Result<Array1<f32>> {
302        let inner = self.inner()?;
303        let is_bigbed = self.header.kind.is_bigbed();
304        // A bigWig reads a single bin per locus, the extraction weighting each
305        // value by the bases it covers, so the reduction runs over the bases
306        // without the bins having to exist. `zoom_bin_size` carries the
307        // resolution the caller asked for past that pinned bin_count.
308        //
309        // A bigBed has no values to weight — its signal is the depth its
310        // entries make, which exists only once they are piled up — so it is
311        // binned at the size the caller asked for and reduced over those bins,
312        // as read_values bins it.
313        let (indexed, batches, level, root, tracker) = self.prepare(
314            inner,
315            &req.common.locs,
316            &req.common,
317            Grid {
318                bin_count: if is_bigbed {
319                    req.common.bin_count
320                } else {
321                    Some(1)
322                },
323                zoom: Some(req.common.bin_size),
324                ..Grid::default()
325            },
326        )?;
327        let ex = self.extraction(inner, &indexed, &batches, level, root, &tracker);
328        let mut stats = if is_bigbed {
329            let pileup = super::extract::entries_pileup(
330                &ex,
331                &inner.executor,
332                &self.auto_sql,
333                req.common.def_value,
334            )?;
335            super::extract::pileup_stats(&indexed, &pileup)
336        } else {
337            super::extract::values_stats(&ex, &inner.executor)?
338        };
339        tracker.done_report();
340
341        // The uncovered part of a locus still takes part, with def_value
342        // standing in for it, unless def_value is NaN — which is how a caller
343        // asks for the covered part alone. Filling it in rather than only
344        // stretching the denominator keeps every reduce mode agreeing on what
345        // those bases hold: min and max see def_value, and a non-zero one
346        // reaches the sum instead of reading as zero.
347        let def = req.common.def_value;
348        if !def.is_nan() {
349            for loc in &indexed.locs {
350                let s = &mut stats[loc.row(indexed.bin_count)];
351                // Counted in whatever unit the extraction accumulates: bases of
352                // the locus for the values, bins of it for the bed pileup —
353                // where nothing is ever missing, the pileup having already put
354                // def_value in the bins no entry reached.
355                let total = if is_bigbed {
356                    (loc.output_end - loc.output_start) as i64
357                } else {
358                    loc.binned_end - loc.binned_start
359                };
360                let missing = total - s.count;
361                if missing <= 0 {
362                    continue;
363                }
364                s.add_repeated(def, missing);
365            }
366        }
367        Ok(Array1::from_vec(
368            stats.iter().map(|s| s.reduce(req.reduce, def)).collect(),
369        ))
370    }
371
372    /// One value per bin, reduced across loci.
373    pub fn profile(&self, req: &ProfileRequest) -> Result<Array1<f32>> {
374        let inner = self.inner()?;
375        let (indexed, batches, level, root, tracker) = self.prepare(
376            inner,
377            &req.common.locs,
378            &req.common,
379            Grid {
380                bin_count: req.common.bin_count,
381                ..Grid::default()
382            },
383        )?;
384        let ex = self.extraction(inner, &indexed, &batches, level, root, &tracker);
385        // `def_value` reaches the pileup, so a bin no entry covered holds what
386        // the caller asked for rather than a bare 0. Dropping it would make
387        // `profile(def_value=NaN)` report 0 over an uncovered region where
388        // read_values reports NaN.
389        let mut stats = if self.header.kind.is_bigbed() {
390            let pileup = super::extract::entries_pileup(
391                &ex,
392                &inner.executor,
393                &self.auto_sql,
394                req.common.def_value,
395            )?;
396            super::extract::pileup_profile(&indexed, &pileup)
397        } else {
398            super::extract::values_profile(&ex, &inner.executor, req.bin_mode)?
399        };
400        tracker.done_report();
401
402        // A locus whose bin held no data still takes part in the profile, with
403        // def_value standing in for it, unless def_value is NaN — which is how
404        // a caller asks for those loci to be left out of the reduction instead.
405        let def = req.common.def_value;
406        if !def.is_nan() {
407            let loc_count = indexed.locs.len() as i64;
408            for s in &mut stats {
409                let missing = loc_count - s.count;
410                if missing <= 0 {
411                    continue;
412                }
413                s.add_repeated(def, missing);
414            }
415        }
416        Ok(Array1::from_vec(
417            stats.iter().map(|s| s.reduce(req.reduce, def)).collect(),
418        ))
419    }
420
421    /// bigBed entries, per locus, in the order the loci were given.
422    ///
423    /// Each locus gets its own full list, so two overlapping loci both report
424    /// the entries they share.
425    pub fn read_entries(&self, req: &EntriesRequest) -> Result<Vec<Vec<super::BedEntry>>> {
426        let inner = self.inner()?;
427        self.require_bigbed("read_entries")?;
428        self.check_col_count(req.col_count, 3)?;
429        let (indexed, batches, level, root, tracker) =
430            self.prepare(inner, &req.common.locs, &req.common, Grid::entries())?;
431        let ex = self.extraction(inner, &indexed, &batches, level, root, &tracker);
432        let out = super::extract::entries(
433            &ex,
434            &inner.executor,
435            &self.auto_sql,
436            &self.chr_names,
437            req.col_count,
438        )?;
439        tracker.done_report();
440        Ok(out)
441    }
442
443    /// Every entry on the named chromosomes, in chromosome then coordinate
444    /// order.
445    pub fn read_all_entries(&self, req: &EntriesRequest) -> Result<Vec<super::BedEntry>> {
446        let inner = self.inner()?;
447        self.require_bigbed("read_all_entries")?;
448        self.check_col_count(req.col_count, 3)?;
449        let locs = Locs::whole_chromosomes(&self.chr_map, &req.common.locs.chr_ids)?;
450        let (indexed, batches, level, root, tracker) =
451            self.prepare(inner, &locs, &req.common, Grid::entries())?;
452        let ex = self.extraction(inner, &indexed, &batches, level, root, &tracker);
453        let by_chr = super::extract::entries(
454            &ex,
455            &inner.executor,
456            &self.auto_sql,
457            &self.chr_names,
458            req.col_count,
459        )?;
460        tracker.done_report();
461        Ok(by_chr.into_iter().flatten().collect())
462    }
463
464    /// A method that only exists for one of the two formats says which.
465    fn require_bigbed(&self, what: &str) -> Result<()> {
466        if self.header.kind.is_bigbed() {
467            Ok(())
468        } else {
469            Err(Error::invalid(format!("{what} only for bigbed")))
470        }
471    }
472
473    fn require_bigwig(&self, what: &str) -> Result<()> {
474        if self.header.kind.is_bigbed() {
475            Err(Error::invalid(format!("{what} only for bigwig")))
476        } else {
477            Ok(())
478        }
479    }
480
481    /// Check a `col_count` against the columns the file holds.
482    ///
483    /// `min` is 3 for the entry readers, whose entries are built from the
484    /// coordinates, and 1 for `to_bed`, which writes as many columns as it is
485    /// told and has no such floor.
486    pub(crate) fn check_col_count(&self, col_count: usize, min: usize) -> Result<()> {
487        if col_count == 0 {
488            return Ok(());
489        }
490        if col_count < min {
491            return Err(Error::invalid(format!(
492                "col_count {col_count} must be 0 or at least {min}"
493            )));
494        }
495        if col_count > self.header.field_count as usize {
496            return Err(Error::invalid(format!(
497                "col_count {col_count} exceeds number of fields {}",
498                self.header.field_count
499            )));
500        }
501        Ok(())
502    }
503
504    /// Bytes of data the file holds at a given resolution.
505    ///
506    /// What a whole-file walk sizes its pieces from: a zoom level holds a
507    /// fraction of what the full data does, so a walk reading one is worth
508    /// splitting far less finely.
509    pub fn data_size(&self, zoom: Option<usize>) -> u64 {
510        match zoom.and_then(|i| self.zoom_headers.get(i)) {
511            Some(z) => z.index_offset.saturating_sub(z.data_offset),
512            None => self
513                .header
514                .full_index_offset
515                .saturating_sub(self.header.full_data_offset),
516        }
517    }
518
519    /// Base pairs the file's chromosomes hold between them.
520    pub fn genome_size(&self) -> i64 {
521        self.chr_map.genome_size()
522    }
523
524    /// Successive windows of values over whole chromosomes.
525    ///
526    /// A plain [`Iterator`], exhausted after one pass. `len()` and `locs()`
527    /// are on the iterator, so
528    /// `iter.locs().iter().zip(iter)` works once, as documented.
529    pub fn iter_all_values(
530        &self,
531        req: &ValuesRequest,
532        window: i64,
533    ) -> Result<super::ValuesWindows<'_>> {
534        super::extract::ValuesWindows::plan(self, req, window)
535    }
536
537    pub fn iter_all_entries(
538        &self,
539        req: &EntriesRequest,
540        window: i64,
541    ) -> Result<super::EntryWindows<'_>> {
542        super::extract::EntryWindows::plan(self, req, window)
543    }
544
545    /// Every bin of the named chromosomes, as `chr\tstart\tend\tvalue`.
546    ///
547    /// The values are the ones [`Self::read_values`] gives for the same
548    /// chromosomes at the same `bin_size`, `bin_mode`, `full_bin`, `def_value`
549    /// and `zoom` — the export bins, it does not copy the file's own intervals.
550    /// A default request therefore writes one interval per base, and a
551    /// `bin_size` of 10,000 writes one per 10,000 bases.
552    ///
553    /// With `merge_bins`, adjacent bins carrying the same value become one
554    /// interval — the shape a bedGraph is usually in, and what keeps a
555    /// whole-genome export from being one line per base wherever the file says
556    /// nothing. Without it every bin is its own interval, so the line count is
557    /// the bin count and the grid is visible in the file.
558    ///
559    /// A `def_value` of NaN is how a caller asks for the covered part alone: a
560    /// bin holding NaN is not written at all, so the gaps stay gaps.
561    ///
562    /// `bin_count` is ignored — the bins are the chromosome's, not a fixed
563    /// number of them.
564    pub fn to_bedgraph(
565        &self,
566        out: &std::path::Path,
567        req: &ValuesRequest,
568        merge_bins: bool,
569    ) -> Result<()> {
570        self.require_bigwig("to_bedgraph")?;
571        self.export_bins(out, req, BedGraphSink::new(merge_bins))
572    }
573
574    /// The same bins as fixedStep WIG sections.
575    ///
576    /// A fixedStep section walks a fixed step from its own start, so anything
577    /// that breaks the run of bins has to open a new one: a change of
578    /// chromosome, a bin left out because its value is NaN, and the shorter
579    /// last bin a `full_bin` export ends a chromosome on.
580    ///
581    /// Bins are written one per line whatever their values, since a fixedStep
582    /// section has no way to say "and again" — the `merge_bins` of
583    /// [`Self::to_bedgraph`] has no equivalent here.
584    pub fn to_wig(&self, out: &std::path::Path, req: &ValuesRequest) -> Result<()> {
585        self.require_bigwig("to_wig")?;
586        self.export_bins(out, req, WigSink::default())
587    }
588
589    /// Every entry as a tab-separated BED line.
590    ///
591    /// `col_count` here has a floor of 1, not 3: this writes as many columns as
592    /// it is told, where the entry readers build an entry from its coordinates
593    /// and cannot go below them.
594    pub fn to_bed(&self, out: &std::path::Path, req: &EntriesRequest) -> Result<()> {
595        let inner = self.inner()?;
596        self.require_bigbed("to_bed")?;
597        self.check_col_count(req.col_count, 1)?;
598        let col_count = if req.col_count == 0 {
599            self.header.field_count as usize
600        } else {
601            req.col_count
602        };
603
604        let locs = Locs::whole_chromosomes(&self.chr_map, &req.common.locs.chr_ids)?;
605        let (indexed, batches, level, root, tracker) =
606            self.prepare(inner, &locs, &req.common, Grid::entries())?;
607        let ex = self.extraction(inner, &indexed, &batches, level, root, &tracker);
608        let wanted = self.walked_chrs(&indexed);
609
610        let mut writer = std::io::BufWriter::new(
611            std::fs::File::create(out).map_err(|e| Error::io(out.to_string_lossy(), e))?,
612        );
613        let mut line = String::new();
614        for batch in &batches {
615            ex.walk_bed_batch(*batch, &self.auto_sql, col_count, |entry, _| {
616                if !wanted.contains(&entry.chr_index) {
617                    return Ok(());
618                }
619                line.clear();
620                line.push_str(self.chr_name(entry.chr_index));
621                if col_count >= 2 {
622                    let _ = write!(line, "\t{}", entry.start);
623                }
624                if col_count >= 3 {
625                    let _ = write!(line, "\t{}", entry.end);
626                }
627                for (_, value) in entry.fields.iter().take(col_count.saturating_sub(3)) {
628                    line.push('\t');
629                    line.push_str(value);
630                }
631                line.push('\n');
632                write_line(&mut writer, &line, out)
633            })?;
634        }
635        flush(&mut writer, out)?;
636        tracker.done_report();
637        Ok(())
638    }
639
640    /// The shared body of `to_bedgraph` and `to_wig`: bin the named
641    /// chromosomes exactly as `read_values` bins them and let `sink` turn each
642    /// bin into text.
643    ///
644    /// The read runs on the pool a window at a time, through the very walk
645    /// `iter_all_values` hands to a caller, which is what makes an export and a
646    /// read of the same chromosomes agree by construction rather than by two
647    /// implementations of the same arithmetic. The formatting and the writing
648    /// are single-threaded on purpose: the output is one file in chromosome
649    /// then coordinate order, so formatted text produced out of order would
650    /// have to be held until its turn came, and a window of it is the memory
651    /// this avoids.
652    fn export_bins(
653        &self,
654        out: &std::path::Path,
655        req: &ValuesRequest,
656        mut sink: impl BinSink,
657    ) -> Result<()> {
658        // Named before the file is created: an export through a closed reader
659        // should not leave an empty file behind.
660        self.inner()?;
661        let walk = super::extract::ValuesWindows::plan(self, req, Self::export_window(req)?)?;
662        let bin_size = walk.bin_size();
663        // Cloned up front: the walk is consumed by the loop below, and each
664        // window needs the chromosome and start coordinate of its own region.
665        let locs = walk.locs().to_vec();
666
667        let mut writer = std::io::BufWriter::new(
668            std::fs::File::create(out).map_err(|e| Error::io(out.to_string_lossy(), e))?,
669        );
670        let mut line = String::new();
671        for (index, values) in walk.enumerate() {
672            let values = values?;
673            let (chr, window_start, window_end) = &locs[index];
674            for (i, &value) in values.iter().enumerate() {
675                // A NaN bin is left out of the file rather than written as the
676                // text "NaN", which no reader of either format accepts. That is
677                // what makes `def_value = NaN` an export of the covered part
678                // alone, and the coordinate discontinuity it leaves is what
679                // both sinks break a run on.
680                if value.is_nan() {
681                    continue;
682                }
683                let start = window_start + i as i64 * bin_size;
684                // The last bin of a `full_bin` chromosome is the short one.
685                let end = (start + bin_size).min(*window_end);
686                line.clear();
687                sink.bin(&mut line, chr, start, end, value);
688                if !line.is_empty() {
689                    write_line(&mut writer, &line, out)?;
690                }
691            }
692        }
693        line.clear();
694        sink.finish(&mut line);
695        if !line.is_empty() {
696            write_line(&mut writer, &line, out)?;
697        }
698        flush(&mut writer, out)
699    }
700
701    /// How much of a chromosome an export reads at a time, in base pairs.
702    ///
703    /// Two caps, and the narrower one wins. [`EXPORT_WINDOW_BINS`] bounds what
704    /// a window costs to hold — 4 MiB of `f32` for a 1 bp export of a human
705    /// chromosome as much as for a 10 kb one — and [`EXPORT_WINDOW_BASES`]
706    /// bounds how much of a chromosome goes by between two progress reports,
707    /// which a bin-only cap would leave at "the whole thing" for any bin wide
708    /// enough. Never below one bin, since a bin may be wider than either cap.
709    ///
710    /// `saturating_mul` because a bin size is any positive whole number of
711    /// bases, including one larger than the genome.
712    fn export_window(req: &ValuesRequest) -> Result<i64> {
713        let bin_size =
714            crate::genomic::BinPlan::new(req.common.bin_size, None, req.common.full_bin)?
715                .whole_bin_size();
716        Ok(bin_size
717            .saturating_mul(EXPORT_WINDOW_BINS)
718            .min(EXPORT_WINDOW_BASES)
719            .max(bin_size))
720    }
721
722    /// The chromosomes a walk was asked for.
723    ///
724    /// A tree item may straddle a chromosome boundary, so a whole-file walk can
725    /// reach data from chromosomes that were never asked for. The exporters
726    /// filter their output on this.
727    fn walked_chrs(&self, locs: &IndexedLocs) -> std::collections::HashSet<u32> {
728        locs.locs.iter().map(|l| l.chr_index as u32).collect()
729    }
730
731    pub(crate) fn chr_name(&self, index: u32) -> &str {
732        self.chr_names
733            .get(index as usize)
734            .map(String::as_str)
735            .unwrap_or("")
736    }
737
738    /// Which zoom level to read, or `None` for the full data.
739    ///
740    /// The automatic choice is the **coarsest level still finer than** the
741    /// output bins, so that each bin averages several summaries rather than
742    /// inheriting the edges of one. `zoom_correction` is how much finer. It may
743    /// well settle on the full data, which is what a file whose coarsest level
744    /// is still wider than the output bins has to be read at.
745    pub(crate) fn select_zoom(&self, bin_size: f64, zoom: Zoom) -> Result<Option<usize>> {
746        let count = self.zoom_headers.len();
747        match zoom {
748            Zoom::Full => Ok(None),
749            Zoom::Level(level) => {
750                if level < count {
751                    Ok(Some(level))
752                } else if count == 0 {
753                    Err(Error::invalid("file has no zoom level"))
754                } else {
755                    Err(Error::invalid(format!(
756                        "requested zoom level {level} exceeds max zoom level {}",
757                        count - 1
758                    )))
759                }
760            }
761            Zoom::Auto => {
762                let threshold = (bin_size * self.zoom_correction).round() as i64;
763                let mut best: Option<usize> = None;
764                let mut best_reduction = 0i64;
765                for (i, zoom) in self.zoom_headers.iter().enumerate() {
766                    let reduction = zoom.reduction_level as i64;
767                    if reduction <= threshold && reduction > best_reduction {
768                        best_reduction = reduction;
769                        best = Some(i);
770                    }
771                }
772                Ok(best)
773            }
774        }
775    }
776}
777
778// ---------------------------------------------------------------------------
779// Requests
780// ---------------------------------------------------------------------------
781//
782// A read takes up to fourteen optional parameters, and Rust has no defaulted
783// parameters, so each request shape gets a builder — shared by the
784// Python layer, the CLI and Rust callers, so there is one place a default is
785// written down.
786
787/// Shared by every read: which loci, at what resolution, from which zoom.
788pub struct ReadCommon {
789    pub locs: Locs,
790    pub bin_size: f64,
791    pub bin_count: Option<usize>,
792    pub full_bin: bool,
793    pub def_value: f32,
794    pub zoom: Zoom,
795    pub progress: Option<ProgressFn>,
796}
797
798impl ReadCommon {
799    pub fn new(locs: Locs) -> Self {
800        Self {
801            locs,
802            bin_size: 1.0,
803            bin_count: None,
804            full_bin: false,
805            def_value: 0.0,
806            zoom: Zoom::Full,
807            progress: None,
808        }
809    }
810}
811
812macro_rules! read_common_builders {
813    ($t:ty) => {
814        impl $t {
815            pub fn bin_size(mut self, v: f64) -> Self {
816                self.common.bin_size = v;
817                self
818            }
819            pub fn bin_count(mut self, v: usize) -> Self {
820                self.common.bin_count = Some(v);
821                self
822            }
823            pub fn full_bin(mut self, v: bool) -> Self {
824                self.common.full_bin = v;
825                self
826            }
827            pub fn def_value(mut self, v: f32) -> Self {
828                self.common.def_value = v;
829                self
830            }
831            pub fn zoom(mut self, v: Zoom) -> Self {
832                self.common.zoom = v;
833                self
834            }
835            pub fn progress(mut self, f: ProgressFn) -> Self {
836                self.common.progress = Some(f);
837                self
838            }
839        }
840    };
841}
842
843pub struct ValuesRequest {
844    pub common: ReadCommon,
845    pub bin_mode: BinMode,
846}
847
848pub struct QuantifyRequest {
849    pub common: ReadCommon,
850    pub reduce: Reduce,
851}
852
853pub struct ProfileRequest {
854    pub common: ReadCommon,
855    pub bin_mode: BinMode,
856    pub reduce: Reduce,
857}
858
859pub struct EntriesRequest {
860    pub common: ReadCommon,
861    /// 0 for all, otherwise at least 3 — or at least 1 for [`BbiReader::to_bed`],
862    /// which writes the columns it is told to rather than building an entry out
863    /// of them. Columns left out are never parsed, so a narrower read is a
864    /// cheaper one.
865    pub col_count: usize,
866}
867
868read_common_builders!(ValuesRequest);
869read_common_builders!(QuantifyRequest);
870read_common_builders!(ProfileRequest);
871read_common_builders!(EntriesRequest);
872
873impl ValuesRequest {
874    pub fn new(locs: Locs) -> Self {
875        Self {
876            common: ReadCommon::new(locs),
877            bin_mode: BinMode::Mean,
878        }
879    }
880    pub fn bin_mode(mut self, v: BinMode) -> Self {
881        self.bin_mode = v;
882        self
883    }
884}
885
886impl QuantifyRequest {
887    pub fn new(locs: Locs) -> Self {
888        Self {
889            common: ReadCommon::new(locs),
890            reduce: Reduce::Mean,
891        }
892    }
893    pub fn reduce(mut self, v: Reduce) -> Self {
894        self.reduce = v;
895        self
896    }
897}
898
899impl ProfileRequest {
900    pub fn new(locs: Locs) -> Self {
901        Self {
902            common: ReadCommon::new(locs),
903            bin_mode: BinMode::Mean,
904            reduce: Reduce::Mean,
905        }
906    }
907    pub fn bin_mode(mut self, v: BinMode) -> Self {
908        self.bin_mode = v;
909        self
910    }
911    pub fn reduce(mut self, v: Reduce) -> Self {
912        self.reduce = v;
913        self
914    }
915}
916
917impl EntriesRequest {
918    pub fn new(locs: Locs) -> Self {
919        Self {
920            common: ReadCommon::new(locs),
921            col_count: 0,
922        }
923    }
924    pub fn col_count(mut self, v: usize) -> Self {
925        self.col_count = v;
926        self
927    }
928}
929
930// ---------------------------------------------------------------------------
931// Export sinks
932// ---------------------------------------------------------------------------
933
934/// Bins per window of the walk the exporters stream through. A window of
935/// `f32`s is 4 MiB, which is what a whole-genome export holds however wide its
936/// bins are.
937const EXPORT_WINDOW_BINS: i64 = 1 << 20;
938
939/// Base pairs per window, the other half of the cap. Progress is reported once
940/// per window, so this is the resolution of a progress bar over an export —
941/// a few hundred steps across a human chromosome — and it is what keeps a
942/// wide-binned export from reading a whole chromosome before saying anything.
943const EXPORT_WINDOW_BASES: i64 = 16 << 20;
944
945/// How an export turns a bin into text.
946///
947/// Both implementations carry a run across bins — an interval being extended,
948/// a section being filled — so a bin often appends nothing and the pending
949/// line comes out one bin later. That is why this is a trait with a `finish`
950/// rather than a closure: the last run has to be written after the last bin.
951trait BinSink {
952    /// Append whatever lines this bin completes. Bins arrive in chromosome
953    /// then coordinate order, and a bin whose value is NaN never arrives.
954    fn bin(&mut self, line: &mut String, chr: &str, start: i64, end: i64, value: f32);
955    /// Append whatever is still pending once the last bin has gone by.
956    fn finish(&mut self, line: &mut String);
957}
958
959/// `chr\tstart\tend\tvalue`, one line per run of equal-valued bins — or per
960/// bin, without `merge`.
961#[derive(Default)]
962struct BedGraphSink {
963    chr: String,
964    start: i64,
965    end: i64,
966    value: f32,
967    /// Whether the three fields above stand for an interval yet. A bare `end`
968    /// of 0 cannot say so: a first bin at `0-1` is a real interval.
969    open: bool,
970    merge: bool,
971}
972
973impl BedGraphSink {
974    fn new(merge: bool) -> Self {
975        Self {
976            merge,
977            ..Self::default()
978        }
979    }
980
981    fn flush(&mut self, line: &mut String) {
982        if !self.open {
983            return;
984        }
985        let _ = write!(line, "{}\t{}\t{}\t", self.chr, self.start, self.end);
986        super::text::push_float(line, self.value);
987        line.push('\n');
988        self.open = false;
989    }
990}
991
992impl BinSink for BedGraphSink {
993    fn bin(&mut self, line: &mut String, chr: &str, start: i64, end: i64, value: f32) {
994        // `start == self.end` is what a skipped NaN bin breaks, and the chromosome
995        // test is what a new chromosome starting at the old one's end would
996        // otherwise slip past.
997        if self.merge && self.open && self.end == start && self.value == value && self.chr == chr {
998            self.end = end;
999            return;
1000        }
1001        self.flush(line);
1002        self.chr.clear();
1003        self.chr.push_str(chr);
1004        self.start = start;
1005        self.end = end;
1006        self.value = value;
1007        self.open = true;
1008    }
1009
1010    fn finish(&mut self, line: &mut String) {
1011        self.flush(line);
1012    }
1013}
1014
1015/// fixedStep sections, one value per line.
1016#[derive(Default)]
1017struct WigSink {
1018    chr: String,
1019    span: i64,
1020    /// Where the next bin has to start for the open section to hold it.
1021    next_start: i64,
1022    open: bool,
1023}
1024
1025impl BinSink for WigSink {
1026    fn bin(&mut self, line: &mut String, chr: &str, start: i64, end: i64, value: f32) {
1027        let span = end - start;
1028        if !self.open || span != self.span || start != self.next_start || self.chr != chr {
1029            // WIG coordinates are 1-based.
1030            let _ = writeln!(
1031                line,
1032                "fixedStep chrom={chr} start={} step={span} span={span}",
1033                start + 1
1034            );
1035            self.chr.clear();
1036            self.chr.push_str(chr);
1037            self.span = span;
1038            self.open = true;
1039        }
1040        super::text::push_float(line, value);
1041        line.push('\n');
1042        self.next_start = start + span;
1043    }
1044
1045    /// Nothing is ever pending: a fixedStep section is closed by the next one
1046    /// or by the end of the file.
1047    fn finish(&mut self, _line: &mut String) {}
1048}
1049
1050/// Write one formatted line, naming the output file if it fails.
1051fn write_line(
1052    writer: &mut std::io::BufWriter<std::fs::File>,
1053    line: &str,
1054    path: &std::path::Path,
1055) -> Result<()> {
1056    writer
1057        .write_all(line.as_bytes())
1058        .map_err(|e| Error::io(path.to_string_lossy(), e))
1059}
1060
1061/// Flush explicitly rather than leaving it to `Drop`, which cannot report a
1062/// failure — and a truncated export that reported success is the worst outcome
1063/// here.
1064fn flush(writer: &mut std::io::BufWriter<std::fs::File>, path: &std::path::Path) -> Result<()> {
1065    writer
1066        .flush()
1067        .map_err(|e| Error::io(path.to_string_lossy(), e))
1068}