backends

navis.nbl.backends.BuiltinBackend

The built-in multiprocessing NBLAST backend.

Source code in navis/nbl/backends/builtin.py

class BuiltinBackend(NblastBackend):
    """The built-in multiprocessing NBLAST backend."""

    name = "builtin"
    priority = 0

    def available(self):
        # Always available - it's pure navis
        return True

    def unsupported(self, operation, **params):
        # The built-in backend supports every operation and parameter
        return super().unsupported(operation)

    # ------------------------------------------------------------------ #
    # Shared helpers
    # ------------------------------------------------------------------ #
    def _map(self, jobs, n_cores, progress, desc="NBLASTing"):
        """Run `jobs` and yield ``(job, result)`` tuples.

        Each *job* is an ``NBlaster`` carrying the picklable attributes read by
        :func:`_run_job` (``_op``, ``_scores`` and the relevant index arrays).
        With a single job (or a single core) work runs inline; otherwise it is
        spread across a spawn-based process pool.

        Override this method to swap in a different dispatcher (joblib, threads,
        serial, ...) while reusing the partitioning and stitching logic.
        """
        from ..nblast_funcs import set_omp_flag, OMP_NUM_THREADS_LIMIT

        multicore = bool(n_cores and n_cores > 1 and len(jobs) > 1)

        # Avoid multiple layers of concurrency (see pykdtree/OMP notes)
        with set_omp_flag(limits=OMP_NUM_THREADS_LIMIT if (n_cores and n_cores > 1) else None):
            if not multicore:
                for this in jobs:
                    yield this, _run_job(this)
                return

            # Note that we're forcing "spawn" instead of "fork" (default on
            # linux) to reduce the memory footprint: "fork" appears to inherit
            # all variables (including all neurons) while "spawn" gets only
            # what's required to run the job.
            with ProcessPoolExecutor(max_workers=n_cores,
                                     mp_context=mp.get_context('spawn')) as pool:
                futures = {}
                for this in jobs:
                    this.progress = False  # no per-block progress bars
                    futures[pool.submit(_run_job, this)] = this

                # We drop the "N / N_total" bit from the progress bar because
                # it's not helpful here.
                fmt = '{desc}: {percentage:3.0f}%|{bar}| [{elapsed}<{remaining}]'
                for f in config.tqdm(as_completed(futures),
                                     desc=desc,
                                     bar_format=fmt,
                                     total=len(futures),
                                     smoothing=0,
                                     disable=not progress,
                                     leave=False):
                    yield futures[f], f.result()

    def _make_blaster(self, use_alpha, normalized, smat, limit_dist, precision,
                      approx_nn, progress, smat_kwargs):
        from ..nblast_funcs import NBlaster
        return NBlaster(use_alpha=use_alpha,
                        normalized=normalized,
                        smat=smat,
                        limit_dist=limit_dist,
                        dtype=precision,
                        approx_nn=approx_nn,
                        progress=progress,
                        smat_kwargs=smat_kwargs)

    def _partition(self, q, t, n_cores, progress, aba=False):
        """Find (n_rows, n_cols) partition of the query/target matrix."""
        from ..nblast_funcs import (find_batch_partition, find_optimal_partition,
                                     JOB_SIZE_MULTIPLIER, JOB_MAX_TIME_SECONDS)

        if not (n_cores and n_cores > 1):
            return 1, 1

        if progress:
            # If progress bar, we need to make smaller mini batches. These mini
            # jobs must not be too small - otherwise the overhead from spawning
            # and sending results between processes slows things down
            # dramatically. Hence we want each job to run for >10s. The run time
            # depends on the system and how big the neurons are, so we run a
            # quick test and extrapolate.
            return find_batch_partition(q, t, T=10 * JOB_SIZE_MULTIPLIER)

        # No progress bar: for a plain query->target NBLAST we aim for each
        # batch to finish in a certain amount of time (to avoid stragglers);
        # for all-by-all we just split evenly across cores.
        if aba:
            return find_optimal_partition(n_cores, q, t)

        n_rows, n_cols = find_batch_partition(q, t, T=JOB_MAX_TIME_SECONDS)
        if (n_rows * n_cols) < n_cores:
            n_rows, n_cols = find_optimal_partition(n_cores, q, t)
        return n_rows, n_cols

    # ------------------------------------------------------------------ #
    # Operations
    # ------------------------------------------------------------------ #
    def nblast(self, query, target, *, scores, normalized, use_alpha, smat,
               limit_dist, approx_nn, precision, n_cores, progress, smat_kwargs):
        """Query -> target NBLAST."""
        query_dps, target_dps = query, target

        n_rows, n_cols = self._partition(query_dps, target_dps, n_cores, progress)

        # Calculate self-hits once for all neurons
        nb = self._make_blaster(use_alpha, normalized, smat, limit_dist,
                                precision, approx_nn, progress, smat_kwargs)
        query_self_hits = np.array([nb.calc_self_hit(n) for n in query_dps])
        target_self_hits = np.array([nb.calc_self_hit(n) for n in target_dps])

        # Build one blaster per block of the score matrix
        jobs = []
        with config.tqdm(desc='Preparing', total=n_rows * n_cols, leave=False,
                         disable=not progress) as pbar:
            for qix in np.array_split(np.arange(len(query_dps)), n_rows):
                for tix in np.array_split(np.arange(len(target_dps)), n_cols):
                    this = self._make_blaster(use_alpha, normalized, smat,
                                              limit_dist, precision, approx_nn,
                                              progress, smat_kwargs)
                    for ix in qix:
                        this.append(query_dps[ix], query_self_hits[ix])
                    for ix in tix:
                        this.append(target_dps[ix], target_self_hits[ix])

                    this.queries = np.arange(len(qix))
                    this.targets = np.arange(len(tix)) + len(qix)
                    this.queries_ix = qix
                    this.targets_ix = tix
                    this.pbar_position = len(jobs) if not utils.is_jupyter() else None
                    this._op = 'multi_query_target'
                    this._scores = scores

                    jobs.append(this)
                    pbar.update()

        # Single block: return its labeled DataFrame directly
        if len(jobs) == 1:
            (this, res), = self._map(jobs, n_cores, progress)
            return res

        # Multiple blocks: stitch results into the big matrix
        out = pd.DataFrame(np.empty((len(query_dps), len(target_dps)),
                                    dtype=nb.dtype),
                           index=query_dps.id, columns=target_dps.id)
        out.index.name = 'query'
        out.columns.name = 'target'
        for this, res in self._map(jobs, n_cores, progress):
            out.iloc[this.queries_ix, this.targets_ix] = res.values

        return out

    def nblast_allbyall(self, x, *, normalized, use_alpha, smat, limit_dist,
                        approx_nn, precision, n_cores, progress, smat_kwargs):
        """All-by-all NBLAST (always forward scores)."""
        dps = x

        n_rows, n_cols = self._partition(dps, dps, n_cores, progress, aba=True)

        # Calculate self-hits once for all neurons
        nb = self._make_blaster(use_alpha, normalized, smat, limit_dist,
                                precision, approx_nn, progress, smat_kwargs)
        self_hits = np.array([nb.calc_self_hit(n) for n in dps])

        jobs = []
        with config.tqdm(desc='Preparing', total=n_rows * n_cols, leave=False,
                         disable=not progress) as pbar:
            for qix in np.array_split(np.arange(len(dps)), n_rows):
                for tix in np.array_split(np.arange(len(dps)), n_cols):
                    this = self._make_blaster(use_alpha, normalized, smat,
                                              limit_dist, precision, approx_nn,
                                              progress, smat_kwargs)

                    # Make sure we don't add the same neuron twice
                    to_add = list(set(qix) | set(tix))
                    ixmap = {}
                    for i, ix in enumerate(to_add):
                        this.append(dps[ix], self_hits[ix])
                        ixmap[ix] = i

                    this.queries = [ixmap[ix] for ix in qix]
                    this.targets = [ixmap[ix] for ix in tix]
                    this.queries_ix = qix
                    this.targets_ix = tix
                    this.pbar_position = len(jobs) if not utils.is_jupyter() else None
                    this._op = 'multi_query_target'
                    this._scores = 'forward'

                    jobs.append(this)
                    pbar.update()

        if len(jobs) == 1:
            return jobs[0].all_by_all()

        out = pd.DataFrame(np.empty((len(dps), len(dps)), dtype=nb.dtype),
                           index=dps.id, columns=dps.id)
        out.index.name = 'query'
        out.columns.name = 'target'
        for this, res in self._map(jobs, n_cores, progress):
            out.iloc[this.queries_ix, this.targets_ix] = res.values

        return out

    def nblast_smart(self, query, target, *, aba, t, criterion, scores,
                     return_mask, normalized, use_alpha, smat, limit_dist,
                     approx_nn, precision, n_cores, progress, smat_kwargs):
        """Smart(er) NBLAST: pre-NBLAST on simplified dotprops, then full."""
        query_dps, target_dps = query, target

        pre_scores = scores
        # For all-by-all's we can compute only forward scores during the
        # pre-NBLAST and produce the mean later.
        if aba and scores == 'mean':
            pre_scores = 'forward'

        try:
            t = int(t)
        except BaseException:
            raise TypeError(f'`t` must be (convertable to) integer - got "{type(t)}"')

        if criterion == 'percentile':
            if (t <= 0 or t >= 100):
                raise ValueError('Expected `t` to be integer between 0 and 100 for '
                                 f'criterion "percentile", got {t}')
        elif criterion == 'N':
            if (t < 0 or t > len(target_dps)):
                raise ValueError('`t` must be between 0 and the total number of '
                                 f'targets ({len(target_dps)}) for criterion "N", '
                                 f'got {t}')

        # Make simplified dotprops
        query_dps_simp = query_dps.downsample(10, inplace=False)
        if not aba:
            target_dps_simp = target_dps.downsample(10, inplace=False)
        else:
            target_dps_simp = query_dps_simp

        # --- Pre-NBLAST on simplified dotprops --- #
        n_rows, n_cols = self._partition(query_dps_simp, target_dps_simp,
                                         n_cores, progress)

        nb = self._make_blaster(use_alpha, normalized, smat, limit_dist,
                                precision, approx_nn, progress, smat_kwargs)
        query_self_hits = np.array([nb.calc_self_hit(n) for n in query_dps_simp])
        target_self_hits = np.array([nb.calc_self_hit(n) for n in target_dps_simp])

        jobs = []
        with config.tqdm(desc='Prep. pre-NBLAST', total=n_rows * n_cols,
                         leave=False, disable=not progress) as pbar:
            for qix in np.array_split(np.arange(len(query_dps_simp)), n_rows):
                for tix in np.array_split(np.arange(len(target_dps_simp)), n_cols):
                    this = self._make_blaster(use_alpha, normalized, smat,
                                              limit_dist, precision, approx_nn,
                                              progress, smat_kwargs)
                    for ix in qix:
                        this.append(query_dps_simp[ix], query_self_hits[ix])
                    for ix in tix:
                        this.append(target_dps_simp[ix], target_self_hits[ix])

                    this.queries = np.arange(len(qix))
                    this.targets = np.arange(len(tix)) + len(qix)
                    this.queries_ix = qix
                    this.targets_ix = tix
                    this.pbar_position = len(jobs) if not utils.is_jupyter() else None
                    this._op = 'multi_query_target'
                    this._scores = pre_scores

                    jobs.append(this)
                    pbar.update()

        if len(jobs) == 1:
            (this, res), = self._map(jobs, n_cores, progress, desc='Pre-NBLASTs')
            scr = res
        else:
            scr = pd.DataFrame(np.empty((len(query_dps_simp),
                                         len(target_dps_simp)), dtype=nb.dtype),
                               index=query_dps_simp.id, columns=target_dps_simp.id)
            scr.index.name = 'query'
            scr.columns.name = 'target'
            for this, res in self._map(jobs, n_cores, progress, desc='Pre-NBLASTs'):
                scr.iloc[this.queries_ix, this.targets_ix] = res.values

        # If this is an all-by-all and we computed only forward scores
        if aba and scores == 'mean':
            scr = (scr + scr.T.values) / 2

        # Now select targets of interest for each query
        if criterion == 'percentile':
            sel = np.percentile(scr, q=t, axis=1)
            mask = scr >= sel.reshape(-1, 1)
        elif criterion == 'score':
            sel = np.full(scr.shape[0], fill_value=t)
            mask = scr >= sel.reshape(-1, 1)
        else:
            srt = np.argsort(scr.values, axis=1)[:, ::-1]
            mask = pd.DataFrame(np.zeros(scr.shape, dtype=bool),
                                columns=scr.columns, index=scr.index)
            _ = np.arange(mask.shape[0])
            for N in range(t):
                mask.iloc[_, srt[:, N]] = True

        # --- Full NBLAST on the selected pairs --- #
        query_self_hits = np.array([nb.calc_self_hit(n) for n in query_dps])
        target_self_hits = np.array([nb.calc_self_hit(n) for n in target_dps])

        jobs = []
        with config.tqdm(desc='Prep. full NBLAST', total=n_rows * n_cols,
                         leave=False, disable=not progress) as pbar:
            for qix in np.array_split(np.arange(len(query_dps)), n_rows):
                for tix in np.array_split(np.arange(len(target_dps)), n_cols):
                    this = self._make_blaster(use_alpha, normalized, smat,
                                              limit_dist, precision, approx_nn,
                                              progress, smat_kwargs)
                    for ix in qix:
                        this.append(query_dps[ix], query_self_hits[ix])
                    for ix in tix:
                        this.append(target_dps[ix], target_self_hits[ix])

                    # Find the pairs to NBLAST in this part of the matrix
                    submask = mask.loc[query_dps[qix].id, target_dps[tix].id]
                    # `pairs` is an array of `[[query, target], [...]]` pairs
                    this.pairs = np.vstack(np.where(submask)).T
                    # Offset the target indices
                    this.pairs[:, 1] += len(qix)

                    # Track this blaster's mask relative to the original big one
                    this.mask = np.zeros(mask.shape, dtype=bool)
                    this.mask[qix[0]:qix[-1] + 1, tix[0]:tix[-1] + 1] = submask

                    this.pbar_position = len(jobs) if not utils.is_jupyter() else None
                    this.desc = 'Full NBLAST'
                    this._op = 'pair_query_target'
                    this._scores = scores

                    jobs.append(this)
                    pbar.update()

        if len(jobs) == 1:
            (this, res), = self._map(jobs, n_cores, progress)
            scr[mask] = res
        else:
            for this, res in self._map(jobs, n_cores, progress):
                scr[this.mask] = res

        if return_mask:
            return scr, mask

        return scr

    def synblast(self, query, target, *, by_type, cn_types, scores, normalized,
                 smat, n_cores, progress):
        """Synapse-based NBLAST (SynBLAST)."""
        from ..synblast_funcs import SynBlaster, find_batch_partition
        from ..nblast_funcs import find_optimal_partition

        def get_connectors(n):
            if cn_types is not None:
                return n.connectors[n.connectors['type'].isin(cn_types)]
            return n.connectors

        # Find a partition that produces batches that each run in ~10s
        if n_cores and n_cores > 1:
            if progress:
                n_rows, n_cols = find_batch_partition(query, target, T=10)
            else:
                n_rows, n_cols = find_optimal_partition(n_cores, query, target)
        else:
            n_rows = n_cols = 1

        # Calculate self-hits once for all neurons
        nb = SynBlaster(normalized=normalized, by_type=by_type, smat=smat,
                        progress=progress)
        query_self_hits = np.array([nb.calc_self_hit(get_connectors(n)) for n in query])
        target_self_hits = np.array([nb.calc_self_hit(get_connectors(n)) for n in target])

        jobs = []
        with config.tqdm(desc='Preparing', total=n_rows * n_cols, leave=False,
                         disable=not progress) as pbar:
            for qix in np.array_split(np.arange(len(query)), n_rows):
                for tix in np.array_split(np.arange(len(target)), n_cols):
                    this = SynBlaster(normalized=normalized, by_type=by_type,
                                      smat=smat, progress=progress)
                    for ix in qix:
                        n = query[ix]
                        this.append(get_connectors(n), id=n.id,
                                    self_hit=query_self_hits[ix])
                    for ix in tix:
                        n = target[ix]
                        this.append(get_connectors(n), id=n.id,
                                    self_hit=target_self_hits[ix])

                    this.queries = np.arange(len(qix))
                    this.targets = np.arange(len(tix)) + len(qix)
                    this.queries_ix = qix
                    this.targets_ix = tix
                    this.pbar_position = len(jobs) if not utils.is_jupyter() else None
                    this._op = 'multi_query_target'
                    this._scores = scores

                    jobs.append(this)
                    pbar.update()

        if len(jobs) == 1:
            (this, res), = self._map(jobs, n_cores, progress)
            return res

        out = pd.DataFrame(np.empty((len(query), len(target)), dtype=nb.dtype),
                           index=query.id, columns=target.id)
        out.index.name = 'query'
        out.columns.name = 'target'
        for this, res in self._map(jobs, n_cores, progress):
            out.iloc[this.queries_ix, this.targets_ix] = res.values

        return out

nblast

Query -> target NBLAST.

Source code in navis/nbl/backends/builtin.py

def nblast(self, query, target, *, scores, normalized, use_alpha, smat,
           limit_dist, approx_nn, precision, n_cores, progress, smat_kwargs):
    """Query -> target NBLAST."""
    query_dps, target_dps = query, target

    n_rows, n_cols = self._partition(query_dps, target_dps, n_cores, progress)

    # Calculate self-hits once for all neurons
    nb = self._make_blaster(use_alpha, normalized, smat, limit_dist,
                            precision, approx_nn, progress, smat_kwargs)
    query_self_hits = np.array([nb.calc_self_hit(n) for n in query_dps])
    target_self_hits = np.array([nb.calc_self_hit(n) for n in target_dps])

    # Build one blaster per block of the score matrix
    jobs = []
    with config.tqdm(desc='Preparing', total=n_rows * n_cols, leave=False,
                     disable=not progress) as pbar:
        for qix in np.array_split(np.arange(len(query_dps)), n_rows):
            for tix in np.array_split(np.arange(len(target_dps)), n_cols):
                this = self._make_blaster(use_alpha, normalized, smat,
                                          limit_dist, precision, approx_nn,
                                          progress, smat_kwargs)
                for ix in qix:
                    this.append(query_dps[ix], query_self_hits[ix])
                for ix in tix:
                    this.append(target_dps[ix], target_self_hits[ix])

                this.queries = np.arange(len(qix))
                this.targets = np.arange(len(tix)) + len(qix)
                this.queries_ix = qix
                this.targets_ix = tix
                this.pbar_position = len(jobs) if not utils.is_jupyter() else None
                this._op = 'multi_query_target'
                this._scores = scores

                jobs.append(this)
                pbar.update()

    # Single block: return its labeled DataFrame directly
    if len(jobs) == 1:
        (this, res), = self._map(jobs, n_cores, progress)
        return res

    # Multiple blocks: stitch results into the big matrix
    out = pd.DataFrame(np.empty((len(query_dps), len(target_dps)),
                                dtype=nb.dtype),
                       index=query_dps.id, columns=target_dps.id)
    out.index.name = 'query'
    out.columns.name = 'target'
    for this, res in self._map(jobs, n_cores, progress):
        out.iloc[this.queries_ix, this.targets_ix] = res.values

    return out

nblast_allbyall

All-by-all NBLAST (always forward scores).

Source code in navis/nbl/backends/builtin.py

def nblast_allbyall(self, x, *, normalized, use_alpha, smat, limit_dist,
                    approx_nn, precision, n_cores, progress, smat_kwargs):
    """All-by-all NBLAST (always forward scores)."""
    dps = x

    n_rows, n_cols = self._partition(dps, dps, n_cores, progress, aba=True)

    # Calculate self-hits once for all neurons
    nb = self._make_blaster(use_alpha, normalized, smat, limit_dist,
                            precision, approx_nn, progress, smat_kwargs)
    self_hits = np.array([nb.calc_self_hit(n) for n in dps])

    jobs = []
    with config.tqdm(desc='Preparing', total=n_rows * n_cols, leave=False,
                     disable=not progress) as pbar:
        for qix in np.array_split(np.arange(len(dps)), n_rows):
            for tix in np.array_split(np.arange(len(dps)), n_cols):
                this = self._make_blaster(use_alpha, normalized, smat,
                                          limit_dist, precision, approx_nn,
                                          progress, smat_kwargs)

                # Make sure we don't add the same neuron twice
                to_add = list(set(qix) | set(tix))
                ixmap = {}
                for i, ix in enumerate(to_add):
                    this.append(dps[ix], self_hits[ix])
                    ixmap[ix] = i

                this.queries = [ixmap[ix] for ix in qix]
                this.targets = [ixmap[ix] for ix in tix]
                this.queries_ix = qix
                this.targets_ix = tix
                this.pbar_position = len(jobs) if not utils.is_jupyter() else None
                this._op = 'multi_query_target'
                this._scores = 'forward'

                jobs.append(this)
                pbar.update()

    if len(jobs) == 1:
        return jobs[0].all_by_all()

    out = pd.DataFrame(np.empty((len(dps), len(dps)), dtype=nb.dtype),
                       index=dps.id, columns=dps.id)
    out.index.name = 'query'
    out.columns.name = 'target'
    for this, res in self._map(jobs, n_cores, progress):
        out.iloc[this.queries_ix, this.targets_ix] = res.values

    return out

nblast_smart

Smart(er) NBLAST: pre-NBLAST on simplified dotprops, then full.

Source code in navis/nbl/backends/builtin.py

def nblast_smart(self, query, target, *, aba, t, criterion, scores,
                 return_mask, normalized, use_alpha, smat, limit_dist,
                 approx_nn, precision, n_cores, progress, smat_kwargs):
    """Smart(er) NBLAST: pre-NBLAST on simplified dotprops, then full."""
    query_dps, target_dps = query, target

    pre_scores = scores
    # For all-by-all's we can compute only forward scores during the
    # pre-NBLAST and produce the mean later.
    if aba and scores == 'mean':
        pre_scores = 'forward'

    try:
        t = int(t)
    except BaseException:
        raise TypeError(f'`t` must be (convertable to) integer - got "{type(t)}"')

    if criterion == 'percentile':
        if (t <= 0 or t >= 100):
            raise ValueError('Expected `t` to be integer between 0 and 100 for '
                             f'criterion "percentile", got {t}')
    elif criterion == 'N':
        if (t < 0 or t > len(target_dps)):
            raise ValueError('`t` must be between 0 and the total number of '
                             f'targets ({len(target_dps)}) for criterion "N", '
                             f'got {t}')

    # Make simplified dotprops
    query_dps_simp = query_dps.downsample(10, inplace=False)
    if not aba:
        target_dps_simp = target_dps.downsample(10, inplace=False)
    else:
        target_dps_simp = query_dps_simp

    # --- Pre-NBLAST on simplified dotprops --- #
    n_rows, n_cols = self._partition(query_dps_simp, target_dps_simp,
                                     n_cores, progress)

    nb = self._make_blaster(use_alpha, normalized, smat, limit_dist,
                            precision, approx_nn, progress, smat_kwargs)
    query_self_hits = np.array([nb.calc_self_hit(n) for n in query_dps_simp])
    target_self_hits = np.array([nb.calc_self_hit(n) for n in target_dps_simp])

    jobs = []
    with config.tqdm(desc='Prep. pre-NBLAST', total=n_rows * n_cols,
                     leave=False, disable=not progress) as pbar:
        for qix in np.array_split(np.arange(len(query_dps_simp)), n_rows):
            for tix in np.array_split(np.arange(len(target_dps_simp)), n_cols):
                this = self._make_blaster(use_alpha, normalized, smat,
                                          limit_dist, precision, approx_nn,
                                          progress, smat_kwargs)
                for ix in qix:
                    this.append(query_dps_simp[ix], query_self_hits[ix])
                for ix in tix:
                    this.append(target_dps_simp[ix], target_self_hits[ix])

                this.queries = np.arange(len(qix))
                this.targets = np.arange(len(tix)) + len(qix)
                this.queries_ix = qix
                this.targets_ix = tix
                this.pbar_position = len(jobs) if not utils.is_jupyter() else None
                this._op = 'multi_query_target'
                this._scores = pre_scores

                jobs.append(this)
                pbar.update()

    if len(jobs) == 1:
        (this, res), = self._map(jobs, n_cores, progress, desc='Pre-NBLASTs')
        scr = res
    else:
        scr = pd.DataFrame(np.empty((len(query_dps_simp),
                                     len(target_dps_simp)), dtype=nb.dtype),
                           index=query_dps_simp.id, columns=target_dps_simp.id)
        scr.index.name = 'query'
        scr.columns.name = 'target'
        for this, res in self._map(jobs, n_cores, progress, desc='Pre-NBLASTs'):
            scr.iloc[this.queries_ix, this.targets_ix] = res.values

    # If this is an all-by-all and we computed only forward scores
    if aba and scores == 'mean':
        scr = (scr + scr.T.values) / 2

    # Now select targets of interest for each query
    if criterion == 'percentile':
        sel = np.percentile(scr, q=t, axis=1)
        mask = scr >= sel.reshape(-1, 1)
    elif criterion == 'score':
        sel = np.full(scr.shape[0], fill_value=t)
        mask = scr >= sel.reshape(-1, 1)
    else:
        srt = np.argsort(scr.values, axis=1)[:, ::-1]
        mask = pd.DataFrame(np.zeros(scr.shape, dtype=bool),
                            columns=scr.columns, index=scr.index)
        _ = np.arange(mask.shape[0])
        for N in range(t):
            mask.iloc[_, srt[:, N]] = True

    # --- Full NBLAST on the selected pairs --- #
    query_self_hits = np.array([nb.calc_self_hit(n) for n in query_dps])
    target_self_hits = np.array([nb.calc_self_hit(n) for n in target_dps])

    jobs = []
    with config.tqdm(desc='Prep. full NBLAST', total=n_rows * n_cols,
                     leave=False, disable=not progress) as pbar:
        for qix in np.array_split(np.arange(len(query_dps)), n_rows):
            for tix in np.array_split(np.arange(len(target_dps)), n_cols):
                this = self._make_blaster(use_alpha, normalized, smat,
                                          limit_dist, precision, approx_nn,
                                          progress, smat_kwargs)
                for ix in qix:
                    this.append(query_dps[ix], query_self_hits[ix])
                for ix in tix:
                    this.append(target_dps[ix], target_self_hits[ix])

                # Find the pairs to NBLAST in this part of the matrix
                submask = mask.loc[query_dps[qix].id, target_dps[tix].id]
                # `pairs` is an array of `[[query, target], [...]]` pairs
                this.pairs = np.vstack(np.where(submask)).T
                # Offset the target indices
                this.pairs[:, 1] += len(qix)

                # Track this blaster's mask relative to the original big one
                this.mask = np.zeros(mask.shape, dtype=bool)
                this.mask[qix[0]:qix[-1] + 1, tix[0]:tix[-1] + 1] = submask

                this.pbar_position = len(jobs) if not utils.is_jupyter() else None
                this.desc = 'Full NBLAST'
                this._op = 'pair_query_target'
                this._scores = scores

                jobs.append(this)
                pbar.update()

    if len(jobs) == 1:
        (this, res), = self._map(jobs, n_cores, progress)
        scr[mask] = res
    else:
        for this, res in self._map(jobs, n_cores, progress):
            scr[this.mask] = res

    if return_mask:
        return scr, mask

    return scr

synblast

Synapse-based NBLAST (SynBLAST).

Source code in navis/nbl/backends/builtin.py

def synblast(self, query, target, *, by_type, cn_types, scores, normalized,
             smat, n_cores, progress):
    """Synapse-based NBLAST (SynBLAST)."""
    from ..synblast_funcs import SynBlaster, find_batch_partition
    from ..nblast_funcs import find_optimal_partition

    def get_connectors(n):
        if cn_types is not None:
            return n.connectors[n.connectors['type'].isin(cn_types)]
        return n.connectors

    # Find a partition that produces batches that each run in ~10s
    if n_cores and n_cores > 1:
        if progress:
            n_rows, n_cols = find_batch_partition(query, target, T=10)
        else:
            n_rows, n_cols = find_optimal_partition(n_cores, query, target)
    else:
        n_rows = n_cols = 1

    # Calculate self-hits once for all neurons
    nb = SynBlaster(normalized=normalized, by_type=by_type, smat=smat,
                    progress=progress)
    query_self_hits = np.array([nb.calc_self_hit(get_connectors(n)) for n in query])
    target_self_hits = np.array([nb.calc_self_hit(get_connectors(n)) for n in target])

    jobs = []
    with config.tqdm(desc='Preparing', total=n_rows * n_cols, leave=False,
                     disable=not progress) as pbar:
        for qix in np.array_split(np.arange(len(query)), n_rows):
            for tix in np.array_split(np.arange(len(target)), n_cols):
                this = SynBlaster(normalized=normalized, by_type=by_type,
                                  smat=smat, progress=progress)
                for ix in qix:
                    n = query[ix]
                    this.append(get_connectors(n), id=n.id,
                                self_hit=query_self_hits[ix])
                for ix in tix:
                    n = target[ix]
                    this.append(get_connectors(n), id=n.id,
                                self_hit=target_self_hits[ix])

                this.queries = np.arange(len(qix))
                this.targets = np.arange(len(tix)) + len(qix)
                this.queries_ix = qix
                this.targets_ix = tix
                this.pbar_position = len(jobs) if not utils.is_jupyter() else None
                this._op = 'multi_query_target'
                this._scores = scores

                jobs.append(this)
                pbar.update()

    if len(jobs) == 1:
        (this, res), = self._map(jobs, n_cores, progress)
        return res

    out = pd.DataFrame(np.empty((len(query), len(target)), dtype=nb.dtype),
                       index=query.id, columns=target.id)
    out.index.name = 'query'
    out.columns.name = 'target'
    for this, res in self._map(jobs, n_cores, progress):
        out.iloc[this.queries_ix, this.targets_ix] = res.values

    return out

navis.nbl.backends.FastcoreBackend

NBLAST backend using navis-fastcore.

Source code in navis/nbl/backends/fastcore.py

class FastcoreBackend(NblastBackend):
    """NBLAST backend using ``navis-fastcore``."""

    name = "fastcore"
    # Preferred over the builtin backend when it can serve the request
    priority = 10

    def available(self):
        from ... import utils
        fc = utils.fastcore
        # Also guard against fastcore builds that predate the NBLAST API: the
        # older `nblast` name resolves to a submodule, not a callable.
        return (fc is not None
                and callable(getattr(fc, 'nblast', None))
                and callable(getattr(fc, 'nblast_allbyall', None)))

    def _smat_ok(self, smat):
        """Whether fastcore can consume this `smat`."""
        if isinstance(smat, str):
            return smat == 'auto'
        if smat is None or isinstance(smat, pd.DataFrame):
            return True
        # A navis Lookup2d (or anything with the expected duck-type)
        return hasattr(smat, 'cells') and hasattr(smat, 'axes')

    def unsupported(self, operation, **params):
        reasons = super().unsupported(operation)
        if reasons:
            return reasons

        reasons = []
        # `nblast_knn` arrived later than the rest of the NBLAST API, so an
        # installed fastcore can be perfectly usable for everything else while
        # lacking this one. Note this is checked here rather than in
        # `available()`/`implements()` so that an old fastcore still serves the
        # other operations, and so that a *missing* fastcore is still reported
        # as "implements it but is not installed".
        if operation == 'nblast_knn':
            from ... import utils
            if not callable(getattr(utils.fastcore, 'nblast_knn', None)):
                reasons.append("the installed navis-fastcore is too old to "
                               "provide `nblast_knn` "
                               "(`pip install -U navis-fastcore`)")
        if params.get('approx_nn', False):
            reasons.append("'approx_nn=True' is not supported by fastcore")
        if params.get('scores', None) == 'both':
            reasons.append("scores='both' is not supported by fastcore")
        if not self._smat_ok(params.get('smat', 'auto')):
            reasons.append("fastcore only supports smat='auto', None, a "
                           "DataFrame or a Lookup2d")
        return reasons

    def _convert_smat(self, smat, use_alpha):
        """Convert navis `smat` to something fastcore understands.

        We resolve ``'auto'`` to the very same FCWB ``Lookup2d`` the built-in
        backend uses, so both backends score against an identical matrix.
        """
        from ..smat import smat_fcwb, Lookup2d
        if isinstance(smat, str) and smat == 'auto':
            return smat_fcwb(use_alpha)
        if isinstance(smat, pd.DataFrame):
            return Lookup2d.from_dataframe(smat)
        # None or a ready Lookup2d
        return smat

    def nblast(self, query, target, *, scores, normalized, use_alpha, smat,
               limit_dist, approx_nn, precision, n_cores, progress, smat_kwargs):
        from ... import utils
        symmetry = None if scores in (None, 'forward') else scores
        M = utils.fastcore.nblast(
            query, target,
            smat=self._convert_smat(smat, use_alpha),
            normalize=normalized,
            symmetry=symmetry,
            use_alpha=use_alpha,
            limit_dist=limit_dist,
            n_cores=n_cores,
            precision=precision,
            progress=progress,
        )
        out = pd.DataFrame(M, index=query.id, columns=target.id)
        out.index.name = 'query'
        out.columns.name = 'target'
        return out

    def nblast_knn(self, query, target, *, k, scores, n_candidates, normalized,
                   use_alpha, smat, limit_dist, precision, n_cores, progress,
                   voxel, n_dirs, splat, format, smat_kwargs):
        from ... import utils
        symmetry = None if scores in (None, 'forward') else scores
        idx, sc = utils.fastcore.nblast_knn(
            query,
            # `None` (rather than `query`) is meaningful: it puts fastcore on
            # the all-by-all path, which excludes each neuron from its own
            # neighbour list. Passing `query` twice would give every row a
            # self-match at 1.0.
            target,
            k=k,
            symmetry=symmetry,
            n_candidates=n_candidates,
            voxel=voxel,
            n_dirs=n_dirs,
            splat=splat,
            smat=self._convert_smat(smat, use_alpha),
            normalize=normalized,
            use_alpha=use_alpha,
            limit_dist=limit_dist,
            n_cores=n_cores,
            precision=precision,
            progress=progress,
        )

        if format == 'arrays':
            return idx, sc

        # Map the (n_query, k) indices back onto neuron IDs. Rows with fewer
        # than `k` candidates are padded with -1/-inf by fastcore; `-1` would
        # silently index the *last* neuron, so mask before taking.
        target_ids = np.asarray((target if target is not None else query).id)
        query_ids = np.asarray(query.id)
        valid = idx >= 0
        matches = target_ids[np.where(valid, idx, 0)]
        # Take the width from the result rather than from `k`: they agree today,
        # but an IndexError here would be a puzzling way to find out otherwise.
        k = idx.shape[1]

        if format == 'long':
            rank = np.broadcast_to(np.arange(1, k + 1), idx.shape)
            out = pd.DataFrame({
                'query': np.repeat(query_ids, k)[valid.ravel()],
                'target': matches[valid],
                'score': sc[valid],
                'rank': rank[valid],
            })
            return out.reset_index(drop=True)

        # 'wide': the layout `navis.nbl.extract_matches` produces, so the two
        # are interchangeable downstream.
        out = pd.DataFrame({'id': query_ids})
        for i in range(k):
            col = np.where(valid[:, i], matches[:, i], None)
            out[f'match_{i + 1}'] = col
            out[f'score_{i + 1}'] = np.where(valid[:, i], sc[:, i], np.nan)
        return out

    def nblast_smart(self, query, target, *, aba, t, criterion, scores,
                     return_mask, normalized, use_alpha, smat, limit_dist,
                     approx_nn, precision, n_cores, progress, smat_kwargs):
        from ... import utils
        symmetry = None if scores in (None, 'forward') else scores
        res = utils.fastcore.nblast_smart(
            query,
            # `None` lets fastcore run its dedicated all-by-all pre-pass
            None if aba else target,
            t=t,
            criterion=criterion,
            smat=self._convert_smat(smat, use_alpha),
            normalize=normalized,
            symmetry=symmetry,
            use_alpha=use_alpha,
            limit_dist=limit_dist,
            n_cores=n_cores,
            precision=precision,
            progress=progress,
            return_mask=return_mask,
        )
        M, mask = res if return_mask else (res, None)

        out = pd.DataFrame(M, index=query.id, columns=target.id)
        out.index.name = 'query'
        out.columns.name = 'target'
        if not return_mask:
            return out

        mask = pd.DataFrame(mask, index=query.id, columns=target.id)
        mask.index.name = 'query'
        mask.columns.name = 'target'
        return out, mask

    def _encode_connectors(self, query, target, by_type, cn_types):
        """Turn navis neurons into fastcore's numeric connector arrays.

        navis stores connectors as DataFrames with string type labels, whereas
        fastcore wants each neuron as an ``(N, 3)`` (``[x, y, z]``) or - when
        ``by_type`` - ``(N, 4)`` (``[x, y, z, type]``) float array. When grouping
        by type we build a single string->int map shared across query and target
        so that like types compare against each other; ``cn_types`` filtering is
        applied here (so fastcore is called with ``cn_types=None``).
        """
        keep = None if cn_types is None else set(cn_types)

        def types_of(nl):
            s = set()
            for n in nl:
                vals = n.connectors['type']
                if keep is not None:
                    vals = vals[vals.isin(keep)]
                s.update(vals.unique().tolist())
            return s

        type_map = {}
        if by_type:
            all_types = sorted(types_of(query) | types_of(target), key=str)
            type_map = {t: i for i, t in enumerate(all_types)}

        def encode(nl):
            out = []
            for n in nl:
                cn = n.connectors
                if keep is not None:
                    cn = cn[cn['type'].isin(keep)]
                pts = cn[['x', 'y', 'z']].to_numpy(dtype=np.float64)
                if by_type:
                    codes = cn['type'].map(type_map).to_numpy(dtype=np.float64)
                    pts = np.hstack([pts, codes.reshape(-1, 1)])
                out.append(np.ascontiguousarray(pts, dtype=np.float64))
            return out

        return encode(query), encode(target)

    def synblast(self, query, target, *, by_type, cn_types, scores, normalized,
                 smat, n_cores, progress):
        from ... import utils
        q_arr, t_arr = self._encode_connectors(query, target, by_type, cn_types)
        symmetry = None if scores in (None, 'forward') else scores
        M = utils.fastcore.synblast(
            q_arr, t_arr,
            by_type=by_type,
            cn_types=None,  # already applied in _encode_connectors
            smat=self._convert_smat(smat, use_alpha=False),
            normalize=normalized,
            symmetry=symmetry,
            n_cores=n_cores,
            precision=64,  # match the built-in SynBlaster's float64 scores
            progress=progress,
        )
        out = pd.DataFrame(M, index=query.id, columns=target.id)
        out.index.name = 'query'
        out.columns.name = 'target'
        return out

    def nblast_allbyall(self, x, *, normalized, use_alpha, smat, limit_dist,
                        approx_nn, precision, n_cores, progress, smat_kwargs):
        from ... import utils
        M = utils.fastcore.nblast_allbyall(
            x,
            smat=self._convert_smat(smat, use_alpha),
            normalize=normalized,
            symmetry=None,  # navis' all-by-all always returns forward scores
            use_alpha=use_alpha,
            limit_dist=limit_dist,
            n_cores=n_cores,
            precision=precision,
            progress=progress,
        )
        out = pd.DataFrame(M, index=x.id, columns=x.id)
        out.index.name = 'query'
        out.columns.name = 'target'
        return out

navis.nbl.backends.NblastBackend

Base class for NBLAST backends.

A backend implements one or more of the operations listed in :data:OPERATIONS as methods of the same name. It does not need to implement all of them.

ATTRIBUTE	DESCRIPTION
name	Unique name used to select this backend (e.g. via the backend parameter or navis.config.default_nblast_backend). TYPE: str
priority	Higher-priority backends are preferred when backend="auto". The built-in backend uses 0; faster backends should use a higher value. TYPE: int

Source code in navis/nbl/backends/base.py

class NblastBackend(ABC):
    """Base class for NBLAST backends.

    A backend implements one or more of the operations listed in
    :data:`OPERATIONS` as methods of the same name. It does not need to
    implement all of them.

    Attributes
    ----------
    name :      str
                Unique name used to select this backend (e.g. via the
                ``backend`` parameter or ``navis.config.default_nblast_backend``).
    priority :  int
                Higher-priority backends are preferred when ``backend="auto"``.
                The built-in backend uses ``0``; faster backends should use a
                higher value.

    """

    name: str = "base"
    priority: int = 0

    def __repr__(self):
        return f"<NblastBackend '{self.name}' (priority={self.priority})>"

    def available(self) -> bool:
        """Whether this backend's dependencies are importable.

        Backends whose optional dependencies are missing should return False
        so they are skipped during ``backend="auto"`` resolution.
        """
        return True

    def implements(self, operation: str) -> bool:
        """Whether this backend implements `operation` at all."""
        return callable(getattr(self, operation, None))

    def unsupported(self, operation: str, **params) -> list:
        """Return reasons why this backend cannot run `operation` with `params`.

        An empty list means the request is supported. Non-empty entries are
        human-readable strings explaining the limitation (used both to skip a
        backend during ``"auto"`` resolution and to build a helpful error
        message when a backend was explicitly requested).

        The default implementation only checks whether the operation is
        implemented at all; subclasses should extend it to reject unsupported
        parameter combinations (e.g. ``approx_nn=True``).
        """
        if not self.implements(operation):
            return [f"backend '{self.name}' does not implement '{operation}'"]
        return []

available

Whether this backend's dependencies are importable.

Backends whose optional dependencies are missing should return False so they are skipped during backend="auto" resolution.

Source code in navis/nbl/backends/base.py

def available(self) -> bool:
    """Whether this backend's dependencies are importable.

    Backends whose optional dependencies are missing should return False
    so they are skipped during ``backend="auto"`` resolution.
    """
    return True

implements

Whether this backend implements operation at all.

Source code in navis/nbl/backends/base.py

def implements(self, operation: str) -> bool:
    """Whether this backend implements `operation` at all."""
    return callable(getattr(self, operation, None))

unsupported

Return reasons why this backend cannot run operation with params.

An empty list means the request is supported. Non-empty entries are human-readable strings explaining the limitation (used both to skip a backend during "auto" resolution and to build a helpful error message when a backend was explicitly requested).

The default implementation only checks whether the operation is implemented at all; subclasses should extend it to reject unsupported parameter combinations (e.g. approx_nn=True).

Source code in navis/nbl/backends/base.py

def unsupported(self, operation: str, **params) -> list:
    """Return reasons why this backend cannot run `operation` with `params`.

    An empty list means the request is supported. Non-empty entries are
    human-readable strings explaining the limitation (used both to skip a
    backend during ``"auto"`` resolution and to build a helpful error
    message when a backend was explicitly requested).

    The default implementation only checks whether the operation is
    implemented at all; subclasses should extend it to reject unsupported
    parameter combinations (e.g. ``approx_nn=True``).
    """
    if not self.implements(operation):
        return [f"backend '{self.name}' does not implement '{operation}'"]
    return []

navis.nbl.backends.available_backends

List backend instances whose dependencies are importable.

Source code in navis/nbl/backends/base.py

def available_backends() -> list:
    """List backend instances whose dependencies are importable."""
    return [b for b in _BACKENDS.values() if b.available()]

navis.nbl.backends.get_backend

Get a registered backend by name.

Source code in navis/nbl/backends/base.py

def get_backend(name: str) -> NblastBackend:
    """Get a registered backend by name."""
    if name not in _BACKENDS:
        raise ValueError(f"Unknown NBLAST backend '{name}'. "
                         f"Available: {list_backends()}")
    return _BACKENDS[name]

navis.nbl.backends.list_backends

List names of all registered backends.

Source code in navis/nbl/backends/base.py

def list_backends() -> list:
    """List names of all registered backends."""
    return list(_BACKENDS)

navis.nbl.backends.register_backend

Register an NBLAST backend.

PARAMETER	DESCRIPTION
backend	The backend instance to register. TYPE: NblastBackend
name	Name to register under. Defaults to ``backend.name``. TYPE: str DEFAULT: None

Source code in navis/nbl/backends/base.py

def register_backend(backend: NblastBackend, name: str = None):
    """Register an NBLAST backend.

    Parameters
    ----------
    backend :   NblastBackend
                The backend instance to register.
    name :      str, optional
                Name to register under. Defaults to ``backend.name``.

    """
    if not isinstance(backend, NblastBackend):
        raise TypeError(f"Expected NblastBackend, got {type(backend)}")
    _BACKENDS[name or backend.name] = backend

navis.nbl.backends.resolve_backend

Select a backend for operation given params.

PARAMETER	DESCRIPTION
operation	The operation to run (e.g. "nblast"). Must be one of :data:`OPERATIONS`. TYPE: str
backend	Either "auto" (pick the highest-priority available backend that supports the request, falling back to "builtin"), or the name of a specific backend, or a backend instance. TYPE: str | NblastBackend DEFAULT: 'auto'
**params	The call parameters (e.g. ``approx_nn``, ``scores``, ``smat``). Used to check whether a backend supports the request. DEFAULT: {}

RETURNS	DESCRIPTION
NblastBackend

Source code in navis/nbl/backends/base.py

def resolve_backend(operation: str, backend="auto", **params) -> NblastBackend:
    """Select a backend for `operation` given `params`.

    Parameters
    ----------
    operation : str
                The operation to run (e.g. "nblast"). Must be one of
                :data:`OPERATIONS`.
    backend :   str | NblastBackend
                Either "auto" (pick the highest-priority available backend that
                supports the request, falling back to "builtin"), or the name of
                a specific backend, or a backend instance.
    **params
                The call parameters (e.g. ``approx_nn``, ``scores``, ``smat``).
                Used to check whether a backend supports the request.

    Returns
    -------
    NblastBackend

    """
    # Allow passing a backend instance directly
    if isinstance(backend, NblastBackend):
        return backend

    if backend in (None, "auto"):
        # Highest priority first; builtin (priority 0) is the fallback
        rejected = {}
        for be in sorted(available_backends(), key=lambda b: -b.priority):
            reasons = be.unsupported(operation, **params)
            if not reasons:
                return be
            rejected[be.name] = reasons

        # Nothing supported the request - fall back to builtin so we can raise
        # a meaningful error (or run it, if builtin does support it).
        builtin = get_backend("builtin")
        if builtin.implements(operation):
            return builtin

        # Not even builtin implements this, i.e. the operation is exclusive to
        # some other backend (currently only `nblast_knn`, which needs
        # navis-fastcore) and that backend is missing, too old or was ruled out
        # by the parameters. Say which, instead of returning a backend that
        # would `AttributeError` on the very next line.
        missing = [b.name for b in _BACKENDS.values()
                   if not b.available() and b.implements(operation)]
        msg = f"No available NBLAST backend can run '{operation}'."
        if missing:
            msg += (f" Backend(s) {missing} implement it but are not available"
                    " - their optional dependencies are not installed"
                    " (`pip install -U navis-fastcore`).")
        if rejected:
            detail = '; '.join(f"{n}: {', '.join(r)}" for n, r in rejected.items())
            msg += f" Rejected: {detail}."
        raise ValueError(msg)

    be = get_backend(backend)

    if not be.available():
        raise ValueError(
            f"NBLAST backend '{be.name}' is not available - its optional "
            "dependencies are not installed."
        )

    reasons = be.unsupported(operation, **params)
    if reasons:
        raise ValueError(
            f"NBLAST backend '{be.name}' cannot run '{operation}' with the "
            f"given parameters: {'; '.join(reasons)}."
        )

    return be