load

lab.connectome.load

The data files the project reads, and the readers for them.

The registry DATASETS lists every file or folder that fetch writes. status reports which are present and usable. The readers are pure: no network and no writes, except that figure_path creates config.FIGURES. Tables come back as polars frames. A missing file raises with the command that writes it.

Dataset dataclass

One data file or folder.

Attributes:

Name	Type	Description
name	str	Short name, as shown in the status report.
source	str	Where it comes from, such as "codex" or "cave".
path	Path	The file or folder, inside config.ROOT / "data".
fetch	str	The command that writes it.
check	callable	Takes path and returns None when it is usable, else what is wrong.
partial	(Path, optional)	A folder of pieces left by an interrupted fetch.

Source code in lab/connectome/load.py

@dataclass(frozen=True)
class Dataset:
    """One data file or folder.

    Attributes
    ----------
    name : str
        Short name, as shown in the status report.
    source : str
        Where it comes from, such as ``"codex"`` or ``"cave"``.
    path : pathlib.Path
        The file or folder, inside ``config.ROOT / "data"``.
    fetch : str
        The command that writes it.
    check : callable
        Takes `path` and returns ``None`` when it is usable, else what
        is wrong.
    partial : pathlib.Path, optional
        A folder of pieces left by an interrupted fetch.
    """

    name: str
    source: str
    path: Path
    fetch: str
    check: Callable[[Path], str | None]
    partial: Path | None = None

check_csv(path, columns)

Check that a CSV file opens and has some columns.

Parameters:

Name	Type	Description	Default
path	Path	The file, gzipped or not.	required
columns	sequence of str	Columns that must be present.	required

Returns:

Type	Description
str or None	None when the file is usable, else what is wrong.

Examples:

>>> import tempfile
>>> folder = Path(tempfile.mkdtemp())
>>> _ = (folder / "a.csv").write_text("id,name\n1,x\n")
>>> check_csv(folder / "a.csv", ["id"])
>>> check_csv(folder / "a.csv", ["size"])
'missing columns: size'

Source code in lab/connectome/load.py

def check_csv(path: Path, columns: Sequence[str]) -> str | None:
    r"""Check that a CSV file opens and has some columns.

    Parameters
    ----------
    path : pathlib.Path
        The file, gzipped or not.
    columns : sequence of str
        Columns that must be present.

    Returns
    -------
    str or None
        ``None`` when the file is usable, else what is wrong.

    Examples
    --------
    >>> import tempfile
    >>> folder = Path(tempfile.mkdtemp())
    >>> _ = (folder / "a.csv").write_text("id,name\n1,x\n")
    >>> check_csv(folder / "a.csv", ["id"])
    >>> check_csv(folder / "a.csv", ["size"])
    'missing columns: size'
    """
    try:
        header = pl.read_csv(path, n_rows=1).columns
    # Any failure to open the file means it is unreadable.
    except Exception as e:  # noqa: BLE001
        return f"unreadable: {type(e).__name__}"
    missing = [c for c in columns if c not in header]
    return f"missing columns: {', '.join(missing)}" if missing else None

check_parquet(path, columns)

Check that a parquet file opens and has some columns.

Only the footer is read, so the check is fast on a large file.

Parameters:

Name	Type	Description	Default
path	Path	The file.	required
columns	sequence of str	Columns that must be present.	required

Returns:

Type	Description
str or None	None when the file is usable, else what is wrong.

Examples:

>>> import tempfile
>>> path = Path(tempfile.mkdtemp()) / "a.parquet"
>>> pl.DataFrame({"count": [1]}).write_parquet(path)
>>> check_parquet(path, ["count"])
>>> check_parquet(path, ["id"])
'missing columns: id'

Source code in lab/connectome/load.py

def check_parquet(path: Path, columns: Sequence[str]) -> str | None:
    """Check that a parquet file opens and has some columns.

    Only the footer is read, so the check is fast on a large file.

    Parameters
    ----------
    path : pathlib.Path
        The file.
    columns : sequence of str
        Columns that must be present.

    Returns
    -------
    str or None
        ``None`` when the file is usable, else what is wrong.

    Examples
    --------
    >>> import tempfile
    >>> path = Path(tempfile.mkdtemp()) / "a.parquet"
    >>> pl.DataFrame({"count": [1]}).write_parquet(path)
    >>> check_parquet(path, ["count"])
    >>> check_parquet(path, ["id"])
    'missing columns: id'
    """
    try:
        header = pl.scan_parquet(path).collect_schema().names()
    # Any failure to open the file means it is unreadable.
    except Exception as e:  # noqa: BLE001
        return f"unreadable: {type(e).__name__}"
    missing = [c for c in columns if c not in header]
    return f"missing columns: {', '.join(missing)}" if missing else None

check_swc_folder(path)

Check that a folder holds skeletons and no unfinished download.

Parameters:

Name	Type	Description	Default
path	Path	The folder of <id>.swc files.	required

Returns:

Type	Description
str or None	None when there is at least one skeleton and no .part file, else what is wrong.

Examples:

>>> import tempfile
>>> folder = Path(tempfile.mkdtemp())
>>> check_swc_folder(folder)
'no skeletons'
>>> _ = (folder / "1.swc").write_text("1 1 0 0 0 1 -1\n")
>>> check_swc_folder(folder)

Source code in lab/connectome/load.py

def check_swc_folder(path: Path) -> str | None:
    r"""Check that a folder holds skeletons and no unfinished download.

    Parameters
    ----------
    path : pathlib.Path
        The folder of ``<id>.swc`` files.

    Returns
    -------
    str or None
        ``None`` when there is at least one skeleton and no ``.part``
        file, else what is wrong.

    Examples
    --------
    >>> import tempfile
    >>> folder = Path(tempfile.mkdtemp())
    >>> check_swc_folder(folder)
    'no skeletons'
    >>> _ = (folder / "1.swc").write_text("1 1 0 0 0 1 -1\n")
    >>> check_swc_folder(folder)
    """
    if not any(path.glob("*.swc")):
        return "no skeletons"
    partial = list(path.glob("*.part"))
    if partial:
        return f"{len(partial)} half-finished downloads"
    return None

figure_path(name, folder=None)

Return the path of a figure file, creating the folder.

Parameters:

Name	Type	Description	Default
name	str	File name inside the folder.	required
folder	Path	Where figures go; config.FIGURES by default.	None

Returns:

Type	Description
Path	folder / name; folder exists afterwards.

Examples:

>>> import tempfile
>>> folder = Path(tempfile.mkdtemp()) / "figures"
>>> figure_path("a.png", folder) == folder / "a.png"
True
>>> folder.is_dir()
True

Source code in lab/connectome/load.py

def figure_path(name: str, folder: Path | None = None) -> Path:
    """Return the path of a figure file, creating the folder.

    Parameters
    ----------
    name : str
        File name inside the folder.
    folder : pathlib.Path, optional
        Where figures go; ``config.FIGURES`` by default.

    Returns
    -------
    pathlib.Path
        ``folder / name``; `folder` exists afterwards.

    Examples
    --------
    >>> import tempfile
    >>> folder = Path(tempfile.mkdtemp()) / "figures"
    >>> figure_path("a.png", folder) == folder / "a.png"
    True
    >>> folder.is_dir()
    True
    """
    folder = config.FIGURES if folder is None else folder
    folder.mkdir(parents=True, exist_ok=True)
    return folder / name

load_census(path=None)

Read the cell-type census.

Only root_id and primary_type are read; the other columns of the file, such as additional_type(s), are left out.

Parameters:

Name	Type	Description	Default
path	Path	The census CSV; config.CENSUS by default.	None

Returns:

Type	Description
DataFrame	Columns root_id and primary_type, one row per typed cell.

Examples:

>>> import tempfile
>>> path = Path(tempfile.mkdtemp()) / "census.csv"
>>> _ = path.write_text(
...     "root_id,primary_type,additional_type(s)\n1,A,\n2,B,x\n"
... )
>>> load_census(path).rows()
[(1, 'A'), (2, 'B')]

Source code in lab/connectome/load.py

def load_census(path: Path | None = None) -> pl.DataFrame:
    r"""Read the cell-type census.

    Only ``root_id`` and ``primary_type`` are read; the other columns
    of the file, such as ``additional_type(s)``, are left out.

    Parameters
    ----------
    path : pathlib.Path, optional
        The census CSV; ``config.CENSUS`` by default.

    Returns
    -------
    polars.DataFrame
        Columns ``root_id`` and ``primary_type``, one row per typed
        cell.

    Examples
    --------
    >>> import tempfile
    >>> path = Path(tempfile.mkdtemp()) / "census.csv"
    >>> _ = path.write_text(
    ...     "root_id,primary_type,additional_type(s)\n1,A,\n2,B,x\n"
    ... )
    >>> load_census(path).rows()
    [(1, 'A'), (2, 'B')]
    """
    return pl.read_csv(
        require(config.CENSUS if path is None else path),
        columns=["root_id", "primary_type"],
        schema_overrides={"root_id": pl.Int64},
    )

load_columns(path=None)

Read the column assignment of the columnar cells.

Parameters:

Name	Type	Description	Default
path	Path	The column assignment CSV; config.COLUMNS by default.	None

Returns:

Type	Description
DataFrame	Columns root_id, hemisphere, type, column_id, x, y, p, and q.

Examples:

>>> import tempfile
>>> path = Path(tempfile.mkdtemp()) / "columns.csv"
>>> _ = path.write_text(
...     "root_id,hemisphere,type,column_id,x,y,p,q\n"
...     "1,right,Tm1,7,2,-3,0,1\n"
... )
>>> load_columns(path).select("root_id", "x", "y").rows()
[(1, 2, -3)]

Source code in lab/connectome/load.py

def load_columns(path: Path | None = None) -> pl.DataFrame:
    r"""Read the column assignment of the columnar cells.

    Parameters
    ----------
    path : pathlib.Path, optional
        The column assignment CSV; ``config.COLUMNS`` by default.

    Returns
    -------
    polars.DataFrame
        Columns ``root_id``, ``hemisphere``, ``type``, ``column_id``,
        ``x``, ``y``, ``p``, and ``q``.

    Examples
    --------
    >>> import tempfile
    >>> path = Path(tempfile.mkdtemp()) / "columns.csv"
    >>> _ = path.write_text(
    ...     "root_id,hemisphere,type,column_id,x,y,p,q\n"
    ...     "1,right,Tm1,7,2,-3,0,1\n"
    ... )
    >>> load_columns(path).select("root_id", "x", "y").rows()
    [(1, 2, -3)]
    """
    return pl.read_csv(
        require(config.COLUMNS if path is None else path),
        schema_overrides={"root_id": pl.Int64},
    )

load_connections_of(root_ids, *, end='pre', min_weight=1, path=None)

Load the stored connections that start or end at some cells.

Parameters:

Name	Type	Description	Default
root_ids	sequence of int	The cells to select.	required
end	(pre, post)	Keep connections that start at ("pre") or end at ("post") one of root_ids.	"pre"
min_weight	int	Smallest weight, in synapses, to keep.	1
path	Path	The connections Parquet file; config.CONNECTIONS by default.	None

Returns:

Type	Description
DataFrame	Columns pre_pt_root_id, post_pt_root_id, and weight.

Raises:

Type	Description
ValueError	For any other end.

Examples:

>>> import tempfile
>>> path = Path(tempfile.mkdtemp()) / "connections.parquet"
>>> pl.DataFrame(
...     {
...         "pre_pt_root_id": [1, 1, 2],
...         "post_pt_root_id": [2, 3, 3],
...         "weight": pl.Series([7, 2, 9], dtype=pl.UInt32),
...     }
... ).write_parquet(path)
>>> load_connections_of([1], min_weight=5, path=path).rows()
[(1, 2, 7)]
>>> load_connections_of([3], end="post", path=path).rows()
[(1, 3, 2), (2, 3, 9)]

Source code in lab/connectome/load.py

def load_connections_of(
    root_ids: Sequence[int],
    *,
    end: str = "pre",
    min_weight: int = 1,
    path: Path | None = None,
) -> pl.DataFrame:
    """Load the stored connections that start or end at some cells.

    Parameters
    ----------
    root_ids : sequence of int
        The cells to select.
    end : {"pre", "post"}, default "pre"
        Keep connections that start at (``"pre"``) or end at
        (``"post"``) one of `root_ids`.
    min_weight : int, default 1
        Smallest weight, in synapses, to keep.
    path : pathlib.Path, optional
        The connections Parquet file; ``config.CONNECTIONS`` by
        default.

    Returns
    -------
    polars.DataFrame
        Columns ``pre_pt_root_id``, ``post_pt_root_id``, and
        ``weight``.

    Raises
    ------
    ValueError
        For any other `end`.

    Examples
    --------
    >>> import tempfile
    >>> path = Path(tempfile.mkdtemp()) / "connections.parquet"
    >>> pl.DataFrame(
    ...     {
    ...         "pre_pt_root_id": [1, 1, 2],
    ...         "post_pt_root_id": [2, 3, 3],
    ...         "weight": pl.Series([7, 2, 9], dtype=pl.UInt32),
    ...     }
    ... ).write_parquet(path)
    >>> load_connections_of([1], min_weight=5, path=path).rows()
    [(1, 2, 7)]
    >>> load_connections_of([3], end="post", path=path).rows()
    [(1, 3, 2), (2, 3, 9)]
    """
    if end not in ("pre", "post"):
        msg = f"end must be 'pre' or 'post', not {end!r}"
        raise ValueError(msg)
    return (
        scan_connections(path)
        .filter(
            pl.col(f"{end}_pt_root_id").is_in(list(root_ids))
            & (pl.col("weight") >= min_weight)
        )
        .collect()
    )

load_proofread_ids(path=None)

Read the list of proofread neurons.

Parameters:

Name	Type	Description	Default
path	Path	The classification CSV; config.CLASSIFICATION by default.	None

Returns:

Type	Description
DataFrame	Column root_id, one row per proofread neuron.

Examples:

>>> import tempfile
>>> path = Path(tempfile.mkdtemp()) / "classification.csv"
>>> _ = path.write_text("root_id,flow\n1,intrinsic\n2,afferent\n")
>>> load_proofread_ids(path).rows()
[(1,), (2,)]

Source code in lab/connectome/load.py

def load_proofread_ids(path: Path | None = None) -> pl.DataFrame:
    r"""Read the list of proofread neurons.

    Parameters
    ----------
    path : pathlib.Path, optional
        The classification CSV; ``config.CLASSIFICATION`` by default.

    Returns
    -------
    polars.DataFrame
        Column ``root_id``, one row per proofread neuron.

    Examples
    --------
    >>> import tempfile
    >>> path = Path(tempfile.mkdtemp()) / "classification.csv"
    >>> _ = path.write_text("root_id,flow\n1,intrinsic\n2,afferent\n")
    >>> load_proofread_ids(path).rows()
    [(1,), (2,)]
    """
    return pl.read_csv(
        require(config.CLASSIFICATION if path is None else path),
        columns=["root_id"],
        schema_overrides={"root_id": pl.Int64},
    )

load_visual_types(path=None)

Read the visual neuron types, with the side of each cell.

Parameters:

Name	Type	Description	Default
path	Path	The visual neuron types CSV; config.VISUAL_TYPES by default.	None

Returns:

Type	Description
DataFrame	Columns root_id, type, side, and any others in the file.

Examples:

>>> import tempfile
>>> path = Path(tempfile.mkdtemp()) / "visual_types.csv"
>>> _ = path.write_text("root_id,type,side\n1,LC9,right\n")
>>> load_visual_types(path).rows()
[(1, 'LC9', 'right')]

Source code in lab/connectome/load.py

def load_visual_types(path: Path | None = None) -> pl.DataFrame:
    r"""Read the visual neuron types, with the side of each cell.

    Parameters
    ----------
    path : pathlib.Path, optional
        The visual neuron types CSV; ``config.VISUAL_TYPES`` by
        default.

    Returns
    -------
    polars.DataFrame
        Columns ``root_id``, ``type``, ``side``, and any others in the
        file.

    Examples
    --------
    >>> import tempfile
    >>> path = Path(tempfile.mkdtemp()) / "visual_types.csv"
    >>> _ = path.write_text("root_id,type,side\n1,LC9,right\n")
    >>> load_visual_types(path).rows()
    [(1, 'LC9', 'right')]
    """
    return pl.read_csv(
        require(config.VISUAL_TYPES if path is None else path),
        schema_overrides={"root_id": pl.Int64},
    )

require(path)

Return a data path, or stop with the command that writes it.

Parameters:

Name	Type	Description	Default
path	Path	A data file or folder, normally one of DATASETS.	required

Returns:

Type	Description
Path	path, when it exists.

Raises:

Type	Description
SystemExit	When path is missing; the message names the command that writes it, when path is one of DATASETS.

Examples:

>>> import tempfile
>>> folder = Path(tempfile.mkdtemp())
>>> require(folder) == folder
True
>>> require(folder / "gone.csv")
Traceback (most recent call last):
SystemExit: .../gone.csv missing

Source code in lab/connectome/load.py

def require(path: Path) -> Path:
    """Return a data path, or stop with the command that writes it.

    Parameters
    ----------
    path : pathlib.Path
        A data file or folder, normally one of `DATASETS`.

    Returns
    -------
    pathlib.Path
        `path`, when it exists.

    Raises
    ------
    SystemExit
        When `path` is missing; the message names the command that
        writes it, when `path` is one of `DATASETS`.

    Examples
    --------
    >>> import tempfile
    >>> folder = Path(tempfile.mkdtemp())
    >>> require(folder) == folder
    True
    >>> require(folder / "gone.csv")  # doctest: +ELLIPSIS
    Traceback (most recent call last):
    SystemExit: .../gone.csv missing
    """
    if not path.exists():
        command = FETCH.get(path)
        msg = f"{path} missing" + (f"; run `{command}`" if command else "")
        raise SystemExit(msg)
    return path

root_ids_of_type(census, cell_type)

List the root ids of one primary type, as Python ints.

The ids are plain Python ints, not numpy integers.

Parameters:

Name	Type	Description	Default
census	DataFrame	Has columns root_id and primary_type, as load_census returns them.	required
cell_type	str	The primary_type to select.	required

Returns:

Type	Description
list of int	The root_id of every row of that type.

Examples:

>>> import polars as pl
>>> census = pl.DataFrame(
...     {
...         "root_id": [1, 2, 3],
...         "primary_type": ["A", "B", "A"],
...     }
... )
>>> root_ids_of_type(census, "A")
[1, 3]

Source code in lab/connectome/load.py

def root_ids_of_type(census: pl.DataFrame, cell_type: str) -> list[int]:
    """List the root ids of one primary type, as Python ints.

    The ids are plain Python ints, not numpy integers.

    Parameters
    ----------
    census : polars.DataFrame
        Has columns ``root_id`` and ``primary_type``, as `load_census`
        returns them.
    cell_type : str
        The ``primary_type`` to select.

    Returns
    -------
    list of int
        The ``root_id`` of every row of that type.

    Examples
    --------
    >>> import polars as pl
    >>> census = pl.DataFrame(
    ...     {
    ...         "root_id": [1, 2, 3],
    ...         "primary_type": ["A", "B", "A"],
    ...     }
    ... )
    >>> root_ids_of_type(census, "A")
    [1, 3]
    """
    return census.filter(pl.col("primary_type") == cell_type)[
        "root_id"
    ].to_list()

scan_connections(path=None)

Scan every stored connection lazily.

Filter before calling collect on the result: the file holds every pair of cells with a proofread end.

Parameters:

Name	Type	Description	Default
path	Path	The connections Parquet file; config.CONNECTIONS by default.	None

Returns:

Type	Description
LazyFrame	Columns pre_pt_root_id, post_pt_root_id, and weight, all Int64.

Examples:

>>> import tempfile
>>> path = Path(tempfile.mkdtemp()) / "connections.parquet"
>>> pl.DataFrame(
...     {
...         "pre_pt_root_id": [1, 2],
...         "post_pt_root_id": [2, 3],
...         "weight": pl.Series([7, 9], dtype=pl.UInt32),
...     }
... ).write_parquet(path)
>>> heavy = scan_connections(path).filter(pl.col("weight") > 8)
>>> heavy.collect().rows()
[(2, 3, 9)]

Source code in lab/connectome/load.py

def scan_connections(path: Path | None = None) -> pl.LazyFrame:
    """Scan every stored connection lazily.

    Filter before calling ``collect`` on the result: the file holds
    every pair of cells with a proofread end.

    Parameters
    ----------
    path : pathlib.Path, optional
        The connections Parquet file; ``config.CONNECTIONS`` by
        default.

    Returns
    -------
    polars.LazyFrame
        Columns ``pre_pt_root_id``, ``post_pt_root_id``, and
        ``weight``, all Int64.

    Examples
    --------
    >>> import tempfile
    >>> path = Path(tempfile.mkdtemp()) / "connections.parquet"
    >>> pl.DataFrame(
    ...     {
    ...         "pre_pt_root_id": [1, 2],
    ...         "post_pt_root_id": [2, 3],
    ...         "weight": pl.Series([7, 9], dtype=pl.UInt32),
    ...     }
    ... ).write_parquet(path)
    >>> heavy = scan_connections(path).filter(pl.col("weight") > 8)
    >>> heavy.collect().rows()
    [(2, 3, 9)]
    """
    return pl.scan_parquet(
        require(config.CONNECTIONS if path is None else path)
    ).with_columns(pl.col("weight").cast(pl.Int64))

status(datasets=DATASETS)

Report whether each dataset is present and usable.

Parameters:

Name	Type	Description	Default
datasets	sequence of Dataset	What to check.	`DATASETS`

Returns:

Type	Description
DataFrame	One row per dataset, with columns: dataset, source As in the registry. state "ok", "missing", or "problem". detail What is wrong, or how many pieces of an interrupted fetch are on disk; empty when state is "ok". size Bytes on disk, null when missing. modified Time of the newest file, null when missing. fetch The command that writes it.

Examples:

>>> import tempfile
>>> folder = Path(tempfile.mkdtemp())
>>> there = Dataset(
...     "a", "x", folder, "f()", lambda p: None
... )
>>> gone = Dataset("b", "x", folder / "gone", "g()", lambda p: None)
>>> status([there, gone]).select("dataset", "state").rows()
[('a', 'ok'), ('b', 'missing')]

Source code in lab/connectome/load.py

def status(datasets: Sequence[Dataset] = DATASETS) -> pl.DataFrame:
    """Report whether each dataset is present and usable.

    Parameters
    ----------
    datasets : sequence of Dataset, default `DATASETS`
        What to check.

    Returns
    -------
    polars.DataFrame
        One row per dataset, with columns:

        ``dataset``, ``source``
            As in the registry.
        ``state``
            ``"ok"``, ``"missing"``, or ``"problem"``.
        ``detail``
            What is wrong, or how many pieces of an interrupted fetch
            are on disk; empty when ``state`` is ``"ok"``.
        ``size``
            Bytes on disk, null when missing.
        ``modified``
            Time of the newest file, null when missing.
        ``fetch``
            The command that writes it.

    Examples
    --------
    >>> import tempfile
    >>> folder = Path(tempfile.mkdtemp())
    >>> there = Dataset(
    ...     "a", "x", folder, "f()", lambda p: None
    ... )
    >>> gone = Dataset("b", "x", folder / "gone", "g()", lambda p: None)
    >>> status([there, gone]).select("dataset", "state").rows()
    [('a', 'ok'), ('b', 'missing')]
    """
    rows = []
    for d in datasets:
        if d.path.exists():
            problem = d.check(d.path)
            files = _files(d.path)
            state = "problem" if problem else "ok"
            detail = problem or ""
            size = sum(f.stat().st_size for f in files)
            newest = max((f.stat().st_mtime for f in files), default=None)
            # The file's modification time in the local zone.
            modified = (
                datetime.fromtimestamp(newest)  # noqa: DTZ006
                if newest
                else None
            )
        else:
            state, size, modified = "missing", None, None
            pieces = (
                len(list(d.partial.glob("*")))
                if d.partial is not None and d.partial.exists()
                else 0
            )
            detail = (
                f"{pieces} pieces of an interrupted fetch" if pieces else ""
            )
        rows.append(
            {
                "dataset": d.name,
                "source": d.source,
                "state": state,
                "detail": detail,
                "size": size,
                "modified": modified,
                "fetch": d.fetch,
            }
        )
    return pl.DataFrame(
        rows,
        schema={
            "dataset": pl.String,
            "source": pl.String,
            "state": pl.String,
            "detail": pl.String,
            "size": pl.Int64,
            "modified": pl.Datetime,
            "fetch": pl.String,
        },
    )