dissertation/where_fi/collection/raytracing/__init__.py

import logging
import time
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Iterator

import numpy as np
import numpy.typing as npt

from where_fi.collection.protocols import CSIProducer, MergedCSI
from where_fi.config import config

C = 299_792_458
MAX_DEPTH = 2
LOSS_EXPONENT = 0.8

logger = logging.getLogger(__name__)


@dataclass
class PathComponent:
    delay: float
    phase: float
    attenuation: float
    frequency: float


class ChannelImpulseResponse:
    def __init__(self, path_components: list[PathComponent]) -> None:
        self.path_components = path_components

    @staticmethod
    def delayed(
        other: "ChannelImpulseResponse", delay: float, reflect: bool = False
    ) -> "ChannelImpulseResponse":
        new_path_components: list[PathComponent] = []
        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
                    + 2 * np.pi * component.frequency * delay
                    + (np.pi if reflect else 0),
                    attenuation=component.attenuation
                    * (np.exp(-distance * LOSS_EXPONENT)),
                    frequency=component.frequency,
                )
            )
        return ChannelImpulseResponse(new_path_components)

    def __add__(self, other: "ChannelImpulseResponse") -> "ChannelImpulseResponse":
        return ChannelImpulseResponse(self.path_components + other.path_components)


class Object:
    def __init__(self, x: float, y: float) -> None:
        self.x = x
        self.y = y
        self.cir = ChannelImpulseResponse([])

    def distance(self, other: "Object") -> float:
        return ((self.x - other.x) ** 2 + (self.y - other.y) ** 2) ** 0.5

    def reflect(
        self,
        incoming_cir: ChannelImpulseResponse,
        target: list["Object"],
        depth: int = 0,
    ) -> None:
        self.cir = self.cir = self.cir + incoming_cir

        if depth > MAX_DEPTH:
            return

        for obj in target:
            if id(obj) == id(self):
                continue
            distance = self.distance(obj)
            assert distance > 0
            delay = distance / C
            obj.reflect(
                ChannelImpulseResponse.delayed(incoming_cir, delay, reflect=True),
                target,
                depth + 1,
            )


class Transmitter(Object):
    def __init__(self, x: float, y: float) -> None:
        super().__init__(x, y)


DELTA_T = 10
GAMMA = np.pi / 4


class Receiver(Object):
    def __init__(self, x: float, y: float, ideal: bool = True) -> None:
        super().__init__(x, y)
        self.delta_t = 0 if ideal else DELTA_T
        self.gamma = 0 if ideal else GAMMA

    def get_cfr(self) -> npt.NDArray[np.complex64]:
        """
        Calculate the CSI matrix for this rx-tx pair. This is computed by the Fourier
        Transform of the Channel Impulse Response (CIR). The CIR is calculated as in
        [1], [2].

        To get the FT of the CIR, we use the sifting property of the Dirac delta

        [1] - https://tns.thss.tsinghua.edu.cn/wst/docs/pre
        [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):
            cfr[i_sub] = sum(
                [
                    component.attenuation
                    * np.exp(1j * component.phase)
                    * np.exp(-1j * subcarrier * component.delay)
                    for component in self.cir.path_components
                ]
            ) * np.exp(
                1j
                * (
                    2
                    * np.pi
                    * (i_sub / len(config.subcarrier_frequencies))
                    * self.delta_t
                    + self.gamma
                )
            )
        return cfr


class Environment:
    def __init__(self, objects: list[Object]) -> 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
        ]

    @staticmethod
    def from_config(filename: Path) -> "Environment":
        """
        Read a config file that includes a scene description and create an envionment
        based on that

        The config file should be a text file where each line corresponds to an object.
        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] = []
        with open(filename, "r") as f:
            for line in f.readlines():
                if line.startswith("#"):
                    continue
                parts = line.split(" ")
                x, y = map(float, parts[1:])
                if parts[0] == "TX":
                    objects.append(Transmitter(x, y))
                elif parts[0] == "RX":
                    objects.append(Receiver(x, y))
                else:
                    objects.append(Object(x, y))

        return Environment(objects)

    def get_csi(self) -> npt.NDArray[np.complex64]:
        """
        Calculate the Channel State Information (CSI) matrix for the simulated
        environment.

        The returned matrix is of shape (num_subcarriers, num_receivers,
                                         num_transmitters)

        This is calculated by finding all paths leading to each receiver, and
        calculating the CFR evaluated at each subcarrier.
        """
        csi = np.zeros(
            (
                len(config.subcarrier_frequencies),
                len(self.receivers),
                len(self.transmitters),
            ),
            dtype=np.complex64,
        )
        for i_tx, transmitter in enumerate(self.transmitters):
            for obj in self.objects + self.receivers + self.transmitters:
                obj.cir = ChannelImpulseResponse([])
            transmitter.reflect(
                ChannelImpulseResponse(
                    [
                        PathComponent(
                            delay=0,
                            phase=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)}")
                csi[:, i_rx, i_tx] = receiver.get_cfr()
        return csi


class SimulatedCSIProducer(CSIProducer):
    pass

    def __init__(self, environment: Environment) -> None:
        self.environment = environment

    def __call__(self) -> Iterator[MergedCSI]:
        while True:
            start = datetime.now()
            csi = self.environment.get_csi()
            yield MergedCSI(
                frames={},
                matrix=csi,
            )
            time.sleep(
                max(
                    0,
                    1 / config.collection_sample_rate
                    - (datetime.now() - start).total_seconds(),
                )
            )