@gmod/gbz-base 0.0.1 → 1.0.0

package/src/gbwt/node.ts

@@ -0,0 +1,64 @@
+export const ENDMARKER = 0
+
+export type Orientation = 'forward' | 'reverse'
+
+export function encodeNode(id: number, orientation: Orientation) {
+  return 2 * id + (orientation === 'reverse' ? 1 : 0)
+}
+
+export function nodeId(handle: number) {
+  return Math.floor(handle / 2)
+}
+
+export function nodeOrientation(handle: number): Orientation {
+  return handle % 2 === 0 ? 'forward' : 'reverse'
+}
+
+export function isReverse(handle: number) {
+  return handle % 2 === 1
+}
+
+export function flipNode(handle: number) {
+  return handle % 2 === 0 ? handle + 1 : handle - 1
+}
+
+export type NodeSide = 'left' | 'right'
+
+export function flipSide(side: NodeSide): NodeSide {
+  return side === 'left' ? 'right' : 'left'
+}
+
+export function entrySide(orientation: Orientation): NodeSide {
+  return orientation === 'forward' ? 'left' : 'right'
+}
+
+export function exitSide(orientation: Orientation): NodeSide {
+  return orientation === 'forward' ? 'right' : 'left'
+}
+
+export function entryOrientation(side: NodeSide): Orientation {
+  return side === 'left' ? 'forward' : 'reverse'
+}
+
+export function exitOrientation(side: NodeSide): Orientation {
+  return side === 'right' ? 'forward' : 'reverse'
+}
+
+export function edgeIsCanonical(from: number, to: number) {
+  const fromId = nodeId(from)
+  const toId = nodeId(to)
+  return isReverse(from)
+    ? toId > fromId || (toId === fromId && !isReverse(to))
+    : toId >= fromId
+}
+
+export function pathIsCanonical(path: number[]) {
+  const first = path[0]
+  const last = path[path.length - 1]
+  if (first === undefined || last === undefined) {
+    return true
+  }
+  return isReverse(first) === isReverse(last)
+    ? !isReverse(first)
+    : edgeIsCanonical(first, last)
+}

package/src/gbwt/record.ts

@@ -0,0 +1,154 @@
+import { ByteCodeReader, RunReader } from './bytecode.ts'
+import { ENDMARKER, flipNode, nodeId } from './node.ts'
+
+export interface Pos {
+  node: number
+  offset: number
+}
+
+export function decompressEdges(bytes: Uint8Array): Pos[] {
+  const reader = new ByteCodeReader(bytes)
+  const sigma = reader.int()
+  if (sigma === undefined || sigma === 0) {
+    return []
+  }
+  const edges: Pos[] = []
+  let prev = 0
+  for (let i = 0; i < sigma; i++) {
+    const delta = reader.int()
+    const offset = reader.int()
+    if (delta === undefined || offset === undefined) {
+      throw new Error('GBWT edge list ends early')
+    }
+    prev += delta
+    edges.push({ node: prev, offset })
+  }
+  return edges
+}
+
+export class GbwtRecord {
+  readonly edges: Pos[]
+  readonly bwt: Uint8Array
+
+  constructor(edges: Pos[], bwt: Uint8Array) {
+    this.edges = edges
+    this.bwt = bwt
+  }
+
+  runs() {
+    return new RunReader(this.bwt, this.edges.length)
+  }
+
+  private edge(rank: number) {
+    const edge = this.edges[rank]
+    if (edge === undefined) {
+      throw new Error(
+        `GBWT run refers to edge rank ${rank} of ${this.edges.length}`,
+      )
+    }
+    return edge
+  }
+
+  lf(i: number): Pos | undefined {
+    const offsets = this.edges.map(e => e.offset)
+    let offset = 0
+    for (const run of this.runs()) {
+      const edge = this.edge(run.value)
+      const soFar = offsets[run.value]!
+      if (offset + run.len > i) {
+        return edge.node === ENDMARKER
+          ? undefined
+          : { node: edge.node, offset: soFar + (i - offset) }
+      }
+      offsets[run.value] = soFar + run.len
+      offset += run.len
+    }
+    return undefined
+  }
+
+  private edgeTo(node: number) {
+    let low = 0
+    let high = this.edges.length
+    while (low < high) {
+      const mid = (low + high) >> 1
+      const edge = this.edges[mid]!
+      if (edge.node === node) {
+        return mid
+      }
+      if (edge.node < node) {
+        low = mid + 1
+      } else {
+        high = mid
+      }
+    }
+    return undefined
+  }
+
+  offsetTo(pos: Pos): number | undefined {
+    if (pos.node === ENDMARKER) {
+      return undefined
+    }
+    const rank = this.edgeTo(pos.node)
+    if (rank === undefined) {
+      return undefined
+    }
+    let succRank = this.edges[rank]!.offset
+    if (succRank > pos.offset) {
+      return undefined
+    }
+    let offset = 0
+    for (const run of this.runs()) {
+      offset += run.len
+      if (run.value !== rank) {
+        continue
+      }
+      succRank += run.len
+      if (succRank > pos.offset) {
+        return offset - (succRank - pos.offset)
+      }
+    }
+    return undefined
+  }
+
+  predecessorAt(i: number): number | undefined {
+    const counts = this.edges.map(e => ({
+      node: e.node === ENDMARKER ? ENDMARKER : flipNode(e.node),
+      count: 0,
+    }))
+    for (const run of this.runs()) {
+      const entry = counts[run.value]
+      if (entry) {
+        entry.count += run.len
+      }
+    }
+    for (let rank = 1; rank < counts.length; rank++) {
+      const prev = counts[rank - 1] as { node: number; count: number }
+      const curr = counts[rank] as { node: number; count: number }
+      if (nodeId(prev.node) === nodeId(curr.node)) {
+        counts[rank - 1] = curr
+        counts[rank] = prev
+      }
+    }
+    let offset = 0
+    for (const entry of counts) {
+      offset += entry.count
+      if (offset > i) {
+        return entry.node === ENDMARKER ? undefined : entry.node
+      }
+    }
+    return undefined
+  }
+
+  decompress(): Pos[] {
+    const offsets = this.edges.map(e => e.offset)
+    const result: Pos[] = []
+    for (const run of this.runs()) {
+      const edge = this.edge(run.value)
+      for (let k = 0; k < run.len; k++) {
+        result.push({ node: edge.node, offset: offsets[run.value]! })
+        offsets[run.value] = offsets[run.value]! + 1
+      }
+    }
+    return result
+  }
+}

package/src/gbwt/sequence.ts

@@ -0,0 +1,50 @@
+const DECODE = ['', 'A', 'C', 'G', 'T', 'N']
+
+export function encodedSequenceLength(encoded: Uint8Array) {
+  const last = encoded[encoded.length - 1]
+  if (last === undefined) {
+    return 0
+  }
+  let value = last
+  let inLast = 0
+  for (let i = 0; i < 3; i++) {
+    if (value % 6 === 0) {
+      break
+    }
+    value = Math.floor(value / 6)
+    inLast += 1
+  }
+  return 3 * (encoded.length - 1) + inLast
+}
+
+export function decodeSequence(encoded: Uint8Array) {
+  let result = ''
+  for (const byte of encoded) {
+    let value = byte
+    for (let i = 0; i < 3; i++) {
+      const base = DECODE[value % 6]!
+      if (base === '') {
+        return result
+      }
+      value = Math.floor(value / 6)
+      result += base
+    }
+  }
+  return result
+}
+
+const COMPLEMENT: Record<string, string> = {
+  A: 'T',
+  C: 'G',
+  G: 'C',
+  T: 'A',
+  N: 'N',
+}
+
+export function reverseComplement(sequence: string) {
+  let result = ''
+  for (let i = sequence.length - 1; i >= 0; i--) {
+    result += COMPLEMENT[sequence[i]!] ?? 'N'
+  }
+  return result
+}

package/src/index.ts

@@ -0,0 +1,31 @@
+export {
+  GBZBase,
+  GENERIC_SAMPLE,
+  GbzRecord,
+  SCHEMA_VERSION,
+  SchemaVersionError,
+  formatPathName,
+} from './db.ts'
+export type { GbzPath, HaplotypeSample, PathName } from './db.ts'
+export type { ByteSource } from './filehandle.ts'
+export type { Pos } from './gbwt/record.ts'
+export { Subgraph } from './subgraph.ts'
+export type {
+  HaplotypeAlignment,
+  HaplotypeOutput,
+  PathIdentity,
+  PathPosition,
+  ReferencePath,
+  SubgraphJson,
+  SubgraphPath,
+  ToJsonOptions,
+} from './subgraph.ts'
+export {
+  subgraphAroundNodes,
+  subgraphAtOffset,
+  subgraphInInterval,
+} from './query.ts'
+export type { PathQuery, QueryOptions } from './query.ts'
+export { SqliteDatabase } from './sqlite/database.ts'
+export { weightedLcs } from './lcs.ts'
+export * as nodes from './gbwt/node.ts'

package/src/lcs.ts

@@ -0,0 +1,282 @@
+interface Point {
+  weight: number
+  a: number
+  b: number
+  matches: number
+}
+
+function prefixSums(sequence: number[], weight: (x: number) => number) {
+  const sums = [0]
+  for (let i = 0; i < sequence.length; i++) {
+    sums.push(sums[i]! + weight(sequence[i]!))
+  }
+  return sums
+}
+
+class MinHeap {
+  private items: number[] = []
+
+  push(value: number) {
+    const items = this.items
+    items.push(value)
+    let i = items.length - 1
+    while (i > 0) {
+      const parent = (i - 1) >> 1
+      if (items[parent]! <= value) {
+        break
+      }
+      items[i] = items[parent]!
+      i = parent
+    }
+    items[i] = value
+  }
+
+  peek() {
+    return this.items[0]
+  }
+
+  pop() {
+    const items = this.items
+    const top = items[0]
+    const last = items.pop()
+    if (items.length > 0 && last !== undefined) {
+      let i = 0
+      for (;;) {
+        const left = 2 * i + 1
+        const right = left + 1
+        let smallest = i
+        let value = last
+        if (left < items.length && items[left]! < value) {
+          smallest = left
+          value = items[left]!
+        }
+        if (right < items.length && items[right]! < value) {
+          smallest = right
+        }
+        if (smallest === i) {
+          break
+        }
+        items[i] = items[smallest]!
+        i = smallest
+      }
+      items[i] = last
+    }
+    return top
+  }
+}
+
+class Matrix {
+  readonly aSums: number[]
+  readonly bSums: number[]
+  readonly a: number[]
+  readonly b: number[]
+  readonly points = new Map<number, Map<number, Point>>()
+  private pendingEdits = new MinHeap()
+
+  constructor(a: number[], b: number[], weight: (x: number) => number) {
+    this.a = a
+    this.b = b
+    this.aSums = prefixSums(a, weight)
+    this.bSums = prefixSums(b, weight)
+    this.set(0, 0, { weight: 0, a: 0, b: 0, matches: 0 })
+  }
+
+  private set(edits: number, diagonal: number, point: Point) {
+    let row = this.points.get(edits)
+    if (!row) {
+      row = new Map()
+      this.points.set(edits, row)
+      this.pendingEdits.push(edits)
+    }
+    row.set(diagonal, point)
+  }
+
+  get(edits: number, diagonal: number) {
+    return this.points.get(edits)?.get(diagonal)
+  }
+
+  private tryInsert(edits: number, diagonal: number, point: Point) {
+    const existing = this.get(edits, diagonal)
+    if (!existing || point.weight > existing.weight) {
+      this.set(edits, diagonal, point)
+    }
+  }
+
+  aWeight(offset: number) {
+    return this.aSums[offset + 1]! - this.aSums[offset]!
+  }
+
+  bWeight(offset: number) {
+    return this.bSums[offset + 1]! - this.bSums[offset]!
+  }
+
+  extend(edits: number): Point | undefined {
+    const row = this.points.get(edits)
+    if (!row) {
+      return undefined
+    }
+    const diagonals = [...row.keys()].sort((x, y) => x - y)
+    for (const diagonal of diagonals) {
+      const found = row.get(diagonal)
+      if (!found) {
+        continue
+      }
+      const point = { ...found }
+      while (
+        point.a < this.a.length &&
+        point.b < this.b.length &&
+        this.a[point.a] === this.b[point.b]
+      ) {
+        point.weight += 2 * this.aWeight(point.a)
+        point.a += 1
+        point.b += 1
+        point.matches += 1
+      }
+      if (point.matches > 0) {
+        row.set(diagonal, point)
+      }
+      if (point.a === this.a.length && point.b === this.b.length) {
+        return point
+      }
+      if (point.a < this.a.length) {
+        const w = this.aWeight(point.a)
+        this.tryInsert(edits + w, diagonal + w, {
+          weight: point.weight,
+          a: point.a + 1,
+          b: point.b,
+          matches: 0,
+        })
+      }
+      if (point.b < this.b.length) {
+        const w = this.bWeight(point.b)
+        this.tryInsert(edits + w, diagonal - w, {
+          weight: point.weight,
+          a: point.a,
+          b: point.b + 1,
+          matches: 0,
+        })
+      }
+    }
+    return undefined
+  }
+
+  nextEdits(edits: number) {
+    while (
+      this.pendingEdits.peek() !== undefined &&
+      this.pendingEdits.peek()! <= edits
+    ) {
+      this.pendingEdits.pop()
+    }
+    return this.pendingEdits.peek()
+  }
+
+  predecessor(
+    a: number,
+    b: number,
+    edits: number,
+  ): [Point, number] | undefined {
+    const diagonal = this.aSums[a]! - this.bSums[b]!
+    const prev =
+      a > 0 && this.aWeight(a - 1) <= edits
+        ? this.get(edits - this.aWeight(a - 1), diagonal - this.aWeight(a - 1))
+        : undefined
+    const next =
+      b > 0 && this.bWeight(b - 1) <= edits
+        ? this.get(edits - this.bWeight(b - 1), diagonal + this.bWeight(b - 1))
+        : undefined
+    if (prev && next) {
+      return prev.weight > next.weight
+        ? [prev, edits - this.aWeight(a - 1)]
+        : [next, edits - this.bWeight(b - 1)]
+    }
+    if (prev) {
+      return [prev, edits - this.aWeight(a - 1)]
+    }
+    if (next) {
+      return [next, edits - this.bWeight(b - 1)]
+    }
+    return undefined
+  }
+}
+
+export function weightedLcs(
+  a: number[],
+  b: number[],
+  weight: (x: number) => number,
+): [pairs: [number, number][], weight: number] {
+  let prefix = 0
+  while (prefix < a.length && prefix < b.length && a[prefix] === b[prefix]) {
+    prefix += 1
+  }
+  let suffix = 0
+  while (
+    suffix < a.length - prefix &&
+    suffix < b.length - prefix &&
+    a[a.length - 1 - suffix] === b[b.length - 1 - suffix]
+  ) {
+    suffix += 1
+  }
+  const pairs: [number, number][] = []
+  let total = 0
+  for (let i = 0; i < prefix; i++) {
+    pairs.push([i, i])
+    total += weight(a[i]!)
+  }
+  const [middle, middleWeight] = weightedLcsCore(
+    a.slice(prefix, a.length - suffix),
+    b.slice(prefix, b.length - suffix),
+    weight,
+  )
+  for (const [i, j] of middle) {
+    pairs.push([i + prefix, j + prefix])
+  }
+  total += middleWeight
+  for (let i = suffix; i > 0; i--) {
+    pairs.push([a.length - i, b.length - i])
+    total += weight(a[a.length - i]!)
+  }
+  return [pairs, total]
+}
+
+function weightedLcsCore(
+  a: number[],
+  b: number[],
+  weight: (x: number) => number,
+): [pairs: [number, number][], weight: number] {
+  if (a.length === 0 || b.length === 0) {
+    return [[], 0]
+  }
+  const matrix = new Matrix(a, b, weight)
+  let edits = 0
+  let point: Point = { weight: 0, a: 0, b: 0, matches: 0 }
+  for (;;) {
+    const end = matrix.extend(edits)
+    if (end) {
+      point = end
+      break
+    }
+    const next = matrix.nextEdits(edits)
+    if (next === undefined) {
+      break
+    }
+    edits = next
+  }
+  const result: [number, number][] = []
+  const finalWeight = point.weight / 2
+  point = { ...point }
+  for (;;) {
+    for (let i = 0; i < point.matches; i++) {
+      point.a -= 1
+      point.b -= 1
+      result.push([point.a, point.b])
+    }
+    const pred = matrix.predecessor(point.a, point.b, edits)
+    if (!pred) {
+      break
+    }
+    point = { ...pred[0] }
+    edits = pred[1]
+  }
+  result.reverse()
+  return [result, finalWeight]
+}

package/src/query.ts

@@ -0,0 +1,79 @@
+import { GENERIC_SAMPLE } from './db.ts'
+import { Subgraph } from './subgraph.ts'
+
+import type { GBZBase, PathName } from './db.ts'
+import type { HaplotypeOutput } from './subgraph.ts'
+
+export interface QueryOptions {
+  context?: number
+  haplotypes?: HaplotypeOutput
+  limit?: number
+}
+
+export interface PathQuery {
+  sample?: string
+  contig: string
+  haplotype?: number
+}
+
+function pathName(query: PathQuery, fragment: number): PathName {
+  return {
+    sample: query.sample ?? GENERIC_SAMPLE,
+    contig: query.contig,
+    haplotype: query.haplotype ?? 0,
+    fragment,
+  }
+}
+
+export async function subgraphAtOffset(
+  db: GBZBase,
+  query: PathQuery,
+  offset: number,
+  opts: QueryOptions = {},
+) {
+  const subgraph = new Subgraph(db)
+  subgraph.limit = opts.limit
+  const reference = await subgraph.pathPosition(pathName(query, offset))
+  await subgraph.aroundPosition(
+    reference.position.handle,
+    reference.position.nodeOffset,
+    opts.context ?? 100,
+  )
+  subgraph.extractPaths(reference, opts.haplotypes ?? 'all')
+  return subgraph
+}
+
+export async function subgraphInInterval(
+  db: GBZBase,
+  query: PathQuery,
+  start: number,
+  end: number,
+  opts: QueryOptions = {},
+) {
+  const subgraph = new Subgraph(db)
+  subgraph.limit = opts.limit
+  const reference = await subgraph.pathPosition(pathName(query, start))
+  await subgraph.aroundInterval(
+    reference.position,
+    end - start,
+    opts.context ?? 100,
+  )
+  subgraph.extractPaths(reference, opts.haplotypes ?? 'all')
+  return subgraph
+}
+
+export async function subgraphAroundNodes(
+  db: GBZBase,
+  nodes: number[],
+  opts: QueryOptions = {},
+) {
+  const haplotypes = opts.haplotypes ?? 'all'
+  if (haplotypes === 'reference-only') {
+    throw new Error('Cannot output a reference path in a node-based query')
+  }
+  const subgraph = new Subgraph(db)
+  subgraph.limit = opts.limit
+  await subgraph.aroundNodes(nodes, opts.context ?? 100)
+  subgraph.extractPaths(undefined, haplotypes)
+  return subgraph
+}