np_paths.py (15485B)
1 # flake8: noqa: E501 # Ignore long lines 2 3 # Implementation following the 4 # "Learning Radio Environments by Differentiable Ray Tracing" 5 # 2024 NVIDIA paper 6 # http://dx.doi.org/10.1109/TMLCN.2024.3474639 7 8 # What this script does: 9 # - Loads the scene and measurements 10 # - Merges seconds and milliseconds (thus getting a total milliseconds dimension, further referred to as "samples") 11 # - Traces paths for each sample (picking random g_num_samples samples) with sionna rt Candidatecalculator and ImageMethod, generating a dataset of g_num_samples PathBuffers 12 # - Trains the scene material parameters according to Algorithm 1 in the paper 13 # - Saves the losses and the trainable parameters (omega) to files (losses.npy and omegas.npy) 14 15 from measurements import load as measurements_load 16 17 import numpy as np 18 import sionna.rt as srt 19 20 from typing import Tuple, List 21 from sys import stdout 22 import os 23 24 # Config 25 # Dataset 26 g_cfg_load_only_round = 11 # Set to None to load all rounds 27 g_cfg_num_samples = 10 # Number of samples to use from the dataset. Picked randomly. 28 g_cfg_batch_size = 3 # Number of samples in each batch (B) 29 # Fit 30 g_cfg_seed = 42 # Random seed for reproducibility 31 g_cfg_num_embeddings = 2 # L in the paper 32 g_cfg_learning_rate = 0.01 # beta in the paper 33 g_cfg_decay = 0.001 # delta in the paper 34 g_cfg_max_iterations = 5 # max number of iterations for the optimization 35 g_cfg_convergence_threshold = 1e-6 36 g_cfg_eps = 1e-6 # For numerical gradient calculation 37 g_cfg_keep_material_params = False # If True, start fitting from the current material parameters 38 # RF 39 g_cfg_central_frequency = 3.75e9 # Set into scene.frequency 40 g_cfg_subcarrier_spacing = 100e6 # Delta_f in the paper 41 g_cfg_sampling_rate = 1. / 1000. 42 # Spacial 43 g_cfg_rx_locations = [ 44 [4, 5, 2.2], 45 [4, 5, 2.2], 46 [20.5, 8, 3], 47 [20.5, 8, 3], 48 [22, 0.5, 2.85], 49 [22, 0.5, 2.85] 50 ] 51 g_cfg_rx_orientations = [ 52 [0, 0, -np.pi/2], 53 [0, 0, 0], 54 [0, 0, -np.pi/2], 55 [0, 0, 0], 56 [0, 0, np.pi/2], 57 [0, 0, 0] 58 ] 59 g_cfg_tx_position_initial = np.array([5.5, 5, 0.5]) # Initial TX position 60 g_cfg_tx_velocity = np.array([0.0005, 0, 0]) # meters per millisecond 61 # Sionna 62 g_cfg_max_depth = 3 63 g_cfg_trace_paths_num_samples = 1000 # Number of samples to trace paths for 64 g_cfg_synthetic_array = False 65 g_cfg_los = True 66 g_cfg_specular_reflection = True 67 g_cfg_diffuse_reflection = True 68 g_cfg_refraction = True 69 # Output 70 g_cfg_record_loss = True 71 g_cfg_record_omegas = True 72 73 # Globals 74 g_scene: srt.Scene 75 g_losses: List[float] # Losses history per iteration. Saved as a numpy array. 76 g_omega: np.ndarray # Shape (2 * 4 * num_materials * num_embeddings) (2 - v/w, 4 - cond/perm/scat/xpd) 77 # Note that the array is flattened, 78 # so the first 4*num_materials*num_embeddings elements are v, and the rest are w 79 # Before saving, it is reshaped to (2, 4, num_materials, num_embeddings) 80 g_cirs: np.ndarray # The dataset of channel impulse responses (CIRs) loaded from the measurements. 81 # Shape (num_samples, num_aps, num_rfs, num_subcarriers) 82 g_trainable_materials: List[srt.RadioMaterial] = [] # List of trainable materials in the scene 83 g_paths_buffers: List[Tuple] # List of g_scene.trace_paths outputs for each sample, 84 # Each element is a tuple of 8: 4 srt.Paths and 4 srt.PathsTmpData 85 # Use each entry as: g_scene.compute_fields(*elem, ...) 86 # (length g_num_samples) 87 g_indices: np.ndarray # Indices of the samples in the g_paths_buffers. Shape (g_num_samples,) 88 g_tx_poss: List[np.ndarray] # List of TX positions for each sample (length g_num_samples) 89 g_bandwidth: float # Bandwidth of the signal (W in the paper) 90 91 92 def calc_material_parameters( 93 w: np.ndarray, 94 v: np.ndarray, 95 ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: 96 """Calculate the material parameters from the read-out vectors and embeddings. 97 98 Args: 99 w: The embeddings (shape: (4, num_materials, num_embeddings)). 100 v: The read-out vectors (shape: (4, num_materials, num_embeddings)). 101 102 Returns: 103 Tuple containing: 104 - cond: Conductivity of the materials. 105 - perm: Relative permittivity of the materials. 106 - scat: Scattering coefficient of the materials. 107 - xpd: XPD coefficient of the materials. 108 """ 109 dot = np.sum(v * w, axis=-1) 110 cond = np.exp(dot[0]) 111 perm = np.exp(dot[1]) + 1. 112 113 def sigmoid(x: np.ndarray) -> np.ndarray: 114 return 1 / (1 + np.exp(-x)) 115 scat = sigmoid(dot[2]) 116 xpd = sigmoid(dot[3]) 117 118 return cond, perm, scat, xpd 119 120 121 def calc_omega( 122 cond: np.ndarray, 123 perm: np.ndarray, 124 scat: np.ndarray, 125 xpd: np.ndarray, 126 ) -> np.ndarray: 127 """Calculate the omega vector from the material parameters. 128 129 Args: 130 cond: Conductivity of the materials. 131 perm: Relative permittivity of the materials. 132 scat: Scattering coefficient of the materials. 133 xpd: XPD coefficient of the materials. 134 135 Returns: 136 The omega vector (shape: (2 * 4 * num_materials * num_embeddings,)). 137 """ 138 num_materials = len(cond) 139 140 dot = np.empty((4, num_materials), dtype=np.float32) 141 dot[0] = np.log(cond) 142 dot[1] = np.log(perm - 1.) 143 144 def invert_sigmoid(x: np.ndarray) -> np.ndarray: 145 return -np.log(1 / x - 1) 146 dot[2] = invert_sigmoid(scat) 147 dot[3] = invert_sigmoid(xpd) 148 149 w = np.random.uniform(-1, 1, (4, num_materials, g_cfg_num_embeddings)).astype(np.float32) 150 v = dot[:, :, np.newaxis] / w 151 152 153 return np.concatenate([v.flatten(), w.flatten()]) 154 155 156 def set_omega() -> None: 157 """Set the trainable parameters in the scene from omega.""" 158 v = g_omega[:4 * len(g_trainable_materials) * g_cfg_num_embeddings].reshape( 159 (4, len(g_trainable_materials), g_cfg_num_embeddings) 160 ) 161 w = g_omega[4 * len(g_trainable_materials) * g_cfg_num_embeddings:].reshape( 162 (4, len(g_trainable_materials), g_cfg_num_embeddings) 163 ) 164 cond, perm, scat, xpd = calc_material_parameters(w, v) 165 for idx, mat in enumerate(g_trainable_materials): 166 g_scene._radio_materials[mat.name].conductivity = cond[idx] 167 g_scene._radio_materials[mat.name].relative_permittivity = perm[idx] 168 g_scene._radio_materials[mat.name].scattering_coefficient = scat[idx] 169 g_scene._radio_materials[mat.name].xpd_coefficient = xpd[idx] 170 171 172 g_p_hat = np.empty((g_cfg_batch_size,), dtype=np.float32) 173 g_tau_hat = np.empty((g_cfg_batch_size,), dtype=np.float32) 174 g_alpha = 6e-9 # Initial α_i 175 def objective(batch_indices) -> float: 176 """Compute the objective function for the optimization. 177 Returns the average loss over the batch as in (38) of the paper. 178 """ 179 set_omega() 180 181 # Compute predicted CIR 182 for i, ind in enumerate(batch_indices): 183 a, tau = g_scene.compute_fields( 184 *g_paths_buffers[ind], 185 check_scene=False, 186 scat_random_phases=False, 187 testing=False, 188 ).cir( 189 los=g_cfg_los, 190 reflection=g_cfg_specular_reflection, 191 diffraction=g_cfg_refraction, 192 scattering=g_cfg_diffuse_reflection, 193 ris=False, 194 cluster_ris_paths=False, 195 num_paths=None, 196 ) 197 a = a.numpy() 198 tau = tau.numpy() 199 a = a[~np.isnan(a)] 200 tau = tau[tau != -1] 201 g_p_hat[i] = np.sum(np.abs(a) ** 2) 202 g_tau_hat[i] = np.mean(tau) 203 204 # Scale the batch: h_b ← √α_i * h_b 205 batch = g_cirs[batch_indices] 206 p = np.sum(np.abs(batch) ** 2, axis=(0, 1, 2)) 207 alpha_hat = np.sum(p[:, None] * g_p_hat) / (np.sum(g_p_hat**2) + 1e-8) 208 alpha = float(g_cfg_decay * g_alpha + (1 - g_cfg_decay) * alpha_hat) 209 batch *= np.sqrt(alpha) 210 p = np.sum(np.abs(batch) ** 2, axis=(0, 1, 2)) 211 212 # Approximate tau_RMS as in (40, 41) 213 total_channel_gain = np.sum(p) 214 p_scaled = p / (total_channel_gain + 1e-8) 215 ls = np.arange(batch.shape[2]) - batch.shape[2] // 2 # Subcarrier indices 216 tau_overline = np.sum(ls[:, None] * p_scaled) 217 tau_rms = (ls - tau_overline) / g_bandwidth 218 tau_rms = np.sqrt(np.sum((tau_rms ** 2)[:, None] * p_scaled)) 219 220 # Compute sample loss L_b as in (38) 221 def smape(x, y): 222 return np.abs(x - y) / (np.abs(x) + np.abs(y) + 1e-8) 223 losses = np.mean(smape(p[:, None], g_p_hat), axis=0) + smape(tau_rms, g_tau_hat) 224 225 # Average loss over the batch: L = (1/B) * sum(L_b) 226 return float(np.mean(losses)) 227 228 229 if __name__ == "__main__": 230 # Initialization 231 # Scene 232 print("Loading scene...") 233 g_scene = srt.load_scene( 234 os.path.join( 235 os.path.dirname(__file__), 236 "..", 237 "scene", 238 "scene.xml") 239 ) 240 g_scene.frequency = g_cfg_central_frequency 241 g_scene.tx_array = srt.PlanarArray( 242 num_rows=1, 243 num_cols=1, 244 vertical_spacing=0.5, 245 horizontal_spacing=0.5, 246 pattern="dipole", 247 polarization="V", 248 ) 249 g_scene.rx_array = srt.PlanarArray( 250 num_rows=1, 251 num_cols=1, 252 vertical_spacing=0.5, 253 horizontal_spacing=0.5, 254 pattern="dipole", 255 polarization="VH", 256 ) 257 for rx_idx, rx_loc in enumerate(g_cfg_rx_locations): 258 rx = srt.Receiver( 259 name=f"rx{rx_idx}", 260 position=rx_loc, 261 orientation=g_cfg_rx_orientations[rx_idx], 262 ) 263 g_scene.add(rx) 264 tx = srt.Transmitter( 265 name="tx", 266 position=g_cfg_tx_position_initial, 267 ) 268 g_scene.add(tx) 269 270 # Make trainable materials 271 for mat in g_scene.radio_materials.values(): 272 if not mat.is_used: 273 continue 274 # Create new trainable material 275 if g_cfg_keep_material_params: 276 # Start from the current material parameters 277 new_mat = srt.RadioMaterial( 278 mat.name + "_train", 279 relative_permittivity=mat.relative_permittivity, 280 conductivity=mat.conductivity, 281 scattering_coefficient=mat.scattering_coefficient, 282 xpd_coefficient=mat.xpd_coefficient 283 ) 284 else: 285 new_mat = srt.RadioMaterial(mat.name + "_train", 286 relative_permittivity=3.0, 287 conductivity=0.1) 288 g_scene.add(new_mat) 289 g_trainable_materials.append(new_mat) 290 291 # Assign trainable materials to the corresponding objects 292 for obj in g_scene.objects.values(): 293 obj.radio_material = obj.radio_material.name + "_train" 294 295 # Load the measurements 296 print("Loading measurements...") 297 data_dirpath = os.path.join( 298 os.path.dirname(__file__), 299 "..", "..", "data" 300 ) 301 g_cirs = measurements_load(data_dirpath, g_cfg_load_only_round) 302 g_bandwidth = g_cirs.shape[-2] * g_cfg_subcarrier_spacing # W in the paper 303 g_cirs = g_cirs.transpose((2, 4, 0, 1, 3)) 304 g_cirs = g_cirs.reshape(( 305 g_cirs.shape[0] * g_cirs.shape[1], 306 g_cirs.shape[2], 307 g_cirs.shape[3], 308 g_cirs.shape[4] 309 )) # Reshape to (num_seconds * 1000, num_aps, num_rfs, num_subcarriers) 310 311 # Trace paths 312 g_paths_buffers = [] # Initialize the paths buffers list 313 g_indices = np.random.randint(g_cirs.shape[0], size=g_cfg_num_samples) 314 for j, i in enumerate(g_indices): # For each sample 315 stdout.write(f"\rTracing paths for sample {j}/{g_cfg_num_samples} ({i})...") 316 stdout.flush() 317 318 tx_pos = g_cfg_tx_position_initial + i * g_cfg_tx_velocity 319 g_scene._transmitters['tx'].position = tx_pos 320 321 g_paths_buffers.append(g_scene.trace_paths( 322 max_depth=g_cfg_max_depth, 323 num_samples=g_cfg_trace_paths_num_samples, 324 los=g_cfg_los, 325 reflection=g_cfg_specular_reflection, 326 diffraction=g_cfg_refraction, 327 scattering=g_cfg_diffuse_reflection, 328 ris=False, 329 edge_diffraction=False, 330 check_scene=False, 331 )) 332 print() 333 g_cirs = g_cirs[g_indices] 334 335 # Fit the scene 336 prev_loss = np.inf # Initialize previous loss for convergence check 337 omega_size = 2 * 4 * len(g_trainable_materials) * g_cfg_num_embeddings 338 g_omega = np.random.uniform(-1, 1, omega_size).astype(np.float32) 339 omega_grad = np.empty_like(g_omega) 340 341 # Adam optimizer parameters 342 m = np.zeros_like(g_omega) 343 v = np.zeros_like(g_omega) 344 beta1 = 0.9 # Exponential decay rate for first moment 345 beta2 = 0.999 # Exponential decay rate for second moment 346 eps = 1e-8 # Small constant for numerical stability 347 348 loss_history = [] 349 omega_history = [] 350 for i in range(g_cfg_max_iterations): 351 print(f"Iteration {i + 1}/{g_cfg_max_iterations}...") 352 353 # Draw a random batch 354 batch_indices = np.random.choice( 355 g_cfg_num_samples, 356 g_cfg_batch_size, 357 replace=False, 358 ) 359 360 # Loss 361 curr_loss = objective(batch_indices) 362 print(f"Current loss: {curr_loss:.6f}") 363 364 # Compute gradient 365 for j in range(g_omega.size): 366 stdout.write(f"\rCalculating gradient for parameter {j + 1}/{g_omega.size}...") 367 stdout.flush() 368 g_omega[j] += g_cfg_eps 369 fun_plus = objective(batch_indices) 370 g_omega[j] -= 2 * g_cfg_eps 371 fun_minus = objective(batch_indices) 372 g_omega[j] += g_cfg_eps 373 omega_grad[j] = (fun_plus - fun_minus) / (2 * g_cfg_eps) 374 375 # Adam optimization step 376 m = beta1 * m + (1 - beta1) * omega_grad 377 v = beta2 * v + (1 - beta2) * (omega_grad ** 2) 378 m_hat = m / (1 - beta1 ** (i + 1)) # Bias correction 379 v_hat = v / (1 - beta2 ** (i + 1)) # Bias correction 380 g_omega = g_omega - g_cfg_learning_rate * m_hat / (np.sqrt(v_hat) + eps) 381 382 # Update the scene parameters with the new omega 383 set_omega() 384 385 # Record the loss and omega 386 if g_cfg_record_loss: 387 loss_history.append(curr_loss) 388 if g_cfg_record_omegas: 389 omega_history.append(np.copy(g_omega)) 390 391 # Check convergence: Stop if change in loss is below threshold 392 if abs(prev_loss - curr_loss) < g_cfg_convergence_threshold: 393 stdout.write(f"Converged at iteration {i}\n") 394 break 395 prev_loss = curr_loss 396 397 msg = f"Iteration {i}, Loss: {curr_loss:.6f}\n" 398 for mat in g_trainable_materials: 399 mat_params = g_scene.get(mat.name) 400 msg += (f"{mat.name}: " 401 f"Cond: {mat_params.conductivity.numpy():.6f}, " 402 f"Perm: {mat_params.relative_permittivity.numpy():.6f}, " 403 f"Scat: {mat_params.scattering_coefficient.numpy():.6f}, " 404 f"XPD: {mat_params.xpd_coefficient.numpy():.6f}\n") 405 print(msg) 406 407 # Final update to scene parameters 408 set_omega() 409 print("Training completed") 410 411 # Save the losses and omegas 412 losses = np.asarray(loss_history) 413 omegas = np.asarray(omega_history) 414 if g_cfg_record_loss: 415 np.save("losses.npy", losses) 416 print("Saved losses to losses.npy: ", losses.shape) 417 if g_cfg_record_omegas: 418 omegas = omegas.reshape( 419 (-1, 2, 4, len(g_trainable_materials), g_cfg_num_embeddings) 420 ) 421 np.save("omegas.npy", omegas) 422 print("Saved omegas to omegas.npy: ", omegas.shape)