commit 6a8c8c9c5760ca6cffb5c0a3480dfc3415dedb7e
parent 7200e95666aacf692d4b0e87f0fab1f16249302d
Author: Egor Achkasov <eaachkasov@gmail.com>
Date: Thu, 16 Oct 2025 13:22:12 +0200
mv load_data.py measurements.py; Improve tf solution; Fix np solution measurements reshape
Diffstat:
6 files changed, 362 insertions(+), 497 deletions(-)
diff --git a/README.md b/README.md
@@ -18,9 +18,9 @@ Find the eact parameters in the code (`src/??_paths.py`). Unfortunately, the mea
# Repository structure
- `scene`: **Simplified** version of the [scene](https://github.com/akdd11/anomaly-detection-calibrated-raytracing/tree/main/scenes/scene0).
-- `src/load_data.py`: loader function for the measurements dataset.
+- `src/measurements.py`: loader function for the measurements dataset.
- `src/np_paths.py`: numpy-based implementation of the method. Note that the gradient is calculated naively, what is a serious performance bottleneck. This solution can be used for testing, but strugles with production of any results in a reasonable time with bigger parameter domain size.
-- `src/tf_paths.py`: same, but the gradient is calculated using Tensorflow's tape, is in the [NVidia material calibration notebook](https://github.com/NVlabs/diff-rt-calibration/blob/main/notebooks/Learned_Materials.ipynb). Slightly more convoluted than the numpy version and a pain to debug.
+- `src/tf_paths.py`: same, but the gradient is calculated using Tensorflow's tape, is in the [NVidia material calibration notebook](https://github.com/NVlabs/diff-rt-calibration/blob/main/notebooks/Learned_Materials.ipynb). Slightly more convoluted than the numpy version and a pain to debug. Currently the graph mode is broken leading to the gradient not being calculated. Finding a way to fix it would solve the np solution's performance issue.
- `src/v.py`: visualization script. Described below.
# Algorithm
diff --git a/src/load_data.py b/src/load_data.py
@@ -1,174 +0,0 @@
-"""
-This file contains an implementation of a data loader for
-the measurement data used from doi: 10.1109/JCS64661.2025.10880629
-
-The data directory should contain files named in the format:
-"Round_<round_idx>_AP_<ap_idx>_RF_<rf_idx>_Sec_<sec_idx>.mat"
-where <round_idx>, <ap_idx>, <rf_idx>, and <sec_idx> are integers.
-The <round_idx> is the round index, <ap_idx> is the access point index,
-<rf_idx> is the RF index, and <sec_idx> is the second index.
-
-Each file must be loadable with scipy.io.loadmat
-and must contain a field named 'cirs' with the shape (512, 1000) each.
-"""
-
-import numpy as np
-from scipy.io import loadmat as scipy_io_loadmat
-
-from typing import Tuple, Dict
-from os import listdir as os_listdir
-from os.path import join as os_path_join
-from functools import reduce as functools_reduce
-from functools import partial as functools_partial
-from itertools import product as itertools_product
-from multiprocessing import Pool as multiprocessing_Pool
-
-
-def load_measurements(
- data_dirpath: str,
- load_only_round: int | None = None
- ) -> Tuple[
- np.ndarray,
- Dict[int, int],
- Dict[int, int],
- Dict[int, int],
- Dict[int, int]
- ]:
- """Load the measurements from the given data directory.
- Args:
- data_dirpath (str): The path to the directory containing the files.
- Returns:
- Tuple containing:
- - cirs_truth (np.ndarray): The channel impulse responses.
- Shape: (num_rounds, num_aps, num_rfs, num_secs, 512, 1000)
- - round_idx_map (dict): Mapping from round to array indices.
- - ap_idx_map (dict): Mapping from access point to array indices.
- - rf_idx_map (dict): Mapping from RF to array indices.
- - sec_idx_map (dict): Mapping from second to array indices.
- Note:
- - The measurements filenames are expected to be in the format:
- "Round_<round_idx>_AP_<ap_idx>_RF_<rf_idx>_Sec_<sec_idx>.mat"
- - In case of missing seconds for some rounds, aps, or rfs,
- inner join is performed on the available seconds.
- Example: round 1 has seconds 0, 1, 2, 3,
- round 2 has seconds 1, 2, 3, 4, then
- the available seconds for both rounds will be 1, 2, 3.
- - Example usage of the maps:
- ```python
- round_idx = 1
- ap_idx = 2
- rf_idx = 0
- sec_idx = 3
- array_idx = (
- round_idx_map[round_idx],
- ap_idx_map[ap_idx],
- rf_idx_map[rf_idx],
- sec_idx_map[sec_idx]
- )
- cfr = cfrs[array_idx]
- ```
- """
-
- # Get the rounds, aps, rfs, secs from the filenames
- rounds = set()
- aps = set()
- rfs = set()
- secs = {}
- for filename in os_listdir(data_dirpath):
- if not filename.endswith(".mat") and not filename.startswith("Round_"):
- continue
- s = filename[6:-4].split("_")
- if len(s) != 7:
- continue
- round_idx, ap_idx, rf_idx, sec_idx = map(
- int,
- (s[0], s[2], s[4], s[6])
- )
- if load_only_round is not None and round_idx != load_only_round:
- continue
- rounds.add(round_idx)
- aps.add(ap_idx)
- rfs.add(rf_idx)
- secs[(round_idx, ap_idx, rf_idx)] = secs.get(
- (round_idx, ap_idx, rf_idx), set())
- secs[(round_idx, ap_idx, rf_idx)].add(sec_idx)
- rounds = sorted(rounds)
- aps = sorted(aps)
- rfs = sorted(rfs)
- # Inner join of all available seconds
- secs = sorted(functools_reduce(
- set.intersection,
- (s for s in secs.values())
- ))
-
- # Get num_rounds, num_aps, num_rfs, num_seconds
- num_rounds = len(rounds)
- num_aps = len(aps)
- num_rfs = len(rfs)
- num_secs = len(secs)
-
- # Make indexes maps for rounds, aps, rfs, secs
- round_idx_map = {r: i for i, r in enumerate(rounds)}
- ap_idx_map = {ap: i for i, ap in enumerate(aps)}
- rf_idx_map = {rf: i for i, rf in enumerate(rfs)}
- sec_idx_map = {sec: i for i, sec in enumerate(secs)}
-
- # Load the measurements
- with multiprocessing_Pool() as pool:
- func = functools_partial(
- load_and_store,
- data_dirpath=data_dirpath,
- round_idx_map=round_idx_map,
- ap_idx_map=ap_idx_map,
- rf_idx_map=rf_idx_map,
- sec_idx_map=sec_idx_map
- )
- results = pool.map(
- func,
- itertools_product(rounds, aps, rfs, secs)
- )
- cirs = np.empty(
- (num_rounds, num_aps, num_rfs, num_secs, 512, 1000),
- dtype=np.complex128
- )
- for i, j, k, l, cir in results:
- cirs[i, j, k, l, :] = cir
-
- # cirs = np.empty(
- # (num_rounds, num_aps, num_rfs, num_secs, 512, 1000),
- # dtype=np.complex128
- # )
- # for r, ap, rf, sec in itertools_product(rounds, aps, rfs, secs):
- # filename = f"Round_{r}_AP_{ap}_RF_{rf}_Sec_{sec}.mat"
- # filepath = os_path_join(data_dirpath, filename)
- # cir = scipy_io_loadmat(filepath)['cirs']
- # assert cir.shape == (512, 1000)
- # cirs[
- # round_idx_map[r],
- # ap_idx_map[ap],
- # rf_idx_map[rf],
- # sec_idx_map[sec],
- # :
- # ] = cir
-
- return cirs, round_idx_map, ap_idx_map, rf_idx_map, sec_idx_map
-
-
-def load_and_store(
- args,
- data_dirpath, round_idx_map, ap_idx_map, rf_idx_map, sec_idx_map
-):
- r, ap, rf, sec = args
- filepath = os_path_join(
- data_dirpath,
- f"Round_{r}_AP_{ap}_RF_{rf}_Sec_{sec}.mat"
- )
- cir = scipy_io_loadmat(filepath)['cirs']
- assert cir.shape == (512, 1000)
- return (
- round_idx_map[r],
- ap_idx_map[ap],
- rf_idx_map[rf],
- sec_idx_map[sec],
- cir
- )
diff --git a/src/measurements.py b/src/measurements.py
@@ -0,0 +1,100 @@
+"""
+This file contains an implementation of a data loader for
+the measurement data used from doi: 10.1109/JCS64661.2025.10880629
+
+The data directory should contain files named in the format:
+"Round_<round_idx>_AP_<ap_idx>_RF_<rf_idx>_Sec_<sec_idx>.mat"
+where <round_idx>, <ap_idx>, <rf_idx>, and <sec_idx> are integers.
+The <round_idx> is the round index, <ap_idx> is the access point index,
+<rf_idx> is the RF index, and <sec_idx> is the second index.
+
+Each file must be loadable with scipy.io.loadmat
+and must contain a field named 'cirs' with the shape (512, 1000) each.
+"""
+
+import numpy as np
+from scipy.io import loadmat as scipy_io_loadmat
+
+from os import listdir as os_listdir
+from os.path import join as os_path_join
+from functools import reduce as functools_reduce
+from functools import partial as functools_partial
+from itertools import product as itertools_product
+from multiprocessing import Pool as multiprocessing_Pool
+
+
+def load(
+ data_dirpath: str,
+ round_idx: int
+) -> np.ndarray:
+ """Load the measurements round from the given data directory.
+ Args:
+ data_dirpath (str): The path to the directory containing the files.
+ round (int): The round index.
+ Returns:
+ cirs_truth (np.ndarray):
+ The channel impulse responses.
+ Shape: (num_aps, num_rfs, num_secs, 512, 1000)
+ Note:
+ - The measurements filenames are expected to be in the format:
+ "Round_<round_idx>_AP_<ap_idx>_RF_<rf_idx>_Sec_<sec_idx>.mat"
+ """
+
+ # Get the rounds, aps, rfs, secs from the filenames
+ aps = set()
+ rfs = set()
+ secs = {}
+ for filename in os_listdir(data_dirpath):
+ if not filename.endswith(".mat") and not filename.startswith("Round_"):
+ continue
+ s = filename[6:-4].split("_")
+ if len(s) != 7:
+ continue
+ if int(s[0]) != round_idx:
+ continue
+ ap_idx, rf_idx, sec_idx = map(int, (s[2], s[4], s[6]))
+ aps.add(ap_idx)
+ rfs.add(rf_idx)
+ secs[(ap_idx, rf_idx)] = secs.get((ap_idx, rf_idx), set())
+ secs[(ap_idx, rf_idx)].add(sec_idx)
+ # Inner join of all available seconds
+ secs = sorted(functools_reduce(
+ set.intersection,
+ (s for s in secs.values())
+ ))
+
+ # Get num_aps, num_rfs, num_seconds
+ num_aps = len(aps)
+ num_rfs = len(rfs)
+ num_secs = len(secs)
+
+ # Load the measurements
+ with multiprocessing_Pool() as pool:
+ func = functools_partial(
+ load_and_store,
+ data_dirpath=data_dirpath,
+ round_idx=round_idx
+ )
+ results = pool.map(
+ func,
+ itertools_product(aps, rfs, secs)
+ )
+ cirs = np.empty(
+ (num_aps, num_rfs, num_secs, 512, 1000),
+ dtype=np.complex128
+ )
+ for j, k, l, cir in results:
+ cirs[j, k, l, :] = cir
+
+ return cirs
+
+
+def load_and_store(args, data_dirpath, round_idx):
+ ap, rf, sec = args
+ filepath = os_path_join(
+ data_dirpath,
+ f"Round_{round_idx}_AP_{ap}_RF_{rf}_Sec_{sec}.mat"
+ )
+ cir = scipy_io_loadmat(filepath)['cirs']
+ assert cir.shape == (512, 1000)
+ return ap-1, rf, sec-1, cir
diff --git a/src/np_paths.py b/src/np_paths.py
@@ -12,7 +12,7 @@
# - Trains the scene material parameters according to Algorithm 1 in the paper
# - Saves the losses and the trainable parameters (omega) to files (losses.npy and omegas.npy)
-from load_data import load_measurements
+from measurements import load as measurements_load
import numpy as np
import sionna.rt as srt
@@ -233,6 +233,7 @@ if __name__ == "__main__":
g_scene = srt.load_scene(
os.path.join(
os.path.dirname(__file__),
+ "..",
"scene",
"scene.xml")
)
@@ -295,18 +296,17 @@ if __name__ == "__main__":
print("Loading measurements...")
data_dirpath = os.path.join(
os.path.dirname(__file__),
- "..", "data"
+ "..", "..", "data"
)
- g_cirs = load_measurements(data_dirpath, g_cfg_load_only_round)[0]
+ g_cirs = measurements_load(data_dirpath, g_cfg_load_only_round)
g_bandwidth = g_cirs.shape[-2] * g_cfg_subcarrier_spacing # W in the paper
- g_cirs = g_cirs[0].reshape(
- (g_cirs.shape[1],
- g_cirs.shape[2],
- g_cirs.shape[4],
- g_cirs.shape[3] * g_cirs.shape[5],)
- ) # Reshape to (num_aps, num_rfs, num_subcarriers, num_samples)
- # Transpose to (num_samples, num_aps, num_rfs, num_subcarriers)
- g_cirs = np.transpose(g_cirs, (3, 0, 1, 2))
+ g_cirs = g_cirs.transpose((2, 4, 0, 1, 3))
+ g_cirs = g_cirs.reshape((
+ g_cirs.shape[0] * g_cirs.shape[1],
+ g_cirs.shape[2],
+ g_cirs.shape[3],
+ g_cirs.shape[4]
+ )) # Reshape to (num_seconds * 1000, num_aps, num_rfs, num_subcarriers)
# Trace paths
g_paths_buffers = [] # Initialize the paths buffers list
diff --git a/src/tf_paths.py b/src/tf_paths.py
@@ -12,26 +12,25 @@
# - Trains the scene material parameters according to Algorithm 1 in the paper
# - Saves the losses and the trainable parameters (omega) to files (losses.npy and omegas.npy)
-from load_data import load_measurements
+from measurements import load as measurements_load
import numpy as np
-import sionna.rt as srt
import tensorflow as tf
+import sionna.rt as srt
+from sionna.constants import PI
from typing import Tuple, List
-from sys import stdout
import os
# Config
# Dataset
-g_cfg_load_only_round = 11 # Set to None to load all rounds
+g_cfg_meas_round = 11 # Measurements round to load
g_cfg_num_samples = 10 # Number of samples to use from the dataset. Picked randomly.
g_cfg_batch_size = 3 # Number of samples in each batch (B)
# Fit
-g_cfg_seed = 42 # Random seed for reproducibility
g_cfg_num_embeddings = 2 # L in the paper
g_cfg_learning_rate = 0.01 # beta in the paper
-g_cfg_decay = 0.001 # delta in the paper
+g_cfg_decay = tf.Variable(0.001, dtype=tf.float32) # delta in the paper
g_cfg_max_iterations = 5 # max number of iterations for the optimization
g_cfg_convergence_threshold = 1e-6
g_cfg_eps = 1e-6 # For numerical gradient calculation
@@ -50,11 +49,11 @@ g_cfg_rx_locations = [
[22, 0.5, 2.85]
]
g_cfg_rx_orientations = [
- [0, 0, -np.pi/2],
+ [0, 0, -PI/2],
[0, 0, 0],
- [0, 0, -np.pi/2],
+ [0, 0, -PI/2],
[0, 0, 0],
- [0, 0, np.pi/2],
+ [0, 0, PI/2],
[0, 0, 0]
]
g_cfg_tx_position_initial = np.array([5.5, 5, 0.5]) # Initial TX position
@@ -73,32 +72,16 @@ g_cfg_record_omegas = True
# Globals
g_scene: srt.Scene
-g_losses: List[float] # Losses history per iteration. Saved as a numpy array.
-g_omega: np.ndarray # Shape (2 * 4 * num_materials * num_embeddings) (2 - v/w, 4 - cond/perm/scat/xpd)
- # Note that the array is flattened,
- # so the first 4*num_materials*num_embeddings elements are v, and the rest are w
- # Before saving, it is reshaped to (2, 4, num_materials, num_embeddings)
-g_cirs: np.ndarray # The dataset of channel impulse responses (CIRs) loaded from the measurements.
- # Shape (num_samples, num_aps, num_rfs, num_subcarriers)
g_trainable_materials: List[srt.RadioMaterial] = [] # List of trainable materials in the scene
-g_paths_buffers: List[Tuple] # List of g_scene.trace_paths outputs for each sample,
- # Each element is a tuple of 8: 4 srt.Paths and 4 srt.PathsTmpData
- # Use each entry as: g_scene.compute_fields(*elem, ...)
- # (length g_num_samples)
-g_indices: np.ndarray # Indices of the samples in the g_paths_buffers. Shape (g_num_samples,)
-g_tx_poss: List[np.ndarray] # List of TX positions for each sample (length g_num_samples)
g_bandwidth: float # Bandwidth of the signal (W in the paper)
-def calc_material_parameters(
- w: np.ndarray,
- v: np.ndarray,
-) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+@tf.function
+def calc_material_parameters(omega: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor, tf.Tensor]:
"""Calculate the material parameters from the read-out vectors and embeddings.
Args:
- w: The embeddings (shape: (4, num_materials, num_embeddings)).
- v: The read-out vectors (shape: (4, num_materials, num_embeddings)).
+ omega: Read-out and embeddings tensor. Shape (2, 4, num_materials, num_embeddings)
Returns:
Tuple containing:
@@ -107,24 +90,21 @@ def calc_material_parameters(
- scat: Scattering coefficient of the materials.
- xpd: XPD coefficient of the materials.
"""
- dot = np.sum(v * w, axis=-1)
- cond = np.exp(dot[0])
- perm = np.exp(dot[1]) + 1.
-
- def sigmoid(x: np.ndarray) -> np.ndarray:
- return 1 / (1 + np.exp(-x))
- scat = sigmoid(dot[2])
- xpd = sigmoid(dot[3])
-
+ dot = tf.reduce_sum(omega[0] * omega[1], axis=-1)
+ cond = tf.exp(dot[0])
+ perm = tf.exp(dot[1]) + 1.
+ scat = tf.sigmoid(dot[2])
+ xpd = tf.sigmoid(dot[3])
return cond, perm, scat, xpd
+@tf.function
def calc_omega(
- cond: np.ndarray,
- perm: np.ndarray,
- scat: np.ndarray,
- xpd: np.ndarray,
-) -> np.ndarray:
+ cond: tf.Tensor,
+ perm: tf.Tensor,
+ scat: tf.Tensor,
+ xpd: tf.Tensor,
+) -> tf.Tensor:
"""Calculate the omega vector from the material parameters.
Args:
@@ -134,271 +114,253 @@ def calc_omega(
xpd: XPD coefficient of the materials.
Returns:
- The omega vector (shape: (2 * 4 * num_materials * num_embeddings,)).
+ The omega vector (shape: (2, 4, num_materials, num_embeddings,)).
"""
- num_materials = len(cond)
-
- dot = np.empty((4, num_materials), dtype=np.float32)
- dot[0] = np.log(cond)
- dot[1] = np.log(perm - 1.)
-
- def invert_sigmoid(x: np.ndarray) -> np.ndarray:
- return -np.log(1 / x - 1)
- dot[2] = invert_sigmoid(scat)
- dot[3] = invert_sigmoid(xpd)
-
- w = np.random.uniform(-1, 1, (4, num_materials, g_cfg_num_embeddings)).astype(np.float32)
- v = dot[:, :, np.newaxis] / w
-
-
- return np.concatenate([v.flatten(), w.flatten()])
-
-
-def set_omega() -> None:
+ num_materials = tf.shape(cond)[0]
+ dot0 = tf.math.log(cond)
+ dot1 = tf.math.log(perm - 1.)
+ dot2 = -tf.math.log(1 / scat - 1)
+ dot3 = -tf.math.log(1 / xpd - 1)
+ dot = tf.stack([dot0, dot1, dot2, dot3], axis=0)
+ w = tf.random.uniform((4, num_materials, g_cfg_num_embeddings), minval=-1, maxval=1, dtype=tf.float32)
+ v = dot[:, :, tf.newaxis] / w
+ return tf.stack([v, w], axis=0)
+
+
+@tf.function
+def set_omega(omega: tf.Tensor) -> None:
"""Set the trainable parameters in the scene from omega."""
- v = g_omega[:4 * len(g_trainable_materials) * g_cfg_num_embeddings].reshape(
- (4, len(g_trainable_materials), g_cfg_num_embeddings)
- )
- w = g_omega[4 * len(g_trainable_materials) * g_cfg_num_embeddings:].reshape(
- (4, len(g_trainable_materials), g_cfg_num_embeddings)
- )
- cond, perm, scat, xpd = calc_material_parameters(w, v)
+ cond, perm, scat, xpd = calc_material_parameters(omega)
for idx, mat in enumerate(g_trainable_materials):
- g_scene._radio_materials[mat.name].conductivity = cond[idx]
- g_scene._radio_materials[mat.name].relative_permittivity = perm[idx]
- g_scene._radio_materials[mat.name].scattering_coefficient = scat[idx]
- g_scene._radio_materials[mat.name].xpd_coefficient = xpd[idx]
+ g_scene._radio_materials[mat.name].conductivity = tf.gather(cond, idx)
+ g_scene._radio_materials[mat.name].relative_permittivity = tf.gather(perm, idx)
+ g_scene._radio_materials[mat.name].scattering_coefficient = tf.gather(scat, idx)
+ g_scene._radio_materials[mat.name].xpd_coefficient = tf.gather(xpd, idx)
-g_p_hat = np.empty((g_cfg_batch_size,), dtype=np.float32)
-g_tau_hat = np.empty((g_cfg_batch_size,), dtype=np.float32)
-g_alpha = 6e-9 # Initial α_i
-def objective(batch_indices) -> float:
+g_alpha = tf.Variable(6e-9, dtype=tf.float32) # Initial α_i
+@tf.function
+def objective(cirs_batch, batch_indices, omega, paths_buffers) -> tf.Variable:
"""Compute the objective function for the optimization.
Returns the average loss over the batch as in (38) of the paper.
"""
- set_omega()
-
- # Compute predicted CIR
- for i, ind in enumerate(batch_indices):
- a, tau = g_scene.compute_fields(
- *g_paths_buffers[ind],
- check_scene=False,
- scat_random_phases=False,
- testing=False,
- ).cir(
- los=g_cfg_los,
- reflection=g_cfg_specular_reflection,
- diffraction=g_cfg_refraction,
- scattering=g_cfg_diffuse_reflection,
- ris=False,
- cluster_ris_paths=False,
- num_paths=None,
+ set_omega(omega)
+ p_hat = tf.TensorArray(dtype=tf.float32, size=tf.shape(batch_indices)[0])
+ tau_hat = tf.TensorArray(dtype=tf.float32, size=tf.shape(batch_indices)[0])
+ for i in tf.range(tf.shape(batch_indices)[0]):
+ ind = batch_indices[i]
+ def get_a_tau(index):
+ return g_scene.compute_fields(
+ *paths_buffers[int(index)],
+ check_scene=False,
+ scat_random_phases=False,
+ testing=False,
+ ).cir(
+ los=g_cfg_los,
+ reflection=g_cfg_specular_reflection,
+ diffraction=g_cfg_refraction,
+ scattering=g_cfg_diffuse_reflection,
+ ris=False,
+ cluster_ris_paths=False,
+ num_paths=None,
+ )
+ a, tau = tf.py_function(
+ get_a_tau,
+ [ind],
+ Tout=[tf.complex64, tf.float32]
)
- a = a.numpy()
- tau = tau.numpy()
- a = a[~np.isnan(a)]
- tau = tau[tau != -1]
- g_p_hat[i] = np.sum(np.abs(a) ** 2)
- g_tau_hat[i] = np.mean(tau)
+ # Record
+ p_val = tf.reduce_sum(tf.abs(a) ** 2)
+ tau_val = tf.reduce_mean(tau)
+ p_hat = p_hat.write(i, p_val)
+ tau_hat = tau_hat.write(i, tau_val)
+ p_hat = p_hat.stack()
+ tau_hat = tau_hat.stack()
# Scale the batch: h_b ← √α_i * h_b
- batch = g_cirs[batch_indices]
- p = np.sum(np.abs(batch) ** 2, axis=(0, 1, 2))
- alpha_hat = np.sum(p[:, None] * g_p_hat) / (np.sum(g_p_hat**2) + 1e-8)
- alpha = float(g_cfg_decay * g_alpha + (1 - g_cfg_decay) * alpha_hat)
- batch *= np.sqrt(alpha)
- p = np.sum(np.abs(batch) ** 2, axis=(0, 1, 2))
+ p = tf.reduce_sum(tf.abs(cirs_batch) ** 2, axis=(0, 1, 2))
+ alpha_hat = tf.reduce_sum(p[..., tf.newaxis] * p_hat) / (tf.reduce_sum(p_hat ** 2) + 1e-8)
+ alpha = g_cfg_decay * g_alpha + (1 - g_cfg_decay) * alpha_hat
+ alpha = tf.cast(alpha, tf.complex64)
+ cirs_scaled = tf.sqrt(alpha) * cirs_batch
+ p = tf.reduce_sum(tf.abs(cirs_scaled) ** 2, axis=(0, 1, 2))
# Approximate tau_RMS as in (40, 41)
- total_channel_gain = np.sum(p)
- p_scaled = p / (total_channel_gain + 1e-8)
- ls = np.arange(batch.shape[2]) - batch.shape[2] // 2 # Subcarrier indices
- tau_overline = np.sum(ls[:, None] * p_scaled)
+ total_channel_gain = tf.reduce_sum(p)
+ p = p / (total_channel_gain + 1e-8)
+ ls = np.arange(512) - 512 // 2 # Subcarrier indices
+ tau_overline = tf.reduce_sum(ls[:, None] * p, axis=0)
tau_rms = (ls - tau_overline) / g_bandwidth
- tau_rms = np.sqrt(np.sum((tau_rms ** 2)[:, None] * p_scaled))
+ tau_rms = tf.sqrt(tf.reduce_sum((tau_rms ** 2)[:, None] * p))
# Compute sample loss L_b as in (38)
def smape(x, y):
- return np.abs(x - y) / (np.abs(x) + np.abs(y) + 1e-8)
- losses = np.mean(smape(p[:, None], g_p_hat), axis=0) + smape(tau_rms, g_tau_hat)
-
+ return tf.abs(x - y) / (tf.abs(x) + tf.abs(y) + 1e-8)
+ losses = tf.reduce_mean(smape(p[..., tf.newaxis], p_hat), axis=0) + smape(tau_rms, tau_hat)
# Average loss over the batch: L = (1/B) * sum(L_b)
- return float(np.mean(losses))
-
-
-################################################################################
-# Scene loading and setup
-################################################################################
-
-print("Loading scene...")
-g_scene = srt.load_scene(
- os.path.join(
- os.path.dirname(__file__),
- "scene",
- "scene.xml")
-)
-g_scene.frequency = g_cfg_central_frequency
-g_scene.tx_array = srt.PlanarArray(
- num_rows=1,
- num_cols=1,
- vertical_spacing=0.5,
- horizontal_spacing=0.5,
- pattern="dipole",
- polarization="V",
-)
-g_scene.rx_array = srt.PlanarArray(
- num_rows=1,
- num_cols=1,
- vertical_spacing=0.5,
- horizontal_spacing=0.5,
- pattern="dipole",
- polarization="VH",
-)
-for rx_idx, rx_loc in enumerate(g_cfg_rx_locations):
- rx = srt.Receiver(
- name=f"rx{rx_idx}",
- position=rx_loc,
- orientation=g_cfg_rx_orientations[rx_idx],
+ return tf.reduce_mean(losses)
+
+
+if __name__ == "__main__":
+ # Initialization
+ # Scene
+ print("Loading scene... ", end='', flush=True)
+ g_scene = srt.load_scene(
+ os.path.join(
+ os.path.dirname(__file__),
+ "..",
+ "scene",
+ "scene.xml")
)
- g_scene.add(rx)
-tx = srt.Transmitter(
- name="tx",
- position=g_cfg_tx_position_initial,
-)
-g_scene.add(tx)
-
-# Make trainable materials
-for mat in g_scene.radio_materials.values():
- if not mat.is_used:
- continue
- # Create new trainable material
- if g_cfg_keep_material_params:
- # Start from the current material parameters
- new_mat = srt.RadioMaterial(
- mat.name + "_train",
- relative_permittivity=mat.relative_permittivity,
- conductivity=mat.conductivity,
- scattering_coefficient=mat.scattering_coefficient,
- xpd_coefficient=mat.xpd_coefficient
+ g_scene.frequency = g_cfg_central_frequency
+ g_scene.tx_array = srt.PlanarArray(
+ num_rows=1,
+ num_cols=1,
+ vertical_spacing=0.5,
+ horizontal_spacing=0.5,
+ pattern="dipole",
+ polarization="V",
+ )
+ g_scene.rx_array = srt.PlanarArray(
+ num_rows=1,
+ num_cols=1,
+ vertical_spacing=0.5,
+ horizontal_spacing=0.5,
+ pattern="dipole",
+ polarization="VH",
+ )
+ for rx_idx, rx_loc in enumerate(g_cfg_rx_locations):
+ rx = srt.Receiver(
+ name=f"rx{rx_idx}",
+ position=rx_loc,
+ orientation=g_cfg_rx_orientations[rx_idx],
)
- else:
- new_mat = srt.RadioMaterial(mat.name + "_train",
- relative_permittivity=3.0,
- conductivity=0.1)
- g_scene.add(new_mat)
- g_trainable_materials.append(new_mat)
-
-# Assign trainable materials to the corresponding objects
-for obj in g_scene.objects.values():
- obj.radio_material = obj.radio_material.name + "_train"
-
-################################################################################
-# Load the measurements
-################################################################################
-
-print("Loading measurements...")
-data_dirpath = os.path.join(
- os.path.dirname(__file__),
- "..", "data"
-)
-g_cirs = load_measurements(data_dirpath, g_cfg_load_only_round)[0]
-g_bandwidth = g_cirs.shape[-2] * g_cfg_subcarrier_spacing # W in the paper
-g_cirs = g_cirs[0].reshape(
- (g_cirs.shape[1],
- g_cirs.shape[2],
- g_cirs.shape[4],
- g_cirs.shape[3] * g_cirs.shape[5],)
-) # Reshape to (num_aps, num_rfs, num_subcarriers, num_samples)
-# Transpose to (num_samples, num_aps, num_rfs, num_subcarriers)
-g_cirs = np.transpose(g_cirs, (3, 0, 1, 2))
-
-################################################################################
-# Trace paths
-################################################################################
-
-g_paths = [] # Initialize the paths buffers list
-g_indices = np.random.randint(g_cirs.shape[0], size=g_cfg_num_samples)
-for j, i in enumerate(g_indices): # For each sample
- stdout.write(f"\rTracing paths for sample {j}/{g_cfg_num_samples} ({i})...")
- stdout.flush()
-
- tx_pos = g_cfg_tx_position_initial + i * g_cfg_tx_velocity
- g_scene._transmitters['tx'].position = tx_pos
-
- g_paths.append(g_scene.trace_paths(
- max_depth=g_cfg_max_depth,
- num_samples=g_cfg_trace_paths_num_samples,
- los=g_cfg_los,
- reflection=g_cfg_specular_reflection,
- diffraction=g_cfg_refraction,
- scattering=g_cfg_diffuse_reflection,
- ris=False,
- edge_diffraction=False,
- check_scene=False,
- ))
-print()
-g_cirs = g_cirs[g_indices]
-
-# Fit the scene
-prev_loss = np.inf # Initialize previous loss for convergence check
-omega_size = 2 * 4 * len(g_trainable_materials) * g_cfg_num_embeddings
-optimizer = tf.keras.optimizers.Adam(learning_rate=g_cfg_learning_rate)
-
-loss_history = []
-omega_history = []
-for i in range(g_cfg_max_iterations):
- print(f"Iteration {i + 1}/{g_cfg_max_iterations}...")
-
- # Draw a random batch
- batch_indices = np.random.choice(
- g_cfg_num_samples,
- g_cfg_batch_size,
- replace=False,
+ g_scene.add(rx)
+ tx = srt.Transmitter(
+ name="tx",
+ position=g_cfg_tx_position_initial,
+ )
+ g_scene.add(tx)
+ print("Done")
+
+ # Make trainable materials
+ for mat in g_scene.radio_materials.values():
+ if not mat.is_used:
+ continue
+ # Create new trainable material
+ if g_cfg_keep_material_params:
+ # Start from the current material parameters
+ new_mat = srt.RadioMaterial(
+ mat.name + "_train",
+ relative_permittivity=mat.relative_permittivity,
+ conductivity=mat.conductivity,
+ scattering_coefficient=mat.scattering_coefficient,
+ xpd_coefficient=mat.xpd_coefficient
+ )
+ else:
+ new_mat = srt.RadioMaterial(mat.name + "_train",
+ relative_permittivity=3.0,
+ conductivity=0.1)
+ g_scene.add(new_mat)
+ g_trainable_materials.append(new_mat)
+
+ # Assign trainable materials to the corresponding objects
+ for obj in g_scene.objects.values():
+ obj.radio_material = obj.radio_material.name + "_train"
+
+ # Load the measurements
+ print("Loading measurements... ", end='', flush=True)
+ data_dirpath = os.path.join(
+ os.path.dirname(__file__),
+ "..", "..", "data"
+ )
+ # [num_aps, num_rfs, num_seconds, num_subcarriers, 1000]
+ cirs = measurements_load(data_dirpath, g_cfg_meas_round)
+ # Cast to complex64
+ cirs = cirs.astype(np.complex64)
+ # Move num_seconds to 0 and 1000 to 1
+ # [num_seconds, 1000, num_aps, num_rfs, num_subcarriers]
+ cirs = cirs.transpose((2, 4, 0, 1, 3))
+ # Unite num_seconds and 1000
+ # [num_seconds * 1000, num_aps, num_rfs, num_subcarriers]
+ cirs = cirs.reshape((cirs.shape[0] * cirs.shape[1], cirs.shape[2], cirs.shape[3], cirs.shape[4]))
+ g_bandwidth = cirs.shape[-2] * g_cfg_subcarrier_spacing # W in the paper
+ print("Done")
+
+ # Make the dataset
+ # Contains:
+ # - CIRs (shape (num_samples, num_aps, num_rfs, num_subcarriers))
+ # - paths indices (shape (num_samples,))
+ indices = np.random.randint(0, cirs.shape[0], (g_cfg_num_samples,))
+ dataset = tf.data.Dataset.from_tensor_slices((cirs[indices], np.arange(g_cfg_num_samples)))
+ assert dataset.cardinality() == g_cfg_num_samples
+
+ # Trace paths
+ print("Tracing paths... ", end="", flush=True)
+ paths_buffers = [] # Initialize the paths buffers list
+ tx_poss = g_cfg_tx_position_initial + indices[:, np.newaxis] * g_cfg_tx_velocity
+ for tx_pos in tx_poss:
+ g_scene._transmitters['tx'].position = tx_pos
+ paths_buffers.append(g_scene.trace_paths(
+ max_depth=g_cfg_max_depth,
+ num_samples=g_cfg_trace_paths_num_samples,
+ los=g_cfg_los,
+ reflection=g_cfg_specular_reflection,
+ diffraction=g_cfg_refraction,
+ scattering=g_cfg_diffuse_reflection,
+ ris=False,
+ edge_diffraction=False,
+ check_scene=False,
+ ))
+ print("Done")
+
+ # Fit the scene
+ prev_loss = tf.Variable(1e10, dtype=tf.float32)
+ omega = tf.random.uniform(
+ (2, 4, len(g_trainable_materials), g_cfg_num_embeddings),
+ minval=-1, maxval=1, dtype=tf.float32
+ )
+ omega = tf.Variable(omega, dtype=tf.float32)
+ optimizer = tf.keras.optimizers.Adam(
+ learning_rate=g_cfg_learning_rate,
+ # beta_1=0.9,
+ # beta_2=0.999,
+ # epsilon=1e-8,
)
- # Loss
- with tf.GradientTape() as tape:
- curr_loss = objective(batch_indices)
- print(f"Current loss: {curr_loss:.6f}")
-
- # Adam optimization step
- optimizer.apply_gradients(zip(
- tape.gradient(
- curr_loss,
- tape.watched_variables,
- unconnected_gradients=tf.UnconnectedGradients.ZERO
- ),
- tape.watched_variables() # type: ignore
- ))
-
- # Update the scene parameters with the new omega
- set_omega()
-
- # Record the loss and omega
+ print("Starting optimization...", flush=True)
+ loss_history = []
+ omega_history = []
+ dataset = dataset.batch(g_cfg_batch_size)
+ for cirs_batch, paths_indices_batch in dataset:
+ with tf.GradientTape() as tape:
+ curr_loss = objective(cirs_batch, paths_indices_batch, omega, paths_buffers)
+ grad = tape.gradient(curr_loss, omega)
+ optimizer.apply_gradients(zip([grad], [omega]))
+
+ if g_cfg_record_loss:
+ loss_history.append(curr_loss)
+ if g_cfg_record_omegas:
+ omega_history.append(omega.numpy().copy())
+
+ # Check convergence: Stop if change in loss is below threshold
+ if abs(prev_loss - curr_loss) < g_cfg_convergence_threshold:
+ print("Converged")
+ break
+ prev_loss = curr_loss
+ print(f"Loss: {curr_loss}", flush=True)
+
+ print("Training completed")
+
+ # Save the losses and omegas
+ losses = np.asarray(loss_history)
+ omegas = np.asarray(omega_history)
if g_cfg_record_loss:
- loss_history.append(curr_loss)
-
- # Check convergence: Stop if change in loss is below threshold
- if abs(prev_loss - curr_loss) < g_cfg_convergence_threshold:
- break
- prev_loss = curr_loss
-
- print(f"Iteration {i}, Loss: {curr_loss:.6f}")
-
-# Final update to scene parameters
-set_omega()
-print("Training completed")
-
-# Save the losses and omegas
-losses = np.asarray(loss_history)
-omegas = np.asarray(omega_history)
-if g_cfg_record_loss:
- np.save("losses.npy", losses)
- print("Saved losses to losses.npy: ", losses.shape)
-if g_cfg_record_omegas:
- omegas = omegas.reshape(
- (-1, 2, 4, len(g_trainable_materials), g_cfg_num_embeddings)
- )
- np.save("omegas.npy", omegas)
- print("Saved omegas to omegas.npy: ", omegas.shape)
+ np.save("losses.npy", losses)
+ print("Saved losses to losses.npy: ", losses.shape)
+ if g_cfg_record_omegas:
+ omegas = omegas.reshape(
+ (-1, 2, 4, len(g_trainable_materials), g_cfg_num_embeddings)
+ )
+ np.save("omegas.npy", omegas)
+ print("Saved omegas to omegas.npy: ", omegas.shape)
diff --git a/src/trace_paths.py b/src/trace_paths.py
@@ -1,23 +0,0 @@
-# flake8: noqa: E501
-
-# provides trace_paths function
-
-from typing import List
-
-import numpy as np
-import sionna.rt as srt
-
-
-def trace_paths(
- scene: srt.Scene,
- tx_poss: np.ndarray,
- trace_paths_params: dict,
-) -> List[srt.Paths]:
- """Trace paths for given transmitter positions in a scene.
- The scene must contain at least one transmitter, named "tx".
- """
- paths = []
- for tx_pos in tx_poss:
- scene._transmitters["tx"].position = tx_pos # type: ignore[attr-defined]
- paths.append(scene.trace_paths(**trace_paths_params))
- return paths
