@gmod/gbz-base 1.0.0 → 2.0.0

package/src/db.ts

@@ -1,23 +1,28 @@
 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 { formatPathName, pathNameFor, toPathQuery } from './pathName.ts'
+import { subgraphInInterval } from './query.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 { PathName, PathRef } from './pathName.ts'
+import type { QueryOptions } from './query.ts'
 import type { PagerOptions } from './sqlite/pager.ts'
 import type { SqlValue } from './sqlite/record.ts'
+import type { HaplotypeAlignment, Subgraph } from './subgraph.ts'
 
-export interface PathName {
-  sample: string
-  contig: string
-  haplotype: number
-  fragment: number
-}
-
-export const GENERIC_SAMPLE = '_gbwt_ref'
 export const SCHEMA_VERSION = 'GBZ-base version 4'
 
+export interface PathFragment {
+  path: GbzPath
+  start: number
+  end: number
+}
+
 export class SchemaVersionError extends Error {
   override name = 'SchemaVersionError'
 
@@ -33,10 +38,6 @@ export class SchemaVersionError extends Error {
   }
 }
 
-export function formatPathName(name: PathName, end: number) {
-  return `${name.sample}#${name.haplotype}#${name.contig}[${name.fragment}-${end}]`
-}
-
 export interface HaplotypeSample {
   node: number
   offset: number
@@ -144,37 +145,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
+  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 +255,182 @@ export class GBZBase {
     return (await this.paths()).filter(p => p.name.sample === sample)
   }
 
+  private async pathsNamed(ref: PathRef) {
+    const name = pathNameFor(toPathQuery(ref), 0)
+    return (await this.paths())
+      .filter(
+        p =>
+          p.name.sample === name.sample &&
+          p.name.contig === name.contig &&
+          p.name.haplotype === name.haplotype,
+      )
+      .sort((a, b) => a.name.fragment - b.name.fragment)
+  }
+
+  async hasPath(ref: PathRef) {
+    return (await this.pathsNamed(ref)).length > 0
+  }
+
+  private pathLengths = new Map<number, Promise<number>>()
+
+  pathLength(handle: number) {
+    let length = this.pathLengths.get(handle)
+    if (!length) {
+      length = this.walkPathLength(handle)
+      this.pathLengths.set(handle, length)
+      length.catch(() => this.pathLengths.delete(handle))
+    }
+    return length
+  }
+
+  private async walkPathLength(handle: number) {
+    const indexed = this.hasHaplotypeIndex
+      ? await this.haplotypeLength(handle)
+      : undefined
+    if (indexed === undefined) {
+      const last = await this.indexedPosition(handle, Number.MAX_SAFE_INTEGER)
+      if (!last) {
+        const path = await this.getPath(handle)
+        throw new Error(
+          `Path ${path ? formatPathName(path.name, path.name.fragment) : handle} has not been indexed for random access`,
+        )
+      }
+      let length = last.pathOffset
+      let pos = last.pos
+      for (;;) {
+        const record = await this.getRecord(pos.node)
+        if (!record) {
+          throw new Error(`Node ${pos.node} does not exist in the graph`)
+        }
+        length += record.sequenceLen
+        const next = record.gbwt().lf(pos.offset)
+        if (!next || next.node === ENDMARKER) {
+          return length
+        }
+        pos = next
+      }
+    }
+    return indexed
+  }
+
+  async pathFragmentsForRange(
+    ref: PathRef,
+    start: number,
+    end: number,
+  ): Promise<PathFragment[]> {
+    if (end <= start) {
+      return []
+    }
+    const ordered = await this.pathsNamed(ref)
+    const covering = ordered.filter(p => p.name.fragment <= start).slice(-1)
+    const within = ordered.filter(
+      p => p.name.fragment > start && p.name.fragment < end,
+    )
+    const fragments = await Promise.all(
+      [...covering, ...within].map(async path => ({
+        path,
+        start: path.name.fragment,
+        end: path.name.fragment + (await this.pathLength(path.handle)),
+      })),
+    )
+    return fragments.filter(f => start < f.end)
+  }
+
+  private async subgraphForFragment(
+    fragment: PathFragment,
+    start: number,
+    end: number,
+    opts: QueryOptions,
+  ) {
+    const lo = Math.max(start, fragment.start)
+    const hi = Math.min(end, fragment.end)
+    let subgraph: Subgraph | undefined
+    if (hi > lo) {
+      const { sample, contig, haplotype } = fragment.path.name
+      subgraph = await subgraphInInterval(
+        this,
+        { sample, contig, haplotype },
+        lo,
+        hi,
+        opts,
+      )
+      const haplotypes = opts.haplotypes ?? 'all'
+      const named = haplotypes === 'all' || haplotypes === 'distinct'
+      if (this.hasHaplotypeIndex && named) {
+        await subgraph.identifyPaths()
+      }
+    }
+    return subgraph
+  }
+
+  async getSubgraphForRange(
+    ref: PathRef,
+    start: number,
+    end: number,
+    opts: QueryOptions = {},
+  ) {
+    const [fragment] = await this.pathFragmentsForRange(ref, start, end)
+    return fragment
+      ? this.subgraphForFragment(fragment, start, end, opts)
+      : undefined
+  }
+
+  async getAlignmentsForRange(
+    ref: PathRef,
+    start: number,
+    end: number,
+    opts: QueryOptions = {},
+  ) {
+    const result: HaplotypeAlignment[] = []
+    for (const fragment of await this.pathFragmentsForRange(ref, start, end)) {
+      const subgraph = await this.subgraphForFragment(
+        fragment,
+        start,
+        end,
+        opts,
+      )
+      if (subgraph) {
+        result.push(...subgraph.alignments())
+      }
+    }
+    return result
+  }
+
+  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 +449,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 +457,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 +469,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 +481,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

@@ -1,31 +1,36 @@
-export {
-  GBZBase,
-  GENERIC_SAMPLE,
-  GbzRecord,
-  SCHEMA_VERSION,
-  SchemaVersionError,
-  formatPathName,
+export { GBZBase, GbzRecord, SCHEMA_VERSION, SchemaVersionError } from './db.ts'
+export type {
+  GbzPath,
+  HaplotypeSample,
+  OpenOptions,
+  PathFragment,
 } from './db.ts'
-export type { GbzPath, HaplotypeSample, PathName } from './db.ts'
+export { GENERIC_SAMPLE, formatPathName, parsePathName } from './pathName.ts'
+export type { PathName, PathQuery, PathRef } from './pathName.ts'
 export type { ByteSource } from './filehandle.ts'
 export type { Pos } from './gbwt/record.ts'
 export { Subgraph } from './subgraph.ts'
 export type {
+  AlignmentSpan,
   HaplotypeAlignment,
   HaplotypeOutput,
   PathIdentity,
   PathPosition,
   ReferencePath,
+  SnarlOutput,
   SubgraphJson,
+  SubgraphOptions,
+  SubgraphOutputOptions,
   SubgraphPath,
-  ToJsonOptions,
 } from './subgraph.ts'
 export {
   subgraphAroundNodes,
   subgraphAtOffset,
+  subgraphBetween,
   subgraphInInterval,
 } from './query.ts'
-export type { PathQuery, QueryOptions } from './query.ts'
+export type { 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/pathName.ts

@@ -0,0 +1,54 @@
+export interface PathName {
+  sample: string
+  contig: string
+  haplotype: number
+  fragment: number
+}
+
+export interface PathQuery {
+  sample?: string
+  contig: string
+  haplotype?: number
+}
+
+export type PathRef = string | PathQuery
+
+export const GENERIC_SAMPLE = '_gbwt_ref'
+
+export function formatPathName(name: PathName, end: number) {
+  return `${name.sample}#${name.haplotype}#${name.contig}[${name.fragment}-${end}]`
+}
+
+const PAN_SN = /^([^#]+)#(\d+)#([^#]+?)(?:\[\d+-\d+\])?$/
+const BARE_CONTIG = /^([^#]+?)(?:\[\d+-\d+\])?$/
+
+export function parsePathName(name: string): PathQuery {
+  const panSn = PAN_SN.exec(name)
+  if (panSn) {
+    return {
+      sample: panSn[1]!,
+      haplotype: Number(panSn[2]),
+      contig: panSn[3]!,
+    }
+  }
+  const bare = BARE_CONTIG.exec(name)
+  if (bare) {
+    return { contig: bare[1]! }
+  }
+  throw new Error(
+    `"${name}" is not a path name; expected "sample#haplotype#contig" or a bare contig`,
+  )
+}
+
+export function toPathQuery(ref: PathRef) {
+  return typeof ref === 'string' ? parsePathName(ref) : ref
+}
+
+export function pathNameFor(query: PathQuery, fragment: number): PathName {
+  return {
+    sample: query.sample ?? GENERIC_SAMPLE,
+    contig: query.contig,
+    haplotype: query.haplotype ?? 0,
+    fragment,
+  }
+}

package/src/query.ts

@@ -1,28 +1,18 @@
-import { GENERIC_SAMPLE } from './db.ts'
+import { pathNameFor } from './pathName.ts'
 import { Subgraph } from './subgraph.ts'
 
-import type { GBZBase, PathName } from './db.ts'
-import type { HaplotypeOutput } from './subgraph.ts'
+import type { GBZBase } from './db.ts'
+import type { PathQuery } from './pathName.ts'
+import type { HaplotypeOutput, SnarlOutput } from './subgraph.ts'
 
-export interface QueryOptions {
-  context?: number
-  haplotypes?: HaplotypeOutput
-  limit?: number
-}
+export type { PathQuery } from './pathName.ts'
 
-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 interface QueryOptions {
+  context?: number | undefined
+  haplotypes?: HaplotypeOutput | undefined
+  snarls?: SnarlOutput | undefined
+  limit?: number | undefined
+  signal?: AbortSignal | undefined
 }
 
 export async function subgraphAtOffset(
@@ -31,14 +21,14 @@ export async function subgraphAtOffset(
   offset: number,
   opts: QueryOptions = {},
 ) {
-  const subgraph = new Subgraph(db)
-  subgraph.limit = opts.limit
-  const reference = await subgraph.pathPosition(pathName(query, offset))
+  const subgraph = new Subgraph(db, opts)
+  const reference = await subgraph.pathPosition(pathNameFor(query, offset))
   await subgraph.aroundPosition(
     reference.position.handle,
     reference.position.nodeOffset,
     opts.context ?? 100,
   )
+  await subgraph.extractSnarls(opts.snarls ?? 'none')
   subgraph.extractPaths(reference, opts.haplotypes ?? 'all')
   return subgraph
 }
@@ -50,14 +40,14 @@ export async function subgraphInInterval(
   end: number,
   opts: QueryOptions = {},
 ) {
-  const subgraph = new Subgraph(db)
-  subgraph.limit = opts.limit
-  const reference = await subgraph.pathPosition(pathName(query, start))
+  const subgraph = new Subgraph(db, opts)
+  const reference = await subgraph.pathPosition(pathNameFor(query, start))
   await subgraph.aroundInterval(
     reference.position,
     end - start,
     opts.context ?? 100,
   )
+  await subgraph.extractSnarls(opts.snarls ?? 'none')
   subgraph.extractPaths(reference, opts.haplotypes ?? 'all')
   return subgraph
 }
@@ -68,12 +58,34 @@ 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')
   }
-  const subgraph = new Subgraph(db)
-  subgraph.limit = opts.limit
+  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, opts)
   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' | 'signal'> = {},
+) {
+  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, opts)
+  await subgraph.betweenNodes(start, end)
   subgraph.extractPaths(undefined, haplotypes)
   return subgraph
 }