Namespaces¶

@requires(Cap.ANNOTATIONS)
def get(
    self,
    x=None,
    *,
    fields: dict | None = None,
    source: str | None = None,
    raw: bool = False,
    version=None,
    units: str = "nm",
    verbose: bool = True,
    **filters,
) -> pd.DataFrame:
    """Annotations, with canonical columns added and every raw column kept.

    Canonical: ``id, type, side, class, nt, status, soma_x/y/z``. ``type`` is
    coalesced from ``spec.fields["type"]`` in priority order; ``side`` is
    mapped to ``left``/``right``/``center``. Pass ``raw=True`` for the
    untouched backend frame.

    ``verbose`` (default ``True``) prints one line saying which source is being
    read and whether it came from the remote server, the local cache, or a
    cache refresh. Set it ``False`` to silence that.
    """
    ds = self._ds
    v = ds._resolve_version_arg(version)
    src = (
        ds.spec.annotation_source(source)
        if source is not None
        else ds._annotation_source
    )
    if src is None:
        raise CapabilityError(
            f"{ds.label} has no annotation source configured."
        )

    table = self._table(src, v, verbose=verbose)

    if raw:
        return table.copy()

    ann = schemas.normalize_annotations(
        table, ds, id_column=src.id_column, fields=fields, version=v, units=units
    )

    if x is not None or filters:
        crit = to_criteria(x) if x is not None else None
        ids = ds.ids(crit, version=version) if crit is not None else None
        if filters:
            from .criteria import NeuronCriteria, resolve_criteria

            extra = resolve_criteria(NeuronCriteria(**filters), ds, version=version)
            ids = extra if ids is None else np.intersect1d(ids, extra)
        ann = ann[ann["id"].isin(ids)].reset_index(drop=True)

    return ann

@requires(Cap.ANNOTATIONS)
def search(self, term: str, *, version=None, regex: bool = True) -> pd.DataFrame:
    """Rows whose `type`, `class` or `instance` matches ``term``."""
    ann = self.get(version=version, verbose=False)
    cols = [c for c in ("type", "class", "instance") if c in ann.columns]
    mask = pd.Series(False, index=ann.index)
    for col in cols:
        s = ann[col].astype("string")
        mask |= (
            s.str.contains(term, case=False, regex=regex, na=False)
            if regex
            else s.str.lower().eq(term.lower())
        )
    return ann[mask].reset_index(drop=True)

@requires(Cap.ANNOTATIONS)
def ids(self, *, version=None) -> np.ndarray:
    """Every annotated neuron in the dataset."""
    got = self.get(version=version, verbose=False)
    return np.unique(got["id"].to_numpy(dtype="int64"))

@requires(Cap.CONNECTIVITY)
def edges(
    self,
    x,
    *,
    upstream: bool = True,
    downstream: bool = True,
    by_roi: bool = False,
    rois=UNSET,
    min_weight: int = 1,
    version=None,
    cache: bool = False,
    extra: bool = False,
    autapses: bool = False,
) -> pd.DataFrame:
    """Edges to/from the given neurons as ``pre, post, weight``.

    A closed schema: identical columns and dtypes on every backend, so
    ``bodyId_pre`` and ``pre_pt_root_id`` both come back as ``pre``, and
    FlyWire's edge view does not smuggle in seventeen extra columns that
    hemibrain has never heard of. ``extra=True`` keeps whatever else the
    backend happened to return.

    ``autapses`` (pre == post) default to **off**. They are nearly always
    segmentation errors, neuPrint omits them entirely, and caveclient's own
    ``synapse_query`` drops them by default - so leaving them in would make the
    same neuron have different connectivity depending on which backend you
    asked. Pass ``autapses=True`` if you want them; you will get them from
    every backend that has them.
    """
    ds = self._ds
    v = ds._resolve_version_arg(version)
    ids = ds.ids(x, version=version)
    if not len(ids):
        return schemas.normalize_edges(
            pd.DataFrame(columns=["pre", "post", "weight"]), ds,
            colmap={}, version=v, extra=extra,
        )

    by_roi = ds._resolve_opt("by_roi", by_roi if by_roi else UNSET,
                             Cap.ROI_CONN, default=False, neutral=False)
    rois = ds._resolve_opt("rois", rois, Cap.ROI_CONN, default=None)

    entry = None
    if cache:
        entry = _cache.CacheEntry(
            ds.name, v, "edges",
            ids.tobytes(), upstream, downstream, by_roi, rois, min_weight,
        )
        if entry.exists():
            return schemas.normalize_edges(
                entry.read(), ds, colmap=ds._edge_colmap, version=v, extra=extra
            )

    frames = []
    if downstream:
        frames.append(
            ds._fetch_edges(pre=ids, post=None, version=v,
                            by_roi=by_roi, min_weight=min_weight, rois=rois)
        )
    if upstream:
        frames.append(
            ds._fetch_edges(pre=None, post=ids, version=v,
                            by_roi=by_roi, min_weight=min_weight, rois=rois)
        )
    if not frames:
        raise ValueError("One of `upstream` or `downstream` must be True.")

    raw = pd.concat(frames, ignore_index=True) if len(frames) > 1 else frames[0]

    if len(frames) > 1:
        # A query neuron connecting to another query neuron appears in both
        # the upstream and the downstream fetch.
        keys = [c for c in raw.columns if c in set(ds._edge_colmap.values())]
        dedup_on = [
            ds._edge_colmap[k] for k in ("pre", "post") if ds._edge_colmap.get(k) in keys
        ]
        if by_roi and ds._edge_colmap.get("roi") in raw.columns:
            dedup_on.append(ds._edge_colmap["roi"])
        if dedup_on:
            raw = raw.drop_duplicates(subset=dedup_on, ignore_index=True)

    if entry is not None:
        entry.write(raw)

    out = schemas.normalize_edges(
        raw, ds, colmap=ds._edge_colmap, version=v, extra=extra
    )
    if not autapses:
        prov = out.attrs["connecto"]
        out = out[out["pre"] != out["post"]].reset_index(drop=True)
        out.attrs["connecto"] = prov  # boolean masking doesn't reliably keep attrs
    return out

@requires(Cap.CONNECTIVITY)
def adjacency(
    self,
    sources,
    targets=None,
    *,
    version=None,
    min_weight: int = 1,
) -> pd.DataFrame:
    """Weight matrix, rows = sources, columns = targets."""
    ds = self._ds
    src = ds.ids(sources, version=version)
    tgt = src if targets is None else ds.ids(targets, version=version)

    edges = self.edges(
        src, upstream=False, downstream=True, min_weight=min_weight, version=version
    )
    edges = edges[edges["post"].isin(tgt)]

    adj = (
        edges.pivot_table(
            index="pre", columns="post", values="weight",
            aggfunc="sum", fill_value=0,
        )
        .reindex(index=src, columns=tgt, fill_value=0)
        .astype("int32")
    )
    adj.index.name, adj.columns.name = "pre", "post"
    return schemas.stamp(adj, ds, query="adjacency", version=ds._resolve_version_arg(version))

@requires(Cap.SYNAPSES)
def synapses(
    self,
    x,
    *,
    pre: bool = True,
    post: bool = True,
    min_score=UNSET,
    transmitters=UNSET,
    rois=UNSET,
    version=None,
    units: str = "nm",
) -> pd.DataFrame:
    """Individual synapses, positions in nanometres.

    ``min_score`` and ``transmitters`` raise on datasets that don't have them,
    rather than being quietly ignored.
    """
    ds = self._ds
    v = ds._resolve_version_arg(version)
    ids = ds.ids(x, version=version)

    min_score = ds._resolve_opt("min_score", min_score, Cap.SYNAPSE_SCORES, default=None)
    transmitters = ds._resolve_opt(
        "transmitters", transmitters, Cap.NT_PER_SYNAPSE, default=False, neutral=False
    )
    rois = ds._resolve_opt("rois", rois, Cap.ROI_CONN, default=None)

    frames = []
    if pre:
        frames.append(ds._fetch_synapses(
            pre=ids, post=None, version=v,
            min_score=min_score, transmitters=transmitters, rois=rois))
    if post:
        frames.append(ds._fetch_synapses(
            pre=None, post=ids, version=v,
            min_score=min_score, transmitters=transmitters, rois=rois))
    if not frames:
        raise ValueError("One of `pre` or `post` must be True.")

    raw = pd.concat(frames, ignore_index=True) if len(frames) > 1 else frames[0]
    idcol = ds._synapse_colmap.get("id")
    if len(frames) > 1 and idcol in raw.columns:
        raw = raw.drop_duplicates(subset=[idcol], ignore_index=True)

    return schemas.normalize_synapses(
        raw, ds, colmap=ds._synapse_colmap, version=v, units=units
    )

@requires(Cap.CONNECTIVITY)
def synapse_counts(self, x, *, by_roi: bool = False, version=None) -> pd.DataFrame:
    """Pre- and post-synapse counts per neuron."""
    ds = self._ds
    ids = ds.ids(x, version=version)

    out = self.edges(ids, version=version, by_roi=by_roi)
    group = ["roi"] if by_roi and "roi" in out.columns else []

    up = (
        out[out["post"].isin(ids)]
        .groupby(["post"] + group, as_index=False)["weight"].sum()
        .rename(columns={"post": "id", "weight": "post"})
    )
    down = (
        out[out["pre"].isin(ids)]
        .groupby(["pre"] + group, as_index=False)["weight"].sum()
        .rename(columns={"pre": "id", "weight": "pre"})
    )
    counts = pd.merge(down, up, on=["id"] + group, how="outer").fillna(0)
    for col in ("pre", "post"):
        counts[col] = counts[col].astype("int32")
    return counts[["id"] + group + ["pre", "post"]]

@requires(Cap.NT_PER_SYNAPSE)
def transmitters(self, x, *, version=None) -> pd.DataFrame:
    """Predicted transmitter per neuron, from its presynapses."""
    syn = self.synapses(x, pre=True, post=False, transmitters=True, version=version)
    if "nt" not in syn.columns:
        raise CapabilityError(f"{self._ds.label} returned no transmitter column.")
    counts = (
        syn.groupby(["pre", "nt"], observed=True).size().rename("n").reset_index()
    )
    top = counts.sort_values("n", ascending=False).drop_duplicates("pre")
    total = counts.groupby("pre", observed=True)["n"].sum().rename("total")
    top = top.merge(total, on="pre")
    top["confidence"] = top["n"] / top["total"]
    return (
        top.rename(columns={"pre": "id"})[["id", "nt", "confidence"]]
        .sort_values("id")
        .reset_index(drop=True)
    )

@requires(Cap.SKELETONS)
def get(self, x, *, output: str = "navis", version=None, **opts):
    """Skeletons for the given neurons -> ``navis.NeuronList``."""
    from . import volume

    ds = self._ds
    v = ds._resolve_version_arg(version)
    ids = ds.ids(x, version=version)

    # A published precomputed skeleton is not backend data: it comes out of a
    # plain HTTPS bucket, canonical and in nanometres, whichever door fetched it.
    # So it is normalised as itself rather than through the backend's column map
    # and voxel scaling. (See PRECOMPUTED_SKELETON_COLMAP for what that cost.)
    precomputed = ds._skeleton_source(v) is not None
    colmap = volume.PRECOMPUTED_SKELETON_COLMAP if precomputed else ds._skeleton_colmap
    source_units = volume.PRECOMPUTED_SKELETON_UNITS if precomputed else None

    out = []
    for nid, raw in ds._fetch_skeletons(ids, v, **opts):
        nodes = schemas.normalize_skeleton(
            raw, ds, colmap=colmap, source_units=source_units
        )
        out.append((nid, nodes))

    if output == "raw":
        return {nid: nodes for nid, nodes in out}

    import navis

    neurons = []
    for nid, nodes in out:
        n = navis.TreeNeuron(nodes, id=nid, units="1 nm")
        n.name = str(nid)
        neurons.append(n)
    return navis.NeuronList(neurons)

@requires(Cap.SKELETONS)
def dotprops(self, x, *, k: int = 5, units: str = "nm", version=None):
    """Dotprops, for NBLAST. Built from the skeleton.

    ``units`` defaults to ``"nm"``, because everything connecto returns is in
    nanometres. **NBLAST is calibrated in microns.** Score it on nanometre dotprops
    and every pair collapses to the floor - two copies of the same cell type come
    back at -0.88 instead of +0.70 - which reads as "nothing is similar" rather than
    as an error. navis notices and logs a warning, but a log line is easy to miss.

    So when the next step is NBLAST, ask for microns::

        dp = ds.skeletons.dotprops(x, units="um")
        navis.nblast(dp, dp)

    If the dataset has an L2 cache, :meth:`~connecto.backends.cave.l2.L2.dotprops`
    gets there without building a skeleton first, and is much cheaper.
    """
    import navis

    if units not in ("nm", "um"):
        raise ValueError(f"`units` must be 'nm' or 'um', got {units!r}.")

    skels = self.get(x, version=version)
    if units == "um":
        skels = skels / 1000
    return navis.make_dotprops(skels, k=k)

@requires(Cap.VOXELS)
def get(
    self,
    x,
    *,
    scale: int | None = None,
    output: str = "navis",
    units: str = "voxel",
    version=None,
    progress: bool = True,
    verbose: bool = False,
    **opts,
):
    """Sparse volumes for the given neurons.

    Parameters
    ----------
    scale :     int, optional
                Pyramid level; each level halves resolution. Defaults to
                something that returns a usable neuron without reading the whole
                brain - **not** 0, which for a hemibrain neuron is 1.17 billion
                voxels. ``ds.voxels.scales()`` lists what is available.
    output :    "navis" | "raw" | "rle"
                ``navis`` gives ``VoxelNeuron``s; ``raw`` a dict of ``(N, 3)``
                arrays; ``rle`` a dict of ``(M, 4)`` runs, which is ~20x smaller
                and is what the wire already carries.
    units :     "voxel" | "nm"
                Coordinate space of ``raw`` output. ``navis`` output is always
                voxels *plus* the units metadata navis expects, so it plots and
                measures in nm regardless.
    """
    if units not in ("voxel", "nm"):
        raise ValueError(f"`units` must be 'voxel' or 'nm', got {units!r}.")
    if output not in ("navis", "raw", "rle"):
        raise ValueError(
            f"`output` must be 'navis', 'raw' or 'rle', got {output!r}."
        )

    from .. import voxels as _voxels

    ds = self._ds
    ds._resolve_version_arg(version)
    ids = ds.ids(x, version=version)
    if scale is None:
        scale = _voxels.default_scale(ds)

    from tqdm.auto import tqdm

    out = []
    for nid in tqdm(
        ids, desc="Voxels", disable=not progress or len(ids) < 2, leave=False
    ):
        runs, resolution, stats = _voxels.fetch(ds, int(nid), scale, **opts)
        if verbose:
            print(f"{nid}: {_voxels.rle.run_voxel_count(runs):,} voxels @ {resolution} nm | {stats}")
        out.append((int(nid), runs, resolution))

    if output == "rle":
        return {nid: runs for nid, runs, _ in out}

    if output == "raw":
        return {
            nid: (
                _voxels.to_nm(_voxels.rle.decode_runs(runs), res)
                if units == "nm"
                else _voxels.rle.decode_runs(runs)
            )
            for nid, runs, res in out
        }

    import navis

    neurons = []
    for nid, runs, res in out:
        # VoxelNeuron carries its own nm-per-voxel, so the coordinates stay in
        # the (compact, integral) voxel grid and navis still measures in nm.
        # Multiplying them out here would inflate the array and lose that.
        #
        # Units go in as a per-axis vector because these grids are routinely
        # anisotropic - aedes is 32x32x45, fish2 at scale 3 is 256x256x240 - and
        # a single scalar would be right about X and wrong about Z.
        n = navis.VoxelNeuron(
            _voxels.rle.decode_runs(runs),
            id=nid,
            units=(
                f"{res[0]} nm"
                if len(set(res)) == 1
                else [f"{r} nm" for r in res]
            ),
        )
        n.name = str(nid)
        neurons.append(n)
    return navis.NeuronList(neurons)

@requires(Cap.VOXELS)
def estimate(self, x, *, scale: int | None = None, version=None) -> pd.DataFrame:
    """What ``get()`` would cost, without doing it.

    Meaningful mainly on the CAVE route, where the answer is a dense read whose
    size the caller controls. The indexed routes just say so.
    """
    from .. import voxels as _voxels

    ds = self._ds
    # Validates the version even though the volume is read from the handle's own
    # source, so a bad `version=` raises here rather than being ignored.
    ds._resolve_version_arg(version)
    ids = ds.ids(x, version=version)
    if scale is None:
        scale = _voxels.default_scale(ds)
    return pd.DataFrame([_voxels.estimate(ds, int(i), scale) for i in ids])

@requires(Cap.VOXELS)
def scales(self) -> tuple[int, ...]:
    """Which scales this dataset serves.

    That a scale is *available* says nothing about whether it is *affordable* -
    every CAVE dataset lists scale 0 and none of them can deliver a whole neuron
    at it. Use :meth:`estimate` for the cost.
    """
    from .. import voxels as _voxels

    return _voxels.scales(self._ds)

@requires(Cap.ROIS)
def hierarchy(self):
    """ROI containment tree as a ``networkx.DiGraph``."""
    return self._ds._fetch_roi_hierarchy()