/ concept-collection / ephys_compression_tests
Sign in
concept-collection / ephys_compression_tests
ephys_compression_tests / web-ui / src / hooks / TimeseriesDataClient.ts
196 lines · 5.6 KBBlameHistoryRaw
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 numChannels: number = 1;
27 private dtype: DType | null = null;
28 private chunkSize: number;
29 private cache: ChunkCache = {};
30 private inProgressFetches: { [key: number]: Promise<SupportedTypedArray> } =
31 {};
32 private datasetJsonUrl: string;
33 private datasetDataUrl: string;
35 constructor(
36 datasetJsonUrl: string,
37 datasetDataUrl: string,
38 chunkSize: number = 100000,
39 ) {
40 this.datasetJsonUrl = datasetJsonUrl;
41 this.datasetDataUrl = datasetDataUrl;
42 this.chunkSize = chunkSize;
43 }
45 static async create(
46 datasetJsonUrl: string,
47 datasetDataUrl: string,
48 chunkSize: number = 1000,
49 ): Promise<TimeseriesDataClient> {
50 const client = new TimeseriesDataClient(
51 datasetJsonUrl,
52 datasetDataUrl,
53 chunkSize,
54 );
55 await client.initialize();
56 return client;
57 }
59 private async initialize() {
60 const infoUrl = this.datasetJsonUrl;
61 const response = await fetch(infoUrl);
62 if (!response.ok) {
63 throw new Error(`Failed to fetch dataset info: ${response.statusText}`);
64 }
65 const info = await response.json();
67 // Handle multi-dimensional shape: [num_timepoints, num_channels]
68 if (Array.isArray(info.shape)) {
69 this.shape = info.shape[0];
70 this.numChannels = info.shape.length > 1 ? info.shape[1] : 1;
71 } else {
72 this.shape = info.shape;
73 this.numChannels = 1;
74 }
76 if (!this.isValidDType(info.dtype)) {
77 throw new Error(`Unsupported data type: ${info.dtype}`);
78 }
79 this.dtype = info.dtype;
80 }
82 private isValidDType(dtype: string): dtype is DType {
83 return dtype in TypedArrayConstructors;
84 }
86 private getChunkIndices(start: number, end: number): number[] {
87 const startChunk = Math.floor(start / this.chunkSize);
88 const endChunk = Math.floor(end / this.chunkSize);
89 const chunks: number[] = [];
90 for (let i = startChunk; i <= endChunk; i++) {
91 chunks.push(i);
92 }
93 return chunks;
94 }
96 private async fetchChunk(chunkIndex: number): Promise<SupportedTypedArray> {
97 // Return cached chunk if available
98 if (this.cache[chunkIndex]) {
99 return this.cache[chunkIndex];
100 }
102 // If this chunk is already being fetched, wait for it to complete
103 const inProgressFetch = this.inProgressFetches[chunkIndex];
104 if (inProgressFetch !== undefined) {
105 return inProgressFetch;
106 }
108 // Start new fetch and track it
109 const fetchPromise = (async () => {
110 if (!this.dtype) {
111 throw new Error("Data type not initialized");
112 }
114 const start = chunkIndex * this.chunkSize;
115 const end = Math.min(start + this.chunkSize, this.shape);
116 const url = this.datasetDataUrl;
117 const itemSize = TypedArrayConstructors[this.dtype].BYTES_PER_ELEMENT;
118 // Account for multi-channel data: each timepoint has numChannels values
119 const byteStart = start * this.numChannels * itemSize;
120 const byteEnd = end * this.numChannels * itemSize;
122 try {
123 const response = await fetch(url, {
124 headers: {
125 Range: `bytes=${byteStart}-${byteEnd - 1}`,
126 },
127 });
128 if (!response.ok) {
129 throw new Error(`Failed to fetch chunk: ${response.statusText}`);
130 }
132 const buffer = await response.arrayBuffer();
133 const ArrayConstructor = TypedArrayConstructors[this.dtype];
134 const data = new ArrayConstructor(buffer);
135 this.cache[chunkIndex] = data;
136 return data;
137 } finally {
138 // Clean up the in-progress fetch regardless of success/failure
139 delete this.inProgressFetches[chunkIndex];
140 }
141 })();
143 // Store the promise for other requests to wait on
144 this.inProgressFetches[chunkIndex] = fetchPromise;
145 return fetchPromise;
146 }
148 async fetchRange(start: number, end: number, channel: number = 0): Promise<SupportedTypedArray> {
149 if (!this.dtype) {
150 throw new Error("Data type not initialized");
151 }
153 if (channel < 0 || channel >= this.numChannels) {
154 throw new Error(`Invalid channel ${channel}. Must be between 0 and ${this.numChannels - 1}`);
155 }
157 const chunkIndices = this.getChunkIndices(start, end);
158 const chunks = await Promise.all(
159 chunkIndices.map((idx) => this.fetchChunk(idx)),
160 );
162 // Calculate total length needed
163 const length = end - start;
164 const ArrayConstructor = TypedArrayConstructors[this.dtype];
165 const result = new ArrayConstructor(length);
167 // Copy data from chunks into result array
168 let resultOffset = 0;
169 for (let i = 0; i < chunks.length; i++) {
170 const chunk = chunks[i];
171 const chunkStart = chunkIndices[i] * this.chunkSize;
172 const copyStart = Math.max(0, start - chunkStart);
173 const copyEnd = Math.min(chunk.length / this.numChannels, end - chunkStart);
175 // Extract the selected channel from interleaved data
176 for (let t = copyStart; t < copyEnd; t++) {
177 const sourceIdx = t * this.numChannels + channel;
178 result[resultOffset++] = chunk[sourceIdx];
179 }
180 }
182 return result;
183 }
185 getShape(): number {
186 return this.shape;
187 }
189 getDType(): DType | null {
190 return this.dtype;
191 }
193 getNumChannels(): number {
194 return this.numChannels;
195 }
moveopenescclose