Merge branch 'observation_slicing' into 'master'

Observation slicing

See merge request !26

resolve/__init__.py

@@ -26,3 +26,4 @@ from .polarization_matrix_exponential import *
 from .response import MfResponse, ResponseDistributor, StokesIResponse, SingleResponse
 from .simple_operators import *
 from .util import *
+from .extra import mpi_load

resolve/data/observation.py

@@ -186,6 +186,7 @@ class SingleDishObservation(BaseObservation):
     freq : numpy.ndarray
         Contains the measured frequencies. Shape (n_channels)
     """
+
     def __init__(self, pointings, data, weight, polarization, freq):
         my_assert_isinstance(pointings, Directions)
         my_assert_isinstance(polarization, Polarization)
@@ -306,7 +307,6 @@ class Observation(BaseObservation):
 
         self._eq_attributes = "_direction", "_polarization", "_freq", "_antpos", "_vis", "_weight"
 
-
     @onlymaster
     def save(self, file_name, compress):
         dct = dict(
@@ -332,17 +332,20 @@ class Observation(BaseObservation):
             if val.size == 0:
                 val = None
             antpos.append(val)
+        antpos = AntennaPositions.from_list(antpos)
         pol = Polarization.from_list(dct["polarization"])
         direction = Direction.from_list(dct["direction"])
-        slc = slice(None) if lo_hi_index is None else slice(*lo_hi_index)
-        return Observation(
-            AntennaPositions.from_list(antpos),
-            dct["vis"][..., slc],
-            dct["weight"][..., slc],
-            pol,
-            dct["freq"][slc],
-            direction,
-        )
+        vis = dct["vis"]
+        wgt = dct["weight"]
+        freq = dct["freq"]
+        if lo_hi_index is not None:
+            slc = slice(*lo_hi_index)
+            # Convert view into its own array
+            vis = vis[..., slc].copy()
+            wgt = wgt[..., slc].copy()
+            freq = freq[slc].copy()
+        del dct
+        return Observation(antpos, vis, wgt, pol, freq, direction)
 
     def flags_to_nan(self):
         if self.fraction_useful == 1.:
@@ -365,32 +368,27 @@ class Observation(BaseObservation):
         if comm is None:
             local_imaging_bands = range(n_imaging_bands)
         else:
-            local_imaging_bands = range(
-                *ift.utilities.shareRange(
-                    n_imaging_bands, comm.Get_size(), comm.Get_rank()
-                )
-            )
+            lo, hi = ift.utilities.shareRange( n_imaging_bands, comm.Get_size(), comm.Get_rank())
+            local_imaging_bands = range(lo, hi)
         full_obs = Observation.load(file_name)
-        obs_list = [
-            full_obs.get_freqs_by_slice(
-                slice(*ift.utilities.shareRange(len(global_freqs), n_imaging_bands, ii))
-            )
-            for ii in local_imaging_bands
-        ]
+        obs_list = []
+        for ii in local_imaging_bands:
+            slc = slice(*ift.utilities.shareRange(len(global_freqs), n_imaging_bands, ii))
+            obs_list.append(full_obs.get_freqs_by_slice(slc))
         nu0 = global_freqs.mean()
         return obs_list, nu0
 
-    def __getitem__(self, slc):
-        return Observation(
-            self._antpos[slc],
-            self._vis[:, slc],
-            self._weight[:, slc],
-            self._polarization,
-            self._freq,
-            self._direction,
-        )
-
-    def get_freqs(self, frequency_list):
+    def __getitem__(self, slc, copy=False):
+        ap = self._antpos[slc]
+        vis = self._vis[slc]
+        wgt = self._weight[:, slc]
+        if copy:
+            ap = ap.copy()
+            vis = vis.copy()
+            wgt = wgt.copy()
+        return Observation(ap, vis, wgt, self._polarization, self._freq, self._direction)
+
+    def get_freqs(self, frequency_list, copy=False):
         """Return observation that contains a subset of the present frequencies
 
         Parameters
@@ -400,17 +398,17 @@ class Observation(BaseObservation):
         """
         mask = np.zeros(self.nfreq, dtype=bool)
         mask[frequency_list] = 1
-        return self.get_freqs_by_slice(mask)
-
-    def get_freqs_by_slice(self, slc):
-        return Observation(
-            self._antpos,
-            self._vis[..., slc],
-            self._weight[..., slc],
-            self._polarization,
-            self._freq[slc],
-            self._direction,
-        )
+        return self.get_freqs_by_slice(mask, copy)
+
+    def get_freqs_by_slice(self, slc, copy=False):
+        vis = self._vis[..., slc]
+        wgt = self._weight[..., slc]
+        freq = self._freq[slc]
+        if copy:
+            vis = vis.copy()
+            wgt = wgt.copy()
+            freq = freq.copy()
+        return Observation( self._antpos, vis, wgt, self._polarization, freq, self._direction)
 
     def average_stokesi(self):
         my_asserteq(self._vis.shape[0], 2)

resolve/extra.py

+# SPDX-License-Identifier: GPL-3.0-or-later
+# Copyright(C) 2019-2021 Max-Planck-Society
+# Author: Philipp Arras
+
+import nifty8 as ift
+
+from .mpi import master
+from .data.observation import Observation
+
+
+def split_data_file(data_path, n_task, target_folder, base_name, n_work, compress):
+    from os import makedirs
+    makedirs(target_folder, exist_ok=True)
+
+    obs = Observation.load(data_path)
+
+    for rank in range(n_task):
+        lo, hi = ift.utilities.shareRange(n_work, n_task, rank)
+        sliced_obs = obs.get_freqs_by_slice(slice(*(lo, hi)))
+        sliced_obs.save(f"{target_folder}/{base_name}_{rank}.npz", compress=compress)
+
+
+def mpi_load(data_folder, base_name, full_data_set, n_work, comm=None, compress=False):
+    if master:
+        from os.path import isdir
+        if not isdir(data_folder):
+            split_data_file(full_data_set, comm.Get_size(), data_folder, base_name, n_work, compress)
+        if comm is None:
+            return Observation.load(full_data_set)
+
+    comm.Barrier()
+    return Observation.load(f"{data_folder}/{base_name}_{comm.Get_rank()}.npz")