diff --git a/.gitignore b/.gitignore index 205839f..2f8b3d3 100644 --- a/.gitignore +++ b/.gitignore @@ -20,3 +20,6 @@ __pycache__/ exclude/ runs/ dev/ +.vscode +models +*.npy \ No newline at end of file diff --git a/curation/app.py b/curation/app.py new file mode 100644 index 0000000..7adec17 --- /dev/null +++ b/curation/app.py @@ -0,0 +1,717 @@ +import os, sys +import numpy as np +import pandas as pd +import panel as pn +from sklearn.cluster import KMeans +import json +from collections import OrderedDict +sys.path.insert(0, 'curation') +from umap_visualiser import UMAPVisualiser +from box_viewer import BoxViewer +from hist_viewer import HistViewer +from sklearn.linear_model import SGDClassifier +import h5py +from umap import UMAP +import time + +pn.extension() + +# Mapping from HistViewer property names to UMAPVisualiser view modes +PROPERTY_TO_VIEW_MODE = { + 'Footprint Size': 'size', + 'ROI Shot Noise': 'snr', + 'Session': 'session', + 'Edge Cells': 'edge', + 'Mean Intensity': 'intensity', + 'Mean Correlation': 'correlation', + 'Probability': 'prob', + 'Peak Value': 'peak', + 'Voxel SNR': 'vox', + 'Contamination Factor': 'contamination' +} + + +class AppOrchestrator: + def __init__(self, umap_file_path, nn_features_path, hdf5_path="data.h5"): + self.umap_file_path = umap_file_path + self.hdf5_path = hdf5_path + self.classifications_file = None + self.sample_size = 50000 + self.use_sampling = True + self.nn_features_path = nn_features_path + + # Shared data state + self.umap_embedding = None + self.nn_features = None + self.full_data = None + self.properties = { + 'Footprint Size': None, + 'Mean Intensity': None, + 'Mean Correlation': None, + 'ROI Shot Noise': None, + 'Edge Cells': None, + 'Session': None, + 'Voxel SNR': None, + 'Probability': None, + 'Peak Value': None, + 'Contamination Factor': None + } + self.UMAPtime = None + self.display_data = None + self.classifications = None + self.sample_indices = None + self.cell_probs = None + self.labelled_features_idx = [] + self.labels = [] + self.available_features = OrderedDict() + self.curation_features_to_use = ['PCA'] + + # Components + self.umap_visualiser = None + self.box_viewer = None + self.hist_viewer = None + self.linear_model = SGDClassifier(loss='log_loss') + + # Shared widgets + self.cluster_slider = pn.widgets.IntSlider( + name='Number of Clusters', + start=3, end=100, value=20, + width=200 + ) + + self.cluster_button = pn.widgets.Button( + name='Update Clustering', + button_type='default', + width=200 + ) + + # Classification buttons + self.classify_cluster_cell_button = pn.widgets.Button( + name='CELL', + button_type='danger', + width=95 + ) + + self.classify_cluster_not_cell_button = pn.widgets.Button( + name='NOT CELL', + button_type='danger', + button_style='outline', + width=95 + ) + + self.reset_button = pn.widgets.Button( + name='Reset Classifications', + button_type='warning', + width=200 + ) + + self.save_button = pn.widgets.Button( + name='Save Classifications', + button_type='success', + width=200 + ) + + self.view_toggle = pn.widgets.ToggleGroup( + name='UMAP view mode', + options={'Cluster': 'cluster', 'Property': 'property'}, + value='cluster', + behavior='radio', + width=200 + ) + + # Status text + self.status_text = pn.pane.Markdown(f"**Status:** Loading {umap_file_path}...", width=200) + + # Set up callbacks + self.cluster_button.on_click(self.update_clusters) + self.classify_cluster_cell_button.on_click(self.classify_cluster_as_cell) + self.classify_cluster_not_cell_button.on_click(self.classify_cluster_as_not_cell) + self.reset_button.on_click(self.reset_classifications) + self.save_button.on_click(self.save_classifications) + + # Initialize + self.load_data() + self.create_components() + + # Set up view toggle callback + self.view_toggle.param.watch(self._on_view_toggle_change, 'value') + + self.UMAPbutton = pn.widgets.Button( + name=f'Rerun UMAP ({self.UMAPtime}s)' if self.UMAPtime is not None else 'Rerun UMAP', + button_type='primary', + width=200 + ) + self.UMAPbutton.on_click(self.recompute_umap) + + def load_data(self): + """Load and prepare all data with sampling coordination""" + print("Loading data...") + try: + if not os.path.exists(self.umap_file_path): + self.status_text.object = f"**Error:** File not found: {self.umap_file_path}" + return + + # Load UMAP embeddings + self.umap_embedding = np.load(self.umap_file_path) + curation_dir = os.path.dirname(self.umap_file_path) + + if self.umap_embedding.ndim != 2 or self.umap_embedding.shape[1] != 2: + self.status_text.object = "**Error:** UMAP file must be 2D with shape (n_points, 2)" + return + + n_points = self.umap_embedding.shape[0] + + # Load NN features + self.nn_features = np.load(self.nn_features_path) + if self.nn_features.shape[0] != n_points: + self.status_text.object = "**Error:** NN features size mismatch with UMAP points" + print("NN features shape:", self.nn_features.shape) + print("UMAP embedding shape:", self.umap_embedding.shape) + return + + pc1_var = np.var(self.nn_features[:, 0]) + + # Check for extra features + # curation_features.npy is a dict of additional features (no need to include the PCA in this as this is loaded by default anyway) + # each key should be the name of the feature, and the value should be a numpy array of shape (n_ROIs, feature_dim) + if os.path.exists(os.path.join(curation_dir, 'curation_features.npy')): + self.feature_dict = np.load(os.path.join(curation_dir, 'curation_features.npy'), allow_pickle=True).item() + for key, value in self.feature_dict.items(): + if value.shape[0] == n_points: + if len(np.unique(value)) <= 1: + print(f"Ignoring {key} as it has 0 variance") + continue + print(f"Found property: {key}") + value_scaled = value - np.mean(value, axis=0, keepdims=True) + value_scaled = value_scaled / (np.std(value_scaled, axis=0, keepdims=True) + 1e-6) * np.sqrt(pc1_var) + self.available_features[key] = np.expand_dims(value_scaled, axis=-1) if value_scaled.ndim == 1 else value_scaled + + # Load the curation features into self.properties for use in UMAP visualiser and hist viewer + self.properties[key] = value + else: + print(f"Ignoring {key} due to size mismatch: {value.shape} vs {n_points} ROIs") + if 'PCA' not in self.available_features: + self.available_features['PCA'] = self.nn_features + self.available_features.move_to_end('PCA', last=False) + self.features_toggle = pn.widgets.MultiChoice( + options=list(self.available_features.keys()), + value=['PCA'], + solid=False, + width=200, + ) + self.features_toggle.param.watch(self.update_curation_features, 'value') + + # Set up classifications file + self.classifications_file = os.path.join(curation_dir, f"roi_classifications.json") + + # Load existing classifications + self.load_existing_classifications(n_points) + + # Determine sampling strategy + if n_points > self.sample_size: + self.use_sampling = True + self.sample_indices = np.random.choice(n_points, self.sample_size, replace=False) + self.sample_indices.sort() + self.status_text.object = f"**Status:** Large dataset ({n_points:,} points) - sampling enabled" + else: + self.use_sampling = False + self.sample_indices = np.arange(n_points) + + # Initial clustering + kmeans = KMeans(n_clusters=20, random_state=42, n_init='auto') + initial_clusters = kmeans.fit_predict(self.nn_features) + + info_file = os.path.join(curation_dir, 'all_sessions_info.npy') + if os.path.exists(info_file): + all_sessions_info = np.load(info_file, allow_pickle=True).item() + self.UMAPtime = all_sessions_info.get('UMAPtime', None) + + # Create full dataset + self.full_data = pd.DataFrame({ + 'umap_x': self.umap_embedding[:, 0], + 'umap_y': self.umap_embedding[:, 1], + 'cluster': initial_clusters, + 'classification': self.classifications, + 'original_index': np.arange(n_points) + }) + + # Prepare display data + self.prepare_display_data() + + display_points = len(self.display_data) + self.status_text.object = f"**Status:** Loaded {n_points:,} points, displaying {display_points:,} with 20 clusters" + + except Exception as e: + self.status_text.object = f"**Error:** Could not load data: {str(e)}" + raise ValueError(f"Could not load data: {str(e)}") + + def prepare_display_data(self): + """Prepare sampled data for display""" + if self.full_data is None: + return + + if self.use_sampling: + self.display_data = self.full_data.iloc[self.sample_indices].copy().reset_index(drop=True) + else: + self.display_data = self.full_data.copy() + + def load_existing_classifications(self, n_points): + """Load existing classifications if available""" + self.classifications = ['unclassified'] * n_points + + if os.path.exists(self.classifications_file): + try: + with open(self.classifications_file, 'r') as f: + saved_data = json.load(f) + if len(saved_data['classifications']) == n_points: + self.classifications = saved_data['classifications'] + self.status_text.object = f"**Status:** Loaded existing classifications" + else: + self.status_text.object = f"**Warning:** Classification file size mismatch, starting fresh" + except Exception as e: + self.status_text.object = f"**Warning:** Could not load classifications: {str(e)}" + + def create_components(self): + """Create and initialize the visualization components""" + self._open_hdf5() + + if self.display_data is not None: + + self.umap_visualiser = UMAPVisualiser(self.display_data, properties=self.properties, sample_indices=self.sample_indices, use_sampling=self.use_sampling) + # Subscribe to selection events + self.umap_visualiser.on_cluster_selected = self.on_cluster_selected + + # For box and hist viewers, we need the full dataset (self.dataset), so we pass this in along with the sampling info + self.box_viewer = BoxViewer(self.dataset, self.sample_indices, self.use_sampling) + self.hist_viewer = HistViewer(self.dataset, self.properties, self.sample_indices, self.use_sampling) + + if self.box_viewer and self.hist_viewer: + self.box_viewer.on_sample_changed = self.hist_viewer.update_individual_sample + + # Connect threshold classification callback + if self.hist_viewer: + self.hist_viewer.on_threshold_classify = self._classify_by_threshold + + # Sync hist viewer property selection to UMAP view mode + if self.hist_viewer and self.umap_visualiser: + self.hist_viewer.property_selector.param.watch(self._on_hist_property_change, 'value') + + def _open_hdf5(self): + """Open HDF5 file for reading (same pattern as BoxViewer)""" + try: + self.hdf5_file = h5py.File(self.hdf5_path, 'r') + + # Find the main dataset (same logic as BoxViewer) + self.dataset = self.hdf5_file.get('data', None) + if self.dataset is None: + self.status_text.object = "**Error:** Couldn't load the dataset from HDF5 file." + else: + self.status_text.object = f"**Status:** Using HDF5 dataset" + + # Verify expected shape + if self.dataset.ndim != 5 or self.dataset.shape[1:] != (3, 5, 20, 20): + self.status_text.object = f"**Error:** Unexpected data shape: {self.dataset.shape}" + self.dataset = None + + # Load shot noise if available + if 'shot_noise' in self.hdf5_file: + shot_noise = self.hdf5_file['shot_noise'][:] + if shot_noise.shape[0] == self.dataset.shape[0]: + self.properties['ROI Shot Noise'] = np.clip(shot_noise, 0, 1) + else: + print(f"Found shot noise but shape mismatch: {shot_noise.shape} vs {self.dataset.shape} (shot noise vs dataset length)") + print("Therefore ignoring shot noise") + else: + print("No shot noise in hdf5") + + # Load edge cells if available + if 'edge_cells' in self.hdf5_file: + edge_cells = self.hdf5_file['edge_cells'][:] + if edge_cells.shape[0] == self.dataset.shape[0]: + self.properties['Edge Cells'] = edge_cells + else: + print(f"Found edge cells but shape mismatch: {edge_cells.shape} vs {self.dataset.shape} (edge cells vs dataset length)") + print("Therefore ignoring edge cells") + else: + print("No edge cells in hdf5") + + # Load session IDs if available + if 'session_id' in self.hdf5_file: + session_id = self.hdf5_file['session_id'][:] + if session_id.shape[0] == self.dataset.shape[0]: + self.properties['Session'] = session_id + else: + print(f"Found session IDs but shape mismatch: {session_id.shape} vs {self.dataset.shape} (session ID vs dataset length)") + print("Therefore ignoring session IDs") + else: + print("No session IDs in hdf5") + + except Exception as e: + self.status_text.object = f"**Error:** Could not open HDF5 file: {str(e)}" + + def on_cluster_selected(self, cluster_id, tapped_idx): + """Handle cluster selection from UMAP visualiser""" + # This will be called when a cluster is selected in the UMAP + # Update status and prepare for classification + cluster_size = len(self.full_data[self.full_data['cluster'] == cluster_id]) + self.status_text.object = f"**Status:** Selected cluster {cluster_id} ({cluster_size:,} points) - use cluster classification buttons" + self.selected_cluster = cluster_id + + # Load cluster data in BoxViewer + if self.box_viewer: + self.box_viewer.load_cluster_data(cluster_id, self.display_data, tapped_idx) + + # also load in HistViewer + if self.hist_viewer: + self.hist_viewer.load_cluster_data(cluster_id, self.display_data) + + def _on_hist_property_change(self, event): + """Sync UMAP view mode to match hist viewer property selection (only when in property mode)""" + if self.view_toggle.value != 'property': + return + property_name = event.new + view_mode = PROPERTY_TO_VIEW_MODE.get(property_name, 'clus') + if self.umap_visualiser: + self.umap_visualiser.set_view_mode(view_mode) + + def _on_view_toggle_change(self, event): + """Handle switching between cluster and property coloring modes""" + if event.new == 'cluster': + if self.umap_visualiser: + self.umap_visualiser.set_view_mode('clus') + else: # property mode + # Use whatever property is currently selected in hist viewer + if self.hist_viewer and self.umap_visualiser: + property_name = self.hist_viewer.property_selector.value + view_mode = PROPERTY_TO_VIEW_MODE.get(property_name, 'clus') + self.umap_visualiser.set_view_mode(view_mode) + + def update_curation_features(self, event=None): + """Update the features used for curation based on toggle selection""" + selected_features = self.features_toggle.value + if not selected_features: + self.curation_features_to_use = ['PCA'] + self.features_toggle.value = ['PCA'] + else: + self.curation_features_to_use = list(selected_features) + + # Re-run clustering with new features + self.update_clusters() + + def update_clusters(self, event=None): + """Update clustering""" + if self.full_data is None: + return + + n_clusters = self.cluster_slider.value + self.status_text.object = f"**Status:** Computing {n_clusters} clusters, please wait..." + + self.curation_features = np.hstack([self.available_features[feat] for feat in self.curation_features_to_use]) + + try: + kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init='auto') + new_clusters = kmeans.fit_predict(self.curation_features) + + # Update full dataset + self.full_data['cluster'] = new_clusters + + # Update display data + self.prepare_display_data() + + # Update UMAP visualiser + if self.umap_visualiser: + self.umap_visualiser.update_data(self.display_data) + self.view_toggle.value = 'cluster' # Reset to cluster view + + # Update BoxViewer + if self.box_viewer: + self.box_viewer.clear_cache() + + # Update HistViewer + if self.hist_viewer: + self.hist_viewer.clear_cache() + + self.status_text.object = f"**Status:** Updated to {n_clusters} clusters" + + except Exception as e: + self.status_text.object = f"**Error:** Clustering failed: {str(e)}" + + def recompute_umap(self, event=None): + """Recompute UMAP embedding""" + if self.curation_features is None: + self.status_text.object = "**Error:** No features selected for UMAP recomputation" + return + + try: + self.status_text.object = "**Status:** Recomputing UMAP, please wait..." + self.UMAPbutton.name = "Running UMAP..." + self.UMAPbutton.button_style = 'outline' + start = time.time() + reducer = UMAP(n_components=2, random_state=42, n_neighbors=30, min_dist=0.1, metric='euclidean') + new_embedding = reducer.fit_transform(self.curation_features) + + # Update UMAP embedding + self.umap_embedding = new_embedding + self.full_data['umap_x'] = new_embedding[:, 0] + self.full_data['umap_y'] = new_embedding[:, 1] + + # Update display data + self.prepare_display_data() + + # Update UMAP visualiser + if self.umap_visualiser: + self.umap_visualiser.update_data(self.display_data, False) + self.view_toggle.value = 'cluster' # Reset to cluster view + + end = time.time() + self.UMAPtime = np.round(end - start, 0) + self.UMAPbutton.name = f"Rerun UMAP ({self.UMAPtime}s)" + self.UMAPbutton.button_style = 'solid' + self.status_text.object = "**Status:** UMAP recomputation complete" + + except Exception as e: + self.status_text.object = f"**Error:** UMAP recomputation failed: {str(e)}" + + def classify_cluster_as_cell(self, event): + """Classify entire cluster as cell""" + self._classify_cluster('cell') + + def classify_cluster_as_not_cell(self, event): + """Classify entire cluster as not cell""" + self._classify_cluster('not_cell') + + def _classify_cluster(self, classification_type): + """Classify entire cluster""" + if not hasattr(self, 'selected_cluster') or self.selected_cluster is None: + self.status_text.object = "**Status:** No cluster selected. Click on a point first." + return + + # Classify all points in the cluster + cluster_mask = self.full_data['cluster'] == self.selected_cluster + clus_idx = np.where(cluster_mask)[0] + cluster_size = cluster_mask.sum() + + self.full_data.loc[cluster_mask, 'classification'] = classification_type + + # Update display data + self.prepare_display_data() + + # Update UMAP visualiser + if self.umap_visualiser: + self.umap_visualiser.update_data(self.display_data, False) + + # Update status + cell_count = sum(1 for c in self.full_data['classification'] if c == 'cell') + not_cell_count = sum(1 for c in self.full_data['classification'] if c == 'not_cell') + + self.status_text.object = "Classifying cluster..." + + self.labelled_features_idx.extend(clus_idx) + self.labels.extend(np.array([1 if classification_type == 'cell' else 0] * clus_idx.shape[0])) + + self.curation_features = np.hstack([self.available_features[feat] for feat in self.curation_features_to_use]) + + if len(np.unique(self.labels)) > 1: + # Fit model if we have positive and negative samples + self.linear_model.fit(self.curation_features[self.labelled_features_idx], self.labels) + class_idx = np.argwhere(self.linear_model.classes_ == 1)[0][0] + self.cell_probs = self.linear_model.predict_proba(self.curation_features[self.sample_indices])[:, class_idx] + self.umap_visualiser.set_probs(self.cell_probs) + + self.hist_viewer.update_property('Probability', self.cell_probs) + + self.status_text.object = f"**Status:** Classified cluster {self.selected_cluster} ({cluster_size:,} points) as '{classification_type}' | Total - Cells: {cell_count:,}, Not cells: {not_cell_count:,}" + + # Reset selection + self.selected_cluster = None + + def _classify_by_threshold(self, classification_type): + """Classify points based on histogram threshold selection (global scope)""" + if self.hist_viewer is None: + return + + # Get the threshold mask from hist_viewer + mask = self.hist_viewer.get_threshold_mask() + if mask is None or not np.any(mask): + self.status_text.object = "**Status:** No points in threshold range" + return + + # Map from population sample indices to original full_data indices + # The mask corresponds to the values in population_cache, which uses sample_indices + if self.use_sampling and self.sample_indices is not None: + threshold_indices = self.sample_indices[mask] + else: + threshold_indices = np.where(mask)[0] + + count = len(threshold_indices) + + # Apply classification to full_data + self.full_data.loc[threshold_indices, 'classification'] = classification_type + + # Update display data + self.prepare_display_data() + + # Update UMAP visualiser + if self.umap_visualiser: + self.umap_visualiser.update_data(self.display_data, False) + + # Update ML model training data + self.labelled_features_idx.extend(threshold_indices) + label_val = 1 if classification_type == 'cell' else 0 + self.labels.extend([label_val] * len(threshold_indices)) + + self.curation_features = np.hstack([self.available_features[feat] for feat in self.curation_features_to_use]) + + if len(np.unique(self.labels)) > 1: + # Fit model if we have positive and negative samples + self.linear_model.fit(self.curation_features[self.labelled_features_idx], self.labels) + class_idx = np.argwhere(self.linear_model.classes_ == 1)[0][0] + self.cell_probs = self.linear_model.predict_proba(self.curation_features[self.sample_indices])[:, class_idx] + self.umap_visualiser.set_probs(self.cell_probs) + self.hist_viewer.update_property('Probability', self.cell_probs) + + # Update status + cell_count = sum(1 for c in self.full_data['classification'] if c == 'cell') + not_cell_count = sum(1 for c in self.full_data['classification'] if c == 'not_cell') + self.status_text.object = f"**Status:** Threshold classified {count:,} points as '{classification_type}' | Total - Cells: {cell_count:,}, Not cells: {not_cell_count:,}" + + def reset_classifications(self, event): + """Reset all classifications""" + if self.full_data is None: + return + + self.full_data['classification'] = 'unclassified' + self.prepare_display_data() + self.linear_model = SGDClassifier(loss='log_loss') + self.labelled_features_idx = [] + self.labels = [] + self.cell_probs = None + + if self.umap_visualiser: + self.umap_visualiser.update_data(self.display_data) + self.view_toggle.value = 'cluster' # Reset to cluster view + + if self.hist_viewer: + self.hist_viewer.remove_property('Probability') + + self.status_text.object = "**Status:** Reset all classifications" + + def save_classifications(self, event): + """Save current classifications to file""" + if self.full_data is None or self.classifications_file is None: + return + + try: + full_classifications = self.full_data['classification'].tolist() + + cell_count = sum(1 for c in full_classifications if c == 'cell') + not_cell_count = sum(1 for c in full_classifications if c == 'not_cell') + unclassified_count = sum(1 for c in full_classifications if c == 'unclassified') + + classifications_data = { + 'filename': self.umap_file_path, + 'n_points': len(self.full_data), + 'classifications': full_classifications, + 'counts': { + 'cell': cell_count, + 'not_cell': not_cell_count, + 'unclassified': unclassified_count + } + } + + with open(self.classifications_file, 'w') as f: + json.dump(classifications_data, f, indent=2) + + self.status_text.object = f"**Status:** Saved - Cells: {cell_count:,}, Not cells: {not_cell_count:,}, Unclassified: {unclassified_count:,}" + + except Exception as e: + self.status_text.object = f"**Error:** Could not save classifications: {str(e)}" + + def get_layout(self): + """Return the complete application layout""" + if self.umap_visualiser is None: + return pn.pane.Markdown("Error loading UMAP visualiser.", width=700) + + # Get UMAP plot + plot_pane = self.umap_visualiser.get_plot_pane() + + # Create stats display + stats_display = pn.pane.Markdown("", width=700, margin=(0, 0, 10, 0)) + + if self.full_data is not None: + total_points = len(self.full_data) + display_points = len(self.display_data) if self.display_data is not None else 0 + + if self.use_sampling: + stats_display.object = f"**Total:** {total_points:,} points | **Displaying:** {display_points:,} (sampled)" + else: + stats_display.object = f"**{total_points:,} points**" + + plot_column = pn.Column( + pn.pane.Markdown("### Suite3D Data Curation", width=700, margin=(0, 0, 10, 0)), + stats_display, + plot_pane, + width=720 + ) + + classification_controls = [ + "### Cluster Classification", + pn.Row(self.classify_cluster_cell_button, self.classify_cluster_not_cell_button, width=200), + pn.Spacer(height=10), + self.reset_button, + pn.Spacer(height=10), + self.save_button, + pn.Spacer(height=10), + self.status_text + ] + + controls = pn.Column( + self.cluster_slider, + pn.Spacer(height=10), + self.cluster_button, + self.features_toggle, + self.UMAPbutton, + self.view_toggle, + *classification_controls, + width=300, + margin=(10, 10) + ) + + top_row_components = [pn.Spacer(width=20), plot_column, pn.Spacer(width=20), controls] + if self.box_viewer: + top_row_components.extend([pn.Spacer(width=20), self.box_viewer.get_layout()]) + + top_row = pn.Row( + *top_row_components, + sizing_mode='stretch_width' + ) + + layout_components = [top_row] + + if self.hist_viewer: + layout_components.extend([ + pn.Spacer(height=20), + self.hist_viewer.get_layout() + ]) + + return pn.Column( + *layout_components, + sizing_mode='stretch_width' + ) + + +def create_app(curation_dir): + """Create the application with orchestrator""" + umap_file = os.path.join(curation_dir, 'umap_2d.npy') + nn_features_path = os.path.join(curation_dir, 'pca_embeddings.npy') + hdf5_path = os.path.join(curation_dir, 'dataset.h5') + orchestrator = AppOrchestrator(umap_file, nn_features_path=nn_features_path, hdf5_path=hdf5_path) + return orchestrator.get_layout() + +if __name__ == "__main__": + + curation_dir = r"C:\Users\suyash\UCL\tinya_data\rois\curation" + app = create_app(curation_dir) + app.servable() + + app.show(port=5007) \ No newline at end of file diff --git a/curation/box_viewer.py b/curation/box_viewer.py new file mode 100644 index 0000000..06596da --- /dev/null +++ b/curation/box_viewer.py @@ -0,0 +1,376 @@ +import numpy as np +import panel as pn +from bokeh.plotting import figure +from bokeh.palettes import Greys256, Blues8 + +class BoxViewer: + """Component for viewing 5D data samples with 3 channels and dual view projections""" + + def __init__(self, dataset, sample_indices=None, use_sampling=False): + self.dataset = dataset + self.sample_indices = sample_indices + self.use_sampling = use_sampling + self.display_dataset = None # Subset of dataset to display (if sample_indices provided) + self.cluster_cache = {} + self.current_cluster_id = None + self.current_sample = None # Index in original (full) dataset of the datapoint to display. Will be a list, with length > 1 only if Mean option is selected + + # Keep track of image renderers to avoid recreation + self.xy_image_renderers = [None, None, None] + self.xz_image_renderers = [None, None, None] + + # UI Controls + self.view_selector = pn.widgets.RadioButtonGroup( + name="View Type", + options=["Selected", "Mean", "Median"], + value="Selected", + width=200 + ) + + self.projection_selector = pn.widgets.RadioButtonGroup( + name="Projection Mode", + options=["Max Projection", "Select Slice"], + value="Max Projection", + width=200 + ) + + self.z_slice_slider = pn.widgets.IntSlider( + name="XY View (Z slice)", + start=0, end=4, value=2, + width=160, + disabled=True # Start disabled + ) + + self.y_slice_slider = pn.widgets.IntSlider( + name="XZ View (Y slice)", + start=0, end=19, value=10, + width=160, + disabled=True # Start disabled + ) + + # Status text + self.status_text = pn.pane.Markdown("**Status:** No cluster selected", width=200) + + # Display plots (3 channels × 2 views each) + self.xy_plots = [self._create_empty_plot(f"XY View (20x20)", small=False) for i in range(3)] + self.xz_plots = [self._create_empty_plot(f"XZ View (5x20)", small=True) for i in range(3)] + + # Set up callbacks + self.view_selector.param.watch(self._on_view_change, 'value') + self.projection_selector.param.watch(self._on_projection_change, 'value') + self.z_slice_slider.param.watch(self._on_z_slice_change, 'value') + self.y_slice_slider.param.watch(self._on_y_slice_change, 'value') + self.on_sample_changed = None + + if self.dataset is not None: + if self.sample_indices is not None and len(self.sample_indices) > 0: + random_sample_idx = np.random.choice(len(self.sample_indices)) + random_orig_idx = self.sample_indices[random_sample_idx] + status_msg = f"Showing random sample {random_orig_idx}" + self.display_dataset = self.dataset[np.sort(self.sample_indices)] + else: + n_samples = self.dataset.shape[0] + random_orig_idx = np.random.choice(n_samples) + status_msg = f"Showing random sample {random_orig_idx}" + self.display_dataset = self.dataset + + self.current_sample = [random_orig_idx] + self._update_plots() + self.status_text.object = status_msg + + def _create_empty_plot(self, title, small=False): + """Create empty bokeh plot for displaying images""" + height = 70 if small else 140 + y_range = (0, 5) if small else (0, 20) + plot = figure( + width=160, height=height, + title=title, + toolbar_location=None, + x_range=(0, 20), y_range=y_range + ) + plot.axis.visible = False + plot.grid.visible = False + plot.title.text_font_size = "10pt" + return plot + + def load_cluster_data(self, cluster_id, display_data, tapped_idx): + """Load and cache cluster statistics from HDF5""" + if self.dataset is None: + self.status_text.object = "**Error:** No HDF5 dataset available" + return + + # Check if already cached and clustering hasn't changed + if cluster_id in self.cluster_cache: + self.current_cluster_id = cluster_id + self.cluster_cache[cluster_id]["selected"] = self.display_dataset[tapped_idx] + if self.sample_indices is not None: + original_idx = self.sample_indices[tapped_idx] + self.cluster_cache[cluster_id]["selected"] = [original_idx] + else: + self.cluster_cache[cluster_id]["selected"] = [tapped_idx] + self._update_current_sample() + self.status_text.object = f"**Status:** Loaded cached cluster {cluster_id}" + return + + try: + # Get indices of points in this cluster + cluster_mask = display_data['cluster'] == cluster_id + cluster_data = display_data[cluster_mask] + cluster_original_indices = cluster_data['original_index'].values + + if len(cluster_original_indices) == 0: + self.status_text.object = f"**Error:** No points in cluster {cluster_id}" + return + + sampling_msg = " (sampled)" if self.use_sampling else "" + self.status_text.object = f"**Status:** Loading {len(cluster_original_indices)} samples{sampling_msg}..." + + sort_order = np.argsort(cluster_original_indices) + sorted_indices = cluster_original_indices[sort_order] + cluster_samples_sorted = self.dataset[sorted_indices] + + # Load cluster data from HDF5 + unsort_order = np.argsort(sort_order) + cluster_samples = cluster_samples_sorted[unsort_order] + + # Verify expected shape + if cluster_samples.ndim != 5 or cluster_samples.shape[1:] != (3, 5, 20, 20): + self.status_text.object = f"**Error:** Unexpected data shape: {cluster_samples.shape}" + return + + # Compute statistics + cluster_umap = cluster_data[['umap_x', 'umap_y']].values + com = np.mean(cluster_umap, axis=0) # center of mass of UMAP cluster + distances = np.linalg.norm(cluster_umap - com, axis=1) + closest_idx = np.argmin(distances) + median_idx = cluster_original_indices[closest_idx] + # median_sample = cluster_samples[closest_idx] # (3, 5, 20, 20) + # random_idx = np.random.choice(len(cluster_original_indices), size=min(10, len(cluster_original_indices)), replace=False) + # random_sample = cluster_samples[random_idx] # (3, 5, 20, 20) + if self.sample_indices is not None: + selected_idx = self.sample_indices[tapped_idx] + else: + selected_idx = tapped_idx + + # Cache results + self.cluster_cache[cluster_id] = { + "selected": [selected_idx], + "mean": sorted_indices, + "median": [median_idx], + } + + self.current_cluster_id = cluster_id + self._update_current_sample() + + self.status_text.object = f"**Status:** Loaded cluster {cluster_id} ({len(cluster_original_indices)} samples)" + + if self.on_sample_changed and self.current_sample is not None: + self.on_sample_changed(self.current_sample) + + except Exception as e: + self.status_text.object = f"**Error:** Failed to load cluster data: {str(e)}" + + def _update_current_sample(self): + """Update current sample based on view selector""" + if self.current_cluster_id is None or self.current_cluster_id not in self.cluster_cache: + return + + view_type = self.view_selector.value.lower() + self.current_sample = self.cluster_cache[self.current_cluster_id][view_type] + self._update_plots() + + if self.on_sample_changed: + self.on_sample_changed(self.current_sample) + + def _on_view_change(self, event): + """Handle view selector changes""" + self._update_current_sample() + + def _on_projection_change(self, event): + """Handle projection mode changes""" + # Update slider enabled/disabled state + use_slices = (event.new == "Select Slice") + self.z_slice_slider.disabled = not use_slices + self.y_slice_slider.disabled = not use_slices + self._update_plots() + + def _on_z_slice_change(self, event): + """Handle Z slice slider changes""" + if self.projection_selector.value == "Select Slice": + self._update_plots(which="top") + + def _on_y_slice_change(self, event): + """Handle Y slice slider changes""" + if self.projection_selector.value == "Select Slice": + self._update_plots(which="bottom") + + def _update_plots(self, which="all"): + """Update all 6 plots based on current sample and settings""" + if self.current_sample is None: + return + if len(self.current_sample) > 1: + sample_volume = np.mean(self.dataset[self.current_sample], axis=0) + else: + sample_volume = self.dataset[self.current_sample[0]] + + projection_mode = self.projection_selector.value + + for channel in range(3): + channel_volume = sample_volume[channel] # (5, 20, 20) - Z, Y, X + + # Generate XY and XZ views + if projection_mode == "Max Projection": + xy_data = np.max(channel_volume, axis=0) # Max over Z → (20, 20) - Y, X + xz_data = np.max(channel_volume, axis=1) # Max over Y → (5, 20) - Z, X + else: # Select Slice + z_slice = self.z_slice_slider.value + y_slice = self.y_slice_slider.value + xy_data = channel_volume[z_slice, :, :] # (20, 20) - Y, X + xz_data = channel_volume[:, y_slice, :] # (5, 20) - Z, X + + # Update plots with proper aspect ratios + if which == "all": + self._update_single_plot(self.xy_plots[channel], xy_data, (20, 20), channel=channel) + self._update_single_plot(self.xz_plots[channel], xz_data, (20, 5), channel=channel) + elif which == "top": + self._update_single_plot(self.xy_plots[channel], xy_data, (20, 20), channel=channel) + elif which == "bottom": + self._update_single_plot(self.xz_plots[channel], xz_data, (20, 5), channel=channel) + + def _update_single_plot(self, plot, data, expected_shape, channel): + """Update a single bokeh plot with 2D data""" + palettes = [Greys256, Greys256, Blues8[::-1]] + + # Update plot ranges to match data + plot.x_range.end = expected_shape[0] + plot.y_range.end = expected_shape[1] + + # Normalize data for display (0-1 range) + if data.max() > data.min(): + data_norm = (data - data.min()) / (data.max() - data.min()) + else: + data_norm = data + + # Flip data vertically for proper image orientation (bokeh displays images upside down) + data_flipped = np.flipud(data_norm) + + # Determine which renderer list to use + if plot in self.xy_plots: + renderer_list = self.xy_image_renderers + plot_index = self.xy_plots.index(plot) + else: + renderer_list = self.xz_image_renderers + plot_index = self.xz_plots.index(plot) + + # If renderer doesn't exist, create it + if renderer_list[plot_index] is None: + renderer_list[plot_index] = plot.image( + image=[data_flipped], + x=0, y=0, + dw=expected_shape[0], dh=expected_shape[1], + palette=palettes[channel] + ) + else: + # Update existing renderer's data source + renderer_list[plot_index].data_source.data = { + 'image': [data_flipped], + 'x': [0], + 'y': [0], + 'dw': [expected_shape[0]], + 'dh': [expected_shape[1]] + } + + def clear_cache(self): + """Clear cluster cache (call when clustering changes)""" + self.cluster_cache = {} + self.current_cluster_id = None + self.current_sample = None + self.status_text.object = "**Status:** Cache cleared - select a cluster" + + # Clear all plots and reset renderer tracking + for plot in self.xy_plots + self.xz_plots: + plot.renderers = [] + + # Reset renderer tracking + self.xy_image_renderers = [None, None, None] + self.xz_image_renderers = [None, None, None] + + def get_layout(self): + """Return the complete BoxViewer layout""" + # Main controls (without sliders) + main_controls = pn.Column( + "### Sample Viewer", + self.view_selector, + pn.Spacer(height=10), + self.projection_selector, + pn.Spacer(height=20), + self.status_text, + width=240, + margin=(10, 10) + ) + + # Create channel columns for plots + channel_columns = [] + for channel in range(3): + if channel == 0: + name = "Image" + elif channel == 1: + name = "Correlation Map" + else: + name = "Footprint" + + channel_header = pn.pane.Markdown(f"**{name}**", + margin=(5, 0, 0, 0), + align='center') + channel_col = pn.Column( + channel_header, + pn.pane.Bokeh(self.xy_plots[channel], sizing_mode='fixed'), + pn.pane.Bokeh(self.xz_plots[channel], sizing_mode='fixed'), + margin=(5, 5) + ) + channel_columns.append(channel_col) + + # Create slider controls positioned next to the rows they control + z_slider_control = pn.Column( + pn.Spacer(height=75), # Align with XY plots (accounting for header) + self.z_slice_slider, + pn.Spacer(height=70), # Space to align with XZ row + margin=(5, 10), + width=180 + ) + + y_slider_control = pn.Column( + self.y_slice_slider, + margin=(5, 10), + width=180 + ) + + # Combine sliders into one column + slider_controls = pn.Column( + z_slider_control, + y_slider_control, + width=180 + ) + + # Create the plots section: 3 channel columns + slider controls + plots_section = pn.Row( + *channel_columns, + slider_controls, + margin=(10, 0) + ) + + return pn.Column( + main_controls, + plots_section, + width=800 + ) + + def __del__(self): + """Clean up HDF5 file handle""" + if self.hdf5_file is not None: + try: + self.hdf5_file.close() + except: + pass + + diff --git a/curation/dataloader.py b/curation/dataloader.py new file mode 100644 index 0000000..b945b20 --- /dev/null +++ b/curation/dataloader.py @@ -0,0 +1,613 @@ +import h5py +import numpy as np +import gc +from pathlib import Path +from typing import Dict, List, Tuple, Optional +import warnings +import os, sys +sys.path.insert(0, os.getcwd()) +from suite3d import quality_metrics + + +class Suite3DProcessor: + """ + Process suite3d neural recording session data and extract cell patches. + """ + + def __init__(self, data_dir: str, output_dir: str, box_size: Tuple[int, int, int] = (5, 20, 20)): + """ + Initialize the processor. + + Args: + data_dir: Directory containing session folders + output_dir: Directory to save processed data + box_size: (nbz, nby, nbx) - size of patches to extract around each cell + """ + self.data_dir = Path(data_dir) + self.output_dir = Path(output_dir) + self.output_dir.mkdir(parents=True, exist_ok=True) + + self.nbz, self.nby, self.nbx = box_size + self.nchannel = 3 # mean_img, correlation map and footprint + + # Track processed sessions + self.processed_sessions = [] + + def load_session_data(self, session_path: Path) -> Dict: + """ + Load all relevant data files from a session folder. + + Args: + session_path: Path to session folder + + Returns: + Dictionary containing loaded data + """ + session_data = {} + + # Load fluorescence data + f_path = session_path / "F.npy" + if f_path.exists(): + session_data['F'] = np.load(f_path) + + # Load info dictionary + info_path = session_path / "info.npy" + if info_path.exists(): + session_data['info'] = np.load(info_path, allow_pickle=True).item() + + # Load cell statistics + stats_path = session_path / "stats.npy" + if stats_path.exists(): + session_data['stats'] = np.load(stats_path, allow_pickle=True) + + return session_data + + def get_correlation_map(self, info: Dict) -> Optional[np.ndarray]: + """ + Get correlation map from info dictionary (handle different naming). + + Args: + info: Info dictionary from info.npy + + Returns: + Correlation map array or None if not found + """ + if 'corrmap' in info: + return info['corrmap'] + elif 'vmap_raw' in info: + return info['vmap_raw'] + elif 'vmap' in info: + return info['vmap'] + else: + warnings.warn("None of 'corrmap', 'vmap_raw' or 'vmap' found in info dictionary") + return None + + def extract_cell_patch(self, volume: np.ndarray, center_coord: Tuple[int, int, int]) -> Tuple[np.ndarray, bool]: + """ + Extract a patch around a cell center from a 3D volume. + + Args: + volume: 3D volume (nz, ny, nx) + center_coord: (z, y, x) center coordinates + + Returns: + Tuple of (extracted patch of size (nbz, nby, nbx), is_edge_cell flag) + """ + nz, ny, nx = volume.shape + cz, cy, cx = center_coord + + # Calculate desired patch boundaries (centered on the cell) + z_start = cz - self.nbz // 2 + z_end = z_start + self.nbz + + y_start = cy - self.nby // 2 + y_end = y_start + self.nby + + x_start = cx - self.nbx // 2 + x_end = x_start + self.nbx + + # Clip boundaries to volume limits (don't adjust to maintain size) + z_start_clipped = max(0, z_start) + z_end_clipped = min(nz, z_end) + + y_start_clipped = max(0, y_start) + y_end_clipped = min(ny, y_end) + + x_start_clipped = max(0, x_start) + x_end_clipped = min(nx, x_end) + + # Extract the available patch + patch = volume[z_start_clipped:z_end_clipped, + y_start_clipped:y_end_clipped, + x_start_clipped:x_end_clipped] + + # Check if we have the full desired patch size + is_edge_cell = (patch.shape[0] != self.nbz or + patch.shape[1] != self.nby or + patch.shape[2] != self.nbx) + + # If it's an edge case, pad to the desired size + if is_edge_cell: + padded_patch = np.zeros((self.nbz, self.nby, self.nbx), dtype=patch.dtype) + + # Calculate where to place the extracted patch within the padded array + # to maintain the original centering as much as possible + + # For each dimension, calculate the offset from the desired start + z_offset = max(0, -z_start) # How much we're offset from desired start + y_offset = max(0, -y_start) + x_offset = max(0, -x_start) + + # If we hit the end boundary, we need to shift the placement + if z_end > nz: + z_offset = self.nbz - patch.shape[0] + if y_end > ny: + y_offset = self.nby - patch.shape[1] + if x_end > nx: + x_offset = self.nbx - patch.shape[2] + + # Place the patch in the padded array + padded_patch[z_offset:z_offset + patch.shape[0], + y_offset:y_offset + patch.shape[1], + x_offset:x_offset + patch.shape[2]] = patch + + return padded_patch, is_edge_cell + + return patch, is_edge_cell + + def create_cell_footprint(self, coords: List[np.ndarray], lam: np.ndarray, + med_coord: Tuple[int, int, int]) -> np.ndarray: + """ + Create cell footprint matrix from coordinates and weights relative to med. + + Args: + coords: List of 3 arrays [z_coords, y_coords, x_coords] of pixel coordinates + lam: Array of weights for each pixel + med_coord: (z, y, x) center coordinates + + Returns: + Footprint matrix of size (nbz, nby, nbx) with weights at relative positions + """ + # Initialize empty footprint + footprint = np.zeros((self.nbz, self.nby, self.nbx), dtype=np.float32) + + # Check if coords and lam are valid + if len(coords) != 3: + warnings.warn("coords should contain exactly 3 arrays (z, y, x)") + return footprint + + z_coords, y_coords, x_coords = coords + + # Check if all arrays have same length + if not (len(z_coords) == len(y_coords) == len(x_coords) == len(lam)): + warnings.warn("coords arrays and lam must have same length") + return footprint + + if len(lam) == 0: + warnings.warn("No weights found, returning empty footprint") + return footprint + + # Convert to relative coordinates (relative to med) + med_z, med_y, med_x = med_coord + + # Calculate patch center + center_z = self.nbz // 2 + center_y = self.nby // 2 + center_x = self.nbx // 2 + + # Convert absolute coords to relative coords within patch + rel_z = z_coords - med_z + center_z + rel_y = y_coords - med_y + center_y + rel_x = x_coords - med_x + center_x + + # Filter coordinates that fall within patch boundaries + valid_mask = ( + (rel_z >= 0) & (rel_z < self.nbz) & + (rel_y >= 0) & (rel_y < self.nby) & + (rel_x >= 0) & (rel_x < self.nbx) + ) + + if np.sum(valid_mask) == 0: + warnings.warn("No cell pixels fall within the patch boundaries") + return footprint + + # Get valid coordinates and weights + valid_z = rel_z[valid_mask].astype(int) + valid_y = rel_y[valid_mask].astype(int) + valid_x = rel_x[valid_mask].astype(int) + valid_weights = lam[valid_mask] + + # Place weights at corresponding positions + footprint[valid_z, valid_y, valid_x] = valid_weights + + return footprint + + def process_and_save_session(self, session_path: Path) -> Dict: + """ + Process a single session and immediately save to disk. + + Args: + session_path: Path to session folder + + Returns: + Session info dictionary (patches are saved to disk, not returned) + """ + print(f"Processing session: {session_path.name}") + + # Load session data + session_data = self.load_session_data(session_path) + + # Check required data exists + required_keys = ['info', 'stats'] + for key in required_keys: + if key not in session_data: + raise ValueError(f"Required file {key}.npy not found in {session_path}") + + info = session_data['info'] + stats = session_data['stats'] + fnpy = session_data.get('F', None) + if fnpy is not None: + shot = quality_metrics.shot_noise_suyash(fnpy, 4) + + # Get mean image and correlation map + if 'mean_img' not in info: + raise ValueError("mean_img not found in info dictionary") + + mean_img = info['mean_img'] + corrmap = self.get_correlation_map(info) + + if corrmap is None: + # If no correlation map, use max_img as second channel or duplicate mean_img + if 'max_img' in info: + corrmap = info['max_img'] + print("Using max_img as second channel (corrmap/vmap not found)") + else: + corrmap = mean_img + print("Using mean_img as second channel (corrmap/vmap not found)") + + n_cells = len(stats) + edge_cells = np.zeros(n_cells, dtype=bool) + footprint_sizes = np.zeros(n_cells, dtype=np.int32) + + # Initialize output array + cell_patches = np.zeros((n_cells, self.nchannel, self.nbz, self.nby, self.nbx), + dtype=np.float32) + + if 'contamination_factor' in stats[0].keys(): + cont_factors = np.zeros(n_cells, dtype=np.float32) + track_contamination = True + else: + track_contamination = False + if 'peak_val' in stats[0].keys(): + peak_vals = np.zeros(n_cells, dtype=np.float32) + track_peak = True + else: + track_peak = False + if 'vox_snrs' in stats[0].keys(): + snrs = np.zeros(n_cells, dtype=np.float32) + track_snr = True + else: + track_snr = False + + # Extract patches for each cell + for i, cell_stat in enumerate(stats): + + # Get cell center coordinates (med field) + if 'med' not in cell_stat: + warnings.warn(f"Cell {i} missing 'med' field, skipping") + continue + + med_coord = cell_stat['med'] # Should be (z, y, x) + + # Extract patches from channels + mean_patch, is_edge_mean = self.extract_cell_patch(mean_img, med_coord) + corr_patch, is_edge_corr = self.extract_cell_patch(corrmap, med_coord) + if 'coords' in cell_stat and 'lam' in cell_stat: + coords = cell_stat['coords'] + lam = cell_stat['lam'] + footprint_patch = self.create_cell_footprint(coords, lam, med_coord) + else: + warnings.warn(f"Cell {i} missing 'coords' or 'lam' field, using zero footprint") + footprint_patch = np.zeros((self.nbz, self.nby, self.nbx), dtype=np.float32) + + cell_patches[i, 0] = mean_patch + cell_patches[i, 1] = corr_patch + cell_patches[i, 2] = footprint_patch + + edge_cells[i] = is_edge_mean or is_edge_corr + footprint_sizes[i] = np.sum(footprint_patch > 0) + + if track_contamination: + cont_factor = cell_stat['contamination_factor'] + cont_factors[i] = cont_factor + if track_peak: + peak_val = cell_stat['peak_val'] + peak_vals[i] = peak_val + if track_snr: + snr = cell_stat['vox_snrs'] + snrs[i] = np.median(snr) + + # Prepare session info + session_info = { + 'session_name': session_path.name, + 'n_cells': n_cells, + 'edge_cells': edge_cells, + 'footprint_size': footprint_sizes + } + + if track_contamination: + session_info['contamination_factors'] = cont_factors + if track_peak: + session_info['peak_vals'] = peak_vals + if fnpy is not None: + session_info['shot_noise'] = shot + if track_snr: + session_info['snrs'] = snrs + + # Save immediately to free memory + session_name = session_path.name + patches_file = self.output_dir / f"{session_name}_patches.npy" + np.save(patches_file, cell_patches) + + # Force garbage collection to free memory + del cell_patches, mean_img, corrmap, session_data + gc.collect() + + return session_info + + def process_all_sessions(self) -> List[Dict]: + """ + Process all sessions in the data directory, saving each one immediately. + + Returns: + List of session_info dictionaries (patches are saved to disk) + """ + if not self.data_dir.exists(): + raise ValueError(f"Data directory {self.data_dir} does not exist") + + # Find all session folders + session_paths = [] + for path in self.data_dir.iterdir(): + if path.is_dir() and (path / "info.npy").exists(): + session_paths.append(path) + + if not session_paths: + raise ValueError(f"No valid session folders found in {self.data_dir}") + + print(f"Found {len(session_paths)} sessions to process") + print(f"Output directory: {self.output_dir}") + + all_session_info = [] + + for i, session_path in enumerate(sorted(session_paths), 1): + try: + print(f"\n[{i}/{len(session_paths)}] ", end="") + session_info = self.process_and_save_session(session_path) + + if session_info is not None: + all_session_info.append(session_info) + self.processed_sessions.append(session_path.name) + print(f"✓ Session {session_path.name} completed") + else: + print(f"⚠ Session {session_path.name} skipped (no cells)") + + except Exception as e: + print(f"✗ Error processing {session_path.name}: {e}") + continue + + return all_session_info + + def create_combined_dataset(self, session_info_list: List[Dict]) -> None: + """ + Create a combined dataset from all processed sessions by loading saved files. + This is memory-efficient as it loads one session at a time. + + Args: + session_info_list: List of session info dictionaries + """ + if not session_info_list: + print("No sessions to combine") + return + + print(f"\nCreating combined dataset from {len(session_info_list)} sessions...") + + # Calculate total number of cells + total_cells = sum(info['n_cells'] for info in session_info_list) + print(f"Total cells across all sessions: {total_cells}") + + if total_cells == 0: + print("No cells to combine") + return + + # Initialize combined array + print("Initializing combined array...") + combined_patches = np.zeros((total_cells, self.nchannel, self.nbz, self.nby, self.nbx), + dtype=np.float32) + + all_shot, all_edge, all_peak, all_contamination, all_snrs, all_footprint_size = [], [], [], [], [], [] + + # Load and concatenate each session + current_idx = 0 + for i, session_info in enumerate(session_info_list): + session_name = session_info['session_name'] + n_cells = session_info['n_cells'] + edge_cells = session_info['edge_cells'] + footprint_size = session_info['footprint_size'] + + all_edge.append(edge_cells) + all_footprint_size.append(footprint_size) + + if n_cells == 0: + continue + + print(f"Loading {session_name}: {n_cells} cells") + + # Load session patches + patches_file = self.output_dir / f"{session_name}_patches.npy" + if patches_file.exists(): + session_patches = np.load(patches_file) + + # Copy to combined array + combined_patches[current_idx:current_idx + n_cells] = session_patches + current_idx += n_cells + + # Free memory immediately + os.remove(patches_file) + del session_patches + gc.collect() + else: + print(f"Warning: {patches_file} not found") + + if 'shot_noise' in session_info: + all_shot.append(session_info['shot_noise']) + else: + all_shot.append(np.zeros(n_cells, dtype=np.float32)) + + if 'peak_vals' in session_info: + all_peak.append(session_info['peak_vals']) + else: + all_peak.append(np.zeros(n_cells, dtype=np.float32)) + + if 'snrs' in session_info: + all_snrs.append(session_info['snrs']) + else: + all_snrs.append(np.zeros(n_cells, dtype=np.float32)) + + if 'contamination_factors' in session_info: + all_contamination.append(session_info['contamination_factors']) + else: + all_contamination.append(np.zeros(n_cells, dtype=np.float32)) + + all_shot = np.concatenate(all_shot) + combined_shot_file = self.output_dir / "all_sessions_shot_noise.npy" + np.save(combined_shot_file, all_shot) + + all_edge = np.concatenate(all_edge) + combined_edge_file = self.output_dir / "all_sessions_edge_cells.npy" + np.save(combined_edge_file, all_edge) + + curation_features_dict = self.output_dir / "curation_features.npy" + curation_features = { + "Footprint Size": np.concatenate(all_footprint_size), + "Mean Intensity": np.mean(combined_patches[:, 0, :, :, :], axis=(1,2,3)), + "Mean Correlation": np.mean(combined_patches[:, 1, :, :, :], axis=(1,2,3)), + "ROI Shot Noise": all_shot + } + if all_peak: + all_peak = np.concatenate(all_peak) + curation_features['Peak Value'] = all_peak + if all_contamination: + all_contamination = np.concatenate(all_contamination) + curation_features['Contamination Factor'] = all_contamination + if all_snrs: + all_snrs = np.concatenate(all_snrs) + curation_features['Voxel SNR'] = all_snrs + np.save(curation_features_dict, curation_features) + + # Save combined dataset + combined_patches_file = self.output_dir / "all_sessions_patches.npy" + print(f"Saving combined patches to {combined_patches_file}") + np.save(combined_patches_file, combined_patches) + + # Create combined info + combined_info = { + 'total_cells': total_cells, + 'n_sessions': len(session_info_list), + 'session_names': [info['session_name'] for info in session_info_list], + 'patch_shape': (self.nbz, self.nby, self.nbx), + 'nchannel': self.nchannel, + 'session_cell_counts': [info['n_cells'] for info in session_info_list] + } + + combined_info_file = self.output_dir / "all_sessions_info.npy" + np.save(combined_info_file, combined_info) + + # Clean up + del combined_patches + gc.collect() + + print(f"✓ Combined dataset saved with {total_cells} cells") + + def run_full_pipeline(self, create_combined: bool = True) -> None: + """ + Run the complete processing pipeline. + + Args: + create_combined: Whether to create combined dataset after individual processing + """ + print("=" * 60) + print("Suite3D Processing Pipeline") + print("=" * 60) + print(f"Data directory: {self.data_dir}") + print(f"Output directory: {self.output_dir}") + print(f"Box size (z,y,x): {(self.nbz, self.nby, self.nbx)}") + + # Process all sessions + session_info_list = self.process_all_sessions() + + # Create combined dataset if requested + if create_combined and session_info_list: + self.create_combined_dataset(session_info_list) + + # Final summary + total_cells = sum(info['n_cells'] for info in session_info_list) + print("\n" + "=" * 60) + print("PROCESSING COMPLETE") + print("=" * 60) + print(f"Sessions processed: {len(session_info_list)}") + print(f"Total cells extracted: {total_cells}") + print(f"Output directory: {self.output_dir}") + if create_combined and total_cells > 0: + print(f" - all_sessions_patches.npy ({total_cells} cells)") + print(f" - all_sessions_info.npy") + print("=" * 60) + + +def main(data_directory: str = None, output_directory: str = None): + """ + Main entry point for processing Suite3D data. + Sessions for processing should each have their own folder in data_directory. + Each folder must contain an info.npy file. + + Args: + data_directory: Path to the directory containing folders for each session to be processed + output_directory: Path to the directory where processed data will be saved + """ + + processor = Suite3DProcessor( + data_dir=data_directory, + output_dir=output_directory, + box_size=(5, 20, 20) # nbz=5, nby=20, nbx=20 + ) + + # Process all sessions + processor.run_full_pipeline(create_combined=True) + + # Now create H5 dataset for curation + data = np.load(os.path.join(output_directory, "all_sessions_patches.npy"), allow_pickle=True) + shot = np.load(os.path.join(output_directory, "all_sessions_shot_noise.npy"), allow_pickle=True) + edge = np.load(os.path.join(output_directory, "all_sessions_edge_cells.npy"), allow_pickle=True) + info = np.load(os.path.join(output_directory, "all_sessions_info.npy"), allow_pickle=True).item() + cell_counts = info['session_cell_counts'] + session_id = [i for i, count in enumerate(cell_counts) for _ in range(count)] + h5_path = os.path.join(output_directory, "dataset.h5") + with h5py.File(h5_path, 'w') as hf: + hf.create_dataset("data", data=data) + hf.create_dataset("shot_noise", data=shot) + hf.create_dataset("edge_cells", data=edge) + hf.create_dataset("session_id", data=np.array(session_id, dtype=np.int32)) + + # Now the npy file can be safely deleted + os.remove(os.path.join(output_directory, "all_sessions_patches.npy")) + os.remove(os.path.join(output_directory, "all_sessions_shot_noise.npy")) + os.remove(os.path.join(output_directory, "all_sessions_edge_cells.npy")) + print(f"H5 dataset created at {h5_path}") + + + +if __name__ == "__main__": + + data_directory = r"\path\to\your\data\directory" + output_directory = r"\path\where\you\want\to\save\output" + + main(data_directory=data_directory, output_directory=output_directory) diff --git a/curation/dataset_visualiser.py b/curation/dataset_visualiser.py new file mode 100644 index 0000000..2f88815 --- /dev/null +++ b/curation/dataset_visualiser.py @@ -0,0 +1,770 @@ +""" +Comprehensive visualisation and sanity checking tool for large suite3d datasets. +Handles datasets with hundreds of thousands of cells efficiently. +""" + +import numpy as np +import matplotlib.pyplot as plt +from pathlib import Path +from typing import Tuple, List, Optional, Dict + + +class Suite3DVisualiser: + """ + Memory-efficient visualisation tool for large suite3d datasets. + """ + + def __init__(self, patches_file: str, info_file: str, use_memmap: bool = True): + """ + Initialize visualiser. + + Args: + patches_file: Path to all_sessions_patches.npy + info_file: Path to all_sessions_info.npy + use_memmap: Use memory mapping to avoid loading full array + """ + self.patches_file = Path(patches_file) + self.info_file = Path(info_file) + self.use_memmap = use_memmap + + if not self.patches_file.exists(): + raise FileNotFoundError(f"Patches file not found: {patches_file}") + if not self.info_file.exists(): + raise FileNotFoundError(f"Info file not found: {info_file}") + + # Load info + print("Loading dataset info...") + self.info = np.load(info_file, allow_pickle=True).item() + + # Load patches (with memory mapping if requested) + print("Loading patches dataset...") + if use_memmap: + print("Using memory mapping (data stays on disk)...") + self.patches = np.load(patches_file, mmap_mode='r') + else: + print("Loading full dataset into memory (may take time)...") + self.patches = np.load(patches_file) + + self.n_cells, self.n_channels, self.nz, self.ny, self.nx = self.patches.shape + + print(f"Dataset loaded!") + print(f" Shape: {self.patches.shape}") + print(f" Memory mapping: {use_memmap}") + print(f" Data type: {self.patches.dtype}") + print(f" File size: {self.patches_file.stat().st_size / 1e9:.2f} GB") + + # Channel names for clarity + self.channel_names = ['Mean Image', 'Correlation Map', 'Cell Footprint'] + + def get_basic_stats(self) -> Dict: + """Get basic dataset statistics (memory efficient).""" + print("\nComputing dataset statistics...") + + # Sample for statistics to avoid memory issues + sample_size = min(10000, self.n_cells) + sample_indices = np.random.choice(self.n_cells, sample_size, replace=False) + sample_data = self.patches[sample_indices] + + stats = { + 'n_cells': self.n_cells, + 'shape': self.patches.shape, + 'dtype': str(self.patches.dtype), + 'sample_size_for_stats': sample_size + } + + for ch in range(self.n_channels): + ch_data = sample_data[:, ch] + stats[f'ch{ch}_mean'] = np.mean(ch_data) + stats[f'ch{ch}_std'] = np.std(ch_data) + stats[f'ch{ch}_min'] = np.min(ch_data) + stats[f'ch{ch}_max'] = np.max(ch_data) + stats[f'ch{ch}_nonzero_frac'] = np.mean(ch_data > 0) + + return stats + + def print_dataset_summary(self): + """Print comprehensive dataset summary.""" + stats = self.get_basic_stats() + + print("\n" + "="*60) + print("DATASET SUMMARY") + print("="*60) + print(f"Total cells: {stats['n_cells']:,}") + print(f"Shape: {stats['shape']}") + print(f"Data type: {stats['dtype']}") + print(f"Channels: {self.n_channels} ({', '.join(self.channel_names)})") + print(f"Patch size: {self.nz} × {self.ny} × {self.nx} (Z × Y × X)") + + if 'session_names' in self.info: + print(f"Sessions: {len(self.info['session_names'])}") + print(f"Session names: {', '.join(self.info['session_names'][:5])}") + if len(self.info['session_names']) > 5: + print(f" ... and {len(self.info['session_names'])-5} more") + + print(f"\nStatistics (based on {stats['sample_size_for_stats']:,} random cells):") + print("-" * 40) + + for ch in range(self.n_channels): + ch_name = self.channel_names[ch] + print(f"{ch_name} (Channel {ch}):") + print(f" Mean: {stats[f'ch{ch}_mean']:.4f}") + print(f" Std: {stats[f'ch{ch}_std']:.4f}") + print(f" Range: [{stats[f'ch{ch}_min']:.4f}, {stats[f'ch{ch}_max']:.4f}]") + print(f" Non-zero: {stats[f'ch{ch}_nonzero_frac']*100:.1f}%") + + print("="*60) + + def visualise_random_samples(self, n_samples: int = 16, figsize: Tuple[int, int] = (20, 12), + save_path: Optional[str] = None): + """ + Visualise random cell samples from all channels. + + Args: + n_samples: Number of cells to visualise + figsize: Figure size + save_path: Optional path to save figure + """ + # Select random cells + cell_indices = np.random.choice(self.n_cells, n_samples, replace=False) + + # Create grid: n_samples columns, 3 channels + 2 more views of 3D image + fig, axes = plt.subplots(5, n_samples, figsize=figsize) + if n_samples == 1: + axes = axes.reshape(-1, 1) + + print(f"\nVisualising {n_samples} random cells (indices: {cell_indices[:10]}{'...' if n_samples > 10 else ''})") + + for i, cell_idx in enumerate(cell_indices): + # Load single cell data + cell_data = self.patches[cell_idx] # Shape: (3, 5, 20, 20) + + # Show each channel (max projection for 3D) + for ch in range(3): + channel_data = cell_data[ch] # Shape: (5, 20, 20) + + # Max projection along Z-axis + max_proj = np.max(channel_data, axis=0) # Shape: (20, 20) + + # Choose colormap based on channel + if ch == 0: # Mean image + cmap = 'gray' + vmin, vmax = None, None + elif ch == 1: # Correlation map + cmap = 'viridis' + vmin, vmax = None, None + else: # Cell footprint + cmap = 'hot' + vmin = 0 + vmax = np.max(max_proj) if np.max(max_proj) > 0 else 1 + + im = axes[ch, i].imshow(max_proj, cmap=cmap, vmin=vmin, vmax=vmax) + + if i == 0: # Add channel labels on first column + axes[ch, i].set_ylabel(f'{self.channel_names[ch]}\n(Max Proj)', fontsize=10) + + axes[ch, i].set_title(f'Cell {cell_idx}' if ch == 0 else '', fontsize=8) + axes[ch, i].axis('off') + + # for composite view + # # Create composite view (all channels overlaid or side by side) + # mean_proj = np.max(cell_data[0], axis=0) + # footprint_proj = np.max(cell_data[2], axis=0) + + # # Normalize for overlay + # if np.max(mean_proj) > 0: + # mean_norm = mean_proj / np.max(mean_proj) + # else: + # mean_norm = mean_proj + + # if np.max(footprint_proj) > 0: + # footprint_norm = footprint_proj / np.max(footprint_proj) + # else: + # footprint_norm = footprint_proj + + # # Create RGB composite: mean in gray, footprint in red + # composite = np.zeros((20, 20, 3)) + # composite[:, :, 0] = footprint_norm # Red channel = cell footprint + # composite[:, :, 1] = mean_norm # Green channel = mean image + # composite[:, :, 2] = mean_norm # Blue channel = mean image + + # axes[3, i].imshow(composite) + # if i == 0: + # axes[3, i].set_ylabel('Composite\n(Gray=Mean, Red=Footprint)', fontsize=10) + # axes[3, i].axis('off') + + channel_data = cell_data[0] # mean image, (5, 20, 20) + middle_z_slice = channel_data[self.nz // 2] # Middle Z slice + middle_y_slice = channel_data[:, self.ny // 2] # Middle Y slice + + axes[3, i].imshow(middle_z_slice, cmap='gray') + if i == 0: + axes[3, i].set_ylabel('Middle Z Slice\n(Gray=Mean, Red=Footprint)', fontsize=10) + axes[3, i].axis('off') + axes[4, i].imshow(middle_y_slice, cmap='gray') + if i == 0: + axes[4, i].set_ylabel('Middle Y Slice\n(Gray=Mean, Red=Footprint)', fontsize=10) + axes[4, i].axis('off') + plt.tight_layout() + + if save_path: + plt.savefig(save_path, dpi=150, bbox_inches='tight') + print(f"Figure saved to {save_path}") + + plt.show() + + def analyze_channel_distributions(self, sample_size: int = 50000, + save_path: Optional[str] = None): + """ + Analyze and plot intensity distributions for all channels. + + Args: + sample_size: Number of cells to sample for statistics + save_path: Optional path to save figure + """ + print(f"\nAnalyzing intensity distributions using {sample_size:,} random cells...") + + # Sample data to avoid memory issues + sample_indices = np.random.choice(self.n_cells, + min(sample_size, self.n_cells), + replace=False) + sample_data = self.patches[sample_indices] + + fig, axes = plt.subplots(2, 3, figsize=(15, 10)) + + # Per-pixel intensity distributions + for ch in range(3): + ch_data = sample_data[:, ch].flatten() + + axes[0, ch].hist(ch_data, bins=100, alpha=0.7, color=f'C{ch}', density=True) + axes[0, ch].set_title(f'{self.channel_names[ch]}\nPer-Pixel Intensities') + axes[0, ch].set_xlabel('Intensity') + axes[0, ch].set_ylabel('Density') + axes[0, ch].set_yscale('log') + + # Add statistics text + stats_text = f'Mean: {np.mean(ch_data):.3f}\n' + stats_text += f'Std: {np.std(ch_data):.3f}\n' + stats_text += f'Non-zero: {np.mean(ch_data > 0)*100:.1f}%' + axes[0, ch].text(0.02, 0.98, stats_text, transform=axes[0, ch].transAxes, + verticalalignment='top', fontsize=9, + bbox=dict(boxstyle='round', facecolor='white', alpha=0.8)) + + # Per-cell mean intensities + cell_means = np.mean(sample_data, axis=(2, 3, 4)) # Average over spatial dimensions + + for ch in range(3): + axes[1, ch].hist(cell_means[:, ch], bins=50, alpha=0.7, color=f'C{ch}', density=True) + axes[1, ch].set_title(f'{self.channel_names[ch]}\nPer-Cell Mean Intensities') + axes[1, ch].set_xlabel('Mean Intensity') + axes[1, ch].set_ylabel('Density') + + plt.tight_layout() + + if save_path: + plt.savefig(save_path, dpi=150, bbox_inches='tight') + print(f"Distribution analysis saved to {save_path}") + + plt.show() + + # Print correlation analysis + print("\nChannel correlation analysis:") + corr_matrix = np.corrcoef(cell_means.T) + for i in range(3): + for j in range(i+1, 3): + corr = corr_matrix[i, j] + print(f" {self.channel_names[i]} vs {self.channel_names[j]}: r = {corr:.3f}") + + def inspect_specific_cells(self, cell_indices: List[int], figsize: Tuple[int, int] = (15, 10), + save_path: Optional[str] = None): + """ + Detailed inspection of specific cells with all z-slices. + + Args: + cell_indices: List of cell indices to inspect + figsize: Figure size + save_path: Optional path to save figure + """ + n_cells = len(cell_indices) + + # Create subplots: n_cells rows, 3 channels * nz slices columns + fig, axes = plt.subplots(n_cells, self.n_channels * self.nz, + figsize=(self.n_channels * self.nz * 2, n_cells * 3)) + + if n_cells == 1: + axes = axes.reshape(1, -1) + + for cell_i, cell_idx in enumerate(cell_indices): + if cell_idx >= self.n_cells: + print(f"Warning: Cell index {cell_idx} out of range (max: {self.n_cells-1})") + continue + + print(f"Loading cell {cell_idx}...") + cell_data = self.patches[cell_idx] + + col = 0 + for ch in range(self.n_channels): + for z in range(self.nz): + slice_data = cell_data[ch, z] + + # Choose colormap + if ch == 0: + cmap = 'gray' + elif ch == 1: + cmap = 'viridis' + else: + cmap = 'hot' + + axes[cell_i, col].imshow(slice_data, cmap=cmap) + + # Add titles + if cell_i == 0: + axes[cell_i, col].set_title(f'{self.channel_names[ch][:4]}\nZ={z}', + fontsize=8) + + axes[cell_i, col].axis('off') + col += 1 + + # Add cell label + axes[cell_i, 0].text(-0.5, 0.5, f'Cell\n{cell_idx}', + transform=axes[cell_i, 0].transAxes, + verticalalignment='center', horizontalalignment='right', + fontsize=10, fontweight='bold') + + plt.tight_layout() + + if save_path: + plt.savefig(save_path, dpi=150, bbox_inches='tight') + print(f"Detailed inspection saved to {save_path}") + + plt.show() + + def validate_footprint_channel(self, n_samples: int = 20): + """ + Validate that the cell footprint channel makes sense. + + Args: + n_samples: Number of cells to validate + """ + print(f"\nValidating cell footprint channel using {n_samples} random cells...") + + sample_indices = np.random.choice(self.n_cells, n_samples, replace=False) + + validation_results = { + 'cells_with_footprint': 0, + 'cells_without_footprint': 0, + 'footprint_sizes': [], + 'footprint_weights_range': [], + 'footprint_center_weights': [] + } + + for cell_idx in sample_indices: + footprint = self.patches[cell_idx, 2] # Channel 2 = footprint + + nonzero_pixels = np.sum(footprint > 0) + + if nonzero_pixels > 0: + validation_results['cells_with_footprint'] += 1 + validation_results['footprint_sizes'].append(nonzero_pixels) + + nonzero_weights = footprint[footprint > 0] + validation_results['footprint_weights_range'].append( + (np.min(nonzero_weights), np.max(nonzero_weights)) + ) + + # Check center pixel weight (should often be high) + center_z, center_y, center_x = self.nz//2, self.ny//2, self.nx//2 + center_weight = footprint[center_z, center_y, center_x] + validation_results['footprint_center_weights'].append(center_weight) + else: + validation_results['cells_without_footprint'] += 1 + + print(f"Footprint validation results:") + print(f" Cells with footprints: {validation_results['cells_with_footprint']}/{n_samples}") + print(f" Cells without footprints: {validation_results['cells_without_footprint']}/{n_samples}") + + if validation_results['footprint_sizes']: + sizes = validation_results['footprint_sizes'] + print(f" Footprint sizes: {np.min(sizes)} - {np.max(sizes)} pixels (mean: {np.mean(sizes):.1f})") + + weights_min = [r[0] for r in validation_results['footprint_weights_range']] + weights_max = [r[1] for r in validation_results['footprint_weights_range']] + print(f" Weight ranges: [{np.min(weights_min):.3f}, {np.max(weights_max):.3f}]") + + center_weights = validation_results['footprint_center_weights'] + center_nonzero = [w for w in center_weights if w > 0] + if center_nonzero: + print(f" Center pixel weights: {len(center_nonzero)}/{len(center_weights)} have weight > 0") + print(f" Mean center weight: {np.mean(center_nonzero):.3f}") + + def interactive_cell_browser(self): + """ + Launch interactive cell browser (works best in Jupyter notebooks). + """ + print("\nLaunching interactive cell browser...") + print("Use slider to browse through cells, or enter specific cell index.") + + # Create interactive plot + fig, axes = plt.subplots(1, 4, figsize=(16, 4)) + + def update_cell(cell_idx): + cell_idx = int(cell_idx) + if cell_idx >= self.n_cells: + cell_idx = self.n_cells - 1 + elif cell_idx < 0: + cell_idx = 0 + + cell_data = self.patches[cell_idx] + + # Clear axes + for ax in axes: + ax.clear() + + # Show each channel + for ch in range(3): + max_proj = np.max(cell_data[ch], axis=0) + + if ch == 0: + cmap = 'gray' + elif ch == 1: + cmap = 'viridis' + else: + cmap = 'hot' + + axes[ch].imshow(max_proj, cmap=cmap) + axes[ch].set_title(f'{self.channel_names[ch]}') + axes[ch].axis('off') + + # Composite + mean_proj = np.max(cell_data[0], axis=0) + footprint_proj = np.max(cell_data[2], axis=0) + + if np.max(mean_proj) > 0: + mean_norm = mean_proj / np.max(mean_proj) + else: + mean_norm = mean_proj + + composite = np.zeros((20, 20, 3)) + composite[:, :, 1] = mean_norm + composite[:, :, 2] = mean_norm + if np.max(footprint_proj) > 0: + composite[:, :, 0] = footprint_proj / np.max(footprint_proj) + + axes[3].imshow(composite) + axes[3].set_title('Composite') + axes[3].axis('off') + + fig.suptitle(f'Cell {cell_idx} / {self.n_cells-1}') + plt.draw() + + # Show first cell + update_cell(0) + + # Add slider (if in interactive environment) + try: + from matplotlib.widgets import Slider + ax_slider = plt.axes([0.1, 0.02, 0.8, 0.03]) + slider = Slider(ax_slider, 'Cell Index', 0, self.n_cells-1, + valinit=0, valfmt='%d') + slider.on_changed(update_cell) + except: + print("Interactive slider not available in this environment") + + plt.tight_layout() + plt.show() + + return update_cell # Return function for manual use + + def find_interesting_cells(self, criteria: str = 'bright', n_cells: int = 10) -> List[int]: + """ + Find cells that meet certain criteria. + + Args: + criteria: 'bright', 'dim', 'large_footprint', 'small_footprint', 'high_corr' + n_cells: Number of cells to return + + Returns: + List of interesting cell indices + """ + print(f"\nFinding {n_cells} cells with criteria: {criteria}") + + # Sample subset for efficiency + sample_size = min(50000, self.n_cells) + sample_indices = np.random.choice(self.n_cells, sample_size, replace=False) + sample_data = self.patches[sample_indices] + + if criteria == 'bright': + # Highest mean intensity in mean image channel + scores = np.mean(sample_data[:, 0], axis=(1, 2, 3)) + selected = np.argsort(scores)[-n_cells:][::-1] + + elif criteria == 'dim': + # Lowest mean intensity in mean image channel + scores = np.mean(sample_data[:, 0], axis=(1, 2, 3)) + selected = np.argsort(scores)[:n_cells] + + elif criteria == 'large_footprint': + # Most pixels in footprint + scores = np.sum(sample_data[:, 2] > 0, axis=(1, 2, 3)) + selected = np.argsort(scores)[-n_cells:][::-1] + + elif criteria == 'small_footprint': + # Fewest pixels in footprint (but > 0) + footprint_sizes = np.sum(sample_data[:, 2] > 0, axis=(1, 2, 3)) + valid_footprints = footprint_sizes > 0 + valid_indices = np.where(valid_footprints)[0] + + if len(valid_indices) >= n_cells: + scores = footprint_sizes[valid_indices] + selected_valid = np.argsort(scores)[:n_cells] + selected = valid_indices[selected_valid] + else: + selected = valid_indices + + elif criteria == 'high_corr': + # Highest max intensity in correlation channel + scores = np.max(sample_data[:, 1], axis=(1, 2, 3)) + selected = np.argsort(scores)[-n_cells:][::-1] + + else: + raise ValueError(f"Unknown criteria: {criteria}") + + # Convert back to original indices + interesting_indices = sample_indices[selected].tolist() + + print(f"Found {len(interesting_indices)} interesting cells:") + for i, idx in enumerate(interesting_indices[:5]): # Show first 5 + print(f" {i+1}. Cell {idx}") + if len(interesting_indices) > 5: + print(f" ... and {len(interesting_indices)-5} more") + + return interesting_indices + + def visualise_cluster_examples(self, cluster_labels_file: str, n_examples: int = 5, + figsize: Tuple[int, int] = (15, 12), save_path: Optional[str] = None): + """ + Visualize typical examples from each cluster. + + Args: + cluster_labels_file: Path to cluster labels numpy file + n_examples: Number of examples to show per cluster (default: 5) + figsize: Figure size (default: (15, 12)) + save_path: Optional path to save figure + """ + print(f"\nLoading cluster labels from {cluster_labels_file}") + + # Load cluster labels + cluster_labels_path = Path(cluster_labels_file) + if not cluster_labels_path.exists(): + raise FileNotFoundError(f"Cluster labels file not found: {cluster_labels_file}") + + cluster_labels = np.load(cluster_labels_file) + + if len(cluster_labels) != self.n_cells: + raise ValueError(f"Cluster labels length ({len(cluster_labels)}) doesn't match dataset size ({self.n_cells})") + + # Get unique clusters, ignore -1 (noise) + unique_clusters = np.unique(cluster_labels) + unique_clusters = [c for c in unique_clusters if c != -1] + n_clusters = len(unique_clusters) + print(f"Found {n_clusters} unique clusters (excluding noise): {unique_clusters}") + + # For each cluster, find typical examples (closest to cluster centroid) + print("Computing cluster centroids and finding representative examples...") + + cluster_examples = {} + for cluster_id in unique_clusters: + # Get all cells in this cluster + cluster_mask = cluster_labels == cluster_id + cluster_indices = np.where(cluster_mask)[0] + n_cells_in_cluster = len(cluster_indices) + + print(f" Cluster {cluster_id}: {n_cells_in_cluster} cells") + + if n_cells_in_cluster == 0: + continue + + # Sample a subset if cluster is too large (for efficiency) + if n_cells_in_cluster > 1000: + sample_indices = np.random.choice(cluster_indices, 1000, replace=False) + else: + sample_indices = cluster_indices + + # Load data for sampled cells + cluster_data = self.patches[sample_indices] + + # Compute centroid (mean across all cells in cluster) + # Shape: (n_channels, nz, ny, nx) + centroid = np.mean(cluster_data, axis=0) + + # Find cells closest to centroid + distances = [] + for i, cell_idx in enumerate(sample_indices): + cell_data = cluster_data[i] + # Compute Euclidean distance to centroid + distance = np.linalg.norm(cell_data - centroid) + distances.append((distance, cell_idx)) + + # Sort by distance and take closest examples + distances.sort(key=lambda x: x[0]) + n_examples_to_show = min(n_examples, len(distances)) + representative_indices = [distances[i][1] for i in range(n_examples_to_show)] + cluster_examples[cluster_id] = representative_indices + + + # Create visualization + # Layout: n_clusters rows, n_examples*3 columns (3 images per sample) + n_cols = n_examples * 3 + fig, axes = plt.subplots(n_clusters, n_cols, figsize=(n_cols*1.2, n_clusters*1.2), squeeze=False) + + print("\nGenerating visualization...") + + for cluster_row, cluster_id in enumerate(unique_clusters): + if cluster_id not in cluster_examples: + continue + examples = cluster_examples[cluster_id] + for example_idx, cell_idx in enumerate(examples): + cell_data = self.patches[cell_idx] + for ch in range(3): + col = example_idx * 3 + ch + if ch == 0: + img = np.max(cell_data[0], axis=0) + cmap = 'gray' + title = 'Mean' + elif ch == 1: + img = np.max(cell_data[1], axis=0) + cmap = 'viridis' + title = 'Corr' + else: + img = np.max(cell_data[2], axis=0) + cmap = 'hot' + title = 'Foot' + axes[cluster_row, col].imshow(img, cmap=cmap) + axes[cluster_row, col].axis('off') + if cluster_row == 0: + axes[cluster_row, col].set_title(f'Ex{example_idx+1}\n{title}', fontsize=7) + # Add cluster label on the left + axes[cluster_row, 0].text(-0.5, 0.5, f'Cluster {cluster_id}\n({len(examples)} shown)', + transform=axes[cluster_row, 0].transAxes, + verticalalignment='center', horizontalalignment='right', + fontsize=8, fontweight='bold', + bbox=dict(boxstyle='round', facecolor='lightblue', alpha=0.7)) + # Hide empty subplots if fewer examples than requested + for example_idx in range(len(examples), n_examples): + for ch in range(3): + col = example_idx * 3 + ch + axes[cluster_row, col].axis('off') + + plt.subplots_adjust(wspace=0.02, hspace=0.08, left=0.03, right=0.99, top=0.95, bottom=0.03) + # plt.suptitle(f'Cluster Examples: {n_examples} representative cells per cluster\n' + # f'Each sample: Mean | Corr | Foot', + # fontsize=12, y=0.99) + + if save_path: + plt.savefig(save_path, dpi=150, bbox_inches='tight') + print(f"Cluster visualization saved to {save_path}") + + plt.show() + + # Print cluster statistics + print("\nCluster Statistics:") + print("-" * 40) + for cluster_id in unique_clusters: + n_cells_in_cluster = np.sum(cluster_labels == cluster_id) + percentage = (n_cells_in_cluster / self.n_cells) * 100 + print(f"Cluster {cluster_id:2d}: {n_cells_in_cluster:6,} cells ({percentage:5.1f}%)") + + def export_summary_report(self, output_file: str): + """ + Export comprehensive summary report. + + Args: + output_file: Path to output text file + """ + print(f"\nExporting summary report to {output_file}") + + stats = self.get_basic_stats() + + with open(output_file, 'w') as f: + f.write("Suite3D Dataset Visualisation Report\n") + f.write("=" * 50 + "\n\n") + + f.write(f"Dataset: {self.patches_file}\n") + f.write(f"Info: {self.info_file}\n\n") + + f.write("Dataset Overview:\n") + f.write(f" Shape: {stats['shape']}\n") + f.write(f" Total cells: {stats['n_cells']:,}\n") + f.write(f" Data type: {stats['dtype']}\n") + f.write(f" File size: {self.patches_file.stat().st_size / 1e9:.2f} GB\n") + f.write(f" Memory mapping used: {self.use_memmap}\n\n") + + f.write("Channels:\n") + for i, name in enumerate(self.channel_names): + f.write(f" {i}: {name}\n") + f.write("\n") + + f.write(f"Statistics (sample size: {stats['sample_size_for_stats']:,}):\n") + for ch in range(self.n_channels): + f.write(f" {self.channel_names[ch]}:\n") + f.write(f" Mean: {stats[f'ch{ch}_mean']:.4f}\n") + f.write(f" Std: {stats[f'ch{ch}_std']:.4f}\n") + f.write(f" Range: [{stats[f'ch{ch}_min']:.4f}, {stats[f'ch{ch}_max']:.4f}]\n") + f.write(f" Non-zero fraction: {stats[f'ch{ch}_nonzero_frac']*100:.1f}%\n\n") + + if hasattr(self, 'info'): + f.write("Session Info:\n") + for key, value in self.info.items(): + f.write(f" {key}: {value}\n") + + print(f"Report saved to {output_file}") + + def visualise_cluster_arrays_side_by_side(self, arr1: np.ndarray, arr2: np.ndarray, figsize: Tuple[int, int] = (10, 2), save_path: Optional[str] = None): + """ + Visualise two arrays of shape (N, 3, 5, 20, 20) side by side. + Each row is a cluster, columns are arr1 and arr2. + Args: + arr1: np.ndarray of shape (N, 3, 5, 20, 20) + arr2: np.ndarray of shape (N, 3, 5, 20, 20) + figsize: Figure size per row (default: (10, 2)) + save_path: Optional path to save figure + """ + assert arr1.shape == arr2.shape, "Arrays must have the same shape." + assert arr1.ndim == 5 and arr1.shape[1] == 3, "Arrays must be of shape (N, 3, 5, 20, 20)." + N = arr1.shape[0] + row_labels = ['Mean', 'Median'] + col_labels = ['Noise'] + [f'Cluster {i}' for i in range(1, N)] + # Make each subplot smaller and tighten layout + fig, axes = plt.subplots(2, N, figsize=(1.8*N, 2.8), squeeze=False) + arrays = [arr1, arr2] + for row in range(2): + for col in range(N): + img = np.max(arrays[row][col, 0], axis=0) + axes[row, col].imshow(img, cmap='gray') + axes[row, col].axis('off') + if row == 0: + axes[row, col].set_title(col_labels[col], fontsize=8) + if col == 0: + axes[row, col].set_ylabel(row_labels[row], fontsize=10, fontweight='bold') + plt.subplots_adjust(wspace=0.01, hspace=0.05, left=0.03, right=0.99, top=0.90, bottom=0.10) + plt.suptitle('Cluster Representatives: Mean (top) and Median (bottom)', fontsize=12, y=0.97) + if save_path: + plt.savefig(save_path, dpi=150, bbox_inches='tight') + print(f"Cluster arrays side-by-side visualization saved to {save_path}") + plt.show() + +def main(): + """Command line interface for dataset visualisation.""" + + # Initialize (uses memory mapping) + viz = Suite3DVisualiser(r"\\path\to\all_sessions_patches.npy", r"\\path\to\all_sessions_info.npy") + + # # Quick overview + # viz.print_dataset_summary() + + # # Visualize random samples + viz.visualise_random_samples(n_samples=20) + + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/curation/dimensionality_reduction.py b/curation/dimensionality_reduction.py new file mode 100644 index 0000000..a81fee4 --- /dev/null +++ b/curation/dimensionality_reduction.py @@ -0,0 +1,466 @@ +from __future__ import annotations +import os +from typing import Iterable, Tuple +import numpy as np +import matplotlib.pyplot as plt +from sklearn.preprocessing import StandardScaler +from sklearn.decomposition import IncrementalPCA, PCA +from sklearn.cluster import HDBSCAN +from umap import UMAP +import json +import h5py +from sklearn.feature_selection import VarianceThreshold +import time + + +def _iter_batches_permod(X: np.ndarray, batch_size: int) -> Iterable[Tuple[np.ndarray, np.ndarray, np.ndarray]]: + """Yield modality-specific batches: returns x0,x1,x2 each shaped (B,2000).""" + N = X.shape[0] + for start in range(0, N, batch_size): + stop = min(N, start + batch_size) + xb = X[start:stop].astype(np.float32) # (B,3,5,20,20) + + # Split modalities and flatten last 3 dims + x0 = xb[:, 0].reshape(xb.shape[0], -1) # (B,2000) + x1 = xb[:, 1].reshape(xb.shape[0], -1) + x2 = xb[:, 2].reshape(xb.shape[0], -1) + + # Clean non-finite values + for arr in (x0, x1, x2): + arr[~np.isfinite(arr)] = 0.0 + + yield x0, x1, x2 + + +def run_permod_pca_umap( + X: np.ndarray, + ncomp_per_mod: int = 32, + batch_size: int = 4096, + out_dir: str = "pca_permod_umap", + whiten: bool = False, + seed: int = 0, + umap_neighbors: int = 30, + umap_min_dist: float = 0.1, + savename: str = "umap_2d", + variance_threshold = 1e-8 +) -> str: + """ + Per-modality IncrementalPCA, concatenate embeddings, then UMAP to 2D. + """ + assert X.ndim == 5 and X.shape[1:] == (3, 5, 20, 20), "X must have shape (N,3,5,20,20)" + os.makedirs(out_dir, exist_ok=True) + + N = X.shape[0] + + # Initialize scalers and PCA per modality + scalers = [StandardScaler() for _ in range(3)] + variance_selectors = [None for _ in range(3)] + ipcas = [IncrementalPCA(n_components=ncomp_per_mod, batch_size=batch_size) for _ in range(3)] + + # Pass 1: fit scalers + print("Fitting scalers...") + for x0, x1, x2 in _iter_batches_permod(X, batch_size): + scalers[0].partial_fit(x0) + scalers[1].partial_fit(x1) + scalers[2].partial_fit(x2) + # Create variance selectors to remove low-variance features + print("Setting up variance filtering...") + + # Use a small sample to determine feature variance + sample_size = min(1000, N) + sample_idx = np.random.choice(N, sample_size, replace=False) + X_sample = X[sample_idx] + + for mod_idx in range(3): + # Get sample data for this modality + x_sample = X_sample[:, mod_idx].reshape(sample_size, -1) + x_sample_scaled = scalers[mod_idx].transform(x_sample) + + # Create variance threshold selector + variance_selectors[mod_idx] = VarianceThreshold(threshold=variance_threshold) + variance_selectors[mod_idx].fit(x_sample_scaled) + + # Determine effective number of components + n_features_remaining = variance_selectors[mod_idx].transform(x_sample_scaled).shape[1] + effective_ncomp = min(ncomp_per_mod, n_features_remaining - 1, sample_size - 1) + + print(f"Modality {mod_idx}: {n_features_remaining} features after variance filtering, " + f"using {effective_ncomp} PCA components") + + # Initialize IncrementalPCA with appropriate number of components + ipcas[mod_idx] = IncrementalPCA( + n_components=effective_ncomp, + batch_size=min(batch_size, sample_size) + ) + + # Calculate total embedding dimension + K = sum(ipca.n_components for ipca in ipcas) + + # Pass 2: fit IncrementalPCA per modality + print("Fitting PCA models...") + for x0, x1, x2 in _iter_batches_permod(X, batch_size): + # Scale and filter variance for each modality + xs_filtered = [] + for mod_idx, x in enumerate([x0, x1, x2]): + x_scaled = scalers[mod_idx].transform(x) + x_filtered = variance_selectors[mod_idx].transform(x_scaled) + xs_filtered.append(x_filtered) + + # Partial fit PCA + ipcas[mod_idx].partial_fit(x_filtered) + + # Pass 3: transform and concatenate + print("Transforming data...") + emb_path = os.path.join(out_dir, "embeddings_permod.npy") + Z = np.lib.format.open_memmap(emb_path, mode="w+", dtype=np.float32, shape=(N, K)) + + i = 0 + final_scaler = None + + for x0, x1, x2 in _iter_batches_permod(X, batch_size): + # Process each modality + embeddings = [] + + for mod_idx, x in enumerate([x0, x1, x2]): + # Scale and filter variance + x_scaled = scalers[mod_idx].transform(x) + x_filtered = variance_selectors[mod_idx].transform(x_scaled) + + # PCA transform + z = ipcas[mod_idx].transform(x_filtered) + + # Optional whitening with numerical stability + if whiten: + explained_var = ipcas[mod_idx].explained_variance_ + # Add small epsilon to prevent division by zero + z = z / np.sqrt(np.maximum(explained_var, 1e-12)) + + embeddings.append(z) + + # Concatenate embeddings from all modalities + Zb = np.concatenate(embeddings, axis=1).astype(np.float32) + + # Apply final scaling + if i == 0: # First batch - fit scaler + final_scaler = StandardScaler() + Zb_scaled = final_scaler.fit_transform(Zb) + else: + Zb_scaled = final_scaler.transform(Zb) + + # Store batch + b = Zb_scaled.shape[0] + Z[i:i+b] = Zb_scaled + i += b + + #save PCA embeddings + np.save(os.path.join(out_dir, "pca_embeddings.npy"), Z) + + # UMAP on concatenated embeddings + print("Running UMAP...") + umap_model = UMAP( + n_neighbors=umap_neighbors, + min_dist=umap_min_dist, + metric="euclidean", + random_state=seed, + ) + + # Convert memmap to array for UMAP + Z_array = np.array(Z) + Y = umap_model.fit_transform(Z_array) + + # Save results + np.save(os.path.join(out_dir, f"{savename}.npy"), Y.astype(np.float32)) + save_scatter(os.path.join(out_dir, f"{savename}.png"), Y) + + print(f"Pipeline complete. Results saved to: {out_dir}") + return out_dir + + +def PCAfunction( + X: np.ndarray, + ncomp: int = 32, + out_dir: str = "pca_permod_umap", + whiten: bool = False, + seed: int = 0, + umap_neighbors: int = 30, + umap_min_dist: float = 0.1, + savename: str = "umap_2d", + visualize_pcs: bool = True, + n_pcs_to_visualize: int = 16, + image_shape: tuple = None, + channel_names: list = None, # e.g., ['slice1', 'slice2', ...] or ['R', 'G', 'B', ...] +) -> str: + """ + PCA reduction followed by UMAP to 2D. + """ + os.makedirs(out_dir, exist_ok=True) + + N = X.shape[0] + X = X.reshape(N, -1).astype(np.float32) + + # Fit PCA on all data at once + print(f"Fitting PCA with {ncomp} components...") + pca = PCA(n_components=ncomp, whiten=whiten, random_state=seed) + Z = pca.fit_transform(X).astype(np.float32) + + print(f"PCA reduced data to {pca.n_components_} components") + print(f"Explained variance ratio: {pca.explained_variance_ratio_.sum():.3f}") + + # Save PCA embeddings + emb_path = os.path.join(out_dir, "pca_embeddings.npy") + np.save(emb_path, Z) + + # Visualize principal components + if visualize_pcs: + print("Visualizing principal components...") + n_viz = min(n_pcs_to_visualize, pca.n_components_) + + # Save raw PC components + pc_components = pca.components_[:n_viz] + np.save(os.path.join(out_dir, "pc_components.npy"), pc_components) + + if image_shape is not None: + if len(image_shape) == 3 and image_shape[0] <= 10: + _visualize_pcs_multichannel(pc_components, image_shape, pca.explained_variance_ratio_, + out_dir, n_viz, channel_names) + + # Always save explained variance plot + _plot_explained_variance(pca, out_dir) + + print(f"PC visualizations saved to: {out_dir}") + + # Run UMAP + print("Running UMAP...") + start_time = time.time() + umap_model = UMAP( + n_neighbors=umap_neighbors, + min_dist=umap_min_dist, + metric="euclidean", + random_state=seed, + ) + Y = umap_model.fit_transform(Z) + end_time = time.time() + UMAPtime = np.round(end_time - start_time, 0) + + # Save UMAP results + np.save(os.path.join(out_dir, f"{savename}.npy"), Y.astype(np.float32)) + + print(f"Pipeline complete. Results saved to: {out_dir}") + return UMAPtime + + +def _visualize_pcs_multichannel(pc_components, image_shape, var_ratios, out_dir, n_viz, channel_names=None): + """ + Visualize PCs for multi-channel/multi-slice data (e.g., shape (5, 20, 20)). + Shows each channel/slice as a separate 2D heatmap. + """ + n_channels, height, width = image_shape + + if channel_names is None: + channel_names = [f'Ch{i+1}' for i in range(n_channels)] + + # Create a large figure showing all PCs and channels + fig = plt.figure(figsize=(n_channels * 3, n_viz * 2.5)) + + for pc_idx in range(n_viz): + pc = pc_components[pc_idx].reshape(image_shape) + + # Normalize across all channels for consistent color scale + vmin, vmax = pc.min(), pc.max() + + for ch_idx in range(n_channels): + ax = plt.subplot(n_viz, n_channels, pc_idx * n_channels + ch_idx + 1) + + im = ax.imshow(pc[ch_idx], cmap='RdBu_r', vmin=vmin, vmax=vmax, aspect='auto') + + # Add title only on top row + if pc_idx == 0: + ax.set_title(channel_names[ch_idx], fontsize=10) + + # Add PC label only on left column + if ch_idx == 0: + ax.set_ylabel(f'PC{pc_idx+1}\n({var_ratios[pc_idx]:.1%})', fontsize=9) + + ax.set_xticks([]) + ax.set_yticks([]) + + plt.tight_layout() + plt.savefig(os.path.join(out_dir, "pc_channels_grid.png"), dpi=150, bbox_inches='tight') + plt.close() + + # Also create individual PC visualizations with colorbars + pc_dir = os.path.join(out_dir, "individual_pcs") + os.makedirs(pc_dir, exist_ok=True) + + for pc_idx in range(min(n_viz, 8)): # Save detailed view for first 8 PCs + pc = pc_components[pc_idx].reshape(image_shape) + + fig, axes = plt.subplots(1, n_channels, figsize=(n_channels * 3, 3)) + if n_channels == 1: + axes = [axes] + + vmin, vmax = pc.min(), pc.max() + + for ch_idx in range(n_channels): + ax = axes[ch_idx] + im = ax.imshow(pc[ch_idx], cmap='RdBu_r', vmin=vmin, vmax=vmax) + ax.set_title(f'{channel_names[ch_idx]}', fontsize=12) + ax.axis('off') + plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) + + fig.suptitle(f'PC{pc_idx+1} ({var_ratios[pc_idx]:.2%} variance explained)', + fontsize=14, y=1.02) + plt.tight_layout() + plt.savefig(os.path.join(pc_dir, f"pc{pc_idx+1}_channels.png"), + dpi=150, bbox_inches='tight') + plt.close() + + +def _plot_explained_variance(pca, out_dir): + """Plot explained variance.""" + import matplotlib.pyplot as plt + + plt.figure(figsize=(10, 4)) + + plt.subplot(1, 2, 1) + plt.bar(range(1, len(pca.explained_variance_ratio_) + 1), pca.explained_variance_ratio_) + plt.xlabel('Principal Component') + plt.ylabel('Explained Variance Ratio') + plt.title('Variance Explained by Each PC') + plt.grid(alpha=0.3) + + plt.subplot(1, 2, 2) + plt.plot(range(1, len(pca.explained_variance_ratio_) + 1), + np.cumsum(pca.explained_variance_ratio_), marker='o') + plt.xlabel('Number of Components') + plt.ylabel('Cumulative Explained Variance') + plt.title('Cumulative Variance Explained') + plt.grid(alpha=0.3) + plt.axhline(y=0.95, color='r', linestyle='--', alpha=0.5, label='95%') + plt.legend() + + plt.tight_layout() + plt.savefig(os.path.join(out_dir, "explained_variance.png"), dpi=150, bbox_inches='tight') + plt.close() + + +def save_scatter(p: str, Y: np.ndarray, labels=None): + """Save a simple scatter plot of UMAP results.""" + plt.figure(figsize=(7, 6), dpi=120) + + # Default case: no coloring + colours = None + show_colorbar = False + + # to colour by session + # if os.path.exists(os.path.join(os.path.dirname(p), "all_sessions_info.npy")): + # info = np.load(os.path.join(os.path.dirname(p), "all_sessions_info.npy"), allow_pickle=True).item() + # cell_counts = info.get('session_cell_counts', []) + # names = info.get('session_names', []) + # if cell_counts and names: + # colours = np.concatenate([np.ones(c) * i for i, c in enumerate(cell_counts)]) + + # Color by cluster labels + if os.path.exists(os.path.join(os.path.dirname(p), "umap_cluster_labels.npy")): + cluster_labels = np.load(os.path.join(os.path.dirname(p), "umap_cluster_labels_within7.npy")) + # if clustering within cluster 7 + # Y = Y[cluster_labels == 7] + # cluster_labels = np.load(os.path.join(os.path.dirname(p), "umap_cluster_labels_within7.npy")) + colours = cluster_labels + show_colorbar = True + + scatter = plt.scatter(Y[:, 0], Y[:, 1], s=0.1, alpha=0.1, c=labels, cmap='viridis') + # plt.legend(*scatter.legend_elements(num=9), title="Session", loc="best", markerscale=1, fontsize='small') + + # Add colorbar for cluster labels + if show_colorbar or labels is not None: + unique_labels = np.unique(labels) + n_clusters = len(unique_labels[unique_labels >= 0]) # Exclude noise points (-1) + plt.colorbar(scatter, label=f'Cluster ID ({n_clusters} clusters)', shrink=0.8, alpha=1) + + plt.xlabel("UMAP-1") + plt.ylabel("UMAP-2") + plt.title("UMAP 2D Embedding") + plt.tight_layout() + plt.savefig(p, dpi=300, bbox_inches='tight') + plt.show() + plt.close() + + +def clustering(OUT_DIR): + + Y = np.load(os.path.join(OUT_DIR, "umap_2d_spatfilt.npy")) + print(f"UMAP shape: {Y.shape}") + # original_labels = np.load(os.path.join(OUT_DIR, "umap_cluster_labels.npy")) + # Y = Y[original_labels == 7] + # Perform clustering on Y + clusterer = HDBSCAN(min_cluster_size=500) + cluster_labels = clusterer.fit_predict(Y) + # Save cluster labels + np.save(os.path.join(OUT_DIR, "umap_cluster_labels_spatfilt.npy"), cluster_labels) + + +def cluster_representatives(umap_2d, labels, full_features, save_path=None, save_format='json'): + """ + For each cluster, compute: + - mean: mean of full_features for cluster + - median: real datapoint closest to cluster centroid in UMAP space + Returns a dict with keys 'mean' and 'median', each a list of arrays (len=n_clusters). + Optionally saves the result to file (json or npy). + """ + import numpy as np + unique_labels = np.unique(labels) + means = [] + medians = [] + for cl in unique_labels: + idx = np.where(labels == cl)[0] + umap_cluster = umap_2d[idx] + features_cluster = full_features[idx] + # Mean: mean of full_features in cluster + mean_feat = features_cluster.mean(axis=0) + means.append(mean_feat.tolist()) + # Median: find point closest to centroid in UMAP space + centroid = umap_cluster.mean(axis=0) + dists = np.linalg.norm(umap_cluster - centroid, axis=1) + median_idx = idx[np.argmin(dists)] + median_feat = full_features[median_idx] + medians.append(median_feat.tolist()) + result = {'mean': means, 'median': medians} + if save_path: + if save_format == 'json': + with open(save_path, 'w') as f: + json.dump(result, f) + elif save_format == 'npy': + np.save(save_path, result) + return result + + + +if __name__ == "__main__": + + + H5_PATH = r"\\znas.cortexlab.net\Lab\Share\Ali\for-suyash\data\dataset.h5" + OUT_DIR = r"\\znas.cortexlab.net\Lab\Share\Ali\for-suyash\output" + + with h5py.File(H5_PATH, 'r') as f: + X = f["data"][:, 2, :, :, :] + + # run_permod_pca_umap( + # X=X, + # ncomp_per_mod=32, + # batch_size=4096, + # out_dir=OUT_DIR, + # whiten=True, + # seed=None, + # umap_neighbors=30, + # umap_min_dist=0.1, + # savename="umap_2d", + # ) + + UMAPtime = PCAfunction( + X=X, + ncomp=16, + out_dir=OUT_DIR, + image_shape=(X.shape[1], X.shape[2], X.shape[3]) + ) + diff --git a/curation/hist_viewer.py b/curation/hist_viewer.py new file mode 100644 index 0000000..8566f80 --- /dev/null +++ b/curation/hist_viewer.py @@ -0,0 +1,643 @@ +import numpy as np +import panel as pn +import h5py +from bokeh.plotting import figure +from bokeh.models import Span, ColumnDataSource, BoxAnnotation +from bokeh.palettes import Blues8, Greens8, Reds8 +from bokeh.layouts import gridplot + +class HistViewer: + """Component for viewing histograms of data properties with population, cluster, and individual sample comparisons""" + + def __init__(self, dataset, properties, sample_indices=None, use_sampling=False): + self.dataset = dataset + self.properties = properties + self.sample_indices = sample_indices + self.use_sampling = use_sampling + + # Cache management + self.population_cache = None # Population histograms + self.cluster_cache = {} # Per-cluster histograms + self.current_cluster_id = None + self.current_individual_properties = None + + # Histogram properties to compute + self.property_names = [k for k, v in self.properties.items() if v is not None and len(np.unique(v)) > 1] + + # UI Controls + self.property_selector = pn.widgets.Select( + name="Property to Display", + options=self.property_names, + value=self.property_names[0], + width=200 + ) + + self.n_bins_slider = pn.widgets.IntSlider( + name="Number of Bins", + start=20, end=100, value=50, + width=200 + ) + + # Status text + self.status_text = pn.pane.Markdown("**Status:** Initializing...", width=200) + + # Threshold state + self._current_threshold_mask = None + self.threshold_box = None # Will be created with plot + + # Histogram plot + self.hist_plot = self._create_empty_plot() + self.hist_sources = { + 'population': ColumnDataSource(data=dict(top=[], left=[], right=[])), + 'cluster': ColumnDataSource(data=dict(top=[], left=[], right=[])), + } + self.individual_span = None # Track the current span to remove old ones + + # Set up callbacks + self.property_selector.param.watch(self._on_property_change, 'value') + self.n_bins_slider.param.watch(self._on_bins_change, 'value') + + # Callback for individual sample changes (set by BoxViewer) + self.on_individual_sample_changed = None + + # Threshold classification widgets + self.threshold_mode = pn.widgets.RadioButtonGroup( + name="Threshold Mode", + options=["Off", "Below", "Above", "Range"], + value="Off", + width=200 + ) + + self.threshold_slider = pn.widgets.FloatSlider( + name="Threshold", + start=0, end=1, value=0.5, + step=0.01, + width=200, + visible=False + ) + + self.threshold_range = pn.widgets.RangeSlider( + name="Range", + start=0, end=1, value=(0.25, 0.75), + step=0.01, + width=200, + visible=False + ) + + self.threshold_preview = pn.pane.Markdown("", width=200) + + self.apply_cell_btn = pn.widgets.Button( + name="CELL", + button_type='danger', + width=95, + visible=False + ) + + self.apply_not_cell_btn = pn.widgets.Button( + name="NOT CELL", + button_type='danger', + button_style='outline', + width=95, + visible=False + ) + + # Threshold callbacks + self.threshold_mode.param.watch(self._on_threshold_mode_change, 'value') + self.threshold_slider.param.watch(self._on_threshold_change, 'value') + self.threshold_range.param.watch(self._on_threshold_change, 'value') + self.apply_cell_btn.on_click(self._on_apply_cell_click) + self.apply_not_cell_btn.on_click(self._on_apply_not_cell_click) + + # Callback for threshold classification (set by AppOrchestrator) + self.on_threshold_classify = None + + # Initialize + if self.dataset is not None: + self._compute_population_histograms() + + def _create_empty_plot(self): + """Create empty bokeh plot for histograms""" + plot = figure( + width=700, height=350, # Adjusted width for side-by-side layout with controls + title="Property Distribution", + toolbar_location="above", + x_axis_label="Property Value", + y_axis_label="Probability Density" + ) + plot.title.text_font_size = "12pt" + plot.yaxis.ticker.desired_num_ticks = 0 + + self.threshold_box = BoxAnnotation( + fill_alpha=0.3, + fill_color='orange', + line_color='orange', + line_width=2, + line_dash='dashed', + visible=False + ) + plot.add_layout(self.threshold_box) + + return plot + + def _compute_sample_properties(self, sample_indices): + """Compute all histogram properties for a single sample or batch of samples + + Args: + sample_data: List of indices to use to get the sample data (indexing into self.dataset, the full unsampled dataset) + + Returns: + dict with property_name -> value(s) + """ + if len(sample_indices) == 1: # Single sample + sample_data = self.dataset[sample_indices[0]] + sample_data = sample_data[np.newaxis, ...] # Add batch dimension + else: + sample_data = self.dataset[sample_indices] + + props = {} + for key, value in self.properties.items(): + if value is not None: + if key == 'Edge Cells': + # Special case: Edge Cells is boolean mask + props[key] = value[sample_indices].astype(np.float32) + else: + props[key] = value[sample_indices] + + return props + + def _compute_population_histograms(self): + """Compute histograms for the entire population (or sample)""" + if self.dataset is None: + return + + self.status_text.object = "**Status:** Computing population histograms..." + + try: + # Use same sampling strategy as other components + if self.use_sampling and self.sample_indices is not None: + indices_to_use = np.sort(self.sample_indices) # Sort indices for HDF5 + n_samples = len(self.sample_indices) + sampling_msg = f" (sampled from {self.dataset.shape[0]:,})" + else: + n_samples = min(10000, self.dataset.shape[0]) # Limit for performance + indices_to_use = np.sort(np.random.choice(self.dataset.shape[0], n_samples, replace=False)) + sampling_msg = f" (random sample of {n_samples:,})" + + # Load data in chunks to manage memory + chunk_size = 5000 + all_properties = {prop: [] for prop in self.property_names} + + for i in range(0, len(indices_to_use), chunk_size): + chunk_indices = indices_to_use[i:i+chunk_size] + chunk_properties = self._compute_sample_properties(chunk_indices) + + for prop_name in self.property_names: + all_properties[prop_name].extend(chunk_properties[prop_name]) + + # Convert to numpy arrays and compute normalized histograms (PDFs) + self.population_cache = {} + for prop_name in self.property_names: + values = np.array(all_properties[prop_name]) + hist, bin_edges = np.histogram(values, bins=self.n_bins_slider.value) + + # Normalize to PDF: divide by (total_count * bin_width) + bin_width = bin_edges[1] - bin_edges[0] + pdf_hist = hist / (np.sum(hist) * bin_width) + + self.population_cache[prop_name] = { + 'hist': pdf_hist, + 'bin_edges': bin_edges, + 'values': values + } + + self.status_text.object = f"**Status:** Population histograms ready ({n_samples:,} samples{sampling_msg})" + self._update_plot() + + except Exception as e: + self.status_text.object = f"**Error:** Failed to compute population histograms: {str(e)}" + + def _compute_single_property_histogram(self, prop_name): + """Compute histogram for a single property (used when adding properties dynamically)""" + if self.dataset is None or self.properties.get(prop_name) is None: + return + + try: + if prop_name != 'Probability': + # Use same sampling strategy as population histograms + if self.use_sampling and self.sample_indices is not None: + indices_to_use = self.sample_indices + else: + indices_to_use = np.arange(min(10000, self.dataset.shape[0])) + + values = self.properties[prop_name][indices_to_use] + else: + values = self.properties[prop_name] + hist, bin_edges = np.histogram(values, bins=self.n_bins_slider.value) + + # Normalize to PDF + bin_width = bin_edges[1] - bin_edges[0] + pdf_hist = hist / (np.sum(hist) * bin_width) if np.sum(hist) > 0 else hist + + if self.population_cache is None: + self.population_cache = {} + + self.population_cache[prop_name] = { + 'hist': pdf_hist, + 'bin_edges': bin_edges, + 'values': values + } + + self._update_plot() + except Exception as e: + self.status_text.object = f"**Error:** Failed to compute histogram for {prop_name}: {str(e)}" + + def load_cluster_data(self, cluster_id, display_data): + """Load and cache cluster histograms""" + if self.dataset is None: + return + + # Check cache + if cluster_id in self.cluster_cache: + self.current_cluster_id = cluster_id + self.status_text.object = f"**Status:** Using cached cluster {cluster_id} histograms" + self._update_plot() + return + + try: + # Get cluster indices (same pattern as BoxViewer) + cluster_mask = display_data['cluster'] == cluster_id + cluster_data = display_data[cluster_mask] + cluster_original_indices = cluster_data['original_index'].values + + if len(cluster_original_indices) == 0: + self.status_text.object = f"**Error:** No points in cluster {cluster_id}" + return + + self.status_text.object = f"**Status:** Computing histograms for cluster {cluster_id} ({len(cluster_original_indices)} samples)..." + + # Sort indices for HDF5 access - we don't care about preserving order for histograms + sorted_indices = np.sort(cluster_original_indices) + + # Compute properties + cluster_properties = self._compute_sample_properties(sorted_indices) + + # Compute histograms using same bins as population, normalized to PDFs + cluster_hist_data = {} + for prop_name in self.property_names: + if self.population_cache and prop_name in self.population_cache: + # Use same bin edges as population for comparison + bin_edges = self.population_cache[prop_name]['bin_edges'] + hist, _ = np.histogram(cluster_properties[prop_name], bins=bin_edges) + else: + # Fallback to auto-binning + hist, bin_edges = np.histogram(cluster_properties[prop_name], bins=self.n_bins_slider.value) + + # Normalize to PDF: divide by (total_count * bin_width) + bin_width = bin_edges[1] - bin_edges[0] + pdf_hist = hist / (np.sum(hist) * bin_width) if np.sum(hist) > 0 else hist + + cluster_hist_data[prop_name] = { + 'hist': pdf_hist, + 'bin_edges': bin_edges, + 'values': cluster_properties[prop_name] + } + + self.cluster_cache[cluster_id] = cluster_hist_data + self.current_cluster_id = cluster_id + + self.status_text.object = f"**Status:** Cluster {cluster_id} histograms ready" + self._update_plot() + + except Exception as e: + self.status_text.object = f"**Error:** Failed to load cluster data: {str(e)}" + + def update_individual_sample(self, sample_data): + """Update individual sample marker from current BoxViewer sample""" + if sample_data is None: + self.current_individual_properties = None + self._update_plot() + return + + try: + # Compute properties for the individual sample + individual_props = self._compute_sample_properties(sample_data) + self.current_individual_properties = individual_props + self._update_plot() + + except Exception as e: + self.status_text.object = f"**Error:** Failed to process individual sample: {str(e)}" + + def remove_property(self, name): + """Dynamically remove a property option""" + if name in self.property_names: + self.property_names.remove(name) + self.property_selector.options = self.property_names + self.property_selector.param.trigger('options') + # Reset to first property if current selection was removed + if self.property_selector.value == name and self.property_names: + self.property_selector.value = self.property_names[0] + if name in self.properties: + self.properties[name] = None + if self.population_cache and name in self.population_cache: + del self.population_cache[name] + self._update_plot() + + def update_property(self, name, values): + """Update an existing property with new values from the app orchestrator + + Args: + name: Name of the property to update + values: New values for the property (array-like) + """ + self.properties[name] = values + + # Add to property names if not already present + if name not in self.property_names: + self.property_names.append(name) + self.property_selector.options = self.property_names + self.property_selector.param.trigger('options') + + # Clear all cached cluster histograms since property values changed + # We clear the entire cache because load_cluster_data checks if cluster_id + # is in cache and returns early - partial cache would cause missing histograms + self.cluster_cache = {} + + # Recompute the population histogram for this property + self._compute_single_property_histogram(name) + + def _update_plot(self): + """Update the histogram plot based on current selections""" + selected_property = self.property_selector.value + + # Remove existing individual span if it exists + if self.individual_span is not None: + try: + for l in [self.hist_plot.center, self.hist_plot.left, self.hist_plot.right, self.hist_plot.above, self.hist_plot.below]: + if self.individual_span in l: + l.remove(self.individual_span) + except (ValueError, AttributeError): + pass # Span was already removed or doesn't exist + + self.individual_span = None + + # Clear existing renderers + self.hist_plot.renderers = [] + + try: + if (hasattr(self.hist_plot, 'legend') and + self.hist_plot.legend is not None and + hasattr(self.hist_plot.legend, 'items') and + len(self.hist_plot.legend.items) > 0): + self.hist_plot.legend.items = [] + except (AttributeError, TypeError): + pass + + # Plot population histogram + if self.population_cache and selected_property in self.population_cache: + pop_data = self.population_cache[selected_property] + hist = pop_data['hist'] + bin_edges = pop_data['bin_edges'] + + # Create bar coordinates + left_edges = bin_edges[:-1] + right_edges = bin_edges[1:] + + self.hist_sources['population'].data = dict( + top=hist, + left=left_edges, + right=right_edges + ) + + self.hist_plot.quad( + top='top', bottom=0, left='left', right='right', + source=self.hist_sources['population'], + alpha=0.3, color=Blues8[2], legend_label="Population" + ) + + # Plot cluster histogram + if (self.current_cluster_id is not None and self.current_cluster_id in self.cluster_cache + and selected_property in self.cluster_cache[self.current_cluster_id]): + + cluster_data = self.cluster_cache[self.current_cluster_id][selected_property] + hist = cluster_data['hist'] + bin_edges = cluster_data['bin_edges'] + + left_edges = bin_edges[:-1] + right_edges = bin_edges[1:] + + self.hist_sources['cluster'].data = dict( + top=hist, + left=left_edges, + right=right_edges + ) + + self.hist_plot.quad( + top='top', bottom=0, left='left', right='right', + source=self.hist_sources['cluster'], + alpha=0.6, color=Greens8[4], legend_label=f"Cluster {self.current_cluster_id}" + ) + + # Plot individual sample marker + if (self.current_individual_properties is not None and + selected_property in self.current_individual_properties): + + individual_value = self.current_individual_properties[selected_property] + if hasattr(individual_value, '__len__') and len(individual_value) > 0: + individual_value = individual_value[0] # Take first if array + + self.individual_span = Span( + location=individual_value, + dimension='height', + line_color=Reds8[1], + line_width=2, + line_dash='dashed' + ) + self.hist_plot.add_layout(self.individual_span) + + # Update plot title and labels + self.hist_plot.title.text = f"Distribution: {selected_property}" + self.hist_plot.xaxis.axis_label = selected_property + + # Update legend + self.hist_plot.legend.location = "top_right" + self.hist_plot.legend.click_policy = "hide" + + def _on_property_change(self, event): + """Handle property selector changes""" + self._update_plot() + # Update threshold bounds when property changes + if self.threshold_mode.value != "Off": + self._update_threshold_bounds() + self._update_threshold_visualization() + + def _on_display_change(self, event): + """Handle display option changes""" + self._update_plot() + + def _on_bins_change(self, event): + """Handle bins slider changes - requires recomputing histograms""" + # Clear caches to force recomputation with new bin count + self.population_cache = None + self.cluster_cache = {} + self.current_cluster_id = None + + # Recompute + if self.dataset is not None: + self._compute_population_histograms() + + def clear_cache(self): + """Clear all caches (call when clustering changes)""" + self.cluster_cache = {} + self.current_cluster_id = None + self.current_individual_properties = None + self.status_text.object = "**Status:** Cache cleared" + self._update_plot() + + def _on_threshold_mode_change(self, event): + """Handle threshold mode changes - show/hide appropriate widgets""" + mode = event.new + + # Show/hide single vs range slider + self.threshold_slider.visible = mode in ["Above", "Below"] + self.threshold_range.visible = mode == "Range" + + # Show/hide action buttons + show_buttons = mode != "Off" + self.apply_cell_btn.visible = show_buttons + self.apply_not_cell_btn.visible = show_buttons + + # Update slider ranges based on current property + if mode != "Off": + self._update_threshold_bounds() + + # Update visualization + self._update_threshold_visualization() + + def _update_threshold_bounds(self): + """Update threshold slider bounds based on current property data""" + selected_property = self.property_selector.value + if self.population_cache and selected_property in self.population_cache: + values = self.population_cache[selected_property]['values'] + data_min, data_max = float(np.nanmin(values)), float(np.nanmax(values)) + + # Add small margin + margin = (data_max - data_min) * 0.02 + + # Update single slider + self.threshold_slider.start = data_min - margin + self.threshold_slider.end = data_max + margin + self.threshold_slider.value = (data_min + data_max) / 2 + self.threshold_slider.step = (data_max - data_min) / 100 + + # Update range slider + self.threshold_range.start = data_min - margin + self.threshold_range.end = data_max + margin + self.threshold_range.value = ( + data_min + (data_max - data_min) * 0.25, + data_min + (data_max - data_min) * 0.75 + ) + self.threshold_range.step = (data_max - data_min) / 100 + + def _on_threshold_change(self, event): + """Handle threshold slider value changes""" + self._update_threshold_visualization() + + def _update_threshold_visualization(self): + """Update BoxAnnotation and preview count based on current threshold""" + mode = self.threshold_mode.value + + if mode == "Off" or not self.population_cache: + self.threshold_box.visible = False + self.threshold_preview.object = "" + self._current_threshold_mask = None + return + + selected_property = self.property_selector.value + if selected_property not in self.population_cache: + self.threshold_box.visible = False + self.threshold_preview.object = "" + self._current_threshold_mask = None + return + + # Get property values + values = self.population_cache[selected_property]['values'] + + # Calculate affected indices based on mode + if mode == "Below": + threshold = self.threshold_slider.value + mask = values < threshold + self.threshold_box.left = None # Extends to left edge + self.threshold_box.right = threshold + elif mode == "Above": + threshold = self.threshold_slider.value + mask = values > threshold + self.threshold_box.left = threshold + self.threshold_box.right = None # Extends to right edge + elif mode == "Range": + low, high = self.threshold_range.value + mask = (values >= low) & (values <= high) + self.threshold_box.left = low + self.threshold_box.right = high + + self.threshold_box.visible = True + self._current_threshold_mask = mask + + # Calculate preview counts + count = int(np.sum(mask)) + total = len(values) + percentage = (count / total) * 100 if total > 0 else 0 + + self.threshold_preview.object = f"**Preview:** {count:,} points ({percentage:.1f}%)" + + def get_threshold_mask(self): + """Return boolean mask of points matching current threshold (for orchestrator)""" + return self._current_threshold_mask + + def _on_apply_cell_click(self, event): + """Handle 'Apply as CELL' button click""" + if self.on_threshold_classify: + self.on_threshold_classify('cell') + + def _on_apply_not_cell_click(self, event): + """Handle 'Apply as NOT CELL' button click""" + if self.on_threshold_classify: + self.on_threshold_classify('not_cell') + + def get_layout(self): + """Return the complete HistViewer layout""" + # Main controls in a vertical layout + controls = pn.Column( + self.property_selector, + pn.Spacer(height=10), + self.n_bins_slider, + "### Threshold Classification", + self.threshold_mode, + self.threshold_slider, + self.threshold_range, + self.threshold_preview, + pn.Row(self.apply_cell_btn, self.apply_not_cell_btn, width=200), + self.status_text, + width=240, + margin=(10, 10) + ) + + plot_pane = pn.pane.Bokeh(self.hist_plot, sizing_mode='fixed', height=350, width=700) + + return pn.Row( + pn.Spacer(width=20), + plot_pane, + pn.Spacer(width=30), + controls, + sizing_mode='stretch_width', + margin=(10, 0) + ) + + def __del__(self): + """Clean up HDF5 file handle""" + if self.hdf5_file is not None: + try: + self.hdf5_file.close() + except: + pass + + diff --git a/curation/run_preprocessing.py b/curation/run_preprocessing.py new file mode 100644 index 0000000..126bb64 --- /dev/null +++ b/curation/run_preprocessing.py @@ -0,0 +1,62 @@ +# This script will run the full curation pipeline. It can be run after running the Suite3D pipeline (as in the demo notebook). +# You will need the relevant job directory. +import os +import h5py +import numpy as np + + + +def preprocess(jobdir): + + from curation.dataloader import main + from curation.dimensionality_reduction import PCAfunction + + output_directory = os.path.join(jobdir, 'rois', 'curation') + os.makedirs(output_directory, exist_ok=True) + main(data_directory=jobdir, output_directory=output_directory) + + + # Step 2: Run the dimensionality reduction on the footprints + h5_path = os.path.join(output_directory, 'dataset.h5') + with h5py.File(h5_path, 'r') as f: + X = f["data"][:, 2, :, :, :] + UMAPtime = PCAfunction( + X=X, + ncomp=8, + out_dir=output_directory, + image_shape=(X.shape[1], X.shape[2], X.shape[3]) + ) + + all_sessions_info = np.load(os.path.join(output_directory, 'all_sessions_info.npy'), allow_pickle=True).item() + all_sessions_info['UMAPtime'] = UMAPtime + np.save(os.path.join(output_directory, 'all_sessions_info.npy'), all_sessions_info) + + +if __name__ == "__main__": + + from dataloader import main + from dimensionality_reduction import PCAfunction + + # Step 1: Run the dataloader to load all the ROIs + jobdir = r"C:\Users\suyash\UCL\tinya_data" + output_directory = os.path.join(jobdir, 'rois', 'curation') + os.makedirs(output_directory, exist_ok=True) + main(data_directory=jobdir, output_directory=output_directory) + + + # Step 2: Run the dimensionality reduction on the footprints + h5_path = os.path.join(output_directory, 'dataset.h5') + with h5py.File(h5_path, 'r') as f: + X = f["data"][:, 2, :, :, :] + UMAPtime = PCAfunction( + X=X, + ncomp=8, + out_dir=output_directory, + image_shape=(X.shape[1], X.shape[2], X.shape[3]) + ) + + all_sessions_info = np.load(os.path.join(output_directory, 'all_sessions_info.npy'), allow_pickle=True).item() + all_sessions_info['UMAPtime'] = UMAPtime + np.save(os.path.join(output_directory, 'all_sessions_info.npy'), all_sessions_info) + + diff --git a/curation/umap_visualiser.py b/curation/umap_visualiser.py new file mode 100644 index 0000000..477160f --- /dev/null +++ b/curation/umap_visualiser.py @@ -0,0 +1,379 @@ +import numpy as np +import panel as pn +from bokeh.plotting import figure +from bokeh.models import ColumnDataSource, HoverTool, TapTool, ColorBar, LinearColorMapper +from bokeh.palettes import Category20, Reds256, Inferno256, Viridis256 +from bokeh.colors import RGB + +class UMAPVisualiser: + """Pure UMAP visualisation component - handles only plot rendering and interactions""" + + def __init__(self, data=None, properties=None, sample_indices=None, use_sampling=False): + """ + Initialize with data DataFrame containing: umap_x, umap_y, cluster, classification + + Args: + data: DataFrame with umap_x, umap_y, cluster, classification columns + properties: dict containing ROI properties for coloring: + - shot_noise: array of shot noise values + - footprint_size: array of footprint sizes + - edge_cells: array of edge cell flags + - session_id: array of session IDs + - mean_intensity: array of mean intensity values + - mean_correlation: array of mean correlation values + - peak_value: array of peak values + - contamination: array of contamination values + sample_indices: indices to sample from properties if use_sampling is True + use_sampling: whether to sample properties using sample_indices + """ + self.data = data + self.properties = properties.copy() + for key, value in properties.items(): + if value is None: + continue + else: + if use_sampling: + self.properties[key] = value[sample_indices] + else: + self.properties[key] = value + + self.source = None + self.plot = None + self.scatter = None + self.cluster_colors = Category20[20] + self.view_mode = 'clus' + self.last_clicked_cluster = None + self.prob_mapper = LinearColorMapper(palette=Reds256[::-1], low=0, high=1) + self.color_bar = None + + # Callback for cluster selection - set by orchestrator + self.on_cluster_selected = None + + # Create plot + self.create_empty_plot() + + # Initial render if data provided + if self.data is not None: + self.update_plot() + + def set_probs(self, probs): + """Set probabilities for coloring points""" + if self.data is not None and probs.shape[0] != self.data.shape[0]: + print(self.data.shape, probs.shape) + raise ValueError("Length of probabilities must match number of data points.") + self.properties['Probability'] = probs + + if self.data is not None: + self.update_plot() + + def set_view_mode(self, mode): + """Set view mode directly with a string value""" + mode = mode.lower() + valid_modes = ['clus', 'prob', 'snr', 'size', 'session', 'edge', 'intensity', 'correlation', 'peak', 'contamination', 'vox'] + if mode not in valid_modes: + raise ValueError(f"View mode must be one of {valid_modes}.") + self.view_mode = mode + self.update_plot() + + def update_data(self, new_data, cluster_mode=True): + """Update with new data and re-render""" + self.data = new_data + if cluster_mode: + self.view_mode = 'clus' # Reset to cluster mode on data update + self.update_plot() + + def get_plot_pane(self): + """Return Panel pane containing the plot""" + return pn.pane.Bokeh(self.plot, sizing_mode='stretch_width') + + def create_empty_plot(self): + """Create empty plot optimized for large datasets""" + base_tools = "pan,wheel_zoom,box_zoom,box_select,lasso_select,reset,save" + + self.plot = figure( + width=700, height=500, + title="", + tools=base_tools, + toolbar_location="right", + output_backend="webgl" # Use WebGL for better performance + ) + + # Style the plot + self.plot.xaxis.axis_label = "UMAP Dimension 1" + self.plot.yaxis.axis_label = "UMAP Dimension 2" + + # Optimize plot settings for performance + self.plot.toolbar.autohide = True + + # Set up initial tools + self.update_plot_tools() + + def update_plot_tools(self): + """Update plot tools based on current mode""" + if self.plot is None: + return + + # Clear existing tools and rebuild + self.plot.toolbar.tools = [] + + # Add basic navigation tools + from bokeh.models import PanTool, WheelZoomTool, BoxZoomTool, ResetTool, SaveTool + self.plot.add_tools(PanTool(), WheelZoomTool(), BoxZoomTool(), ResetTool(), SaveTool()) + + # Add tap tool for cluster selection + tap_tool = TapTool() + self.plot.add_tools(tap_tool) + + def on_point_tap(self, attr, old, new): + """Handle point tap in cluster mode""" + if not new or self.data is None: + return + + # Get the cluster of the tapped point + tapped_idx = new[0] + if tapped_idx < len(self.data): + self.last_clicked_cluster = self.data.iloc[tapped_idx]['cluster'] + + # Notify orchestrator if callback is set + if self.on_cluster_selected: + self.on_cluster_selected(self.last_clicked_cluster, tapped_idx) + + def get_point_colors_and_properties(self): + """Color assignment using vectorized operations""" + if self.data is None or len(self.data) == 0: + return [], [] + + n_points = len(self.data) + + # Pre-allocate arrays + colors = np.empty(n_points, dtype=object) + alphas = np.full(n_points, 0.4) + + # Get classification and cluster arrays + classifications = self.data['classification'].values + clusters = self.data['cluster'].values + + # Vectorized color assignment + cell_mask = classifications == 'cell' + not_cell_mask = classifications == 'not_cell' + unclassified_mask = ~(cell_mask | not_cell_mask) + + if self.view_mode == "prob" and self.properties.get('Probability', None) is not None: + # For all ROIs, use probabilities to set colors + for i in range(n_points): + prob = self.properties['Probability'][i] + colors[i] = RGB(r=255, g=int(255 * (1 - prob)), b=int(255 * (1 - prob))) + elif self.view_mode == "snr": + # For all ROIs, use shot noise to set colors + shot = self.properties.get('ROI Shot Noise', None) + norm_noise = (shot - np.min(shot)) / (np.max(shot) - np.min(shot) + 1e-8) + for i in range(n_points): + noise_val = norm_noise[i] + colors[i] = Inferno256[int(noise_val * 255)] + elif self.view_mode == "size": + # For all ROIs, use footprint size to set colors + size = self.properties.get('Footprint Size', None) + norm_size = (size - np.min(size)) / (np.max(size) - np.min(size) + 1e-8) + for i in range(n_points): + size_val = norm_size[i] + colors[i] = Viridis256[int(size_val * 255)] + elif self.view_mode == "session": + # For all ROIs, use session ID to set colors + session_id = self.properties.get('Session', None) + unique_sessions = np.unique(session_id) + session_color_map = {sess: Category20[20][i % 20] for i, sess in enumerate(unique_sessions)} + for i in range(n_points): + sess_id = session_id[i] + colors[i] = session_color_map[sess_id] + elif self.view_mode == "edge": + # For all ROIs, use edge cell status to set colors + for i in range(n_points): + if self.properties['Edge Cells'][i]: + colors[i] = 'red' + else: + colors[i] = 'blue' + elif self.view_mode == "intensity": + # For all ROIs, use mean intensity to set colors + mean_intensity = self.properties.get('Mean Intensity', None) + norm_val = (mean_intensity - np.min(mean_intensity)) / (np.max(mean_intensity) - np.min(mean_intensity) + 1e-8) + for i in range(n_points): + colors[i] = Viridis256[int(norm_val[i] * 255)] + elif self.view_mode == "correlation": + # For all ROIs, use mean correlation to set colors + corr = self.properties.get('Mean Correlation', None) + norm_val = (corr - np.min(corr)) / (np.max(corr) - np.min(corr) + 1e-8) + for i in range(n_points): + colors[i] = Inferno256[int(norm_val[i] * 255)] + elif self.view_mode == "peak": + # For all ROIs, use peak value to set colors + peak = self.properties.get('Peak Value', None) + norm_val = (peak - np.min(peak)) / (np.max(peak) - np.min(peak) + 1e-8) + for i in range(n_points): + colors[i] = Viridis256[int(norm_val[i] * 255)] + elif self.view_mode == "contamination": + # For all ROIs, use contamination to set colors + contamination = self.properties.get('Contamination Factor', None) + norm_val = (contamination - np.min(contamination)) / (np.max(contamination) - np.min(contamination) + 1e-8) + for i in range(n_points): + colors[i] = Reds256[int(norm_val[i] * 255)] + elif self.view_mode == "vox": + # For all ROIs, use voxel SNR to set colors + vox_snr = self.properties.get('Voxel SNR', None) + norm_val = (vox_snr - np.min(vox_snr)) / (np.max(vox_snr) - np.min(vox_snr) + 1e-8) + for i in range(n_points): + colors[i] = Inferno256[int(norm_val[i] * 255)] + else: + # Cluster mode + colors[cell_mask] = 'black' + colors[not_cell_mask] = 'gray' + # For unclassified, use cluster colors + for i in np.where(unclassified_mask)[0]: + colors[i] = self.cluster_colors[clusters[i] % len(self.cluster_colors)] + + # Set alphas + alphas[cell_mask] = 0.8 + alphas[not_cell_mask] = 0.05 + + return colors.tolist(), alphas.tolist() + + def update_plot(self): + """Update the plot with current data""" + if self.data is None or len(self.data) == 0: + return + + # Get colors efficiently + colors, alphas = self.get_point_colors_and_properties() + + # Prepare data source + self.source = ColumnDataSource(data=dict( + x=self.data['umap_x'].values, + y=self.data['umap_y'].values, + cluster=self.data['cluster'].values, + classification=self.data['classification'].values, + colors=colors, + alphas=alphas + )) + + # Clear existing renderers + self.plot.renderers = [] + + # Add hover tool for smaller datasets + if len(self.data) < 50000: + hover = HoverTool(tooltips=[ + ("Cluster", "@cluster"), + ("Classification", "@classification") + ]) + # Remove existing hover tools first + self.plot.toolbar.tools = [tool for tool in self.plot.toolbar.tools if not isinstance(tool, HoverTool)] + self.plot.add_tools(hover) + + # Create scatter plot + self.scatter = self.plot.scatter( + 'x', 'y', + source=self.source, + size=3, + color='colors', + alpha='alphas', + selection_color="red", + nonselection_alpha=0.1 + ) + + no_color_bar = ['clus', 'session', 'edge'] + + if self.view_mode in no_color_bar and self.color_bar is not None and self.color_bar in self.plot.right: + # if in cluster mode, remove color bar if present + self.plot.right.remove(self.color_bar) + elif self.view_mode in no_color_bar: + self.color_bar = None + else: + if self.color_bar is not None and self.color_bar in self.plot.right: + self.plot.right.remove(self.color_bar) + + if self.view_mode == 'prob': + self.color_bar = ColorBar( + color_mapper=self.prob_mapper, + label_standoff=12, + border_line_color=None, + location=(0, 0), + title="Probability" + ) + + elif self.view_mode == 'snr': + self.color_bar = ColorBar( + color_mapper=LinearColorMapper(palette=Inferno256, low=np.min(self.properties['ROI Shot Noise']), high=np.max(self.properties['ROI Shot Noise'])), + label_standoff=12, + border_line_color=None, + location=(0, 0), + title="Shot Noise" + ) + + elif self.view_mode == 'size': + self.color_bar = ColorBar( + color_mapper=LinearColorMapper(palette=Viridis256, low=np.min(self.properties['Footprint Size']), high=np.max(self.properties['Footprint Size'])), + label_standoff=12, + border_line_color=None, + location=(0, 0), + title="Footprint Size" + ) + + elif self.view_mode == 'intensity': + self.color_bar = ColorBar( + color_mapper=LinearColorMapper(palette=Viridis256, low=np.min(self.properties['Mean Intensity']), high=np.max(self.properties['Mean Intensity'])), + label_standoff=12, + border_line_color=None, + location=(0, 0), + title="Mean Intensity" + ) + + elif self.view_mode == 'correlation': + self.color_bar = ColorBar( + color_mapper=LinearColorMapper(palette=Inferno256, low=np.min(self.properties['Mean Correlation']), high=np.max(self.properties['Mean Correlation'])), + label_standoff=12, + border_line_color=None, + location=(0, 0), + title="Mean Correlation" + ) + + elif self.view_mode == 'peak': + self.color_bar = ColorBar( + color_mapper=LinearColorMapper(palette=Viridis256, low=np.min(self.properties['Peak Value']), high=np.max(self.properties['Peak Value'])), + label_standoff=12, + border_line_color=None, + location=(0, 0), + title="Peak Value" + ) + + elif self.view_mode == 'contamination': + self.color_bar = ColorBar( + color_mapper=LinearColorMapper(palette=Reds256, low=np.min(self.properties['Contamination Factor']), high=np.max(self.properties['Contamination Factor'])), + label_standoff=12, + border_line_color=None, + location=(0, 0), + title="Contamination" + ) + + elif self.view_mode == 'vox': + self.color_bar = ColorBar( + color_mapper=LinearColorMapper(palette=Inferno256, low=np.min(self.properties['Voxel SNR']), high=np.max(self.properties['Voxel SNR'])), + label_standoff=12, + border_line_color=None, + location=(0, 0), + title="Voxel SNR" + ) + + self.plot.add_layout(self.color_bar, 'right') + + # Set up tap callback for cluster mode + if self.source: + self.source.selected.on_change('indices', self.on_point_tap) + + def get_selected_indices(self): + """Get currently selected point indices""" + if self.source is None: + return [] + return self.source.selected.indices + + def clear_selection(self): + """Clear current selection""" + if self.source is not None: + self.source.selected.indices = [] \ No newline at end of file diff --git a/demos/Demo-Standard Data.ipynb b/demos/Demo-Standard Data.ipynb index 9461753..04783b5 100644 --- a/demos/Demo-Standard Data.ipynb +++ b/demos/Demo-Standard Data.ipynb @@ -2,23 +2,13 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "id": "f57652f9", "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "/home/ali/anaconda3/envs/suite3d-gpu/lib/python3.8/site-packages/paramiko/pkey.py:100: CryptographyDeprecationWarning: TripleDES has been moved to cryptography.hazmat.decrepit.ciphers.algorithms.TripleDES and will be removed from this module in 48.0.0.\n", - " \"cipher\": algorithms.TripleDES,\n", - "/home/ali/anaconda3/envs/suite3d-gpu/lib/python3.8/site-packages/paramiko/transport.py:258: CryptographyDeprecationWarning: TripleDES has been moved to cryptography.hazmat.decrepit.ciphers.algorithms.TripleDES and will be removed from this module in 48.0.0.\n", - " \"class\": algorithms.TripleDES,\n" - ] - } - ], + "outputs": [], "source": [ - "import os \n", + "import os, sys\n", + "sys.path.insert(0, os.getcwd())\n", "from datetime import date\n", "from matplotlib import pyplot as plt\n", "import numpy as np\n", @@ -30,52 +20,25 @@ "\n", "from suite3d.job import Job\n", "from suite3d import io\n", - "from suite3d import plot_utils as plot" + "from suite3d import plot_utils as plot\n", + "from curation import run_preprocessing" ] }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "id": "f78c8707-e058-4f54-88a1-e24b3ceb0220", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "/mnt/md0/data/demo/lbm/2024-08-10_2_AH012_2P_00001_00001.tif\n", - "/mnt/md0/data/demo/lbm/2024-08-10_2_AH012_2P_00001_00002.tif\n", - "/mnt/md0/data/demo/lbm/2024-08-10_2_AH012_2P_00001_00003.tif\n", - "/mnt/md0/data/demo/lbm/2024-08-10_2_AH012_2P_00001_00004.tif\n", - "/mnt/md0/data/demo/lbm/2024-08-10_2_AH012_2P_00001_00005.tif\n", - "/mnt/md0/data/demo/lbm/2024-08-10_2_AH012_2P_00001_00006.tif\n", - "/mnt/md0/data/demo/lbm/2024-08-10_2_AH012_2P_00001_00007.tif\n", - "/mnt/md0/data/demo/lbm/2024-08-10_2_AH012_2P_00001_00008.tif\n", - "/mnt/md0/data/demo/lbm/2024-08-10_2_AH012_2P_00001_00009.tif\n", - "/mnt/md0/data/demo/lbm/2024-08-10_2_AH012_2P_00001_00010.tif\n" - ] - } - ], - "source": [ - "# update this to point to the demo data! \n", - "tifs = io.get_tif_paths('/mnt/md0/data/demo/lbm')[:10]\n", - "for tif in tifs: print(tif)" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "id": "9f0d542b", - "metadata": {}, "outputs": [], "source": [ "# update this to point to the demo data! \n", - "tifs = io.get_tif_paths('/mnt/md0/data/demo/standard-2p')[:10]" + "tifs = io.get_tif_paths(r'C:\\Users\\suyash\\UCL\\s3d_demo_data\\standard-2p')[:10]\n", + "for tif in tifs: print(tif)" ] }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "id": "2b3e38e5", "metadata": {}, "outputs": [], @@ -117,19 +80,19 @@ }, { "cell_type": "code", - "execution_count": 23, + "execution_count": null, "id": "e02910a1", "metadata": {}, "outputs": [], "source": [ "# Create the job\n", - "job = Job(r'/mnt/md0/runs','demo-std', tifs = tifs,\n", + "job = Job(r'C:\\Users\\suyash\\UCL\\s3d_demo_data','demo-std', tifs = tifs,\n", " params=params, create=True, overwrite=True, verbosity = 3)" ] }, { "cell_type": "code", - "execution_count": 22, + "execution_count": null, "id": "c5310b6e", "metadata": { "scrolled": true @@ -141,7 +104,7 @@ }, { "cell_type": "code", - "execution_count": 9, + "execution_count": null, "id": "d45d99a0", "metadata": {}, "outputs": [], @@ -152,7 +115,7 @@ }, { "cell_type": "code", - "execution_count": 11, + "execution_count": null, "id": "90a72b8c-042a-49c1-988f-73d16e8295e1", "metadata": {}, "outputs": [], @@ -162,15 +125,15 @@ "ref_img = summary['ref_img_3d']\n", "\n", "# # view 1 plane at a time\n", - "# plot.show_img(ref_img[3], figsize=(3,4))\n", + "plot.show_img(ref_img[3], figsize=(3,4))\n", "\n", "# # interactive 3D viewer\n", - "# plot.VolumeViewer(ref_img)\n" + "plot.VolumeViewer(ref_img)\n" ] }, { "cell_type": "code", - "execution_count": 21, + "execution_count": null, "id": "e465e68a", "metadata": { "scrolled": true @@ -182,7 +145,7 @@ }, { "cell_type": "code", - "execution_count": 20, + "execution_count": null, "id": "24f514a1", "metadata": { "scrolled": true @@ -194,7 +157,7 @@ }, { "cell_type": "code", - "execution_count": 14, + "execution_count": null, "id": "3339cfba", "metadata": {}, "outputs": [], @@ -205,7 +168,7 @@ }, { "cell_type": "code", - "execution_count": 18, + "execution_count": null, "id": "920c16f8", "metadata": { "scrolled": true @@ -220,46 +183,14 @@ { "cell_type": "code", "execution_count": null, - "id": "fba5b0c7", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": 19, "id": "81d8ee5c", "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - " Updated main params file\n", - " Movie shape: (5, 4000, 514, 513)\n", - "1736\n", - " Extracting 1736 valid cells, and saving cell flags to /mnt/md0/runs/s3d-demo-std/rois/iscell_extracted.npy\n", - " Extracting activity\n", - " Will extract in 8 batches of 500\n", - " Saving intermediate results to /mnt/md0/runs/s3d-demo-std/rois\n", - " Deconvolving\n", - " Saving to /mnt/md0/runs/s3d-demo-std/rois\n" - ] - } - ], + "outputs": [], "source": [ "job.compute_npil_masks()\n", "traces = job.extract_and_deconvolve()" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "eaf5f023", - "metadata": {}, - "outputs": [], - "source": [] - }, { "cell_type": "code", "execution_count": null, @@ -267,17 +198,9 @@ "metadata": {}, "outputs": [], "source": [ - "job.export_results('path/to/output',result_dir_name='rois')" + "job.export_results(r\"C:\\Users\\suyash\\UCL\\s3d_demo_data\\output\",result_dir_name='rois')" ] }, - { - "cell_type": "code", - "execution_count": null, - "id": "3a9a7fdb", - "metadata": {}, - "outputs": [], - "source": [] - }, { "cell_type": "markdown", "id": "eeaeac6f-0ba5-4f65-95cf-0212a378297f", @@ -295,54 +218,36 @@ "id": "ef0e33fd", "metadata": {}, "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "1fc0c864", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "96c1b3c6", - "metadata": {}, - "outputs": [], - "source": [] - }, - { - "cell_type": "code", - "execution_count": null, - "id": "afbf1e16", - "metadata": {}, - "outputs": [], - "source": [] + "source": [ + "# Run the curation pipeline\n", + "run_preprocessing.preprocess(job.job_dir)" + ] }, { - "cell_type": "code", - "execution_count": null, - "id": "7ff8a7f1", + "cell_type": "markdown", + "id": "28070053", "metadata": {}, - "outputs": [], - "source": [] + "source": [ + "Now you can run the UI for the whole Suite3D pipeline.\n", + "This should open in your browser automatically." + ] }, { "cell_type": "code", "execution_count": null, - "id": "5938038f", + "id": "8aa49d09", "metadata": {}, "outputs": [], - "source": [] + "source": [ + "! panel serve webui/webui.py --show --port 5006" + ] } ], "metadata": { "kernelspec": { - "display_name": "Python [conda env:suite3d-gpu]", + "display_name": "suite3d", "language": "python", - "name": "conda-env-suite3d-gpu-py" + "name": "python3" }, "language_info": { "codemirror_mode": { @@ -354,7 +259,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.8.20" + "version": "3.12.12" } }, "nbformat": 4, diff --git a/neural_network/README.md b/neural_network/README.md new file mode 100644 index 0000000..2f93f64 --- /dev/null +++ b/neural_network/README.md @@ -0,0 +1,8 @@ +## Neural network code + + +This folder contains some of the neural network-related code we experimented with. +We ended up opting for PCA for the dimensionality reduction part of the curation pipeline, so this code was not used in the final method. We include it here in case anyone wants to experiment with nonlinear dimensionality reduction - this may give better curation results but will take a bit of work to get things up and running! + +The idea would be to drop in a neural network's latent representations in place of the PCA embeddings currently used. +There are a couple of different architectures and training functions provided (autoencoder, multi-channel CNN, contrastive learning, augmentation module, etc.). diff --git a/neural_network/autoencoder.py b/neural_network/autoencoder.py new file mode 100644 index 0000000..15a5613 --- /dev/null +++ b/neural_network/autoencoder.py @@ -0,0 +1,180 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.data import DataLoader +from model import HDF5Dataset, Augmentation3D, InfoNCELoss +from tqdm import tqdm +import os +import numpy as np +from umap import UMAP + + +class Autoencoder3D(nn.Module): + def __init__(self, latent_dim=16): + super().__init__() + + # Encoder: (1, 5, 20, 20) -> latent_dim + self.encoder = nn.Sequential( + nn.Conv3d(1, 16, kernel_size=3, stride=(1, 2, 2), padding=1), # -> (16, 5, 10, 10) + nn.ReLU(), + nn.Conv3d(16, 32, kernel_size=3, stride=(1, 2, 2), padding=1), # -> (32, 5, 5, 5) + nn.ReLU(), + ) + + # Bottleneck + self.flatten_size = 32 * 5 * 5 * 5 # 4000 + self.fc_encode = nn.Linear(self.flatten_size, 128) + self.fc_decode = nn.Linear(128, self.flatten_size) + + # Projection head for contrastive learning + self.projection_head = nn.Sequential( + nn.Linear(128, 64), + nn.ReLU(), + nn.Linear(64, latent_dim), + ) + + # Decoder: latent_dim -> (1, 5, 20, 20) + self.decoder = nn.Sequential( + nn.ConvTranspose3d(32, 16, kernel_size=3, stride=(1, 2, 2), padding=1, output_padding=(0, 1, 1)), # -> (16, 5, 10, 10) + nn.ReLU(), + nn.ConvTranspose3d(16, 1, kernel_size=3, stride=(1, 2, 2), padding=1, output_padding=(0, 1, 1)), # -> (1, 5, 20, 20) + ) + + self.latent_dim = latent_dim + + def encode(self, x): + """Encode to latent representation""" + x = self.encoder(x) + x = x.view(x.size(0), -1) + x = self.fc_encode(x) + return x + + def decode(self, z): + """Decode from latent representation""" + x = self.fc_decode(z) + x = x.view(x.size(0), 32, 5, 5, 5) + x = self.decoder(x) + return x + + def forward(self, x): + """Full autoencoder forward pass""" + z = self.encode(x) + x_recon = self.decode(z) + return x_recon, z + + def save_ckpt(self, path): + torch.save({ + 'model_state_dict': self.state_dict(),}, path) + + +def train_autoencoder(model, dataloader, num_epochs=50, lr=1e-3, device='cuda'): + """Stage 1: Train as standard autoencoder""" + model = model.to(device) + print(f"Using device: {device}") + optimizer = torch.optim.Adam(model.parameters(), lr=lr) + criterion = nn.MSELoss() + + model.train() + for epoch in range(num_epochs): + total_loss = 0 + for aug1, aug2 in tqdm(dataloader): + # Assuming batch shape: (B, 5, 20, 20) + images = aug1.to(device) + + optimizer.zero_grad() + recon, _ = model(images) + loss = criterion(recon, images) + loss.backward() + optimizer.step() + + total_loss += loss.item() + + avg_loss = total_loss / len(dataloader) + print(f"Epoch {epoch+1}/{num_epochs}, Reconstruction Loss: {avg_loss:.6f}") + + +def train_contrastive(model, dataloader, loss_func, num_epochs=50, lr=1e-4, device='cuda'): + """Stage 2: Train with contrastive learning""" + model = model.to(device) + optimizer = torch.optim.Adam(model.parameters(), lr=lr) + + model.train() + for epoch in range(num_epochs): + total_loss = 0 + for aug1, aug2 in dataloader: + aug1, aug2 = aug1.to(device), aug2.to(device) + + optimizer.zero_grad() + + # Encode both views + z1 = model.encode(aug1) + z1 = model.projection_head(z1) + z2 = model.encode(aug2) + z2 = model.projection_head(z2) + + # Contrastive loss + loss = loss_func(z1, z2) + loss.backward() + optimizer.step() + + total_loss += loss.item() + + avg_loss = total_loss / len(dataloader) + print(f"Epoch {epoch+1}/{num_epochs}, Contrastive Loss: {avg_loss:.6f}") + + +def extract_features(model, dataloader, device): + """Extract features for downstream analysis (UMAP, etc.)""" + model.eval() + all_features = [] + + with torch.no_grad(): + for view1, _ in tqdm(dataloader, desc='Extracting features'): + view1 = view1.to(device) + features = model.encode(view1) + features = model.projection_head(features) + all_features.append(features.cpu().numpy()) + + return np.concatenate(all_features, axis=0) + + +if __name__ == "__main__": + + model = Autoencoder3D(latent_dim=128) + + hdf5_path = r"\\path\to\dataset.h5" + AEdataset = HDF5Dataset(hdf5_path, dataset_key='data', augmentation=Augmentation3D(prob=0, rotation_range=0), normalize=False) + CLdataset = HDF5Dataset(hdf5_path, dataset_key='data', augmentation=Augmentation3D(), normalize=False) + AEloader = DataLoader(AEdataset, batch_size=256, shuffle=True, num_workers=4, pin_memory=True) + CLloader = DataLoader(CLdataset, batch_size=256, shuffle=True, num_workers=4, pin_memory=True) + + loss = InfoNCELoss() + + model = Autoencoder3D() + print(f"Model has {sum(p.numel() for p in model.parameters()):,} parameters.") + + # Stage 1: Autoencoder training + train_autoencoder(model, AEloader, num_epochs=30, lr=1e-3) + + # Stage 2: Contrastive learning + train_contrastive(model, CLloader, loss, num_epochs=50, lr=1e-3) + + model.save_ckpt(os.path.join("models", "AE16.pth")) + feat = extract_features(model, AEloader, device='cuda') + np.save("AE16_features.npy", feat) + + print("Extracted features shape:", feat.shape) + + umap_model = UMAP( + n_neighbors=30, + min_dist=0.1, + metric="euclidean", + random_state=0, + ) + Y = umap_model.fit_transform(feat) + + out_dir = r"\\path\to\save\output" + + # Save UMAP results + np.save(os.path.join(out_dir, f"AE16_umap.npy"), Y.astype(np.float32)) + diff --git a/neural_network/model.py b/neural_network/model.py new file mode 100644 index 0000000..478f65a --- /dev/null +++ b/neural_network/model.py @@ -0,0 +1,394 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.optim as optim +from torch.utils.data import Dataset, DataLoader +import h5py +import numpy as np +from tqdm import tqdm +from torch.utils.tensorboard import SummaryWriter +from datetime import datetime +from scipy.ndimage import rotate + + +class MultiChannelCNN(nn.Module): + """3D CNN with separate processing for each of the 3 channels""" + + def __init__(self, input_channels=3, channel_feature_dim=64, feature_dim=128): + super(MultiChannelCNN, self).__init__() + + assert input_channels == 3, "This architecture is designed for exactly 3 channels" + + # Create separate CNNs for each channel + self.channel_cnn_0 = SingleChannelCNN(channel_feature_dim) + self.channel_cnn_1 = SingleChannelCNN(channel_feature_dim) + self.channel_cnn_2 = SingleChannelCNN(channel_feature_dim) + + # Fusion layers to combine features from all channels + self.fusion_fc1 = nn.Linear(3 * channel_feature_dim, 256) + self.fusion_bn = nn.BatchNorm1d(256) + self.fusion_fc2 = nn.Linear(256, feature_dim) + + def forward(self, x): + # x shape: (batch, 3, 5, 20, 20) + batch_size = x.size(0) + + # Split into individual channels + channel_0 = x[:, 0:1, :, :, :] # (batch, 1, 5, 20, 20) + channel_1 = x[:, 1:2, :, :, :] # (batch, 1, 5, 20, 20) + channel_2 = x[:, 2:3, :, :, :] # (batch, 1, 5, 20, 20) + + # Process each channel independently + feat_0 = self.channel_cnn_0(channel_0) # (batch, channel_feature_dim) + feat_1 = self.channel_cnn_1(channel_1) # (batch, channel_feature_dim) + feat_2 = self.channel_cnn_2(channel_2) # (batch, channel_feature_dim) + + # Concatenate channel features + combined_features = torch.cat([feat_0, feat_1, feat_2], dim=1) # (batch, 3*channel_feature_dim) + + # Fusion layers + x = F.relu(self.fusion_bn(self.fusion_fc1(combined_features))) + x = self.fusion_fc2(x) # (batch, feature_dim) + + return x + + +class SingleChannelCNN(nn.Module): + def __init__(self, channel_feature_dim=64, hidden_dim=32): + super(SingleChannelCNN, self).__init__() + + self.conv_in = nn.Conv3d(1, hidden_dim, kernel_size=3, padding=1) + self.bn_in = nn.BatchNorm3d(hidden_dim) + + self.conv1 = nn.Conv3d(hidden_dim, hidden_dim, kernel_size=3, padding=1) + self.bn1 = nn.BatchNorm3d(hidden_dim) + + self.conv2 = nn.Conv3d(hidden_dim, hidden_dim, kernel_size=3, + stride=(1, 2, 2), padding=1) + self.bn2 = nn.BatchNorm3d(hidden_dim) + + self.downsample_skip = nn.AvgPool3d(kernel_size=(1, 2, 2), stride=(1, 2, 2)) + self.global_pool = nn.AdaptiveAvgPool3d(1) + + def forward(self, x): + + x = F.relu(self.bn_in(self.conv_in(x))) + + identity = x + out = self.conv1(x) + out = self.bn1(out) + out += identity + x = F.relu(out) + + identity = self.downsample_skip(x) + out = self.conv2(x) + out = self.bn2(out) + out += identity + out = F.relu(out) + + out = self.global_pool(out) + out = out.view(out.size(0), -1) + + return out + + +class ProjectionHead(nn.Module): + """MLP projection head for contrastive learning""" + + def __init__(self, input_dim=128, hidden_dim=256, output_dim=128): + super(ProjectionHead, self).__init__() + self.projection = nn.Sequential( + nn.Linear(input_dim, hidden_dim), + nn.ReLU(inplace=True), + nn.Linear(hidden_dim, output_dim) + ) + + def forward(self, x): + return self.projection(x) + + +class ContrastiveModel(nn.Module): + """Complete contrastive learning model""" + + def __init__(self, backbone_feature_dim=128, output_dim=32): + super(ContrastiveModel, self).__init__() + self.backbone = SingleChannelCNN(channel_feature_dim=32) + self.projection_head = ProjectionHead( + input_dim=32, + output_dim=output_dim + ) + + def forward(self, x): + features = self.backbone(x) + projections = self.projection_head(features) + return features, projections + + +class Augmentation3D: + """ + Augmentation class for 3D volumes (3, 5, 20, 20). + Apply augmentations suitable for small 3D data. + """ + + def __init__(self, + rotation_range=15, # degrees + prob=0.5, + gaussian_noise_std=0.1, + shot_noise_std=0.1, + brightness_range=0.2 + ): + + self.rotation_range = rotation_range + self.prob = prob + self.gaussian_noise_std = gaussian_noise_std + self.shot_noise_std = shot_noise_std + self.brightness_range = brightness_range + + def __call__(self, volume): + """ + Apply random augmentations to a 3D volume. + + Args: + volume: numpy array of shape (3, 5, 20, 20) + + Returns: + Augmented volume of same shape + """ + volume = volume.copy() + + # Random rotations + if self.rotation_range > 0: + # for channel in range(volume.shape[0]): + # # Rotation in XY plane (around Z axis) + # angle_z = np.random.uniform(-self.rotation_range, self.rotation_range) + # volume[channel] = rotate(volume[channel], angle_z, axes=(1, 2), reshape=False, order=1) + volume = rotate(volume, np.random.uniform(-self.rotation_range, self.rotation_range), + axes=(1, 2), reshape=False, order=1) + + # Random flip along z-axis (not same as 180-degree rotation) + if np.random.random() < self.prob: + # Flip along depth axis (axis=1 after channel) + volume = np.flip(volume, axis=1).copy() + + # Gaussian noise + if np.random.random() < self.prob and self.gaussian_noise_std > 0: + noise = np.random.normal(0, self.gaussian_noise_std, volume.shape) + volume += noise + + # # Shot noise + # if np.random.random() < self.prob and self.shot_noise_std > 0: + # noise = np.random.normal(0, self.shot_noise_std, volume.shape) + # volume += volume * noise # element-wise multiplication + + # # Brightness adjustment + # if np.random.random() < self.prob and self.brightness_range > 0: + # brightness_factor = np.random.uniform(-self.brightness_range, self.brightness_range) + # volume[:2, :, :] += brightness_factor + + # Shift within plane - this wraps around values that go out of bounds + if np.random.random() * 2 < self.prob: + shift_x = np.random.randint(-5, 5) + shift_y = np.random.randint(-5, 5) + shift_z = np.random.randint(-2, 3) + volume = np.roll(volume, shift=(shift_z, shift_x, shift_y), axis=(0, 1, 2)) + + return volume.astype(np.float32) + + +class HDF5Dataset(Dataset): + """Dataset class for loading 3D data from HDF5 file with augmentations""" + + def __init__(self, hdf5_path, dataset_key='data', augmentation=None, normalize=True): + self.hdf5_path = hdf5_path + self.dataset_key = dataset_key + self.augmentation = augmentation + self.normalize = normalize + + # Get dataset length without loading all data + with h5py.File(hdf5_path, 'r') as f: + self.length = len(f[dataset_key]) + + # Compute dataset statistics for normalization (sample-based) + if self.normalize: + # Sample a subset to compute stats (avoid loading full dataset) + n_samples = min(1000, self.length) + indices = np.random.choice(self.length, n_samples, replace=False) + samples = [f[dataset_key][i, 2, :, :, :] for i in indices] + all_samples = np.stack(samples) + + self.mean = all_samples.mean() + self.std = all_samples.std() + print(f"Dataset statistics: mean={self.mean:.4f}, std={self.std:.4f}") + + def __len__(self): + return self.length + + def __getitem__(self, idx): + with h5py.File(self.hdf5_path, 'r') as f: + # sample = f[self.dataset_key][idx] # Shape: (3, 5, 20, 20) + sample = f[self.dataset_key][idx, 2, :, :, :] # Shape: (5, 20, 20) + sample = np.expand_dims(sample, axis=0) # Shape: (1, 5, 20, 20) + + sample = sample.astype(np.float32) + + if self.normalize: + sample = (sample - self.mean) / (self.std + 1e-8) + + # Create two augmented views for contrastive learning + if self.augmentation is not None: + view1 = self.augmentation(sample) + view2 = self.augmentation(sample) + else: + view1 = sample.copy() + view2 = sample.copy() + + return torch.from_numpy(view1), torch.from_numpy(view2) + + +class InfoNCELoss(nn.Module): + """InfoNCE loss for contrastive learning""" + + def __init__(self, temperature=0.1): + super(InfoNCELoss, self).__init__() + self.temperature = temperature + + def forward(self, z1, z2): + """ + Args: + z1, z2: projections of shape (batch_size, projection_dim) + """ + batch_size = z1.shape[0] + + z1 = F.normalize(z1, dim=1) + z2 = F.normalize(z2, dim=1) + + z = torch.cat([z1, z2], dim=0) # (2 * batch_size, projection_dim) + + # Compute similarity matrix + sim_matrix = torch.mm(z, z.t()) / self.temperature # (2 * batch_size, 2 * batch_size) + + labels = torch.arange(2 * batch_size, device=z.device) + labels[:batch_size] += batch_size + labels[batch_size:] -= batch_size + + # Mask out self-similarities + mask = torch.eye(2 * batch_size, device=z.device).bool() + sim_matrix.masked_fill_(mask, -float('inf')) + + # Compute cross-entropy loss + loss = F.cross_entropy(sim_matrix, labels) + + return loss + + +class ContrastiveLearner: + """Training class for contrastive learning""" + + def __init__(self, model, device='cuda', lr=1e-3, temperature=0.1, log_dir=None, experiment_name=None): + self.model = model.to(device) + self.device = device + self.optimizer = optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4) + self.criterion = InfoNCELoss(temperature=temperature) + self.scheduler = optim.lr_scheduler.CosineAnnealingLR( + self.optimizer, T_max=100, eta_min=1e-6 + ) + + if log_dir is None: + timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") + exp_name = experiment_name or f"contrastive_3d_{timestamp}" + log_dir = f"runs/{exp_name}" + + self.writer = SummaryWriter(log_dir) + self.global_step = 0 + print(f"TensorBoard logging to: {log_dir}") + print(f"View logs with: tensorboard --logdir={log_dir}") + print(f"Using device: {device}") + + def train_epoch(self, dataloader, epoch): + self.model.train() + total_loss = 0 + num_batches = 0 + batch_losses = [] + + pbar = tqdm(dataloader, desc=f'Epoch {epoch+1}') + for batch_idx, (view1, view2) in enumerate(pbar): + view1, view2 = view1.to(self.device), view2.to(self.device) + + features1, projections1 = self.model(view1) + features2, projections2 = self.model(view2) + + loss = self.criterion(projections1, projections2) + + self.optimizer.zero_grad() + loss.backward() + self.optimizer.step() + + # Track metrics + batch_loss = loss.item() + total_loss += batch_loss + batch_losses.append(batch_loss) + num_batches += 1 + self.global_step += 1 + + # Log to TensorBoard every 50 batches + if batch_idx % 50 == 0: + self.writer.add_scalar('Loss/Batch', batch_loss, self.global_step) + self.writer.add_scalar('Learning_Rate', + self.optimizer.param_groups[0]['lr'], + self.global_step) + + pbar.set_postfix({ + 'loss': f'{batch_loss:.4f}', + 'avg_loss': f'{total_loss/num_batches:.4f}' + }) + + avg_epoch_loss = total_loss / num_batches + + # Log epoch-level metrics + self.writer.add_scalar('Loss/Epoch', avg_epoch_loss, epoch) + self.writer.add_scalar('Loss/Epoch_Std', np.std(batch_losses), epoch) + self.writer.add_histogram('Loss/Batch_Distribution', np.array(batch_losses), epoch) + + # Log model parameters and gradients every 10 epochs + if epoch % 10 == 0: + for name, param in self.model.named_parameters(): + if param.grad is not None: + self.writer.add_histogram(f'Gradients/{name}', param.grad, epoch) + self.writer.add_histogram(f'Parameters/{name}', param, epoch) + + return avg_epoch_loss + + def extract_features(self, dataloader): + """Extract features for downstream analysis (UMAP, etc.)""" + self.model.eval() + all_features = [] + + with torch.no_grad(): + for view1, _ in tqdm(dataloader, desc='Extracting features'): + view1 = view1.to(self.device) + features, projections = self.model(view1) + all_features.append(projections.cpu().numpy()) + + return np.concatenate(all_features, axis=0) + + def save_model(self, path): + torch.save({ + 'model_state_dict': self.model.state_dict(), + 'optimizer_state_dict': self.optimizer.state_dict(), + 'scheduler_state_dict': self.scheduler.state_dict(), + 'global_step': self.global_step, + }, path) + + def load_model(self, path): + checkpoint = torch.load(path, map_location=self.device) + self.model.load_state_dict(checkpoint['model_state_dict']) + self.optimizer.load_state_dict(checkpoint['optimizer_state_dict']) + if 'scheduler_state_dict' in checkpoint: + self.scheduler.load_state_dict(checkpoint['scheduler_state_dict']) + if 'global_step' in checkpoint: + self.global_step = checkpoint['global_step'] + + def close_tensorboard(self): + """Close TensorBoard writer""" + self.writer.close() + diff --git a/neural_network/train.py b/neural_network/train.py new file mode 100644 index 0000000..ddf13d5 --- /dev/null +++ b/neural_network/train.py @@ -0,0 +1,109 @@ +from neural_network.model import * +import os + + +def train_contrastive_model(hdf5_path, dataset_key='data', batch_size=256, output_dim=32, num_epochs=500, device='cuda', + save_path='contrastive_model.pth', log_dir=None, experiment_name=None, augmentation_config=None): + """Main training function""" + + # Create dataset and dataloader + dataset = HDF5Dataset(hdf5_path, dataset_key=dataset_key, augmentation=augmentation_config) + dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True) + + model = ContrastiveModel(backbone_feature_dim=128, output_dim=output_dim) + + if os.path.exists(os.path.dirname(save_path)): + checkpoint = torch.load(os.path.join(os.path.dirname(save_path), "ckpt_best.pth"), map_location=device) + model.load_state_dict(checkpoint['model_state_dict']) + print(f"Loaded checkpoint from {os.path.join(os.path.dirname(save_path), 'ckpt_best.pth')}") + + learner = ContrastiveLearner( + model, + device=device, + lr=1e-3, + log_dir=log_dir, + experiment_name=experiment_name + ) + + learner.writer.add_text('Model/Architecture', 'Single Channel CNN on Footprint Only') + learner.writer.add_text('Model/Total_Parameters', f'{370880:,} parameters') + learner.writer.add_text('Training/Batch_Size', str(batch_size)) + learner.writer.add_text('Training/Learning_Rate', str(1e-3)) + learner.writer.add_text('Training/Dataset_Size', f'{len(dataset):,} samples') + + os.makedirs(save_path, exist_ok=True) + os.makedirs(log_dir, exist_ok=True) + + # Training loop + best_loss = float('inf') + try: + for epoch in range(num_epochs): + # Train epoch + avg_loss = learner.train_epoch(dataloader, epoch) + learner.scheduler.step() + + # Print progress + current_lr = learner.optimizer.param_groups[0]['lr'] + print(f"Epoch {epoch+1}/{num_epochs}") + print(f" Average Loss: {avg_loss:.4f}") + print(f" Learning Rate: {current_lr:.6f}") + + # Save best model + if avg_loss < best_loss: + best_loss = avg_loss + learner.save_model(os.path.join(save_path, "ckpt_best.pth")) + learner.writer.add_scalar('Loss/Best', best_loss, epoch) + + # Save checkpoint every 10 epochs + if (epoch + 1) % 10 == 0: + checkpoint_path = os.path.join(save_path, f"ckpt_epoch_{epoch+1}.pth") + learner.save_model(checkpoint_path) + print(f" Saved checkpoint: {checkpoint_path}") + + # Log validation metrics + # learner.writer.add_scalar('Loss/Validation', val_loss, epoch) + + except KeyboardInterrupt: + print("\nTraining interrupted by user") + + finally: + learner.save_model(os.path.join(save_path, f"ckpt_epoch_{epoch+1}.pth")) + print(f"Training completed! Final model saved to {os.path.join(save_path, f'ckpt_epoch_{epoch+1}.pth')}") + print(f"Best model saved to {os.path.join(save_path, 'ckpt_best.pth')} (loss: {best_loss:.4f})") + + learner.close_tensorboard() + + return learner + + +if __name__ == "__main__": + + hdf5_path = r"\\path\to\dataset.h5" + + exp_name = "footprint_only_contrastive_16d_augmented" + + augmentations = Augmentation3D() # default augmentation settings + + learner = train_contrastive_model( + hdf5_path=hdf5_path, + dataset_key='data', + batch_size=256, + output_dim=16, + num_epochs=300, + device='cuda', + save_path=os.path.join(r"path\to\save\models", exp_name, "ckpt"), + log_dir=os.path.join(r"path\to\save\models", exp_name, "logs"), + experiment_name=exp_name, + augmentation_config=augmentations + ) + + # Extract trained model embeddings + dataset = HDF5Dataset(hdf5_path, dataset_key='data') + dataloader = DataLoader(dataset, batch_size=256, shuffle=False, num_workers=4) + + features = learner.extract_features(dataloader) + print(f"Extracted features shape: {features.shape}") + + # Save features for later use + np.save(f'{exp_name}_features.npy', features) + print(f"Features saved to {exp_name}_features.npy") diff --git a/suite3d/quality_metrics.py b/suite3d/quality_metrics.py index e880d99..77d6142 100644 --- a/suite3d/quality_metrics.py +++ b/suite3d/quality_metrics.py @@ -1,4 +1,5 @@ import numpy as n +import numba def volume_quality(volume, pct_high = 99.99, pct_low = 25.0): @@ -45,6 +46,19 @@ def shot_noise_pct(fs, frate_hz): return noise_level +@numba.jit(nopython=True, parallel=True) +def shot_noise_suyash(fs, frate_hz): + npix, nt = fs.shape + noise_level = n.empty(npix) + + for i in numba.prange(npix): + row = fs[i] + mean_f = n.mean(row) + df = n.diff(row) / mean_f + noise_level[i] = n.nanmedian(n.abs(n.diff(df))) + + return noise_level / frate_hz + def choose_top_pix(vol, pct = 98): nz, ny, nx = vol.shape pcts = n.array([n.percentile(vol[i].flatten(), pct) for i in range(nz)]) diff --git a/suite3d/register_gpu.py b/suite3d/register_gpu.py index daa2b7d..047ff29 100644 --- a/suite3d/register_gpu.py +++ b/suite3d/register_gpu.py @@ -5,8 +5,10 @@ try: from mkl_fft import fft2, ifft2 + using_scipy_fft = False except ImportError: from scipy.fft import fft2, ifft2 + using_scipy_fft = True try: import cupy as cp @@ -337,10 +339,16 @@ def convolve_2d_gpu(mov, ref_f, axes=(1,2)): return mov def convolve_2d_cpu(mov, ref_f): - mov[:] = fft2(mov, axes=(1,2), overwrite_x=True) + if using_scipy_fft: + mov[:] = fft2(mov, axes=(1,2), overwrite_x=True) + else: + mov[:] = fft2(mov, axes=(1,2)) mov /= n.abs(mov) + n.complex64(1e-5) mov *= ref_f - mov[:] = ifft2(mov, axes=(1,2), overwrite_x=True) + if using_scipy_fft: + mov[:] = ifft2(mov, axes=(1,2), overwrite_x=True) + else: + mov[:] = ifft2(mov, axes=(1,2)) return mov #TODO xpad/ypad should be integer ? diff --git a/webui/curation_panel.py b/webui/curation_panel.py new file mode 100644 index 0000000..7786f6d --- /dev/null +++ b/webui/curation_panel.py @@ -0,0 +1,75 @@ +import panel as pn +import param +import numpy as np +import h5py +from pathlib import Path +import sys + +sys.path.insert(0, '.') +from curation.app import AppOrchestrator + + +class CurationPanel(param.Parameterized): + """Panel wrapper for the curation app to integrate with webui""" + + def __init__(self, max_height=800): + super().__init__() + self.max_height = max_height + self.orchestrator = None + self.job = None + + # Create placeholder layout + self.layout = pn.Column( + pn.pane.Markdown("## Curation", width=700), + pn.pane.Markdown("Load a job to begin curation.", width=700), + name="Curation", + max_height=max_height + ) + + def load_job(self, job_interface): + """Load data from Suite3D job and initialize curation app""" + try: + self.job = job_interface.job + jobdir = Path(job_interface.job_data['jobdir']) + + # Check if curation data exists + curation_dir = jobdir / 'rois' / 'curation' + if not curation_dir.exists(): + self.layout[:] = [ + pn.pane.Markdown("## Curation", width=700), + pn.pane.Markdown( + "**No curation data found.**\n\n", + width=700 + ) + ] + return + + # Load curation files + umap_file = str(curation_dir / 'umap_2d.npy') + features_file = str(curation_dir / 'pca_embeddings.npy') + hdf5_file = str(curation_dir / 'dataset.h5') + + # Check if files exist + if not all(Path(f).exists() for f in [umap_file, features_file, hdf5_file]): + missing = [f for f in [umap_file, features_file, hdf5_file] if not Path(f).exists()] + self.layout[:] = [ + pn.pane.Markdown("## Curation", width=700), + pn.pane.Markdown( + f"**Missing curation files:**\n\n" + "\n".join(f"- {Path(f).name}" for f in missing), + width=700 + ) + ] + return + + # Create orchestrator with job data + self.orchestrator = AppOrchestrator(umap_file, features_file, hdf5_file) + + # Replace placeholder with actual curation interface + self.layout[:] = [self.orchestrator.get_layout()] + + except Exception as e: + self.layout[:] = [ + pn.pane.Markdown("## Curation", width=700), + pn.pane.Markdown(f"**Error loading curation data:**\n\n{str(e)}", width=700) + ] + print(f"Error in CurationPanel.load_job: {e}") diff --git a/webui/webui.py b/webui/webui.py index ee0d306..451376c 100644 --- a/webui/webui.py +++ b/webui/webui.py @@ -1,9 +1,8 @@ -import numpy as n import os +os.environ["OMP_NUM_THREADS"] = "6" import sys from pathlib import Path - import panel as pn from bokeh.io import curdoc @@ -24,6 +23,7 @@ from webui.init_pass_panel import InitPanel from webui.corrmap_panel import CorrmapPanel from webui.registration_panel import RegistrationPanel +from webui.curation_panel import CurationPanel pn.extension(design="native") @@ -36,9 +36,10 @@ init_panel = InitPanel(max_height=800) reg_panel = RegistrationPanel(max_height=800) corrmap_panel = CorrmapPanel(max_height=800) +curation_panel = CurationPanel(max_height=800) -ui = pn.Tabs(job_interface.job_widget, init_panel.layout, reg_panel.layout, corrmap_panel.layout)#, volume_vis_panel) +ui = pn.Tabs(job_interface.job_widget, init_panel.layout, reg_panel.layout, corrmap_panel.layout, curation_panel.layout) def job_load_callback(value): if value: @@ -54,6 +55,10 @@ def job_load_callback(value): corrmap_panel.load_job(job_interface) except: print("Could not load corrmap panel") + try: + curation_panel.load_job(job_interface) + except: + print("Could not load curation panel") pn.bind(job_load_callback, job_interface.param.job_loaded, watch=True)