Skip to main content

sva_samples/collapse/
plan.rs

1// Concern: which row of the collapse table a form takes, and what that row costs | Non-concern: running the row (collapse.rs) | IO: (&SpectralSum, rate, Horizon) -> Plan, flops
2
3use sva_formula::closed_form::{Part, map_children};
4use sva_formula::spectral_sum::atom::SpectralAtom;
5use sva_formula::{Body, ClosedForm, Lane, Line, SpectralSum, Var, normalize_closed_form};
6
7use super::truncate::Audible;
8use super::{Horizon, atoms, lines, point, truncate};
9use crate::error::CollapseError;
10use crate::label::Rule;
11use crate::profile::Profile;
12
13/// Both routes place a line exactly, so cost decides; the direct one hands a lone line back
14/// unchanged.
15pub fn direct_is_cheaper(lines: usize, n: usize, len: usize) -> bool {
16    (lines as f64) * (len as f64) <= (n as f64) * (n as f64).log2()
17}
18
19pub struct LinePlan {
20    pub placed: Vec<Vec<Line>>,
21    pub summed: Vec<Vec<Line>>,
22    pub dropped: Vec<Vec<Line>>,
23    pub bins: usize,
24    pub tail_db: Option<f64>,
25}
26
27/// FORMAT 9.2's window route.
28pub struct Group {
29    pub factor: SpectralAtom,
30    pub placed: Vec<Line>,
31    pub summed: Vec<Line>,
32}
33
34pub enum LanePlan {
35    Grouped { groups: Vec<Group>, bins: usize },
36    Sweep { atoms: usize, samples: usize },
37}
38
39pub struct Sampled {
40    pub sum: Box<SpectralSum>,
41    pub rule: Rule,
42    pub lanes: Vec<LanePlan>,
43}
44
45/// The row of FORMAT 9.1 that runs, decided once for the collapse and the count alike.
46pub enum Plan {
47    Spectrum(Box<SpectralSum>),
48    Lines(Box<LinePlan>),
49    Sampled(Box<Sampled>),
50    /// Row 4 over the written closed form: every node the truncation leaves, at every instant.
51    Point {
52        written: Box<ClosedForm>,
53        nodes: usize,
54        width: usize,
55    },
56    /// A sum whose addends name different rows, each taking its own.
57    Added(Vec<Plan>),
58}
59
60/// The rows in FORMAT 9.1's order; the first guard that holds decides.
61pub fn of(
62    sum: &SpectralSum,
63    rate: u32,
64    horizon: Horizon,
65    profile: &Profile,
66    len: usize,
67) -> Result<Plan, CollapseError> {
68    let ceiling = profile.ceiling(rate);
69    if sum.var == Var::F {
70        return Ok(Plan::Spectrum(Box::new(sum.clone())));
71    }
72    if let Some(found) = line_plan(sum, rate, horizon, profile, len, ceiling)? {
73        return Ok(Plan::Lines(Box::new(found)));
74    }
75    let truncated = truncate::spectral_sum(sum, Audible::of(profile, rate))?;
76    let rule = match () {
77        () if atoms::band_limited(&truncated, ceiling, profile) => Rule::BandLimited,
78        () if atoms::windowed(&truncated) => Rule::CroppedPair,
79        () => Rule::PointSampled,
80    };
81    let lanes = truncated
82        .lanes
83        .iter()
84        .map(|lane| lane_plan(lane, rate, horizon, len))
85        .collect();
86    Ok(Plan::Sampled(Box::new(Sampled {
87        sum: Box::new(truncated),
88        rule,
89        lanes,
90    })))
91}
92
93/// The row a closed form with no spectral sum takes: one per addend where they differ, else 4.
94pub fn of_written(
95    form: &ClosedForm,
96    rate: u32,
97    horizon: Horizon,
98    profile: &Profile,
99    len: usize,
100) -> Result<Plan, CollapseError> {
101    let Some(addends) = addends(form) else {
102        return point_plan(form, rate, profile);
103    };
104    let mut parts = Vec::with_capacity(addends.len());
105    for addend in &addends {
106        match of_term(addend, rate, horizon, profile, len)? {
107            Plan::Added(inner) => parts.extend(inner),
108            part => parts.push(part),
109        }
110    }
111    // Addends that all name row 4 are one row 4 over the whole sum, measured once.
112    match parts.iter().all(|part| matches!(part, Plan::Point { .. })) {
113        true => point_plan(form, rate, profile),
114        false => Ok(Plan::Added(parts)),
115    }
116}
117
118/// A row the addend alone cannot take refuses nothing; the written row still stands for it.
119/// A nesting past the bound is the one exception, refused wherever it is written.
120fn of_term(
121    form: &ClosedForm,
122    rate: u32,
123    horizon: Horizon,
124    profile: &Profile,
125    len: usize,
126) -> Result<Plan, CollapseError> {
127    let Ok(sum) = normalize_closed_form(form) else {
128        return of_written(form, rate, horizon, profile, len);
129    };
130    match of(&sum, rate, horizon, profile, len) {
131        Err(nested @ CollapseError::NestedSeries { .. }) => Err(nested),
132        Err(_) => of_written(form, rate, horizon, profile, len),
133        held => held,
134    }
135}
136
137fn addends(form: &ClosedForm) -> Option<Vec<ClosedForm>> {
138    if let Body::Add(parts) = &form.body {
139        return Some(
140            parts
141                .iter()
142                .map(|part| ClosedForm {
143                    var: form.var,
144                    body: (*part.body).clone(),
145                    origin: part.origin,
146                })
147                .collect(),
148        );
149    }
150    let under = linear_over(&form.body)?;
151    let held = addends(&ClosedForm {
152        body: under,
153        ..form.clone()
154    })?;
155    Some(
156        held.into_iter()
157            .map(|addend| ClosedForm {
158                body: map_children(&form.body, |_| {
159                    Part::new(addend.origin, addend.body.clone())
160                }),
161                ..addend
162            })
163            .collect(),
164    )
165}
166
167/// Each of these is linear, so over a sum it is the sum of itself over every addend.
168fn linear_over(f: &Body) -> Option<Body> {
169    match f {
170        Body::Crop { of, .. }
171        | Body::Shift { of, .. }
172        | Body::Deriv { of, .. }
173        | Body::Channel(of, _) => Some((*of.body).clone()),
174        _ => None,
175    }
176}
177
178fn point_plan(form: &ClosedForm, rate: u32, profile: &Profile) -> Result<Plan, CollapseError> {
179    let written = ClosedForm {
180        body: truncate::written(&form.body, Audible::of(profile, rate))?,
181        ..form.clone()
182    };
183    let width = point::width_of(&written.body, &point::NoRefs).max(1);
184    Ok(Plan::Point {
185        nodes: point_nodes(&written.body),
186        written: Box::new(written),
187        width,
188    })
189}
190
191/// What one instant's evaluation walks over every component: a join walks that component's branch.
192pub fn point_nodes(f: &Body) -> usize {
193    let width = point::width_of(f, &point::NoRefs).max(1);
194    (0..width).map(|c| nodes_at(f, c)).sum()
195}
196
197fn nodes_at(f: &Body, component: usize) -> usize {
198    let branch = |part: &Part, c: usize| nodes_at(&part.body, c);
199    match f {
200        Body::Join(parts) => {
201            let widths: Vec<usize> = parts
202                .iter()
203                .map(|p| point::width_of(&p.body, &point::NoRefs))
204                .collect();
205            match point::lane_of(&widths, component) {
206                Some((at, inner)) => 1 + branch(&parts[at], inner),
207                None => 1,
208            }
209        }
210        Body::Channel(of, k) => 1 + branch(of, usize::from(*k)),
211        other => {
212            1 + sva_formula::closed_form::children(other)
213                .iter()
214                .map(|part| branch(part, component))
215                .sum::<usize>()
216        }
217    }
218}
219
220/// Rows one and two: every atom a line, none of them windowed.
221fn line_plan(
222    sum: &SpectralSum,
223    rate: u32,
224    horizon: Horizon,
225    profile: &Profile,
226    len: usize,
227    ceiling: f64,
228) -> Result<Option<LinePlan>, CollapseError> {
229    let mut per_lane = Vec::with_capacity(sum.lanes.len());
230    let mut grids: Vec<f64> = Vec::new();
231    let mut tail: Option<f64> = None;
232    for lane in &sum.lanes {
233        match lines::of_lane(lane, ceiling, profile.floor(ceiling)) {
234            Some(found) => {
235                if let Some(left) = found.tail_db {
236                    tail = Some(tail.map_or(left, |held: f64| held.max(left)));
237                }
238                grids.extend(found.grids);
239                per_lane.push(found.lines);
240            }
241            None => return Ok(None),
242        }
243    }
244
245    let (mut kept, mut dropped): (Vec<Vec<Line>>, Vec<Vec<Line>>) = (Vec::new(), Vec::new());
246    for found in &per_lane {
247        let (here, gone): (Vec<Line>, Vec<Line>) = found.iter().partition(|l| l.hz.abs() < ceiling);
248        kept.push(here);
249        dropped.push(gone);
250    }
251    // A form with no line at all is silence, not a band every line sat above.
252    if kept.iter().all(Vec::is_empty) && !dropped.iter().all(Vec::is_empty) {
253        let lowest = dropped
254            .iter()
255            .flatten()
256            .map(|l| l.hz.abs())
257            .fold(f64::INFINITY, f64::min);
258        return Err(CollapseError::EmptyBand { ceiling, lowest });
259    }
260
261    let bins = lines::grid(&kept, &grids, rate, horizon.span(), len);
262    let split: Vec<(Vec<Line>, Vec<Line>)> = kept
263        .iter()
264        .map(|k| match direct_is_cheaper(k.len(), bins, len) {
265            true => (Vec::new(), k.clone()),
266            false => lines::split(k, bins, rate),
267        })
268        .collect();
269    Ok(Some(LinePlan {
270        placed: split.iter().map(|(on, _)| on.clone()).collect(),
271        summed: split.into_iter().map(|(_, off)| off).collect(),
272        dropped,
273        bins,
274        tail_db: tail,
275    }))
276}
277
278fn lane_plan(lane: &Lane, rate: u32, horizon: Horizon, len: usize) -> LanePlan {
279    let bins = lines::bins(horizon.span(), rate);
280    let sweep = LanePlan::Sweep {
281        atoms: lane.atoms.len(),
282        samples: super::span::nonzero(lane, horizon, rate, len)
283            .iter()
284            .map(|(from, to)| to - from)
285            .sum(),
286    };
287    let Some(found) = lines::grouped(lane) else {
288        return sweep;
289    };
290    let groups: Vec<Group> = found
291        .into_iter()
292        .map(|(factor, held)| {
293            let (placed, summed) = match direct_is_cheaper(held.len(), bins, len) {
294                true => (Vec::new(), held),
295                false => lines::split(&held, bins, rate),
296            };
297            Group {
298                factor,
299                placed,
300                summed,
301            }
302        })
303        .collect();
304    let grouped = LanePlan::Grouped { groups, bins };
305    match grouped.flops(len) <= sweep.flops(len) {
306        true => grouped,
307        false => sweep,
308    }
309}
310
311/// `n*log2(n)` butterflies; a length that is no power of two is Bluestein's three
312/// transforms over the next one past `2n-1`.
313pub fn transform_flops(n: usize) -> u128 {
314    let stages = |m: usize| m as u128 * (m.max(2).trailing_zeros() as u128);
315    match n.is_power_of_two() {
316        true => stages(n),
317        false => 3 * stages((2 * n - 1).next_power_of_two()),
318    }
319}
320
321impl LinePlan {
322    pub fn flops(&self, len: usize) -> u128 {
323        let transforms: u128 = self
324            .placed
325            .iter()
326            .filter(|p| !p.is_empty())
327            .map(|_| transform_flops(self.bins))
328            .sum();
329        let direct: u128 = self
330            .summed
331            .iter()
332            .map(|s| s.len() as u128 * len as u128)
333            .sum();
334        transforms + direct
335    }
336
337    pub fn rule(&self) -> Rule {
338        match (lines::distinct(&self.placed), lines::distinct(&self.summed)) {
339            (_, 0) => Rule::LineSpectrumExact,
340            (0, _) => Rule::LineSpectrumSummed,
341            _ => Rule::LineSpectrumMixed,
342        }
343    }
344}
345
346impl LanePlan {
347    pub fn flops(&self, len: usize) -> u128 {
348        match self {
349            LanePlan::Grouped { groups, bins } => groups
350                .iter()
351                .map(|g| {
352                    let placed = match g.placed.is_empty() {
353                        true => 0,
354                        false => transform_flops(*bins),
355                    };
356                    placed + (g.summed.len() as u128 + 1) * len as u128
357                })
358                .sum(),
359            LanePlan::Sweep { atoms, samples } => *atoms as u128 * *samples as u128,
360        }
361    }
362
363    fn swept_flops(&self, len: u128) -> u128 {
364        let terms: u128 = match self {
365            LanePlan::Grouped { groups, .. } => groups
366                .iter()
367                .map(|g| (g.placed.len() + g.summed.len() + 1) as u128)
368                .sum(),
369            LanePlan::Sweep { atoms, .. } => *atoms as u128,
370        };
371        terms * len
372    }
373}
374
375impl Plan {
376    pub fn flops(&self, len: usize) -> u128 {
377        match self {
378            Plan::Spectrum(_) => transform_flops(len),
379            Plan::Lines(found) => found.flops(len),
380            Plan::Sampled(held) => held.lanes.iter().map(|lane| lane.flops(len)).sum(),
381            Plan::Point { nodes, .. } => *nodes as u128 * len as u128,
382            Plan::Added(parts) => parts.iter().map(|part| part.flops(len)).sum(),
383        }
384    }
385
386    pub fn alias_flops(&self, len: usize) -> u128 {
387        self.alias_flops_at(len, super::ALIAS_OVERSAMPLE)
388    }
389
390    /// The multiple a reading names need not be the label's `ALIAS_OVERSAMPLE`. Every component
391    /// is scored, so every component's reference is paid for.
392    pub fn alias_flops_at(&self, len: usize, oversample: usize) -> u128 {
393        let reference = (len * oversample) as u128;
394        match self {
395            Plan::Point { nodes, .. } => *nodes as u128 * reference,
396            Plan::Sampled(held) if held.rule == Rule::PointSampled => held
397                .lanes
398                .iter()
399                .map(|lane| lane.swept_flops(reference))
400                .sum(),
401            Plan::Added(parts) => parts
402                .iter()
403                .map(|part| part.alias_flops_at(len, oversample))
404                .sum(),
405            Plan::Spectrum(_) | Plan::Lines(_) | Plan::Sampled(_) => 0,
406        }
407    }
408
409    pub fn rule(&self) -> Rule {
410        match self {
411            Plan::Spectrum(_) => Rule::InverseSpectrum,
412            Plan::Lines(found) => found.rule(),
413            Plan::Sampled(held) => held.rule,
414            Plan::Point { .. } => Rule::PointSampled,
415            Plan::Added(_) => Rule::Added,
416        }
417    }
418}