hermespy-rt

commit 8d658d7c8d04df7712e58f56bbb4b8df292ce501
parent 379a405fb71653adf06384ff3aee8f337651b5af
Author: Egor Achkasov <eaachkasov@edu.hse.ru>
Date:   Wed, 12 Mar 2025 00:57:55 +0100

Add PathInfo.freq_shift field

Diffstat:
M
compute_paths_pybind11.cpp
 | 
22
++++++++++++++++++----
M
inc/compute_paths.h
 | 
1
+
M
test.c
 | 
6
++++--
M
test.py
 | 
32
++++++++++++++++++++++++++------

4 files changed, 49 insertions(+), 12 deletions(-)
diff --git a/compute_paths_pybind11.cpp b/compute_paths_pybind11.cpp
@@ -46,6 +46,7 @@ public:
     py::array_t<std::complex<float>> a_te;
     py::array_t<std::complex<float>> a_tm;
     py::array_t<float> tau;
+    py::array_t<float> freq_shift;
 
     PathsInfoPython(const PathsInfo& paths, size_t num_rx, size_t num_tx) {
         num_paths = paths.num_paths;
@@ -66,7 +67,17 @@ public:
         a_te = make_complex_array(paths.a_te_re, paths.a_te_im, num_rx, num_tx, paths.num_paths);
         a_tm = make_complex_array(paths.a_tm_re, paths.a_tm_im, num_rx, num_tx, paths.num_paths);
 
-        tau = py::array_t<float>({num_rx, num_tx, (size_t)paths.num_paths}, paths.tau, py::capsule(paths.tau, capsule_deleter));
+        tau = py::array_t<float>(
+            {num_rx, num_tx, (size_t)paths.num_paths},
+            paths.tau,
+            py::capsule(paths.tau, capsule_deleter)
+        );
+
+        freq_shift = py::array_t<float>(
+            {num_rx, num_tx, (size_t)paths.num_paths},
+            paths.freq_shift,
+            py::capsule(paths.freq_shift, capsule_deleter)
+        );
     }
 };
 
@@ -98,7 +109,8 @@ compute_paths_wrapper(
         .a_te_im = new float[num_rx * num_tx],
         .a_tm_re = new float[num_rx * num_tx],
         .a_tm_im = new float[num_rx * num_tx],
-        .tau = new float[num_rx * num_tx]
+        .tau = new float[num_rx * num_tx],
+        .freq_shift = new float[num_rx * num_tx]
     };
     PathsInfo scatter = {
         .num_paths = (uint32_t)(num_bounces * num_paths),
@@ -108,7 +120,8 @@ compute_paths_wrapper(
         .a_te_im = new float[num_rx * num_tx * num_bounces * num_paths],
         .a_tm_re = new float[num_rx * num_tx * num_bounces * num_paths],
         .a_tm_im = new float[num_rx * num_tx * num_bounces * num_paths],
-        .tau = new float[num_rx * num_tx * num_bounces * num_paths]
+        .tau = new float[num_rx * num_tx * num_bounces * num_paths],
+        .freq_shift = new float[num_rx * num_tx * num_bounces * num_paths]
     };
 
     // Call the C function
@@ -141,7 +154,8 @@ PYBIND11_MODULE(rt, m) {
         .def_readonly("directions_tx", &PathsInfoPython::directions_tx)
         .def_readonly("a_te", &PathsInfoPython::a_te)
         .def_readonly("a_tm", &PathsInfoPython::a_tm)
-        .def_readonly("tau", &PathsInfoPython::tau);
+        .def_readonly("tau", &PathsInfoPython::tau)
+        .def_readonly("freq_shift", &PathsInfoPython::freq_shift);
 
     // TODO write a proper docstring
     m.def("compute_paths", &compute_paths_wrapper, "Compute gains and delays",
diff --git a/inc/compute_paths.h b/inc/compute_paths.h
@@ -17,6 +17,7 @@ typedef struct {
     float *a_tm_re; /* shape (num_rx, num_tx, num_paths) */
     float *a_tm_im; /* shape (num_rx, num_tx, num_paths) */
     float *tau; /* shape (num_rx, num_tx, num_paths) */
+    float *freq_shift; /* shape (num_rx, num_tx, num_paths) */
 } PathsInfo;
 
 /** Compute gains and delays between tx and rx in a 3D scene.
diff --git a/test.c b/test.c
@@ -24,7 +24,8 @@ int main(int argc, char **argv)
     .a_te_im = (float*)malloc(g_numRx * g_numTx * sizeof(float)),
     .a_tm_re = (float*)malloc(g_numRx * g_numTx * sizeof(float)),
     .a_tm_im = (float*)malloc(g_numRx * g_numTx * sizeof(float)),
-    .tau = (float*)malloc(g_numRx * g_numTx * sizeof(float))
+    .tau = (float*)malloc(g_numRx * g_numTx * sizeof(float)),
+    .freq_shift = (float*)malloc(g_numRx * g_numTx * sizeof(float))
   };
   PathsInfo scatter = {
     .num_paths = g_numBounces * g_numPaths,
@@ -34,7 +35,8 @@ int main(int argc, char **argv)
     .a_te_im = (float*)malloc(g_numRx * g_numTx * g_numBounces * g_numPaths * sizeof(float)),
     .a_tm_re = (float*)malloc(g_numRx * g_numTx * g_numBounces * g_numPaths * sizeof(float)),
     .a_tm_im = (float*)malloc(g_numRx * g_numTx * g_numBounces * g_numPaths * sizeof(float)),
-    .tau = (float*)malloc(g_numRx * g_numTx * g_numBounces * g_numPaths * sizeof(float))
+    .tau = (float*)malloc(g_numRx * g_numTx * g_numBounces * g_numPaths * sizeof(float)),
+    .freq_shift = (float*)malloc(g_numRx * g_numTx * g_numBounces * g_numPaths * sizeof(float))
   };
   compute_paths(
     argv[1],
diff --git a/test.py b/test.py
@@ -54,11 +54,31 @@ print(f"scatter.a_tm.real min, max: {np.min(scatter.a_tm.real)}, {np.max(scatter
 print(f"scatter.a_tm.imag min, max: {np.min(scatter.a_tm.imag)}, {np.max(scatter.a_tm.imag)}")
 
 # Asssert shapes
+# LoS
 assert los.num_paths == 1
+assert (
+    (num_rx, num_tx, 1, 3)
+    == los.directions_rx.shape
+    == los.directions_tx.shape
+)
+assert (
+    (num_rx, num_tx, 1)
+    == los.a_te.shape
+    == los.a_tm.shape
+    == los.tau.shape
+    == los.freq_shift.shape
+)
+# Scatter
 assert scatter.num_paths == num_bounces * num_paths
-assert los.directions_rx.shape == (num_rx, num_tx, 1, 3)
-assert los.directions_tx.shape == (num_rx, num_tx, 1, 3)
-assert scatter.directions_rx.shape == (num_rx, num_tx, scatter.num_paths, 3)
-assert scatter.directions_tx.shape == (num_rx, num_tx, scatter.num_paths, 3)
-assert los.a_te.shape == los.a_tm.shape == los.tau.shape == (num_rx, num_tx, 1)
-assert scatter.a_te.shape == scatter.a_tm.shape == scatter.tau.shape == (num_rx, num_tx, scatter.num_paths)
+assert (
+    (num_rx, num_tx, scatter.num_paths, 3)
+    == scatter.directions_rx.shape
+    == scatter.directions_tx.shape
+)
+assert (
+    (num_rx, num_tx, scatter.num_paths)
+    == scatter.a_te.shape
+    == scatter.a_tm.shape
+    == scatter.tau.shape
+    == scatter.freq_shift.shape
+)