@jax-data-science/sci-image-visualizer 0.8.2 → 0.9.0

package/fesm2022/onnx-sam.worker.js

@@ -1,14 +1,41 @@
 // src/lib/toolbar/segmentation/onnx-sam.worker.ts
+import { LruCache } from "napari-js";
 import * as ort2 from "onnxruntime-web";
 
 // src/lib/toolbar/segmentation/sam-onnx-core.ts
 import * as ort from "onnxruntime-web";
 var PIXEL_MEAN = [123.675, 116.28, 103.53];
 var PIXEL_STD = [58.395, 57.12, 57.375];
-var MODEL_CACHE = "sam-onnx";
-async function fetchModel(url, onProgress) {
-  const cache = typeof caches !== "undefined" ? await caches.open(MODEL_CACHE).catch(() => null) : null;
-  const hit = cache ? await cache.match(url) : null;
+var MODEL_CACHE_PREFIX = "sam-onnx";
+var MODEL_CACHE = `${MODEL_CACHE_PREFIX}-v2`;
+var staleCachesPurged = null;
+function purgeStaleModelCaches() {
+  staleCachesPurged ?? (staleCachesPurged = caches.keys().then(
+    (names) => Promise.all(
+      names.filter((n) => n.startsWith(MODEL_CACHE_PREFIX) && n !== MODEL_CACHE).map((n) => caches.delete(n))
+    )
+  ).then(() => void 0).catch(() => void 0));
+  return staleCachesPurged;
+}
+async function openModelCache() {
+  if (typeof caches === "undefined")
+    return null;
+  await purgeStaleModelCaches();
+  return caches.open(MODEL_CACHE).catch(() => null);
+}
+function modelCacheKey(url, revision) {
+  if (!revision)
+    return url;
+  return `${url}${url.includes("?") ? "&" : "?"}siv-model-rev=${encodeURIComponent(revision)}`;
+}
+async function cacheModel(cache, url, key, body) {
+  await cache.delete(url, { ignoreSearch: true }).catch(() => void 0);
+  await cache.put(key, body).catch(() => void 0);
+}
+async function fetchModel(url, onProgress, revision) {
+  const key = modelCacheKey(url, revision);
+  const cache = await openModelCache();
+  const hit = cache ? await cache.match(key) : null;
   if (hit) {
     onProgress?.(1);
     return await hit.arrayBuffer();
@@ -21,7 +48,7 @@ async function fetchModel(url, onProgress) {
     const buf = await resp.arrayBuffer();
     onProgress?.(1);
     if (cache)
-      await cache.put(url, new Response(buf)).catch(() => void 0);
+      await cacheModel(cache, url, key, new Response(buf));
     return buf;
   }
   const reader = resp.body.getReader();
@@ -42,7 +69,7 @@ async function fetchModel(url, onProgress) {
     off += c.length;
   }
   if (cache) {
-    await cache.put(url, new Response(out.slice(), { headers: { "content-length": String(received) } })).catch(() => void 0);
+    await cacheModel(cache, url, key, new Response(out, { headers: { "content-length": String(received) } }));
   }
   onProgress?.(1);
   return out.buffer;
@@ -149,7 +176,7 @@ var encoder = null;
 var decoder = null;
 var inputSize = 1024;
 var nextToken = 1;
-var embeddings = /* @__PURE__ */ new Map();
+var embeddings = new LruCache(3);
 function post(msg, transfer = []) {
   self.postMessage(msg, transfer);
 }
@@ -158,8 +185,12 @@ async function onLoad(msg) {
   inputSize = msg.inputSize || 1024;
   const hasGpu = typeof navigator !== "undefined" && "gpu" in navigator;
   const eps = msg.encoderProviders && msg.encoderProviders.length ? msg.encoderProviders : hasGpu ? ["webgpu", "wasm"] : ["wasm"];
-  const encBuf = await fetchModel(msg.encoderUrl, (f) => post({ id: msg.id, type: "progress", fraction: f }));
-  const decBuf = await fetchModel(msg.decoderUrl);
+  const encBuf = await fetchModel(
+    msg.encoderUrl,
+    (f) => post({ id: msg.id, type: "progress", fraction: f }),
+    msg.revision
+  );
+  const decBuf = await fetchModel(msg.decoderUrl, void 0, msg.revision);
   encoder = await ort2.InferenceSession.create(encBuf, { executionProviders: eps });
   decoder = await ort2.InferenceSession.create(decBuf, { executionProviders: ["wasm"] });
   post({ id: msg.id, type: "loaded" });
@@ -171,9 +202,15 @@ async function onEmbed(msg) {
   const e = await runEncoder(encoder, data, msg.width, msg.height, inputSize);
   const token = nextToken++;
   embeddings.set(token, e);
-  if (embeddings.size > 3)
-    embeddings.delete(embeddings.keys().next().value);
-  post({ id: msg.id, type: "embedded", token, scale: e.scale, imageWidth: e.imageWidth, imageHeight: e.imageHeight, dims: e.dims });
+  post({
+    id: msg.id,
+    type: "embedded",
+    token,
+    scale: e.scale,
+    imageWidth: e.imageWidth,
+    imageHeight: e.imageHeight,
+    dims: e.dims
+  });
 }
 async function onDecode(msg) {
   if (!decoder)
@@ -182,10 +219,9 @@ async function onDecode(msg) {
   if (!e)
     throw new Error("SAM embedding expired; re-encode the image.");
   const r = await runDecoder(decoder, e, msg.prompt);
-  post(
-    { id: msg.id, type: "decoded", width: r.width, height: r.height, iou: r.iou, buffer: r.mask.buffer },
-    [r.mask.buffer]
-  );
+  post({ id: msg.id, type: "decoded", width: r.width, height: r.height, iou: r.iou, buffer: r.mask.buffer }, [
+    r.mask.buffer
+  ]);
 }
 self.onmessage = async (ev) => {
   const msg = ev.data;

package/fesm2022/spatial-math.worker.js

@@ -0,0 +1,471 @@
+// src/lib/spatial/spatial-density.ts
+function gaussianKernel(sigma) {
+  const radius = Math.max(1, Math.ceil(sigma * 3));
+  const kernel = new Float32Array(radius * 2 + 1);
+  let sum = 0;
+  for (let i = -radius; i <= radius; i++) {
+    const v = Math.exp(-(i * i) / (2 * sigma * sigma));
+    kernel[i + radius] = v;
+    sum += v;
+  }
+  for (let i = 0; i < kernel.length; i++)
+    kernel[i] /= sum;
+  return { kernel, radius };
+}
+function blurVolumeAxis(field, w, h, d, axis, sigmaVox) {
+  if (!(sigmaVox > 0.01))
+    return;
+  const { kernel, radius } = gaussianKernel(sigmaVox);
+  const len = axis === 0 ? w : axis === 1 ? h : d;
+  const stride = axis === 0 ? 1 : axis === 1 ? w : w * h;
+  const lines = axis === 0 ? h * d : axis === 1 ? w * d : w * h;
+  const line = new Float32Array(len);
+  for (let l = 0; l < lines; l++) {
+    let base;
+    if (axis === 0)
+      base = l % h * w + Math.floor(l / h) * w * h;
+    else if (axis === 1)
+      base = l % w + Math.floor(l / w) * w * h;
+    else
+      base = l;
+    for (let i = 0; i < len; i++)
+      line[i] = field[base + i * stride];
+    for (let i = 0; i < len; i++) {
+      let acc = 0;
+      for (let k = -radius; k <= radius; k++) {
+        const j = Math.min(len - 1, Math.max(0, i + k));
+        acc += line[j] * kernel[k + radius];
+      }
+      field[base + i * stride] = acc;
+    }
+  }
+}
+function sampledPlanes(obs, grid) {
+  const cover = new Float32Array(grid.depth);
+  const z = obs.z;
+  if (!z)
+    return { cover: cover.fill(1), first: 0, last: grid.depth - 1 };
+  let first = grid.depth;
+  let last = -1;
+  for (let i = 0; i < obs.count; i++) {
+    const k = Math.floor(z[i] / grid.voxelSize[2]);
+    if (k < 0 || k >= grid.depth)
+      continue;
+    cover[k] = 1;
+    if (k < first)
+      first = k;
+    if (k > last)
+      last = k;
+  }
+  return { cover, first, last };
+}
+function blur1d(profile, sigmaVox) {
+  const out = profile.slice();
+  blurVolumeAxis(out, 1, 1, out.length, 2, sigmaVox);
+  return out;
+}
+function rasterizeDensity(obs, grid, opts) {
+  const { width: w, height: h, depth: d, voxelSize } = grid;
+  const field = new Float32Array(w * h * d);
+  const idx = opts.indices;
+  const n = idx ? idx.length : obs.count;
+  const z = obs.z;
+  let placed = 0;
+  for (let k = 0; k < n; k++) {
+    const i = idx ? idx[k] : k;
+    const vx = Math.floor(obs.x[i] / voxelSize[0]);
+    const vy = Math.floor(obs.y[i] / voxelSize[1]);
+    const vz = z ? Math.floor(z[i] / voxelSize[2]) : 0;
+    if (vx < 0 || vx >= w || vy < 0 || vy >= h || vz < 0 || vz >= d)
+      continue;
+    field[(vz * h + vy) * w + vx] += 1;
+    placed++;
+  }
+  if (!placed)
+    return null;
+  const sigmaVox = [
+    opts.sigma[0] / voxelSize[0],
+    opts.sigma[1] / voxelSize[1],
+    opts.sigma[2] / voxelSize[2]
+  ];
+  blurVolumeAxis(field, w, h, d, 0, sigmaVox[0]);
+  blurVolumeAxis(field, w, h, d, 1, sigmaVox[1]);
+  blurVolumeAxis(field, w, h, d, 2, sigmaVox[2]);
+  const { cover, first, last } = sampledPlanes(obs, grid);
+  const smoothedCover = blur1d(cover, sigmaVox[2]);
+  let max = 0;
+  const plane = w * h;
+  for (let k = 0; k < d; k++) {
+    if (k < first || k > last) {
+      field.fill(0, k * plane, (k + 1) * plane);
+      continue;
+    }
+    const scale = 1 / Math.max(smoothedCover[k], 0.5);
+    for (let i = k * plane; i < (k + 1) * plane; i++) {
+      const v = field[i] * scale;
+      field[i] = v;
+      if (v > max)
+        max = v;
+    }
+  }
+  if (!(max > 0))
+    return null;
+  const out = new Uint8Array(field.length);
+  for (let i = 0; i < field.length; i++)
+    out[i] = Math.round(field[i] / max * 255);
+  return out;
+}
+
+// src/lib/spatial/stats.ts
+function quantiles(values, ps, opts = {}) {
+  const positive = opts.filter === "positive";
+  const where = opts.where;
+  const n = values.length;
+  const keep = (i) => {
+    const v = values[i];
+    if (positive ? !(v > 0) : !Number.isFinite(v))
+      return false;
+    return !where || where[i] > 0;
+  };
+  let count = 0;
+  for (let i = 0; i < n; i++)
+    if (keep(i))
+      count++;
+  if (count === 0)
+    return null;
+  const step = opts.sampleSize && count >= 2 * opts.sampleSize ? Math.floor(count / opts.sampleSize) : 1;
+  const sample = new Float64Array(Math.ceil(count / step));
+  let seen = 0;
+  let k = 0;
+  for (let i = 0; i < n && k < sample.length; i++) {
+    if (!keep(i))
+      continue;
+    if (seen++ % step === 0)
+      sample[k++] = values[i];
+  }
+  const sorted = sample.subarray(0, k).sort();
+  const last = sorted.length - 1;
+  return ps.map((p) => {
+    const f = Math.max(0, Math.min(1, Number.isFinite(p) ? p : 0));
+    return sorted[Math.round(f * last)];
+  });
+}
+function quantile(values, p, opts = {}) {
+  return quantiles(values, [p], opts)?.[0] ?? null;
+}
+
+// src/lib/spatial/spatial-expression.ts
+function expressionField(obs, opts) {
+  const { width: w, height: h, step, sigma, values, indices } = opts;
+  if (w <= 0 || h <= 0)
+    return null;
+  const [sx, sy] = opts.ref?.scale ?? [1, 1];
+  const [tx, ty] = opts.ref?.translate ?? [0, 0];
+  const num = new Float32Array(w * h);
+  const den = new Float32Array(w * h);
+  const n = indices ? indices.length : obs.count;
+  let placed = 0;
+  for (let k = 0; k < n; k++) {
+    const i = indices ? indices[k] : k;
+    const e = values[i];
+    if (!Number.isFinite(e))
+      continue;
+    const px = Math.floor((obs.x[i] * sx + tx) / step);
+    const py = Math.floor((obs.y[i] * sy + ty) / step);
+    if (px < 0 || px >= w || py < 0 || py >= h)
+      continue;
+    const at = py * w + px;
+    num[at] += e;
+    den[at] += 1;
+    placed++;
+  }
+  if (!placed)
+    return null;
+  for (const axis of [0, 1]) {
+    blurVolumeAxis(num, w, h, 1, axis, sigma);
+    blurVolumeAxis(den, w, h, 1, axis, sigma);
+  }
+  const mean = new Float32Array(w * h);
+  let lo = Infinity;
+  let hi = -Infinity;
+  for (let i = 0; i < mean.length; i++) {
+    const d = den[i];
+    if (d <= 0)
+      continue;
+    const m = num[i] / d;
+    mean[i] = m;
+    if (m < lo)
+      lo = m;
+    if (m > hi)
+      hi = m;
+  }
+  if (!Number.isFinite(lo))
+    return null;
+  const supportScale = quantile(den, 0.5, { filter: "positive" }) || 0;
+  return {
+    width: w,
+    height: h,
+    mean,
+    support: den,
+    supportScale,
+    range: hi > lo ? [lo, hi] : [lo, lo + 1]
+  };
+}
+function expressionVolume(obs, grid, opts) {
+  const { width: w, height: h, depth: d, voxelSize } = grid;
+  if (w <= 0 || h <= 0 || d <= 0)
+    return null;
+  const num = new Float32Array(w * h * d);
+  const den = new Float32Array(w * h * d);
+  const idx = opts.indices;
+  const n = idx ? idx.length : obs.count;
+  const z = obs.z;
+  let placed = 0;
+  let first = d;
+  let last = -1;
+  for (let k = 0; k < n; k++) {
+    const i = idx ? idx[k] : k;
+    const e = opts.values[i];
+    if (!Number.isFinite(e))
+      continue;
+    const vx = Math.floor(obs.x[i] / voxelSize[0]);
+    const vy = Math.floor(obs.y[i] / voxelSize[1]);
+    const vz = z ? Math.floor(z[i] / voxelSize[2]) : 0;
+    if (vx < 0 || vx >= w || vy < 0 || vy >= h || vz < 0 || vz >= d)
+      continue;
+    const at = (vz * h + vy) * w + vx;
+    num[at] += e;
+    den[at] += 1;
+    placed++;
+    if (vz < first)
+      first = vz;
+    if (vz > last)
+      last = vz;
+  }
+  if (!placed)
+    return null;
+  const sigmaVox = [
+    opts.sigma[0] / voxelSize[0],
+    opts.sigma[1] / voxelSize[1],
+    opts.sigma[2] / voxelSize[2]
+  ];
+  for (const axis of [0, 1]) {
+    blurVolumeAxis(num, w, h, d, axis, sigmaVox[axis]);
+    blurVolumeAxis(den, w, h, d, axis, sigmaVox[axis]);
+  }
+  if (opts.interpolate) {
+    blurVolumeAxis(num, w, h, d, 2, sigmaVox[2]);
+    blurVolumeAxis(den, w, h, d, 2, sigmaVox[2]);
+    const plane = w * h;
+    for (let k = 0; k < d; k++) {
+      if (k >= first && k <= last)
+        continue;
+      num.fill(0, k * plane, (k + 1) * plane);
+      den.fill(0, k * plane, (k + 1) * plane);
+    }
+  }
+  const mean = new Float32Array(num.length);
+  let lo = Infinity;
+  let hi = -Infinity;
+  for (let i = 0; i < mean.length; i++) {
+    const support = den[i];
+    if (support <= 0)
+      continue;
+    const m = num[i] / support;
+    mean[i] = m;
+    if (m < lo)
+      lo = m;
+    if (m > hi)
+      hi = m;
+  }
+  if (!Number.isFinite(lo))
+    return null;
+  return {
+    width: w,
+    height: h,
+    depth: d,
+    mean,
+    support: den,
+    range: hi > lo ? [lo, hi] : [lo, lo + 1]
+  };
+}
+
+// src/lib/contracts/spatial-dataset.contract.ts
+var NO_CATEGORY = 65535;
+
+// src/lib/spatial/spatial-heatmap.ts
+var HEATMAP_MIN_CELLS = 3;
+function heatmapMatrix(genes, groups, opts = {}) {
+  const nGroups = groups.categories.length;
+  if (genes.length === 0 || nGroups === 0)
+    return null;
+  const minCells = opts.minCells ?? HEATMAP_MIN_CELLS;
+  const indices = opts.indices;
+  const n = indices ? indices.length : groups.codes.length;
+  const groupCount = new Uint32Array(nGroups);
+  for (let k = 0; k < n; k++) {
+    const i = indices ? indices[k] : k;
+    const code = groups.codes[i];
+    if (code === NO_CATEGORY || code >= nGroups)
+      continue;
+    groupCount[code]++;
+  }
+  const keep = [];
+  for (let g = 0; g < nGroups; g++)
+    if (groupCount[g] >= minCells)
+      keep.push(g);
+  if (keep.length === 0)
+    return null;
+  const rows = genes.map((g) => g.name);
+  const cols = keep.map((g) => groups.categories[g]);
+  const counts = Uint32Array.from(keep, (g) => groupCount[g]);
+  const values = new Float32Array(rows.length * cols.length).fill(NaN);
+  for (let r = 0; r < genes.length; r++) {
+    const gene = genes[r];
+    const sum = new Float64Array(nGroups);
+    const measured = new Uint32Array(nGroups);
+    for (let k = 0; k < n; k++) {
+      const i = indices ? indices[k] : k;
+      const code = groups.codes[i];
+      if (code === NO_CATEGORY || code >= nGroups)
+        continue;
+      const v = gene.values[i];
+      if (!Number.isFinite(v))
+        continue;
+      sum[code] += v;
+      measured[code]++;
+    }
+    for (let c = 0; c < keep.length; c++) {
+      const g = keep[c];
+      if (measured[g] === 0)
+        continue;
+      values[r * cols.length + c] = sum[g] / measured[g];
+    }
+  }
+  if (opts.zScore !== false)
+    zScoreRows(values, rows.length, cols.length);
+  let hiddenCols = 0;
+  let outCols = cols;
+  let outValues = values;
+  let outCounts = counts;
+  if (opts.maxCols !== void 0 && cols.length > opts.maxCols) {
+    const strength = cols.map((_, c) => {
+      let peak = 0;
+      for (let r = 0; r < rows.length; r++) {
+        const v = Math.abs(values[r * cols.length + c]);
+        if (Number.isFinite(v) && v > peak)
+          peak = v;
+      }
+      return { c, peak };
+    });
+    strength.sort((a, b) => b.peak - a.peak || a.c - b.c);
+    const picked = strength.slice(0, opts.maxCols).map((s) => s.c).sort((a, b) => a - b);
+    hiddenCols = cols.length - picked.length;
+    outCols = picked.map((c) => cols[c]);
+    outCounts = Uint32Array.from(picked, (c) => counts[c]);
+    outValues = new Float32Array(rows.length * picked.length);
+    for (let r = 0; r < rows.length; r++) {
+      for (let k = 0; k < picked.length; k++) {
+        outValues[r * picked.length + k] = values[r * cols.length + picked[k]];
+      }
+    }
+  }
+  let lo = Infinity;
+  let hi = -Infinity;
+  for (let i = 0; i < outValues.length; i++) {
+    const v = outValues[i];
+    if (!Number.isFinite(v))
+      continue;
+    if (v < lo)
+      lo = v;
+    if (v > hi)
+      hi = v;
+  }
+  if (!Number.isFinite(lo))
+    return null;
+  return {
+    rows,
+    hiddenCols,
+    cols: outCols,
+    values: outValues,
+    counts: outCounts,
+    range: hi > lo ? [lo, hi] : [lo, lo + 1]
+  };
+}
+function zScoreRows(values, nRows, nCols) {
+  for (let r = 0; r < nRows; r++) {
+    const base = r * nCols;
+    let sum = 0;
+    let count = 0;
+    for (let c = 0; c < nCols; c++) {
+      const v = values[base + c];
+      if (!Number.isFinite(v))
+        continue;
+      sum += v;
+      count++;
+    }
+    if (count === 0)
+      continue;
+    const mean = sum / count;
+    let variance = 0;
+    for (let c = 0; c < nCols; c++) {
+      const v = values[base + c];
+      if (!Number.isFinite(v))
+        continue;
+      variance += (v - mean) * (v - mean);
+    }
+    const sd = Math.sqrt(variance / count);
+    for (let c = 0; c < nCols; c++) {
+      const v = values[base + c];
+      if (!Number.isFinite(v))
+        continue;
+      values[base + c] = sd > 0 ? (v - mean) / sd : 0;
+    }
+  }
+}
+
+// src/lib/workers/spatial-math-ops.ts
+function runSpatialMath(request) {
+  const r = request;
+  switch (r.op) {
+    case "expressionField":
+      return expressionField(r.obs, r.opts);
+    case "expressionVolume":
+      return expressionVolume(r.obs, r.grid, r.opts);
+    case "rasterizeDensity":
+      return rasterizeDensity(r.obs, r.grid, r.opts);
+    case "heatmapMatrix":
+      return heatmapMatrix(r.genes, r.groups, r.opts);
+    default:
+      throw new Error(`unknown spatial-math op "${r.op}"`);
+  }
+}
+function resultTransferables(result) {
+  const found = /* @__PURE__ */ new Set();
+  const visit = (v) => {
+    if (ArrayBuffer.isView(v)) {
+      if (v.buffer instanceof ArrayBuffer)
+        found.add(v.buffer);
+    } else if (v && typeof v === "object") {
+      for (const value of Object.values(v))
+        visit(value);
+    }
+  };
+  visit(result);
+  return [...found];
+}
+
+// src/lib/workers/spatial-math.worker.ts
+self.onmessage = (event) => {
+  const { id } = event.data;
+  let reply;
+  let transfer = [];
+  try {
+    const result = runSpatialMath(event.data);
+    reply = { id, ok: true, result };
+    transfer = resultTransferables(result);
+  } catch (err) {
+    reply = { id, ok: false, error: err?.message ?? String(err) };
+  }
+  self.postMessage(reply, transfer);
+};

package/fesm2022/tsne.worker.js

@@ -4810,14 +4810,14 @@ function reductionTilePlan(kernel, strideMap, canStorePartial = true) {
   const tileSize = commonTileSize(kernel.size, strideMap, 4, 4);
   if (tileSize === null)
     return null;
-  const tileRows = tiledRows(kernel.size, tileSize);
+  const tileRows2 = tiledRows(kernel.size, tileSize);
   const tileVectors = tiledColumns(tileSize, 4);
   return {
     tileSize,
-    tileRows,
+    tileRows: tileRows2,
     tileVectors,
     tileK: canStorePartial ? divisorAtMost(kernel.reduction.size, TILED_SIMD_K) : kernel.reduction.size,
-    microRows: microTile(tileRows, Math.floor(kernel.size / tileSize), TILED_SIMD_MICRO_ROWS),
+    microRows: microTile(tileRows2, Math.floor(kernel.size / tileSize), TILED_SIMD_MICRO_ROWS),
     microVectors: microTileWithTail(tileVectors, TILED_SIMD_MICRO_VECTORS)
   };
 }
@@ -4830,13 +4830,13 @@ function reductionKTilePlan(kernel, strideMap) {
   if (tileSize === null)
     return null;
   const kUnroll = divisorAtMost(kernel.reduction.size / 4, K_SIMD_MAX_UNROLL);
-  const tileRows = tiledRows(kernel.size, tileSize);
+  const tileRows2 = tiledRows(kernel.size, tileSize);
   const tileCols = tiledColumns(tileSize, 1);
   return {
     tileSize,
-    tileRows,
+    tileRows: tileRows2,
     tileCols,
-    microRows: microTile(tileRows, Math.floor(kernel.size / tileSize), K_SIMD_MICRO_ROWS),
+    microRows: microTile(tileRows2, Math.floor(kernel.size / tileSize), K_SIMD_MICRO_ROWS),
     microCols: microTile(tileCols, tileSize, K_SIMD_MICRO_COLS),
     kUnroll
   };
@@ -5838,9 +5838,9 @@ tune.exp: ${tune.exp}`);
       cg.i32.gt_u();
       cg.br_if(0);
     }, emitBody);
-    const emitRowTileLoop = (rowOffset, rowTileBase, tileRows, tileSize, emitBody) => emitLoopWithBreaks(() => {
+    const emitRowTileLoop = (rowOffset, rowTileBase, tileRows2, tileSize, emitBody) => emitLoopWithBreaks(() => {
       cg.local.get(rowOffset);
-      cg.i32.const(tileRows);
+      cg.i32.const(tileRows2);
       cg.i32.ge_u();
       cg.br_if(0);
       cg.local.get(rowTileBase);
@@ -18752,7 +18752,13 @@ var init_dist = __esm({
 });
 
 // src/lib/spatial/tsne-gpu.ts
-var TILE = 2048;
+var MAX_TILE_ROWS = 2048;
+var MIN_TILE_ROWS = 64;
+var TILE_BUDGET_BYTES = MAX_TILE_ROWS * 5e3 * 3 * 4 * 3;
+function tileRows(nObs, dims) {
+  const perRow = Math.max(1, nObs * dims * 4 * 3);
+  return Math.max(MIN_TILE_ROWS, Math.min(MAX_TILE_ROWS, Math.floor(TILE_BUDGET_BYTES / perRow)));
+}
 async function createGpuRepulsion() {
   const jax = await Promise.resolve().then(() => (init_dist(), dist_exports));
   const np = jax.numpy;
@@ -18774,13 +18780,11 @@ async function createGpuRepulsion() {
       const Y = np.array(flat, { shape: [nObs, dims], dtype: jax.DType.Float32 });
       const rep2 = new Float64Array(nObs * dims);
       let z = 0;
-      for (let s = 0; s < nObs; s += TILE) {
-        const rows = Math.min(TILE, nObs - s);
+      const tile2 = tileRows(nObs, dims);
+      for (let s = 0; s < nObs; s += tile2) {
+        const rows = Math.min(tile2, nObs - s);
         const yb = Y.ref.slice([s, s + rows]);
-        const diff2 = np.subtract(
-          np.reshape(yb, [rows, 1, dims]),
-          np.reshape(Y.ref, [1, nObs, dims])
-        );
+        const diff2 = np.subtract(np.reshape(yb, [rows, 1, dims]), np.reshape(Y.ref, [1, nObs, dims]));
         const w = np.divide(1, np.add(1, np.sum(np.square(diff2.ref), 2)));
         z += Number(await np.sum(w.ref).item()) - rows;
         const weighted = np.multiply(np.reshape(np.square(w), [rows, nObs, 1]), diff2);
@@ -18972,8 +18976,21 @@ async function tsneEmbed(x, nObs, nDims, options = {}) {
   } = options;
   const k = Math.max(1, Math.min(nObs - 1, Math.round(NEIGHBOUR_FACTOR * perplexity)));
   const effective = Math.min(perplexity, Math.max(1, Math.floor(k / NEIGHBOUR_FACTOR)));
+  const stopped = () => ({
+    embedding: new Float64Array(nObs * dims),
+    perplexity: effective,
+    neighbours: k,
+    completed: false
+  });
+  if (shouldStop?.())
+    return stopped();
   const knn = knnGraph(x, nObs, nDims, k);
   const { rowPtr, colIdx, val } = affinities(knn, nObs, effective);
+  if (shouldStop) {
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    if (shouldStop())
+      return stopped();
+  }
   const random = rng(seed);
   const y = new Float64Array(nObs * dims);
   for (let i = 0; i < y.length; i++)