From 6b5c4148c8820415fd6531a3535e9c3796c145b9 Mon Sep 17 00:00:00 2001
From: Xavier Arteaga <xavier@softwareradiosystems.com>
Date: Wed, 11 Nov 2020 18:49:00 +0100
Subject: [PATCH] DL Channel estimation results accepts nof RE

---
 lib/include/srslte/phy/ch_estimation/chest_dl.h | 1 +
 lib/src/phy/ch_estimation/chest_dl.c            | 9 +++++++--
 2 files changed, 8 insertions(+), 2 deletions(-)

diff --git a/lib/include/srslte/phy/ch_estimation/chest_dl.h b/lib/include/srslte/phy/ch_estimation/chest_dl.h
index 6ff10f549..19cf7c764 100644
--- a/lib/include/srslte/phy/ch_estimation/chest_dl.h
+++ b/lib/include/srslte/phy/ch_estimation/chest_dl.h
@@ -141,6 +141,7 @@ SRSLTE_API int srslte_chest_dl_init(srslte_chest_dl_t* q, uint32_t max_prb, uint
 SRSLTE_API void srslte_chest_dl_free(srslte_chest_dl_t* q);
 
 SRSLTE_API int srslte_chest_dl_res_init(srslte_chest_dl_res_t* q, uint32_t max_prb);
+SRSLTE_API int srslte_chest_dl_res_init_re(srslte_chest_dl_res_t* q, uint32_t nof_re);
 
 SRSLTE_API void srslte_chest_dl_res_set_identity(srslte_chest_dl_res_t* q);
 
diff --git a/lib/src/phy/ch_estimation/chest_dl.c b/lib/src/phy/ch_estimation/chest_dl.c
index f7c0f31ab..caca92fd3 100644
--- a/lib/src/phy/ch_estimation/chest_dl.c
+++ b/lib/src/phy/ch_estimation/chest_dl.c
@@ -207,9 +207,14 @@ void srslte_chest_dl_free(srslte_chest_dl_t* q)
 }
 
 int srslte_chest_dl_res_init(srslte_chest_dl_res_t* q, uint32_t max_prb)
+{
+  return srslte_chest_dl_res_init_re(q, SRSLTE_SF_LEN_RE(max_prb, SRSLTE_CP_NORM));
+}
+
+int srslte_chest_dl_res_init_re(srslte_chest_dl_res_t* q, uint32_t nof_re)
 {
   bzero(q, sizeof(srslte_chest_dl_res_t));
-  q->nof_re = SRSLTE_SF_LEN_RE(max_prb, SRSLTE_CP_NORM);
+  q->nof_re = nof_re;
   for (uint32_t i = 0; i < SRSLTE_MAX_PORTS; i++) {
     for (uint32_t j = 0; j < SRSLTE_MAX_PORTS; j++) {
       q->ce[i][j] = srslte_vec_cf_malloc(q->nof_re);
@@ -217,7 +222,7 @@ int srslte_chest_dl_res_init(srslte_chest_dl_res_t* q, uint32_t max_prb)
         perror("malloc");
         return -1;
       }
-      srslte_vec_cf_zero(q->ce[i][j], SRSLTE_SF_LEN_RE(max_prb, SRSLTE_CP_NORM));
+      srslte_vec_cf_zero(q->ce[i][j], nof_re);
     }
   }
   return 0;