Electrophysiology Pipeline with Object Storage¶

import datajoint as dj
import numpy as np
import matplotlib.pyplot as plt
from pathlib import Path
import tempfile

# Create a temporary directory for object storage
STORE_PATH = tempfile.mkdtemp(prefix='dj_ephys_')

# Configure object storage with partitioning by mouse_id, session_date, and neuron_id
dj.config.stores['ephys'] = {
    'protocol': 'file',
    'location': STORE_PATH,
    'partition_pattern': '{mouse_id}/{session_date}/{neuron_id}',  # Partition by subject, session, and neuron
}

schema = dj.Schema('tutorial_electrophysiology_npy')

# Data directory (relative to this notebook)
DATA_DIR = Path('./data')

print(f"Store configured at: {STORE_PATH}")
print("Partitioning: {mouse_id}/{session_date}/{neuron_id}")

# Insert mice
Mouse.insert([
    {'mouse_id': 0, 'dob': '2017-03-01', 'sex': 'M'},
    {'mouse_id': 5, 'dob': '2016-12-25', 'sex': 'F'},
    {'mouse_id': 100, 'dob': '2017-05-12', 'sex': 'F'},
], skip_duplicates=True)

# Insert sessions (matching our data files)
Session.insert([
    {'mouse_id': 0, 'session_date': '2017-05-15', 'experimenter': 'Alice'},
    {'mouse_id': 0, 'session_date': '2017-05-19', 'experimenter': 'Alice'},
    {'mouse_id': 5, 'session_date': '2017-01-05', 'experimenter': 'Bob'},
    {'mouse_id': 100, 'session_date': '2017-05-25', 'experimenter': 'Carol'},
    {'mouse_id': 100, 'session_date': '2017-06-01', 'experimenter': 'Carol'},
], skip_duplicates=True)

Session()

Neuron.populate(display_progress=True)

Neuron()

# Fetch returns NpyRef, not the array
key = {'mouse_id': 0, 'session_date': '2017-05-15', 'neuron_id': 0}
ref = (Neuron & key).fetch1('activity')

print(f"Type: {type(ref).__name__}")
print(f"Shape: {ref.shape} (no download!)")
print(f"Dtype: {ref.dtype}")
print(f"Is loaded: {ref.is_loaded}")

# Explicitly load when ready
activity = ref.load()
print(f"Loaded: {activity.shape}, is_loaded: {ref.is_loaded}")

# Memory-mapped loading - efficient for large arrays
key = {'mouse_id': 0, 'session_date': '2017-05-15', 'neuron_id': 0}
ref = (Neuron & key).fetch1('activity')

# Load as memory-mapped array (read-only)
mmap_arr = ref.load(mmap_mode='r')

print(f"Type: {type(mmap_arr).__name__}")
print(f"Shape: {mmap_arr.shape}")

# Random access only reads the needed portion from disk
slice_data = mmap_arr[100:200]
print(f"Slice [100:200]: {slice_data[:5]}...")

fig, axes = plt.subplots(2, 3, figsize=(12, 6))

for ax, key in zip(axes.ravel(), Neuron.keys()):
    # fetch1 returns NpyRef, but plotting works via __array__
    activity = (Neuron & key).fetch1('activity')
    ax.plot(np.asarray(activity))  # Convert to array for plotting
    ax.set_title(f"Mouse {key['mouse_id']}, {key['session_date']}")
    ax.set_xlabel('Time bin')
    ax.set_ylabel('Activity')

# Hide unused subplot
axes[1, 2].axis('off')
plt.tight_layout()

ActivityStats.populate(display_progress=True)
ActivityStats()

@schema
class SpikeParams(dj.Lookup):
    definition = """
    spike_param_id : int16
    ---
    threshold : float32    # spike detection threshold
    """
    
    contents = [
        {'spike_param_id': 1, 'threshold': 0.5},
        {'spike_param_id': 2, 'threshold': 0.9},
    ]

SpikeParams()

Spikes.populate(display_progress=True)

# View spike counts for each neuron x parameter combination
Spikes.proj('spike_count')

# Pick one neuron to visualize
neuron_key = {'mouse_id': 0, 'session_date': '2017-05-15', 'neuron_id': 0}
activity = (Neuron & neuron_key).fetch1('activity').load()  # Explicit load

fig, axes = plt.subplots(2, 1, figsize=(12, 6), sharex=True)

for ax, param_id in zip(axes, [1, 2]):
    key = {**neuron_key, 'spike_param_id': param_id}
    spike_times_ref = (Spikes & key).fetch1('spike_times')
    spike_times = spike_times_ref.load()  # Load spike times
    spike_count = (Spikes & key).fetch1('spike_count')
    threshold = (SpikeParams & {'spike_param_id': param_id}).fetch1('threshold')
    
    ax.plot(activity, 'b-', alpha=0.7, label='Activity')
    ax.axhline(threshold, color='r', linestyle='--', label=f'Threshold={threshold}')
    ax.scatter(spike_times, activity[spike_times], color='red', s=20, zorder=5)
    ax.set_title(f'Threshold={threshold}: {spike_count} spikes detected')
    ax.set_ylabel('Activity')
    ax.legend(loc='upper right')

axes[1].set_xlabel('Time bin')
plt.tight_layout()

# Get waveforms for one neuron with threshold=0.5
key = {'mouse_id': 100, 'session_date': '2017-05-25', 'neuron_id': 0, 'spike_param_id': 1}

# Fetch waveforms NpyRef directly from Spikes table
waveforms_ref = (Spikes & key).fetch1('waveforms')

# Load the waveform matrix
waveform_matrix = waveforms_ref.load()

if len(waveform_matrix) > 0:
    plt.figure(figsize=(8, 4))
    # Plot individual waveforms (light)
    for wf in waveform_matrix:
        plt.plot(wf, 'b-', alpha=0.2)
    # Plot mean waveform (bold)
    plt.plot(waveform_matrix.mean(axis=0), 'r-', linewidth=2, label='Mean waveform')
    plt.axvline(40, color='k', linestyle='--', alpha=0.5, label='Spike time')
    plt.xlabel('Sample (relative to spike)')
    plt.ylabel('Activity')
    plt.title(f'Spike Waveforms (n={len(waveform_matrix)})')
    plt.legend()
else:
    print("No waveforms found for this key")

# Fetch all neurons - returns NpyRefs, NOT arrays
all_neurons = Neuron.to_dicts()
print(f"Fetched {len(all_neurons)} neurons\n")

# Inspect metadata without downloading
for neuron in all_neurons:
    ref = neuron['activity']
    print(f"Mouse {neuron['mouse_id']}, {neuron['session_date']}: "
          f"shape={ref.shape}, loaded={ref.is_loaded}")

dj.Diagram(schema)

# Find neurons with high spike counts (threshold=0.5)
(Spikes & 'spike_param_id = 1' & 'spike_count > 20').proj('spike_count')

# Join with Mouse to see which mice have most spikes
(Mouse * Session * Spikes & 'spike_param_id = 1').proj('spike_count')

# Explore the store directory structure
from pathlib import Path

print(f"Store location: {STORE_PATH}\n")
print("Directory structure (one subject shown in full):")

def print_tree(directory, prefix="", max_depth=7, current_depth=0, limit_items=True):
    """Print directory tree with limited depth."""
    if current_depth >= max_depth:
        return
    
    try:
        entries = sorted(Path(directory).iterdir())
    except PermissionError:
        return
    
    dirs = [e for e in entries if e.is_dir()]
    files = [e for e in entries if e.is_file()]
    
    # At depth 0 (root), only show first subject in detail
    if current_depth == 0 and limit_items:
        dirs = dirs[:1]  # Only show first mouse_id
    
    # Show directories
    for i, d in enumerate(dirs):
        is_last_dir = (i == len(dirs) - 1) and len(files) == 0
        connector = "└── " if is_last_dir else "├── "
        print(f"{prefix}{connector}{d.name}/")
        
        extension = "    " if is_last_dir else "│   "
        print_tree(d, prefix + extension, max_depth, current_depth + 1, limit_items=False)
    
    if current_depth == 0 and limit_items and len(sorted(Path(directory).iterdir(), key=lambda x: x.is_file())) > 1:
        total_mice = len([e for e in Path(directory).iterdir() if e.is_dir()])
        if total_mice > 1:
            print(f"{prefix}└── ... and {total_mice - 1} more subjects (mouse_id=5, mouse_id=100)")
    
    # Show files
    for i, f in enumerate(files):
        is_last = i == len(files) - 1
        connector = "└── " if is_last else "├── "
        size_kb = f.stat().st_size / 1024
        print(f"{prefix}{connector}{f.name} ({size_kb:.1f} KB)")

print_tree(STORE_PATH)

# Get the actual path from an NpyRef
key = {'mouse_id': 0, 'session_date': '2017-05-15', 'neuron_id': 0}
ref = (Neuron & key).fetch1('activity')

print(f"NpyRef path (relative): {ref.path}")
print(f"Full path: {Path(STORE_PATH) / ref.path}\n")

# Load directly with NumPy - bypass the database!
direct_path = Path(STORE_PATH) / ref.path
activity_direct = np.load(direct_path)
print(f"Loaded array: shape={activity_direct.shape}, dtype={activity_direct.dtype}")
print(f"First 5 values: {activity_direct[:5]}")

# Find all .npy files for a specific mouse using filesystem tools
# This works without any database query!

store_path = Path(STORE_PATH)
all_npy_files = list(store_path.rglob("*.npy"))

print(f"All .npy files in store ({len(all_npy_files)} total):\n")
for f in sorted(all_npy_files)[:8]:
    rel_path = f.relative_to(STORE_PATH)
    size_kb = f.stat().st_size / 1024
    print(f"  {rel_path} ({size_kb:.1f} KB)")
if len(all_npy_files) > 8:
    print(f"  ... and {len(all_npy_files) - 8} more files")