labdata 0.1.17__tar.gz → 0.1.18__tar.gz

{labdata-0.1.17 → labdata-0.1.18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: labdata
-Version: 0.1.17
+Version: 0.1.18
 Summary: Package to manage data in experimental neuroscience labs
 Author-email: Joao Couto <jpcouto@gmail.com>
 Project-URL: Homepage, https://github.com/jcouto/labdata

{labdata-0.1.17 → labdata-0.1.18}/labdata/__init__.py

@@ -1,4 +1,4 @@
-VERSION = "0.1.17"
+VERSION = "0.1.18"
 
 from .utils import *
 from .copy import copy_to_upload_server

{labdata-0.1.17 → labdata-0.1.18}/labdata/cli.py

@@ -625,16 +625,21 @@ Available analysis are
             return
         from .compute import parse_analysis, run_analysis
         # parse analysis will check if the analysis is defined
-        jobids,obj = parse_analysis(analysis = args.analysis,
-                                    job_id = args.job, 
-                                    subject = args.subject,
-                                    session = args.session,
-                                    datatype = args.datatype,
-                                    secondary_args = secondary_args,
-                                    multisession = args.multisession,
-                                    force_submit = args.force_submit,
-                                    full_command = ' '.join(sys.argv[1:]),
-                                    project = args.project)
+        try:
+            jobids,obj = parse_analysis(analysis = args.analysis,
+                                        job_id = args.job, 
+                                        subject = args.subject,
+                                        session = args.session,
+                                        datatype = args.datatype,
+                                        secondary_args = secondary_args,
+                                        multisession = args.multisession,
+                                        force_submit = args.force_submit,
+                                        full_command = ' '.join(sys.argv[1:]),
+                                        project = args.project)
+        except (ValueError, NotImplementedError) as err:
+            # what the request asked for is not there: say what, and stop
+            print(f'{tcolor["r"]("Nothing to run:")} {err}')
+            sys.exit(1)
         if not len(jobids):
             print('Nothing to run.')
             return

{labdata-0.1.17 → labdata-0.1.18}/labdata/compute/pose.py

@@ -159,7 +159,19 @@ class DeeplabcutCompute(BaseCompute):
             raise(NotImplementedError('Specifying no session is not yet implemented'))
         datasets = []
         for k in keys:
-            datasets += (self.schema.DatasetVideo()& k).fetch(as_dict = True)
+            videos = (self.schema.DatasetVideo() & k).fetch(as_dict = True)
+            if not len(videos):
+                raise ValueError(f'Session {k["session_name"]} of {k["subject_name"]} has no video '
+                                 '(nothing in DatasetVideo); videos are registered when the '
+                                 'session is uploaded.')
+            video_name = self.parameters.get('video_name')
+            if video_name is not None:
+                names = sorted({v['video_name'] for v in videos})
+                videos = [v for v in videos if v['video_name'] == video_name]
+                if not len(videos):
+                    raise ValueError(f'Session {k["session_name"]} of {k["subject_name"]} has no video '
+                                     f'named {video_name}; it has {", ".join(names)}.')
+            datasets += videos
         datasets = [dict(subject_name = d['subject_name'],
                          session_name = d['session_name'],
                          dataset_name = d['dataset_name']) for d in datasets]

{labdata-0.1.17 → labdata-0.1.18}/labdata/compute/utils.py

@@ -1,13 +1,36 @@
 from ..utils import *
 import traceback
 
-def load_analysis_object(analysis):
+def available_analyses(project=None):
+    '''The compute classes of the preference file that are loaded, by analysis name.
+
+    An entry whose class is not there, because it comes from a plugin that is
+    not loaded, or is not a compute class at all, is left out.  With
+    ``project`` the project's plugins are loaded first.
+    '''
+    import labdata
+    if project is not None:
+        load_project_schema(project)
+    analyses = {}
+    for name, path in prefs['compute']['analysis'].items():
+        try:
+            obj = eval(path)
+        except (AttributeError, NameError):
+            continue
+        if isinstance(obj, type) and issubclass(obj, BaseCompute):
+            analyses[name] = obj
+    return analyses
+
+
+def load_analysis_object(analysis, project=None):
     '''Load a compute class by its analysis name from user preferences.
 
     Parameters
     ----------
     analysis : str
         Key in ``prefs['compute']['analysis']`` (e.g. ``'spks'``, ``'suite2p'``).
+    project : str or None, optional
+        Project whose plugins to load before looking for the class.
 
     Returns
     -------
@@ -17,23 +40,22 @@ def load_analysis_object(analysis):
     Raises
     ------
     SystemExit
-        If the analysis name is not found in preferences, prints a help message
-        listing available analyses and exits.
+        If the analysis is not among the loaded ones, prints a help message
+        listing the analyses that are and exits.
     '''
-    if not analysis in prefs['compute']['analysis'].keys():
+    analyses = available_analyses(project)
+    if analysis not in analyses:
+        path = prefs['compute']['analysis'].get(analysis)
+        where = f' for project {project}' if project else ''
+        because = (f' ({path} is not loaded{where}, or is not a compute class)' if path else '')
         print(f'''
-Could not find {tcolor['r'](analysis)} analysis.
-
-    The analysis are {tcolor['g'](','.join(prefs["compute"]["analysis"].keys()))}
-
+Could not find {tcolor['r'](analysis)} analysis{because}.
+    The analyses are {tcolor['g'](','.join(analyses.keys()))}
 Add the analysis to the "{tcolor['y']("compute")}" section of the preference file {tcolor['y']("{analysis_name:analysis_object}")}.
-
 ''')
-        sys.exit()        
-        raise ValueError('Add the analysis to the "compute" section of the preference file {analysis_name:analysis_object}.\n\n')
-        
-    import labdata
-    return eval(prefs['compute']['analysis'][analysis])
+        sys.exit(1)
+    return analyses[analysis]
+
 
 def handle_compute(job_id, project=None):
     '''Instantiate and return the correct compute object for a queued job.
@@ -61,7 +83,7 @@ def handle_compute(job_id, project=None):
     jobinfo = jobinfo.iloc[0]
     if jobinfo.task_waiting == 0:
         print(f'Task {job_id} is running on {jobinfo.task_host}')
-    obj = load_analysis_object(jobinfo.task_name)(jobinfo.job_id,project)
+    obj = load_analysis_object(jobinfo.task_name, project=project)(jobinfo.job_id,project)
     return obj
 
 def parse_analysis(analysis, job_id=None,
@@ -115,7 +137,7 @@ def parse_analysis(analysis, job_id=None,
     obj : BaseCompute
         Instantiated compute object.
     '''
-    obj = load_analysis_object(analysis)(job_id, project=project)
+    obj = load_analysis_object(analysis, project=project)(job_id, project=project)
     obj.secondary_parse(secondary_args, parameter_number)
     if obj.job_id is None:
         # then we have to create jobs and assign
@@ -123,6 +145,9 @@ def parse_analysis(analysis, job_id=None,
             schema = load_project_schema(project)
         else:
             schema = obj.schema
+        # a subject or session that is not in the project is the usual reason
+        # for finding nothing; say so, with the closest names, before looking
+        check_subject_and_session(schema, subject, session, project = project)
          # first check if there is a task that has already been submitted with the exact same command.
          # this has a caveat: if the order of the arguments is switched, it wont work..
         if not full_command is None:
@@ -135,6 +160,12 @@ def parse_analysis(analysis, job_id=None,
                 return [], None
         if not subject is None or not session is None:
             datasets = obj.find_datasets(subject_name = subject, session_name = session)
+            if datasets is not None and not len(datasets):
+                asked = ' '.join(f'{name} {", ".join(_as_names(v))}' for name, v in
+                                 (('subject', subject), ('session', session)) if _as_names(v))
+                print(tcolor['r'](f'No dataset for {analysis} matches {asked}.')
+                      + _datasets_on_record(schema, subject, session))
+                return [], obj
         else:
             datasets = None
         job_ids = obj.place_tasks_in_queue(datasets,
@@ -144,6 +175,60 @@ def parse_analysis(analysis, job_id=None,
         # now we have the job ids, need to figure out how to launch the jobs
         return job_ids, obj # returns the job ids and the task
 
+def _as_names(value):
+    '''The names in a command line argument: a string, a list of strings, or None.'''
+    if value is None:
+        return []
+    if isinstance(value, str):
+        value = [value]
+    return [str(v) for v in value if str(v).strip()]
+
+
+def _datasets_on_record(schema, subject, session):
+    '''One line saying what datasets the requested sessions do have, or nothing.'''
+    keys = [dict(subject_name = a, session_name = s) for a in _as_names(subject) for s in _as_names(session)]
+    if not keys:
+        return ''
+    rows = (schema.Dataset() & keys).fetch('session_name', 'dataset_name', 'dataset_type', as_dict = True)
+    if not rows:
+        return ' The session has no dataset on record.'
+    have = sorted({f'{r["dataset_name"]} ({r["dataset_type"] or "no type"})' for r in rows})
+    return f' On record: {", ".join(have)}.'
+
+
+def check_subject_and_session(schema, subject=None, session=None, project=None):
+    '''Raise a ValueError naming what is not in the project, before a compute looks for datasets.
+
+    A subject that is not in ``Subject`` is reported with the closest names
+    on record; a session that the subject does not have is reported with the
+    subject's latest sessions.  Nothing happens when neither is given.
+    '''
+    import difflib
+    where = f'project {project}' if project else 'the project'
+    subjects = _as_names(subject)
+    sessions = _as_names(session)
+    if not subjects and not sessions:
+        return
+    known = list(schema.Subject().fetch('subject_name'))
+    for name in subjects:
+        if name not in known:
+            close = difflib.get_close_matches(name, known, n = 5, cutoff = 0.6)
+            hint = f' Close names: {", ".join(close)}.' if close else ''
+            raise ValueError(f'Subject {name} is not in {where}.{hint}')
+    if sessions and not subjects:
+        raise ValueError('A session needs its subject: pass the subject name as well.')
+    for name in subjects:
+        have = list((schema.Session() & dict(subject_name = name)).fetch(
+            'session_name', order_by = 'session_datetime DESC'))
+        for ses in sessions:
+            if ses not in have:
+                close = difflib.get_close_matches(ses, have, n = 3, cutoff = 0.6)
+                hint = (f' Close names: {", ".join(close)}.' if close else
+                        f' Its latest sessions: {", ".join(have[:5])}.' if have else
+                        ' It has no sessions.')
+                raise ValueError(f'Subject {name} has no session {ses} in {where}.{hint}')
+
+
 def check_archived(jobid, check_local=False, schema=None):
     '''Check whether files assigned to a compute job are in S3 Glacier.
 
@@ -591,7 +676,9 @@ class BaseCompute():
                     idx += list(filter(lambda x: not x is None,[i if f in s else None for i,s in enumerate(
                         files.file_path.values)]))
                 if len(idx) == 0:
-                    raise ValueError(f'Could not find valid Dataset.DataFiles for {dataset}')
+                    found = ', '.join(Path(f).name for f in files.file_path.values[:8]) if len(files) else 'no files at all'
+                    raise ValueError(f'{self.name} needs a file of type {", ".join(self.file_filters)} '
+                                     f'and dataset {dataset} has {found}.')
                 files = files.iloc[idx]
                 if type(dataset) is dict:
                     key = dict(dataset,task_name = self.name)

{labdata-0.1.17 → labdata-0.1.18}/labdata/dashboard/admin.py

@@ -3,6 +3,7 @@ from datetime import datetime, timedelta
 
 import pandas as pd
 import streamlit as st
+from writes import log
 import plotly.express as px
 
 
@@ -285,8 +286,10 @@ def _add_user_dialog(project, user):
             add_user_to_project(project, user)
             get_project_users.clear()
             st.session_state['admin_msg'] = ('success', f'Added {user} to {project}.')
+            log(f'added user {user} to project {project}')
         except Exception as e:
             st.session_state['admin_msg'] = ('error', str(e))
+            log(f'could not add user {user} to project {project}: {e}')
         st.rerun()
     if col2.button("Cancel", key="add_cancel"):
         st.rerun()
@@ -304,8 +307,10 @@ def _remove_user_dialog(project, user):
             remove_user_from_project(project, user)
             get_project_users.clear()
             st.session_state['admin_msg'] = ('success', f'Removed {user} from {project}.')
+            log(f'removed user {user} from project {project}')
         except Exception as e:
             st.session_state['admin_msg'] = ('error', str(e))
+            log(f'could not remove user {user} from project {project}: {e}')
         st.rerun()
     if col2.button("Cancel", key="remove_cancel"):
         st.rerun()

{labdata-0.1.17 → labdata-0.1.18}/labdata/dashboard/compute.py

@@ -1,7 +1,7 @@
 import pandas as pd
 import numpy as np
 import streamlit as st
-from writes import write, show_flash
+from writes import write, show_flash, log
 
 STATES = ['COMPLETED', 'WORKING', 'WAITING', 'FAILED', 'CANCELED','ON SERVER']
 STATE_COLORS = ['#01D186', '#48CAE4', '#F3A21C', '#FF5733', '#F3A21C','#F3A21C']
@@ -268,9 +268,13 @@ def _jobs_panel(table, status_field, target_field, cols, delete_key,
                 st.session_state.pop(f'{delete_key}_resubmit_pending', None)
                 if errors:
                     st.session_state[f'{delete_key}_resubmit_errors'] = errors
+                    log(f'could not re-submit {len(errors)} of {len(jobs_to_resubmit)} job(s): '
+                        + '; '.join(str(e) for e in errors))
                 else:
                     st.session_state[f'{delete_key}_msg'] = (
                         f'Re-submitted {len(jobs_to_resubmit)} job(s).')
+                    log(f're-submitted job(s) {", ".join(str(j) for j in jobs_to_resubmit)}'
+                        + (f' to {override}' if override else ''))
                 st.rerun()
             if col_no.button('Cancel', key=f'{delete_key}_resubmit_no'):
                 st.session_state.pop(f'{delete_key}_resubmit_pending', None)

labdata-0.1.18/labdata/dashboard/drift.py

@@ -0,0 +1,331 @@
+'''The drift map and the waveforms over depth of one probe, a view of the Spike Sorting tab.
+
+The spikes come straight from ``SpikeSorting.Unit`` and the waveforms from
+``SpikeSorting.Waveforms``; what was read stays cached, so a session already
+seen comes back at once.  The criteria filter keeps the read and the plots
+to the single units, and a click on either plot selects the unit for the
+whole tab.
+'''
+import numpy as np
+import pandas as pd
+import streamlit as st
+
+SAMPLE_MAX = 500_000            # spikes kept per probe read
+DRAW_OPTIONS = [20_000, 50_000, 100_000, 200_000, 500_000]
+MAX_CHANNELS = 12               # electrodes drawn per unit waveform
+NEIGHBOURHOOD = 100.0           # µm around the peak electrode, without active electrodes on record
+UNIT_COLOURS = ['#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd',
+                '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf']
+LABEL_ORDER = [None, 'noise', 'MUA', 'SUA']      # unlabelled, noise, multi unit, single unit
+LABEL_COLOURS = {None: '#c8c8c8', 'noise': '#222222', 'MUA': '#ff7f0e', 'SUA': '#2ca02c'}
+
+
+def _discrete_scale(colours):
+    '''A plotly colorscale that maps the integers 0..n-1 to ``colours``.'''
+    n = len(colours)
+    scale = []
+    for i, c in enumerate(colours):
+        scale += [[i / n, c], [(i + 1) / n, c]]
+    return scale
+
+
+@st.cache_data(show_spinner=False)
+def _load_probe(_schema, subject_name, session_name, dataset_name, parameter_set_num, probe_num, unit_ids):
+    '''The spikes and waveforms of the units of one probe.
+
+    ``unit_ids`` restricts the units, or is None for all of them.  The spikes
+    are shuffled once, with a fixed seed, and at most ``SAMPLE_MAX`` kept, so
+    the first n of them are a uniform sample of any size.
+
+    Returns a dict: ``times`` (seconds), ``depths`` (µm), ``amplitudes`` and
+    ``spike_units`` for the sample; ``n_spikes`` in the read; the electrode
+    ``coords``; a ``units`` frame (unit_id, depth, shank, n_spikes); and
+    ``waveforms``, unit_id -> (electrode indices, median waveform on them).
+    '''
+    key = dict(subject_name=subject_name, session_name=session_name, dataset_name=dataset_name,
+               parameter_set_num=int(parameter_set_num), probe_num=int(probe_num))
+    empty = dict(times=np.zeros(0, np.float32), depths=np.zeros(0, np.float32),
+                 amplitudes=np.zeros(0, np.float32), spike_units=np.zeros(0, np.int32),
+                 n_spikes=0, coords=np.zeros((0, 2), np.float32),
+                 units=pd.DataFrame(columns=['unit_id', 'depth', 'shank', 'n_spikes']), waveforms={})
+    if unit_ids is not None and not len(unit_ids):
+        return empty
+    restriction = ({} if unit_ids is None
+                   else f'unit_id IN ({",".join(str(int(u)) for u in unit_ids)})')
+    units = (_schema.SpikeSorting.Unit & key & restriction).get_spike_times(
+        extra_keys=['spike_amplitudes', 'spike_positions'], warn=False)
+
+    times, depths, amplitudes, spike_units, counts = [], [], [], [], {}
+    for u in units:
+        t = np.asarray(u['spike_times'], dtype=np.float32).ravel()
+        n = len(t)
+        counts[int(u['unit_id'])] = n
+        if not n:
+            continue
+        pos = None if u['spike_positions'] is None else np.asarray(u['spike_positions'], dtype=np.float32)
+        d = (pos[:, 1] if pos is not None and pos.ndim == 2 and pos.shape[1] > 1 and len(pos) == n
+             else np.full(n, np.nan, dtype=np.float32))
+        a = None if u['spike_amplitudes'] is None else np.asarray(u['spike_amplitudes'], dtype=np.float32).ravel()
+        a = a if a is not None and len(a) == n else np.full(n, np.nan, dtype=np.float32)
+        times.append(t)
+        depths.append(d)
+        amplitudes.append(a)
+        spike_units.append(np.full(n, int(u['unit_id']), dtype=np.int32))
+    n_spikes = int(sum(len(t) for t in times))
+    if n_spikes:
+        order = np.random.default_rng(0).permutation(n_spikes)[:SAMPLE_MAX]
+        times, depths, amplitudes, spike_units = (np.concatenate(x)[order]
+                                                  for x in (times, depths, amplitudes, spike_units))
+    else:
+        times, depths, amplitudes, spike_units = (empty[k] for k in ('times', 'depths', 'amplitudes', 'spike_units'))
+
+    coords = np.asarray((_schema.SpikeSorting & key).fetch1('sorting_channel_coords'), dtype=np.float32)
+    metrics = {int(m['unit_id']): m for m in (_schema.UnitMetrics & key & restriction).fetch(
+        'unit_id', 'depth', 'shank', 'active_electrodes', as_dict=True)}
+    waveforms = {}
+    peaks = {}
+    for w in (_schema.SpikeSorting.Waveforms & key & restriction).fetch('unit_id', 'waveform_median', as_dict=True):
+        if w['waveform_median'] is None:
+            continue
+        wf = np.asarray(w['waveform_median'], dtype=np.float32)
+        if wf.ndim != 2 or not wf.size:
+            continue
+        uid = int(w['unit_id'])
+        n_channels = min(wf.shape[1], len(coords))
+        peak = int(np.argmax(np.abs(wf[:, :n_channels]).max(axis=0)))
+        peaks[uid] = peak
+        channels = np.zeros(0, dtype=int)
+        active = metrics.get(uid, {}).get('active_electrodes')
+        if active is not None:
+            channels = np.unique(np.asarray(active).ravel().astype(int))
+            channels = channels[(channels >= 0) & (channels < n_channels)]
+        if not len(channels):
+            dist = np.linalg.norm(coords[:n_channels] - coords[peak], axis=1)
+            channels = np.where(dist <= NEIGHBOURHOOD)[0]
+        dist = np.linalg.norm(coords[channels] - coords[peak], axis=1)
+        channels = channels[np.argsort(dist)][:MAX_CHANNELS]
+        waveforms[uid] = (channels, wf[:, channels])
+
+    rows = []
+    for uid, n in counts.items():
+        m = metrics.get(uid, {})
+        depth = m.get('depth')
+        if depth is None:
+            own = depths[spike_units == uid]
+            depth = (float(np.nanmedian(own)) if len(own) and np.isfinite(own).any()
+                     else float(coords[peaks[uid], 1]) if uid in peaks else np.nan)
+        shank = m.get('shank')
+        rows.append(dict(unit_id=uid, depth=float(depth), shank=-1 if shank is None else int(shank), n_spikes=n))
+    units_df = (pd.DataFrame(rows).sort_values(['depth', 'unit_id']).reset_index(drop=True)
+                if rows else empty['units'])
+    return dict(times=times, depths=depths, amplitudes=amplitudes, spike_units=spike_units,
+                n_spikes=n_spikes, coords=coords, units=units_df, waveforms=waveforms)
+
+
+def _only_units(data, unit_ids):
+    '''The part of a probe read that belongs to ``unit_ids``.'''
+    keep = np.isin(data['spike_units'], list(unit_ids))
+    return dict(data,
+                times=data['times'][keep], depths=data['depths'][keep],
+                amplitudes=data['amplitudes'][keep], spike_units=data['spike_units'][keep],
+                units=data['units'][data['units']['unit_id'].isin(unit_ids)].reset_index(drop=True),
+                waveforms={u: w for u, w in data['waveforms'].items() if u in unit_ids})
+
+
+def _drift_figure(data, n_draw, colour_by, label_of, highlight, colours, revision):
+    '''The drift map: spike depth over time, coloured by amplitude, unit or label.'''
+    import plotly.graph_objects as go
+    n = min(int(n_draw), len(data['times']))
+    t, d, a, u = (data[k][:n] for k in ('times', 'depths', 'amplitudes', 'spike_units'))
+    hover = ('unit %{customdata[0]}<br>%{x:.2f} s · %{y:.0f} µm<br>'
+             'amplitude %{customdata[1]:.0f}<extra></extra>')
+    if colour_by == 'unit':
+        rank = {uid: i for i, uid in enumerate(data['units']['unit_id'])}
+        colour = np.array([rank.get(int(x), 0) % len(UNIT_COLOURS) for x in u], dtype=np.int16)
+        marker = dict(color=colour, colorscale=_discrete_scale(UNIT_COLOURS), cmin=0, cmax=len(UNIT_COLOURS))
+    elif colour_by == 'label':
+        code = {uid: LABEL_ORDER.index(label if label in LABEL_ORDER else None)
+                for uid, label in label_of.items()}
+        colour = np.array([code.get(int(x), 0) for x in u], dtype=np.int16)
+        marker = dict(color=colour, colorscale=_discrete_scale([LABEL_COLOURS[k] for k in LABEL_ORDER]),
+                      cmin=0, cmax=len(LABEL_ORDER))
+    else:
+        finite = a[np.isfinite(a)]
+        top = float(np.percentile(finite, 95)) if len(finite) else 1.0
+        marker = dict(color=a, colorscale='Spectral', reversescale=True, cmin=0, cmax=top)
+    fig = go.Figure(go.Scattergl(
+        x=t, y=d, mode='markers', name='spikes',
+        marker=dict(size=3, opacity=0.35, showscale=False, **marker),
+        customdata=np.column_stack([u, a]).tolist(), hovertemplate=hover))
+    for uid in highlight:
+        sel = u == uid
+        fig.add_trace(go.Scattergl(
+            x=t[sel], y=d[sel], mode='markers', name=f'unit {uid}',
+            marker=dict(size=5, opacity=0.9, color=colours.get(uid, '#d62728')),
+            customdata=np.column_stack([u[sel], a[sel]]).tolist(), hovertemplate=hover))
+    fig.update_layout(height=430, margin=dict(l=60, r=10, t=10, b=40), template='simple_white',
+                      xaxis_title='Time (s)', yaxis_title='Depth (µm)', uirevision=revision,
+                      showlegend=bool(highlight), legend=dict(orientation='h', y=1.02, x=0))
+    return fig
+
+
+def _waveform_figure(data, highlight, colours, scale, spacing, revision):
+    '''The waveforms over depth: every unit's median waveform at its electrodes, in depth order.'''
+    import plotly.graph_objects as go
+    units, coords, waveforms = data['units'], data['coords'], data['waveforms']
+    xs, ys = [], []
+    lit = {uid: ([], []) for uid in highlight}
+    marks = []
+    for j, row in enumerate(units.itertuples()):
+        uid = int(row.unit_id)
+        if uid not in waveforms:
+            continue
+        channels, wf = waveforms[uid]
+        n_samples = wf.shape[0]
+        tt = np.linspace(-1, 1, n_samples, dtype=np.float32) * 15
+        cx, cy = coords[channels, 0], coords[channels, 1]
+        x0 = j * spacing
+        centre = float(cx.mean())
+        target = lit[uid] if uid in lit else (xs, ys)
+        for c in range(len(channels)):
+            target[0].append(x0 + tt + 1.2 * (cx[c] - centre))
+            target[0].append([np.nan])
+            target[1].append(cy[c] + wf[:, c] * scale)
+            target[1].append([np.nan])
+        marks.append((x0, float(row.depth), uid, int(row.shank), int(row.n_spikes)))
+    fig = go.Figure()
+    if xs:
+        fig.add_trace(go.Scattergl(x=np.concatenate(xs), y=np.concatenate(ys), mode='lines',
+                                   line=dict(color='black', width=1), hoverinfo='skip', showlegend=False))
+    for uid in highlight:
+        lx, ly = lit.get(uid, ([], []))
+        if lx:
+            fig.add_trace(go.Scattergl(x=np.concatenate(lx), y=np.concatenate(ly), mode='lines',
+                                       line=dict(color=colours.get(uid, '#d62728'), width=2),
+                                       hoverinfo='skip', name=f'unit {uid}'))
+    if marks:
+        m = pd.DataFrame(marks, columns=['x', 'depth', 'unit_id', 'shank', 'n_spikes'])
+        fig.add_trace(go.Scattergl(
+            x=m['x'], y=m['depth'], mode='markers+text', text=m['unit_id'].astype(str),
+            textposition='top center', textfont=dict(color='darkorange', size=10),
+            marker=dict(size=14, opacity=0.35, color=m['shank'], colorscale='Viridis', showscale=False),
+            customdata=m[['unit_id', 'shank', 'n_spikes', 'depth']].values.tolist(),
+            hovertemplate=('unit %{customdata[0]}<br>shank %{customdata[1]}<br>'
+                           '%{customdata[2]} spikes · depth %{customdata[3]:.0f} µm<extra></extra>'),
+            showlegend=False, name='units'))
+    fig.update_layout(height=720, margin=dict(l=60, r=10, t=10, b=10), template='simple_white',
+                      xaxis=dict(visible=False), yaxis_title='Depth (µm)', uirevision=revision,
+                      showlegend=bool(highlight), legend=dict(orientation='h', y=1.02, x=0),
+                      dragmode='pan')
+    return fig
+
+
+def _fresh_units(event, memory_key):
+    '''The units under the points of a plotly selection that are new since the last run.
+
+    The selection of a chart stays as it is across reruns, so the points seen
+    are remembered under ``memory_key`` and only the ones added since count.
+    '''
+    points = (getattr(event, 'selection', None) or {}).get('points', []) or []
+    signature = frozenset((p.get('curve_number'), p.get('point_index')) for p in points)
+    seen = st.session_state.get(memory_key, frozenset())
+    st.session_state[memory_key] = signature
+    units = []
+    for p in points:
+        if (p.get('curve_number'), p.get('point_index')) in seen:
+            continue
+        uid = _unit_of(p)
+        if uid is not None and uid not in units:
+            units.append(uid)
+    return units
+
+
+def _unit_of(point):
+    '''The unit id a plotly point carries first in its customdata, or None.
+
+    The browser hands the row back as a list, or as a dict with string keys
+    when plotly sent it as a typed array, or as a bare value.
+    '''
+    custom = point.get('customdata')
+    if isinstance(custom, dict):
+        custom = custom.get('0', custom.get(0))
+    elif isinstance(custom, (list, tuple)):
+        custom = custom[0] if len(custom) else None
+    try:
+        return int(custom)
+    except (TypeError, ValueError):
+        return None
+
+
+def drift_section(schema, sel_key, probe, label_of, sua_units, selected, colours, on_pick):
+    '''Draw the drift map and the waveforms over depth of one probe.
+
+    ``sel_key`` names the session and parameter set, ``probe`` the probe;
+    ``label_of`` maps unit ids to their label under the current criteria
+    ('SUA', 'MUA', 'noise' or None) and ``sua_units`` are the ids that pass
+    it as single units.  ``selected`` are the selected unit ids on the probe,
+    drawn in ``colours`` (unit id -> colour) on both plots, and
+    ``on_pick(unit_ids)`` is called with the units clicked on either plot.
+    '''
+    c_sua, c_draw, c_colour = st.columns([1.3, 1.8, 1.3], vertical_alignment='bottom')
+    sua_only = c_sua.toggle('Single units only', value=True, key='srt_drift_sua_only',
+                            help='Read and draw only the units that pass the criteria as single units.')
+    n_draw = c_draw.select_slider('Spikes drawn', DRAW_OPTIONS, value=100_000, key='srt_drift_n_draw',
+                                  format_func=lambda n: f'{n // 1000}k')
+    colour_by = c_colour.selectbox('Colour by', ['amplitude', 'unit', 'label'], key='srt_drift_colour')
+    unit_filter = tuple(sorted(int(u) for u in sua_units)) if sua_only else None
+    if sua_only and not unit_filter:
+        st.caption('No single units on this probe under the criteria; turn the filter off to read every unit.')
+        return
+    with st.spinner(f'Reading the spikes of probe {probe}…'):
+        data = _load_probe(schema, sel_key['subject_name'], sel_key['session_name'], sel_key['dataset_name'],
+                           int(sel_key['parameter_set_num']), int(probe), unit_filter)
+    shanks = sorted(int(x) for x in data['units']['shank'].unique())
+    chosen = shanks
+    if len(shanks) > 1:
+        # the shanks are filtered after the read, so switching them costs nothing.
+        # a choice of every shank stays "every shank" when the set changes with the
+        # probe or the filter; a narrower choice keeps what still applies
+        stored = st.session_state.get('srt_drift_shanks')
+        against = st.session_state.get('_srt_drift_shanks_all')
+        if stored is None or (against != shanks and sorted(stored) == against):
+            st.session_state['srt_drift_shanks'] = shanks
+        else:
+            valid = [x for x in stored if x in shanks]
+            st.session_state['srt_drift_shanks'] = valid or shanks
+        st.session_state['_srt_drift_shanks_all'] = shanks
+        chosen = st.segmented_control('Shanks', shanks, selection_mode='multi', key='srt_drift_shanks',
+                                      format_func=lambda x: f'shank {x}' if x >= 0 else 'no shank')
+        chosen = sorted(chosen) if chosen else shanks
+        if chosen != shanks:
+            data = _only_units(data, set(data['units'].loc[data['units']['shank'].isin(chosen), 'unit_id']))
+    n_units = len(data['units'])
+    st.caption(f'{n_units} units · {len(data["times"]):,} spikes'
+               + (f' of {data["n_spikes"]:,}' if len(data['times']) != data['n_spikes'] else '')
+               + f', {min(n_draw, len(data["times"])):,} drawn'
+               + (' · single units only' if unit_filter is not None else '')
+               + (f' · shank {", ".join(str(x) for x in chosen)}' if chosen != shanks else ''))
+    if not n_units:
+        return
+    if not np.isfinite(data['depths']).any():
+        st.warning('The spikes of this sorting have no positions, so the drift map is empty.')
+    unit_ids = {int(u) for u in data['units']['unit_id']}
+    highlight = [int(u) for u in selected if int(u) in unit_ids]
+    revision = (f'{sel_key["subject_name"]}/{sel_key["session_name"]}/{sel_key["dataset_name"]}/'
+                f'{sel_key["parameter_set_num"]}/{probe}/{"sua" if unit_filter is not None else "all"}/'
+                + '-'.join(str(x) for x in chosen))
+
+    event = st.plotly_chart(_drift_figure(data, n_draw, colour_by, label_of, highlight, colours, revision),
+                            key=f'srt_drift_map_{revision}', on_select='rerun',
+                            selection_mode=('points', 'box', 'lasso'), width='stretch')
+    picked = _fresh_units(event, f'_srt_drift_map_seen_{revision}')
+    c_scale, c_spacing = st.columns(2)
+    scale = c_scale.slider('Waveform scale (µm per µV)', 0.02, 0.5, 0.1, 0.01, key='srt_drift_wf_scale')
+    spacing = c_spacing.slider('Unit spacing (µm)', 40, 300, 120, 10, key='srt_drift_wf_spacing')
+    event_wf = st.plotly_chart(_waveform_figure(data, highlight, colours, scale, spacing, revision),
+                               key=f'srt_drift_wf_{revision}', on_select='rerun',
+                               selection_mode=('points',), width='stretch')
+    picked += [u for u in _fresh_units(event_wf, f'_srt_drift_wf_seen_{revision}') if u not in picked]
+    fresh = [u for u in picked if u in unit_ids and u not in highlight]
+    if fresh:
+        on_pick(fresh)