Skeleton cache (connectomekg.skeleton_cache)

connectomekg.skeleton_cache

Simplified skeletons and somas, cached per dataset from one pass over the SWC download.

The Codex skeleton download is 31 GB of SWC text, one file per neuron (:func:connectomekg.skeletons.skeleton_path). Two things in it are worth keeping and the rest is not:

• The soma. A neuron node's x/y/z is FlyWire's marked point, an anchor that can sit tens of microns from the cell body. The soma is the Label 1 row of the skeleton, and it is what every fly atlas draws. :func:write_somas copies it into the neuron node's metadata as soma_x/soma_y/soma_z with a has_soma flag, so the graph answers for it without the download and so does every query over it.
• The shape, simplified. At :data:DEFAULT_CACHE_STEP a skeleton keeps about 29 % of its points and all of its topology (see :func:connectomekg.skeletons.segments), which is what the 3-D circuit view draws anyway. The cache is a directory of Parquet shards holding every neuron's kept points -- 212 M of them in 2.5 GB on FAFB v783, against the download's 31 GB -- so a render needs neither the download nor an SWC parse.

Both come out of the same pass, so :func:build_skeleton_cache does both and connkg skeletons writes both. The pass reads every file once: 19m 30s on one core for FAFB v783, or 2m 39s across twelve, since parsing SWC is pure Python and splits cleanly across processes. Nothing else in the module needs the download afterwards.

The cache carries each kept point's radius, which is what lets the circuit view taper a drawn neuron instead of piping it at one width. Its labels are reconstructed, not stored: a cached skeleton's only label is the soma's. Everything else about it round-trips, so :func:connectomekg.skeletons.segments, :func:connectomekg.skeletons.polylines and :func:connectomekg.skeletons.soma read a cached skeleton exactly as they read a parsed one.

A format-1 cache, written before the radius column existed, still loads: its radii read back as zeros and its skeletons draw at a constant width, which is what they did when it was written. Re-running connkg skeletons is what adds the taper.

CacheReport(n_cached, n_missing, n_unreadable, n_soma, n_points, step, bytes_written, n_shards=1) dataclass

What one :func:build_skeleton_cache pass found.

Parameters:

Name	Type	Description	Default
n_cached	int	Neurons written to the cache.	required
n_missing	int	Requested neurons with no SWC file.	required
n_unreadable	int	Neurons whose SWC file failed to parse.	required
n_soma	int	Cached neurons with a real Label 1 soma row.	required
n_points	int	Kept points across every cached skeleton.	required
step	int	The stride the cache was built at.	required
bytes_written	int	Total size of the cache's Parquet shards.	required
n_shards	int	Shards written, one per worker.	1

build_skeleton_cache(data_dir, root_ids, dest, *, step=DEFAULT_CACHE_STEP, jobs=1, progress=None)

Read every neuron's SWC file once, writing the cache and collecting somas.

Parsing an SWC file is pure Python and holds the GIL, so the pass is bound by one core until it is split across processes: jobs of them, each taking a contiguous range of root ids and writing its own shard, after proteusPy's DisulfideExtractor_mp. The cache is the directory of those shards, which :func:load_cached_skeletons reads as one.

Parameters:

Name	Type	Description	Default
data_dir	str | Path	Root of the skeleton download (default layout: fafb_v783).	required
root_ids	Iterable[int]	Neuron root ids to read, in any order; the cache is written in ascending root-id order regardless.	required
dest	str | Path	The cache directory to write, usually <dataset dir>/.connectomekg/skeletons.	required
step	int	Simplification stride, at least 1.	DEFAULT_CACHE_STEP
jobs	int	Worker processes; 1 runs in this process and starts no pool. Above the number of neurons it is reduced to that.	1
progress	Callable[[str], None] | None	Called with a short message as each shard finishes.	None

Returns:

Type	Description
tuple[CacheReport, dict[int, tuple[ndarray, bool]]]	(report, somas) -- somas maps root id to (point_nm, is_soma) as :func:connectomekg.skeletons.soma defines it.

Raises:

Type	Description
ValueError	If step or jobs is below 1.

Source code in src/connectomekg/skeleton_cache.py

def build_skeleton_cache(
    data_dir: str | Path,
    root_ids: Iterable[int],
    dest: str | Path,
    *,
    step: int = DEFAULT_CACHE_STEP,
    jobs: int = 1,
    progress: Callable[[str], None] | None = None,
) -> tuple[CacheReport, dict[int, tuple[np.ndarray, bool]]]:
    """Read every neuron's SWC file once, writing the cache and collecting somas.

    Parsing an SWC file is pure Python and holds the GIL, so the pass is bound
    by one core until it is split across processes: ``jobs`` of them, each
    taking a contiguous range of root ids and writing its own shard, after
    ``proteusPy``'s ``DisulfideExtractor_mp``. The cache is the directory of
    those shards, which :func:`load_cached_skeletons` reads as one.

    :param data_dir: Root of the skeleton download (default layout: ``fafb_v783``).
    :param root_ids: Neuron root ids to read, in any order; the cache is
        written in ascending root-id order regardless.
    :param dest: The cache directory to write, usually
        ``<dataset dir>/.connectomekg/skeletons``.
    :param step: Simplification stride, at least 1.
    :param jobs: Worker processes; 1 runs in this process and starts no pool.
        Above the number of neurons it is reduced to that.
    :param progress: Called with a short message as each shard finishes.
    :return: ``(report, somas)`` -- ``somas`` maps root id to
        ``(point_nm, is_soma)`` as :func:`connectomekg.skeletons.soma` defines it.
    :raises ValueError: If ``step`` or ``jobs`` is below 1.
    """
    step = int(step)
    if step < 1:
        raise ValueError(f"step must be at least 1, got {step}")
    if int(jobs) < 1:
        raise ValueError(f"jobs must be at least 1, got {jobs}")
    dest = Path(dest)

    ids = sorted({int(r) for r in root_ids})
    chunks = _chunk(ids, jobs)
    # Built beside the destination and renamed over it at the end, so an
    # interrupted pass never leaves a half-written cache where the renderer
    # would find one, and a rebuild does not read its own leftovers.
    partial = dest.with_name(dest.name + ".part")
    shutil.rmtree(partial, ignore_errors=True)
    partial.mkdir(parents=True)

    totals = _ShardResult(0, 0, 0, 0, 0, {})
    finished = False
    try:
        work = [
            (str(data_dir), chunk, str(partial / f"part-{k:04d}.parquet"), step)
            for k, chunk in enumerate(chunks)
        ]
        done = 0
        for result in _run_shards(work, jobs):
            done += 1
            totals = _ShardResult(
                totals.n_cached + result.n_cached,
                totals.n_missing + result.n_missing,
                totals.n_unreadable + result.n_unreadable,
                totals.n_soma + result.n_soma,
                totals.n_points + result.n_points,
                {**totals.somas, **result.somas},
            )
            if progress is not None:
                progress(
                    f"shard {done} of {len(work)} done "
                    f"({totals.n_cached} of {len(ids)} skeletons cached)"
                )
        if dest.exists():
            shutil.rmtree(dest)
        partial.replace(dest)
        finished = True
    finally:
        # Not `except Exception`: Ctrl-C is the likeliest way a long read ends
        # early, and it must not leave gigabytes of shards behind either.
        if not finished:
            shutil.rmtree(partial, ignore_errors=True)

    return (
        CacheReport(
            n_cached=totals.n_cached,
            n_missing=totals.n_missing,
            n_unreadable=totals.n_unreadable,
            n_soma=totals.n_soma,
            n_points=totals.n_points,
            step=step,
            bytes_written=sum(f.stat().st_size for f in dest.glob("*.parquet")),
            n_shards=len(chunks),
        ),
        totals.somas,
    )

cache_info(path)

A cache's stride and neuron count, read from its Parquet metadata alone.

Every shard records the same stride, so the stride comes from the first of them; the neuron count is summed across all of them, since each shard counts only its own.

Parameters:

Name	Type	Description	Default
path	str | Path	The cache directory.	required

Returns:

Type	Description
tuple[int, int]	(step, n_neurons).

Raises:

Type	Description
ValueError	If the directory holds no shard this module wrote.

Source code in src/connectomekg/skeleton_cache.py

def cache_info(path: str | Path) -> tuple[int, int]:
    """A cache's stride and neuron count, read from its Parquet metadata alone.

    Every shard records the same stride, so the stride comes from the first of
    them; the neuron count is summed across all of them, since each shard
    counts only its own.

    :param path: The cache directory.
    :return: ``(step, n_neurons)``.
    :raises ValueError: If the directory holds no shard this module wrote.
    """
    shards = sorted(Path(path).glob("*.parquet"))
    if not shards:
        raise ValueError(f"{path}: no {_FORMAT} shards here")
    step = 0
    n_neurons = 0
    for shard in shards:
        # Key-value metadata added at close lands in the file footer, not on
        # the Arrow schema -- `schema_arrow.metadata` reads back None for it.
        metadata = pq.ParquetFile(shard).metadata.metadata or {}

        def value(key: str, metadata: dict = metadata) -> str:
            # Only this module's own keys are decoded: pyarrow's own
            # ARROW:schema entry beside them is not text.
            return (metadata.get(key.encode()) or b"").decode()

        if value("connectomekg.format") not in _READABLE_FORMATS:
            raise ValueError(f"{shard}: not a {_FORMAT} skeleton cache")
        step = int(value("connectomekg.step"))
        n_neurons += int(value("connectomekg.n_neurons"))
    return step, n_neurons

effective_step(requested, cache_step)

The stride to draw a cached skeleton with, given the one the caller asked for.

A cached skeleton has already been simplified at cache_step, so passing the caller's stride straight to :func:connectomekg.skeletons.segments would simplify it a second time and thin it to their product. Dividing instead makes the drawn stride approximately the requested one, and a request at or below cache_step draws the cache as it stands -- the cache's own stride is the finest detail it holds.

Parameters:

Name	Type	Description	Default
requested	int	The caller's skeleton_step.	required
cache_step	int	The stride the cache was built at.	required

Returns:

Type	Description
int	The stride to pass to :func:connectomekg.skeletons.segments, at least 1.

Source code in src/connectomekg/skeleton_cache.py

def effective_step(requested: int, cache_step: int) -> int:
    """The stride to draw a *cached* skeleton with, given the one the caller asked for.

    A cached skeleton has already been simplified at ``cache_step``, so passing
    the caller's stride straight to :func:`connectomekg.skeletons.segments`
    would simplify it a second time and thin it to their product. Dividing
    instead makes the drawn stride approximately the requested one, and a
    request at or below ``cache_step`` draws the cache as it stands -- the
    cache's own stride is the finest detail it holds.

    :param requested: The caller's ``skeleton_step``.
    :param cache_step: The stride the cache was built at.
    :return: The stride to pass to :func:`connectomekg.skeletons.segments`, at least 1.
    """
    return max(1, int(requested) // max(1, int(cache_step)))

load_cached_skeletons(path, root_ids)

Read whichever of root_ids the cache holds, reporting the rest.

Mirrors :func:connectomekg.skeletons.load_skeletons, so a caller can use either source.

The cache is a directory of shards; pyarrow reads it as one table, and because each shard holds a disjoint span of root ids their row-group statistics prune a filtered read down to the few shards that can match.

Parameters:

Name	Type	Description	Default
path	str | Path	The cache directory.	required
root_ids	Iterable[int]	Neuron root ids to read.	required

Returns:

Type	Description
tuple[dict[int, Skeleton], list[int]]	(skeletons, missing) -- the skeletons the cache holds, already simplified at its own stride, and the root ids it does not.

Source code in src/connectomekg/skeleton_cache.py

def load_cached_skeletons(
    path: str | Path, root_ids: Iterable[int]
) -> tuple[dict[int, Skeleton], list[int]]:
    """Read whichever of ``root_ids`` the cache holds, reporting the rest.

    Mirrors :func:`connectomekg.skeletons.load_skeletons`, so a caller can use
    either source.

    The cache is a directory of shards; pyarrow reads it as one table, and
    because each shard holds a disjoint span of root ids their row-group
    statistics prune a filtered read down to the few shards that can match.

    :param path: The cache directory.
    :param root_ids: Neuron root ids to read.
    :return: ``(skeletons, missing)`` -- the skeletons the cache holds, already
        simplified at its own stride, and the root ids it does not.
    """
    wanted = [int(r) for r in root_ids]
    if not wanted or not any(Path(path).glob("*.parquet")):
        return {}, wanted
    table = pq.read_table(
        path, filters=[("root_id", "in", set(wanted))], schema=_SCHEMA, use_threads=True
    )
    loaded: dict[int, Skeleton] = {}
    for row in table.to_pylist():
        x = np.asarray(row["x"], dtype=np.float64)
        y = np.asarray(row["y"], dtype=np.float64)
        z = np.asarray(row["z"], dtype=np.float64)
        labels = np.zeros(len(x), dtype=np.int64)
        if row["soma_index"] >= 0:
            labels[row["soma_index"]] = 1
        loaded[int(row["root_id"])] = Skeleton(
            root_id=int(row["root_id"]),
            points=np.stack([x, y, z], axis=1),
            radius=(
                np.zeros(len(x), dtype=np.float64)
                if row["radius"] is None
                else np.asarray(row["radius"], dtype=np.float64)
            ),
            labels=labels,
            parent=np.asarray(row["parent"], dtype=np.int64),
        )
    return loaded, [r for r in wanted if r not in loaded]

skeleton_cache_path(db_path)

Where a dataset's skeleton cache lives: beside its graph.

Parameters:

Name	Type	Description	Default
db_path	str | Path	The dataset's graph.sqlite.	required

Returns:

Type	Description
Path	<dataset dir>/.connectomekg/skeletons/, which may not exist.

Source code in src/connectomekg/skeleton_cache.py

def skeleton_cache_path(db_path: str | Path) -> Path:
    """Where a dataset's skeleton cache lives: beside its graph.

    :param db_path: The dataset's ``graph.sqlite``.
    :return: ``<dataset dir>/.connectomekg/skeletons/``, which may not exist.
    """
    return Path(db_path).parent / SKELETON_CACHE

write_somas(store, somas)

Copy somas into their neuron nodes' metadata, in one transaction.

Adds soma_x, soma_y, soma_z and has_soma to each neuron node that has a soma here, leaving the marked point (x/y/z) and every other key as they were. Running again overwrites the four keys, so a second pass over a corrected download is safe.

Parameters:

Name	Type	Description	Default
store	GraphStore	The open store to update.	required
somas	dict[int, tuple[ndarray, bool]]	{root_id: (point_nm, is_soma)}, as :func:build_skeleton_cache returns.	required

Returns:

Type	Description
int	Neuron nodes updated.

Source code in src/connectomekg/skeleton_cache.py

def write_somas(store: GraphStore, somas: dict[int, tuple[np.ndarray, bool]]) -> int:
    """Copy somas into their neuron nodes' metadata, in one transaction.

    Adds ``soma_x``, ``soma_y``, ``soma_z`` and ``has_soma`` to each neuron
    node that has a soma here, leaving the marked point (``x``/``y``/``z``) and
    every other key as they were. Running again overwrites the four keys, so a
    second pass over a corrected download is safe.

    :param store: The open store to update.
    :param somas: ``{root_id: (point_nm, is_soma)}``, as
        :func:`build_skeleton_cache` returns.
    :return: Neuron nodes updated.
    """
    con = store.con
    by_root = {
        int(root_id): node_id
        for node_id, root_id in con.execute(
            "SELECT id, json_extract(metadata,'$.root_id') FROM nodes WHERE kind = 'neuron'"
        )
        if root_id is not None
    }
    updates = [
        (
            json.dumps(
                {
                    "soma_x": float(point[0]),
                    "soma_y": float(point[1]),
                    "soma_z": float(point[2]),
                    "has_soma": bool(is_soma),
                }
            ),
            by_root[root_id],
        )
        for root_id, (point, is_soma) in somas.items()
        if root_id in by_root
    ]
    with con:
        con.executemany("UPDATE nodes SET metadata = json_patch(metadata, ?) WHERE id = ?", updates)
    return len(updates)