crystal.types

 1from typing import Union, Iterator, Tuple, List
 2from itertools import product
 3
 4class Point:
 5    """
 6    Represents a single point in one dimension with an associated unit.
 7    """
 8    def __init__(self, value: float, unit: str) -> None:
 9        self.value = value
10        self.unit = unit
11
12    def __repr__(self) -> str:
13        return f"Point({self.value!r}, unit={self.unit!r})"
14
15
16class Interval:
17    """
18    Represents a 1-D interval divided into a fixed number of equally spaced points.
19    """
20    def __init__(self, start: float, end: float, num_points: int, unit: str) -> None:
21        if num_points < 2:
22            raise ValueError("num_points must be >= 2 to define an interval.")
23        self.start = start
24        self.end = end
25        self.num_points = num_points
26        self.unit = unit
27
28    @property
29    def points(self) -> List[Point]:
30        """
31        Generate the list of Points spanning the interval.
32        """
33        step = (self.end - self.start) / (self.num_points - 1)
34        return [Point(self.start + i * step, unit=self.unit) for i in range(self.num_points)]
35
36    def __iter__(self) -> Iterator[Point]:
37        return iter(self.points)
38
39    def __repr__(self) -> str:
40        return (
41            f"Interval({self.start!r}, {self.end!r}, {self.num_points!r}, "
42            f"unit={self.unit!r})"
43        )
44
45# Allow either a Point or an Interval on each axis
46Axis = Union[Point, Interval]
47
48class BoxOfPoints:
49    """
50    A 3-D “box” defined by three axes (x, y, z), each of which is either
51    a single Point or an Interval. Iterating over a BoxOfPoints yields every
52    Point in the 3D grid (Cartesian product of axes).
53    """
54    def __init__(self,
55                 x_axis: Axis,
56                 y_axis: Axis,
57                 z_axis: Axis) -> None:
58        self.axes: Tuple[Axis, Axis, Axis] = (x_axis, y_axis, z_axis)
59        for ax in self.axes:
60            if not isinstance(ax, (Point, Interval)):
61                raise TypeError(
62                    f"Each axis must be a Point or Interval, got {type(ax)}"
63                )
64
65    def __iter__(self) -> Iterator[Point]:
66        # Convert each axis into a sequence of Points
67        seqs = [([axis] if isinstance(axis, Point) else axis.points)
68                for axis in self.axes]
69        # Cartesian product over x, y, z
70        for px, py, pz in product(*seqs):
71            # Preserve the unit from the first axis
72            yield Point(px.value, py.value, pz.value, unit=px.unit)
73
74    def __repr__(self) -> str:
75        return (
76            f"BoxOfPoints(x_axis={self.axes[0]!r}, "
77            f"y_axis={self.axes[1]!r}, z_axis={self.axes[2]!r})"
78        )