""" Build the augmented figure-scoped master units table. This builder preserves the base unit-table columns from build_master_unit_table.py, then adds only the derived columns confirmed by docs/unit_column_usage_audit.md: pulse_high_mean_evoked_rate_zscored pulse_high_first_spike_latency pulse_high_first_spike_jitter raised_cosine_high_fraction_time_responsive SST, VIP, RS, FS brain_division, no_anomalies, cluster_labels_new Legacy derived columns such as old `cluster_labels`, unused optotagging metrics, per-image responsiveness columns, and RF metrics are intentionally omitted. The output path is configurable; use the legacy `master_units_with_responsiveness.csv` filename when running existing notebooks. """ import argparse from pathlib import Path import numpy as np import pandas as pd DEFAULT_MANIFEST = "visual-behavior-neuropixels_project_manifest_v0.5.0.json" OPTO_COLUMNS = [ "pulse_high_mean_evoked_rate_zscored", "pulse_high_first_spike_latency", "pulse_high_first_spike_jitter", "raised_cosine_high_fraction_time_responsive", ] DERIVED_COLUMNS_TO_DROP = { "uid", "unit_id.1", "cluster_labels", "responds_to_at_least_one_nonchange_image", } DERIVED_PREFIXES_TO_DROP = ( "pulse_low_", "pulse_med_", "raised_cosine_low_", "raised_cosine_med_", ) RESPONSIVENESS_PATTERNS_TO_DROP = ( "_nonchange_", "change_response_pval", "change_positive_modulation", "change_mean_evoked", "change_peak_evoked", ) RF_COLUMNS_TO_DROP = { "azimuth_rf", "elevation_rf", "width_rf", "height_rf", "area_rf", "p_value_rf", "on_screen_rf", "is_inverted", } def load_cache(cache_dir, manifest=DEFAULT_MANIFEST, cache_source="local"): """Load a VisualBehaviorNeuropixelsProjectCache without import-time side effects.""" from allensdk.brain_observatory.behavior.behavior_project_cache.\ behavior_neuropixels_project_cache import VisualBehaviorNeuropixelsProjectCache if cache_source == "s3": cache = VisualBehaviorNeuropixelsProjectCache.from_s3_cache(cache_dir=cache_dir) else: cache = VisualBehaviorNeuropixelsProjectCache.from_local_cache(cache_dir=cache_dir) if manifest: cache.load_manifest(manifest) return cache def _reset_index_column(table, column_name): """Return a copy with `column_name` as a plain column.""" table = table.copy() if column_name in table.columns: return table table = table.reset_index() if column_name in table.columns: return table first = table.columns[0] return table.rename(columns={first: column_name}) def _read_csv_without_unnamed(path): table = pd.read_csv(path) unnamed = [col for col in table.columns if str(col).startswith("Unnamed")] return table.drop(columns=unnamed) def load_base_units(cache=None, base_unit_table_file=None): """Load the base units table, preserving all non-bookkeeping columns.""" if base_unit_table_file is not None: return _read_csv_without_unnamed(base_unit_table_file) if cache is None: raise ValueError("Either cache or base_unit_table_file is required.") return _reset_index_column(cache.get_unit_table(), "unit_id") def load_session_table(cache=None, session_table_file=None): if session_table_file is not None: return _read_csv_without_unnamed(session_table_file) if cache is None: return None return _reset_index_column( cache.get_ecephys_session_table(filter_abnormalities=False), "ecephys_session_id", ) def merge_session_metadata(units, sessions): """Merge all session metadata columns, preserving existing unit-table columns.""" if sessions is None or "ecephys_session_id" not in units.columns: return units session_cols = [ col for col in sessions.columns if col == "ecephys_session_id" or col not in units.columns ] return units.merge( sessions[session_cols], on="ecephys_session_id", how="left", ) def add_no_anomalies(units): """Add no_anomalies from abnormality flags if it is not already present.""" if "no_anomalies" in units.columns: return units if {"abnormal_activity", "abnormal_histology"}.issubset(units.columns): units["no_anomalies"] = ( units["abnormal_activity"].isna() & units["abnormal_histology"].isna() ) return units def _list_parents(area, structure_tree): if pd.isna(area) or area == "root": return [] parent_list = [] current_area = area parent_structure_id = 0.0 while not pd.isna(parent_structure_id): row = structure_tree[structure_tree["acronym"] == current_area] if len(row) == 0: break parent_list.append(row["name"].values[0]) parent_structure_id = row["parent_structure_id"].values[0] parent = structure_tree[structure_tree["id"] == parent_structure_id] if len(parent) == 0: break current_area = parent["acronym"].values[0] return parent_list def _brain_division_for_area(area, structure_tree): divisions = [ "Isocortex", "Hippocampal formation", "Thalamus", "Midbrain", "Hypothalamus", "Striatum", ] parents = _list_parents(area, structure_tree) intersection = np.intersect1d(parents, divisions) if len(intersection) > 0: return intersection[0] return "not in list" def add_brain_division(units, structure_tree_file=None): """Add brain_division if missing and a CCF structure tree is provided.""" if "brain_division" in units.columns: return units if structure_tree_file is None: return units structure_tree = pd.read_csv(structure_tree_file) units["brain_division"] = units["structure_acronym"].apply( lambda area: _brain_division_for_area(area, structure_tree) ) return units def merge_opto_metrics(units, opto_metrics_file=None): """Merge only audited optotagging columns required for cell-type labels.""" if opto_metrics_file is None: for col in OPTO_COLUMNS: if col not in units.columns: units[col] = np.nan return units opto = _read_csv_without_unnamed(opto_metrics_file) if "unit_id" not in opto.columns and "uid" in opto.columns: opto = opto.rename(columns={"uid": "unit_id"}) missing = [col for col in ["unit_id"] + OPTO_COLUMNS if col not in opto.columns] if missing: raise ValueError(f"Opto metrics file is missing columns: {missing}") keep = ["unit_id"] + OPTO_COLUMNS units = units.drop(columns=[col for col in OPTO_COLUMNS if col in units.columns]) return units.merge(opto[keep], on="unit_id", how="left") def derive_cell_type_columns(units): """Derive SST, VIP, RS, and FS from audited optotagging thresholds.""" for col in OPTO_COLUMNS: if col not in units.columns: units[col] = np.nan genotype = units.get("genotype", pd.Series("", index=units.index)).fillna("") waveform_duration = units.get( "waveform_duration", pd.Series(np.nan, index=units.index), ) opto_responsive = ( (units["pulse_high_mean_evoked_rate_zscored"] > 2) & (units["pulse_high_first_spike_latency"] < 0.008) & (units["pulse_high_first_spike_jitter"] < 0.002) & (units["raised_cosine_high_fraction_time_responsive"] > 0.3) ) units["SST"] = genotype.str.contains("Sst", na=False) & opto_responsive units["VIP"] = genotype.str.contains("Vip", na=False) & opto_responsive units["RS"] = (waveform_duration > 0.4) & ~units["SST"] & ~units["VIP"] units["FS"] = (waveform_duration < 0.4) & ~units["SST"] & ~units["VIP"] return units def merge_cluster_labels(units, cluster_labels_file=None): """Merge cluster_labels_new and drop legacy old-cluster columns.""" units = units.drop(columns=[col for col in ("cluster_labels",) if col in units.columns]) if cluster_labels_file is None: if "cluster_labels_new" not in units.columns: units["cluster_labels_new"] = np.nan return units clusters = _read_csv_without_unnamed(cluster_labels_file) if "cluster_labels_new_weak_cluster_12" in clusters.columns: clusters = clusters.drop(columns=["cluster_labels_new"], errors="ignore") clusters = clusters.rename( columns={"cluster_labels_new_weak_cluster_12": "cluster_labels_new"} ) elif "cluster_labels_new" not in clusters.columns and "cluster_labels" in clusters.columns: clusters = clusters.rename(columns={"cluster_labels": "cluster_labels_new"}) missing = [col for col in ("unit_id", "cluster_labels_new") if col not in clusters.columns] if missing: raise ValueError(f"Cluster labels file is missing columns: {missing}") units = units.drop(columns=["cluster_labels_new"], errors="ignore") return units.merge( clusters[["unit_id", "cluster_labels_new"]], on="unit_id", how="left", ) def drop_unneeded_derived_columns(units): """Prune unused derived columns while preserving base/cache columns.""" drop_cols = set() for col in units.columns: if col in DERIVED_COLUMNS_TO_DROP: drop_cols.add(col) if col in RF_COLUMNS_TO_DROP: drop_cols.add(col) if any(col.startswith(prefix) for prefix in DERIVED_PREFIXES_TO_DROP): drop_cols.add(col) if col.startswith("pulse_high_") and col not in OPTO_COLUMNS: drop_cols.add(col) if col.startswith("raised_cosine_high_") and col not in OPTO_COLUMNS: drop_cols.add(col) if any(pattern in col for pattern in RESPONSIVENESS_PATTERNS_TO_DROP): drop_cols.add(col) return units.drop(columns=sorted(drop_cols), errors="ignore") def validate_required_columns(units): """Raise a clear error for columns required by current figure notebooks.""" required = [ "unit_id", "ecephys_session_id", "structure_acronym", "cortical_layer", "waveform_duration", "quality", "isi_violations", "amplitude_cutoff", "presence_ratio", "abnormal_activity", "abnormal_histology", "no_anomalies", "genotype", "image_set", "session_number", "experience_level", "brain_division", "cluster_labels_new", "SST", "VIP", "RS", "FS", ] + OPTO_COLUMNS missing = [col for col in required if col not in units.columns] if missing: raise ValueError( "Output is missing required figure/helper columns: " f"{missing}. Provide a richer base unit table, session table, " "structure tree, opto metrics file, or cluster labels file. " "For cortical_layer, use build_master_unit_table.py first." ) def build_augmented_master_units_table( output_file, cache=None, base_unit_table_file=None, session_table_file=None, opto_metrics_file=None, cluster_labels_file=None, structure_tree_file=None, validate=True, ): """Build and save the figure-scoped master units table.""" units = load_base_units(cache=cache, base_unit_table_file=base_unit_table_file) sessions = load_session_table(cache=cache, session_table_file=session_table_file) units = merge_session_metadata(units, sessions) units = add_no_anomalies(units) units = add_brain_division(units, structure_tree_file=structure_tree_file) units = merge_opto_metrics(units, opto_metrics_file=opto_metrics_file) units = derive_cell_type_columns(units) units = merge_cluster_labels(units, cluster_labels_file=cluster_labels_file) units = drop_unneeded_derived_columns(units) if validate: validate_required_columns(units) output_file = Path(output_file) output_file.parent.mkdir(parents=True, exist_ok=True) units.to_csv(output_file, index=False) print(f"Saved {len(units):,} rows and {len(units.columns):,} columns to {output_file}") return units def parse_args(): parser = argparse.ArgumentParser( description="Build the augmented figure-scoped master units table." ) parser.add_argument("--output-file", "--output_file", dest="output_file", required=True) parser.add_argument("--cache-dir", "--cache_dir", dest="cache_dir", default=None, help="Path to the VBN cache directory. Required unless base/session CSVs are supplied.") parser.add_argument("--base-unit-table-file", "--base_unit_table_file", dest="base_unit_table_file", default=None, help="Optional base unit table CSV, normally build_master_unit_table.py output. All non-bookkeeping base columns are preserved.") parser.add_argument("--session-table-file", "--session_table_file", dest="session_table_file", default=None, help="Optional ecephys session table CSV.") parser.add_argument("--opto-metrics-file", "--opto_metrics_file", dest="opto_metrics_file", default=None, help="Optional unit_opto_metrics.csv path.") parser.add_argument("--cluster-labels-file", "--cluster_labels_file", dest="cluster_labels_file", default=None, help="Optional cluster labels CSV containing unit_id and cluster_labels_new.") parser.add_argument("--structure-tree-file", "--structure_tree_file", dest="structure_tree_file", default=None, help="Optional CCF structure tree CSV used to add brain_division.") parser.add_argument("--manifest", default=DEFAULT_MANIFEST, help=f"Cache manifest to load. Defaults to {DEFAULT_MANIFEST}.") parser.add_argument("--cache-source", choices=("local", "s3"), default="local", help="Use from_local_cache or from_s3_cache. Defaults to local.") parser.add_argument("--skip-validation", action="store_true", help="Write output even if audited figure columns are missing.") return parser.parse_args() def main(): args = parse_args() cache = None if args.cache_dir is not None: cache = load_cache( cache_dir=args.cache_dir, manifest=args.manifest, cache_source=args.cache_source, ) build_augmented_master_units_table( output_file=args.output_file, cache=cache, base_unit_table_file=args.base_unit_table_file, session_table_file=args.session_table_file, opto_metrics_file=args.opto_metrics_file, cluster_labels_file=args.cluster_labels_file, structure_tree_file=args.structure_tree_file, validate=not args.skip_validation, ) if __name__ == "__main__": main()