style: apply black 23.1.0

draco/analysis/beam.py

@@ -248,13 +248,11 @@ def _evaluate_beam(self, data):
         # Loop over local frequencies and polarisations and evaluate the beam
         # by calling the telescopes beam method.
         for ff, freq in enumerate(local_freq_index):
-
             if not local_freq_flag[ff]:
                 weight[ff] = 0.0
                 continue
 
             for pp, pol in enumerate(pol_pairs):
-
                 bii = self.telescope.beam(map_pol_to_feed[pol[0]], freq, angpos)
 
                 if pol[0] != pol[1]:

draco/analysis/beamform.py

@@ -111,7 +111,6 @@ def setup(self, manager):
         # Ensure that if we are using variable time tracking,
         # then we are also collapsing over hour angle.
         if self.variable_timetrack:
-
             if self.collapse_ha:
                 self.log.info(
                     "Tracking source for declination dependent amount of time "
@@ -197,7 +196,6 @@ def process(self):
 
         # For each source, beamform and populate container.
         for src in range(self.nsource):
-
             if src % 1000 == 0:
                 self.log.info(f"Source {src}/{self.nsource}")
 
@@ -269,7 +267,6 @@ def process(self):
 
             # Loop over polarisations
             for pol, pol_str in enumerate(self.process_pol):
-
                 primary_beam = self._beamfunc(pol_str, dec, ha_array)
 
                 # Fringestop and sum over products
@@ -452,7 +449,6 @@ def _initialize_beam_with_data(self):
         # Find the index of the local frequencies in
         # the frequency axis of the telescope instance
         if not self.no_beam_model:
-
             self.freq_local_telescope_index = np.array(
                 [
                     np.argmin(np.abs(nu - self.telescope.frequencies))
@@ -492,7 +488,6 @@ def _beamfunc(self, pol, dec, ha):
         primary_beam = np.zeros((nfreq, ha.size), dtype=np.float64)
 
         for ff, freq in enumerate(self.freq_local_telescope_index):
-
             bii = self.telescope.beam(self.map_pol_feed[pol[0]], freq, angpos)
 
             if pol[0] != pol[1]:

draco/analysis/dayenu.py

@@ -93,7 +93,6 @@ def process(self, stream):
 
         # Loop over products
         for bb, bcut in enumerate(cutoff):
-
             t0 = time.time()
 
             # Flag frequencies and times with zero weight
@@ -144,7 +143,6 @@ def process(self, stream):
         return stream
 
     def _get_cut(self, prod):
-
         baselines = (
             self.telescope.feedpositions[prod["input_a"], :]
             - self.telescope.feedpositions[prod["input_b"], :]
@@ -204,7 +202,6 @@ def setup(self):
         """Create the function used to determine the delay cutoff."""
 
         if self.filename is not None:
-
             fcut = containers.DelayCutoff.from_file(self.filename, distributed=False)
             kind = fcut.attrs.get("kind", "linear")
 
@@ -215,7 +212,6 @@ def setup(self):
 
             self._cut_interpolator = {}
             for pp, pol in enumerate(fcut.pol):
-
                 self._cut_interpolator[pol] = scipy.interpolate.interp1d(
                     fcut.el,
                     fcut.cutoff[pp],
@@ -270,13 +266,11 @@ def process(self, ringmap):
 
         # Loop over beam and polarisation
         for ind in np.ndindex(*lshp):
-
             wind = ind[1:]
 
             kwargs = {ax: ringmap.index_map[ax][ii] for ax, ii in zip(axes, ind)}
 
             for ee, el in enumerate(els):
-
                 t0 = time.time()
 
                 slc = ind + (slice(None), slice(None), ee)
@@ -319,12 +313,10 @@ def process(self, ringmap):
 
                 # Apply the filter
                 if self.single_mask:
-
                     rm[slc] = np.matmul(NF[0], erm)
                     weight[wslc] = tools.invert_no_zero(np.matmul(NF[0] ** 2, evar))
 
                     if self.atten_threshold > 0.0:
-
                         diag = np.diag(NF[0])
                         med_diag = np.median(diag[diag > 0.0])
 
@@ -333,15 +325,13 @@ def process(self, ringmap):
                         weight[wslc] *= flag_low[:, np.newaxis].astype(np.float32)
 
                 else:
-
                     for ii, rr in enumerate(index):
                         rm[ind][:, rr, ee] = np.matmul(NF[ii], erm[:, rr])
                         weight[wind][:, rr, ee] = tools.invert_no_zero(
                             np.matmul(NF[ii] ** 2, evar[:, rr])
                         )
 
                         if self.atten_threshold > 0.0:
-
                             diag = np.diag(NF[ii])
                             med_diag = np.median(diag[diag > 0.0])
 
@@ -449,7 +439,6 @@ def process(self, stream):
 
         # Loop over frequencies
         for ff, nu in enumerate(freq):
-
             t0 = time.time()
 
             # The next several lines determine the mask as a function of time
@@ -488,7 +477,6 @@ def process(self, stream):
 
             # Loop over E-W baselines
             for uu, ub in enumerate(uniqb):
-
                 iub = np.flatnonzero(indexb == uu)
 
                 visfb = np.ascontiguousarray(vis[ff, iub])
@@ -512,7 +500,6 @@ def process(self, stream):
         return stream
 
     def _get_cut(self, freq, xsep):
-
         lmbda = units.c / (freq * 1e6)
         u = xsep / lmbda
         m = instantaneous_m(

draco/analysis/delay.py

@@ -92,7 +92,6 @@ def process(self, ss):
         baselines = tel.feedpositions[ia] - tel.feedpositions[ib]
 
         for lbi, bi in ss.vis[:].enumerate(axis=1):
-
             # Select the baseline length to use
             baseline = baselines[bi]
             if self.telescope_orientation == "NS":
@@ -277,7 +276,6 @@ def process(self, ss: FreqContainerType) -> FreqContainerType:
         )
 
         for lbi, bi in ss.datasets[dset][:].enumerate(axis=dist_axis_pos):
-
             # Extract the part of the array that we are processing, and
             # transpose/reshape to make a 2D array with frequency as axis=0
             vis_local = _take_view(ssv, lbi, dist_axis_pos)
@@ -448,7 +446,6 @@ def process(self, ss):
 
         # Iterate over all baselines and use the Gibbs sampler to estimate the spectrum
         for lbi, bi in delay_spec.spectrum[:].enumerate(axis=0):
-
             self.log.debug("Delay transforming baseline %i/%i", bi, len(baselines))
 
             # Get the local selections
@@ -680,7 +677,6 @@ def process(self, ss: FreqContainerType) -> containers.DelaySpectrum:
 
         # Iterate over all baselines and use the Gibbs sampler to estimate the spectrum
         for lbi, bi in delay_spec.spectrum[:].enumerate(axis=0):
-
             self.log.debug(f"Delay transforming baseline {bi}/{nbase}")
 
             # Get the local selections
@@ -772,7 +768,6 @@ def stokes_I(sstream, tel):
     # Cache beamclass as it's regenerated every call
     beamclass = tel.beamclass[:]
     for ii, ui in enumerate(uinv):
-
         # Skip if not all polarisations were included
         if ucount[ui] < 4:
             continue
@@ -1047,7 +1042,6 @@ def delay_spectrum_gibbs(
 
     # Window the frequency data
     if window is not None:
-
         # Construct the window function
         x = fsel * 1.0 / total_freq
         w = window_generalised(x, window=window)
@@ -1173,7 +1167,6 @@ def _draw_ps_sample(d):
     # Perform the Gibbs sampling iteration for a given number of loops and
     # return the power spectrum output of them.
     for ii in range(niter):
-
         d_samp = _draw_signal_sample(S_samp)
         S_samp = _draw_ps_sample(d_samp)
 

draco/analysis/fgfilter.py

@@ -82,7 +82,6 @@ def _forward(self, mmodes):
 
         # Iterate over local m's, project mode and save to disk.
         for lm, mi in mmodes.vis[:].enumerate(axis=0):
-
             tm = mmodes.vis[mi].transpose((1, 0, 2)).reshape(tel.nfreq, 2 * tel.npairs)
             svdm = bt.project_vector_telescope_to_svd(mi, tm)
 
@@ -128,7 +127,6 @@ def _backward(self, svdmodes):
 
         # Iterate over local m's, project mode and save to disk.
         for lm, mi in mmodes.vis[:].enumerate(axis=0):
-
             svdm = svdmodes.vis[mi]
             tm = bt.project_vector_svd_to_telescope(mi, svdm)
 
@@ -190,7 +188,6 @@ def _forward(self, svdmodes):
 
         # Iterate over local m's and project mode into KL basis
         for lm, mi in svdmodes.vis[:].enumerate(axis=0):
-
             sm = svdmodes.vis[mi][: svdmodes.nmode[mi]]
             klm = kl.project_vector_svd_to_kl(mi, sm, threshold=self.threshold)
 
@@ -227,7 +224,6 @@ def _backward(self, klmodes):
 
         # Iterate over local m's and project mode into KL basis
         for lm, mi in klmodes.vis[:].enumerate(axis=0):
-
             klm = klmodes.vis[mi][: klmodes.nmode[mi]]
             sm = kl.project_vector_kl_to_svd(mi, klm, threshold=self.threshold)
 

draco/analysis/flagging.py

@@ -486,7 +486,6 @@ def process(self, data):
         med_weight = np.zeros(npol, dtype=np.float32)
 
         for pp in range(npol):
-
             wlocal = data.weight[:, pp]
             wglobal = np.zeros(wlocal.global_shape, dtype=wlocal.dtype)
 
@@ -616,7 +615,6 @@ def process(self, data):
         weight_local = weight.local_array
 
         for lfi, gfi in weight.enumerate(axis=0):
-
             # MPIArray takes the local index, returns a local np.ndarray
             # Find values equal to zero to preserve them in final weights
             zeromask = weight_local[lfi] == 0.0
@@ -807,7 +805,6 @@ def process(self, sensitivity):
         stmask[:] = False
 
         for li, ii in madmask.enumerate(axis=0):
-
             # Only process this polarisation if we should be including it,
             # otherwise skip and let it be implicitly set to False (i.e. not
             # masked)
@@ -1033,7 +1030,6 @@ def process(
 
         # Get the rank with stack to create the new mask
         if sstream.comm.rank == rank_with_ind:
-
             # Cut out the right section
             wf = ssv.local_array[:, self.stack_ind - lstart]
             ww = ssw.local_array[:, self.stack_ind - lstart]
@@ -1271,7 +1267,6 @@ def _bad_freq_mask(self, nfreq: int) -> np.ndarray:
         mask = np.zeros(nfreq, dtype=bool)
 
         for s in self.bad_freq_ind:
-
             if isinstance(s, int):
                 if s < nfreq:
                     mask[s] = True
@@ -1470,7 +1465,6 @@ def tv_channels_flag(x, freq, sigma=5, f=0.5, debug=False):
     freq_end = freq + 0.5 * df
 
     for i in range(67):
-
         # Find all frequencies that lie wholly or partially within the TV channel
         fs = tvstart_freq + i * tvwidth_freq
         fe = fs + tvwidth_freq
@@ -1635,7 +1629,6 @@ def process(self, data):
 
         # Find the median offset between the stack and the daily data
         if self.match_median:
-
             # Find the parts of the both the stack and the daily data that are both
             # measured
             mask = (
@@ -1689,7 +1682,6 @@ def process(self, data):
             weight *= weight_stack
 
         else:
-
             # Perform a weighted average of the data to fill in missing samples
             dset *= weight
             dset += weight_stack * self.frac * (dset_stack + stack_offset)
@@ -1811,7 +1803,6 @@ def process(
         # If only flagging co-pol baselines, make separate mask to select those,
         # and multiply into low-weight mask
         if self.pols_to_flag == "copol":
-
             # Get local section of stack axis
             local_stack = stream.stack[stream.weight[:].local_bounds]
 

draco/analysis/mapmaker.py

@@ -77,7 +77,6 @@ def process(self, mmodes):
 
         # Loop over all m's and solve from m-mode visibilities to alms.
         for mi, m in m_array.enumerate(axis=0):
-
             self.log.debug(
                 "Processing m=%i (local %i/%i)", m, mi + 1, m_array.local_shape[0]
             )
@@ -156,7 +155,6 @@ class DirtyMapMaker(BaseMapMaker):
     """
 
     def _solve_m(self, m, f, v, Ni):
-
         bt = self.beamtransfer
 
         # Massage the arrays into shape
@@ -188,7 +186,6 @@ class MaximumLikelihoodMapMaker(BaseMapMaker):
     """
 
     def _solve_m(self, m, f, v, Ni):
-
         bt = self.beamtransfer
 
         # Massage the arrays into shape
@@ -241,7 +238,6 @@ class WienerMapMaker(BaseMapMaker):
     bt_cache = None
 
     def _solve_m(self, m, f, v, Ni):
-
         import scipy.linalg as la
 
         bt = self.beamtransfer