@gmod/gbz-base 1.0.0 → 1.1.0

package/src/sqlite/pager.ts

@@ -33,6 +33,23 @@ export class Pager {
     this.maxBlocks = opts.maxBlocks ?? 256
   }
 
+  private async read(length: number, start: number) {
+    let attempt = 0
+    for (;;) {
+      try {
+        const bytes = await this.source.read(length, start)
+        this.bytesFetched += bytes.length
+        return bytes
+      } catch (error) {
+        attempt += 1
+        if (attempt >= 3) {
+          throw error
+        }
+        await new Promise(resolve => setTimeout(resolve, 200 * attempt))
+      }
+    }
+  }
+
   seed(index: number, bytes: Uint8Array) {
     if (
       bytes.length ===
@@ -51,19 +68,72 @@ export class Pager {
     }
     const start = index * this.blockSize
     const length = Math.min(this.blockSize, this.fileSize - start)
-    const pending = this.source.read(length, start).then(bytes => {
-      this.bytesFetched += bytes.length
-      return bytes
-    })
+    const pending = this.read(length, start)
     this.fetches += 1
     this.blocks.set(index, pending)
-    if (this.blocks.size > this.maxBlocks) {
+    this.forgetOnFailure(pending, [index])
+    this.evict()
+    return pending
+  }
+
+  prefetch(pageNumbers: number[]) {
+    const wanted = [
+      ...new Set(
+        pageNumbers.map(pageNumber =>
+          Math.floor(((pageNumber - 1) * this.pageSize) / this.blockSize),
+        ),
+      ),
+    ]
+      .filter(index => !this.blocks.has(index))
+      .sort((a, b) => a - b)
+    let runStart = 0
+    while (runStart < wanted.length) {
+      let runEnd = runStart + 1
+      while (
+        runEnd < wanted.length &&
+        wanted[runEnd] === wanted[runEnd - 1]! + 1
+      ) {
+        runEnd += 1
+      }
+      this.fetchRun(wanted[runStart]!, runEnd - runStart)
+      runStart = runEnd
+    }
+  }
+
+  private fetchRun(firstIndex: number, count: number) {
+    const start = firstIndex * this.blockSize
+    const length = Math.min(count * this.blockSize, this.fileSize - start)
+    const pending = this.read(length, start)
+    this.fetches += 1
+    const indexes: number[] = []
+    for (let i = 0; i < count; i++) {
+      const within = i * this.blockSize
+      indexes.push(firstIndex + i)
+      this.blocks.set(
+        firstIndex + i,
+        pending.then(bytes => bytes.subarray(within, within + this.blockSize)),
+      )
+    }
+    this.forgetOnFailure(pending, indexes)
+    this.evict()
+  }
+
+  private forgetOnFailure(pending: Promise<Uint8Array>, indexes: number[]) {
+    pending.catch(() => {
+      for (const index of indexes) {
+        this.blocks.delete(index)
+      }
+    })
+  }
+
+  private evict() {
+    while (this.blocks.size > this.maxBlocks) {
       const oldest = this.blocks.keys().next().value
-      if (oldest !== undefined) {
-        this.blocks.delete(oldest)
+      if (oldest === undefined) {
+        break
       }
+      this.blocks.delete(oldest)
     }
-    return pending
   }
 
   async page(pageNumber: number) {

package/src/subgraph.ts

@@ -7,12 +7,14 @@ import {
   entrySide,
   exitOrientation,
   exitSide,
+  flipNode,
   flipSide,
   isReverse,
   nodeId,
   nodeOrientation,
   pathIsCanonical,
 } from './gbwt/node.ts'
+import { gfaHeaderLines, sha256Hex, subgraphName } from './graphName.ts'
 import { weightedLcs } from './lcs.ts'
 
 import type {
@@ -27,6 +29,13 @@ import type { Pos } from './gbwt/record.ts'
 
 export type HaplotypeOutput = 'all' | 'distinct' | 'reference-only' | 'none'
 
+export type SnarlOutput = 'none' | 'contained' | 'overlapping'
+
+type HandleType =
+  | { kind: 'snarl-exit'; snarl: [number, number] }
+  | { kind: 'chain' }
+  | { kind: 'regular' }
+
 export interface PathPosition {
   seqOffset: number
   handle: number
@@ -95,31 +104,63 @@ export interface ToJsonOptions {
   names?: 'anonymous' | 'resolved'
 }
 
+function sideBefore(
+  a: [number, number, NodeSide],
+  b: [number, number, NodeSide],
+) {
+  return (
+    a[0] < b[0] ||
+    (a[0] === b[0] && (a[1] < b[1] || (a[1] === b[1] && a[2] < b[2])))
+  )
+}
+
 class SideQueue {
-  private items: [number, number, NodeSide][] = []
+  private heap: [number, number, NodeSide][] = []
 
   push(distance: number, node: number, side: NodeSide) {
-    this.items.push([distance, node, side])
+    const heap = this.heap
+    heap.push([distance, node, side])
+    let i = heap.length - 1
+    while (i > 0) {
+      const parent = (i - 1) >> 1
+      if (sideBefore(heap[i]!, heap[parent]!)) {
+        ;[heap[i], heap[parent]] = [heap[parent]!, heap[i]!]
+        i = parent
+      } else {
+        break
+      }
+    }
   }
 
   pop() {
-    let best = 0
-    for (let i = 1; i < this.items.length; i++) {
-      const a = this.items[i]!
-      const b = this.items[best]!
-      if (
-        a[0] < b[0] ||
-        (a[0] === b[0] && (a[1] < b[1] || (a[1] === b[1] && a[2] < b[2])))
-      ) {
-        best = i
+    const heap = this.heap
+    const top = heap[0]
+    const last = heap.pop()
+    if (heap.length > 0 && last !== undefined) {
+      heap[0] = last
+      let i = 0
+      for (;;) {
+        const left = 2 * i + 1
+        const right = left + 1
+        let smallest = i
+        if (left < heap.length && sideBefore(heap[left]!, heap[smallest]!)) {
+          smallest = left
+        }
+        if (right < heap.length && sideBefore(heap[right]!, heap[smallest]!)) {
+          smallest = right
+        }
+        if (smallest === i) {
+          break
+        }
+        ;[heap[i], heap[smallest]] = [heap[smallest]!, heap[i]!]
+        i = smallest
       }
     }
-    const [item] = this.items.splice(best, 1)
-    return item
+    return top
   }
 
   get size() {
-    return this.items.length
+    return this.heap.length
   }
 }
 
@@ -394,6 +435,150 @@ export class Subgraph {
     return { inserted, removed: toRemove.size }
   }
 
+  async betweenNodes(start: number, end: number) {
+    this.clearPaths()
+    const active = [start, flipNode(end)]
+    const visited = new Set([nodeId(start), nodeId(end)])
+    let inserted = 0
+    while (active.length > 0) {
+      const curr = active.pop()!
+      const id = nodeId(curr)
+      if (!this.hasNode(id)) {
+        await this.addNode(id)
+        inserted += 1
+      }
+      for (const successor of this.record(curr).successors()) {
+        const successorId = nodeId(successor)
+        if (!visited.has(successorId)) {
+          active.push(successor, flipNode(successor))
+          visited.add(successorId)
+        }
+      }
+    }
+    return inserted
+  }
+
+  async extractSnarls(snarls: SnarlOutput) {
+    let inserted = 0
+    for (const [start, end] of await this.overlappingSnarls(snarls)) {
+      inserted += await this.betweenNodes(start, end)
+    }
+    return inserted
+  }
+
+  private async overlappingSnarls(snarls: SnarlOutput) {
+    const result: [number, number][] = []
+    if (snarls !== 'none') {
+      let foundLink = false
+      for (const handle of this.sortedHandles()) {
+        const record = this.record(handle)
+        const next = record.next
+        if (next !== undefined) {
+          foundLink = true
+          if (this.hasHandle(next)) {
+            if (edgeIsCanonical(handle, next)) {
+              result.push([handle, next])
+            }
+          } else if (
+            snarls === 'overlapping' &&
+            this.isSnarlEntryInSubgraph(record)
+          ) {
+            result.push([handle, next])
+          }
+        }
+      }
+      if (
+        !foundLink &&
+        snarls === 'overlapping' &&
+        (await this.db.hasChainLinks())
+      ) {
+        const covering = await this.findCoveringSnarl()
+        if (covering) {
+          result.push(covering)
+        }
+      }
+    }
+    return result
+  }
+
+  private isSnarlEntryInSubgraph(record: GbzRecord) {
+    const successors = record.successors()
+    const first = successors.find(handle => this.hasHandle(handle))
+    return first === undefined
+      ? false
+      : successors.length > 1 ||
+          this.record(flipNode(first)).successors().length > 1
+  }
+
+  private recordReader(onFetch?: () => void) {
+    const outside = new Map<number, GbzRecord>()
+    return async (handle: number) => {
+      const inside = this.records.get(handle)
+      if (inside) {
+        return inside
+      }
+      let record = outside.get(handle)
+      if (!record) {
+        record = await this.db.getRecord(handle)
+        onFetch?.()
+        if (!record) {
+          throw new Error(`Node record ${handle} is missing from the database`)
+        }
+        outside.set(handle, record)
+      }
+      return record
+    }
+  }
+
+  private async findCoveringSnarl() {
+    const read = this.recordReader()
+    const isSnarlEntry = async (record: GbzRecord) => {
+      const successors = record.successors()
+      const first = successors[0]
+      return first === undefined
+        ? false
+        : successors.length > 1 ||
+            (await read(flipNode(first))).successors().length > 1
+    }
+    const classify = async (handle: number): Promise<HandleType> => {
+      const reverse = await read(flipNode(handle))
+      return reverse.next !== undefined
+        ? (await isSnarlEntry(reverse))
+          ? { kind: 'snarl-exit', snarl: [flipNode(handle), reverse.next] }
+          : { kind: 'chain' }
+        : (await read(handle)).next !== undefined
+          ? { kind: 'chain' }
+          : { kind: 'regular' }
+    }
+    const visited = new Set<number>()
+    const queue = this.sortedHandles().flatMap(handle =>
+      this.record(handle).successors(),
+    )
+    let result: [number, number] | undefined
+    let done = false
+    while (!done && queue.length > 0) {
+      const handle = queue.shift()!
+      const id = nodeId(handle)
+      if (!this.hasHandle(handle) && !visited.has(id)) {
+        visited.add(id)
+        const type = await classify(handle)
+        if (type.kind === 'snarl-exit') {
+          result = type.snarl
+          done = true
+        } else if (type.kind === 'chain') {
+          done = true
+        } else {
+          for (const orientation of ['forward', 'reverse'] as const) {
+            queue.push(
+              ...(await read(encodeNode(id, orientation))).successors(),
+            )
+          }
+        }
+      }
+    }
+    return result
+  }
+
   extractPaths(reference: ReferencePath | undefined, output: HaplotypeOutput) {
     this.clearPaths()
     if (output === 'none') {
@@ -405,72 +590,69 @@ export class Subgraph {
     const handles = this.sortedHandles()
     const successors = new Map<
       number,
-      { next: Pos; hasPredecessor: boolean }[]
+      { nodes: Int32Array; offsets: Int32Array; hasPredecessor: Uint8Array }
     >()
     for (const handle of handles) {
-      successors.set(
-        handle,
-        this.record(handle)
-          .gbwt()
-          .decompress()
-          .map(next => ({ next, hasPredecessor: false })),
-      )
+      const { nodes, offsets } = this.record(handle).gbwt().decompressArrays()
+      successors.set(handle, {
+        nodes,
+        offsets,
+        hasPredecessor: new Uint8Array(nodes.length),
+      })
     }
     for (const handle of handles) {
-      for (const { next } of successors.get(handle) as { next: Pos }[]) {
-        const entry = successors.get(next.node)?.[next.offset]
+      const { nodes, offsets } = successors.get(handle)!
+      for (let i = 0; i < nodes.length; i++) {
+        const entry = successors.get(nodes[i]!)
         if (entry) {
-          entry.hasPredecessor = true
+          entry.hasPredecessor[offsets[i]!] = 1
         }
       }
     }
     let refOffset: number | undefined
     for (const handle of handles) {
-      const entries = successors.get(handle) as {
-        next: Pos
-        hasPredecessor: boolean
-      }[]
-      entries.forEach((entry, offset) => {
-        if (entry.hasPredecessor) {
-          return
-        }
-        let curr: Pos | undefined = { node: handle, offset }
-        let isRef = false
-        const path: number[] = []
-        const positions: Pos[] = []
-        let len = 0
-        while (curr) {
-          if (
-            curr.node === refPos?.handle &&
-            curr.offset === refPos.gbwtOffset
-          ) {
-            this.refId = this.paths.length
-            refOffset = path.length
-            isRef = true
+      const entries = successors.get(handle)!
+      for (let offset = 0; offset < entries.nodes.length; offset++) {
+        if (entries.hasPredecessor[offset] === 0) {
+          let currNode: number | undefined = handle
+          let currOffset = offset
+          let isRef = false
+          const path: number[] = []
+          const positions: Pos[] = []
+          let len = 0
+          while (currNode !== undefined) {
+            if (
+              currNode === refPos?.handle &&
+              currOffset === refPos.gbwtOffset
+            ) {
+              this.refId = this.paths.length
+              refOffset = path.length
+              isRef = true
+            }
+            path.push(currNode)
+            positions.push({ node: currNode, offset: currOffset })
+            len += this.record(currNode).sequenceLen
+            const step = successors.get(currNode)!
+            const nextNode: number = step.nodes[currOffset]!
+            const nextOffset: number = step.offsets[currOffset]!
+            if (nextNode !== ENDMARKER && successors.has(nextNode)) {
+              currNode = nextNode
+              currOffset = nextOffset
+            } else {
+              currNode = undefined
+            }
+          }
+          if (isRef || pathIsCanonical(path)) {
+            this.paths.push({
+              path,
+              positions,
+              len,
+              weight: undefined,
+              identity: undefined,
+            })
           }
-          path.push(curr.node)
-          positions.push(curr)
-          len += this.record(curr.node).sequenceLen
-          const step: { next: Pos } | undefined = successors.get(curr.node)?.[
-            curr.offset
-          ]
-          curr =
-            step &&
-            step.next.node !== ENDMARKER &&
-            successors.has(step.next.node)
-              ? step.next
-              : undefined
-        }
-        if (isRef || pathIsCanonical(path)) {
-          this.paths.push({
-            path,
-            positions,
-            len,
-            weight: undefined,
-            identity: undefined,
-          })
         }
-      })
+      }
     }
     if (refPos) {
       if (refOffset === undefined || this.refId === undefined) {
@@ -551,23 +733,9 @@ export class Subgraph {
         starts.set(posKey(first), index)
       }
     })
-    const outside = new Map<number, GbzRecord>()
-    const recordAt = async (handle: number) => {
-      const inside = this.records.get(handle)
-      if (inside) {
-        return inside
-      }
-      let record = outside.get(handle)
-      if (!record) {
-        record = await this.db.getRecord(handle)
-        this.stats.identificationFetches += 1
-        if (!record) {
-          throw new Error(`Node record ${handle} is missing from the database`)
-        }
-        outside.set(handle, record)
-      }
-      return record
-    }
+    const recordAt = this.recordReader(() => {
+      this.stats.identificationFetches += 1
+    })
     const sampleAt = async (pos: Pos) => {
       if (pos.node >= minHandle && pos.node <= maxHandle) {
         return samples.get(posKey(pos))
@@ -957,6 +1125,125 @@ export class Subgraph {
     return result
   }
 
+  private canonicalEdges(id: number) {
+    const edges: [number, number, number][] = []
+    for (const orientation of ['forward', 'reverse'] as const) {
+      const handle = encodeNode(id, orientation)
+      for (const successor of this.record(handle).successors()) {
+        if (this.hasHandle(successor) && edgeIsCanonical(handle, successor)) {
+          edges.push([
+            orientation === 'reverse' ? 1 : 0,
+            nodeId(successor),
+            isReverse(successor) ? 1 : 0,
+          ])
+        }
+      }
+    }
+    edges.sort((a, b) => a[0] - b[0] || a[1] - b[1] || a[2] - b[2])
+    return edges.filter(
+      (edge, i) =>
+        i === 0 ||
+        edge[0] !== edges[i - 1]![0] ||
+        edge[1] !== edges[i - 1]![1] ||
+        edge[2] !== edges[i - 1]![2],
+    )
+  }
+
+  async stableName() {
+    const encoder = new TextEncoder()
+    const chunks: Uint8Array[] = []
+    for (const handle of this.sortedHandles()) {
+      if (!isReverse(handle)) {
+        const id = nodeId(handle)
+        let text = `S\t${id}\t${this.record(handle).sequence}\n`
+        for (const [fromReverse, toId, toReverse] of this.canonicalEdges(id)) {
+          text += `L\t${id}\t${fromReverse ? '-' : '+'}\t${toId}\t${toReverse ? '-' : '+'}\n`
+        }
+        chunks.push(encoder.encode(text))
+      }
+    }
+    return sha256Hex(chunks)
+  }
+
+  async toGFA(cigar: boolean, opts: ToJsonOptions = {}) {
+    const lines = [
+      this.refPath ? `H\tVN:Z:1.1\tRS:Z:${this.refPath.sample}` : 'H\tVN:Z:1.1',
+      ...gfaHeaderLines(
+        subgraphName(await this.stableName(), await this.db.graphName()),
+      ),
+    ]
+    const handles = this.sortedHandles()
+    for (const handle of handles) {
+      if (!isReverse(handle)) {
+        lines.push(`S\t${nodeId(handle)}\t${this.record(handle).sequence}`)
+      }
+    }
+    const sign = (handle: number) => (isReverse(handle) ? '-' : '+')
+    for (const handle of handles) {
+      for (const successor of this.record(handle).successors()) {
+        if (this.hasHandle(successor) && edgeIsCanonical(handle, successor)) {
+          lines.push(
+            `L\t${nodeId(handle)}\t${sign(handle)}\t${nodeId(successor)}\t${sign(successor)}\t0M`,
+          )
+        }
+      }
+    }
+    const walk = (
+      info: PathInfo,
+      name: PathName,
+      end: number,
+      cigarString: string | undefined,
+    ) => {
+      const steps = info.path
+        .map(handle => `${isReverse(handle) ? '<' : '>'}${nodeId(handle)}`)
+        .join('')
+      const weight = info.weight === undefined ? '' : `\tWT:i:${info.weight}`
+      const cg = cigarString === undefined ? '' : `\tCG:Z:${cigarString}`
+      return `W\t${name.sample}\t${name.haplotype}\t${name.contig}\t${name.fragment}\t${end}\t${steps}${weight}${cg}`
+    }
+    const contig = this.refPath?.contig ?? 'unknown'
+    if (this.refId !== undefined && this.refPath && this.refInterval) {
+      lines.push(
+        walk(
+          this.paths[this.refId]!,
+          {
+            ...this.refPath,
+            fragment: this.refPath.fragment + this.refInterval[0],
+          },
+          this.refPath.fragment + this.refInterval[1],
+          undefined,
+        ),
+      )
+    }
+    let haplotype = 1
+    this.paths.forEach((info, index) => {
+      if (index !== this.refId) {
+        const resolved = opts.names === 'resolved' ? info.identity : undefined
+        const cigarString = cigar ? this.alignToRef(index) : undefined
+        lines.push(
+          resolved
+            ? walk(
+                info,
+                {
+                  ...resolved.name,
+                  fragment: resolved.name.fragment + resolved.hapStart,
+                },
+                resolved.name.fragment + resolved.hapEnd,
+                cigarString,
+              )
+            : walk(
+                info,
+                { sample: 'unknown', contig, haplotype, fragment: 0 },
+                info.len,
+                cigarString,
+              ),
+        )
+        haplotype += 1
+      }
+    })
+    return `${lines.join('\n')}\n`
+  }
+
   toJSON(cigar: boolean, opts: ToJsonOptions = {}): SubgraphJson {
     const handles = this.sortedHandles()
     const nodes = handles