Simu: Antenna encapsulate a Digitizer

simulations/lib/location/antenna.py

@@ -1,7 +1,7 @@
 from functools import partial
 
 from .location import Location
-from ..signals import Signal
+from ..sampling import Digitizer
 
 class Antenna(Location):
     """
@@ -9,23 +9,59 @@ class Antenna(Location):
 
     Either emitting or receiving.
 
-    Optionally applies a transformation to the traced signal.
+    Optionally uses digitizer to transform the signal
+    when receiving.
     """
+    def __init__(self, x, digitizer=None):
+        super().__init__(x)
+
+        self.digitizer = digitizer
+
     def __repr__(self):
-        return "Antenna({})".format(repr(self.x))
+        return "Antenna({}, {})".format(repr(self.x), repr(self.x))
 
-    def emit(self, signal: Signal) -> callable:
-        return partial(signal, x_0=self.x)
+    def _digitise_partial(self, signal: callable, *args, digitise=True, **kwargs) -> callable:
+        """
+        A wrapper around functools.partial to support optionally
+        digitising the returned signal.
+        """
+        if self.digitizer and digitise:
+            signal = self.digitizer.digitise(signal)
 
-    def recv(self, signal: Signal) -> callable:
+        return partial(signal, *args, **kwargs)
+
+
+    def emit(self, signal: callable, digitise=False) -> callable:
+        """
+        Return a function that emits a signal from the antenna's location
+        """
+        return self._digitise_partial(signal, x_0=self.x, digitise=digitise)
+
+    def recv(self, signal: callable, digitise=True) -> callable:
         """
         Return a function that traces the signal as a function of time
         at the antenna's location
         """
-        return partial(signal, x_f=self.x)
+        return self._digitise_partial(signal, x_f=self.x, digitise=digitise)
 
     receive = recv
 
+    # math
+    def __add__(self, other):
+        if isinstance(other, Location):
+            other = other.x
+
+        return self.__class__(self.x + other, self.digitizer)
+
+    def __sub__(self, other):
+        if isinstance(other, Location):
+            other = other.x
+
+        return self.__class__(self.x - other, self.digitizer)
+
+    def __mul__(self, other):
+        return self.__class__(self.x * other, self.digitizer)
+
 class Receiver(Antenna):
     """
     An antenna which main purpose is to trace a signal over time.