neuroglancer.py

"""
Code for generating and working with neuroglancer links.
"""

import ast
import json
from pathlib import Path

import neuroglancer
from neuroglancer import CoordinateSpace

from hipct_data_tools.data_model import HiPCTDataSet

NEUROGLANCER_INSTANCE = "https://neuroglancer-demo.appspot.com"
# Version of the nueroglancer dict. If *anything* changes in the dict,
# increment this version and document the changes here.
#
# v1: Initial version
VERSION: int = 1


def dataset_to_layer(
    dataset: HiPCTDataSet, *, add_transform: bool
) -> neuroglancer.ManagedLayer:
    invlerp_params = neuroglancer.InvlerpParameters(
        range=(dataset.contrast_low, dataset.contrast_high)
    )
    shader_controls = {"normalized": invlerp_params.to_json()}

    source_url = f"n5://gs://{dataset.gcs_bucket}/{dataset.gcs_path}"

    if add_transform:
        coord_space = CoordinateSpace(
            names=["x", "y", "z"], units=["um"] * 3, scales=[dataset.resolution_um] * 3
        )
        coord_transform = neuroglancer.CoordinateSpaceTransform(
            output_dimensions=coord_space, matrix=get_transform_matrix(dataset)
        )
        source = neuroglancer.LayerDataSource(url=source_url, transform=coord_transform)
    else:
        source = neuroglancer.LayerDataSource(url=source_url)

    layer = neuroglancer.ManagedLayer(
        name=dataset.name,
        layer=neuroglancer.ImageLayer(shader_controls=shader_controls, source=source),
    )
    layer.visible = True
    return layer


def individual_neuroglancer_state(dataset: HiPCTDataSet) -> neuroglancer.ViewerState:
    """
    Get the viewer state for an individual dataset.
    """
    dimensions = neuroglancer.CoordinateSpace(
        names=["x", "y", "z"], units=["um"] * 3, scales=[dataset.resolution_um] * 3
    )
    layer = dataset_to_layer(dataset, add_transform=False)
    selected_layer = neuroglancer.SelectedLayerState(layer=layer.name, visible=True)

    viewer_state = neuroglancer.ViewerState(
        dimensions=dimensions,
        layout="4panel",
        projection_orientation=(0.3, 0.2, 0, -0.9),
        projectionScale=4096,
        selectedLayer=selected_layer.to_json(),
    )
    viewer_state.layers.append(layer)
    return viewer_state


def neuroglancer_link(dataset: HiPCTDataSet, *, multires: bool = False) -> str | None:
    """
    Nueroglancer dict for an individual dataset.

    Parameters
    ----------
    dataset :
        Dataset to create link for.
    multires :
        If `True`, create a multiresolution link that contains child datasets.
        Only works for complete organ datasets.

    Notes
    -----
    Returns `None` if dataset isn't present in a GCS bucket.
    """
    if dataset.gcs_bucket == "none":
        return None

    state = individual_neuroglancer_state(dataset)
    return f"{NEUROGLANCER_INSTANCE}/#!{json.dumps(state.to_json(), separators=(',', ':'))}"


def get_transform_matrix(dataset: HiPCTDataSet) -> list[list[float]]:
    """
    Get transform_matrix for a dataset.

    Returns
    -------
    matrix :
        Transformation matrix for neuroglancer.
    """
    # Neuroglancer requires the model coordinate space to be in nm
    resolution_nm = dataset.resolution_um / 1000
    if (log_path := dataset.registration_log_path) is None:
        raise ValueError(f"No registration log file found for {dataset.name}")

    return _get_transform_matrix_from_log(log_path, resolution_nm)


def _multires_nueroglancer_state(
    dataset: HiPCTDataSet,
) -> neuroglancer.ViewerState | None:
    """
    Neuroglancer state for a complete organ dataset with multi-resolution child datasets
    included.

    Notes
    -----
    Returns `None` if dataset isn't present in a GCS bucket.
    """
    if dataset.gcs_bucket == "none":
        return None
    if not dataset.is_complete_organ:
        raise ValueError(
            "Can only make multiresolution link for complete organ dataests."
        )

    state = individual_neuroglancer_state(dataset)
    for child in dataset.child_datasets():
        layer = dataset_to_layer(child, add_transform=True)
        state.layers.append(layer)

    return state


def neuroglancer_link_multires(dataset: HiPCTDataSet) -> str | None:
    """
    Nueroglancer link for a complete organ dataset with multi-resolution child datasets
    included.

    Notes
    -----
    Returns `None` if dataset isn't present in a GCS bucket.
    """
    if (state := _multires_nueroglancer_state(dataset)) is not None:
        return f"{NEUROGLANCER_INSTANCE}/#!{json.dumps(state.to_json(), separators=(',', ':'))}"
    return None


def _get_transform_matrix_from_log(
    log_path: Path, resolution_nm: float
) -> list[list[float]]:
    """
    Get transform_matrix from a registration log file.

    Parameters
    ----------
    registration_log_file :
        Path to registration log file.
    resolution_m :
        Voxel size in micro meters.

    Returns
    -------
    matrix :
        Transformation matrix for neuroglancer.
    """
    with open(log_path, newline="\n") as f:
        lines = [l.rstrip("\n") for l in f.readlines()]

    rotation_matrix = None
    for l in lines:
        if "Translation (Tx,Ty,Tz)" in l:
            translation = [
                float(n) / resolution_nm for n in l.split("= ")[-1].split(",")
            ]
        if "Inverse rotation matrix" in l:
            rotation_matrix = list(ast.literal_eval(l.split("= ")[-1]))
            break

    if rotation_matrix is None:
        raise RuntimeError(f"Did not find rotation matrix in file {log_path}")

    return [
        rotation_matrix[0:3] + [translation[0]],
        rotation_matrix[3:6] + [translation[1]],
        rotation_matrix[6:9] + [translation[2]],
    ]