census.rs

1use crate::factor::mobius;23// FAMILIES45pub struct Family {6    pub q: u64,7    pub digits: Vec<u64>,8    pub label: String,9    pub lmax: usize,10    pub children: Vec<u64>,11}1213pub struct Outcome {14    pub counts: Vec<u64>,15    pub meter: Vec<i64>,16    pub mmax: Vec<u64>,17    pub twisted: Vec<(u64, Vec<i64>)>,18}1920fn depth(q: u64, k: usize) -> usize {21    match (q, k) {22        (3, 2) => 24,23        (4, 2) => 22,24        (4, 3) => 14,25        (5, 2) => 21,26        (5, 3) => 13,27        (5, 4) => 11,28        _ => panic!("no depth for q={q} k={k}"),29    }30}3132fn gcd(mut a: u64, mut b: u64) -> u64 {33    while b != 0 {34        let t = a % b;35        a = b;36        b = t;37    }38    a39}4041pub fn digit_gcd(digits: &[u64]) -> u64 {42    digits.iter().fold(0, |g, &d| gcd(g, d))43}4445pub fn make_family(q: u64, digits: Vec<u64>) -> Family {46    let lmax = depth(q, digits.len());47    make_family_depth(q, digits, lmax)48}4950pub fn make_family_depth(q: u64, digits: Vec<u64>, lmax: usize) -> Family {51    let label: String = digits.iter().map(|d| d.to_string()).collect();52    let mut children = Vec::new();53    if digit_gcd(&digits) == 1 {54        let top = *digits.iter().max().unwrap();55        let mut a = 2;56        while a * top <= q - 1 {57            children.push(a);58            a += 1;59        }60    }61    Family {62        q,63        digits,64        label,65        lmax,66        children,67    }68}6970pub fn families() -> Vec<Family> {71    let mut out = Vec::new();72    for q in [3u64, 4, 5] {73        for k in 2..=(q as usize - 1) {74            for mask in 0u64..(1 << q) {75                if mask.count_ones() as usize != k {76                    continue;77                }78                let digits: Vec<u64> = (0..q).filter(|d| mask >> d & 1 == 1).collect();79                out.push(make_family(q, digits));80            }81        }82    }83    out84}8586// SWEEP8788pub fn sweep(q: u64, digits: &[u64], lmax: usize, visit: &mut impl FnMut(u64, usize)) {89    fn rec(90        v: u64,91        len: usize,92        q: u64,93        digits: &[u64],94        lmax: usize,95        visit: &mut impl FnMut(u64, usize),96    ) {97        visit(v, len);98        if len < lmax {99            for &d in digits {100                rec(v * q + d, len + 1, q, digits, lmax, visit);101            }102        }103    }104    for &d in digits {105        if d != 0 {106            rec(d, 1, q, digits, lmax, visit);107        }108    }109}110111fn twist(a: u64, v: u64, mu_v: i8) -> i64 {112    match a {113        2 => {114            if v % 2 == 0 {115                0116            } else {117                -(mu_v as i64)118            }119        }120        3 => {121            if v % 3 == 0 {122                0123            } else {124                -(mu_v as i64)125            }126        }127        4 => 0,128        _ => panic!("no twist for a={a}"),129    }130}131132// ORDERED SWEEP133134pub fn sweep_length(q: u64, digits: &[u64], length: usize, visit: &mut impl FnMut(u64)) {135    fn rec(v: u64, len: usize, q: u64, digits: &[u64], length: usize, visit: &mut impl FnMut(u64)) {136        if len == length {137            visit(v);138            return;139        }140        for &d in digits {141            rec(v * q + d, len + 1, q, digits, length, visit);142        }143    }144    for &d in digits {145        if d != 0 {146            rec(d, 1, q, digits, length, visit);147        }148    }149}150151// METER152153pub fn run_family(fam: &Family, primes: &[u64]) -> Outcome {154    let l = fam.lmax;155    let has01 = fam.digits.contains(&0) && fam.digits.contains(&1);156    let mut counts = vec![0u64; l + 1];157    let mut meter = vec![0i64; l + 1];158    let mut mmax = vec![0u64; l + 1];159    let mut twisted: Vec<(u64, Vec<i64>)> = fam160        .children161        .iter()162        .map(|&a| (a, vec![0i64; l + 1]))163        .collect();164    let mut run = 0i64;165    let mut peak = 0u64;166    let mut total = 0u64;167    let mut boundary = 1u64;168    for lev in 1..=l {169        let skip = boundary;170        boundary *= fam.q;171        let mut tw_len: Vec<i64> = vec![0; fam.children.len()];172        sweep_length(fam.q, &fam.digits, lev, &mut |v| {173            let mu = mobius(v, primes);174            for (slot, a) in tw_len.iter_mut().zip(fam.children.iter()) {175                *slot += twist(*a, v, mu);176            }177            if has01 && lev > 1 && v == skip {178                return;179            }180            total += 1;181            run += mu as i64;182            peak = peak.max(run.unsigned_abs());183        });184        if has01 {185            total += 1;186            run += mobius(boundary, primes) as i64;187            peak = peak.max(run.unsigned_abs());188        }189        counts[lev] = total;190        meter[lev] = run;191        mmax[lev] = peak;192        for (slot, tl) in twisted.iter_mut().zip(tw_len.iter()) {193            slot.1[lev] = slot.1[lev - 1] + tl;194        }195    }196    Outcome {197        counts,198        meter,199        mmax,200        twisted,201    }202}