scattering_pattern.py (19363B)
1 # 2 # SPDX-FileCopyrightText: Copyright (c) 2021-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 # SPDX-License-Identifier: Apache-2.0 4 # 5 """ 6 Classes and functions related to Scattering patterns for the Sionna RT Module 7 """ 8 9 import tensorflow as tf 10 from scipy.special import binom 11 from abc import ABC 12 import numpy as np 13 import matplotlib.pyplot as plt 14 from matplotlib import cm 15 16 from .utils import theta_phi_from_unit_vec, r_hat, dot 17 from sionna.constants import PI 18 19 class ScatteringPattern(ABC): 20 # pylint: disable=line-too-long 21 r""" 22 Abstract class defining a scattering pattern for diffuse reflections 23 24 This class implements a mix of the Backscattering, Directive, and 25 Lambertian scattering models as described in [Degli-Esposti07]_. 26 27 Parameters 28 ---------- 29 alpha_r : int, [0,1,2,...] 30 Parameter related to the width of the scattering lobe in the 31 direction of the specular reflection. 32 A value of 0 indicates Lambertian scattering. 33 34 alpha_i : int, [1,2,...] 35 Parameter related to the width of the scattering lobe in the 36 incoming direction. Only plays a role if ``alpha_r``>0. 37 38 lambda_ : float, [0,1] 39 Parameter determining the percentage of the diffusely 40 reflected energy in the lobe around the specular reflection. 41 42 dtype : tf.complex64 or tf.complex128 43 Datatype used for all computations. 44 Defaults to `tf.complex64`. 45 46 Input 47 ----- 48 k_i : [batch_size, 3], dtype.real_dtype 49 Tensor of incoming directions 50 51 k_s : [batch_size,3], dtype.real_dtype 52 Tensor of outgoing directions 53 54 n_hat : [batch_size, 3], dtype.real_dtype 55 Tensor of surface normals 56 57 Output 58 ------ 59 pattern : [batch_size], dtype.real_dtype 60 Scattering pattern 61 """ 62 def __init__(self, 63 alpha_r, 64 alpha_i, 65 lambda_, 66 dtype=tf.complex64): 67 if dtype not in (tf.complex64, tf.complex128): 68 msg = "`dtype` must be `tf.complex64` or `tf.complex128`" 69 raise ValueError(msg) 70 self._dtype = dtype 71 self._rdtype = dtype.real_dtype 72 self.alpha_r = alpha_r 73 self.alpha_i = alpha_i 74 self.lambda_ = lambda_ 75 76 ### 77 ### Static properties and methods 78 ### 79 _alpha_max = -1 # Maximum value of alpha_r, alpha_i 80 _params_32 = None # Dictionary of precomputed tensors in 32bit precision 81 _params_64 = None # Dictionary of precomputed tensors in 64bit precision 82 83 @staticmethod 84 def _update_params(alpha_max): 85 """Precomputes several tensors needed for the __call__""" 86 alpha_max = max(alpha_max, 1) 87 ScatteringPattern._alpha_max = alpha_max 88 j_even = tf.range(0, alpha_max+1, delta=2, dtype=tf.float64) 89 j_odd = tf.range(1, alpha_max+1, delta=2, dtype=tf.float64) 90 i_j_even = 2*PI/(j_even+1) 91 n_max = (j_odd[-1]-1)/2 92 w_range = tf.range(0, n_max+1, dtype=tf.float64) 93 w_2 = 2*w_range 94 binom_2 = tf.constant(binom(w_2, w_range), tf.float64) 95 binom_1 = np.zeros([alpha_max+1, alpha_max+1]) 96 for alpha in range(alpha_max+1): 97 binom_1[alpha, :alpha+1] = binom(alpha, np.arange(alpha+1))/2**alpha 98 binom_1_even = tf.cast(binom_1[:,::2], tf.float64) 99 binom_1_odd = tf.cast(binom_1[:,1::2], tf.float64) 100 f_summands_even = tf.reduce_sum(binom_1_even*i_j_even, axis=-1) 101 ScatteringPattern._params_64 = { 102 "binom_1_odd" : binom_1_odd, 103 "binom_2" : binom_2, 104 "f_summands_even" : f_summands_even, 105 "j_odd" : j_odd, 106 "n_max" : int(n_max), 107 "w_2" : w_2 108 } 109 ScatteringPattern._params_32 = { 110 "binom_1_odd" : tf.cast(binom_1_odd, tf.float32), 111 "binom_2" : tf.cast(binom_2, tf.float32), 112 "f_summands_even" : tf.cast(f_summands_even, tf.float32), 113 "j_odd" : tf.cast(j_odd, tf.float32), 114 "n_max" : int(n_max), 115 "w_2" : tf.cast(w_2, tf.float32) 116 } 117 118 @staticmethod 119 def f_alpha(cos_theta_i, alphas): 120 """Compute the normalization factor F_{alpha_R}""" 121 cos_theta_i = tf.expand_dims(cos_theta_i, -1) 122 sin_theta_i = tf.sqrt(1-cos_theta_i**2) 123 dtype = cos_theta_i.dtype 124 if dtype==tf.float32: 125 params = ScatteringPattern._params_32 126 else: 127 params = ScatteringPattern._params_64 128 129 # Compute K_n 130 series = params["binom_2"]*(sin_theta_i/2)**(params["w_2"]) 131 n = tf.constant(0, tf.int32) 132 n_max = tf.cast(params["n_max"], tf.int32) 133 k_n = tf.TensorArray(dtype=dtype, size=n_max+1) 134 135 def condition(n, n_max, k_n, series): # pylint: disable=unused-argument 136 return n <= n_max 137 138 # Iteratively compute the sums of the first n elements of series 139 def body(n, n_max, k_n, series): 140 k_n = k_n.write(n, tf.reduce_sum(series[...,:n+1], axis=-1)) 141 return n+1, n_max, k_n, series 142 143 _, _, k_n, _ = tf.while_loop(condition, body, 144 loop_vars=[n, n_max, k_n, series], 145 parallel_iterations=params["n_max"]+1) 146 k_n = k_n.stack() 147 k_n = tf.transpose(k_n, perm=[1, 0]) 148 149 # Compute I_j for odd values of j 150 i_j_odd = 2*PI/(params["j_odd"]+1)*cos_theta_i*k_n 151 152 # Compute the sum over odd values of j 153 f_summands_odd = tf.gather(params["binom_1_odd"], alphas)*i_j_odd 154 f_odd = tf.reduce_sum(f_summands_odd, -1) 155 156 # Get the precomputed sum over even values of j 157 f_even = tf.gather(params["f_summands_even"], alphas) 158 159 # Compute the final normalization factor F_alpha 160 f = f_odd+f_even 161 162 return f 163 164 @staticmethod 165 def pattern(k_i, k_s, n_hat, alpha_r, alpha_i, lambda_): 166 """Compute the scattering pattern 167 168 The function always computes the BackscatteringPattern as well 169 as the LambertianPattern and then selects which one is applied, 170 depending on the values of alpha_r (alpha_r=0 means Lambertian). 171 The DirectivePattern is a special case for lambda_=0. 172 173 This design choice has been made to allow the computation 174 of a different scattering pattern for each pair k_i, k_s. 175 """ 176 dtype = k_i.dtype 177 cos_theta_i = -dot(k_i, n_hat, clip=True) 178 cos_theta_s = dot(k_s, n_hat, clip=True) 179 k_r = k_i + 2*tf.expand_dims(cos_theta_i, -1)*n_hat 180 181 # Compute backscattering pattern 182 f_alpha_r = ScatteringPattern.f_alpha(cos_theta_i, alpha_r) 183 f_alpha_i = ScatteringPattern.f_alpha(cos_theta_i, alpha_i) 184 f = lambda_*f_alpha_r + (1-lambda_)*f_alpha_i 185 pattern_bs = lambda_*tf.pow((1+dot(k_r, k_s, clip=True))/2, 186 tf.cast(alpha_r, k_r.dtype)) 187 pattern_bs += (1-lambda_)*tf.pow((1-dot(k_i, k_s, clip=True))/2, 188 tf.cast(alpha_i, k_r.dtype)) 189 pattern_bs /= f 190 191 # Compute Lambertian pattern 192 pattern_l = cos_theta_s/tf.cast(PI, k_i.dtype) 193 pattern_l = tf.where(pattern_l<0., tf.cast(0, dtype), pattern_l) 194 195 # Select one of the two models depending on alpha_r 196 pattern = tf.where(alpha_r==0, pattern_l, pattern_bs) 197 198 return pattern 199 200 ### 201 ### Instance properties and methods 202 ### 203 @property 204 def alpha_r(self): 205 """ 206 bool : Get/set ``alpha_r`` 207 """ 208 return self._alpha_r 209 210 @alpha_r.setter 211 def alpha_r(self, value): 212 if not isinstance(value, (int, np.integer)): 213 raise TypeError("alpha_r must be a positive integer") 214 215 if isinstance(self, LambertianPattern): 216 value = 0 217 else: 218 if value < 1: 219 raise ValueError("alpha must be >=1") 220 if value > ScatteringPattern._alpha_max: 221 ScatteringPattern._update_params(value) 222 self._alpha_r = value 223 224 @property 225 def alpha_i(self): 226 """ 227 bool : Get/set ``alpha_i`` 228 """ 229 return self._alpha_i 230 231 @alpha_i.setter 232 def alpha_i(self, value): 233 if not isinstance(value, (int, np.integer)): 234 raise TypeError("alpha_i must be a positive integer") 235 if value < 1: 236 raise ValueError("alpha msu be >=1") 237 if value > ScatteringPattern._alpha_max: 238 ScatteringPattern._update_params(value) 239 self._alpha_i = value 240 241 @property 242 def lambda_(self): 243 """ 244 bool : Get/set ``lambda_`` 245 """ 246 return self._lambda_ 247 248 @lambda_.setter 249 def lambda_(self, value): 250 if isinstance(value, tf.Variable): 251 if value.dtype != self._rdtype: 252 msg = f"`lambda_` must have dtype={self._rdtype}" 253 raise TypeError(msg) 254 else: 255 self._lambda_ = value 256 else: 257 self._lambda_ = tf.cast(value, self._rdtype) 258 259 def __call__(self, k_i, k_s, n_hat): 260 return ScatteringPattern.pattern(k_i, k_s, n_hat, self.alpha_r, 261 self.alpha_i, self._lambda_) 262 263 def visualize(self, k_i=(0.7071, 0., -0.7071), show_directions=False): 264 r""" 265 Visualizes the scattering pattern 266 267 It is assumed that the surface normal points toward the 268 positive z-axis. 269 270 Input 271 ----- 272 k_i : [3], array_like 273 Incoming direction 274 275 show_directions : bool 276 If `True`, the incoming and specular reflection directions 277 are shown. 278 Defaults to `False`. 279 280 Output 281 ------ 282 : :class:`matplotlib.pyplot.Figure` 283 3D visualization of the scattering pattern 284 285 : :class:`matplotlib.pyplot.Figure` 286 Visualization of the incident plane cut through 287 the scattering pattern 288 """ 289 k_i_in = k_i 290 291 ### 292 ### 3D visualization 293 ### 294 theta = np.linspace(0.0, PI/2, 50) 295 phi = np.linspace(-PI, PI, 100) 296 theta_grid, phi_grid = np.meshgrid(theta, phi, indexing='ij') 297 k_s = tf.cast(r_hat(theta_grid, phi_grid), self._dtype.real_dtype) 298 k_i = tf.cast(tf.broadcast_to(k_i_in, k_s.shape), self._dtype.real_dtype) 299 k_s = tf.reshape(k_s, [-1, 3]) 300 k_i = tf.reshape(k_i, [-1, 3]) 301 n_hat = tf.constant([[0,0,1]], dtype=k_i.dtype) 302 n_hat = tf.repeat(n_hat, k_i.shape[0], 0) 303 pattern = self(k_i, k_s, n_hat) 304 pattern = tf.reshape(pattern, [50, 100]) 305 x = pattern * np.sin(theta_grid) * np.cos(phi_grid) 306 y = pattern * np.sin(theta_grid) * np.sin(phi_grid) 307 z = pattern * np.cos(theta_grid) 308 p_min = np.min(pattern) 309 p_max = np.max(pattern) 310 311 def norm(x): 312 """Maps input to [0,1] range""" 313 x_min = np.min(x) 314 x_max = np.max(x) 315 if x_min==x_max: 316 x = np.ones_like(x) 317 else: 318 x -= x_min 319 x /= np.abs(x_max-x_min) 320 return x 321 322 fig_3d = plt.figure() 323 ax = fig_3d.add_subplot(1,1,1, projection='3d', computed_zorder=False) 324 325 ax.plot_surface(x, y, z, rstride=1, cstride=1, linewidth=0, 326 antialiased=False, alpha=0.7, 327 facecolors=cm.turbo(norm(pattern))) 328 329 ax.set_box_aspect((np.ptp(x), np.ptp(y), np.ptp(z))) 330 331 if show_directions: 332 r = np.max(pattern)*1.1 333 k_i = k_i[0,0] 334 uvw = -k_i.numpy() 335 plt.quiver(*r*uvw, *-uvw, length=r, linestyle='dashed', color="black", arrow_length_ratio=0.07, alpha=0.5) 336 ax.text(r*uvw[0], r*uvw[1], r*uvw[2], r"$\hat{\mathbf{k}}_\mathrm{i}$") 337 338 theta_i, phi_i = theta_phi_from_unit_vec(-k_i) 339 theta_r, phi_r = theta_i, phi_i + tf.cast(PI, phi_i.dtype) 340 k_r = r_hat(theta_r, phi_r).numpy() 341 plt.quiver(*[0,0,0], *k_r, length=r, linestyle='dashed', color="black", arrow_length_ratio=0.07, alpha=0.5) 342 ax.text(r*k_r[0], r*k_r[1], r*k_r[2], r"$\hat{\mathbf{k}}_\mathrm{r}$") 343 344 sm = cm.ScalarMappable(cmap=plt.cm.turbo) 345 sm.set_array([]) 346 cbar = plt.colorbar(sm, ax=ax, orientation="vertical", location="right", 347 shrink=0.7, pad=0.15) 348 349 xticks = cbar.ax.get_yticks() 350 xticklabels = cbar.ax.get_yticklabels() 351 xticklabels = p_min + xticks*(p_max-p_min) 352 xticklabels = [f"{z:.2f}" for z in xticklabels] 353 cbar.ax.set_yticks(xticks) 354 cbar.ax.set_yticklabels(xticklabels) 355 ax.view_init(elev=30., azim=-60) 356 plt.xlabel("x") 357 plt.ylabel("y") 358 ax.set_zlabel("z") 359 ax.set_aspect('auto') 360 plt.suptitle(r"3D visualization of the scattering pattern $f_\mathrm{s}(\hat{\mathbf{k}}_\mathrm{i}, \hat{\mathbf{k}}_\mathrm{s})$") 361 362 ### 363 ### Incident plane cut through the scattering pattern 364 ### 365 k_i = tf.constant(k_i_in, self._dtype.real_dtype) 366 theta_i, phi_i = theta_phi_from_unit_vec(-k_i) 367 theta_r, phi_r = theta_i, phi_i + tf.cast(PI, phi_i.dtype) 368 369 # Pattern around reflected direction 370 theta_s = tf.cast(tf.linspace(0.0, PI/2, 100), dtype=self._dtype.real_dtype) 371 phi_s = tf.broadcast_to(phi_r, theta_s.shape) 372 k_s = r_hat(theta_s, phi_s) 373 k_i_1 = tf.broadcast_to(k_i, k_s.shape) 374 n_hat = tf.constant([[0,0,1]], dtype=k_i.dtype) 375 n_hat = tf.broadcast_to(n_hat, k_i_1.shape) 376 pattern = self(k_i_1, k_s, n_hat) 377 378 # Pattern around incident direction 379 k_s = r_hat(theta_s, phi_s+PI) 380 pattern2 = self(k_i_1, k_s, n_hat) 381 382 fig_cut = plt.figure() 383 plt.polar(theta_s, pattern, color='C0') 384 plt.polar(2*PI-theta_s , pattern2, color='C0') 385 386 ax = fig_cut.axes[0] 387 ax.set_theta_zero_location("N") 388 ax.set_theta_direction(-1) 389 ax.set_thetamin(-90) 390 ax.set_thetamax(90) 391 392 if show_directions: 393 xticks = list(ax.get_xticks()) 394 if not theta_i.numpy() in xticks: 395 xticks += [theta_i.numpy()] 396 if not -theta_i.numpy() in xticks: 397 xticks += [-theta_i.numpy()] 398 ax.set_xticks(xticks) 399 ax.text(-theta_i-10*PI/180, ax.get_yticks()[-1]*2/3, r"$\hat{\mathbf{k}}_\mathrm{i}$", horizontalalignment='center') 400 ax.text(theta_i+10*PI/180, ax.get_yticks()[-1]*2/3, r"$\hat{\mathbf{k}}_\mathrm{r}$", horizontalalignment='center') 401 plt.quiver([0], [0], [np.sin(theta_i)], [np.cos(theta_i)], scale=1., color="grey",) 402 plt.quiver([0], [0], [-np.sin(theta_i)], [np.cos(theta_i)], scale=1., color="grey",) 403 404 plt.title(r"Incident plane cut through the scattering pattern $f_\mathrm{s}(\hat{\mathbf{k}}_\mathrm{i}, \hat{\mathbf{k}}_\mathrm{s})$ ($\phi_\mathrm{s}=\phi_\mathrm{i}+\pi$)") 405 plt.tight_layout() 406 407 return fig_3d, fig_cut 408 409 class LambertianPattern(ScatteringPattern): 410 # pylint: disable=line-too-long 411 r""" 412 Lambertian scattering model from [Degli-Esposti07]_ as given in :eq:`lambertian_model` 413 414 Parameters 415 ---------- 416 dtype : tf.complex64 or tf.complex128 417 Datatype used for all computations. 418 Defaults to `tf.complex64`. 419 420 Input 421 ----- 422 k_i : [batch_size, 3], dtype.real_dtype 423 Incoming directions 424 425 k_s : [batch_size,3], dtype.real_dtype 426 Outgoing directions 427 428 Output 429 ------ 430 pattern : [batch_size], dtype.real_dtype 431 Scattering pattern 432 433 Example 434 ------- 435 >>> LambertianPattern().visualize() 436 437 .. figure:: ../figures/lambertian_pattern_3d.png 438 :align: center 439 440 .. figure:: ../figures/lambertian_pattern_cut.png 441 :align: center 442 """ 443 def __init__(self, dtype=tf.complex64): 444 super().__init__(alpha_r=0, alpha_i=1, lambda_=1, dtype=dtype) 445 446 class DirectivePattern(ScatteringPattern): 447 # pylint: disable=line-too-long 448 r""" 449 Directive scattering model from [Degli-Esposti07]_ as given in :eq:`directive_model` 450 451 Parameters 452 ---------- 453 alpha_r : int, [1,2,...] 454 Parameter related to the width of the scattering lobe in the 455 direction of the specular reflection. 456 457 dtype : tf.complex64 or tf.complex128 458 Datatype used for all computations. 459 Defaults to `tf.complex64`. 460 461 Input 462 ----- 463 k_i : [batch_size, 3], dtype.real_dtype 464 Incoming directions 465 466 k_s : [batch_size,3], dtype.real_dtype 467 Outgoing directions 468 469 Output 470 ------ 471 pattern : [batch_size], dtype.real_dtype 472 Scattering pattern 473 474 Example 475 ------- 476 >>> DirectivePattern(alpha_r=10).visualize() 477 478 .. figure:: ../figures/directive_pattern_3d.png 479 :align: center 480 481 .. figure:: ../figures/directive_pattern_cut.png 482 :align: center 483 """ 484 def __init__(self, 485 alpha_r, 486 dtype=tf.complex64): 487 super().__init__(alpha_r=alpha_r, alpha_i=1, lambda_=1, dtype=dtype) 488 489 class BackscatteringPattern(ScatteringPattern): 490 # pylint: disable=line-too-long 491 r""" 492 Backscattering model from [Degli-Esposti07]_ as given in :eq:`backscattering_model` 493 494 The parameter ``lambda_`` can be assigned to a TensorFlow variable 495 or tensor. In the latter case, the tensor can be the output of a callable, such as 496 a Keras layer implementing a neural network. 497 In the former case, it can be set to a trainable variable: 498 499 .. code-block:: Python 500 501 sp = BackscatteringPattern(alpha_r=3, 502 alpha_i=5, 503 lambda_=tf.Variable(0.3, dtype=tf.float32)) 504 505 Parameters 506 ---------- 507 alpha_r : int, [1,2,...] 508 Parameter related to the width of the scattering lobe in the 509 direction of the specular reflection. 510 511 alpha_i : int, [1,2,...] 512 Parameter related to the width of the scattering lobe in the 513 incoming direction. 514 515 lambda_ : float, [0,1] 516 Parameter determining the percentage of the diffusely 517 reflected energy in the lobe around the specular reflection. 518 519 dtype : tf.complex64 or tf.complex128 520 Datatype used for all computations. 521 Defaults to `tf.complex64`. 522 523 Input 524 ----- 525 k_i : [batch_size, 3], dtype.real_dtype 526 Incoming directions 527 528 k_s : [batch_size,3], dtype.real_dtype 529 Outgoing directions 530 531 Output 532 ------ 533 pattern : [batch_size], dtype.real_dtype 534 Scattering pattern 535 536 Example 537 ------- 538 >>> BackscatteringPattern(alpha_r=20, alpha_i=30, lambda_=0.7).visualize() 539 540 .. figure:: ../figures/backscattering_pattern_3d.png 541 :align: center 542 543 .. figure:: ../figures/backscattering_pattern_cut.png 544 :align: center 545 """ 546 def __init__(self, 547 alpha_r, 548 alpha_i, 549 lambda_, 550 dtype=tf.complex64): 551 super().__init__(alpha_r=alpha_r, alpha_i=alpha_i, lambda_=lambda_, 552 dtype=dtype)