swcgeom 0.15.0__py3-none-any.whl → 0.17.0__py3-none-any.whl

swcgeom/transforms/images.py

@@ -1,17 +1,31 @@
 """Image stack related transform."""
 
-
+import warnings
 from typing import Tuple
 
 import numpy as np
 import numpy.typing as npt
 
-from swcgeom.transforms.base import Transform
+from swcgeom.transforms.base import Identity, Transform
+
+__all__ = [
+    "ImagesCenterCrop",
+    "ImagesScale",
+    "ImagesClip",
+    "ImagesFlip",
+    "ImagesFlipY",
+    "ImagesNormalizer",
+    "ImagesMeanVarianceAdjustment",
+    "ImagesScaleToUnitRange",
+    "ImagesHistogramEqualization",
+    "Center",  # legacy
+]
 
-__all__ = ["Center"]
 
+NDArrayf32 = npt.NDArray[np.float32]
 
-class Center(Transform[npt.NDArray[np.float32], npt.NDArray[np.float32]]):
+
+class ImagesCenterCrop(Transform[NDArrayf32, NDArrayf32]):
     """Get image stack center."""
 
     def __init__(self, shape_out: int | Tuple[int, int, int]):
@@ -22,7 +36,7 @@ class Center(Transform[npt.NDArray[np.float32], npt.NDArray[np.float32]]):
             else (shape_out, shape_out, shape_out)
         )
 
-    def __call__(self, x: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
         diff = np.subtract(x.shape[:3], self.shape_out)
         s = diff // 2
         e = np.add(s, self.shape_out)
@@ -30,3 +44,160 @@ class Center(Transform[npt.NDArray[np.float32], npt.NDArray[np.float32]]):
 
     def extra_repr(self) -> str:
         return f"shape_out=({','.join(str(a) for a in self.shape_out)})"
+
+
+class Center(ImagesCenterCrop):
+    """Get image stack center.
+
+    .. deprecated:: 0.16.0
+        Use :class:`ImagesCenterCrop` instead.
+    """
+
+    def __init__(self, shape_out: int | Tuple[int, int, int]):
+        warnings.warn(
+            "`Center` is deprecated, use `ImagesCenterCrop` instead",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        super().__init__(shape_out)
+
+
+class ImagesScale(Transform[NDArrayf32, NDArrayf32]):
+    def __init__(self, scaler: float) -> None:
+        super().__init__()
+        self.scaler = scaler
+
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        return self.scaler * x
+
+    def extra_repr(self) -> str:
+        return f"scaler={self.scaler}"
+
+
+class ImagesClip(Transform[NDArrayf32, NDArrayf32]):
+    def __init__(self, vmin: float = 0, vmax: float = 1, /) -> None:
+        super().__init__()
+        self.vmin, self.vmax = vmin, vmax
+
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        return np.clip(x, self.vmin, self.vmax)
+
+    def extra_repr(self) -> str:
+        return f"vmin={self.vmin}, vmax={self.vmax}"
+
+
+class ImagesFlip(Transform[NDArrayf32, NDArrayf32]):
+    """Flip image stack along axis."""
+
+    def __init__(self, axis: int, /) -> None:
+        super().__init__()
+        self.axis = axis
+
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        return np.flip(x, axis=self.axis)
+
+    def extra_repr(self) -> str:
+        return f"axis={self.axis}"
+
+
+class ImagesFlipY(ImagesFlip):
+    """Flip image stack along Y-axis.
+
+    See Also
+    --------
+    ~.images.io.TeraflyImageStack:
+        Terafly and Vaa3d use a especial right-handed coordinate system
+        (with origin point in the left-top and z-axis points front),
+        but we flip y-axis to makes it a left-handed coordinate system
+        (with orgin point in the left-bottom and z-axis points front).
+        If you need to use its coordinate system, remember to FLIP
+        Y-AXIS BACK.
+    """
+
+    def __init__(self, axis: int = 1, /) -> None:
+        super().__init__(axis)  # (X, Y, Z, C)
+
+
+class ImagesNormalizer(Transform[NDArrayf32, NDArrayf32]):
+    """Normalize image stack."""
+
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        mean = np.mean(x)
+        variance = np.var(x)
+        return (x - mean) / variance
+
+
+class ImagesMeanVarianceAdjustment(Transform[NDArrayf32, NDArrayf32]):
+    """Adjust image stack mean and variance.
+
+    See Also
+    --------
+    ~swcgeom.images.ImageStackFolder.stat
+    """
+
+    def __init__(self, mean: float, variance: float) -> None:
+        super().__init__()
+        self.mean = mean
+        self.variance = variance
+
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        return (x - self.mean) / self.variance
+
+    def extra_repr(self) -> str:
+        return f"mean={self.mean}, variance={self.variance}"
+
+
+class ImagesScaleToUnitRange(Transform[NDArrayf32, NDArrayf32]):
+    """Scale image stack to unit range."""
+
+    def __init__(self, vmin: float, vmax: float, *, clip: bool = True) -> None:
+        """Scale image stack to unit range.
+
+        Parameters
+        ----------
+        vmin : float
+            Minimum value.
+        vmax : float
+            Maximum value.
+        clip : bool, default True
+            Clip values to [0, 1] to avoid numerical issues.
+        """
+
+        super().__init__()
+        self.vmin = vmin
+        self.vmax = vmax
+        self.diff = vmax - vmin
+        self.clip = clip
+        self.post = ImagesClip(0, 1) if self.clip else Identity()
+
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        return self.post((x - self.vmin) / self.diff)
+
+    def extra_repr(self) -> str:
+        return f"vmin={self.vmin}, vmax={self.vmax}, clip={self.clip}"
+
+
+class ImagesHistogramEqualization(Transform[NDArrayf32, NDArrayf32]):
+    """Image histogram equalization.
+
+    References
+    ----------
+    http://www.janeriksolem.net/histogram-equalization-with-python-and.html
+    """
+
+    def __init__(self, bins: int = 256) -> None:
+        super().__init__()
+        self.bins = bins
+
+    def __call__(self, x: NDArrayf32) -> NDArrayf32:
+        # get image histogram
+        hist, bin_edges = np.histogram(x.flatten(), self.bins, density=True)
+        cdf = hist.cumsum()  # cumulative distribution function
+        cdf = cdf / cdf[-1]  # normalize
+
+        # use linear interpolation of cdf to find new pixel values
+        equalized = np.interp(x.flatten(), bin_edges[:-1], cdf)
+        return equalized.reshape(x.shape).astype(np.float32)
+
+    def extra_repr(self) -> str:
+        return f"bins={self.bins}"

swcgeom/transforms/mst.py

@@ -23,11 +23,11 @@ class PointsToCuntzMST(Transform[npt.NDArray[np.float32], Tree]):
 
     References
     ----------
-    [1] Cuntz, H., Forstner, F., Borst, A. & Häusser, M. One Rule to
-    Grow Them Al: A General Theory of Neuronal Branching and Its
-    Practical Application. PLOS Comput Biol 6, e1000877 (2010).
-    [2] Cuntz, H., Borst, A. & Segev, I. Optimization principles of
-    dendritic structure. Theor Biol Med Model 4, 21 (2007).
+    .. [1] Cuntz, H., Forstner, F., Borst, A. & Häusser, M. One Rule to
+       Grow Them Al: A General Theory of Neuronal Branching and Its
+       Practical Application. PLOS Comput Biol 6, e1000877 (2010).
+    .. [2] Cuntz, H., Borst, A. & Segev, I. Optimization principles of
+       dendritic structure. Theor Biol Med Model 4, 21 (2007).
     """
 
     def __init__(

swcgeom/transforms/neurolucida_asc.py

@@ -0,0 +1,495 @@
+"""Neurolucida related transformation."""
+
+import os
+import re
+from enum import Enum, auto
+from io import TextIOBase
+from typing import Any, List, NamedTuple, Optional, cast
+
+import numpy as np
+
+from swcgeom.core import Tree
+from swcgeom.core.swc_utils import SWCNames, SWCTypes, get_names, get_types
+from swcgeom.transforms.base import Transform
+
+__all__ = ["NeurolucidaAscToSwc"]
+
+
+class NeurolucidaAscToSwc(Transform[str, Tree]):
+    """Convert neurolucida asc format to swc format."""
+
+    def __call__(self, x: str) -> Tree:
+        return self.convert(x)
+
+    @classmethod
+    def convert(cls, fname: str) -> Tree:
+        with open(fname, "r") as f:
+            tree = cls.from_stream(f, source=os.path.abspath(fname))
+
+        return tree
+
+    @classmethod
+    def from_stream(cls, x: TextIOBase, *, source: str = "") -> Tree:
+        parser = Parser(x, source=source)
+        ast = parser.parse()
+        tree = cls.from_ast(ast)
+        return tree
+
+    @staticmethod
+    def from_ast(
+        ast: "AST",
+        *,
+        names: Optional[SWCNames] = None,
+        types: Optional[SWCTypes] = None,
+    ) -> Tree:
+        names = get_names(names)
+        types = get_types(types)
+        ndata = {n: [] for n in names.cols()}
+
+        next_id = 0
+        typee = [types.undefined]
+
+        def walk_ast(root: ASTNode, pid: int = -1) -> None:
+            nonlocal next_id, typee
+            match root.type:
+                case ASTType.ROOT:
+                    for n in root.children:
+                        walk_ast(n)
+
+                case ASTType.TREE:
+                    match root.value:
+                        case "AXON":
+                            typee.append(types.axon)
+                        case "DENDRITE":
+                            typee.append(types.basal_dendrite)
+
+                    for n in root.children:
+                        walk_ast(n)
+
+                    typee.pop()
+
+                case ASTType.NODE:
+                    x, y, z, r = root.value
+                    idx = next_id
+                    next_id += 1
+
+                    ndata[names.id].append(idx)
+                    ndata[names.type].append(typee[-1])
+                    ndata[names.x].append(x)
+                    ndata[names.y].append(y)
+                    ndata[names.z].append(z)
+                    ndata[names.r].append(r)
+                    ndata[names.pid].append(pid)
+
+                    for n in root.children:
+                        walk_ast(n, pid=idx)
+
+        walk_ast(ast)
+        tree = Tree(
+            next_id,
+            source=ast.source,
+            names=names,
+            **ndata,  # type: ignore
+        )
+        return tree
+
+
+# -----------------
+# ASC format parser
+# -----------------
+
+# AST
+
+
+class ASTType(Enum):
+    ROOT = auto()
+    TREE = auto()
+    NODE = auto()
+    COLOR = auto()
+    COMMENT = auto()
+
+
+class ASTNode:
+    parent: "ASTNode | None" = None
+
+    def __init__(
+        self,
+        type: ASTType,
+        value: Any = None,
+        tokens: Optional[List["Token"]] = None,
+        children: Optional[List["ASTNode"]] = None,
+    ):
+        self.type = type
+        self.value = value
+        self.tokens = tokens or []
+        self.children = children or []
+        for child in self.children:
+            child.parent = self
+
+    def add_child(self, child: "ASTNode") -> None:
+        self.children.append(child)
+        child.parent = self
+        if child.tokens is not None:
+            self.tokens.extend(child.tokens)
+
+    def __eq__(self, __value: object) -> bool:
+        """
+        Compare two ASTNode objects.
+
+        Notes
+        -----
+        The `parent`, `tokens` attribute is not compared.
+        """
+        return (
+            isinstance(__value, ASTNode)
+            and self.type == __value.type
+            and self.value == __value.value
+            and self.children == __value.children
+        )
+
+
+class AST(ASTNode):
+    def __init__(self, children: Optional[List[ASTNode]] = None, source: str = ""):
+        super().__init__(ASTType.ROOT, children=children)
+        self.source = source
+
+
+# ASC values
+
+
+class ASCNode(NamedTuple):
+    x: float
+    y: float
+    z: float
+    r: float
+
+
+class ASCColor(NamedTuple):
+    color: str
+
+    def __eq__(self, __value: object) -> bool:
+        return (
+            isinstance(__value, ASCColor)
+            and self.color.upper() == __value.color.upper()
+        )
+
+
+class ASCComment(NamedTuple):
+    comment: str
+
+
+# Error
+
+
+class TokenTypeError(ValueError):
+    def __init__(self, token: "Token", expected: str):
+        super().__init__(
+            f"Unexpected token {token.type.name} `{token.value}` at {token.lineno}:{token.column}, expected {expected}"
+        )
+
+
+class LiteralTokenError(ValueError):
+    def __init__(self, token: "Token", expected: str):
+        super().__init__(
+            f"Unexpected LITERAL token {token.value} at {token.lineno}:{token.column}, expected {expected}"
+        )
+
+
+class AssertionTokenTypeError(Exception):
+    pass
+
+
+# Parser
+
+
+class Parser:
+    def __init__(self, r: TextIOBase, *, source: str = ""):
+        self.lexer = Lexer(r)
+        self.next_token = None
+        self.source = source
+        self._read_token()
+
+    def parse(self) -> AST:
+        try:
+            return self._parse()
+        except AssertionTokenTypeError as assertion_err:
+            msg = (
+                f"Error parsing {self.source}" if self.source != "" else "Error parsing"
+            )
+            original_error = assertion_err.__cause__
+            err = ValueError(msg)
+            if original_error is None:
+                raise err
+
+            ignores = ["_assert_and_cunsume", "_assert"]
+            current = assertion_err.__traceback__
+            while current is not None:
+                if (
+                    current.tb_next is not None
+                    and current.tb_next.tb_frame.f_code.co_name in ignores
+                ):
+                    current.tb_next = None
+                else:
+                    current = current.tb_next
+
+            original_error.__traceback__ = assertion_err.__traceback__
+
+            raise err from original_error
+        except Exception as original_error:
+            msg = f"Error parsing {self.source}" if self.source else "Error parsing"
+            raise ValueError(msg) from original_error
+
+    def _parse(self) -> AST:
+        root = AST(source=self.source)
+
+        token = self._assert_and_cunsume(TokenType.BRACKET_LEFT)
+        root.tokens.append(token)
+
+        while (token := self.next_token) is not None:
+            if token.type == TokenType.BRACKET_RIGHT:
+                break
+
+            if token.type != TokenType.BRACKET_LEFT:
+                raise TokenTypeError(token, "BRACKET_LEFT, BRACKET_RIGHT")
+
+            root.tokens.append(token)
+            self._consume()
+
+            token = self._assert(self.next_token, TokenType.LITERAL)
+            match str.upper(token.value):
+                case "AXON" | "DENDRITE":
+                    self._parse_tree(root)
+
+                case "COLOR":
+                    self._parse_color(root)  # TODO: bug
+
+                case _:
+                    raise LiteralTokenError(token, "AXON, DENDRITE, COLOR")
+
+        token = self._assert(self.next_token, TokenType.BRACKET_RIGHT)
+        token = self._assert_and_cunsume(TokenType.BRACKET_RIGHT)
+        root.tokens.append(token)
+        return root
+
+    def _parse_tree(self, root: ASTNode) -> None:
+        t1 = self._assert_and_cunsume(TokenType.LITERAL)
+        node = ASTNode(ASTType.TREE, str.upper(t1.value), tokens=[t1])
+
+        t2 = self._assert_and_cunsume(TokenType.BRACKET_RIGHT)
+        node.tokens.append(t2)
+
+        t3 = self._assert_and_cunsume(TokenType.BRACKET_LEFT)
+        node.tokens.append(t3)
+
+        self._parse_subtree(node)
+        root.add_child(node)
+
+    def _parse_subtree(self, root: ASTNode) -> None:
+        flag = True  # flag to check if the brachet_left can be consumed
+        current = root
+        while (token := self.next_token) is not None:
+            match token.type:
+                case TokenType.BRACKET_LEFT:
+                    self._read_token()
+                    if flag:
+                        flag = False
+                    else:
+                        self._parse_subtree(current)
+
+                case TokenType.BRACKET_RIGHT:
+                    break
+
+                case TokenType.FLOAT:
+                    current = self._parse_node(current)
+                    flag = True
+
+                case TokenType.LITERAL:
+                    match str.upper(token.value):
+                        case "COLOR":
+                            self._parse_color(current)
+                        case _:
+                            raise LiteralTokenError(token, "COLOR")
+
+                    flag = True
+
+                case TokenType.OR:
+                    current = root
+                    self._read_token()
+                    flag = True
+
+                case TokenType.COMMENT:
+                    self._parse_comment(current)
+
+                case _:
+                    excepted = (
+                        "BRACKET_LEFT, BRACKET_RIGHT, LITERAL, FLOAT, OR, COMMENT"
+                    )
+                    raise TokenTypeError(token, excepted)
+
+            current.tokens.append(token)
+
+    def _parse_node(self, root: ASTNode) -> ASTNode:
+        # FLOAT FLOAT FLOAT FLOAT )
+        t1 = self._assert_and_cunsume(TokenType.FLOAT)
+        t2 = self._assert(self.next_token, TokenType.FLOAT)
+        self._read_token()
+        t3 = self._assert(self.next_token, TokenType.FLOAT)
+        self._read_token()
+        t4 = self._assert(self.next_token, TokenType.FLOAT)
+        self._read_token()
+        t5 = self._assert_and_cunsume(TokenType.BRACKET_RIGHT)
+
+        x, y, z, r = t1.value, t2.value, t3.value, t4.value
+        node = ASTNode(ASTType.NODE, ASCNode(x, y, z, r), tokens=[t1, t2, t3, t4, t5])
+        root.add_child(node)
+        return node
+
+    def _parse_color(self, root: ASTNode) -> ASTNode:
+        # COLOR COLOR_VALUE )
+        t1 = self._assert_and_cunsume(TokenType.LITERAL)
+        t2 = self._assert_and_cunsume(TokenType.LITERAL)
+        t3 = self._assert_and_cunsume(TokenType.BRACKET_RIGHT)
+
+        node = ASTNode(ASTType.COLOR, ASCColor(t2.value), tokens=[t1, t2, t3])
+        root.add_child(node)
+        return node
+
+    def _parse_comment(self, root: ASTNode) -> ASTNode:
+        # ; COMMENT
+        t1 = self._assert_and_cunsume(TokenType.COMMENT)
+        node = ASTNode(ASTType.COMMENT, ASCComment(t1.value), tokens=[t1])
+        root.add_child(node)  # ? where the comment should be added
+        return node
+
+    def _read_token(self) -> None:
+        self.next_token = next(self.lexer, None)
+
+    def _assert_and_cunsume(self, type: "TokenType") -> "Token":
+        token = self._consume()
+        token = self._assert(token, type)
+        return cast(Token, token)
+
+    def _assert(self, token: "Token | None", type: "TokenType") -> "Token":
+        if token is None:
+            raise AssertionTokenTypeError() from ValueError("Unexpected EOF")
+
+        if token.type != type:
+            raise AssertionTokenTypeError() from TokenTypeError(token, type.name)
+
+        return token
+
+    def _consume(self) -> "Token | None":
+        token = self.next_token
+        self._read_token()
+        return token
+
+
+# -----------------
+# ASC format lexer
+# -----------------
+
+
+class TokenType(Enum):
+    BRACKET_LEFT = auto()
+    BRACKET_RIGHT = auto()
+    COMMENT = auto()
+    OR = auto()
+    FLOAT = auto()
+    LITERAL = auto()
+
+
+class Token:
+    def __init__(self, type: TokenType, value: Any, lineno: int, column: int):
+        self.type = type
+        self.value = value
+        self.lineno = lineno
+        self.column = column
+
+    def __repr__(self) -> str:
+        return f"Token({self.type.name}, {self.value}, Position={self.lineno}:{self.column})"
+
+
+RE_FLOAT = re.compile(r"[-+]?[0-9]*\.?[0-9]+([eE][-+]?[0-9]+)?")
+
+
+class Lexer:
+    def __init__(self, r: TextIOBase):
+        self.r = r
+        self.lineno = 1
+        self.column = 1
+        self.next_char = self.r.read(1)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self) -> Token:
+        match (word := self._read_word()):
+            case "":
+                raise StopIteration
+
+            case "(":
+                return self._token(TokenType.BRACKET_LEFT, word)
+
+            case ")":
+                return self._token(TokenType.BRACKET_RIGHT, word)
+
+            case ";":
+                return self._token(TokenType.COMMENT, self._read_line())
+
+            case "|":
+                return self._token(TokenType.OR, word)
+
+            case _ if RE_FLOAT.match(word) is not None:
+                return self._token(TokenType.FLOAT, float(word))
+
+            case _:
+                return self._token(TokenType.LITERAL, word)
+
+    def _read_char(self) -> bool:
+        self.next_char = self.r.read(1)
+        if self.next_char == "":
+            return False
+
+        if self.next_char == "\n":
+            self.lineno += 1
+            self.column = 1
+        else:
+            self.column += 1
+        return True
+
+    def _read_word(self) -> str:
+        # skip leading spaces
+        while self.next_char != "" and self.next_char in " \t\n":
+            self._read_char()
+
+        token = ""
+        while self.next_char != "" and self.next_char not in " \t\n();|":
+            token += self.next_char
+            self._read_char()
+
+        if token != "":
+            return token
+
+        if self.next_char == "":
+            return ""
+
+        ch = self.next_char
+        self._read_char()
+        return ch
+
+    def _read_line(self) -> str:
+        if self.next_char != "\n":
+            line = self.r.readline()
+            line = self.next_char + line
+            if line.endswith("\n"):
+                line = line[:-1]
+        else:
+            line = ""
+
+        self.lineno += 1
+        self.column = 1
+        self.next_char = self.r.read(1)
+        return line
+
+    def _token(self, type: TokenType, value: Any) -> Token:
+        return Token(type, value, self.lineno, self.column)