puzlink 0.1.0

package/src/lib/logCounter.ts

@@ -0,0 +1,74 @@
+import { LogNum } from "./logNum.js";
+
+/** A map from items to log counts. */
+export class LogCounter<T extends PropertyKey> {
+  private readonly counts: ReadonlyMap<T, LogNum>;
+
+  constructor(counts: ReadonlyMap<T, LogNum>, totalCache?: LogNum) {
+    this.counts = counts;
+    this.totalCache = totalCache;
+  }
+
+  static from(data: string): LogCounter<string>;
+  static from<T extends PropertyKey>(data: Iterable<T>): LogCounter<T>;
+  static from<T extends PropertyKey, U>(
+    data: Iterable<U>,
+    map: (item: U) => T,
+  ): LogCounter<T>;
+  static from(
+    data: string | Iterable<PropertyKey>,
+    map?: (item: PropertyKey) => PropertyKey,
+  ) {
+    const counts = new Map<PropertyKey, number>();
+    let total = 0;
+
+    for (const item_ of data) {
+      const item = map ? map(item_) : item_;
+      counts.set(item, (counts.get(item) ?? 0) + 1);
+      total++;
+    }
+
+    return new LogCounter(
+      new Map(
+        Array.from(counts).map(([item, count]) => [item, LogNum.from(count)]),
+      ),
+      LogNum.from(total),
+    );
+  }
+
+  /** The number of distinct items. */
+  get distinct(): LogNum {
+    return LogNum.from(this.counts.size);
+  }
+
+  private totalCache: LogNum | undefined;
+
+  /** The total number of all items. */
+  get total(): LogNum {
+    return (this.totalCache ??= LogNum.sum(Array.from(this.counts.values())));
+  }
+
+  /** The log count of the given item. */
+  get(item: T): LogNum {
+    return this.counts.get(item) ?? LogNum.from(0);
+  }
+
+  /** Returns an iterable of [item, log count] pairs. */
+  entries(): IterableIterator<[T, LogNum]> {
+    return this.counts.entries();
+  }
+
+  /** Returns a list of items that satisfy the given predicate. */
+  filterKeys(fn: (item: T, count: LogNum) => boolean): T[] {
+    return Array.from(this.counts.entries())
+      .filter(([item, count]) => fn(item, count))
+      .map(([item]) => item);
+  }
+
+  /** Returns an iterable of [item, log probability] pairs. */
+  *frequencies(): IterableIterator<[T, LogNum]> {
+    for (const [item, count] of this.counts) {
+      yield [item, count.div(this.total)];
+    }
+  }
+}

package/src/lib/logNum.ts

@@ -0,0 +1,193 @@
+import { zipfToLogProb } from "cromulence";
+import { gammaln } from "simple-statistics";
+import { memoize } from "./memoize.js";
+import { interval } from "./util.js";
+
+/**
+ * A numerically-stable-ish log(1 - exp(x)).
+ *
+ * Follows the algorithm in <https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf>.
+ */
+function log1mExp(x: number): number {
+  return x > -Math.log(2) ? Math.log1p(-Math.exp(x)) : Math.log(-Math.expm1(x));
+}
+
+/**
+ * A numerically-stable-ish log(1 + exp(x)).
+ *
+ * Follows the algorithm in <https://cran.r-project.org/web/packages/Rmpfr/vignettes/log1mexp-note.pdf>.
+ */
+function log1pExp(x: number): number {
+  return x <= 18 ? Math.log1p(Math.exp(x)) : x + Math.exp(-x);
+}
+
+/**
+ * A numerically-stable-ish log of a non-negative number. Operations represent
+ * operations on the numbers, not the log numbers.
+ *
+ * @example LogNum.from(2).add(LogNum.from(3)).toNum() // 5
+ */
+export class LogNum {
+  private readonly data: number;
+
+  /** Construct from the actual log number. */
+  constructor(data: number) {
+    this.data = data;
+  }
+
+  static from(value: number): LogNum {
+    return new LogNum(Math.log(value));
+  }
+
+  static fromExp(value: number): LogNum {
+    return new LogNum(value);
+  }
+
+  static fromFraction(numerator: number, denominator: number): LogNum {
+    return new LogNum(Math.log(numerator) - Math.log(denominator));
+  }
+
+  @memoize(2)
+  static fromBinomial(n: number, k: number): LogNum {
+    if (n === 0 || k === 0 || k === n) {
+      return LogNum.from(1);
+    }
+
+    return LogNum.fromBinomial(n - 1, k - 1).add(LogNum.fromBinomial(n - 1, k));
+  }
+
+  static fromZipf(zipf: number): LogNum {
+    return new LogNum(zipfToLogProb(zipf));
+  }
+
+  static fromFactorial(n: number): LogNum {
+    return new LogNum(gammaln(n + 1));
+  }
+
+  toNum(): number {
+    return Math.exp(this.data);
+  }
+
+  toLog(): number {
+    return this.data;
+  }
+
+  exp(): LogNum {
+    return new LogNum(Math.exp(this.data));
+  }
+
+  log(): LogNum {
+    return LogNum.from(this.data);
+  }
+
+  pow(power: number): LogNum {
+    return new LogNum(this.data * power);
+  }
+
+  mul(other: LogNum): LogNum {
+    return new LogNum(this.data + other.data);
+  }
+
+  div(other: LogNum): LogNum {
+    return new LogNum(this.data - other.data);
+  }
+
+  add(other: LogNum): LogNum {
+    let [max, min] = [this.data, other.data];
+    if (max < min) {
+      [max, min] = [min, max];
+    }
+    if (min === -Infinity) {
+      return new LogNum(max);
+    }
+    return new LogNum(max + log1pExp(min - max));
+  }
+
+  absSub(other: LogNum): LogNum {
+    let [max, min] = [this.data, other.data];
+    if (max < min) {
+      [max, min] = [min, max];
+    }
+    return new LogNum(max + log1mExp(min - max));
+  }
+
+  sub(other: LogNum): LogNum {
+    if (this.lt(other)) {
+      throw new Error(
+        `log underflow: ${this.data.toString()} - ${other.data.toString()}`,
+      );
+    }
+    const [max, min] = [this.data, other.data];
+    return new LogNum(max + log1mExp(min - max));
+  }
+
+  gt(other: LogNum): boolean {
+    return this.data > other.data;
+  }
+
+  lt(other: LogNum): boolean {
+    return this.data < other.data;
+  }
+
+  static max(values: LogNum[]): LogNum {
+    return new LogNum(Math.max(...values.map((x) => x.data)));
+  }
+
+  static min(values: LogNum[]): LogNum {
+    return new LogNum(Math.min(...values.map((x) => x.data)));
+  }
+
+  static sum(values: LogNum[]): LogNum {
+    // Strip zeroes first:
+    values = values.filter((x) => x.data !== -Infinity);
+    const max = Math.max(...values.map((x) => x.data));
+    if (max === Infinity) {
+      return new LogNum(Infinity);
+    }
+    const expSum = values.reduce((acc, x) => acc + Math.exp(x.data - max), 0);
+    return new LogNum(max + Math.log(expSum));
+  }
+
+  static prod(values: LogNum[]): LogNum {
+    return new LogNum(values.reduce((acc, x) => acc + x.data, 0));
+  }
+
+  static binomialProb(
+    successes: number,
+    trials: number,
+    frequency: LogNum,
+  ): LogNum {
+    return LogNum.fromBinomial(trials, successes).mul(
+      frequency.pow(successes).mul(
+        LogNum.from(1)
+          .sub(frequency)
+          .pow(trials - successes),
+      ),
+    );
+  }
+
+  static binomialPValue(
+    successes: number,
+    trials: number,
+    frequency: LogNum,
+  ): LogNum {
+    const expected = trials * frequency.toNum();
+    const probs = [];
+
+    if (successes > expected) {
+      for (const i of interval(successes, trials)) {
+        probs.push(LogNum.binomialProb(i, trials, frequency));
+      }
+    } else {
+      for (const i of interval(0, successes)) {
+        probs.push(LogNum.binomialProb(i, trials, frequency));
+      }
+    }
+
+    return LogNum.sum(probs);
+  }
+
+  [Symbol.for("nodejs.util.inspect.custom")](): string {
+    return `LogNum(${this.data.toFixed(3)}: ${this.toNum().toString()})`;
+  }
+}

package/src/lib/memoize.ts

@@ -0,0 +1,136 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+
+type AnyFn = (...args: readonly any[]) => unknown;
+
+type Cache<Args extends readonly any[], Return> = {
+  get: (key: Args) => Return | null;
+  set: (key: Args, value: Return) => void;
+};
+
+function memoizeWith<Fn extends AnyFn>(
+  fn: Fn,
+  makeCache: () => Cache<Parameters<Fn>, ReturnType<Fn>>,
+): Fn {
+  const cache = makeCache();
+
+  const memoized = function (
+    this: any,
+    ...args: Parameters<Fn>
+  ): ReturnType<Fn> {
+    const cached = cache.get(args);
+    if (cached !== null) {
+      return cached;
+    }
+    const result = fn.apply(this, args) as ReturnType<Fn>;
+    cache.set(args, result);
+    return result;
+  } as Fn;
+
+  for (const property of Reflect.ownKeys(fn)) {
+    if (
+      property === "length" ||
+      property === "prototype" ||
+      property === "argunments" ||
+      property === "prototype"
+    ) {
+      continue;
+    }
+    const fnDesc = Object.getOwnPropertyDescriptor(fn, property)!;
+    const mmDesc = Object.getOwnPropertyDescriptor(memoized, property);
+    if (
+      mmDesc === undefined ||
+      mmDesc.configurable ||
+      (mmDesc.writable === fnDesc.writable &&
+        mmDesc.enumerable === fnDesc.enumerable &&
+        mmDesc.configurable === fnDesc.configurable &&
+        (mmDesc.writable || mmDesc.value === fnDesc.value))
+    ) {
+      Object.defineProperty(memoized, property, fnDesc);
+    }
+  }
+
+  Object.setPrototypeOf(memoized, Object.getPrototypeOf(fn) as object);
+
+  return memoized;
+}
+
+function memoizeDecorator<Fn extends AnyFn>(
+  makeCache: () => Cache<Parameters<Fn>, ReturnType<Fn>>,
+) {
+  return (target: Fn): Fn => {
+    const instanceMap = new WeakMap<any, Fn>();
+    return function (this: any, ...args: Parameters<Fn>): ReturnType<Fn> {
+      let memoized = instanceMap.get(this);
+      if (memoized === undefined) {
+        memoized = memoizeWith(target, makeCache);
+        instanceMap.set(this, memoized);
+      }
+      return memoized.apply(this, args) as ReturnType<Fn>;
+    } as Fn;
+  };
+}
+
+function memoize1<Fn extends (Arg: any) => unknown>() {
+  return memoizeDecorator<Fn>(() => {
+    const cache = new Map<any, unknown>();
+    return {
+      get: ([key]) => {
+        return cache.get(key) ?? null;
+      },
+      set: ([key], value) => {
+        return cache.set(key, value);
+      },
+    } as Cache<Parameters<Fn>, ReturnType<Fn>>;
+  });
+}
+
+function memoize2<Fn extends (Arg1: any, Arg2: any) => unknown>() {
+  return memoizeDecorator<Fn>(() => {
+    const cache = new Map<any, Map<any, unknown>>();
+    return {
+      get: ([outerKey, innerKey]) => {
+        return cache.get(outerKey)?.get(innerKey) ?? null;
+      },
+      set: ([outerKey, innerKey], value) => {
+        let inner = cache.get(outerKey);
+        if (inner === undefined) {
+          inner = new Map();
+          cache.set(outerKey, inner);
+        }
+        inner.set(innerKey, value);
+      },
+    } as Cache<Parameters<Fn>, ReturnType<Fn>>;
+  });
+}
+
+function memoize3<Fn extends (Arg1: any, Arg2: any, Arg3: any) => unknown>() {
+  return memoizeDecorator<Fn>(() => {
+    const cache = new Map<any, Map<any, Map<any, unknown>>>();
+    return {
+      get: ([outerKey, middleKey, innerKey]) => {
+        return cache.get(outerKey)?.get(middleKey)?.get(innerKey) ?? null;
+      },
+      set: ([outerKey, middleKey, innerKey], value) => {
+        let middle = cache.get(outerKey);
+        if (middle === undefined) {
+          middle = new Map();
+          cache.set(outerKey, middle);
+        }
+        let inner = middle.get(middleKey);
+        if (inner === undefined) {
+          inner = new Map();
+          middle.set(middleKey, inner);
+        }
+        inner.set(innerKey, value);
+      },
+    } as Cache<Parameters<Fn>, ReturnType<Fn>>;
+  });
+}
+
+/** Memoize a class method that takes up to three arguments. */
+export function memoize<const Fn extends AnyFn>(
+  levels?: 1 | 2 | 3,
+): (target: Fn) => Fn;
+export function memoize(levels?: 1 | 2 | 3) {
+  return levels === 3 ? memoize3() : levels === 2 ? memoize2() : memoize1();
+}

package/src/lib/testUtils.ts

@@ -0,0 +1 @@
+export const speedTest = process.env["SPEED_TEST"] === "true";

package/src/lib/util.ts

@@ -0,0 +1,150 @@
+/* eslint-disable @typescript-eslint/no-explicit-any */
+
+/** Capitalizes the letters of a slug at the given indices. */
+export function capitalizeAt(slug: string, indices: number[]): string {
+  const capitalized = [];
+  for (let i = 0; i < slug.length; i++) {
+    capitalized.push(indices.includes(i) ? slug[i]!.toUpperCase() : slug[i]!);
+  }
+  return capitalized.join("");
+}
+
+/** Prints the index of a slug at the given indices. */
+export function printIndexSlug(slug: string, indices: number[]): string {
+  indices = Array.from(new Set(indices)).sort((a, b) => a - b);
+  const isInterval =
+    indices.length > 2 &&
+    Math.max(...indices) - Math.min(...indices) === indices.length - 1;
+  const indexString = isInterval
+    ? `${Math.min(...indices).toString()}..${Math.max(...indices).toString()}`
+    : `${indices.slice(0, 5).join(", ")}${indices.length > 5 ? ", ..." : ""}`;
+  return `index(${slug}, ${indexString}) = ${capitalizeAt(slug, indices)}`;
+}
+
+/** Returns an array of numbers from start to end (inclusive). */
+export function interval(start: number, end: number, step = 1): number[] {
+  const result = [];
+  for (let i = start; i <= end; i += step) {
+    result.push(i);
+  }
+  return result;
+}
+
+type Product<A extends Iterable<any>[]> = A extends [infer First, ...infer Rest]
+  ? First extends Iterable<infer F>
+    ? Rest extends Iterable<any>[]
+      ? [F, ...Product<Rest>]
+      : never
+    : never
+  : [];
+
+/** Cartesian product of iterables. */
+function* product<const Args extends Iterable<any>[]>(
+  ...args: Args
+): Generator<Product<Args>> {
+  if (args.length === 1) {
+    for (const arg of args[0]!) {
+      yield [arg] as Product<Args>;
+    }
+    return;
+  }
+  const [first, ...rest] = args;
+  for (const a of first!) {
+    for (const b of product(...rest)) {
+      yield [a, ...b] as Product<Args>;
+    }
+  }
+}
+
+type IterMap<A extends any[]> = A extends [infer First, ...infer Rest]
+  ? [Iterable<First>, ...IterMap<Rest>]
+  : [];
+
+/** Map a function to a product of iterators. */
+export function* mapProduct<const Args extends any[], R>(
+  fn: (...args: Args) => R,
+  ...args: IterMap<Args>
+): Generator<R> {
+  for (const arg of product(...args)) {
+    yield fn(...(arg as unknown as Args));
+  }
+}
+
+/** Returns the ordinal suffix of the given number. */
+export function ordinal(n: number): string {
+  const suffix =
+    Math.abs(n) % 100 >= 11 && Math.abs(n) % 100 <= 14
+      ? "th"
+      : Math.abs(n) % 10 === 1
+        ? "st"
+        : Math.abs(n) % 10 === 2
+          ? "nd"
+          : Math.abs(n) % 10 === 3
+            ? "rd"
+            : "th";
+  return `${n.toString()}${suffix}`;
+}
+
+const aCharCode = "a".charCodeAt(0);
+
+/** Caesar shift a slug. */
+export function caesar(slug: string, n: number): string {
+  n %= 26;
+  return Array.from(slug, (c) => {
+    return String.fromCharCode(
+      ((c.charCodeAt(0) - aCharCode + n + 26) % 26) + aCharCode,
+    );
+  }).join("");
+}
+
+/** Returns an iterator over [index, item] pairs. */
+export function* enumerate<T>(iter: Iterable<T>): Generator<[number, T]> {
+  let i = 0;
+  for (const item of iter) {
+    yield [i, item];
+    i += 1;
+  }
+}
+
+/** Returns an iterator over windows of the given size. */
+export function windows<T>(iter: Iterable<T>, size: 2): Generator<[T, T]>;
+export function windows<T>(iter: Iterable<T>, size: 3): Generator<[T, T, T]>;
+export function windows<T>(iter: Iterable<T>, size: number): Generator<T[]>;
+export function* windows(iter: Iterable<any>, size: number) {
+  const buffer = [];
+  for (const item of iter) {
+    if (buffer.length === size) {
+      buffer.shift();
+    }
+    buffer.push(item);
+    if (buffer.length === size) {
+      yield buffer.slice();
+    }
+  }
+}
+
+/**
+ * If the given array is a non-constant arithmetic sequence, returns the start,
+ * step, and last. Else, returns null.
+ */
+export function getArithmeticSequenceInfo(terms: number[]): {
+  start: number;
+  step: number;
+  last: number;
+} | null {
+  if (terms.length < 2) {
+    return null;
+  }
+  terms.sort((a, b) => a - b);
+  const start = terms[0]!;
+  const step = terms[1]! - start;
+  if (step === 0) {
+    return null;
+  }
+  for (let i = 2; i < terms.length; i++) {
+    if (terms[i]! - terms[i - 1]! !== step) {
+      return null;
+    }
+  }
+  return { start, step, last: terms.at(-1)! };
+}