Compute fields at multiple frequencies in parallel in one ray tracing

[Anplus]

Description

I have revised the frequency check in scene.py and the frequency update for all ITU materials‘ properties in itu_material.py using TensorFlow. The old code was implemented in raw Python and NumPy, which could not be used to construct the computation graph needed for parallel computation. This enhancement enables the computation of fields at multiple frequencies in parallel within a single ray tracing process. Specifically, we can first trace all paths, and then use tf.map_fn(process_freq, freq, dtype=tf.complex64) to construct a parallel computation for field computation at multiple frequencies, where the process_freq function updates the frequency and recomputes the field with the given traced paths.

Additionally, I have added a Jupyter notebook as an example to illustrate this feature.

• Fixes a bug?

NA.

• Adds a new feature?

This feature allows for parallel computation of fields at multiple frequencies during ray tracing, enhancing computation efficiency. I think this is a good update for the discussion: #527

• Introduces API changes?

No API changes.

• Other contributions

Checklist

• Detailed description
• Added references to issues and discussions
• Added / modified documentation as needed
• Added / modified unit tests as needed
• Passes all tests
• Lint the code
• Performed a self review
• Ensure you Signed-off the commits. Required to accept contributions!
• Co-authored with someone? Add Co-authored-by: user@domain and ensure they signed off their commits too.

[jhoydis]

Hi @Anplus,

Thank you for this PR which seems like a useful feature. We'll review it and come back to you. This might take a bit of time though.

[ZitongLan]

Hi Zhenlin @Anplus , I have been recently back to this sionna related project. Thank you for your commit that relates to my previosu discussion. And I found this feature very useful.

However, currently I am working towards to build a differiable rendering for multiple different frequencies, so that my materials should also have different properties (variables) at differet discrete frequenceis. I achieve this by using frequency_call_back function to init the material properties. However, I just found that these variables are not involve in the computational graph. Do you have any idea of what is wrong?

below is a minimum code that I am experimenting with

freq_grid_hz = tf.constant(np.arange(3e9, 3e9 + 300e6, 50e6, dtype=np.float32)) # shape (N,)
conductivity_table = tf.Variable(100*np.ones_like(freq_grid_hz, np.float32), trainable=True)

def custom_properties(f_hz):
    N = freq_grid_hz.shape[0]

    f_hz = tf.convert_to_tensor(f_hz, tf.float32)

    # map each frequency to its index in freq_grid_hz
    idx = tf.where(freq_grid_hz == f_hz)[0][0]  # Get the index of the frequency in freq_grid_hz

    relative_permittivities = tf.constant(1, dtype=tf.float32)
    conductivities = tf.gather(conductivity_table, idx) 

    return (relative_permittivities, conductivities)

for mat in scene.radio_materials.values():
    if mat.is_used:        
        new_mat = RadioMaterial(mat.name + "_train", frequency_update_callback=custom_properties)

        scene.add(new_mat)
        trainable_materials.append(new_mat)
        original_materials.append(mat)