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}