terms.rs

12.7 kB · rust · 425 lines

1use crate::design::Design;2use mrlynum::classics::primes;3use mrlynum::factor::{gcd, mobius_sieve};4use std::sync::atomic::{AtomicU64, Ordering};5use std::time::Instant;67const CHUNK: u32 = 9;8const CHUNK_SPAN: u64 = 19683;9const FINE: u64 = 1 << 16;10const BLOCK: u64 = 1 << 13;1112fn digit_table(invert: bool) -> Vec<u32> {13    (0..CHUNK_SPAN)14        .map(|value| {15            let mut mask = 0u32;16            let mut rest = value;17            for place in 0..CHUNK {18                if (rest % 3 == 1) != invert {19                    mask |= 1 << place;20                }21                rest /= 3;22            }23            mask24        })25        .collect()26}2728fn mask_of(mut x: u64, table: &[u32], chunks: u32, full: u32) -> u32 {29    let mut mask = 0u32;30    for chunk in 0..chunks {31        mask |= table[(x % CHUNK_SPAN) as usize] << (chunk * CHUNK);32        x /= CHUNK_SPAN;33    }34    mask & full35}3637fn mobius_block(lo: u64, hi: u64, base: &[usize], sign: &mut Vec<i8>, rest: &mut Vec<u64>) {38    let width = (hi - lo) as usize;39    sign.clear();40    sign.resize(width, 1);41    rest.clear();42    rest.extend(lo..hi);43    for &p in base {44        let p = p as u64;45        if p * p >= hi {46            break;47        }48        let mut m = lo.div_ceil(p) * p;49        while m < hi {50            let slot = (m - lo) as usize;51            sign[slot] = -sign[slot];52            rest[slot] /= p;53            m += p;54        }55        let mut m = lo.div_ceil(p * p) * p * p;56        while m < hi {57            sign[(m - lo) as usize] = 0;58            m += p * p;59        }60    }61    for slot in 0..width {62        if sign[slot] != 0 && rest[slot] > 1 {63            sign[slot] = -sign[slot];64        }65    }66}6768fn subset_sums(buf: &mut [u64], level: u32) {69    for place in 0..level {70        let bit = 1usize << place;71        for base in (0..buf.len()).step_by(bit << 1) {72            let (low, high) = buf[base..base + (bit << 1)].split_at_mut(bit);73            for index in 0..bit {74                high[index] += low[index];75            }76        }77    }78}7980struct Lattice {81    level: u32,82    size: usize,83    full: u32,84    chunks: u32,85    dimension: usize,86    count: Vec<u32>,87    items: Vec<(u32, u32)>,88    sub: Vec<(u32, u32)>,89    sums: Vec<u64>,90    ranked: Vec<u64>,91    choose: Vec<Vec<u64>>,92}9394impl Lattice {95    fn new(level: u32, dimension: usize) -> Lattice {96        let size = 1usize << level;97        let ranks = level as usize + 1;98        let mut choose = vec![vec![0u64; ranks]; ranks];99        for row in 0..ranks {100            choose[row][0] = 1;101            for column in 1..=row {102                choose[row][column] = choose[row - 1][column - 1] + choose[row - 1][column];103            }104        }105        Lattice {106            level,107            size,108            full: (1u32 << level) - 1,109            chunks: level.div_ceil(CHUNK).max(1),110            dimension,111            count: vec![0; size],112            items: Vec::new(),113            sub: Vec::new(),114            sums: vec![0; size],115            ranked: if dimension == 3 {116                vec![0; ranks * size]117            } else {118                Vec::new()119            },120            choose,121        }122    }123124    fn gather(&mut self, modulus: u64, span: u64, table: &[u32]) {125        self.items.clear();126        for x in (0..span).step_by(modulus as usize) {127            let mask = mask_of(x, table, self.chunks, self.full) as usize;128            if self.count[mask] == 0 {129                self.items.push((mask as u32, 0));130            }131            self.count[mask] += 1;132        }133        for item in self.items.iter_mut() {134            item.1 = self.count[item.0 as usize];135        }136    }137138    fn release(&mut self) {139        for &(mask, _) in self.items.iter() {140            self.count[mask as usize] = 0;141        }142    }143144    fn spread(&mut self) {145        for &(mask, weight) in self.items.iter() {146            self.sums[mask as usize] = weight as u64;147        }148        subset_sums(&mut self.sums, self.level);149    }150151    fn direct2(&self) -> u128 {152        let mut total = 0u128;153        for &(a, ca) in self.items.iter() {154            let inner: u64 = self155                .items156                .iter()157                .filter(|(b, _)| a & b == 0)158                .map(|(_, cb)| *cb as u64)159                .sum();160            total += ca as u128 * inner as u128;161        }162        total163    }164165    fn zeta2(&mut self) -> u128 {166        self.spread();167        let mut total = 0u128;168        for &(a, ca) in self.items.iter() {169            total += ca as u128 * self.sums[(!a & self.full) as usize] as u128;170        }171        self.sums.fill(0);172        total173    }174175    fn direct3(&mut self) -> u128 {176        let mut total = 0u128;177        for &(a, ca) in self.items.iter() {178            self.sub.clear();179            self.sub180                .extend(self.items.iter().filter(|(b, _)| a & b == 0));181            let mut inner = 0u128;182            for &(b, cb) in self.sub.iter() {183                let reach: u64 = self184                    .sub185                    .iter()186                    .filter(|(c, _)| b & c == 0)187                    .map(|(_, cc)| *cc as u64)188                    .sum();189                inner += cb as u128 * reach as u128;190            }191            total += ca as u128 * inner;192        }193        total194    }195196    fn zeta3(&mut self) -> u128 {197        self.spread();198        let width = self.items.len() as u64;199        let mut total = 0u128;200        for &(a, ca) in self.items.iter() {201            let free = !a & self.full;202            let mut inner = 0u64;203            if 1u64 << free.count_ones() <= width {204                let mut b = free;205                loop {206                    inner += self.count[b as usize] as u64 * self.sums[(free ^ b) as usize];207                    if b == 0 {208                        break;209                    }210                    b = (b - 1) & free;211                }212            } else {213                for &(b, cb) in self.items.iter() {214                    if a & b == 0 {215                        inner += cb as u64 * self.sums[(free & !b) as usize];216                    }217                }218            }219            total += ca as u128 * inner as u128;220        }221        self.sums.fill(0);222        total223    }224225    fn convolve3(&mut self) -> u128 {226        let top = self.level as usize;227        let size = self.size;228        self.ranked.fill(0);229        for &(mask, weight) in self.items.iter() {230            let rank = mask.count_ones() as usize;231            self.ranked[rank * size + mask as usize] = weight as u64;232        }233        for rank in 0..=top {234            subset_sums(&mut self.ranked[rank * size..(rank + 1) * size], self.level);235        }236        let mut poly = vec![0u128; top + 1];237        let mut square = vec![0u128; top + 1];238        let mut total = 0i128;239        for set in 0..size {240            let rank = (set as u32).count_ones() as usize;241            for degree in 0..=rank {242                poly[degree] = self.ranked[degree * size + set] as u128;243            }244            let reach = (2 * rank).min(top);245            square[..=reach].fill(0);246            for left in 0..=rank.min(reach) {247                for right in 0..=rank.min(reach - left) {248                    square[left + right] += poly[left] * poly[right];249                }250            }251            let mut acc = 0i128;252            for degree in rank..=(3 * rank).min(top) {253                let mut cube = 0u128;254                for left in degree.saturating_sub(rank)..=reach.min(degree) {255                    cube += square[left] * poly[degree - left];256                }257                let weight = self.choose[top - rank][degree - rank] as i128 * cube as i128;258                if (degree - rank) % 2 == 0 {259                    acc += weight;260                } else {261                    acc -= weight;262                }263            }264            total += acc;265        }266        assert!(total >= 0);267        total as u128268    }269270    fn count(&mut self, modulus: u64, span: u64, table: &[u32]) -> u128 {271        self.gather(modulus, span, table);272        let width = self.items.len();273        let sweep = self.level as usize * self.size;274        let found = if self.dimension == 2 {275            if width * width <= sweep {276                self.direct2()277            } else {278                self.zeta2()279            }280        } else {281            let direct = width * width * width / 40;282            let pair = sweep + width * width;283            let ranked = 2 * self.level as usize * sweep;284            if direct <= pair.min(ranked) {285                self.direct3()286            } else if pair <= ranked {287                self.zeta3()288            } else {289                self.convolve3()290            }291        };292        self.release();293        found294    }295}296297fn weight(298    design: &Design,299    level: u32,300    threads: usize,301    table: &[u32],302    small: &[i8],303    base: &[usize],304) -> i128 {305    let span = 3u64.pow(level);306    let fine = span.min(FINE);307    let origin: i128 = if design.origin_filled() { 1 } else { 0 };308    let cursor = AtomicU64::new(0);309    std::thread::scope(|scope| {310        let handles: Vec<_> = (0..threads)311            .map(|_| {312                let cursor = &cursor;313                scope.spawn(move || {314                    let mut lattice = Lattice::new(level, design.dimension);315                    let mut sign = Vec::new();316                    let mut rest = Vec::new();317                    let mut acc = 0i128;318                    loop {319                        let task = cursor.fetch_add(1, Ordering::Relaxed);320                        let (lo, hi) = if task < fine {321                            (task + 1, task + 2)322                        } else {323                            let lo = fine + 1 + (task - fine) * BLOCK;324                            (lo, lo + BLOCK)325                        };326                        if lo >= span {327                            break;328                        }329                        let hi = hi.min(span);330                        if task >= fine {331                            mobius_block(lo, hi, base, &mut sign, &mut rest);332                        }333                        for modulus in lo..hi {334                            let mu = if task < fine {335                                small[modulus as usize]336                            } else {337                                sign[(modulus - lo) as usize]338                            };339                            if mu == 0 || modulus % 3 == 0 {340                                continue;341                            }342                            let found = lattice.count(modulus, span, table) as i128;343                            acc += mu as i128 * (found - origin);344                        }345                    }346                    acc347                })348            })349            .collect();350        handles.into_iter().map(|h| h.join().unwrap()).sum()351    })352}353354pub struct Term {355    pub level: u32,356    pub value: i128,357    pub seconds: f64,358}359360pub fn terms(design: &Design, top: u32, threads: usize) -> Vec<Term> {361    let table = digit_table(design.invert);362    let small = mobius_sieve(FINE as usize);363    let base = primes((3f64.powi(top as i32)).sqrt() as usize + 2);364    let mut previous = 0i128;365    let mut out = Vec::new();366    for level in 1..=top {367        let clock = Instant::now();368        let current = weight(design, level, threads, &table, &small, &base);369        let value = if design.origin_filled() {370            current - previous371        } else {372            current373        };374        previous = current;375        out.push(Term {376            level,377            value,378            seconds: clock.elapsed().as_secs_f64(),379        });380    }381    out382}383384fn walk(depth: u32, coords: &mut [u64], corners: &[Vec<u64>], found: &mut u64) {385    if depth == 0 {386        let common = coords.iter().fold(0usize, |g, &c| gcd(g, c as usize));387        if common == 1 {388            *found += 1;389        }390        return;391    }392    for corner in corners {393        for (slot, digit) in coords.iter_mut().zip(corner) {394            *slot = *slot * 3 + digit;395        }396        walk(depth - 1, coords, corners, found);397        for (slot, digit) in coords.iter_mut().zip(corner) {398            *slot = (*slot - digit) / 3;399        }400    }401}402403pub fn brute(design: &Design, level: u32) -> u64 {404    let corners = design.corners();405    let mut coords = vec![0u64; design.dimension];406    let mut found = 0;407    walk(level, &mut coords, &corners, &mut found);408    found409}410411pub fn stored(design: &Design) -> Vec<(u32, i128)> {412    let path = format!(413        "{}/terms/{}_bfile.txt",414        env!("CARGO_MANIFEST_DIR"),415        design.name416    );417    std::fs::read_to_string(path)418        .unwrap_or_default()419        .lines()420        .filter_map(|line| {421            let mut parts = line.split_whitespace();422            Some((parts.next()?.parse().ok()?, parts.next()?.parse().ok()?))423        })424        .collect()425}