@gmod/gbz-base 1.0.0 → 2.0.0

package/dist/subgraph.js

@@ -1,26 +1,58 @@
-import { formatPathName } from "./db.js";
-import { ENDMARKER, edgeIsCanonical, encodeNode, entryOrientation, entrySide, exitOrientation, exitSide, flipSide, isReverse, nodeId, nodeOrientation, pathIsCanonical, } from "./gbwt/node.js";
+import { ENDMARKER, edgeIsCanonical, encodeNode, entryOrientation, entrySide, exitOrientation, exitSide, flipNode, flipSide, isReverse, nodeId, nodeOrientation, pathIsCanonical, } from "./gbwt/node.js";
+import { gfaHeaderLines, sha256Hex, subgraphName } from "./graphName.js";
 import { weightedLcs } from "./lcs.js";
+import { formatPathName } from "./pathName.js";
+function sideBefore(a, b) {
+    return (a[0] < b[0] ||
+        (a[0] === b[0] && (a[1] < b[1] || (a[1] === b[1] && a[2] < b[2]))));
+}
 class SideQueue {
-    items = [];
+    heap = [];
     push(distance, node, side) {
-        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;
     }
 }
 function posKey(pos) {
@@ -35,7 +67,6 @@ export class Subgraph {
     refInterval;
     refIndexCache;
     refPrefixCache;
-    limit;
     stats = {
         orderedAlignments: 0,
         lcsAlignments: 0,
@@ -43,8 +74,12 @@ export class Subgraph {
         identificationFetches: 0,
     };
     db;
-    constructor(db) {
+    limit;
+    signal;
+    constructor(db, opts = {}) {
         this.db = db;
+        this.limit = opts.limit;
+        this.signal = opts.signal;
     }
     get nodeCount() {
         return this.records.size / 2;
@@ -78,6 +113,7 @@ export class Subgraph {
         return [...this.records.keys()].sort((a, b) => a - b);
     }
     async addNode(id) {
+        this.signal?.throwIfAborted();
         if (this.limit !== undefined && this.nodeCount >= this.limit) {
             throw new Error(`Subgraph size limit of ${this.limit} nodes exceeded`);
         }
@@ -242,6 +278,139 @@ export class Subgraph {
         }
         return { inserted, removed: toRemove.size };
     }
+    async betweenNodes(start, end) {
+        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) {
+        let inserted = 0;
+        for (const [start, end] of await this.overlappingSnarls(snarls)) {
+            inserted += await this.betweenNodes(start, end);
+        }
+        return inserted;
+    }
+    async overlappingSnarls(snarls) {
+        const result = [];
+        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;
+    }
+    isSnarlEntryInSubgraph(record) {
+        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;
+    }
+    recordReader(onFetch) {
+        const outside = new Map();
+        return async (handle) => {
+            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;
+        };
+    }
+    async findCoveringSnarl() {
+        const read = this.recordReader();
+        const isSnarlEntry = async (record) => {
+            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) => {
+            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();
+        const queue = this.sortedHandles().flatMap(handle => this.record(handle).successors());
+        let result;
+        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']) {
+                        queue.push(...(await read(encodeNode(id, orientation))).successors());
+                    }
+                }
+            }
+        }
+        return result;
+    }
     extractPaths(reference, output) {
         this.clearPaths();
         if (output === 'none') {
@@ -253,59 +422,65 @@ export class Subgraph {
         const handles = this.sortedHandles();
         const successors = new Map();
         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)) {
-                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;
         for (const handle of handles) {
             const entries = successors.get(handle);
-            entries.forEach((entry, offset) => {
-                if (entry.hasPredecessor) {
-                    return;
-                }
-                let curr = { node: handle, offset };
-                let isRef = false;
-                const path = [];
-                const positions = [];
-                let len = 0;
-                while (curr) {
-                    if (curr.node === refPos?.handle &&
-                        curr.offset === refPos.gbwtOffset) {
-                        this.refId = this.paths.length;
-                        refOffset = path.length;
-                        isRef = true;
+            for (let offset = 0; offset < entries.nodes.length; offset++) {
+                if (entries.hasPredecessor[offset] === 0) {
+                    let currNode = handle;
+                    let currOffset = offset;
+                    let isRef = false;
+                    const path = [];
+                    const positions = [];
+                    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 = step.nodes[currOffset];
+                        const nextOffset = 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 = 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) {
@@ -380,23 +555,9 @@ export class Subgraph {
                 starts.set(posKey(first), index);
             }
         });
-        const outside = new Map();
-        const recordAt = async (handle) => {
-            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) => {
             if (pos.node >= minHandle && pos.node <= maxHandle) {
                 return samples.get(posKey(pos));
@@ -451,6 +612,7 @@ export class Subgraph {
             let current = start;
             let pos;
             while (anchor === undefined) {
+                this.signal?.throwIfAborted();
                 if (current !== undefined) {
                     if (visited.has(current)) {
                         break;
@@ -729,23 +891,16 @@ export class Subgraph {
                 last -= 1;
             }
             const identity = info.identity;
-            const strand = info.path.some(handle => isReverse(handle)) &&
+            const walkStrand = info.path.some(handle => isReverse(handle)) &&
                 !info.path.some(handle => !isReverse(handle))
                 ? '-'
                 : '+';
-            result.push({
-                pathHandle: identity?.pathHandle,
-                name: identity?.name,
+            const span = {
                 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
@@ -754,11 +909,115 @@ 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, opts = {}) {
+    canonicalEdges(id) {
+        const edges = [];
+        for (const orientation of ['forward', 'reverse']) {
+            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 = [];
+        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 = {}) {
+        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) => (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, name, end, cigarString) => {
+            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 = {}) {
+        const cigar = opts.cigar ?? false;
         const handles = this.sortedHandles();
         const nodes = handles
             .filter(handle => !isReverse(handle))