dissertation/where_fi/config/models.py
2025-05-15 16:51:47 +01:00

169 lines
4.8 KiB
Python

from functools import cached_property
from typing import Literal, Self
from pydantic import BaseModel, model_validator
Host = tuple[str, int]
class Preprocessing(BaseModel):
class Bandpass(BaseModel):
lowcut: int
highcut: int
@property
def bounds(self) -> tuple[int, int]:
return (self.lowcut, self.highcut)
bandpass: Bandpass | None = None
subcarrier_step: int = 1
denoising: Literal["none", "median", "mean"] = "median"
denoising_period: float = 1
steps: list[
Literal[
"fill_pilots",
"skip_subcarriers",
"remove_agc",
"remove_sfo",
"remove_sto",
"bandpass",
]
] = ["fill_pilots", "skip_subcarriers", "remove_agc", "remove_sfo"]
@model_validator(mode="after")
def skip_in_steps(self) -> Self:
if "skip_subcarriers" not in self.steps and self.subcarrier_step != 1:
raise ValueError(
"subcarrier_step must be 1 if skip_subcarriers is not in steps"
)
return self
@model_validator(mode="after")
def bandpass_in_steps(self) -> Self:
if "bandpass" not in self.steps and self.bandpass is not None:
raise ValueError(
"bandpass must be in preprocessing steps if bandpass configuration is "
"provided"
)
if "bandpass" in self.steps and self.bandpass is None:
raise ValueError(
"bandpass configuration must be provided if bandpass is in "
"preprocessing steps"
)
return self
class MUSIC(BaseModel):
eigval_threshold: float
window_size: int
class Heatmap(BaseModel):
theta_resolution: int
tof_resolution: int
tof_max: float
heatmap: Heatmap
class Antennas(BaseModel):
spacing: float
order: list[tuple[Host, int]]
@property
def count(self) -> int:
return len(self.order)
class Config(BaseModel):
receive_hosts: list[Host]
transmit_host: Host
antennas: Antennas
collection_sample_rate: int
processing_sample_rate: int
central_freq: int
channel_width: Literal[20, 40, 80, 160]
frame_format: Literal["NOHT", "HT", "VHT", "HESU"]
preprocessing: Preprocessing
music: MUSIC
@cached_property
def central_freq_hz(self) -> int:
return self.central_freq * 1_000_000
@cached_property
def band(self) -> Literal["2.4", "5", "6"]:
if self.central_freq in range(2412, 2484):
return "2.4"
if self.central_freq in range(5180, 5320):
return "5"
if self.central_freq in range(5955, 7115):
return "6"
raise ValueError(f"{self.central_freq} is not a valid Wi-Fi channel")
@cached_property
def num_guards(self) -> int:
if self.channel_width == 20:
return 7
return 11
@cached_property
def subcarriers(self) -> int:
nulls = {
20: 1,
40: 3,
}
return len(self.subcarrier_frequencies) + nulls[self.channel_width]
@cached_property
def _delta_f_no_skipping(self) -> int:
if self.frame_format == "HESU":
return 78_125
return 312_500
@cached_property
def delta_f(self) -> int:
return self._delta_f_no_skipping * self.preprocessing.subcarrier_step
@cached_property
def _subcarriers_no_skipping(self) -> list[int]:
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 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
if self.band == "2.4" and self.channel_width not in [20, 40]:
raise ValueError(
f"2.4GHz channel {self.central_freq} must "
"have a channel width of 20 or 40 MHz"
)
if (self.band == "2.4" and (self.central_freq - band_start) % 5 != 0) or (
self.band in ["5", "6"]
and ((self.central_freq - band_start) % self.channel_width != 0)
):
raise ValueError(
f"Central frequency {self.central_freq} must be a channel as in https://en.wikipedia.org/wiki/List_of_WLAN_channels"
)
return self