Inconsistent dimensions when concatenating traced paths of UEs sampled batch-wise from coverage map.

[naperman]

This issue is a follow-up on the discussion in this thread.

I wrote a concat function as follows:

def concatenate_traced_paths(traced_paths, traced_paths_batch):
    for path_type_idx in range(len(traced_paths)):
        traced_paths[path_type_idx].a = tf.concat([traced_paths[path_type_idx].a, traced_paths_batch[path_type_idx].a], axis=0)
        traced_paths[path_type_idx].doppler = tf.concat([traced_paths[path_type_idx].doppler, traced_paths_batch[path_type_idx].doppler], axis=0)
        traced_paths[path_type_idx].mask = tf.concat([traced_paths[path_type_idx].mask, traced_paths_batch[path_type_idx].mask], axis=0)
        traced_paths[path_type_idx].objects = tf.concat([traced_paths[path_type_idx].objects, traced_paths_batch[path_type_idx].objects], axis=1)  # Concatenate along second axis
        traced_paths[path_type_idx].phi_r = tf.concat([traced_paths[path_type_idx].phi_r, traced_paths_batch[path_type_idx].phi_r], axis=0)
        traced_paths[path_type_idx].phi_t = tf.concat([traced_paths[path_type_idx].phi_t, traced_paths_batch[path_type_idx].phi_t], axis=0)
        traced_paths[path_type_idx].targets = tf.concat([traced_paths[path_type_idx].targets, traced_paths_batch[path_type_idx].targets], axis=0)
        traced_paths[path_type_idx].targets_sources_mask = tf.concat([traced_paths[path_type_idx].targets_sources_mask, traced_paths_batch[path_type_idx].targets_sources_mask], axis=0)
        traced_paths[path_type_idx].tau = tf.concat([traced_paths[path_type_idx].tau, traced_paths_batch[path_type_idx].tau], axis=0)
        traced_paths[path_type_idx].theta_r = tf.concat([traced_paths[path_type_idx].theta_r, traced_paths_batch[path_type_idx].theta_r], axis=0)
        traced_paths[path_type_idx].theta_t = tf.concat([traced_paths[path_type_idx].theta_t, traced_paths_batch[path_type_idx].theta_t], axis=0)
        traced_paths[path_type_idx].vertices = tf.concat([traced_paths[path_type_idx].vertices, traced_paths_batch[path_type_idx].vertices], axis=1) # Concatenate along second axis

I set doppler as 0 for all UEs in all batches, but i ran into the following error when the loop reached the second batch:

---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
Cell In[7], line 45
     34 else:
     35     print("---------- GT Scene ----------")
     36     (
     37         traced_paths_gt,
     38         a_gt,
     39         tau_gt,
     40         ue_pos_gt,
     41         ue_pos_uncut_gt,
     42         scene_objects_materials_gt,
     43         scene_objects_positions_gt,
     44         scene_objects_heights_gt,
---> 45     ) = scene_ops(path=scene_path_gt, num_sample_ues=num_UEs, num_ues_per_batch=100, ue_positions=None, debug=True, prune_cir=False)
     46     print(f"$ a_gt shape            = {a_gt.shape}")
     47     print(f"$ tau_gt shape          = {tau_gt.shape}")

Cell In[5], line 250, in scene_ops(path, num_sample_ues, num_ues_per_batch, ue_positions, debug, return_scene, prune_cir)
    247 else:
    248     # Concatenate along the num_tx dimension
    249     ue_pos = np.concatenate([ue_pos, ue_pos_batch], axis=0)
--> 250     traced_paths = concatenate_traced_paths(traced_paths, traced_paths_batch) 
    251     a = np.concatenate([a, a_batch], axis=3)
    252     tau = np.concatenate([tau, tau_batch], axis=3)  # axis=2

Cell In[4], line 6, in concatenate_traced_paths(traced_paths, traced_paths_batch)
      4 for path_type_idx in range(len(traced_paths)):
      5     traced_paths[path_type_idx].a = tf.concat([traced_paths[path_type_idx].a, traced_paths_batch[path_type_idx].a], axis=0)
----> 6     traced_paths[path_type_idx].doppler = tf.concat([traced_paths[path_type_idx].doppler, traced_paths_batch[path_type_idx].doppler], axis=0)
      7     traced_paths[path_type_idx].mask = tf.concat([traced_paths[path_type_idx].mask, traced_paths_batch[path_type_idx].mask], axis=0)
      8     traced_paths[path_type_idx].objects = tf.concat([traced_paths[path_type_idx].objects, traced_paths_batch[path_type_idx].objects], axis=1)  # Concatenate along second axis

File ~/.local/lib/python3.10/site-packages/tensorflow/python/util/traceback_utils.py:153, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    151 except Exception as e:
    152   filtered_tb = _process_traceback_frames(e.__traceback__)
--> 153   raise e.with_traceback(filtered_tb) from None
    154 finally:
    155   del filtered_tb

File ~/.local/lib/python3.10/site-packages/tensorflow/python/framework/ops.py:5883, in raise_from_not_ok_status(e, name)
   5881 def raise_from_not_ok_status(e, name) -> NoReturn:
   5882   e.message += (" name: " + str(name if name is not None else ""))
-> 5883   raise core._status_to_exception(e) from None

InvalidArgumentError: {{function_node __wrapped__ConcatV2_N_2_device_/job:localhost/replica:0/task:0/device:GPU:0}} ConcatOp : Dimension 2 in both shapes must be equal: shape[0] = [100,4,39] vs. shape[1] = [100,4,38] [Op:ConcatV2] name: concat

One potential solution i tried was to call apply_doppler on the traced paths rather than the final paths(traced+fields). This resulted in the following error:

[233](vscode-notebook-cell:?execution_count=5&line=233) for traced_paths_component in traced_paths_batch:
--> [234](vscode-notebook-cell:?execution_count=5&line=234)     traced_paths_component.apply_doppler(  
    [235](vscode-notebook-cell:?execution_count=5&line=235)         sampling_frequency=subcarrier_spacing,
    [236](vscode-notebook-cell:?execution_count=5&line=236)         num_time_steps=1,
    [237](vscode-notebook-cell:?execution_count=5&line=237)         tx_velocities=[0.0, 0.0, 0],
    [238](vscode-notebook-cell:?execution_count=5&line=238)         rx_velocities=[0.0, 0.0, 0],
    [239](vscode-notebook-cell:?execution_count=5&line=239)     )
    [241](vscode-notebook-cell:?execution_count=5&line=241) paths = scene.compute_fields(*traced_paths_batch)
    [243](vscode-notebook-cell:?execution_count=5&line=243) # Transform paths into channel impulse responses

File ~/src/sionna/sionna/rt/paths.py:543, in Paths.apply_doppler(self, sampling_frequency, num_time_steps, tx_velocities, rx_velocities)
    [541](https://file+.vscode-resource.vscode-cdn.net/home/nperiasamy/src/notebooks-sionnart/~/src/sionna/sionna/rt/paths.py:541)     ds += two_pi*self.doppler[..., tf.newaxis, :, tf.newaxis, :]
    [542](https://file+.vscode-resource.vscode-cdn.net/home/nperiasamy/src/notebooks-sionnart/~/src/sionna/sionna/rt/paths.py:542) else:
--> [543](https://file+.vscode-resource.vscode-cdn.net/home/nperiasamy/src/notebooks-sionnart/~/src/sionna/sionna/rt/paths.py:543)     ds += two_pi*self.doppler
    [545](https://file+.vscode-resource.vscode-cdn.net/home/nperiasamy/src/notebooks-sionnart/~/src/sionna/sionna/rt/paths.py:545) # Expand for the time sample dimension
    [546](https://file+.vscode-resource.vscode-cdn.net/home/nperiasamy/src/notebooks-sionnart/~/src/sionna/sionna/rt/paths.py:546) # [batch_dim, num_rx, num_rx_ant, num_tx, num_tx_ant, max_num_paths, 1]
    [547](https://file+.vscode-resource.vscode-cdn.net/home/nperiasamy/src/notebooks-sionnart/~/src/sionna/sionna/rt/paths.py:547) # or
    [548](https://file+.vscode-resource.vscode-cdn.net/home/nperiasamy/src/notebooks-sionnart/~/src/sionna/sionna/rt/paths.py:548) # [batch_dim, num_rx, 1, num_tx, 1, max_num_paths, 1]
    [549](https://file+.vscode-resource.vscode-cdn.net/home/nperiasamy/src/notebooks-sionnart/~/src/sionna/sionna/rt/paths.py:549) ds = tf.expand_dims(ds, axis=-1)

File ~/.local/lib/python3.10/site-packages/tensorflow/python/util/traceback_utils.py:153, in filter_traceback.<locals>.error_handler(*args, **kwargs)
    [151](https://file+.vscode-resource.vscode-cdn.net/home/nperiasamy/src/notebooks-sionnart/~/.local/lib/python3.10/site-packages/tensorflow/python/util/traceback_utils.py:151) except Exception as e:
    [152](https://file+.vscode-resource.vscode-cdn.net/home/nperiasamy/src/notebooks-sionnart/~/.local/lib/python3.10/site-packages/tensorflow/python/util/traceback_utils.py:152)   filtered_tb = _process_traceback_frames(e.__traceback__)
--> [153](https://file+.vscode-resource.vscode-cdn.net/home/nperiasamy/src/notebooks-sionnart/~/.local/lib/python3.10/site-packages/tensorflow/python/util/traceback_utils.py:153)   raise e.with_traceback(filtered_tb) from None
    [154](https://file+.vscode-resource.vscode-cdn.net/home/nperiasamy/src/notebooks-sionnart/~/.local/lib/python3.10/site-packages/tensorflow/python/util/traceback_utils.py:154) finally:
    [155](https://file+.vscode-resource.vscode-cdn.net/home/nperiasamy/src/notebooks-sionnart/~/.local/lib/python3.10/site-packages/tensorflow/python/util/traceback_utils.py:155)   del filtered_tb

File ~/.local/lib/python3.10/site-packages/tensorflow/python/framework/ops.py:5883, in raise_from_not_ok_status(e, name)
   [5881](https://file+.vscode-resource.vscode-cdn.net/home/nperiasamy/src/notebooks-sionnart/~/.local/lib/python3.10/site-packages/tensorflow/python/framework/ops.py:5881) def raise_from_not_ok_status(e, name) -> NoReturn:
   [5882](https://file+.vscode-resource.vscode-cdn.net/home/nperiasamy/src/notebooks-sionnart/~/.local/lib/python3.10/site-packages/tensorflow/python/framework/ops.py:5882)   e.message += (" name: " + str(name if name is not None else ""))
-> [5883](https://file+.vscode-resource.vscode-cdn.net/home/nperiasamy/src/notebooks-sionnart/~/.local/lib/python3.10/site-packages/tensorflow/python/framework/ops.py:5883)   raise core._status_to_exception(e) from None

InvalidArgumentError: {{function_node __wrapped__AddV2_device_/job:localhost/replica:0/task:0/device:GPU:0}} Input for dynamic binary or n-ary op lowering was of a rank greater than 5 [Op:AddV2] name:

I'm not sure what is causing the above error. It is valid to call apply_doppler on traced paths rather than final paths? If so, would it make any difference?

I can understand that the number of traced paths can be slightly different for different batches of sampled UEs, but it would very useful to have an extra parameter to specify number of paths in scene.trace_paths() just like we can do with paths.cir(num_paths=some_number), where the excess paths can have nan or zero in the necessary fields.

If there is any other workaround for this than to reshape the properties of the traced_paths and fill in empty values, please suggest.