higlass-zarr-datafetchers
Version:
Fetch HiGlass track data from Zarr stores
242 lines (205 loc) • 9.16 kB
JavaScript
import {
open as zarrOpen,
root as zarrRoot,
get as zarrGet,
slice,
FetchStore,
} from 'zarrita';
function multivecChunksToTileDenseArray(chunks, tileShape, isRow) {
// Allocate a Float32Array for the tile (with length tile_size).
const fullTileLength = (isRow ? tileShape[1] : tileShape[0] * tileShape[1]);
const fullTileArray = new Float32Array(fullTileLength);
// Fill in the data for each sample and chunk.
let offset = 0;
if(isRow) {
// Single row, no need to iterate over samples.
for (const chunk of chunks) {
const chunkData = chunk.data;
fullTileArray.set(chunkData, offset);
offset += chunkData.length;
}
} else {
// Multi-row, need to iterate over samples.
const numSamples = tileShape[0];
for (let sampleI = 0; sampleI < numSamples; sampleI++) {
for (const chunk of chunks) {
// Zarrita returns strided arrays.
const chunkData = chunk.data.subarray(sampleI * chunk.stride[0], (sampleI+1) * chunk.stride[0]);
fullTileArray.set(chunkData, offset);
offset += chunkData.length;
}
}
}
return fullTileArray;
}
const ZarrMultivecDataFetcher = function ZarrMultivecDataFetcher(HGC, ...args) {
if (!new.target) {
throw new Error(
'Uncaught TypeError: Class constructor cannot be invoked without "new"',
);
}
const { slugid } = HGC.libraries;
const {
absToChr,
parseChromsizesRows,
genomicRangeToChromosomeChunks,
DenseDataExtrema1D,
minNonZero,
maxNonZero,
} = HGC.utils;
class ZarrMultivecDataFetcherClass {
constructor(dataConfig) {
this.dataConfig = dataConfig;
this.trackUid = slugid.nice();
if (dataConfig.hasStoreRoot) {
// We cannot pass the storeRoot directly, since HiGlass
// stringifies the dataConfig.
// Reference: https://github.com/higlass/higlass/blob/e6ddd547ec9e6b2298e2c42c8dbf664039cbf098/app/scripts/data-fetchers/DataFetcher.js#L108
// Instead, we store the storeRoot on a global, keyed by its URL.
this.storeRoot = Promise.resolve(ZarrMultivecDataFetcher.urlToStoreRoot[dataConfig.url]);
} else if (dataConfig.url) {
// console.assert(dataConfig.url.endsWith('.zarr'));
// S3 bucket must have a CORS policy to allow reading from any origin.
const { url, options = {} } = dataConfig;
this.store = new FetchStore(url, options);
this.storeRoot = Promise.resolve(zarrRoot(this.store));
}
if(dataConfig.row !== undefined) {
this.row = dataConfig.row;
}
}
tilesetInfo(callback) {
this.tilesetInfoLoading = true;
// Use the tileset_info stored as JSON in file.zarr/.zattrs
return this.storeRoot
.then(root => zarrOpen(root))
.then(grp => {
const attrs = grp.attrs;
this.tilesetInfoLoading = false;
const chromSizes = attrs.multiscales.map(d => ([d.name, d.metadata.chromsize]));
const finalChrom = attrs.multiscales[attrs.multiscales.length - 1];
const maxPos = finalChrom.metadata.chromoffset + finalChrom.metadata.chromsize;
const tileSize = attrs.shape[1];
const retVal = {
...attrs,
shape: [attrs.shape[1], attrs.shape[0]],
chromSizes,
tile_size: tileSize,
max_width: maxPos,
min_pos: [0],
max_pos: [maxPos],
max_zoom: Math.ceil(Math.log(maxPos / tileSize) / Math.log(2)),
};
if (callback) {
callback(retVal);
}
return retVal;
})
.catch(err => {
this.tilesetInfoLoading = false;
if (callback) {
callback({
error: `Error parsing zarr multivec: ${err}`,
});
}
});
}
fetchTilesDebounced(receivedTiles, tileIds) {
const tiles = {};
const validTileIds = [];
const tilePromises = [];
for (const tileId of tileIds) {
const parts = tileId.split('.');
const z = parseInt(parts[0], 10);
const x = parseInt(parts[1], 10);
if (Number.isNaN(x) || Number.isNaN(z)) {
console.warn('Invalid tile zoom or position:', z, x);
continue;
}
validTileIds.push(tileId);
tilePromises.push(this.tile(z, x, tileId));
}
Promise.all(tilePromises).then(values => {
for (let i = 0; i < values.length; i++) {
const validTileId = validTileIds[i];
tiles[validTileId] = values[i];
tiles[validTileId].tilePositionId = validTileId;
}
receivedTiles(tiles);
});
return tiles;
}
tile(z, x, tileId) {
const { storeRoot } = this;
return this.tilesetInfo().then(tsInfo => {
// const multiscales = tsInfo.multiscales;
// Adapted from clodius.tiles.multivec.get_single_tile
// Reference: https://github.com/higlass/clodius/blob/develop/clodius/tiles/multivec.py#L66
// z is the index of the resolution that should be selected.
// Resolution is size of each bin (except for the last bin in each chromosome).
const resolution = +tsInfo.resolutions[z];
const tileSize = +tsInfo.tile_size;
const binSize = resolution;
// Where in the data does the tile start and end?
const tileStart = x * tileSize * resolution;
const tileEnd = tileStart + tileSize * resolution;
// chromSizes is an array of "tuples" [ ["chr1", 1000], ... ]
const chromSizes = tsInfo.chromSizes;
// Adapted from clodius.tiles.multivec.get_tile
// Reference: https://github.com/higlass/clodius/blob/develop/clodius/tiles/multivec.py#L110
const chromInfo = parseChromsizesRows(chromSizes);
const [chrStart, chrStartPos] = absToChr(tileStart, chromInfo);
const [chrEnd, chrEndPos] = absToChr(tileEnd, chromInfo);
const genomicStart = { chr: chrStart, pos: chrStartPos };
const genomicEnd = { chr: chrEnd, pos: chrEndPos };
// Using the [genomicStart, genomicEnd] range, get an array of "chromosome chunks",
// where each chunk range starts and ends with the same chromosome.
// Start a new chromosome chunk at each chromosome boundary.
const chrChunks = genomicRangeToChromosomeChunks(
chromSizes,
genomicStart,
genomicEnd,
binSize,
tileSize,
);
// Get the zarr data for each chromosome chunk,
// since data for each chromosome is stored in a separate zarr array.
return Promise.all(
chrChunks.map(([chrName, zStart, zEnd]) => {
return storeRoot
.then(root => zarrOpen(root.resolve(`/chromosomes/${chrName}/${resolution}/`), { kind: "array" }))
.then(arr => (this.row !== undefined
? zarrGet(arr, [this.row, slice(zStart, zEnd)])
: zarrGet(arr, [null, slice(zStart, zEnd)])
));
}),
).then(chunks => {
const dense = multivecChunksToTileDenseArray(chunks, [tsInfo.shape[1], tsInfo.shape[0]], this.row !== undefined);
return Promise.resolve({
dense,
denseDataExtrema: new DenseDataExtrema1D(dense),
dtype: 'float32',
min_value: Math.min.apply(null, dense),
max_value: Math.max.apply(null, dense),
minNonZero: minNonZero(dense),
maxNonZero: maxNonZero(dense),
server: null,
size: 1,
shape: tsInfo.shape,
tileId,
tilePos: [x],
tilePositionId: tileId,
tilesetUid: null,
zoomLevel: z,
});
});
});
}
} // end class
return new ZarrMultivecDataFetcherClass(...args);
} // end function wrapper
ZarrMultivecDataFetcher.config = {
type: 'zarr-multivec',
};
ZarrMultivecDataFetcher.urlToStoreRoot = {};
export default ZarrMultivecDataFetcher;