@gmod/gbz-base 1.0.0 → 1.1.0

package/src/cli.ts

@@ -1,13 +1,15 @@
 import { LocalFile, RemoteFile } from 'generic-filehandle2'
 
 import { GBZBase, formatPathName } from './db.ts'
+import { encodeNode } from './gbwt/node.ts'
 import {
   subgraphAroundNodes,
   subgraphAtOffset,
+  subgraphBetween,
   subgraphInInterval,
 } from './query.ts'
 
-import type { HaplotypeOutput } from './subgraph.ts'
+import type { HaplotypeOutput, SnarlOutput } from './subgraph.ts'
 
 const USAGE = `Usage: gbz-base-query [options] graph.gbz.db
 
@@ -17,12 +19,17 @@ const USAGE = `Usage: gbz-base-query [options] graph.gbz.db
   -o, --offset INT     sequence offset
   -i, --interval A..B  half-open sequence interval
   -n, --node INT       node identifier (may repeat)
+  -b, --between A:B    subgraph between two chain boundary handles, each INT[+-]
   --context INT        context length in bp (default: 100)
+  --snarls             extend the subgraph with contained top-level snarls
+  --extend-snarls      extend the subgraph with overlapping top-level snarls
   --limit INT          safety limit for the number of nodes
   --haplotypes SEL     all, distinct, reference-only or none (default: all)
   --cigar              output CIGAR strings for the haplotypes
+  --format FMT         json (default) or gfa
   --resolve            name haplotypes from the HaplotypeSamples table
   --alignments         print one alignment record per haplotype fragment instead of the subgraph
+  --haplotype-index F  companion database written by gbz-haplotype-index --output
   --block-size INT     bytes fetched per range request (default: 65536)
   --stats              print fetch statistics to stderr
 `
@@ -35,24 +42,39 @@ interface Args {
   offset?: number
   interval?: [number, number]
   nodes: number[]
+  between?: [number, number]
   context: number
+  snarls: SnarlOutput
   limit?: number
   haplotypes: HaplotypeOutput
   cigar: boolean
+  format: 'json' | 'gfa'
   resolve: boolean
   alignments: boolean
   blockSize: number
+  haplotypeIndex?: string
   stats: boolean
 }
 
+function parseHandle(text: string) {
+  const orientation = text.endsWith('-') ? 'reverse' : 'forward'
+  const digits = /[+-]$/.test(text) ? text.slice(0, -1) : text
+  if (!/^\d+$/.test(digits)) {
+    throw new Error(`Failed to parse oriented node ${text}`)
+  }
+  return encodeNode(Number(digits), orientation)
+}
+
 function parseArgs(argv: string[]): Args {
   const args: Args = {
     file: '',
     haplotype: 0,
     nodes: [],
     context: 100,
+    snarls: 'none',
     haplotypes: 'all',
     cigar: false,
+    format: 'json',
     resolve: false,
     alignments: false,
     blockSize: 65536,
@@ -91,9 +113,25 @@ function parseArgs(argv: string[]): Args {
       case '--node':
         args.nodes.push(Number(next(i++)))
         break
+      case '-b':
+      case '--between': {
+        const [a, b, extra] = next(i++).split(':')
+        if (a === undefined || b === undefined || extra !== undefined) {
+          throw new Error(`--between needs two oriented nodes, like 14+:17-`)
+        }
+        args.between = [parseHandle(a), parseHandle(b)]
+        break
+      }
       case '--context':
         args.context = Number(next(i++))
         break
+      case '--snarls':
+        args.snarls =
+          args.snarls === 'overlapping' ? 'overlapping' : 'contained'
+        break
+      case '--extend-snarls':
+        args.snarls = 'overlapping'
+        break
       case '--limit':
         args.limit = Number(next(i++))
         break
@@ -113,12 +151,20 @@ function parseArgs(argv: string[]): Args {
       case '--block-size':
         args.blockSize = Number(next(i++))
         break
+      case '--haplotype-index':
+        args.haplotypeIndex = next(i++)
+        break
       case '--stats':
         args.stats = true
         break
-      case '--format':
-        next(i++)
+      case '--format': {
+        const format = next(i++)
+        if (format !== 'json' && format !== 'gfa') {
+          throw new Error(`Unknown output format ${format}`)
+        }
+        args.format = format
         break
+      }
       case '-h':
       case '--help':
         process.stdout.write(USAGE)
@@ -139,13 +185,18 @@ function parseArgs(argv: string[]): Args {
 
 export async function main(argv: string[]) {
   const args = parseArgs(argv)
-  const source = /^https?:\/\//.test(args.file)
-    ? new RemoteFile(args.file)
-    : new LocalFile(args.file)
-  const db = await GBZBase.open(source, { blockSize: args.blockSize })
+  const open = (file: string) =>
+    /^https?:\/\//.test(file) ? new RemoteFile(file) : new LocalFile(file)
+  const db = await GBZBase.open(open(args.file), {
+    blockSize: args.blockSize,
+    ...(args.haplotypeIndex === undefined
+      ? {}
+      : { haplotypeIndex: open(args.haplotypeIndex) }),
+  })
   const opts = {
     context: args.context,
     haplotypes: args.haplotypes,
+    snarls: args.snarls,
     ...(args.limit === undefined ? {} : { limit: args.limit }),
   }
   const query = {
@@ -153,8 +204,9 @@ export async function main(argv: string[]) {
     haplotype: args.haplotype,
     ...(args.sample === undefined ? {} : { sample: args.sample }),
   }
-  const subgraph =
-    args.nodes.length > 0
+  const subgraph = args.between
+    ? await subgraphBetween(db, args.between[0], args.between[1], opts)
+    : args.nodes.length > 0
       ? await subgraphAroundNodes(db, args.nodes, opts)
       : args.interval
         ? await subgraphInInterval(
@@ -169,12 +221,13 @@ export async function main(argv: string[]) {
           : undefined
   if (!subgraph) {
     throw new Error(
-      'Query type must be specified using --offset, --interval or --node',
+      'Query type must be specified using --offset, --interval, --node or --between',
     )
   }
   if (args.resolve) {
     await subgraph.identifyPaths()
   }
+  const names = args.resolve ? 'resolved' : 'anonymous'
   const output = args.alignments
     ? subgraph.alignments().map(a => ({
         ...a,
@@ -184,10 +237,12 @@ export async function main(argv: string[]) {
             : undefined,
         start: undefined,
       }))
-    : subgraph.toJSON(args.cigar, {
-        names: args.resolve ? 'resolved' : 'anonymous',
-      })
-  process.stdout.write(`${JSON.stringify(output)}\n`)
+    : subgraph.toJSON(args.cigar, { names })
+  process.stdout.write(
+    args.format === 'gfa' && !args.alignments
+      ? await subgraph.toGFA(args.cigar, { names })
+      : `${JSON.stringify(output)}\n`,
+  )
   if (args.stats) {
     const { fetches, bytesFetched } = db.sqlite.pager
     const {

package/src/db.ts

@@ -1,10 +1,12 @@
 import { ENDMARKER, nodeId, nodeOrientation } from './gbwt/node.ts'
 import { GbwtRecord, decompressEdges } from './gbwt/record.ts'
 import { decodeSequence, encodedSequenceLength } from './gbwt/sequence.ts'
+import { graphNameFromTags } from './graphName.ts'
 import { SqliteDatabase } from './sqlite/database.ts'
 
 import type { ByteSource } from './filehandle.ts'
 import type { Pos } from './gbwt/record.ts'
+import type { GraphName } from './graphName.ts'
 import type { PagerOptions } from './sqlite/pager.ts'
 import type { SqlValue } from './sqlite/record.ts'
 
@@ -144,37 +146,63 @@ function rowToPath(rowid: number, values: SqlValue[]): GbzPath {
   }
 }
 
+export interface OpenOptions extends PagerOptions {
+  haplotypeIndex?: ByteSource
+}
+
+async function readTags(sqlite: SqliteDatabase) {
+  const tags = new Map<string, string>()
+  for await (const { values } of sqlite.scan('Tags')) {
+    tags.set(str(values[0], 'Tags.key'), str(values[1], 'Tags.value'))
+  }
+  return tags
+}
+
 export class GBZBase {
   private tagCache: Promise<Map<string, string>> | undefined
   private pathCache: Promise<GbzPath[]> | undefined
+  private indexTags: Map<string, string> | undefined
 
   readonly sqlite: SqliteDatabase
+  private readonly index: SqliteDatabase
 
-  private constructor(sqlite: SqliteDatabase) {
+  private constructor(sqlite: SqliteDatabase, index: SqliteDatabase) {
     this.sqlite = sqlite
+    this.index = index
   }
 
-  static async open(source: ByteSource, opts: PagerOptions = {}) {
-    const sqlite = await SqliteDatabase.open(source, opts)
+  static async open(source: ByteSource, opts: OpenOptions = {}) {
+    const { haplotypeIndex, ...pagerOptions } = opts
+    const sqlite = await SqliteDatabase.open(source, pagerOptions)
     for (const table of ['Tags', 'Nodes', 'Paths', 'ReferenceIndex']) {
       sqlite.rootPage(table)
     }
-    const db = new GBZBase(sqlite)
+    const index = haplotypeIndex
+      ? await SqliteDatabase.open(haplotypeIndex, pagerOptions)
+      : sqlite
+    const db = new GBZBase(sqlite, index)
     const version = await db.tag('version')
     if (version !== SCHEMA_VERSION) {
       throw new SchemaVersionError(version)
     }
+    if (haplotypeIndex) {
+      for (const table of ['Tags', 'HaplotypeSamples', 'HaplotypeLengths']) {
+        index.rootPage(table)
+      }
+      db.indexTags = await readTags(index)
+      const indexed = db.indexTags.get('haplotype_index_paths')
+      const paths = await db.tag('paths')
+      if (indexed !== paths) {
+        throw new Error(
+          `haplotype index was built for ${indexed ?? 'an unknown number of'} paths but the graph has ${paths}`,
+        )
+      }
+    }
     return db
   }
 
   tags() {
-    this.tagCache ??= (async () => {
-      const tags = new Map<string, string>()
-      for await (const { values } of this.sqlite.scan('Tags')) {
-        tags.set(str(values[0], 'Tags.key'), str(values[1], 'Tags.value'))
-      }
-      return tags
-    })()
+    this.tagCache ??= readTags(this.sqlite)
     return this.tagCache
   }
 
@@ -228,14 +256,41 @@ export class GBZBase {
     return (await this.paths()).filter(p => p.name.sample === sample)
   }
 
+  async graphName() {
+    const gbzTags = new Map<string, string>()
+    for (const [key, value] of await this.tags()) {
+      if (key.startsWith('gbz_')) {
+        gbzTags.set(key.slice('gbz_'.length), value)
+      }
+    }
+    let name: GraphName = {
+      name: undefined,
+      subgraph: new Map(),
+      translation: new Map(),
+    }
+    try {
+      name = graphNameFromTags(gbzTags)
+    } catch {
+      // upstream falls back to an empty name when the tags do not parse
+    }
+    return name
+  }
+
+  async hasChainLinks() {
+    const links = await this.tag('chain_links')
+    return links !== undefined && Number(links) > 0
+  }
+
   get hasHaplotypeIndex() {
     return (
-      this.sqlite.has('HaplotypeSamples') && this.sqlite.has('HaplotypeLengths')
+      this.index.has('HaplotypeSamples') && this.index.has('HaplotypeLengths')
     )
   }
 
   async haplotypeSampleInterval() {
-    const value = await this.tag('haplotype_index_interval')
+    const value = this.indexTags
+      ? this.indexTags.get('haplotype_index_interval')
+      : await this.tag('haplotype_index_interval')
     return value === undefined ? undefined : Number(value)
   }
 
@@ -254,7 +309,7 @@ export class GBZBase {
 
   async haplotypeSamplesInRange(minHandle: number, maxHandle: number) {
     const samples: HaplotypeSample[] = []
-    for await (const key of this.sqlite.indexScanFrom('HaplotypeSamples', [
+    for await (const key of this.index.indexScanFrom('HaplotypeSamples', [
       minHandle,
       0,
     ])) {
@@ -262,7 +317,7 @@ export class GBZBase {
       if (node > maxHandle) {
         break
       }
-      const row = await this.sqlite.byRowid(
+      const row = await this.index.byRowid(
         'HaplotypeSamples',
         num(key[2], 'HaplotypeSamples rowid'),
       )
@@ -274,14 +329,11 @@ export class GBZBase {
   }
 
   async haplotypeSampleAt(node: number, offset: number) {
-    const key = await this.sqlite.indexSeekLE('HaplotypeSamples', [
-      node,
-      offset,
-    ])
+    const key = await this.index.indexSeekLE('HaplotypeSamples', [node, offset])
     if (key?.[0] !== node || key[1] !== offset) {
       return undefined
     }
-    const row = await this.sqlite.byRowid(
+    const row = await this.index.byRowid(
       'HaplotypeSamples',
       num(key[2], 'HaplotypeSamples rowid'),
     )
@@ -289,7 +341,7 @@ export class GBZBase {
   }
 
   async haplotypeLength(pathHandle: number) {
-    const row = await this.sqlite.byRowid('HaplotypeLengths', pathHandle)
+    const row = await this.index.byRowid('HaplotypeLengths', pathHandle)
     return row ? num(row[1], 'HaplotypeLengths.length') : undefined
   }
 

package/src/gbwt/record.ts

@@ -139,6 +139,29 @@ export class GbwtRecord {
     return undefined
   }
 
+  decompressArrays() {
+    const offsets = this.edges.map(e => e.offset)
+    let total = 0
+    for (const run of this.runs()) {
+      total += run.len
+    }
+    const nodes = new Int32Array(total)
+    const nextOffsets = new Int32Array(total)
+    let i = 0
+    for (const run of this.runs()) {
+      const edge = this.edge(run.value)
+      let offset = offsets[run.value]!
+      for (let k = 0; k < run.len; k++) {
+        nodes[i] = edge.node
+        nextOffsets[i] = offset
+        offset += 1
+        i += 1
+      }
+      offsets[run.value] = offset
+    }
+    return { nodes, offsets: nextOffsets }
+  }
+
   decompress(): Pos[] {
     const offsets = this.edges.map(e => e.offset)
     const result: Pos[] = []

package/src/graphName.ts

@@ -0,0 +1,94 @@
+export interface GraphName {
+  name: string | undefined
+  subgraph: Map<string, Set<string>>
+  translation: Map<string, Set<string>>
+}
+
+function parseRelationships(field: string | undefined, what: string) {
+  const relationships = new Map<string, Set<string>>()
+  if (field !== undefined) {
+    for (const rel of field.split(';')) {
+      const parts = rel.split(',')
+      const [from, to] = parts
+      if (parts.length !== 2 || !from || !to) {
+        throw new Error(`Invalid ${what} relationship: ${rel}`)
+      }
+      let targets = relationships.get(from)
+      if (!targets) {
+        targets = new Set()
+        relationships.set(from, targets)
+      }
+      targets.add(to)
+    }
+  }
+  return relationships
+}
+
+export function graphNameFromTags(tags: Map<string, string>): GraphName {
+  return {
+    name: tags.get('pggname'),
+    subgraph: parseRelationships(tags.get('subgraph'), 'subgraph'),
+    translation: parseRelationships(tags.get('translation'), 'translation'),
+  }
+}
+
+function merge(into: Map<string, Set<string>>, from: Map<string, Set<string>>) {
+  for (const [key, values] of from) {
+    let targets = into.get(key)
+    if (!targets) {
+      targets = new Set()
+      into.set(key, targets)
+    }
+    for (const value of values) {
+      targets.add(value)
+    }
+  }
+}
+
+export function subgraphName(name: string, parent: GraphName): GraphName {
+  const result: GraphName = {
+    name,
+    subgraph: new Map(),
+    translation: new Map(),
+  }
+  if (parent.name !== undefined) {
+    result.subgraph.set(name, new Set([parent.name]))
+    merge(result.subgraph, parent.subgraph)
+    merge(result.translation, parent.translation)
+  }
+  return result
+}
+
+function sortedEntries(relationships: Map<string, Set<string>>) {
+  return [...relationships.entries()]
+    .sort(([a], [b]) => (a < b ? -1 : a > b ? 1 : 0))
+    .flatMap(([from, tos]) => [...tos].sort().map(to => [from, to] as const))
+}
+
+export function gfaHeaderLines(graphName: GraphName) {
+  const lines: string[] = []
+  if (graphName.name !== undefined) {
+    lines.push(`H\tNM:Z:${graphName.name}`)
+  }
+  for (const [subgraph, supergraph] of sortedEntries(graphName.subgraph)) {
+    lines.push(`H\tSG:Z:${subgraph},${supergraph}`)
+  }
+  for (const [from, to] of sortedEntries(graphName.translation)) {
+    lines.push(`H\tTL:Z:${from},${to}`)
+  }
+  return lines
+}
+
+export async function sha256Hex(chunks: Uint8Array[]) {
+  const total = chunks.reduce((sum, chunk) => sum + chunk.length, 0)
+  const bytes = new Uint8Array(total)
+  let offset = 0
+  for (const chunk of chunks) {
+    bytes.set(chunk, offset)
+    offset += chunk.length
+  }
+  const digest = await globalThis.crypto.subtle.digest('SHA-256', bytes)
+  return [...new Uint8Array(digest)]
+    .map(byte => byte.toString(16).padStart(2, '0'))
+    .join('')
+}

package/src/index.ts

@@ -6,7 +6,7 @@ export {
   SchemaVersionError,
   formatPathName,
 } from './db.ts'
-export type { GbzPath, HaplotypeSample, PathName } from './db.ts'
+export type { GbzPath, HaplotypeSample, OpenOptions, PathName } from './db.ts'
 export type { ByteSource } from './filehandle.ts'
 export type { Pos } from './gbwt/record.ts'
 export { Subgraph } from './subgraph.ts'
@@ -16,6 +16,7 @@ export type {
   PathIdentity,
   PathPosition,
   ReferencePath,
+  SnarlOutput,
   SubgraphJson,
   SubgraphPath,
   ToJsonOptions,
@@ -23,9 +24,11 @@ export type {
 export {
   subgraphAroundNodes,
   subgraphAtOffset,
+  subgraphBetween,
   subgraphInInterval,
 } from './query.ts'
 export type { PathQuery, QueryOptions } from './query.ts'
 export { SqliteDatabase } from './sqlite/database.ts'
+export type { GraphName } from './graphName.ts'
 export { weightedLcs } from './lcs.ts'
 export * as nodes from './gbwt/node.ts'

package/src/query.ts

@@ -2,11 +2,12 @@ import { GENERIC_SAMPLE } from './db.ts'
 import { Subgraph } from './subgraph.ts'
 
 import type { GBZBase, PathName } from './db.ts'
-import type { HaplotypeOutput } from './subgraph.ts'
+import type { HaplotypeOutput, SnarlOutput } from './subgraph.ts'
 
 export interface QueryOptions {
   context?: number
   haplotypes?: HaplotypeOutput
+  snarls?: SnarlOutput
   limit?: number
 }
 
@@ -39,6 +40,7 @@ export async function subgraphAtOffset(
     reference.position.nodeOffset,
     opts.context ?? 100,
   )
+  await subgraph.extractSnarls(opts.snarls ?? 'none')
   subgraph.extractPaths(reference, opts.haplotypes ?? 'all')
   return subgraph
 }
@@ -58,6 +60,7 @@ export async function subgraphInInterval(
     end - start,
     opts.context ?? 100,
   )
+  await subgraph.extractSnarls(opts.snarls ?? 'none')
   subgraph.extractPaths(reference, opts.haplotypes ?? 'all')
   return subgraph
 }
@@ -68,12 +71,36 @@ export async function subgraphAroundNodes(
   opts: QueryOptions = {},
 ) {
   const haplotypes = opts.haplotypes ?? 'all'
+  const snarls = opts.snarls ?? 'none'
   if (haplotypes === 'reference-only') {
     throw new Error('Cannot output a reference path in a node-based query')
   }
+  if (snarls === 'overlapping' && nodes.length > 1) {
+    throw new Error(
+      'Overlapping snarls cannot be extracted for a node-based query with multiple nodes',
+    )
+  }
   const subgraph = new Subgraph(db)
   subgraph.limit = opts.limit
   await subgraph.aroundNodes(nodes, opts.context ?? 100)
+  await subgraph.extractSnarls(snarls)
+  subgraph.extractPaths(undefined, haplotypes)
+  return subgraph
+}
+
+export async function subgraphBetween(
+  db: GBZBase,
+  start: number,
+  end: number,
+  opts: Pick<QueryOptions, 'haplotypes' | 'limit'> = {},
+) {
+  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.betweenNodes(start, end)
   subgraph.extractPaths(undefined, haplotypes)
   return subgraph
 }

package/src/sqlite/btree.ts

@@ -167,10 +167,15 @@ export class BTree {
         }
       }
     } else {
+      const children: number[] = []
       for (let i = 0; i < header.cellCount; i++) {
-        yield* this.tableScan(readUint32(page, cellOffset(page, header, i)))
+        children.push(readUint32(page, cellOffset(page, header, i)))
+      }
+      children.push(header.rightChild)
+      this.pager.prefetch(children)
+      for (const child of children) {
+        yield* this.tableScan(child)
       }
-      yield* this.tableScan(header.rightChild)
     }
   }