@gmod/gbz-base 1.0.0 → 2.0.0

package/src/subgraph.ts

@@ -1,4 +1,3 @@
-import { formatPathName } from './db.ts'
 import {
   ENDMARKER,
   edgeIsCanonical,
@@ -7,26 +6,31 @@ 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 { formatPathName } from './pathName.ts'
 
-import type {
-  GBZBase,
-  GbzPath,
-  GbzRecord,
-  HaplotypeSample,
-  PathName,
-} from './db.ts'
+import type { GBZBase, GbzPath, GbzRecord, HaplotypeSample } from './db.ts'
 import type { NodeSide, Orientation } from './gbwt/node.ts'
 import type { Pos } from './gbwt/record.ts'
+import type { PathName } from './pathName.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
@@ -77,12 +81,8 @@ export interface SubgraphJson {
   paths: SubgraphPath[]
 }
 
-export interface HaplotypeAlignment {
-  pathHandle: number | undefined
-  name: PathName | undefined
+export interface AlignmentSpan {
   strand: '+' | '-'
-  hapStart: number | undefined
-  hapEnd: number | undefined
   refStart: number
   refEnd: number
   cigar: string
@@ -91,35 +91,86 @@ export interface HaplotypeAlignment {
   start: Pos
 }
 
-export interface ToJsonOptions {
-  names?: 'anonymous' | 'resolved'
+export type HaplotypeAlignment = AlignmentSpan &
+  (
+    | {
+        resolved: true
+        name: PathName
+        label: string
+        pathHandle: number
+        hapStart: number
+        hapEnd: number
+      }
+    | { resolved: false }
+  )
+
+export interface SubgraphOutputOptions {
+  cigar?: boolean | undefined
+  names?: 'anonymous' | 'resolved' | undefined
+}
+
+export interface SubgraphOptions {
+  limit?: number | undefined
+  signal?: AbortSignal | undefined
+}
+
+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
   }
 }
 
@@ -142,7 +193,6 @@ export class Subgraph {
   private refInterval: [number, number] | undefined
   private refIndexCache: Map<number, number[]> | undefined
   private refPrefixCache: number[] | undefined
-  limit: number | undefined
   readonly stats = {
     orderedAlignments: 0,
     lcsAlignments: 0,
@@ -151,9 +201,13 @@ export class Subgraph {
   }
 
   private db: GBZBase
+  private readonly limit: number | undefined
+  private readonly signal: AbortSignal | undefined
 
-  constructor(db: GBZBase) {
+  constructor(db: GBZBase, opts: SubgraphOptions = {}) {
     this.db = db
+    this.limit = opts.limit
+    this.signal = opts.signal
   }
 
   get nodeCount() {
@@ -195,6 +249,7 @@ export class Subgraph {
   }
 
   private async addNode(id: number) {
+    this.signal?.throwIfAborted()
     if (this.limit !== undefined && this.nodeCount >= this.limit) {
       throw new Error(`Subgraph size limit of ${this.limit} nodes exceeded`)
     }
@@ -394,6 +449,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 +604,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 +747,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))
@@ -632,6 +814,7 @@ export class Subgraph {
       let current: number | undefined = start
       let pos: Pos | undefined
       while (anchor === undefined) {
+        this.signal?.throwIfAborted()
         if (current !== undefined) {
           if (visited.has(current)) {
             break
@@ -926,24 +1109,17 @@ export class Subgraph {
         last -= 1
       }
       const identity = info.identity
-      const strand =
+      const walkStrand =
         info.path.some(handle => isReverse(handle)) &&
         !info.path.some(handle => !isReverse(handle))
           ? '-'
           : '+'
-      result.push({
-        pathHandle: identity?.pathHandle,
-        name: identity?.name,
+      const span: AlignmentSpan = {
         strand: identity
           ? identity.orientation === 'forward'
             ? '+'
             : '-'
-          : strand,
-        hapStart: identity
-          ? identity.name.fragment + identity.hapStart
-          : undefined,
-        hapEnd: identity ? identity.name.fragment + identity.hapEnd : undefined,
-        start: info.positions[0]!,
+          : walkStrand,
         refStart: reference.start + leading,
         refEnd: reference.start + refTotal - trailing,
         cigar: edits
@@ -952,12 +1128,152 @@ export class Subgraph {
           .join(''),
         weight: info.weight,
         path: info.path,
-      })
+        start: info.positions[0]!,
+      }
+      if (identity) {
+        const hapStart = identity.name.fragment + identity.hapStart
+        const hapEnd = identity.name.fragment + identity.hapEnd
+        result.push({
+          ...span,
+          resolved: true,
+          name: identity.name,
+          label: formatPathName(
+            { ...identity.name, fragment: hapStart },
+            hapEnd,
+          ),
+          pathHandle: identity.pathHandle,
+          hapStart,
+          hapEnd,
+        })
+      } else {
+        result.push({ ...span, resolved: false })
+      }
     })
     return result
   }
 
-  toJSON(cigar: boolean, opts: ToJsonOptions = {}): SubgraphJson {
+  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(opts: SubgraphOutputOptions = {}) {
+    const cigar = opts.cigar ?? false
+    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`
+  }
+
+  toSubgraphJson(opts: SubgraphOutputOptions = {}): SubgraphJson {
+    const cigar = opts.cigar ?? false
     const handles = this.sortedHandles()
     const nodes = handles
       .filter(handle => !isReverse(handle))