From c9f6a5f988ca65c3ac5eccf20abcf0a5b98f9336 Mon Sep 17 00:00:00 2001
From: Nicolas Fernandez <cornhundred@users.noreply.github.com>
Date: Wed, 3 Dec 2025 17:25:36 -0500
Subject: [PATCH 1/8] testing opus dataset update

---
 js/obs_store/obs_store.js  |   3 +
 js/ui/dataset_dropdown.js  |  94 ++++++++++++
 js/ui/switch_dataset.js    | 286 +++++++++++++++++++++++++++++++++++++
 js/ui/ui_containers.js     |   7 +
 js/viz/landscape_ist.js    |   8 +-
 js/widget.js               |   6 +-
 src/celldega/viz/widget.py |  63 +++++++-
 7 files changed, 460 insertions(+), 7 deletions(-)
 create mode 100644 js/ui/dataset_dropdown.js
 create mode 100644 js/ui/switch_dataset.js

diff --git a/js/obs_store/obs_store.js b/js/obs_store/obs_store.js
index bc5cb596..cf4c70af 100644
--- a/js/obs_store/obs_store.js
+++ b/js/obs_store/obs_store.js
@@ -37,6 +37,9 @@ export const create_obs_store = () => {
     landscape_view: Observable('spatial'),
     umap_state: Observable(false),
     scale_bar_view_state: Observable(null),
+    // Dataset switching observables
+    current_dataset_index: Observable(0),
+    dataset_switching: Observable(false),
     // to do utilize for setProps
     deck_check: Observable({
       background_layer: true,
diff --git a/js/ui/dataset_dropdown.js b/js/ui/dataset_dropdown.js
new file mode 100644
index 00000000..0cf31dda
--- /dev/null
+++ b/js/ui/dataset_dropdown.js
@@ -0,0 +1,94 @@
+import * as d3 from 'd3';
+
+import { switch_dataset } from './switch_dataset';
+
+/**
+ * Creates a small dataset dropdown selector for switching between datasets.
+ * @param {Object} viz_state - The visualization state object
+ * @param {Object} deck_ist - The deck.gl instance
+ * @param {Object} layers_obj - The layers object
+ * @returns {HTMLElement} The dropdown container element
+ */
+export const make_dataset_dropdown = (viz_state, deck_ist, layers_obj) => {
+  const base_urls = viz_state.base_urls || [];
+
+  // Don't create dropdown if there's only one or no datasets
+  if (base_urls.length <= 1) {
+    return null;
+  }
+
+  const container = document.createElement('div');
+  container.className = 'dataset-dropdown-container';
+  container.style.display = 'inline-flex';
+  container.style.alignItems = 'center';
+  container.style.marginLeft = '5px';
+  container.style.marginRight = '5px';
+
+  const select = document.createElement('select');
+  select.className = 'dataset-dropdown';
+  select.style.width = '55px';
+  select.style.height = '18px';
+  select.style.fontSize = '9px';
+  select.style.padding = '1px 2px';
+  select.style.border = '1px solid #d3d3d3';
+  select.style.borderRadius = '3px';
+  select.style.backgroundColor = 'white';
+  select.style.cursor = 'pointer';
+  select.style.outline = 'none';
+  select.style.fontFamily = '-apple-system, BlinkMacSystemFont, "San Francisco", "Helvetica Neue", Helvetica, Arial, sans-serif';
+  select.title = 'Switch dataset';
+
+  // Add focus/hover styling
+  select.addEventListener('focus', () => {
+    select.style.borderColor = '#8797ff';
+  });
+  select.addEventListener('blur', () => {
+    select.style.borderColor = '#d3d3d3';
+  });
+
+  // Add options for each dataset
+  base_urls.forEach((dataset, index) => {
+    const option = document.createElement('option');
+    option.value = index;
+    // Truncate label if too long
+    const label = dataset.label || `Dataset ${index + 1}`;
+    option.textContent = label.length > 7 ? label.substring(0, 6) + '…' : label;
+    option.title = label; // Full label on hover
+    select.appendChild(option);
+  });
+
+  // Set initial value from obs_store
+  select.value = viz_state.obs_store.current_dataset_index.get();
+
+  // Handle change event
+  select.addEventListener('change', async (event) => {
+    const newIndex = parseInt(event.target.value, 10);
+    const currentIndex = viz_state.obs_store.current_dataset_index.get();
+
+    if (newIndex !== currentIndex && !viz_state.obs_store.dataset_switching.get()) {
+      // Show loading state
+      select.disabled = true;
+      select.style.opacity = '0.5';
+
+      try {
+        await switch_dataset(newIndex, viz_state, deck_ist, layers_obj);
+      } catch (error) {
+        console.error('Error switching dataset:', error);
+        // Revert selection on error
+        select.value = currentIndex;
+      } finally {
+        select.disabled = false;
+        select.style.opacity = '1';
+      }
+    }
+  });
+
+  // Subscribe to dataset index changes to keep dropdown in sync
+  viz_state.obs_store.current_dataset_index.subscribe((index) => {
+    select.value = index;
+  });
+
+  container.appendChild(select);
+
+  return container;
+};
diff --git a/js/ui/switch_dataset.js b/js/ui/switch_dataset.js
new file mode 100644
index 00000000..f7cb895a
--- /dev/null
+++ b/js/ui/switch_dataset.js
@@ -0,0 +1,286 @@
+import * as d3 from 'd3';
+
+import { ini_background_layer } from '../deck-gl/layers/background_layer';
+import { make_image_layers } from '../deck-gl/layers/image_layers';
+import { get_layers_list } from '../deck-gl/utils/layers_ist';
+import { set_cell_cats, set_dict_cell_cats } from '../global_variables/cat';
+import { set_cell_names_array, set_cell_name_to_index_map } from '../global_variables/cell_names_array';
+import { set_color_dict_gene } from '../global_variables/color_dict_gene';
+import { options } from '../global_variables/fetch_options';
+import { set_global_base_url } from '../global_variables/global_base_url';
+import { set_dimensions } from '../global_variables/image_dimensions';
+import { set_image_info, set_image_layer_colors, set_image_format } from '../global_variables/image_info';
+import { set_landscape_parameters } from '../global_variables/landscape_parameters';
+import { set_cluster_metadata } from '../global_variables/meta_cluster';
+import { set_meta_gene } from '../global_variables/meta_gene';
+import { get_arrow_table } from '../read_parquet/get_arrow_table';
+import { get_scatter_data } from '../read_parquet/get_scatter_data';
+import { scale_umap_data } from '../umap/scale_umap_data';
+
+import { set_image_layer_sliders } from './sliders';
+
+/**
+ * Switch to a different dataset by updating layers and state without rebuilding deck.gl.
+ * This approach avoids accumulating WebGL contexts.
+ *
+ * @param {number} newIndex - The index of the dataset to switch to
+ * @param {Object} viz_state - The visualization state object
+ * @param {Object} deck_ist - The deck.gl instance
+ * @param {Object} layers_obj - The layers object
+ */
+export const switch_dataset = async (newIndex, viz_state, deck_ist, layers_obj) => {
+  const base_urls = viz_state.base_urls || [];
+
+  if (newIndex < 0 || newIndex >= base_urls.length) {
+    console.error('Invalid dataset index:', newIndex);
+    return;
+  }
+
+  // Mark that we're switching datasets
+  viz_state.obs_store.dataset_switching.set(true);
+
+  // Temporarily disable deck rendering updates
+  viz_state.obs_store.deck_check.set({
+    ...viz_state.obs_store.deck_check.get(),
+    background_layer: false,
+    image_layers: false,
+    cell_layer: false,
+    path_layer: false,
+    trx_layer: false,
+  });
+
+  try {
+    const newDataset = base_urls[newIndex];
+    const new_base_url = newDataset.url;
+
+    // Update global base URL
+    set_global_base_url(viz_state, new_base_url);
+
+    // Load new landscape parameters
+    await set_landscape_parameters(viz_state.img, new_base_url, viz_state.aws);
+
+    const tmp_image_info = viz_state.img.landscape_parameters.image_info;
+    const image_name_for_dim = tmp_image_info[0].name;
+
+    // Update image format and info
+    set_image_format(viz_state.img, viz_state.img.landscape_parameters.image_format);
+    set_image_info(viz_state.img, tmp_image_info);
+    set_image_layer_sliders(viz_state.img);
+    set_image_layer_colors(viz_state.img.image_layer_colors, viz_state.img.image_info);
+
+    // Update dimensions
+    const tech = viz_state.img.landscape_parameters.technology;
+    if (tech !== 'Chromium' && tech !== 'point-cloud') {
+      await set_dimensions(viz_state, new_base_url, image_name_for_dim);
+    }
+
+    // Load new meta_gene data
+    viz_state.genes.gene_counts = [];
+    viz_state.genes.meta_gene = {};
+    await set_meta_gene(viz_state.genes, new_base_url, viz_state.seg.version, viz_state.aws);
+    viz_state.genes.top_gene_counts = viz_state.genes.gene_counts.slice(0, 100);
+
+    // Update gene bar data
+    viz_state.obs_store.new_gene_bar_data.set(viz_state.genes.top_gene_counts);
+
+    // Load new color dict for genes
+    await set_color_dict_gene(viz_state.genes, new_base_url, viz_state.seg.version, viz_state.aws);
+
+    // Load new cell metadata
+    let cell_url;
+    if (viz_state.seg.version === 'default') {
+      cell_url = `${new_base_url}/cell_metadata.parquet`;
+    } else {
+      cell_url = `${new_base_url}/cell_metadata_${viz_state.seg.version}.parquet`;
+    }
+
+    const cell_arrow_table = await get_arrow_table(cell_url, options.fetch, viz_state.aws);
+    set_cell_names_array(viz_state.cats, cell_arrow_table);
+    viz_state.spatial.cell_scatter_data = get_scatter_data(cell_arrow_table);
+
+    set_cell_name_to_index_map(viz_state.cats);
+
+    // Load cluster data
+    if (viz_state.cats.has_meta_cell) {
+      const inst_index = viz_state.cats.meta_cell_attr.indexOf(viz_state.cats.inst_cell_attr);
+      viz_state.cats.cell_cats = viz_state.cats.cell_names_array.map((name) => {
+        const attrs = viz_state.cats.meta_cell[name];
+        return attrs?.[inst_index] ?? 'N.A.';
+      });
+    } else {
+      const cluster_arrow_table = await get_arrow_table(
+        `${new_base_url}/cell_clusters${viz_state.seg.version && viz_state.seg.version !== 'default' ? `_${viz_state.seg.version}` : ''}/cluster.parquet`,
+        options.fetch,
+        viz_state.aws
+      );
+      set_cell_cats(viz_state.cats, cluster_arrow_table, 'cluster');
+    }
+
+    set_dict_cell_cats(viz_state.cats);
+
+    // Reset cluster metadata and counts
+    viz_state.cats.color_dict_cluster = {};
+    viz_state.cats.cluster_counts = [];
+    await set_cluster_metadata(viz_state);
+
+    // Update cell bar data
+    viz_state.obs_store.new_cell_bar_data.set(viz_state.cats.cluster_counts);
+
+    // Rebuild cell scatter data objects
+    const new_cell_names_array = cell_arrow_table.getChild('name').toArray();
+    const flatCoordinateArray = viz_state.spatial.cell_scatter_data.attributes.getPosition.value;
+    const dim = viz_state.spatial.cell_scatter_data.attributes.getPosition.size || 2;
+
+    // Update combo_data.cell
+    viz_state.combo_data.cell = new_cell_names_array.map((name, index) => ({
+      name,
+      cat: viz_state.cats.dict_cell_cats[name],
+      x: flatCoordinateArray[index * dim],
+      y: flatCoordinateArray[index * dim + 1],
+      z: dim === 3 ? flatCoordinateArray[index * dim + 2] : 0,
+    }));
+
+    // Build cell scatter data objects
+    let cell_scatter_data_objects;
+    if (viz_state.umap.has_umap) {
+      const flatCoordinateArray_umap = new Float64Array(
+        viz_state.cats.cell_names_array.flatMap((cell_id) => {
+          let coords = viz_state.umap.umap[cell_id];
+          if (!coords) {
+            coords = [0, 0];
+          }
+          return coords;
+        })
+      );
+
+      const numRows = viz_state.spatial.cell_scatter_data.length;
+      cell_scatter_data_objects = Array.from({ length: numRows }, (_, i) => ({
+        name: viz_state.cats.cell_names_array[i],
+        position:
+          dim === 3
+            ? [
+                flatCoordinateArray[i * dim],
+                flatCoordinateArray[i * dim + 1],
+                flatCoordinateArray[i * dim + 2],
+              ]
+            : [flatCoordinateArray[i * dim], flatCoordinateArray[i * dim + 1]],
+        umap: [
+          flatCoordinateArray_umap[i * 2],
+          flatCoordinateArray_umap[i * 2 + 1],
+        ],
+      }));
+
+      cell_scatter_data_objects = scale_umap_data(viz_state, cell_scatter_data_objects);
+    } else {
+      const numRows = viz_state.spatial.cell_scatter_data.length;
+      cell_scatter_data_objects = Array.from({ length: numRows }, (_, i) => ({
+        name: viz_state.cats.cell_names_array[i],
+        position:
+          dim === 3
+            ? [
+                flatCoordinateArray[i * dim],
+                flatCoordinateArray[i * dim + 1],
+                flatCoordinateArray[i * dim + 2],
+              ]
+            : [flatCoordinateArray[i * dim], flatCoordinateArray[i * dim + 1]],
+      }));
+    }
+
+    viz_state.spatial.cell_scatter_data_objects = cell_scatter_data_objects;
+
+    // Update spatial bounds
+    viz_state.spatial.x_min = d3.min(cell_scatter_data_objects.map((d) => d.position[0]));
+    viz_state.spatial.x_max = d3.max(cell_scatter_data_objects.map((d) => d.position[0]));
+    viz_state.spatial.y_min = d3.min(cell_scatter_data_objects.map((d) => d.position[1]));
+    viz_state.spatial.y_max = d3.max(cell_scatter_data_objects.map((d) => d.position[1]));
+    if (dim === 3) {
+      viz_state.spatial.z_min = d3.min(cell_scatter_data_objects.map((d) => d.position[2]));
+      viz_state.spatial.z_max = d3.max(cell_scatter_data_objects.map((d) => d.position[2]));
+    }
+
+    viz_state.spatial.center_x = (viz_state.spatial.x_max + viz_state.spatial.x_min) / 2;
+    viz_state.spatial.center_y = (viz_state.spatial.y_max + viz_state.spatial.y_min) / 2;
+    viz_state.spatial.data_width = viz_state.spatial.x_max - viz_state.spatial.x_min;
+    viz_state.spatial.data_height = viz_state.spatial.y_max - viz_state.spatial.y_min;
+
+    // Update cell layer with new data (clone, don't recreate)
+    // Disable transitions for instant dataset switching
+    layers_obj.cell_layer = layers_obj.cell_layer.clone({
+      id: `cell-layer-dataset-${newIndex}`,
+      data: cell_scatter_data_objects,
+      transitions: false,
+      updateTriggers: {
+        getPosition: [viz_state.obs_store.umap_state.get(), newIndex],
+        getFillColor: [viz_state.selection_token, newIndex],
+      },
+    });
+
+    // Update image layers by creating new ones with new base_url
+    const new_image_layers = await make_image_layers(viz_state);
+    layers_obj.image_layers = new_image_layers;
+
+    // Update background layer extent if needed
+    if (viz_state.dimensions) {
+      layers_obj.background_layer = ini_background_layer(viz_state);
+    }
+
+    // Clear transcript and path data for the new dataset
+    viz_state.genes.trx_data = [];
+    viz_state.genes.trx_names_array = [];
+    viz_state.cats.polygon_cell_names = [];
+    viz_state.combo_data.trx = [];
+
+    layers_obj.trx_layer = layers_obj.trx_layer.clone({
+      id: `trx-layer-dataset-${newIndex}`,
+      data: [],
+    });
+
+    layers_obj.path_layer = layers_obj.path_layer.clone({
+      id: `path-layer-dataset-${newIndex}`,
+      data: [],
+    });
+
+    // Clear cache for new dataset
+    if (viz_state.cache?.cell) {
+      viz_state.cache.cell.clear();
+    }
+    if (viz_state.cache?.trx) {
+      viz_state.cache.trx.clear();
+    }
+
+    // Reset selections
+    viz_state.cats.selected_cats = [];
+    viz_state.genes.selected_genes = [];
+    viz_state.obs_store.selected_cats.set([]);
+    viz_state.obs_store.selected_genes.set([]);
+
+    // Update the layers reference
+    viz_state.layers_obj = layers_obj;
+
+    // Update the current dataset index
+    viz_state.obs_store.current_dataset_index.set(newIndex);
+
+    // Re-enable deck rendering
+    viz_state.obs_store.deck_check.set({
+      ...viz_state.obs_store.deck_check.get(),
+      background_layer: true,
+      image_layers: true,
+      cell_layer: true,
+      path_layer: true,
+      trx_layer: true,
+      trx_data: true,
+      path_data: true,
+    });
+
+    // Force deck to update with new layers
+    const layers_list = get_layers_list(layers_obj, viz_state.close_up);
+    deck_ist.setProps({ layers: layers_list });
+
+  } catch (error) {
+    console.error('Error during dataset switch:', error);
+    throw error;
+  } finally {
+    // Mark switching as complete
+    viz_state.obs_store.dataset_switching.set(false);
+  }
+};
diff --git a/js/ui/ui_containers.js b/js/ui/ui_containers.js
index 49e02b60..5a97f953 100644
--- a/js/ui/ui_containers.js
+++ b/js/ui/ui_containers.js
@@ -34,6 +34,7 @@ import {
   make_bar_container,
   bar_callback_gene,
 } from './bar_plot';
+import { make_dataset_dropdown } from './dataset_dropdown';
 import { set_gene_search } from './gene_search';
 import { logo } from './logo';
 import {
@@ -444,6 +445,12 @@ export const make_ist_ui_container = (
 
     viz_state.containers.image.appendChild(spatial_toggle_container);
 
+    // Add dataset dropdown if multiple datasets are available
+    const dataset_dropdown = make_dataset_dropdown(viz_state, deck_ist, layers_obj);
+    if (dataset_dropdown) {
+      spatial_toggle_container.appendChild(dataset_dropdown);
+    }
+
     const get_slider_by_name = (img, name) => {
       return img.image_layer_sliders.filter((slider) => slider.name === name);
     };
diff --git a/js/viz/landscape_ist.js b/js/viz/landscape_ist.js
index bbc39406..b1c75f42 100644
--- a/js/viz/landscape_ist.js
+++ b/js/viz/landscape_ist.js
@@ -203,7 +203,9 @@ export const landscape_ist = async (
   rotation_x = 0,
   rotate = 0,
   max_tiles_to_view = 50,
-  scale_bar_microns_per_pixel = null
+  scale_bar_microns_per_pixel = null,
+  base_urls = [],
+  cell_name_prefix = false
 ) => {
   if (width === 0) {
     width = '100%';
@@ -261,6 +263,10 @@ export const landscape_ist = async (
 
   set_global_base_url(viz_state, base_url);
 
+  // Store multi-dataset configuration
+  viz_state.base_urls = base_urls;
+  viz_state.cell_name_prefix = cell_name_prefix;
+
   viz_state.close_up = false;
   viz_state.model = ini_model;
 
diff --git a/js/widget.js b/js/widget.js
index 790cf1fa..e04ec217 100644
--- a/js/widget.js
+++ b/js/widget.js
@@ -21,6 +21,8 @@ const render_landscape_ist = async ({ model, el }) => {
   const ini_z = model.get('ini_z');
   const ini_zoom = model.get('ini_zoom');
   const base_url = model.get('base_url');
+  const base_urls = model.get('base_urls') || [];
+  const cell_name_prefix = model.get('cell_name_prefix') || false;
   const dataset_name = model.get('dataset_name');
   const width = model.get('width');
   const height = model.get('height');
@@ -88,7 +90,9 @@ const render_landscape_ist = async ({ model, el }) => {
     rotation_x,
     rotate,
     max_tiles_to_view,
-    scale_bar_microns_per_pixel
+    scale_bar_microns_per_pixel,
+    base_urls,
+    cell_name_prefix
   );
 };
 
diff --git a/src/celldega/viz/widget.py b/src/celldega/viz/widget.py
index cd35b46b..94806e01 100644
--- a/src/celldega/viz/widget.py
+++ b/src/celldega/viz/widget.py
@@ -49,11 +49,16 @@ class Landscape(anywidget.AnyWidget):
         rotation_x (float, optional): Rotating angle around X axis for
             point-cloud views.
         token (str): The token traitlet.
-        base_url (str): The base URL for the widget.
+        base_url (str or list): The base URL(s) for the widget. Can be a single string
+            or a list of dicts with 'url' and 'label' keys for multiple datasets.
+            Example: [{'url': 'http://...', 'label': 'Dataset1'}, ...]
         rotate (float, optional): Degrees to rotate the 2D landscape visualization.
         AnnData (AnnData, optional): AnnData object to derive metadata from.
         dataset_name (str, optional): The name of the dataset to visualize. This
             will show up in the user interface bar.
+        cell_name_prefix (bool, optional): If True, cell names in adata.obs.index
+            are assumed to have a dataset prefix (e.g., "dataset-name_cell-name")
+            that should be trimmed when mapping to LandscapeFiles. Default: False.
 
     The AnnData input automatically extracts cell attributes (e.g., ``leiden``
     clusters), the corresponding colors (or derives them when missing), and any
@@ -66,6 +71,9 @@ class Landscape(anywidget.AnyWidget):
 
     technology = traitlets.Unicode("Xenium").tag(sync=True)
     base_url = traitlets.Unicode("").tag(sync=True)
+    # List of dataset configurations: [{'url': str, 'label': str}, ...]
+    base_urls = traitlets.List(trait=traitlets.Dict(), default_value=[]).tag(sync=True)
+    cell_name_prefix = traitlets.Bool(False).tag(sync=True)
     token = traitlets.Unicode("").tag(sync=True)
     creds = traitlets.Dict({}).tag(sync=True)
     max_tiles_to_view = traitlets.Int(50).tag(sync=True)
@@ -128,6 +136,33 @@ def __init__(self, **kwargs):
         if nbhd_gdf is not None and nbhd_edit:
             raise ValueError("nbhd_edit cannot be True when nbhd data is provided")
 
+        # Handle base_url which can be a string, list of strings, or list of dicts
+        # Also accept base_urls directly for convenience
+        raw_base_url = kwargs.pop("base_urls", None) or kwargs.get("base_url", "")
+        base_urls_list = []
+
+        if isinstance(raw_base_url, list):
+            # Convert list to standardized format
+            for i, item in enumerate(raw_base_url):
+                if isinstance(item, dict):
+                    # Already in dict format with 'url' and optionally 'label'
+                    url = item.get("url", "")
+                    label = item.get("label", f"Dataset {i + 1}")
+                    base_urls_list.append({"url": url, "label": label})
+                else:
+                    # Just a string URL, create a label from the index
+                    base_urls_list.append({"url": str(item), "label": f"Dataset {i + 1}"})
+
+            # Set the first URL as the primary base_url
+            if base_urls_list:
+                kwargs["base_url"] = base_urls_list[0]["url"]
+            kwargs["base_urls"] = base_urls_list
+        else:
+            # Single string URL
+            if raw_base_url:
+                base_urls_list = [{"url": raw_base_url, "label": "Dataset 1"}]
+            kwargs["base_urls"] = base_urls_list
+
         base_path = (kwargs.get("base_url") or "") + "/"
         path_transformation_matrix = base_path + "micron_to_image_transform.csv"
 
@@ -164,6 +199,9 @@ def _df_to_bytes(df):
             pq.write_table(pa.Table.from_pandas(df), buf, compression="zstd")
             return buf.getvalue()
 
+        # Get cell_name_prefix setting
+        cell_name_prefix_setting = kwargs.get("cell_name_prefix", False)
+
         if adata is not None:
             # if cell_id is in the adata.obs, use it as index
             if "cell_id" in adata.obs.columns:
@@ -174,6 +212,15 @@ def _df_to_bytes(df):
             if meta_cell_df.index.name is None:
                 meta_cell_df.index.name = "cell_id"
 
+            # If cell_name_prefix is True, trim the prefix from cell names
+            # This allows mapping to LandscapeFiles which have short names
+            if cell_name_prefix_setting:
+                # Trim prefix before first underscore from index
+                new_index = meta_cell_df.index.map(
+                    lambda x: x.split("_", 1)[1] if "_" in str(x) else x
+                )
+                meta_cell_df.index = new_index
+
             pq_meta_cell = _df_to_bytes(meta_cell_df)
 
             if "leiden" in adata.obs.columns:
@@ -202,10 +249,16 @@ def _df_to_bytes(df):
                 pq_meta_cluster = _df_to_bytes(meta_cluster_df)
 
             if "X_umap" in adata.obsm:
-                umap_df = (
-                    pd.DataFrame(adata.obsm["X_umap"], index=adata.obs.index)
-                    .reset_index()
-                    .rename(columns={"index": "cell_id", 0: "umap_0", 1: "umap_1"})
+                umap_df = pd.DataFrame(adata.obsm["X_umap"], index=adata.obs.index)
+
+                # If cell_name_prefix is True, trim the prefix from cell names
+                if cell_name_prefix_setting:
+                    umap_df.index = umap_df.index.map(
+                        lambda x: x.split("_", 1)[1] if "_" in str(x) else x
+                    )
+
+                umap_df = umap_df.reset_index().rename(
+                    columns={"index": "cell_id", 0: "umap_0", 1: "umap_1"}
                 )
                 pq_umap = _df_to_bytes(umap_df)
 

From e52dbd0ad23c79bdda49294b2975302a0e9cb186 Mon Sep 17 00:00:00 2001
From: Nicolas Fernandez <cornhundred@users.noreply.github.com>
Date: Wed, 3 Dec 2025 17:40:32 -0500
Subject: [PATCH 2/8] Enhance image layer management by disposing of old layers
 and forcing recreation with unique IDs. Update make_image_layers and
 make_simple_image_layer functions to include dataset index and cache key for
 improved layer handling.

---
 js/deck-gl/layers/image_layers.js       | 18 ++++++++++++------
 js/deck-gl/layers/simple_image_layer.js |  9 ++++++---
 js/ui/switch_dataset.js                 | 17 +++++++++++++++--
 3 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/js/deck-gl/layers/image_layers.js b/js/deck-gl/layers/image_layers.js
index e66af3e0..1350e31f 100644
--- a/js/deck-gl/layers/image_layers.js
+++ b/js/deck-gl/layers/image_layers.js
@@ -9,18 +9,21 @@ import {
 
 import { make_simple_image_layer } from './simple_image_layer';
 
-const make_image_layer = (viz_state, info) => {
+const make_image_layer = (viz_state, info, datasetIndex = 0, cacheKey = '') => {
   const { max_pyramid_zoom } = viz_state.img.landscape_parameters;
 
   const opacity = 5;
 
+  // Include dataset index and cache key in ID to force complete layer recreation
+  const layerId = `${info.button_name}-ds${datasetIndex}${cacheKey ? `-${cacheKey}` : ''}`;
+
   const image_layer = new TileLayer({
-    id: info.button_name,
+    id: layerId,
     tileSize: viz_state.dimensions.tileSize,
     refinementStrategy: 'no-overlap',
     minZoom: -7,
     maxZoom: 0,
-    maxCacheSize: 20,
+    maxCacheSize: 0, // Disable internal tile caching
     extent: [0, 0, viz_state.dimensions.width, viz_state.dimensions.height],
     getTileData: create_get_tile_data(
       viz_state.global_base_url,
@@ -40,19 +43,22 @@ const make_image_layer = (viz_state, info) => {
   return image_layer;
 };
 
-export const make_image_layers = async (viz_state) => {
+export const make_image_layers = async (viz_state, datasetIndex = 0) => {
   const { image_info } = viz_state.img;
 
+  // Generate a unique cache key to force complete layer recreation
+  const cacheKey = Date.now().toString(36);
+
   if (
     image_info.length === 1 &&
     (image_info[0].name === 'h_and_e' || image_info[0].name === 'h&e')
   ) {
-    const layer = await make_simple_image_layer(viz_state, image_info[0]);
+    const layer = await make_simple_image_layer(viz_state, image_info[0], datasetIndex, cacheKey);
     return [layer];
   }
 
   const image_layers = image_info.map((info) =>
-    make_image_layer(viz_state, info)
+    make_image_layer(viz_state, info, datasetIndex, cacheKey)
   );
   return image_layers;
 };
diff --git a/js/deck-gl/layers/simple_image_layer.js b/js/deck-gl/layers/simple_image_layer.js
index 3851c7bf..11feaa07 100644
--- a/js/deck-gl/layers/simple_image_layer.js
+++ b/js/deck-gl/layers/simple_image_layer.js
@@ -7,19 +7,22 @@ import {
   create_get_tile_data,
 } from '../utils/tiles';
 
-export const make_simple_image_layer = async (viz_state, info) => {
+export const make_simple_image_layer = async (viz_state, info, datasetIndex = 0, cacheKey = '') => {
   const { global_base_url } = viz_state;
   const { dimensions } = viz_state;
   const { landscape_parameters } = viz_state.img;
   const { image_format } = viz_state.img.landscape_parameters;
 
+  // Include dataset index and cache key in ID to force complete layer recreation
+  const layerId = `global-simple-image-layer-ds${datasetIndex}${cacheKey ? `-${cacheKey}` : ''}`;
+
   const simple_image_layer = new TileLayer({
-    id: 'global-simple-image-layer',
+    id: layerId,
     tileSize: dimensions.tileSize,
     refinementStrategy: 'no-overlap',
     minZoom: -7,
     maxZoom: 0,
-    maxCacheSize: 20,
+    maxCacheSize: 0, // Disable internal tile caching
     extent: [0, 0, dimensions.width, dimensions.height],
     getTileData: create_get_tile_data(
       global_base_url,
diff --git a/js/ui/switch_dataset.js b/js/ui/switch_dataset.js
index f7cb895a..9453d1c8 100644
--- a/js/ui/switch_dataset.js
+++ b/js/ui/switch_dataset.js
@@ -215,8 +215,21 @@ export const switch_dataset = async (newIndex, viz_state, deck_ist, layers_obj)
       },
     });
 
-    // Update image layers by creating new ones with new base_url
-    const new_image_layers = await make_image_layers(viz_state);
+    // Dispose of old image layers to clear any cached tile data
+    if (layers_obj.image_layers && Array.isArray(layers_obj.image_layers)) {
+      layers_obj.image_layers.forEach((layer) => {
+        if (layer && typeof layer.finalize === 'function') {
+          try {
+            layer.finalize();
+          } catch (e) {
+            // Ignore finalize errors
+          }
+        }
+      });
+    }
+
+    // Create completely new image layers with unique IDs to force fresh tile fetching
+    const new_image_layers = await make_image_layers(viz_state, newIndex);
     layers_obj.image_layers = new_image_layers;
 
     // Update background layer extent if needed

From 230743d37b77b450d6a76b7e725cdaa6c5fe4939 Mon Sep 17 00:00:00 2001
From: Nicolas Fernandez <cornhundred@users.noreply.github.com>
Date: Wed, 3 Dec 2025 18:17:53 -0500
Subject: [PATCH 3/8] fixing statefulness

---
 js/ui/dataset_dropdown.js  |  40 +++++++---
 js/ui/switch_dataset.js    | 154 +++++++++++++++++++++++++++++++++----
 src/celldega/viz/widget.py |  26 ++++++-
 3 files changed, 188 insertions(+), 32 deletions(-)

diff --git a/js/ui/dataset_dropdown.js b/js/ui/dataset_dropdown.js
index 0cf31dda..e78d4d1f 100644
--- a/js/ui/dataset_dropdown.js
+++ b/js/ui/dataset_dropdown.js
@@ -1,5 +1,3 @@
-import * as d3 from 'd3';
-
 import { switch_dataset } from './switch_dataset';
 
 /**
@@ -23,6 +21,7 @@ export const make_dataset_dropdown = (viz_state, deck_ist, layers_obj) => {
   container.style.alignItems = 'center';
   container.style.marginLeft = '5px';
   container.style.marginRight = '5px';
+  container.style.position = 'relative';
 
   const select = document.createElement('select');
   select.className = 'dataset-dropdown';
@@ -36,24 +35,41 @@ export const make_dataset_dropdown = (viz_state, deck_ist, layers_obj) => {
   select.style.cursor = 'pointer';
   select.style.outline = 'none';
   select.style.fontFamily = '-apple-system, BlinkMacSystemFont, "San Francisco", "Helvetica Neue", Helvetica, Arial, sans-serif';
+  select.style.transition = 'width 0.15s ease';
   select.title = 'Switch dataset';
 
-  // Add focus/hover styling
+  // Calculate max width needed for full labels
+  const max_label_length = Math.max(...base_urls.map(d => (d.label || '').length));
+  const expanded_width = Math.min(Math.max(max_label_length * 7 + 20, 80), 150);
+
+  // Add focus/hover styling - expand on focus/mousedown
+  select.addEventListener('mousedown', () => {
+    select.style.width = `${expanded_width}px`;
+  });
   select.addEventListener('focus', () => {
     select.style.borderColor = '#8797ff';
+    select.style.width = `${expanded_width}px`;
   });
   select.addEventListener('blur', () => {
     select.style.borderColor = '#d3d3d3';
+    select.style.width = '55px';
+  });
+  select.addEventListener('change', () => {
+    // Collapse after selection
+    setTimeout(() => {
+      select.style.width = '55px';
+    }, 100);
   });
 
   // Add options for each dataset
   base_urls.forEach((dataset, index) => {
     const option = document.createElement('option');
     option.value = index;
-    // Truncate label if too long
-    const label = dataset.label || `Dataset ${index + 1}`;
-    option.textContent = label.length > 7 ? label.substring(0, 6) + '…' : label;
-    option.title = label; // Full label on hover
+    // Use short_label if available, otherwise use label, otherwise default
+    const label = dataset.short_label || dataset.label || `DS-${index + 1}`;
+    const full_label = dataset.label || `Dataset ${index + 1}`;
+    option.textContent = label;
+    option.title = full_label; // Full label on hover
     select.appendChild(option);
   });
 
@@ -62,20 +78,20 @@ export const make_dataset_dropdown = (viz_state, deck_ist, layers_obj) => {
 
   // Handle change event
   select.addEventListener('change', async (event) => {
-    const newIndex = parseInt(event.target.value, 10);
-    const currentIndex = viz_state.obs_store.current_dataset_index.get();
+    const new_index = parseInt(event.target.value, 10);
+    const current_index = viz_state.obs_store.current_dataset_index.get();
 
-    if (newIndex !== currentIndex && !viz_state.obs_store.dataset_switching.get()) {
+    if (new_index !== current_index && !viz_state.obs_store.dataset_switching.get()) {
       // Show loading state
       select.disabled = true;
       select.style.opacity = '0.5';
 
       try {
-        await switch_dataset(newIndex, viz_state, deck_ist, layers_obj);
+        await switch_dataset(new_index, viz_state, deck_ist, layers_obj);
       } catch (error) {
         console.error('Error switching dataset:', error);
         // Revert selection on error
-        select.value = currentIndex;
+        select.value = current_index;
       } finally {
         select.disabled = false;
         select.style.opacity = '1';
diff --git a/js/ui/switch_dataset.js b/js/ui/switch_dataset.js
index 9453d1c8..90eab490 100644
--- a/js/ui/switch_dataset.js
+++ b/js/ui/switch_dataset.js
@@ -3,7 +3,8 @@ import * as d3 from 'd3';
 import { ini_background_layer } from '../deck-gl/layers/background_layer';
 import { make_image_layers } from '../deck-gl/layers/image_layers';
 import { get_layers_list } from '../deck-gl/utils/layers_ist';
-import { set_cell_cats, set_dict_cell_cats } from '../global_variables/cat';
+import { set_cell_cats, set_dict_cell_cats, update_cat, update_selected_cats } from '../global_variables/cat';
+import { update_cell_exp_array } from '../global_variables/cell_exp_array';
 import { set_cell_names_array, set_cell_name_to_index_map } from '../global_variables/cell_names_array';
 import { set_color_dict_gene } from '../global_variables/color_dict_gene';
 import { options } from '../global_variables/fetch_options';
@@ -13,29 +14,131 @@ import { set_image_info, set_image_layer_colors, set_image_format } from '../glo
 import { set_landscape_parameters } from '../global_variables/landscape_parameters';
 import { set_cluster_metadata } from '../global_variables/meta_cluster';
 import { set_meta_gene } from '../global_variables/meta_gene';
+import { update_selected_genes } from '../global_variables/selected_genes';
 import { get_arrow_table } from '../read_parquet/get_arrow_table';
 import { get_scatter_data } from '../read_parquet/get_scatter_data';
 import { scale_umap_data } from '../umap/scale_umap_data';
 
 import { set_image_layer_sliders } from './sliders';
 
+/**
+ * Save the current visualization state before switching datasets
+ * @param {Object} viz_state - The visualization state object
+ * @returns {Object} Saved state object
+ */
+const save_persistent_state = (viz_state) => {
+  return {
+    selected_cats: [...viz_state.cats.selected_cats],
+    selected_genes: [...viz_state.genes.selected_genes],
+    cat: viz_state.cats.cat,
+    viz_image_layers: viz_state.obs_store.viz_image_layers.get(),
+    landscape_view: viz_state.obs_store.landscape_view.get(),
+  };
+};
+
+/**
+ * Restore visualization state after switching datasets
+ * Only restores selections that exist in the new dataset
+ * @param {Object} viz_state - The visualization state object
+ * @param {Object} layers_obj - The layers object
+ * @param {Object} saved_state - Previously saved state
+ */
+const restore_persistent_state = async (viz_state, layers_obj, saved_state) => {
+  // Restore landscape view (UMAP vs spatial)
+  if (viz_state.umap.has_umap) {
+    viz_state.obs_store.landscape_view.set(saved_state.landscape_view);
+  }
+
+  // Check if selected genes exist in new dataset and restore
+  const valid_genes = saved_state.selected_genes.filter(
+    (gene) => viz_state.genes.gene_names.includes(gene)
+  );
+
+  if (valid_genes.length > 0) {
+    const inst_gene = valid_genes[0];
+
+    // Load gene expression data for the new dataset FIRST
+    // This ensures cell_exp_array is populated before we update the layer
+    await update_cell_exp_array(
+      viz_state.cats,
+      viz_state.genes,
+      viz_state.global_base_url,
+      inst_gene,
+      viz_state.seg.version,
+      viz_state.vector_name_integer,
+      viz_state.aws
+    );
+
+    // Now update cat and selections AFTER expression data is loaded
+    update_cat(viz_state.cats, inst_gene);
+    update_selected_genes(viz_state.genes, valid_genes, viz_state.obs_store);
+    update_selected_cats(viz_state.cats, valid_genes, viz_state.obs_store);
+
+    // When a gene is selected, images should be hidden (same behavior as clicking a gene)
+    viz_state.obs_store.viz_image_layers.set(false);
+
+    // Force cell layer to refresh with new expression data
+    viz_state.selection_token = (viz_state.selection_token || 0) + 1;
+    layers_obj.cell_layer = layers_obj.cell_layer.clone({
+      id: `cell-layer-gene-${inst_gene}-${viz_state.selection_token}`,
+      updateTriggers: {
+        getFillColor: [viz_state.selection_token, inst_gene],
+      },
+    });
+    viz_state.layers_obj = layers_obj;
+
+  } else {
+    // No valid gene selection, check for cluster selection
+    // Get available clusters in the new dataset
+    const available_clusters = new Set(
+      viz_state.cats.cluster_counts.map((c) => c.name)
+    );
+
+    // Filter to only clusters that exist in the new dataset
+    const valid_cats = saved_state.selected_cats.filter((cat) =>
+      available_clusters.has(cat)
+    );
+
+    if (valid_cats.length > 0) {
+      update_cat(viz_state.cats, 'cluster');
+      update_selected_cats(viz_state.cats, valid_cats, viz_state.obs_store);
+      update_selected_genes(viz_state.genes, [], viz_state.obs_store);
+    }
+
+    // Restore image layer visibility only when NOT showing gene expression
+    viz_state.obs_store.viz_image_layers.set(saved_state.viz_image_layers);
+  }
+
+  // Update persistent_state observable
+  viz_state.obs_store.persistent_state.set({
+    selected_cats: viz_state.cats.selected_cats,
+    selected_genes: viz_state.genes.selected_genes,
+    cat: viz_state.cats.cat,
+    viz_image_layers: viz_state.obs_store.viz_image_layers.get(),
+    landscape_view: saved_state.landscape_view,
+  });
+};
+
 /**
  * Switch to a different dataset by updating layers and state without rebuilding deck.gl.
  * This approach avoids accumulating WebGL contexts.
  *
- * @param {number} newIndex - The index of the dataset to switch to
+ * @param {number} new_index - The index of the dataset to switch to
  * @param {Object} viz_state - The visualization state object
  * @param {Object} deck_ist - The deck.gl instance
  * @param {Object} layers_obj - The layers object
  */
-export const switch_dataset = async (newIndex, viz_state, deck_ist, layers_obj) => {
+export const switch_dataset = async (new_index, viz_state, deck_ist, layers_obj) => {
   const base_urls = viz_state.base_urls || [];
 
-  if (newIndex < 0 || newIndex >= base_urls.length) {
-    console.error('Invalid dataset index:', newIndex);
+  if (new_index < 0 || new_index >= base_urls.length) {
+    console.error('Invalid dataset index:', new_index);
     return;
   }
 
+  // Save current state before switching (for persistence across datasets)
+  const saved_state = save_persistent_state(viz_state);
+
   // Mark that we're switching datasets
   viz_state.obs_store.dataset_switching.set(true);
 
@@ -50,8 +153,8 @@ export const switch_dataset = async (newIndex, viz_state, deck_ist, layers_obj)
   });
 
   try {
-    const newDataset = base_urls[newIndex];
-    const new_base_url = newDataset.url;
+    const new_dataset = base_urls[new_index];
+    const new_base_url = new_dataset.url;
 
     // Update global base URL
     set_global_base_url(viz_state, new_base_url);
@@ -206,12 +309,12 @@ export const switch_dataset = async (newIndex, viz_state, deck_ist, layers_obj)
     // Update cell layer with new data (clone, don't recreate)
     // Disable transitions for instant dataset switching
     layers_obj.cell_layer = layers_obj.cell_layer.clone({
-      id: `cell-layer-dataset-${newIndex}`,
+      id: `cell-layer-dataset-${new_index}`,
       data: cell_scatter_data_objects,
       transitions: false,
       updateTriggers: {
-        getPosition: [viz_state.obs_store.umap_state.get(), newIndex],
-        getFillColor: [viz_state.selection_token, newIndex],
+        getPosition: [viz_state.obs_store.umap_state.get(), new_index],
+        getFillColor: [viz_state.selection_token, new_index],
       },
     });
 
@@ -229,7 +332,7 @@ export const switch_dataset = async (newIndex, viz_state, deck_ist, layers_obj)
     }
 
     // Create completely new image layers with unique IDs to force fresh tile fetching
-    const new_image_layers = await make_image_layers(viz_state, newIndex);
+    const new_image_layers = await make_image_layers(viz_state, new_index);
     layers_obj.image_layers = new_image_layers;
 
     // Update background layer extent if needed
@@ -244,12 +347,12 @@ export const switch_dataset = async (newIndex, viz_state, deck_ist, layers_obj)
     viz_state.combo_data.trx = [];
 
     layers_obj.trx_layer = layers_obj.trx_layer.clone({
-      id: `trx-layer-dataset-${newIndex}`,
+      id: `trx-layer-dataset-${new_index}`,
       data: [],
     });
 
     layers_obj.path_layer = layers_obj.path_layer.clone({
-      id: `path-layer-dataset-${newIndex}`,
+      id: `path-layer-dataset-${new_index}`,
       data: [],
     });
 
@@ -261,17 +364,16 @@ export const switch_dataset = async (newIndex, viz_state, deck_ist, layers_obj)
       viz_state.cache.trx.clear();
     }
 
-    // Reset selections
+    // Temporarily reset selections (will be restored below)
     viz_state.cats.selected_cats = [];
     viz_state.genes.selected_genes = [];
-    viz_state.obs_store.selected_cats.set([]);
-    viz_state.obs_store.selected_genes.set([]);
+    viz_state.cats.cat = 'cluster';
 
     // Update the layers reference
     viz_state.layers_obj = layers_obj;
 
     // Update the current dataset index
-    viz_state.obs_store.current_dataset_index.set(newIndex);
+    viz_state.obs_store.current_dataset_index.set(new_index);
 
     // Re-enable deck rendering
     viz_state.obs_store.deck_check.set({
@@ -289,6 +391,24 @@ export const switch_dataset = async (newIndex, viz_state, deck_ist, layers_obj)
     const layers_list = get_layers_list(layers_obj, viz_state.close_up);
     deck_ist.setProps({ layers: layers_list });
 
+    // Restore persistent state (selected clusters, genes, UMAP view, image visibility)
+    // This allows users to compare the same selection across different datasets
+    await restore_persistent_state(viz_state, layers_obj, saved_state);
+
+    // Force deck to update with restored layers (especially if gene expression was restored)
+    const final_layers_list = get_layers_list(viz_state.layers_obj, viz_state.close_up);
+    deck_ist.setProps({ layers: final_layers_list });
+
+    // Trigger a final layer update to reflect restored state
+    viz_state.obs_store.deck_check.set({
+      ...viz_state.obs_store.deck_check.get(),
+      cell_layer: false,
+    });
+    viz_state.obs_store.deck_check.set({
+      ...viz_state.obs_store.deck_check.get(),
+      cell_layer: true,
+    });
+
   } catch (error) {
     console.error('Error during dataset switch:', error);
     throw error;
diff --git a/src/celldega/viz/widget.py b/src/celldega/viz/widget.py
index 94806e01..fdc1082c 100644
--- a/src/celldega/viz/widget.py
+++ b/src/celldega/viz/widget.py
@@ -52,6 +52,10 @@ class Landscape(anywidget.AnyWidget):
         base_url (str or list): The base URL(s) for the widget. Can be a single string
             or a list of dicts with 'url' and 'label' keys for multiple datasets.
             Example: [{'url': 'http://...', 'label': 'Dataset1'}, ...]
+            You can also pass a simple list of URL strings.
+        dataset_names (list, optional): Short names for the datasets to display in
+            the dropdown selector. Should match the length of base_urls.
+            Example: ['Brain', 'Kidney'] for two datasets.
         rotate (float, optional): Degrees to rotate the 2D landscape visualization.
         AnnData (AnnData, optional): AnnData object to derive metadata from.
         dataset_name (str, optional): The name of the dataset to visualize. This
@@ -139,6 +143,8 @@ def __init__(self, **kwargs):
         # Handle base_url which can be a string, list of strings, or list of dicts
         # Also accept base_urls directly for convenience
         raw_base_url = kwargs.pop("base_urls", None) or kwargs.get("base_url", "")
+        # Optional dataset_names for short display names in dropdown
+        dataset_names = kwargs.pop("dataset_names", None)
         base_urls_list = []
 
         if isinstance(raw_base_url, list):
@@ -148,10 +154,24 @@ def __init__(self, **kwargs):
                     # Already in dict format with 'url' and optionally 'label'
                     url = item.get("url", "")
                     label = item.get("label", f"Dataset {i + 1}")
-                    base_urls_list.append({"url": url, "label": label})
+                    short_label = item.get("short_label", f"DS-{i + 1}")
+                    base_urls_list.append({"url": url, "label": label, "short_label": short_label})
                 else:
                     # Just a string URL, create a label from the index
-                    base_urls_list.append({"url": str(item), "label": f"Dataset {i + 1}"})
+                    base_urls_list.append({
+                        "url": str(item),
+                        "label": f"Dataset {i + 1}",
+                        "short_label": f"DS-{i + 1}"
+                    })
+
+            # Apply dataset_names if provided (overrides short_label)
+            if dataset_names and isinstance(dataset_names, list):
+                for i, name in enumerate(dataset_names):
+                    if i < len(base_urls_list) and name:
+                        base_urls_list[i]["short_label"] = str(name)
+                        # Also use as label if label is default
+                        if base_urls_list[i]["label"] == f"Dataset {i + 1}":
+                            base_urls_list[i]["label"] = str(name)
 
             # Set the first URL as the primary base_url
             if base_urls_list:
@@ -160,7 +180,7 @@ def __init__(self, **kwargs):
         else:
             # Single string URL
             if raw_base_url:
-                base_urls_list = [{"url": raw_base_url, "label": "Dataset 1"}]
+                base_urls_list = [{"url": raw_base_url, "label": "Dataset 1", "short_label": "DS-1"}]
             kwargs["base_urls"] = base_urls_list
 
         base_path = (kwargs.get("base_url") or "") + "/"

From 6a437f7f41b7f168a638f039f844e3479e9a3f98 Mon Sep 17 00:00:00 2001
From: Nicolas Fernandez <cornhundred@users.noreply.github.com>
Date: Wed, 3 Dec 2025 19:45:50 -0500
Subject: [PATCH 4/8] camel_case

---
 js/obs_store/obs_store.js | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/js/obs_store/obs_store.js b/js/obs_store/obs_store.js
index cf4c70af..4f5a6c50 100644
--- a/js/obs_store/obs_store.js
+++ b/js/obs_store/obs_store.js
@@ -40,6 +40,15 @@ export const create_obs_store = () => {
     // Dataset switching observables
     current_dataset_index: Observable(0),
     dataset_switching: Observable(false),
+    // Persistent state across dataset switches
+    // This allows users to compare the same cluster/gene across datasets
+    persistent_state: Observable({
+      selected_cats: [],
+      selected_genes: [],
+      cat: 'cluster', // 'cluster' or gene name
+      viz_image_layers: true,
+      landscape_view: 'spatial', // 'spatial' or 'umap'
+    }),
     // to do utilize for setProps
     deck_check: Observable({
       background_layer: true,

From 0c94766793cdf9b8cb5adb8d0dafe99be5ebab20 Mon Sep 17 00:00:00 2001
From: Nicolas Fernandez <cornhundred@users.noreply.github.com>
Date: Wed, 3 Dec 2025 20:00:52 -0500
Subject: [PATCH 5/8] fixed first gene click obs_store state

---
 .../update_ist_landscape_from_cgm.js             | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/js/widget_interactions/update_ist_landscape_from_cgm.js b/js/widget_interactions/update_ist_landscape_from_cgm.js
index 0f8fafa6..c02f1daa 100644
--- a/js/widget_interactions/update_ist_landscape_from_cgm.js
+++ b/js/widget_interactions/update_ist_landscape_from_cgm.js
@@ -37,13 +37,9 @@ export const update_ist_landscape_from_cgm = async (
 
       update_cat(viz_state.cats, new_cat);
       update_selected_genes(viz_state.genes, [inst_gene], viz_state.obs_store);
-      // update_selected_cats(viz_state.cats, [], viz_state.obs_store);
-      update_selected_cats(
-        viz_state.cats,
-        new_cat === 'cluster' ? [] : [inst_gene],
-        viz_state.obs_store
-      );
 
+      // Load gene expression data BEFORE updating selected_cats
+      // This ensures cell_exp_array is populated before the cell layer refreshes
       await update_cell_exp_array(
         viz_state.cats,
         viz_state.genes,
@@ -54,6 +50,14 @@ export const update_ist_landscape_from_cgm = async (
         viz_state.aws
       );
 
+      // Update selected_cats after cell_exp_array has been populated
+      // This triggers the subscription that refreshes the cell layer
+      update_selected_cats(
+        viz_state.cats,
+        new_cat === 'cluster' ? [] : [inst_gene],
+        viz_state.obs_store
+      );
+
       refresh_layer(viz_state, layers_obj, 'cell_layer');
     } else if (click_type === 'col_label') {
       inst_gene = 'cluster';

From e9a64cd2dd27a80a0e9b69fd135780d5174dcfaa Mon Sep 17 00:00:00 2001
From: Nicolas Fernandez <cornhundred@users.noreply.github.com>
Date: Wed, 10 Dec 2025 16:38:19 -0500
Subject: [PATCH 6/8] ruff format

---
 src/celldega/viz/widget.py | 16 ++++++++++------
 1 file changed, 10 insertions(+), 6 deletions(-)

diff --git a/src/celldega/viz/widget.py b/src/celldega/viz/widget.py
index fdc1082c..ce4abe4e 100644
--- a/src/celldega/viz/widget.py
+++ b/src/celldega/viz/widget.py
@@ -158,11 +158,13 @@ def __init__(self, **kwargs):
                     base_urls_list.append({"url": url, "label": label, "short_label": short_label})
                 else:
                     # Just a string URL, create a label from the index
-                    base_urls_list.append({
-                        "url": str(item),
-                        "label": f"Dataset {i + 1}",
-                        "short_label": f"DS-{i + 1}"
-                    })
+                    base_urls_list.append(
+                        {
+                            "url": str(item),
+                            "label": f"Dataset {i + 1}",
+                            "short_label": f"DS-{i + 1}",
+                        }
+                    )
 
             # Apply dataset_names if provided (overrides short_label)
             if dataset_names and isinstance(dataset_names, list):
@@ -180,7 +182,9 @@ def __init__(self, **kwargs):
         else:
             # Single string URL
             if raw_base_url:
-                base_urls_list = [{"url": raw_base_url, "label": "Dataset 1", "short_label": "DS-1"}]
+                base_urls_list = [
+                    {"url": raw_base_url, "label": "Dataset 1", "short_label": "DS-1"}
+                ]
             kwargs["base_urls"] = base_urls_list
 
         base_path = (kwargs.get("base_url") or "") + "/"

From 0f9aedfe09f3f50bab811a884e1e615eccd84547 Mon Sep 17 00:00:00 2001
From: Nicolas Fernandez <cornhundred@users.noreply.github.com>
Date: Wed, 10 Dec 2025 21:59:50 -0500
Subject: [PATCH 7/8] js lint

---
 js/ui/dataset_dropdown.js | 2 +-
 js/ui/switch_dataset.js   | 5 +----
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/js/ui/dataset_dropdown.js b/js/ui/dataset_dropdown.js
index e78d4d1f..8c150a6b 100644
--- a/js/ui/dataset_dropdown.js
+++ b/js/ui/dataset_dropdown.js
@@ -89,9 +89,9 @@ export const make_dataset_dropdown = (viz_state, deck_ist, layers_obj) => {
       try {
         await switch_dataset(new_index, viz_state, deck_ist, layers_obj);
       } catch (error) {
-        console.error('Error switching dataset:', error);
         // Revert selection on error
         select.value = current_index;
+        void error;
       } finally {
         select.disabled = false;
         select.style.opacity = '1';
diff --git a/js/ui/switch_dataset.js b/js/ui/switch_dataset.js
index 90eab490..2623835d 100644
--- a/js/ui/switch_dataset.js
+++ b/js/ui/switch_dataset.js
@@ -132,7 +132,6 @@ export const switch_dataset = async (new_index, viz_state, deck_ist, layers_obj)
   const base_urls = viz_state.base_urls || [];
 
   if (new_index < 0 || new_index >= base_urls.length) {
-    console.error('Invalid dataset index:', new_index);
     return;
   }
 
@@ -326,6 +325,7 @@ export const switch_dataset = async (new_index, viz_state, deck_ist, layers_obj)
             layer.finalize();
           } catch (e) {
             // Ignore finalize errors
+            void e;
           }
         }
       });
@@ -409,9 +409,6 @@ export const switch_dataset = async (new_index, viz_state, deck_ist, layers_obj)
       cell_layer: true,
     });
 
-  } catch (error) {
-    console.error('Error during dataset switch:', error);
-    throw error;
   } finally {
     // Mark switching as complete
     viz_state.obs_store.dataset_switching.set(false);

From c69f6eaab98460e2a2bcf556f0382502af9ed1b8 Mon Sep 17 00:00:00 2001
From: Nicolas Fernandez <cornhundred@users.noreply.github.com>
Date: Wed, 10 Dec 2025 22:04:31 -0500
Subject: [PATCH 8/8] js format

---
 js/deck-gl/layers/image_layers.js       |   7 +-
 js/deck-gl/layers/simple_image_layer.js |   7 +-
 js/ui/dataset_dropdown.js               |  12 ++-
 js/ui/switch_dataset.js                 | 117 ++++++++++++++++++------
 js/ui/ui_containers.js                  |   6 +-
 5 files changed, 115 insertions(+), 34 deletions(-)

diff --git a/js/deck-gl/layers/image_layers.js b/js/deck-gl/layers/image_layers.js
index 1350e31f..931d3226 100644
--- a/js/deck-gl/layers/image_layers.js
+++ b/js/deck-gl/layers/image_layers.js
@@ -53,7 +53,12 @@ export const make_image_layers = async (viz_state, datasetIndex = 0) => {
     image_info.length === 1 &&
     (image_info[0].name === 'h_and_e' || image_info[0].name === 'h&e')
   ) {
-    const layer = await make_simple_image_layer(viz_state, image_info[0], datasetIndex, cacheKey);
+    const layer = await make_simple_image_layer(
+      viz_state,
+      image_info[0],
+      datasetIndex,
+      cacheKey
+    );
     return [layer];
   }
 
diff --git a/js/deck-gl/layers/simple_image_layer.js b/js/deck-gl/layers/simple_image_layer.js
index 11feaa07..da5bcd96 100644
--- a/js/deck-gl/layers/simple_image_layer.js
+++ b/js/deck-gl/layers/simple_image_layer.js
@@ -7,7 +7,12 @@ import {
   create_get_tile_data,
 } from '../utils/tiles';
 
-export const make_simple_image_layer = async (viz_state, info, datasetIndex = 0, cacheKey = '') => {
+export const make_simple_image_layer = async (
+  viz_state,
+  info,
+  datasetIndex = 0,
+  cacheKey = ''
+) => {
   const { global_base_url } = viz_state;
   const { dimensions } = viz_state;
   const { landscape_parameters } = viz_state.img;
diff --git a/js/ui/dataset_dropdown.js b/js/ui/dataset_dropdown.js
index 8c150a6b..dbf3aa80 100644
--- a/js/ui/dataset_dropdown.js
+++ b/js/ui/dataset_dropdown.js
@@ -34,12 +34,15 @@ export const make_dataset_dropdown = (viz_state, deck_ist, layers_obj) => {
   select.style.backgroundColor = 'white';
   select.style.cursor = 'pointer';
   select.style.outline = 'none';
-  select.style.fontFamily = '-apple-system, BlinkMacSystemFont, "San Francisco", "Helvetica Neue", Helvetica, Arial, sans-serif';
+  select.style.fontFamily =
+    '-apple-system, BlinkMacSystemFont, "San Francisco", "Helvetica Neue", Helvetica, Arial, sans-serif';
   select.style.transition = 'width 0.15s ease';
   select.title = 'Switch dataset';
 
   // Calculate max width needed for full labels
-  const max_label_length = Math.max(...base_urls.map(d => (d.label || '').length));
+  const max_label_length = Math.max(
+    ...base_urls.map((d) => (d.label || '').length)
+  );
   const expanded_width = Math.min(Math.max(max_label_length * 7 + 20, 80), 150);
 
   // Add focus/hover styling - expand on focus/mousedown
@@ -81,7 +84,10 @@ export const make_dataset_dropdown = (viz_state, deck_ist, layers_obj) => {
     const new_index = parseInt(event.target.value, 10);
     const current_index = viz_state.obs_store.current_dataset_index.get();
 
-    if (new_index !== current_index && !viz_state.obs_store.dataset_switching.get()) {
+    if (
+      new_index !== current_index &&
+      !viz_state.obs_store.dataset_switching.get()
+    ) {
       // Show loading state
       select.disabled = true;
       select.style.opacity = '0.5';
diff --git a/js/ui/switch_dataset.js b/js/ui/switch_dataset.js
index 2623835d..db381316 100644
--- a/js/ui/switch_dataset.js
+++ b/js/ui/switch_dataset.js
@@ -3,14 +3,26 @@ import * as d3 from 'd3';
 import { ini_background_layer } from '../deck-gl/layers/background_layer';
 import { make_image_layers } from '../deck-gl/layers/image_layers';
 import { get_layers_list } from '../deck-gl/utils/layers_ist';
-import { set_cell_cats, set_dict_cell_cats, update_cat, update_selected_cats } from '../global_variables/cat';
+import {
+  set_cell_cats,
+  set_dict_cell_cats,
+  update_cat,
+  update_selected_cats,
+} from '../global_variables/cat';
 import { update_cell_exp_array } from '../global_variables/cell_exp_array';
-import { set_cell_names_array, set_cell_name_to_index_map } from '../global_variables/cell_names_array';
+import {
+  set_cell_names_array,
+  set_cell_name_to_index_map,
+} from '../global_variables/cell_names_array';
 import { set_color_dict_gene } from '../global_variables/color_dict_gene';
 import { options } from '../global_variables/fetch_options';
 import { set_global_base_url } from '../global_variables/global_base_url';
 import { set_dimensions } from '../global_variables/image_dimensions';
-import { set_image_info, set_image_layer_colors, set_image_format } from '../global_variables/image_info';
+import {
+  set_image_info,
+  set_image_layer_colors,
+  set_image_format,
+} from '../global_variables/image_info';
 import { set_landscape_parameters } from '../global_variables/landscape_parameters';
 import { set_cluster_metadata } from '../global_variables/meta_cluster';
 import { set_meta_gene } from '../global_variables/meta_gene';
@@ -50,8 +62,8 @@ const restore_persistent_state = async (viz_state, layers_obj, saved_state) => {
   }
 
   // Check if selected genes exist in new dataset and restore
-  const valid_genes = saved_state.selected_genes.filter(
-    (gene) => viz_state.genes.gene_names.includes(gene)
+  const valid_genes = saved_state.selected_genes.filter((gene) =>
+    viz_state.genes.gene_names.includes(gene)
   );
 
   if (valid_genes.length > 0) {
@@ -86,7 +98,6 @@ const restore_persistent_state = async (viz_state, layers_obj, saved_state) => {
       },
     });
     viz_state.layers_obj = layers_obj;
-
   } else {
     // No valid gene selection, check for cluster selection
     // Get available clusters in the new dataset
@@ -128,7 +139,12 @@ const restore_persistent_state = async (viz_state, layers_obj, saved_state) => {
  * @param {Object} deck_ist - The deck.gl instance
  * @param {Object} layers_obj - The layers object
  */
-export const switch_dataset = async (new_index, viz_state, deck_ist, layers_obj) => {
+export const switch_dataset = async (
+  new_index,
+  viz_state,
+  deck_ist,
+  layers_obj
+) => {
   const base_urls = viz_state.base_urls || [];
 
   if (new_index < 0 || new_index >= base_urls.length) {
@@ -165,10 +181,16 @@ export const switch_dataset = async (new_index, viz_state, deck_ist, layers_obj)
     const image_name_for_dim = tmp_image_info[0].name;
 
     // Update image format and info
-    set_image_format(viz_state.img, viz_state.img.landscape_parameters.image_format);
+    set_image_format(
+      viz_state.img,
+      viz_state.img.landscape_parameters.image_format
+    );
     set_image_info(viz_state.img, tmp_image_info);
     set_image_layer_sliders(viz_state.img);
-    set_image_layer_colors(viz_state.img.image_layer_colors, viz_state.img.image_info);
+    set_image_layer_colors(
+      viz_state.img.image_layer_colors,
+      viz_state.img.image_info
+    );
 
     // Update dimensions
     const tech = viz_state.img.landscape_parameters.technology;
@@ -179,14 +201,24 @@ export const switch_dataset = async (new_index, viz_state, deck_ist, layers_obj)
     // Load new meta_gene data
     viz_state.genes.gene_counts = [];
     viz_state.genes.meta_gene = {};
-    await set_meta_gene(viz_state.genes, new_base_url, viz_state.seg.version, viz_state.aws);
+    await set_meta_gene(
+      viz_state.genes,
+      new_base_url,
+      viz_state.seg.version,
+      viz_state.aws
+    );
     viz_state.genes.top_gene_counts = viz_state.genes.gene_counts.slice(0, 100);
 
     // Update gene bar data
     viz_state.obs_store.new_gene_bar_data.set(viz_state.genes.top_gene_counts);
 
     // Load new color dict for genes
-    await set_color_dict_gene(viz_state.genes, new_base_url, viz_state.seg.version, viz_state.aws);
+    await set_color_dict_gene(
+      viz_state.genes,
+      new_base_url,
+      viz_state.seg.version,
+      viz_state.aws
+    );
 
     // Load new cell metadata
     let cell_url;
@@ -196,7 +228,11 @@ export const switch_dataset = async (new_index, viz_state, deck_ist, layers_obj)
       cell_url = `${new_base_url}/cell_metadata_${viz_state.seg.version}.parquet`;
     }
 
-    const cell_arrow_table = await get_arrow_table(cell_url, options.fetch, viz_state.aws);
+    const cell_arrow_table = await get_arrow_table(
+      cell_url,
+      options.fetch,
+      viz_state.aws
+    );
     set_cell_names_array(viz_state.cats, cell_arrow_table);
     viz_state.spatial.cell_scatter_data = get_scatter_data(cell_arrow_table);
 
@@ -204,7 +240,9 @@ export const switch_dataset = async (new_index, viz_state, deck_ist, layers_obj)
 
     // Load cluster data
     if (viz_state.cats.has_meta_cell) {
-      const inst_index = viz_state.cats.meta_cell_attr.indexOf(viz_state.cats.inst_cell_attr);
+      const inst_index = viz_state.cats.meta_cell_attr.indexOf(
+        viz_state.cats.inst_cell_attr
+      );
       viz_state.cats.cell_cats = viz_state.cats.cell_names_array.map((name) => {
         const attrs = viz_state.cats.meta_cell[name];
         return attrs?.[inst_index] ?? 'N.A.';
@@ -230,8 +268,10 @@ export const switch_dataset = async (new_index, viz_state, deck_ist, layers_obj)
 
     // Rebuild cell scatter data objects
     const new_cell_names_array = cell_arrow_table.getChild('name').toArray();
-    const flatCoordinateArray = viz_state.spatial.cell_scatter_data.attributes.getPosition.value;
-    const dim = viz_state.spatial.cell_scatter_data.attributes.getPosition.size || 2;
+    const flatCoordinateArray =
+      viz_state.spatial.cell_scatter_data.attributes.getPosition.value;
+    const dim =
+      viz_state.spatial.cell_scatter_data.attributes.getPosition.size || 2;
 
     // Update combo_data.cell
     viz_state.combo_data.cell = new_cell_names_array.map((name, index) => ({
@@ -272,7 +312,10 @@ export const switch_dataset = async (new_index, viz_state, deck_ist, layers_obj)
         ],
       }));
 
-      cell_scatter_data_objects = scale_umap_data(viz_state, cell_scatter_data_objects);
+      cell_scatter_data_objects = scale_umap_data(
+        viz_state,
+        cell_scatter_data_objects
+      );
     } else {
       const numRows = viz_state.spatial.cell_scatter_data.length;
       cell_scatter_data_objects = Array.from({ length: numRows }, (_, i) => ({
@@ -291,19 +334,35 @@ export const switch_dataset = async (new_index, viz_state, deck_ist, layers_obj)
     viz_state.spatial.cell_scatter_data_objects = cell_scatter_data_objects;
 
     // Update spatial bounds
-    viz_state.spatial.x_min = d3.min(cell_scatter_data_objects.map((d) => d.position[0]));
-    viz_state.spatial.x_max = d3.max(cell_scatter_data_objects.map((d) => d.position[0]));
-    viz_state.spatial.y_min = d3.min(cell_scatter_data_objects.map((d) => d.position[1]));
-    viz_state.spatial.y_max = d3.max(cell_scatter_data_objects.map((d) => d.position[1]));
+    viz_state.spatial.x_min = d3.min(
+      cell_scatter_data_objects.map((d) => d.position[0])
+    );
+    viz_state.spatial.x_max = d3.max(
+      cell_scatter_data_objects.map((d) => d.position[0])
+    );
+    viz_state.spatial.y_min = d3.min(
+      cell_scatter_data_objects.map((d) => d.position[1])
+    );
+    viz_state.spatial.y_max = d3.max(
+      cell_scatter_data_objects.map((d) => d.position[1])
+    );
     if (dim === 3) {
-      viz_state.spatial.z_min = d3.min(cell_scatter_data_objects.map((d) => d.position[2]));
-      viz_state.spatial.z_max = d3.max(cell_scatter_data_objects.map((d) => d.position[2]));
+      viz_state.spatial.z_min = d3.min(
+        cell_scatter_data_objects.map((d) => d.position[2])
+      );
+      viz_state.spatial.z_max = d3.max(
+        cell_scatter_data_objects.map((d) => d.position[2])
+      );
     }
 
-    viz_state.spatial.center_x = (viz_state.spatial.x_max + viz_state.spatial.x_min) / 2;
-    viz_state.spatial.center_y = (viz_state.spatial.y_max + viz_state.spatial.y_min) / 2;
-    viz_state.spatial.data_width = viz_state.spatial.x_max - viz_state.spatial.x_min;
-    viz_state.spatial.data_height = viz_state.spatial.y_max - viz_state.spatial.y_min;
+    viz_state.spatial.center_x =
+      (viz_state.spatial.x_max + viz_state.spatial.x_min) / 2;
+    viz_state.spatial.center_y =
+      (viz_state.spatial.y_max + viz_state.spatial.y_min) / 2;
+    viz_state.spatial.data_width =
+      viz_state.spatial.x_max - viz_state.spatial.x_min;
+    viz_state.spatial.data_height =
+      viz_state.spatial.y_max - viz_state.spatial.y_min;
 
     // Update cell layer with new data (clone, don't recreate)
     // Disable transitions for instant dataset switching
@@ -396,7 +455,10 @@ export const switch_dataset = async (new_index, viz_state, deck_ist, layers_obj)
     await restore_persistent_state(viz_state, layers_obj, saved_state);
 
     // Force deck to update with restored layers (especially if gene expression was restored)
-    const final_layers_list = get_layers_list(viz_state.layers_obj, viz_state.close_up);
+    const final_layers_list = get_layers_list(
+      viz_state.layers_obj,
+      viz_state.close_up
+    );
     deck_ist.setProps({ layers: final_layers_list });
 
     // Trigger a final layer update to reflect restored state
@@ -408,7 +470,6 @@ export const switch_dataset = async (new_index, viz_state, deck_ist, layers_obj)
       ...viz_state.obs_store.deck_check.get(),
       cell_layer: true,
     });
-
   } finally {
     // Mark switching as complete
     viz_state.obs_store.dataset_switching.set(false);
diff --git a/js/ui/ui_containers.js b/js/ui/ui_containers.js
index 5a97f953..defc3b7c 100644
--- a/js/ui/ui_containers.js
+++ b/js/ui/ui_containers.js
@@ -446,7 +446,11 @@ export const make_ist_ui_container = (
     viz_state.containers.image.appendChild(spatial_toggle_container);
 
     // Add dataset dropdown if multiple datasets are available
-    const dataset_dropdown = make_dataset_dropdown(viz_state, deck_ist, layers_obj);
+    const dataset_dropdown = make_dataset_dropdown(
+      viz_state,
+      deck_ist,
+      layers_obj
+    );
     if (dataset_dropdown) {
       spatial_toggle_container.appendChild(dataset_dropdown);
     }