✨ New arithmetic operators on Dimension

src/tqec/templates/scale.py

@@ -1,6 +1,8 @@
 import typing as ty
 from dataclasses import dataclass
 
+from tqec.exceptions import TQECException
+
 
 @dataclass(frozen=True)
 class LinearFunction:
@@ -13,6 +15,33 @@ def __call__(self, x: int) -> int:
     def to_dict(self) -> dict[str, ty.Any]:
         return {"type": type(self).__name__, "slope": self.slope, "offset": self.offset}
 
+    def __add__(self, other: "LinearFunction") -> "LinearFunction":
+        return LinearFunction(self.slope + other.slope, self.offset + other.offset)
+
+    def __sub__(self, other: "LinearFunction") -> "LinearFunction":
+        return LinearFunction(self.slope - other.slope, self.offset - other.offset)
+
+    def __mul__(self, other: int) -> "LinearFunction":
+        return self.__rmul__(other)
+
+    def __rmul__(self, other: int) -> "LinearFunction":
+        return LinearFunction(other * self.slope, other * self.offset)
+
+    def invert(self, value: int) -> int:
+        if self.slope == 0:
+            raise TQECException(
+                f"Cannot invert {self}: the linear function is constant and so "
+                f"has an infinite set of inverse for {self.offset} and no inverse "
+                "for any other value."
+            )
+        offset_value = value - self.offset
+        if offset_value % self.slope != 0:
+            raise TQECException(
+                f"{self} cannot invert exactly {value} as this would lead to "
+                f"a non-integer value ({offset_value / self.slope})."
+            )
+        return (value - self.offset) // self.slope
+
 
 class Dimension:
     def __init__(
@@ -64,6 +93,27 @@ def to_dict(self) -> dict[str, ty.Any]:
             "scaling_function": self._scaling_function.to_dict(),
         }
 
+    def __add__(self, other: "Dimension") -> "Dimension":
+        return Dimension(
+            self._scaling_function.invert(self.value),
+            self._scaling_function + other._scaling_function,
+        )
+
+    def __sub__(self, other: "Dimension") -> "Dimension":
+        return Dimension(
+            self._scaling_function.invert(self.value),
+            self._scaling_function - other._scaling_function,
+        )
+
+    def __mul__(self, other: int) -> "Dimension":
+        return other * self
+
+    def __rmul__(self, other: int) -> "Dimension":
+        return Dimension(
+            self._scaling_function.invert(self.value),
+            self._scaling_function * other,
+        )
+
 
 class FixedDimension(Dimension):
     def __init__(self, value: int) -> None:

src/tqec/templates/scale_test.py

@@ -1,4 +1,6 @@
 import pytest
+
+from tqec.exceptions import TQECException
 from tqec.templates.scale import Dimension, FixedDimension, LinearFunction
 
 
@@ -13,6 +15,31 @@ def test_linear_function(slope: int, offset: int):
     assert f(-4) == slope * -4 + offset
 
 
+def test_linear_function_operators():
+    a, b = LinearFunction(2, 5), LinearFunction(3, 1)
+    assert (a + b)(10) == a(10) + b(10)
+    assert (a - b)(3) == a(3) - b(3)
+    assert (3 * a)(54) == 3 * a(54)
+    assert (a * 3)(54) == 3 * a(54)
+
+
+def test_linear_function_invert():
+    a = LinearFunction(2, 5)
+    assert a.invert(a(11)) == 11
+
+
+def test_linear_function_constant_invert():
+    a = LinearFunction(0, 5)
+    with pytest.raises(TQECException):
+        a.invert(5)
+
+
+def test_linear_function_non_exact_invert():
+    a = LinearFunction(2, 0)
+    with pytest.raises(TQECException):
+        a.invert(5)
+
+
 def test_dimension_init():
     dim = Dimension(5, LinearFunction())
     assert dim.value == 5
@@ -38,3 +65,12 @@ def test_fixed_dimension():
     assert dim.value == 3
     assert dim.scale_to(10).value == 3
     assert dim.scale_to(421).value == 3
+
+
+def test_dimension_operators():
+    sfa, sfb = LinearFunction(2, 5), LinearFunction(3, 1)
+    da, db = Dimension(2, sfa), Dimension(2, sfb)
+    assert (da + db).scale_to(10).value == sfa(10) + sfb(10)
+    assert (da - db).scale_to(3).value == sfa(3) - sfb(3)
+    assert (3 * da).scale_to(54).value == 3 * sfa(54)
+    assert (da * 3).scale_to(54).value == 3 * sfa(54)