From 363d2875dc4b4b3f987ac733e75f85d4bb21d1e3 Mon Sep 17 00:00:00 2001
From: landmanbester <lbester@sarao.ac.za>
Date: Mon, 11 Nov 2024 17:16:52 +0200
Subject: [PATCH] generalise band mapping for multi-ms and multi-spw data

---
 pfb/utils/misc.py   |  6 ------
 pfb/workers/init.py | 39 ++++++++++++++++++++++++++++++---------
 2 files changed, 30 insertions(+), 15 deletions(-)

diff --git a/pfb/utils/misc.py b/pfb/utils/misc.py
index eef8bf80..b500eb95 100644
--- a/pfb/utils/misc.py
+++ b/pfb/utils/misc.py
@@ -461,12 +461,6 @@ def construct_mappings(ms_name,
                 raise RuntimeError("Something went wrong constructing the "
                                    "frequency mapping. sum(fchunks != nchan)")
 
-            # nfreq_chunks = nchan_in // cpit
-            # freq_chunks = (cpit,)*nfreq_chunks
-            # rem = nchan_in - nfreq_chunks * cpit
-            # if rem:
-            #     freq_chunks += (rem,)
-
             ms_chunks[ms].append({'row': row_chunks,
                                   'chan': freq_chunks})
 
diff --git a/pfb/workers/init.py b/pfb/workers/init.py
index 312a037f..7381b516 100644
--- a/pfb/workers/init.py
+++ b/pfb/workers/init.py
@@ -207,20 +207,41 @@ def _init(**kw):
     else:
         print(f"No weights provided, using unity weights", file=log)
 
-
-    # band mapping
-    ddid2bid = {}
-    b0 = 0
+    # distinct freq groups
+    igroup = 0
+    sgroup = 0
+    freq_groups = {}
     for ms in opts.ms:
-        for idt, fmap in freq_mapping[ms].items():
+        for idt, freq in freqs[ms].items():
             ilo = idt.find('DDID') + 4
             ihi = idt.rfind('_')
             ddid = int(idt[ilo:ihi])
             if (opts.ddids is not None) and (ddid not in opts.ddids):
                 continue
-            elif ddid not in ddid2bid.keys():
-                ddid2bid[ddid] = b0
-                b0 += fmap['counts'].size
+            if not len(freq_groups.keys()):
+                freq_groups[igroup] = {}
+                freq_groups[igroup]['freq'] = freq
+                freq_groups[igroup]['sgroup'] = sgroup
+                igroup += 1
+                sgroup += freq_mapping[ms][idt]['counts'].size
+            else:
+                for i, fs in freq_groups.items():
+                    if not np.all(freq == fs['freq']):
+                        freq_groups[igroup] = {}
+                        freq_groups[igroup]['freq'] = freq
+                        freq_groups[igroup]['sgroup'] = sgroup
+                        igroup += 1
+                        sgroup += freq_mapping[ms][idt]['counts'].size
+
+    # band mapping
+    msddid2bid = {}
+    for ms in opts.ms:
+        msddid2bid[ms] = {}
+        for idt, freq in freqs[ms].items():
+            # find group where it matches
+            for igroup, fs in freq_groups.items():
+                if np.all(freq == fs['freq']):
+                    msddid2bid[ms][idt] = fs['sgroup']
 
     # a flat list to use with as_completed
     datasets = []
@@ -264,7 +285,7 @@ def _init(**kw):
 
                 fitr = enumerate(zip(freq_mapping[ms][idt]['start_indices'],
                                      freq_mapping[ms][idt]['counts']))
-                b0 = ddid2bid[ddid]
+                b0 = msddid2bid[ms][idt]
                 for fi, (flow, fcounts) in fitr:
                     Inu = slice(flow, flow + fcounts)