optimise raytracer

where_fi/application.py

@@ -6,7 +6,7 @@ from typing import Any, Callable, NamedTuple
 
 import numpy.typing as npt
 
-from where_fi.collection import CSIMatrix, ingest
+from where_fi.collection import CSIMatrix, NoopCSIProducer, ingest
 from where_fi.collection.csi_frame import CSI
 from where_fi.collection.protocols import CSIProducer, MergedCSI
 from where_fi.config import config
@@ -70,10 +70,10 @@ class CSIApplication:
                         mostly used as a scheduler).
     """
 
-    def __init__(self, producer: CSIProducer, visualise_raw: bool = False) -> None:
+    def __init__(self, visualise_raw: bool = False) -> None:
         self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")
 
-        self.producer = producer
+        self.producer = NoopCSIProducer()
 
         self.webapp_queue: "mp.Queue[visualise.VisualiserData]" = mp.Queue(
             config.processing_sample_rate
@@ -94,6 +94,13 @@ class CSIApplication:
             visualise.figures.FigureId, visualise.figures.SpecificFigure
         ] = {}
 
+    def set_producer(self, producer: CSIProducer) -> None:
+        """
+        Set the producer for the application. This is used to change the data source
+        at runtime.
+        """
+        self.producer = producer
+
     def register_figure(
         self,
         figure_id: visualise.figures.FigureId,
@@ -109,7 +116,6 @@ class CSIApplication:
         Update a visualisation with the given data. The available visualisations are as
         per visualise.server.all_figures.
         """
-        if not self.webapp_queue.full():
         self.webapp_queue.put(visualise.VisualiserData(data, dtype))
 
     def on_raw(self, func: Callable[[CSIMatrix], None]) -> Callable[[CSIMatrix], None]:

where_fi/collection/raytracing/__init__.py

@@ -21,9 +21,8 @@ logger = logging.getLogger(__name__)
 @dataclass
 class PathComponent:
     delay: float
-    phase: float
+    phase_offset: float
     attenuation: float
-    frequency: float
 
 
 class ChannelImpulseResponse:
@@ -38,16 +37,14 @@ class ChannelImpulseResponse:
         distance = delay * C
 
         for component in other.path_components:
-            wavelength = C / component.frequency
             new_path_components.append(
                 PathComponent(
                     delay=component.delay + delay,
-                    phase=component.phase
+                    phase_offset=component.phase_offset + (np.pi if reflect else 0),
-                    + 2 * np.pi * component.frequency * delay
-                    + (np.pi if reflect else 0),
                     attenuation=component.attenuation
-                    * (np.exp(-distance * LOSS_EXPONENT)),
+                    * (1 / (1 + distance) ** LOSS_EXPONENT),
-                    frequency=component.frequency,
+                    # attenuation=component.attenuation
+                    # * (np.exp(-distance * LOSS_EXPONENT)),
                 )
             )
         return ChannelImpulseResponse(new_path_components)
@@ -56,22 +53,25 @@ class ChannelImpulseResponse:
         return ChannelImpulseResponse(self.path_components + other.path_components)
 
 
-class Object:
+class PathObject:
-    def __init__(self, x: float, y: float) -> None:
+    def __init__(self, x: float, y: float, z: float) -> None:
         self.x = x
         self.y = y
+        self.z = z
         self.cir = ChannelImpulseResponse([])
 
-    def distance(self, other: "Object") -> float:
+    def distance(self, other: "PathObject") -> float:
-        return ((self.x - other.x) ** 2 + (self.y - other.y) ** 2) ** 0.5
+        return (
+            (self.x - other.x) ** 2 + (self.y - other.y) ** 2 + (self.z - other.z) ** 2
+        ) ** 0.5
 
     def reflect(
         self,
         incoming_cir: ChannelImpulseResponse,
-        target: list["Object"],
+        target: list["PathObject"],
         depth: int = 0,
     ) -> None:
-        self.cir = self.cir = self.cir + incoming_cir
+        self.cir = self.cir + incoming_cir
 
         if depth > MAX_DEPTH:
             return
@@ -89,18 +89,18 @@ class Object:
             )
 
 
-class Transmitter(Object):
+class Transmitter(PathObject):
-    def __init__(self, x: float, y: float) -> None:
+    def __init__(self, x: float, y: float, z: float) -> None:
-        super().__init__(x, y)
+        super().__init__(x, y, z)
 
 
 DELTA_T = 10
 GAMMA = np.pi / 4
 
 
-class Receiver(Object):
+class Receiver(PathObject):
-    def __init__(self, x: float, y: float, ideal: bool = True) -> None:
+    def __init__(self, x: float, y: float, z: float, ideal: bool = True) -> None:
-        super().__init__(x, y)
+        super().__init__(x, y, z)
         self.delta_t = 0 if ideal else DELTA_T
         self.gamma = 0 if ideal else GAMMA
 
@@ -116,12 +116,12 @@ class Receiver(Object):
         [2] - https://dl.acm.org/doi/10.1145/2543581.2543592, Equation 4
         """
         cfr = np.zeros(len(config.subcarrier_frequencies), dtype=np.complex64)
-        for i_sub, subcarrier in enumerate(config.subcarrier_frequencies):
+        for i_sub, f_sub in enumerate(config.subcarrier_frequencies):
             cfr[i_sub] = sum(
                 [
                     component.attenuation
-                    * np.exp(1j * component.phase)
+                    * np.exp(1j * component.phase_offset)
-                    * np.exp(-1j * subcarrier * component.delay)
+                    * np.exp(-1j * 2 * np.pi * f_sub * component.delay)
                     for component in self.cir.path_components
                 ]
             ) * np.exp(
@@ -138,12 +138,20 @@ class Receiver(Object):
 
 
 class Environment:
-    def __init__(self, objects: list[Object]) -> None:
+    def __init__(self, objects: list[PathObject]) -> None:
         self.transmitters = [obj for obj in objects if isinstance(obj, Transmitter)]
         self.receivers = [obj for obj in objects if isinstance(obj, Receiver)]
         self.objects = [
             obj for obj in objects if obj not in self.transmitters + self.receivers
         ]
+        self.count = 0
+
+    def move(self) -> None:
+        self.count += 1
+        if self.count > 300:
+            print("Moving objects")
+            for obj in self.objects:
+                obj.x += np.sin(self.count / 100 * 2 * np.pi) * 0.1
 
     @staticmethod
     def from_config(filename: Path) -> "Environment":
@@ -155,22 +163,33 @@ class Environment:
         The first word of each line should be the type of object (TX/RX/OBJ), followed
         by the x and y coordinates of the object.
         """
-        objects: list[Object] = []
+        objects: list[PathObject] = []
         with open(filename, "r") as f:
             for line in f.readlines():
                 if line.startswith("#"):
                     continue
                 parts = line.split(" ")
-                x, y = map(float, parts[1:])
+                x, y, z = map(float, parts[1:])
                 if parts[0] == "TX":
-                    objects.append(Transmitter(x, y))
+                    objects.append(Transmitter(x, y, z))
                 elif parts[0] == "RX":
-                    objects.append(Receiver(x, y))
+                    objects.append(Receiver(x, y, z))
                 else:
-                    objects.append(Object(x, y))
+                    objects.append(PathObject(x, y, z))
 
         return Environment(objects)
 
+    def add_awgn(
+        self, signal: npt.NDArray[np.complex64], snr_dB: float
+    ) -> npt.NDArray[np.complex64]:
+        signal_power = np.mean(np.abs(signal) ** 2)
+        snr_linear = 10 ** (snr_dB / 10)
+        noise_power = signal_power / snr_linear
+        noise = np.sqrt(noise_power / 2) * (
+            np.random.randn(*signal.shape) + 1j * np.random.randn(*signal.shape)
+        )
+        return signal + noise.astype(np.complex64)
+
     def get_csi(self) -> npt.NDArray[np.complex64]:
         """
         Calculate the Channel State Information (CSI) matrix for the simulated
@@ -198,19 +217,17 @@ class Environment:
                     [
                         PathComponent(
                             delay=0,
-                            phase=0,
+                            phase_offset=0,
                             # attenuation=100000000,
                             attenuation=100,
-                            frequency=subcarrier,
                         )
-                        for subcarrier in config.subcarrier_frequencies
                     ]
                 ),
                 self.objects + self.receivers,
             )
             for i_rx, receiver in enumerate(self.receivers):
-                logger.info(f"RX {i_rx} paths: {len(receiver.cir.path_components)}")
+                logger.debug(f"RX {i_rx} paths: {len(receiver.cir.path_components)}")
-                csi[:, i_rx, i_tx] = receiver.get_cfr()
+                csi[:, i_rx, i_tx] = self.add_awgn(receiver.get_cfr(), snr_dB=30)
         return csi
 
 
@@ -219,9 +236,11 @@ class SimulatedCSIProducer(CSIProducer):
 
     def __init__(self, environment: Environment) -> None:
         self.environment = environment
+        self.active = True
 
     def __call__(self) -> Iterator[MergedCSI]:
-        while True:
+        while self.active:
+            self.environment.move()
             start = datetime.now()
             csi = self.environment.get_csi()
             yield MergedCSI(
@@ -235,3 +254,6 @@ class SimulatedCSIProducer(CSIProducer):
                     - (datetime.now() - start).total_seconds(),
                 )
             )
+
+    def stop(self) -> None:
+        self.active = False

where_fi/config/models.py

@@ -1,3 +1,4 @@
+from functools import cached_property
 from typing import Literal, Self
 
 from pydantic import BaseModel, model_validator
@@ -90,11 +91,11 @@ class Config(BaseModel):
     preprocessing: Preprocessing
     music: MUSIC
 
-    @property
+    @cached_property
     def central_freq_hz(self) -> int:
         return self.central_freq * 1_000_000
 
-    @property
+    @cached_property
     def band(self) -> Literal["2.4", "5", "6"]:
         if self.central_freq in range(2412, 2484):
             return "2.4"
@@ -104,36 +105,50 @@ class Config(BaseModel):
             return "6"
         raise ValueError(f"{self.central_freq} is not a valid Wi-Fi channel")
 
-    @property
+    @cached_property
     def num_guards(self) -> int:
         if self.channel_width == 20:
             return 7
         return 11
 
-    @property
+    @cached_property
     def subcarriers(self) -> int:
-        return (
+        nulls = {
-            int((self.channel_width * 1e6) // self.delta_f) - self.num_guards
+            20: 1,
-        ) // self.preprocessing.subcarrier_step
+            40: 3,
+        }
+        return len(self.subcarrier_frequencies) + nulls[self.channel_width]
 
-    @property
+    @cached_property
     def _delta_f_no_skipping(self) -> int:
         if self.frame_format == "HESU":
             return 78_125
         return 312_500
 
-    @property
+    @cached_property
     def delta_f(self) -> int:
         return self._delta_f_no_skipping * self.preprocessing.subcarrier_step
 
-    @property
+    @cached_property
-    def subcarrier_frequencies(self) -> list[int]:
+    def _subcarriers_no_skipping(self) -> list[int]:
-        num_subcarriers = self.channel_width * 1_000_000 // self._delta_f_no_skipping
+        used = {
+            20: (1, 29),
+            40: (2, 59),
+        }
+        subcarrier_indices = list(
+            range(-used[self.channel_width][1] + 1, -used[self.channel_width][0] + 1)
+        ) + list(range(used[self.channel_width][0], used[self.channel_width][1]))
+        print(subcarrier_indices)
+
         return [
             self.central_freq_hz + i * self._delta_f_no_skipping
-            for i in range(-num_subcarriers // 2, num_subcarriers // 2)
+            for i in subcarrier_indices
         ]
 
+    @cached_property
+    def subcarrier_frequencies(self) -> list[int]:
+        return self._subcarriers_no_skipping[:: self.preprocessing.subcarrier_step]
+
     @model_validator(mode="after")
     def channels(self) -> Self:
         band_start = 2412 if self.band == "2.4" else 5180 if self.band == "5" else 5955