/ concept-collection / ephys_compression_tests
Sign in
concept-collection / ephys_compression_tests
ephys_compression_tests / web-ui / src / hooks / TimeseriesDataClient.ts
175 lines · 4.9 KBCodeBlameHistory
552a4baadd web-uiJeremy Magland 1export type SupportedTypedArray =
2 | Uint8Array
3 | Uint16Array
4 | Uint32Array
5 | Int16Array
6 | Int32Array
7 | Float32Array;
9interface ChunkCache {
10 [key: number]: SupportedTypedArray;
13type DType = "uint8" | "uint16" | "uint32" | "int16" | "int32" | "float32";
15const TypedArrayConstructors = {
16 uint8: Uint8Array,
17 uint16: Uint16Array,
18 uint32: Uint32Array,
19 int16: Int16Array,
20 int32: Int32Array,
21 float32: Float32Array,
22} as const;
24export class TimeseriesDataClient {
25 private shape: number = 0;
26 private dtype: DType | null = null;
27 private chunkSize: number;
28 private cache: ChunkCache = {};
29 private inProgressFetches: { [key: number]: Promise<SupportedTypedArray> } =
30 {};
31 private datasetJsonUrl: string;
32 private datasetDataUrl: string;
34 constructor(
35 datasetJsonUrl: string,
36 datasetDataUrl: string,
37 chunkSize: number = 100000,
38 ) {
39 this.datasetJsonUrl = datasetJsonUrl;
40 this.datasetDataUrl = datasetDataUrl;
41 this.chunkSize = chunkSize;
42 }
44 static async create(
45 datasetJsonUrl: string,
46 datasetDataUrl: string,
47 chunkSize: number = 1000,
48 ): Promise<TimeseriesDataClient> {
49 const client = new TimeseriesDataClient(
50 datasetJsonUrl,
51 datasetDataUrl,
52 chunkSize,
53 );
54 await client.initialize();
55 return client;
56 }
58 private async initialize() {
59 const infoUrl = this.datasetJsonUrl;
60 const response = await fetch(infoUrl);
61 if (!response.ok) {
62 throw new Error(`Failed to fetch dataset info: ${response.statusText}`);
63 }
64 const info = await response.json();
65 this.shape = info.shape[0];
67 if (!this.isValidDType(info.dtype)) {
68 throw new Error(`Unsupported data type: ${info.dtype}`);
69 }
70 this.dtype = info.dtype;
71 }
73 private isValidDType(dtype: string): dtype is DType {
74 return dtype in TypedArrayConstructors;
75 }
77 private getChunkIndices(start: number, end: number): number[] {
78 const startChunk = Math.floor(start / this.chunkSize);
79 const endChunk = Math.floor(end / this.chunkSize);
80 const chunks: number[] = [];
81 for (let i = startChunk; i <= endChunk; i++) {
82 chunks.push(i);
83 }
84 return chunks;
85 }
87 private async fetchChunk(chunkIndex: number): Promise<SupportedTypedArray> {
88 // Return cached chunk if available
89 if (this.cache[chunkIndex]) {
90 return this.cache[chunkIndex];
91 }
93 // If this chunk is already being fetched, wait for it to complete
94 const inProgressFetch = this.inProgressFetches[chunkIndex];
95 if (inProgressFetch !== undefined) {
96 return inProgressFetch;
97 }
99 // Start new fetch and track it
100 const fetchPromise = (async () => {
101 if (!this.dtype) {
102 throw new Error("Data type not initialized");
103 }
105 const start = chunkIndex * this.chunkSize;
106 const end = Math.min(start + this.chunkSize, this.shape);
107 const url = this.datasetDataUrl;
108 const itemSize = TypedArrayConstructors[this.dtype].BYTES_PER_ELEMENT;
109 const byteStart = start * itemSize;
110 const byteEnd = end * itemSize;
112 try {
113 const response = await fetch(url, {
114 headers: {
115 Range: `bytes=${byteStart}-${byteEnd - 1}`,
116 },
117 });
118 if (!response.ok) {
119 throw new Error(`Failed to fetch chunk: ${response.statusText}`);
120 }
122 const buffer = await response.arrayBuffer();
123 const ArrayConstructor = TypedArrayConstructors[this.dtype];
124 const data = new ArrayConstructor(buffer);
125 this.cache[chunkIndex] = data;
126 return data;
127 } finally {
128 // Clean up the in-progress fetch regardless of success/failure
129 delete this.inProgressFetches[chunkIndex];
130 }
131 })();
133 // Store the promise for other requests to wait on
134 this.inProgressFetches[chunkIndex] = fetchPromise;
135 return fetchPromise;
136 }
138 async fetchRange(start: number, end: number): Promise<SupportedTypedArray> {
139 if (!this.dtype) {
140 throw new Error("Data type not initialized");
141 }
143 const chunkIndices = this.getChunkIndices(start, end);
144 const chunks = await Promise.all(
145 chunkIndices.map((idx) => this.fetchChunk(idx)),
146 );
148 // Calculate total length needed
149 const length = end - start;
150 const ArrayConstructor = TypedArrayConstructors[this.dtype];
151 const result = new ArrayConstructor(length);
153 // Copy data from chunks into result array
154 let resultOffset = 0;
155 for (let i = 0; i < chunks.length; i++) {
156 const chunk = chunks[i];
157 const chunkStart = chunkIndices[i] * this.chunkSize;
158 const copyStart = Math.max(0, start - chunkStart);
159 const copyEnd = Math.min(chunk.length, end - chunkStart);
160 const copyLength = copyEnd - copyStart;
161 result.set(chunk.subarray(copyStart, copyEnd), resultOffset);
162 resultOffset += copyLength;
163 }
165 return result;
166 }
168 getShape(): number {
169 return this.shape;
170 }
172 getDType(): DType | null {
173 return this.dtype;
174 }
moveopenescclose