""" Build the base master_unit_table.csv used by downstream unit-table builders. The AllenSDK VBN cache unit table does not include cortical layer labels. This builder preserves the cache unit columns, merges probe/no-anomaly metadata, and adds the layer columns by indexing each unit's CCF coordinate into the CCF annotation volume from AllenSDK ReferenceSpaceCache. The CCF lookup converts AP/DV/LR microns to annotation-volume indices, gets the structure acronym at that voxel, and extracts the cortical-layer suffix from the acronym. """ import argparse import re from pathlib import Path import numpy as np import pandas as pd DEFAULT_MANIFEST = "visual-behavior-neuropixels_project_manifest_v0.5.0.json" DEFAULT_CCF_REFERENCE_SPACE_KEY = "annotation/ccf_2017" COORD_COLUMNS = [ "anterior_posterior_ccf_coordinate", "dorsal_ventral_ccf_coordinate", "left_right_ccf_coordinate", ] 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): 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_units(cache=None, unit_table_file=None): if unit_table_file is not None: return _read_csv_without_unnamed(unit_table_file) if cache is None: raise ValueError("Either cache or unit_table_file is required.") return _reset_index_column(cache.get_unit_table(), "unit_id") def load_probe_table(cache=None, probe_table_file=None): if probe_table_file is not None: return _read_csv_without_unnamed(probe_table_file) if cache is None: return None if not hasattr(cache, "get_probe_table"): return None return _reset_index_column(cache.get_probe_table(), "ecephys_probe_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_probe_metadata(units, probes): """Merge probe metadata columns without overwriting unit-table columns.""" if probes is None or "ecephys_probe_id" not in units.columns: return units probe_cols = [ col for col in probes.columns if col == "ecephys_probe_id" or col not in units.columns ] return units.merge(probes[probe_cols], on="ecephys_probe_id", how="left") def add_no_anomalies(units, sessions): """Add no_anomalies from session abnormality flags if needed.""" if "no_anomalies" in units.columns: return units required = {"ecephys_session_id", "abnormal_activity", "abnormal_histology"} if sessions is None or not required.issubset(set(sessions.columns) | set(units.columns)): return units if {"abnormal_activity", "abnormal_histology"}.issubset(units.columns): abnormal_activity = units["abnormal_activity"] abnormal_histology = units["abnormal_histology"] else: session_flags = sessions[ ["ecephys_session_id", "abnormal_activity", "abnormal_histology"] ] merged = units[["ecephys_session_id"]].merge( session_flags, on="ecephys_session_id", how="left", ) abnormal_activity = merged["abnormal_activity"] abnormal_histology = merged["abnormal_histology"] units["no_anomalies"] = abnormal_activity.isna() & abnormal_histology.isna() return units def load_nrrd_volume(path): import nrrd volume, _ = nrrd.read(path) return volume def load_reference_space_cache( ccf_manifest_path="manifest.json", ccf_resolution=10, ccf_reference_space_key=DEFAULT_CCF_REFERENCE_SPACE_KEY, ): from allensdk.core.reference_space_cache import ReferenceSpaceCache return ReferenceSpaceCache( ccf_resolution, ccf_reference_space_key, manifest=ccf_manifest_path, ) def _structure_map_from_tree_dataframe(structure_tree): required = {"id", "acronym"} missing = required - set(structure_tree.columns) if missing: raise ValueError(f"Structure tree is missing columns: {sorted(missing)}") return dict(zip(structure_tree["id"].astype(int), structure_tree["acronym"])) def _structure_map_from_reference_tree(structure_tree, structure_ids): ids = sorted(int(structure_id) for structure_id in set(structure_ids) if structure_id > 0) if not ids: return {} structures = structure_tree.get_structures_by_id(ids) return {structure["id"]: structure["acronym"] for structure in structures} def _structure_map_from_source(structure_source, structure_ids): if isinstance(structure_source, dict): return structure_source return _structure_map_from_reference_tree(structure_source, structure_ids) def _voxel_indices(units, volume_shape, resolution): missing = [col for col in COORD_COLUMNS if col not in units.columns] if missing: raise ValueError(f"Unit table is missing CCF coordinate columns: {missing}") coords = units[COORD_COLUMNS].to_numpy(dtype=float) valid = np.isfinite(coords).all(axis=1) & (coords >= 0).all(axis=1) indices = np.zeros(coords.shape, dtype=int) indices[valid] = (coords[valid] / resolution).astype(int) for axis, axis_size in enumerate(volume_shape): valid &= indices[:, axis] < axis_size return indices, valid def get_structure_ids_for_units(units, annotation_volume, resolution): """Return CCF structure ids for each unit coordinate.""" indices, valid = _voxel_indices(units, annotation_volume.shape, resolution) structure_ids = np.zeros(len(units), dtype=int) structure_ids[valid] = annotation_volume[ indices[valid, 0], indices[valid, 1], indices[valid, 2], ] return structure_ids def get_values_for_units(units, volume, resolution, default_value=np.nan): indices, valid = _voxel_indices(units, volume.shape, resolution) values = np.full(len(units), default_value, dtype=float) values[valid] = volume[ indices[valid, 0], indices[valid, 1], indices[valid, 2], ] return values def get_layer_name(acronym): """Extract the cortical-layer suffix from a CCF structure acronym.""" if pd.isna(acronym) or acronym in ("CA1", "CA2", "CA3", "outside_brain"): return np.nan match = re.search(r"\d", str(acronym)) if match is None: return np.nan return str(acronym)[match.start():] def assign_cortical_layers( units, annotation_volume, structure_source, resolution=10, ): """Add cortical_layer and structure_with_layer from unit CCF coordinates.""" units = units.copy() structure_ids = get_structure_ids_for_units(units, annotation_volume, resolution) structure_id_to_acronym = _structure_map_from_source(structure_source, structure_ids) layer_acronyms = [ structure_id_to_acronym.get(int(structure_id), np.nan) for structure_id in structure_ids ] units["cortical_layer"] = [get_layer_name(acronym) for acronym in layer_acronyms] units["structure_with_layer"] = [ f"{area}{layer}" if pd.notna(layer) and str(layer) != "" else area for area, layer in zip(units["structure_acronym"], units["cortical_layer"]) ] return units def add_cortical_depth( units, streamlines_volume=None, resolution=10, ): if streamlines_volume is None: return units units = units.copy() units["cortical_depth"] = get_values_for_units( units, streamlines_volume, resolution=resolution, default_value=0, ) return units def build_ccf_resources( annotation_volume_file=None, structure_tree_file=None, streamlines_file=None, ccf_manifest_path="manifest.json", ccf_resolution=10, ccf_reference_space_key=DEFAULT_CCF_REFERENCE_SPACE_KEY, include_cortical_depth=False, ): """Load CCF annotation/tree resources from AllenSDK or local files.""" if annotation_volume_file is not None: annotation_volume = load_nrrd_volume(annotation_volume_file) if structure_tree_file is None: raise ValueError("structure_tree_file is required with annotation_volume_file.") structure_tree = pd.read_csv(structure_tree_file) structure_source = _structure_map_from_tree_dataframe(structure_tree) streamlines_volume = ( load_nrrd_volume(streamlines_file) if include_cortical_depth and streamlines_file is not None else None ) return annotation_volume, structure_source, streamlines_volume rspc = load_reference_space_cache( ccf_manifest_path=ccf_manifest_path, ccf_resolution=ccf_resolution, ccf_reference_space_key=ccf_reference_space_key, ) annotation_volume, _ = rspc.get_annotation_volume() structure_tree = rspc.get_structure_tree(structure_graph_id=1) streamlines_volume = ( load_nrrd_volume(streamlines_file) if include_cortical_depth and streamlines_file is not None else None ) return annotation_volume, structure_tree, streamlines_volume def validate_master_unit_table(units): required = [ "unit_id", "ecephys_channel_id", "ecephys_probe_id", "ecephys_session_id", "structure_acronym", "cortical_layer", "structure_with_layer", "waveform_duration", "quality", "isi_violations", "amplitude_cutoff", "presence_ratio", "no_anomalies", ] missing = [col for col in required if col not in units.columns] if missing: raise ValueError( "Output is missing required base unit-table columns: " f"{missing}. Provide unit, probe, session, and CCF inputs." ) def build_master_unit_table( output_file, cache=None, unit_table_file=None, probe_table_file=None, session_table_file=None, annotation_volume_file=None, structure_tree_file=None, streamlines_file=None, ccf_manifest_path="manifest.json", ccf_resolution=10, ccf_reference_space_key=DEFAULT_CCF_REFERENCE_SPACE_KEY, include_cortical_depth=False, validate=True, ): units = load_units(cache=cache, unit_table_file=unit_table_file) probes = load_probe_table(cache=cache, probe_table_file=probe_table_file) sessions = load_session_table(cache=cache, session_table_file=session_table_file) annotation_volume, structure_source, streamlines_volume = build_ccf_resources( annotation_volume_file=annotation_volume_file, structure_tree_file=structure_tree_file, streamlines_file=streamlines_file, ccf_manifest_path=ccf_manifest_path, ccf_resolution=ccf_resolution, ccf_reference_space_key=ccf_reference_space_key, include_cortical_depth=include_cortical_depth, ) units = assign_cortical_layers( units, annotation_volume=annotation_volume, structure_source=structure_source, resolution=ccf_resolution, ) units = add_cortical_depth( units, streamlines_volume=streamlines_volume, resolution=ccf_resolution, ) units = add_no_anomalies(units, sessions) units = merge_probe_metadata(units, probes) if validate: validate_master_unit_table(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 master_unit_table.csv with CCF-derived cortical layers." ) 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 CSV inputs are supplied.") parser.add_argument("--unit-table-file", "--unit_table_file", dest="unit_table_file", default=None, help="Optional cache units CSV. Preserves all unit columns.") parser.add_argument("--probe-table-file", "--probe_table_file", dest="probe_table_file", default=None, help="Optional cache probes CSV for probe metadata.") parser.add_argument("--session-table-file", "--session_table_file", dest="session_table_file", default=None, help="Optional ecephys sessions CSV for no_anomalies.") parser.add_argument("--annotation-volume-file", "--annotation_volume_file", dest="annotation_volume_file", default=None, help="Optional local CCF annotation NRRD. If omitted, AllenSDK ReferenceSpaceCache is used.") parser.add_argument("--structure-tree-file", "--structure_tree_file", dest="structure_tree_file", default=None, help="CCF structure tree CSV required with --annotation-volume-file.") parser.add_argument("--streamlines-file", "--streamlines_file", dest="streamlines_file", default=None, help="Optional laplacian_10.nrrd streamlines file used with --include-cortical-depth.") parser.add_argument("--ccf-manifest-path", "--ccf_manifest_path", dest="ccf_manifest_path", default="manifest.json", help="AllenSDK ReferenceSpaceCache manifest path. Defaults to manifest.json.") parser.add_argument("--ccf-reference-space-key", "--ccf_reference_space_key", dest="ccf_reference_space_key", default=DEFAULT_CCF_REFERENCE_SPACE_KEY, help=f"ReferenceSpaceCache key. Defaults to {DEFAULT_CCF_REFERENCE_SPACE_KEY}.") parser.add_argument("--ccf-resolution", "--ccf_resolution", dest="ccf_resolution", type=int, default=10, help="CCF annotation resolution in microns. Defaults to 10.") 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("--include-cortical-depth", action="store_true", help="Add cortical_depth from --streamlines-file.") parser.add_argument("--skip-validation", action="store_true", help="Write output even if required base 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_master_unit_table( output_file=args.output_file, cache=cache, unit_table_file=args.unit_table_file, probe_table_file=args.probe_table_file, session_table_file=args.session_table_file, annotation_volume_file=args.annotation_volume_file, structure_tree_file=args.structure_tree_file, streamlines_file=args.streamlines_file, ccf_manifest_path=args.ccf_manifest_path, ccf_resolution=args.ccf_resolution, ccf_reference_space_key=args.ccf_reference_space_key, include_cortical_depth=args.include_cortical_depth, validate=not args.skip_validation, ) if __name__ == "__main__": main()