Skip to content

Commit

Permalink
Coordinate quality of life improvements (#235)
Browse files Browse the repository at this point in the history
* Add `__repr__` methods to `CoordinateTransform` methods
* Add `from_axis_names` creation method to `CoordinateSpace`
  • Loading branch information
jp-dark authored Oct 14, 2024
1 parent edd8ab2 commit eee3130
Show file tree
Hide file tree
Showing 2 changed files with 48 additions and 1 deletion.
41 changes: 40 additions & 1 deletion python-spec/src/somacore/coordinates.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,13 @@

import abc
import collections.abc
from typing import Optional, Sequence, Tuple, Union
import itertools
from typing import Iterable, Optional, Sequence, Tuple, Union

import attrs
import numpy as np
import numpy.typing as npt
from typing_extensions import Self

from .types import str_or_seq_length
from .types import to_string_tuple
Expand Down Expand Up @@ -37,6 +39,10 @@ class CoordinateSpace(collections.abc.Sequence[Axis]):

axes: Tuple[Axis, ...] = attrs.field(converter=tuple)

@classmethod
def from_axis_names(cls, axis_names: Sequence[str,]) -> Self:
return cls(tuple(Axis(name) for name in axis_names)) # type: ignore[misc]

@axes.validator
def _validate(self, _, axes: Tuple[Axis, ...]) -> None:
if not axes:
Expand Down Expand Up @@ -95,6 +101,25 @@ def _check_rmatmul_inner_axes(self, other: "CoordinateTransform"):
f"{type(self).__name__}."
)

@abc.abstractmethod
def _contents_lines(self) -> Iterable[str]:
return
yield

def _my_repr(self) -> Iterable[str]:
yield f"{type(self).__name__}"
yield f" input axes: {self._input_axes}"
yield f" output axes: {self._output_axes}"

def __repr__(self) -> str:
content = self._contents_lines
lines = (
self._my_repr()
if content is None
else itertools.chain(self._my_repr(), self._contents_lines())
)
return "<" + "\n".join(lines) + ">"

@abc.abstractmethod
def __matmul__(self, other: object) -> "CoordinateTransform":
raise NotImplementedError()
Expand Down Expand Up @@ -185,6 +210,10 @@ def __init__(
f"Unexpected shape {self._matrix.shape} for the input affine matrix."
)

def _contents_lines(self) -> Iterable[str]:
yield " augmented matrix:"
yield " " + str(self._matrix).replace("\n", "\n ")

def __matmul__(self, other: object) -> CoordinateTransform:
if not isinstance(other, CoordinateTransform):
raise NotImplementedError(
Expand Down Expand Up @@ -259,6 +288,9 @@ def __init__(

super().__init__(input_axes, output_axes, np.diag(self._scale_factors))

def _contents_lines(self) -> Iterable[str]:
yield f" scales: {self._scale_factors}"

def __matmul__(self, other: object) -> CoordinateTransform:
if not isinstance(other, CoordinateTransform):
raise NotImplementedError(
Expand Down Expand Up @@ -312,6 +344,9 @@ def __init__(
rank = str_or_seq_length(input_axes)
super().__init__(input_axes, output_axes, rank * [self._scale])

def _contents_lines(self) -> Iterable[str]:
yield f" scale: {self._scale}"

def __matmul__(self, other: object) -> CoordinateTransform:
if not isinstance(other, CoordinateTransform):
raise NotImplementedError(
Expand Down Expand Up @@ -361,6 +396,10 @@ def __init__(
):
super().__init__(input_axes, output_axes, 1)

def _contents_lines(self) -> Iterable[str]:
return
yield

def __matmul__(self, other: object) -> CoordinateTransform:
if not isinstance(other, CoordinateTransform):
raise NotImplementedError(
Expand Down
8 changes: 8 additions & 0 deletions python-spec/testing/test_coordinates.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,14 @@ def test_coordinate_space():
assert coord_space[0] == Axis("x", unit="nanometer")


def test_coordiante_space_from_axis_names():
coord_space = CoordinateSpace.from_axis_names(["alpha", "beta"])
assert len(coord_space) == 2
assert coord_space.axis_names == ("alpha", "beta")
assert coord_space[0] == Axis("alpha", unit=None)
assert coord_space[1] == Axis("beta", unit=None)


@pytest.mark.parametrize(
("input", "expected"),
[
Expand Down

0 comments on commit eee3130

Please sign in to comment.