Commits
Benjamin Bean authored 5866e3b4aa2
1 1 | import copy |
2 2 | import numpy as np |
3 3 | import os |
4 4 | |
5 5 | from casatasks.private.casa_transition import is_CASA6 |
6 - | from casampi.MPIEnvironment import MPIEnvironment |
7 6 | if is_CASA6: |
8 7 | from casatasks import casalog |
9 8 | else: |
10 9 | from taskinit import * |
11 10 | |
12 11 | class SummaryMinor(dict): |
13 12 | """A return object from tclean's minor cycles. Gathers the information together |
14 13 | in a way that makes it easier to query for exactly what you want. |
15 14 | |
16 15 | The structure for this dictionary is: |
192 191 | chans = list(np.sort(np.unique(matrix[oldChanIdx]))) |
193 192 | chans = [int(x) for x in chans] |
194 193 | if uss: |
195 194 | pols = [0] |
196 195 | else: |
197 196 | pols = list(np.sort(np.unique(matrix[oldPolIdx]))) |
198 197 | pols = [int(x) for x in pols] |
199 198 | ncycles = 0 |
200 199 | if len(chans) > 0 and len(pols) > 0: |
201 200 | ncycles = int( ncols / (len(chans)*len(pols)) ) |
202 - | if uss and MPIEnvironment.is_mpi_enabled: |
203 - | # This is necessary because we may have an odd number of "channels" due to each process getting only a subchunk. |
204 - | # Example: |
205 - | # Process 1 gets polarities 0-1, process 2 gets polarity 2 |
206 - | # Each of them assigns channel id = chan + pol * nsubpols |
207 - | # Process 1 assigns channel ids [0,2], Process 2 assigns channel id 0. |
208 - | # This hack is not needed when not using a small summary minor because we have the extra knowledge of the polarities, instead of mapping polarities + channels onto chunks. |
209 - | chanslist = matrix[oldChanIdx].tolist() |
210 - | for chan in chans: |
211 - | singlechan_occurances = list(filter(lambda x: x == chan, chanslist)) |
212 - | ncycles = max(ncycles, len(singlechan_occurances)) |
201 + | if uss: |
202 + | try: |
203 + | from casampi.MPIEnvironment import MPIEnvironment |
204 + | if MPIEnvironment.is_mpi_enabled: |
205 + | # This is necessary because we may have an odd number of "channels" due to each process getting only a subchunk. |
206 + | # Example: |
207 + | # Process 1 gets polarities 0-1, process 2 gets polarity 2 |
208 + | # Each of them assigns channel id = chan + pol * nsubpols |
209 + | # Process 1 assigns channel ids [0,2], Process 2 assigns channel id 0. |
210 + | # This hack is not needed when not using a small summary minor because we have the extra knowledge of the polarities, instead of mapping polarities + channels onto chunks. |
211 + | chanslist = matrix[oldChanIdx].tolist() |
212 + | for chan in chans: |
213 + | singlechan_occurances = list(filter(lambda x: x == chan, chanslist)) |
214 + | ncycles = max(ncycles, len(singlechan_occurances)) |
215 + | except ModuleNotFoundError as e: |
216 + | raise |
213 217 | |
214 218 | # ret is the return dictionary[chans][pols][rows][cycles] |
215 219 | # cummulativeCnt counts how many cols we've read for each channel/polarity/row |
216 220 | ret = {desc:[0]*ncycles for desc in SummaryMinor._getRowDescriptions(uss)} |
217 221 | # channel and polarity information is in the indexing, don't need to add it to the return dict |
218 222 | for desc in ["chan", "pol"]: |
219 223 | if desc in ret: |
220 224 | del ret[desc] |
221 225 | ret = {pol:copy.deepcopy(ret) for pol in pols} |
222 226 | ret = {chan:copy.deepcopy(ret) for chan in chans} |