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.

simulations/lib/sampling/digitizer.py

@@ -1,4 +1,5 @@
 import numpy as np
+from functools import wraps
 
 from . import sampling as smp
 from .sampler import Sampler
@@ -8,7 +9,7 @@ class Digitizer(Sampler):
     Digitizer that takes in a signal and resamples and quantises the signal.
     """
 
-    def __init__(self, resolution=0.1, sampling_frequency=None):
+    def __init__(self, resolution=0.1, bias=0, sampling_frequency=None):
         """
 
         Parameters
@@ -21,6 +22,7 @@ class Digitizer(Sampler):
         super().__init__(sampling_frequency)
 
         self.resolution = resolution
+        self.bias = bias
 
     def digitise(self, signal, signal_sample_frequency=None):
         """
@@ -28,9 +30,23 @@ class Digitizer(Sampler):
 
         Effectively resamples signal
         """
+
+        if callable(signal):
+            @wraps(signal)
+            def func(*args, **kwargs):
+                return smp.quantise(
+                        self.sample(signal(*args, **kwargs), signal_sample_frequency),
+                        self.resolution,
+                        self.bias
+                        )
+
+            return func
+
+        else:
             signal = np.asarray(signal)
 
             return smp.quantise(
                 self.sample(signal, signal_sample_frequency),
-            self.resolution
+                self.resolution,
+                self.bias
             )