diff --git a/hydra/__init__.py b/hydra/__init__.py
index adce64b..bc3c178 100644
--- a/hydra/__init__.py
+++ b/hydra/__init__.py
@@ -1,5 +1,13 @@
-
-from . import beam_sampler, cosmo_sampler, diffuse_sampler, gain_sampler, \
-              pspec_sampler, ptsrc_sampler, region_sampler, sh_sampler, vis_sampler
+from . import (
+    beam_sampler,
+    cosmo_sampler,
+    diffuse_sampler,
+    gain_sampler,
+    pspec_sampler,
+    ptsrc_sampler,
+    region_sampler,
+    sh_sampler,
+    vis_sampler,
+)
 from . import config, example, linear_solver, plot, sparse_beam, utils, vis_simulator
 from .utils import *
diff --git a/hydra/beam_sampler.py b/hydra/beam_sampler.py
index d365e9d..8ef37aa 100644
--- a/hydra/beam_sampler.py
+++ b/hydra/beam_sampler.py
@@ -29,12 +29,14 @@ def split_real_imag(arr, kind):
         new_arr (array_like):
             Array that has been split into real and imaginary parts.
     """
-    valid_kinds = ['op', 'vec']
+    valid_kinds = ["op", "vec"]
     nax = len(arr.shape)
     if kind not in valid_kinds:
-        raise ValueError("kind must be 'op' or 'vec' when splitting complex "
-                         "arrays into real and imaginary components")
-    if kind == 'op':
+        raise ValueError(
+            "kind must be 'op' or 'vec' when splitting complex "
+            "arrays into real and imaginary components"
+        )
+    if kind == "op":
         new_arr = np.zeros((2, 2) + arr.shape, dtype=float)
         new_arr[0, 0] = arr.real
         new_arr[0, 1] = -arr.imag
@@ -44,12 +46,14 @@ def split_real_imag(arr, kind):
         # Prepare to put these axes at the end
         axes = list(range(2, nax + 2)) + [0, 1]
 
-    elif kind == 'vec':
+    elif kind == "vec":
         new_arr = np.zeros((2,) + arr.shape, dtype=float)
         new_arr[0], new_arr[1] = (arr.real, arr.imag)
 
         # Prepare to put this axis at the end
-        axes = list(range(1, nax + 1)) + [0, ]
+        axes = list(range(1, nax + 1)) + [
+            0,
+        ]
 
     # Rearrange axes (they are always expected at the end)
     new_arr = np.transpose(new_arr, axes=axes)
@@ -80,13 +84,22 @@ def reshape_data_arr(arr, Nfreqs, Ntimes, Nants, Npol):
     """
 
     arr_trans = np.transpose(arr, (0, 1, 3, 4, 2))
-    arr_beam = np.zeros([Npol, Npol, Nfreqs, Ntimes, Nants, Nants], dtype=arr_trans.dtype)
+    arr_beam = np.zeros(
+        [Npol, Npol, Nfreqs, Ntimes, Nants, Nants], dtype=arr_trans.dtype
+    )
     for pol_ind1 in range(Npol):
         for pol_ind2 in range(Npol):
             for freq_ind in range(Nfreqs):
                 for time_ind in range(Ntimes):
-                        triu_inds = np.triu_indices(Nants, k=1)
-                        arr_beam[pol_ind1, pol_ind2, freq_ind, time_ind, triu_inds[0], triu_inds[1]] = arr_trans[pol_ind1, pol_ind2, freq_ind, time_ind]
+                    triu_inds = np.triu_indices(Nants, k=1)
+                    arr_beam[
+                        pol_ind1,
+                        pol_ind2,
+                        freq_ind,
+                        time_ind,
+                        triu_inds[0],
+                        triu_inds[1],
+                    ] = arr_trans[pol_ind1, pol_ind2, freq_ind, time_ind]
 
     return arr_beam
 
@@ -112,8 +125,6 @@ def get_bess_matr(nmodes, mmodes, rho, phi):
             Fourier-Bessel transformation matrix.
     """
 
-
-
     unique_n, ninv = np.unique(nmodes, return_inverse=True)
     nmax = np.amax(unique_n)
     bzeros = jn_zeros(0, nmax)
@@ -125,18 +136,20 @@ def get_bess_matr(nmodes, mmodes, rho, phi):
 
     unique_m, minv = np.unique(mmodes, return_inverse=True)
     # Shape Ntimes, Nptsrc, len(unique_m)
-    az_modes = np.exp(1.j * unique_m[np.newaxis, np.newaxis, :] * phi[:, :, np.newaxis])
+    az_modes = np.exp(
+        1.0j * unique_m[np.newaxis, np.newaxis, :] * phi[:, :, np.newaxis]
+    )
     # Shape Ntimes, Nptsrc, len(mmodes)
-    az_vals = az_modes[:, :, minv] / np.sqrt(np.pi) # making orthonormal
+    az_vals = az_modes[:, :, minv] / np.sqrt(np.pi)  # making orthonormal
 
     bess_matr = bess_vals * az_vals
 
-
     return bess_matr
 
 
-def fit_bess_to_beam(beam, freqs, nmodes, mmodes, rho, phi, polarized=False,
-                     force_spw_index=False):
+def fit_bess_to_beam(
+    beam, freqs, nmodes, mmodes, rho, phi, polarized=False, force_spw_index=False
+):
     """
     Get the best fit Fourier-Bessel coefficients for a beam based on its value at a
     a set direction cosines. A least-squares algorithm is used to perform the
@@ -172,25 +185,30 @@ def fit_bess_to_beam(beam, freqs, nmodes, mmodes, rho, phi, polarized=False,
     rho_un = np.unique(rho)
     phi_un = np.unique(phi)
 
-
     drho = rho_un[1] - rho_un[0]
     dphi = phi_un[1] - phi_un[0]
 
     # Before indexing conventions enforced
-    rhs_full = beam.interp(az_array=phi.flatten(),
-                           za_array=np.arccos(1 - rho**2).flatten(),
-                           freq_array=freqs)[0]
-    
+    rhs_full = beam.interp(
+        az_array=phi.flatten(),
+        za_array=np.arccos(1 - rho**2).flatten(),
+        freq_array=freqs,
+    )[0]
+
     if polarized:
         if spw_axis_present or force_spw_index:
-            rhs = rhs_full[:, 0] # Will have shape Nfeed, Naxes_vec, Nfreq, Nrho * Nphi
+            rhs = rhs_full[:, 0]  # Will have shape Nfeed, Naxes_vec, Nfreq, Nrho * Nphi
         else:
             rhs = rhs_full
     else:
         if spw_axis_present or force_spw_index:
-            rhs = rhs_full[1:, 0, :1] # FIXME: analyticbeam gives nans and zeros for all other indices
+            rhs = rhs_full[
+                1:, 0, :1
+            ]  # FIXME: analyticbeam gives nans and zeros for all other indices
         else:
-            rhs = rhs_full[1:, :1] # FIXME: analyticbeam gives nans and zeros for all other indices
+            rhs = rhs_full[
+                1:, :1
+            ]  # FIXME: analyticbeam gives nans and zeros for all other indices
 
     Npol = 1 + polarized
 
@@ -198,18 +216,20 @@ def fit_bess_to_beam(beam, freqs, nmodes, mmodes, rho, phi, polarized=False,
     Nfreqs = len(freqs)
     fit_beam = np.zeros((Nfreqs, ncoeff, Npol, Npol), dtype=complex)
 
-    BT = (bess_matr.conj())
+    BT = bess_matr.conj()
     lhs_op = np.tensordot(BT, bess_matr, axes=((0, 1), (0, 1)))
     BT_res = BT.reshape(rho.size, ncoeff)
     for freq_ind in range(Nfreqs):
         for feed_ind in range(Npol):
             for pol_ind in range(Npol):
-                rhs_vec = (BT_res * rhs[feed_ind, pol_ind, freq_ind, :, np.newaxis]).sum(axis=0)
-                soln = solve(lhs_op, rhs_vec, assume_a='her')
+                rhs_vec = (
+                    BT_res * rhs[feed_ind, pol_ind, freq_ind, :, np.newaxis]
+                ).sum(axis=0)
+                soln = solve(lhs_op, rhs_vec, assume_a="her")
                 fit_beam[freq_ind, :, feed_ind, pol_ind] = soln
 
     # Reshape coefficients array
-    fit_beam = np.array(fit_beam) # Has shape Nfreqs, ncoeffs, Npol, Npol
+    fit_beam = np.array(fit_beam)  # Has shape Nfreqs, ncoeffs, Npol, Npol
 
     return fit_beam
 
@@ -258,37 +278,53 @@ def get_bess_outer(bess_matr):
 
     return bess_matr[:, :, np.newaxis] * bess_matr.conj()[:, :, :, np.newaxis]
 
-def get_bess_sky_contraction(bess_outer, ants, fluxes, ra, dec, freqs, lsts,
-                             polarized=False, 
-                             precision=1, latitude=-30.7215 * np.pi / 180.0, 
-                             use_feed="x", multiprocess=True):
-    
+
+def get_bess_sky_contraction(
+    bess_outer,
+    ants,
+    fluxes,
+    ra,
+    dec,
+    freqs,
+    lsts,
+    polarized=False,
+    precision=1,
+    latitude=-30.7215 * np.pi / 180.0,
+    use_feed="x",
+    multiprocess=True,
+):
+
     Npol = 2 if polarized else 1
     Nfreqs = len(freqs)
     Ncoeff = bess_outer.shape[-1]
     Ntimes = len(lsts)
     Nants = len(ants)
     contract_shape = [Npol, Npol, Nfreqs, Ntimes, Nants, Nants, Ncoeff, Ncoeff]
-    
+
     # tsb,qQftaAs,tsB -> qQftaAbB
     # inner loop is already over frequency, so just loop over that to save mem
     bess_sky_contraction = np.zeros(contract_shape, dtype=complex)
-    beams = [AnalyticBeam("uniform") for ant_ind in range(len(ants))] 
-    
-    #for freq_ind, freq in enumerate(freqs):
-        
-        #FIXME: can do away with freq loop if we use sparse_beam here but it needs to be rewritten
+    beams = [AnalyticBeam("uniform") for ant_ind in range(len(ants))]
+
+    # for freq_ind, freq in enumerate(freqs):
+
+    # FIXME: can do away with freq loop if we use sparse_beam here but it needs to be rewritten
     for time_ind in range(Ntimes):
-        sky_amp_phase = simulate_vis_per_source(ants, fluxes,
-                                                ra, dec, freqs, lsts[time_ind:time_ind + 1],
-                                                beams=beams, polarized=polarized,
-                                                precision=precision,
-                                                latitude=latitude,
-                                                use_feed=use_feed,
-                                                multiprocess=multiprocess)
-
-            
-            
+        sky_amp_phase = simulate_vis_per_source(
+            ants,
+            fluxes,
+            ra,
+            dec,
+            freqs,
+            lsts[time_ind : time_ind + 1],
+            beams=beams,
+            polarized=polarized,
+            precision=precision,
+            latitude=latitude,
+            use_feed=use_feed,
+            multiprocess=multiprocess,
+        )
+
         if not polarized:
             sky_amp_phase = sky_amp_phase[np.newaxis, np.newaxis, :]
         # Need this conjugation since only lower half of array is filled
@@ -296,30 +332,45 @@ def get_bess_sky_contraction(bess_outer, ants, fluxes, ra, dec, freqs, lsts,
         # makes compute later easier to think about
         sky_amp_phase = sky_amp_phase + sky_amp_phase.swapaxes(4, 5).conj()
 
-        bess_sky_contraction[:, :, :, time_ind] = np.tensordot(sky_amp_phase[:, :, :, 0],
-                                                               bess_outer[time_ind],
-                                                               axes=((-1, ), (0, )))
-
+        bess_sky_contraction[:, :, :, time_ind] = np.tensordot(
+            sky_amp_phase[:, :, :, 0], bess_outer[time_ind], axes=((-1,), (0,))
+        )
 
     return bess_sky_contraction
 
-def get_bess_to_vis_from_contraction(bess_sky_contraction, beam_coeffs, ants,
-                                     ant_samp_ind):
+
+def get_bess_to_vis_from_contraction(
+    bess_sky_contraction, beam_coeffs, ants, ant_samp_ind
+):
     Nants = len(ants)
     ant_inds = get_ant_inds(ant_samp_ind, Nants)
-    beam_res = (beam_coeffs.transpose((2, 3, 1, 0, 4)))[ant_inds] # bfApQ -> ApfbQ
-    bess_trans = np.einsum("ApfbQ,qQftAbB->pqftAB", 
-                           beam_res.conj(), 
-                           bess_sky_contraction[:, :, :, :,  ant_inds, ant_samp_ind],
-                           optimize=True)
-    
+    beam_res = (beam_coeffs.transpose((2, 3, 1, 0, 4)))[ant_inds]  # bfApQ -> ApfbQ
+    bess_trans = np.einsum(
+        "ApfbQ,qQftAbB->pqftAB",
+        beam_res.conj(),
+        bess_sky_contraction[:, :, :, :, ant_inds, ant_samp_ind],
+        optimize=True,
+    )
+
     return bess_trans
 
 
-def get_bess_to_vis(bess_matr, ants, fluxes, ra, dec, freqs, lsts,
-                    beam_coeffs, ant_samp_ind, polarized=False, precision=1,
-                    latitude=-30.7215 * np.pi / 180.0, use_feed="x",
-                    multiprocess=True):
+def get_bess_to_vis(
+    bess_matr,
+    ants,
+    fluxes,
+    ra,
+    dec,
+    freqs,
+    lsts,
+    beam_coeffs,
+    ant_samp_ind,
+    polarized=False,
+    precision=1,
+    latitude=-30.7215 * np.pi / 180.0,
+    use_feed="x",
+    multiprocess=True,
+):
     """
     Compute the matrices that act as the quadratic forms by which visibilities
     are made. Calls simulate_vis_per_source to get the Fourier operator, then
@@ -376,13 +427,21 @@ def get_bess_to_vis(bess_matr, ants, fluxes, ra, dec, freqs, lsts,
 
     # Use uniform beams so that we just get the Fourier operator.
     beams = [AnalyticBeam("uniform") for ant_ind in range(len(ants))]
-    sky_amp_phase = simulate_vis_per_source(ants, fluxes, ra, dec, freqs, lsts,
-                                            beams=beams, polarized=polarized,
-                                            precision=precision,
-                                            latitude=latitude,
-                                            use_feed=use_feed,
-                                            multiprocess=multiprocess,
-                                            subarr_ant=ant_samp_ind)
+    sky_amp_phase = simulate_vis_per_source(
+        ants,
+        fluxes,
+        ra,
+        dec,
+        freqs,
+        lsts,
+        beams=beams,
+        polarized=polarized,
+        precision=precision,
+        latitude=latitude,
+        use_feed=use_feed,
+        multiprocess=multiprocess,
+        subarr_ant=ant_samp_ind,
+    )
     sky_amp_phase = sky_amp_phase[:, :, ant_inds]
     if not polarized:
         sky_amp_phase = sky_amp_phase[np.newaxis, np.newaxis, :]
@@ -391,23 +450,28 @@ def get_bess_to_vis(bess_matr, ants, fluxes, ra, dec, freqs, lsts,
     # Determined optimal contraction order with opt_einsum
     # Implementing steps in faster way using tdot
     # aqfBQ,tsB->aqfQts->Qqftas
-    beam_res = (beam_coeffs.transpose((2, 3, 1, 0, 4)))[ant_inds] # BfaqQ -> aqfBQ
-    beam_on_sky = np.tensordot(beam_res.conj(), bess_matr.conj(),
-                               axes=((3,), (2,))).transpose((3,1,2,4,0,5))
+    beam_res = (beam_coeffs.transpose((2, 3, 1, 0, 4)))[ant_inds]  # BfaqQ -> aqfBQ
+    beam_on_sky = np.tensordot(
+        beam_res.conj(), bess_matr.conj(), axes=((3,), (2,))
+    ).transpose((3, 1, 2, 4, 0, 5))
 
     # Qqftas,QPftas->qPftas
     # reassign to save memory
     # Qqftas -> Qq_ftas; QPftas->Q_Pftas; Qq_ftas,Q_Pftas->qPftas
-    sky_amp_phase = (beam_on_sky[:, :, np.newaxis] * sky_amp_phase[:, np.newaxis]).sum(axis=0)
+    sky_amp_phase = (beam_on_sky[:, :, np.newaxis] * sky_amp_phase[:, np.newaxis]).sum(
+        axis=0
+    )
 
     # qPftas,tsb->qPftab
-    bess_trans = np.zeros((Npol, Npol, freqs.size, lsts.size, nants - 1, beam_res.shape[-2]),
-                          dtype=sky_amp_phase.dtype)
+    bess_trans = np.zeros(
+        (Npol, Npol, freqs.size, lsts.size, nants - 1, beam_res.shape[-2]),
+        dtype=sky_amp_phase.dtype,
+    )
 
     for time_ind in range(lsts.size):
-        bess_trans[:,:, :, time_ind, :, :] = np.tensordot(sky_amp_phase[:, :, :, time_ind],
-                                                     bess_matr[time_ind],
-                                                     axes=((-1, ), (0, )))
+        bess_trans[:, :, :, time_ind, :, :] = np.tensordot(
+            sky_amp_phase[:, :, :, time_ind], bess_matr[time_ind], axes=((-1,), (0,))
+        )
     return bess_trans
 
 
@@ -440,26 +504,32 @@ def get_cov_Qdag_Ninv_Q(inv_noise_var, bess_trans, cov_tuple):
     # Should not affect the vis at other times/freqs
 
     # qpfta,qPftab->qpPftab
-    Ninv_Q = inv_noise_var[:, :, np.newaxis, :, :, :, np.newaxis] * bess_trans[:, np.newaxis, :, :, :, :, :]
+    Ninv_Q = (
+        inv_noise_var[:, :, np.newaxis, :, :, :, np.newaxis]
+        * bess_trans[:, np.newaxis, :, :, :, :, :]
+    )
 
     # qRftab,qpPFtaB->RfbpPFB
     # Actually just want diagonals in frequency but don't need to save memory here
-    Qdag_Ninv_Q = np.tensordot(bess_trans.conj(), Ninv_Q,
-                              axes=((0, 3, 4), (0, 4, 5)))
+    Qdag_Ninv_Q = np.tensordot(bess_trans.conj(), Ninv_Q, axes=((0, 3, 4), (0, 4, 5)))
     # Get the diagonals RfbpPFB-> fRbpPB
     Qdag_Ninv_Q = Qdag_Ninv_Q[:, range(Nfreqs), :, :, :, range(Nfreqs)]
     # Factor of 2 because 1/2 the variance for each complex component
-    Qdag_Ninv_Q = 2*split_real_imag(Qdag_Ninv_Q, kind='op') # fRbpPBcC
-    Qdag_Ninv_Q = Qdag_Ninv_Q.transpose((1,2,3,4,5,7,6,0)) # fRbpPBcC->RbpPBCcf
+    Qdag_Ninv_Q = 2 * split_real_imag(Qdag_Ninv_Q, kind="op")  # fRbpPBcC
+    Qdag_Ninv_Q = Qdag_Ninv_Q.transpose((1, 2, 3, 4, 5, 7, 6, 0))  # fRbpPBcC->RbpPBCcf
 
     # c,fF->cfF->fcF
     cov_matr = np.swapaxes(comp_matr[:, np.newaxis, np.newaxis] * freq_matr, 0, 1)
     # fcF,RbpPBCcf->fRbpPBCcF
-    cov_Qdag_Ninv_Q = cov_matr[:, np.newaxis, np.newaxis, np.newaxis, np.newaxis, np.newaxis, np.newaxis] * Qdag_Ninv_Q[np.newaxis]
-    #fRbpPBCcF->fRbpPBcCF
+    cov_Qdag_Ninv_Q = (
+        cov_matr[
+            :, np.newaxis, np.newaxis, np.newaxis, np.newaxis, np.newaxis, np.newaxis
+        ]
+        * Qdag_Ninv_Q[np.newaxis]
+    )
+    # fRbpPBCcF->fRbpPBcCF
     cov_Qdag_Ninv_Q = np.swapaxes(cov_Qdag_Ninv_Q, -2, -3)
 
-
     return cov_Qdag_Ninv_Q
 
 
@@ -482,13 +552,14 @@ def apply_operator(x, cov_Qdag_Ninv_Q):
     Npol = x.shape[2]
     # Linear so we can split the LHS multiply
     # fRbpPBcCF,BFpPC->fRbpcp->pRfbc->pfbRc
-    Ax1 = np.tensordot(cov_Qdag_Ninv_Q, x, axes=((4, 5, 7,8), (3, 0, 4, 1)))
-    Ax1 = (Ax1[:, :, :, range(Npol), :, range(Npol)]).transpose((0,2,3,1,4))
+    Ax1 = np.tensordot(cov_Qdag_Ninv_Q, x, axes=((4, 5, 7, 8), (3, 0, 4, 1)))
+    Ax1 = (Ax1[:, :, :, range(Npol), :, range(Npol)]).transpose((0, 2, 3, 1, 4))
 
     # Second term is identity due to preconditioning
     Ax = (Ax1 + x).transpose((1, 2, 0, 3, 4))
     return Ax
 
+
 def get_std_norm(shape):
     """
     Get an array of complex standard normal samples.
@@ -499,12 +570,14 @@ def get_std_norm(shape):
     Returns:
         std_norm (array_like): desired complex samples.
     """
-    std_norm = (np.random.normal(size=shape) + 1.0j * np.random.normal(size=shape)) / np.sqrt(2)
+    std_norm = (
+        np.random.normal(size=shape) + 1.0j * np.random.normal(size=shape)
+    ) / np.sqrt(2)
 
     return std_norm
 
-def construct_rhs(vis, inv_noise_var, mu, bess_trans,
-                  cov_tuple, cho_tuple, flx=True):
+
+def construct_rhs(vis, inv_noise_var, mu, bess_trans, cov_tuple, cho_tuple, flx=True):
     """
     Construct the right hand side of the Gaussian Constrained Realization (GCR)
     equation.
@@ -543,32 +616,35 @@ def construct_rhs(vis, inv_noise_var, mu, bess_trans,
         flx0 = np.zeros(flx0_shape)
         flx1 = np.zeros(flx1_shape)
 
-
-
     Ninv_d = inv_noise_var * vis
     Ninv_sqrt_flx1 = np.sqrt(inv_noise_var) * flx1
 
     # qPftab,qpfta->bfpP
     # qPftab->bPqfta
-    bess_trans_use = bess_trans.transpose((5,1,0,2,3,4))
+    bess_trans_use = bess_trans.transpose((5, 1, 0, 2, 3, 4))
     # Weird factors of sqrt(2) etc since we will split these in a sec
     # bPqfta,pqfta->bpPf->bfpP
 
-    Qdag_terms = np.sum(bess_trans_use.conj()[:,np.newaxis] * (2 * Ninv_d + np.sqrt(2) * Ninv_sqrt_flx1).transpose((1, 0, 2, 3, 4))[:, :, np.newaxis],
-                        axis=(3,5,6)).transpose((0,3,1,2))
-    Qdag_terms = split_real_imag(Qdag_terms, kind='vec')
+    Qdag_terms = np.sum(
+        bess_trans_use.conj()[:, np.newaxis]
+        * (2 * Ninv_d + np.sqrt(2) * Ninv_sqrt_flx1).transpose((1, 0, 2, 3, 4))[
+            :, :, np.newaxis
+        ],
+        axis=(3, 5, 6),
+    ).transpose((0, 3, 1, 2))
+    Qdag_terms = split_real_imag(Qdag_terms, kind="vec")
 
     freq_matr, comp_matr = cov_tuple
     # c,bfpPc->bfpPc Ff,bfpPc->FbpPc
-    cov_Qdag_terms = np.tensordot(freq_matr, (comp_matr * Qdag_terms),
-                                  axes=((1,), (1,)))
+    cov_Qdag_terms = np.tensordot(
+        freq_matr, (comp_matr * Qdag_terms), axes=((1,), (1,))
+    )
 
     freq_cho, comp_cho = cho_tuple
-    flx0 = split_real_imag(flx0, kind='vec')
-    mu_real_imag = split_real_imag(mu, kind='vec')
+    flx0 = split_real_imag(flx0, kind="vec")
+    mu_real_imag = split_real_imag(mu, kind="vec")
 
-    flx0_add = np.tensordot(freq_cho, (comp_cho * flx0),
-                                  axes=((1,), (1,)))
+    flx0_add = np.tensordot(freq_cho, (comp_cho * flx0), axes=((1,), (1,)))
 
     rhs = cov_Qdag_terms + np.swapaxes(mu_real_imag, 0, 1) + flx0_add
 
@@ -593,12 +669,14 @@ def non_norm_gauss(A, sig, x):
             Values of the Gaussian function at positions given by x
     """
 
-    gvals = A * np.exp(-x**2 / (2 * sig**2))
+    gvals = A * np.exp(-(x**2) / (2 * sig**2))
 
     return gvals
 
-def make_prior_cov(freqs, times, ncoeff, std, sig_freq,
-                   constrain_phase=False, constraint=1e-4, ridge=0):
+
+def make_prior_cov(
+    freqs, times, ncoeff, std, sig_freq, constrain_phase=False, constraint=1e-4, ridge=0
+):
     """
     Make a prior covariance for the beam coefficients.
 
@@ -625,14 +703,14 @@ def make_prior_cov(freqs, times, ncoeff, std, sig_freq,
     freq_col[0] += ridge
     freq_matr = toeplitz(freq_col)
     comp_matr = np.ones(2)
-    if constrain_phase: # Make the imaginary variance small compared to the real one
+    if constrain_phase:  # Make the imaginary variance small compared to the real one
         comp_matr[1] = constraint
 
-
     cov_tuple = (freq_matr, comp_matr)
 
     return cov_tuple
 
+
 def do_cov_cho(cov_tuple, check_op=False):
     """
     Returns the Cholesky decomposition of a matrix factorable over
@@ -650,14 +728,16 @@ def do_cov_cho(cov_tuple, check_op=False):
     """
     freq_matr, comp_matr = cov_tuple
     freq_cho = cholesky(freq_matr, lower=True)
-    comp_cho = np.sqrt(comp_matr) # Currently always diagonal
+    comp_cho = np.sqrt(comp_matr)  # Currently always diagonal
 
     cho_tuple = (freq_cho, comp_cho)
     if check_op:
         prod = freq_cho @ freq_cho.T.conj()
         allclose = np.allclose(prod, freq_matr)
-        print(f"Successful cholesky factorization of beam for frequency covariance: "
-              f"{allclose}")
+        print(
+            f"Successful cholesky factorization of beam for frequency covariance: "
+            f"{allclose}"
+        )
         if not allclose:
             raise LinAlgError(f"Cholesky factorization failed for frequency covariance")
 
@@ -684,18 +764,27 @@ def get_beam_from_FB_coeff(beam_coeffs, za, az, nmodes, mmodes):
         beam (array_like): Beam evaluated at a grid of za, az
     """
 
-
     rho = np.sqrt(1 - np.cos(za))
     Rho, Az = np.meshgrid(rho, az)
     B = get_bess_matr(nmodes, mmodes, Rho, Az)
 
-    beam = B@beam_coeffs
+    beam = B @ beam_coeffs
 
     return beam
 
 
-def plot_FB_beam(beam, za, az, vmin=-1, vmax=1, norm=SymLogNorm, linthresh=1e-3, 
-                 cmap="Spectral", save=False, **kwargs):
+def plot_FB_beam(
+    beam,
+    za,
+    az,
+    vmin=-1,
+    vmax=1,
+    norm=SymLogNorm,
+    linthresh=1e-3,
+    cmap="Spectral",
+    save=False,
+    **kwargs,
+):
     """
     Plots a Fourier_Bessel beam at specified zenith angles and azimuths.
 
@@ -716,12 +805,22 @@ def plot_FB_beam(beam, za, az, vmin=-1, vmax=1, norm=SymLogNorm, linthresh=1e-3,
 
     Az, Za = np.meshgrid(az, za)
 
-    fig, ax = plt.subplots(ncols=2, subplot_kw={'projection': 'polar'}, figsize=(16, 8))
-    cax = ax[0].pcolormesh(Az, Za, beam.real,
-                           norm=norm(vmin=vmin, vmax=vmax, linthresh=linthresh, **kwargs), cmap=cmap)
+    fig, ax = plt.subplots(ncols=2, subplot_kw={"projection": "polar"}, figsize=(16, 8))
+    cax = ax[0].pcolormesh(
+        Az,
+        Za,
+        beam.real,
+        norm=norm(vmin=vmin, vmax=vmax, linthresh=linthresh, **kwargs),
+        cmap=cmap,
+    )
     ax[0].set_title("Real Component")
-    ax[1].pcolormesh(Az, Za, beam.imag,
-                     norm=norm(vmin=vmin, vmax=vmax, linthresh=linthresh, **kwargs), cmap=cmap)
+    ax[1].pcolormesh(
+        Az,
+        Za,
+        beam.imag,
+        norm=norm(vmin=vmin, vmax=vmax, linthresh=linthresh, **kwargs),
+        cmap=cmap,
+    )
     ax[1].set_title("Imaginary Component")
     fig.colorbar(cax, ax=ax.ravel().tolist())
     if save:
@@ -744,17 +843,19 @@ def get_zernike_rad(r, n, m):
     Returns:
         rad: radial polynomial of degree (n,m) evaluated at (theta,r)
     """
-    if (n - m) % 2: # odd difference, return 0
-        raise ValueError("Difference between n and m must be even. "
-                         f"n={n} and m={m}.")
+    if (n - m) % 2:  # odd difference, return 0
+        raise ValueError(
+            "Difference between n and m must be even. " f"n={n} and m={m}."
+        )
     else:
         nm_diff = (n - m) // 2
         nm_sum = (n + m) // 2
-        prefac = (-1)**nm_diff * comb(nm_sum, m) * r**m
+        prefac = (-1) ** nm_diff * comb(nm_sum, m) * r**m
         rad = prefac * hyp2f1(1 + nm_sum, -nm_diff, 1 + m, r**2)
 
     return rad
 
+
 def get_zernike_azim(theta, m):
     """
     Get the azimuthal part of a Zernike function
@@ -770,7 +871,7 @@ def get_zernike_azim(theta, m):
         azim = np.cos(m * theta)
     elif m < 0:
         azim = np.sin(np.abs(m) * theta)
-    return(azim)
+    return azim
 
 
 def get_zernike_matrix(nmax, theta, r):
@@ -794,9 +895,22 @@ def get_zernike_matrix(nmax, theta, r):
 
     return zern_matr.transpose((1, 2, 0))
 
-def get_pert_beam(seed, beam_file, trans_std=1e-2, rot_std_deg=1., 
-                  stretch_std=1e-2, mmax=45, nmax=80, sqrt=True, Nfeeds=2,
-                  num_modes_comp=32, save=False, outdir="", load=False):
+
+def get_pert_beam(
+    seed,
+    beam_file,
+    trans_std=1e-2,
+    rot_std_deg=1.0,
+    stretch_std=1e-2,
+    mmax=45,
+    nmax=80,
+    sqrt=True,
+    Nfeeds=2,
+    num_modes_comp=32,
+    save=False,
+    outdir="",
+    load=False,
+):
     """
     Get a perturbed sparse_beam instance.
 
@@ -816,7 +930,7 @@ def get_pert_beam(seed, beam_file, trans_std=1e-2, rot_std_deg=1.,
         nmax (int):
             The maximum radial mode number to use in the FB basis.
         sqrt (bool):
-            Whether to take the square root of the unperturbed beam before 
+            Whether to take the square root of the unperturbed beam before
             fitting. Used for power beams.
         Nfeeds (int):
             Number of feeds. Set to None if using E-field beam, 2 for power beam.
@@ -838,13 +952,22 @@ def get_pert_beam(seed, beam_file, trans_std=1e-2, rot_std_deg=1.,
     sin_pert_coeffs = np.random.normal(size=8)
 
     mmodes = np.arange(-mmax, mmax + 1)
-    sb = sparse_beam(beam_file, nmax, mmodes, Nfeeds=Nfeeds,
-                     num_modes_comp=num_modes_comp, sqrt=sqrt, 
-                     perturb=True, trans_x=trans_x, trans_y=trans_y, 
-                     rot=rot, stretch_x=stretch_x, 
-                     stretch_y=stretch_y, 
-                     sin_pert_coeffs=sin_pert_coeffs)
-    
+    sb = sparse_beam(
+        beam_file,
+        nmax,
+        mmodes,
+        Nfeeds=Nfeeds,
+        num_modes_comp=num_modes_comp,
+        sqrt=sqrt,
+        perturb=True,
+        trans_x=trans_x,
+        trans_y=trans_y,
+        rot=rot,
+        stretch_x=stretch_x,
+        stretch_y=stretch_y,
+        sin_pert_coeffs=sin_pert_coeffs,
+    )
+
     beam_outfile = f"{outdir}/perturbed_beam_beamvals_seed_{seed}.npy"
     if load:
         pert_beam = np.load(beam_outfile)
@@ -855,19 +978,16 @@ def get_pert_beam(seed, beam_file, trans_std=1e-2, rot_std_deg=1.,
 
     if save:
         np.save(beam_outfile, pert_beam)
-        np.save(f"{outdir}/perturbed_beam_fit_coeffs_seed_{seed}.npy", 
-                fit_coeffs)
-    
+        np.save(f"{outdir}/perturbed_beam_fit_coeffs_seed_{seed}.npy", fit_coeffs)
+
     return sb
 
 
 def init_beam_sampler(beams, freqs, beam_mmax, beam_nmax):
-    """
-
-    """
+    """ """
     beam_nmodes, beam_mmodes = np.meshgrid(
-                                    np.arange(1, beam_nmax + 1), 
-                                    np.arange(-beam_mmax, beam_mmax + 1))
+        np.arange(1, beam_nmax + 1), np.arange(-beam_mmax, beam_mmax + 1)
+    )
     beam_nmodes = beam_nmodes.flatten()
     beam_mmodes = beam_mmodes.flatten()
 
@@ -876,29 +996,21 @@ def init_beam_sampler(beams, freqs, beam_mmax, beam_nmax):
     phi_fit = np.linspace(0, 2 * np.pi, num=360)
     PHI, RHO = np.meshgrid(phi_fit, rho_fit)
 
-    bess_matr_fit = get_bess_matr(beam_nmodes,
-                                                     beam_mmodes,
-                                                     RHO, PHI)
+    bess_matr_fit = get_bess_matr(beam_nmodes, beam_mmodes, RHO, PHI)
 
     beam_coeffs_fit = fit_bess_to_beam(
-                                      beams[0],
-                                      1e6 * freqs,
-                                      beam_nmodes,
-                                      beam_mmodes,
-                                      RHO, PHI,
-                                      force_spw_index=True)
-    
+        beams[0], 1e6 * freqs, beam_nmodes, beam_mmodes, RHO, PHI, force_spw_index=True
+    )
+
     print("\tBeam best fit dynamic range:")
     print("\t", np.amax(np.abs(beam_coeffs_fit)), np.amin(np.abs(beam_coeffs_fit)))
-    
+
     txs, tys, tzs = convert_to_tops(ra, dec, times, array_latitude)
 
     # area-preserving
     rho = np.sqrt(1 - tzs) / args.rho_const
     phi = np.arctan2(tys, txs)
-    bess_matr = hydra.beam_sampler.get_bess_matr(beam_nmodes,
-                                                 beam_mmodes,
-                                                 rho, phi)
+    bess_matr = hydra.beam_sampler.get_bess_matr(beam_nmodes, beam_mmodes, rho, phi)
 
     # All the same, so just repeat (for now)
     beam_coeffs = np.array(Nants * [beam_coeffs_fit])
@@ -908,24 +1020,29 @@ def init_beam_sampler(beams, freqs, beam_mmax, beam_nmax):
     ncoeffs = beam_coeffs.shape[0]
 
     if PLOTTING:
-        plot_beam_cross(beam_coeffs, 0, 0, '_best_fit')
+        plot_beam_cross(beam_coeffs, 0, 0, "_best_fit")
 
     amp_use = x_soln if SAMPLE_PTSRC_AMPS else ptsrc_amps
     flux_use = get_flux_from_ptsrc_amp(amp_use, freqs, beta_ptsrc)
 
     # Hardcoded parameters. Make variations smooth in time/freq.
     sig_freq = 0.5 * (freqs[-1] - freqs[0])
-    cov_tuple = hydra.beam_sampler.make_prior_cov(freqs, times, ncoeffs,
-                                                  args.beam_prior_std, sig_freq,
-                                                  ridge=1e-6)
+    cov_tuple = hydra.beam_sampler.make_prior_cov(
+        freqs, times, ncoeffs, args.beam_prior_std, sig_freq, ridge=1e-6
+    )
     cho_tuple = hydra.beam_sampler.do_cov_cho(cov_tuple, check_op=False)
-    cov_tuple_0 = hydra.beam_sampler.make_prior_cov(freqs, times, ncoeffs,
-                                                  args.beam_prior_std, sig_freq,
-                                                  ridge=1e-6,
-                                                  constrain_phase=True,
-                                                  constraint=1)
+    cov_tuple_0 = hydra.beam_sampler.make_prior_cov(
+        freqs,
+        times,
+        ncoeffs,
+        args.beam_prior_std,
+        sig_freq,
+        ridge=1e-6,
+        constrain_phase=True,
+        constraint=1,
+    )
     cho_tuple_0 = hydra.beam_sampler.do_cov_cho(cov_tuple, check_op=False)
-    
+
     # Be lazy and just use the initial guess.
     coeff_mean = beam_coeffs[:, :, 0]
-    bess_outer = hydra.beam_sampler.get_bess_outer(bess_matr)
\ No newline at end of file
+    bess_outer = hydra.beam_sampler.get_bess_outer(bess_matr)
diff --git a/hydra/config.py b/hydra/config.py
index 818077f..406f6fa 100644
--- a/hydra/config.py
+++ b/hydra/config.py
@@ -1,181 +1,415 @@
-
 import argparse
 import numpy as np
 
+
 def get_config():
     """
     Parse commandline arguments to get configuration settings.
     """
-    description = "Example Gibbs sampling of the joint posterior of several analysis " \
-                  "parameters in 21-cm power spectrum estimation from a simulated " \
-                  "visibility data set"
+    description = (
+        "Example Gibbs sampling of the joint posterior of several analysis "
+        "parameters in 21-cm power spectrum estimation from a simulated "
+        "visibility data set"
+    )
     parser = argparse.ArgumentParser(description=description)
-    parser.add_argument("--seed", type=int, action="store", default=0,
-                        required=False, dest="seed",
-                        help="Set the random seed.")
+    parser.add_argument(
+        "--seed",
+        type=int,
+        action="store",
+        default=0,
+        required=False,
+        dest="seed",
+        help="Set the random seed.",
+    )
 
     # Samplers
-    parser.add_argument("--gains", action="store_true",
-                        required=False, dest="sample_gains",
-                        help="Sample gains.")
-    parser.add_argument("--cosmo", action="store_true",
-                        required=False, dest="sample_cosmo_field",
-                        help="Sample cosmo field.")
-    #parser.add_argument("--vis", action="store_true",
+    parser.add_argument(
+        "--gains",
+        action="store_true",
+        required=False,
+        dest="sample_gains",
+        help="Sample gains.",
+    )
+    parser.add_argument(
+        "--cosmo",
+        action="store_true",
+        required=False,
+        dest="sample_cosmo_field",
+        help="Sample cosmo field.",
+    )
+    # parser.add_argument("--vis", action="store_true",
     #                    required=False, dest="sample_vis",
     #                    help="Sample visibilities in general.")
-    parser.add_argument("--ptsrc", action="store_true",
-                        required=False, dest="sample_ptsrc",
-                        help="Sample point source amplitudes.")
-    parser.add_argument("--regions", action="store_true",
-                        required=False, dest="sample_regions",
-                        help="Sample amplitudes of regions of a diffuse map.")
-    parser.add_argument("--beam", action="store_true",
-                        required=False, dest="sample_beam",
-                        help="Sample beams.")
-    parser.add_argument("--sh", action="store_true",
-                        required=False, dest="sample_sh",
-                        help="Sample spherical harmonic modes.")
-    parser.add_argument("--cl", action="store_true",
-                        required=False, dest="sample_sh_pspec",
-                        help="Sample spherical harmonic angular power spectrum.")
-    parser.add_argument("--pspec", action="store_true",
-                        required=False, dest="sample_pspec",
-                        help="Sample 21cm power spectrum.")
+    parser.add_argument(
+        "--ptsrc",
+        action="store_true",
+        required=False,
+        dest="sample_ptsrc",
+        help="Sample point source amplitudes.",
+    )
+    parser.add_argument(
+        "--regions",
+        action="store_true",
+        required=False,
+        dest="sample_regions",
+        help="Sample amplitudes of regions of a diffuse map.",
+    )
+    parser.add_argument(
+        "--beam",
+        action="store_true",
+        required=False,
+        dest="sample_beam",
+        help="Sample beams.",
+    )
+    parser.add_argument(
+        "--sh",
+        action="store_true",
+        required=False,
+        dest="sample_sh",
+        help="Sample spherical harmonic modes.",
+    )
+    parser.add_argument(
+        "--cl",
+        action="store_true",
+        required=False,
+        dest="sample_sh_pspec",
+        help="Sample spherical harmonic angular power spectrum.",
+    )
+    parser.add_argument(
+        "--pspec",
+        action="store_true",
+        required=False,
+        dest="sample_pspec",
+        help="Sample 21cm power spectrum.",
+    )
 
     # Debug mode
-    parser.add_argument("--debug", action="store_true",
-                        required=False, dest="debug",
-                        help="Whether to enable debug mode, which shows extra output.")
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        required=False,
+        dest="debug",
+        help="Whether to enable debug mode, which shows extra output.",
+    )
 
     # Output options
-    parser.add_argument("--stats", action="store_true",
-                        required=False, dest="calculate_stats",
-                        help="Calcultae statistics about the sampling results.")
-    parser.add_argument("--diagnostics", action="store_true",
-                        required=False, dest="output_diagnostics",
-                        help="Output diagnostics.") # This will be ignored
-    parser.add_argument("--timing", action="store_true", required=False,
-                        dest="save_timing_info", help="Save timing info.")
-    parser.add_argument("--plotting", action="store_true",
-                        required=False, dest="plotting",
-                        help="Output plots.")
+    parser.add_argument(
+        "--stats",
+        action="store_true",
+        required=False,
+        dest="calculate_stats",
+        help="Calcultae statistics about the sampling results.",
+    )
+    parser.add_argument(
+        "--diagnostics",
+        action="store_true",
+        required=False,
+        dest="output_diagnostics",
+        help="Output diagnostics.",
+    )  # This will be ignored
+    parser.add_argument(
+        "--timing",
+        action="store_true",
+        required=False,
+        dest="save_timing_info",
+        help="Save timing info.",
+    )
+    parser.add_argument(
+        "--plotting",
+        action="store_true",
+        required=False,
+        dest="plotting",
+        help="Output plots.",
+    )
 
     # Array and data shape options
-    parser.add_argument("--latitude", type=float, action="store", default=-30.7215,
-                        required=False, dest="latitude",
-                        help="Latitude of the array, in degrees.")
-    parser.add_argument('--hex-array', type=int, action="store", default=(3,4),
-                        required=False, nargs='+', dest="hex_array",
-                        help="Hex array layout, specified as the no. of antennas "
-                             "in the 1st and middle rows, e.g. '--hex-array 3 4'.")
-    parser.add_argument("--Nptsrc", type=int, action="store", default=100,
-                        required=False, dest="Nptsrc",
-                        help="Number of point sources to use in simulation (and model).")
-    parser.add_argument("--Ntimes", type=int, action="store", default=30,
-                        required=False, dest="Ntimes",
-                        help="Number of times to use in the simulation.")
-    parser.add_argument("--Nfreqs", type=int, action="store", default=60,
-                        required=False, dest="Nfreqs",
-                        help="Number of frequencies to use in the simulation.")
-    parser.add_argument("--Niters", type=int, action="store", default=100,
-                        required=False, dest="Niters",
-                        help="Number of joint samples to gather.")
+    parser.add_argument(
+        "--latitude",
+        type=float,
+        action="store",
+        default=-30.7215,
+        required=False,
+        dest="latitude",
+        help="Latitude of the array, in degrees.",
+    )
+    parser.add_argument(
+        "--hex-array",
+        type=int,
+        action="store",
+        default=(3, 4),
+        required=False,
+        nargs="+",
+        dest="hex_array",
+        help="Hex array layout, specified as the no. of antennas "
+        "in the 1st and middle rows, e.g. '--hex-array 3 4'.",
+    )
+    parser.add_argument(
+        "--Nptsrc",
+        type=int,
+        action="store",
+        default=100,
+        required=False,
+        dest="Nptsrc",
+        help="Number of point sources to use in simulation (and model).",
+    )
+    parser.add_argument(
+        "--Ntimes",
+        type=int,
+        action="store",
+        default=30,
+        required=False,
+        dest="Ntimes",
+        help="Number of times to use in the simulation.",
+    )
+    parser.add_argument(
+        "--Nfreqs",
+        type=int,
+        action="store",
+        default=60,
+        required=False,
+        dest="Nfreqs",
+        help="Number of frequencies to use in the simulation.",
+    )
+    parser.add_argument(
+        "--Niters",
+        type=int,
+        action="store",
+        default=100,
+        required=False,
+        dest="Niters",
+        help="Number of joint samples to gather.",
+    )
 
     # Noise level
-    parser.add_argument("--sigma-noise", type=float, action="store",
-                        default=0.05, required=False, dest="sigma_noise",
-                        help="Standard deviation of the noise, in the same units "
-                             "as the visibility data.")
-
-    parser.add_argument("--solver", type=str, action="store",
-                        default='cg', required=False, dest="solver_name",
-                        help="Which linear solver to use ('cg' or 'gmres' or 'mpicg').")
-    #parser.add_argument("--mpicg-split", type=int, action="store",
+    parser.add_argument(
+        "--sigma-noise",
+        type=float,
+        action="store",
+        default=0.05,
+        required=False,
+        dest="sigma_noise",
+        help="Standard deviation of the noise, in the same units "
+        "as the visibility data.",
+    )
+
+    parser.add_argument(
+        "--solver",
+        type=str,
+        action="store",
+        default="cg",
+        required=False,
+        dest="solver_name",
+        help="Which linear solver to use ('cg' or 'gmres' or 'mpicg').",
+    )
+    # parser.add_argument("--mpicg-split", type=int, action="store",
     #                    default=1, required=False, dest="mpicg_split",
     #                    help="If the MPI CG solver is being used, how many blocks to split the linear system into.")
-    parser.add_argument("--output-dir", type=str, action="store",
-                        default="./output", required=False, dest="output_dir",
-                        help="Output directory.")
-    #parser.add_argument("--multiprocess", action="store_true", dest="multiprocess",
+    parser.add_argument(
+        "--output-dir",
+        type=str,
+        action="store",
+        default="./output",
+        required=False,
+        dest="output_dir",
+        help="Output directory.",
+    )
+    # parser.add_argument("--multiprocess", action="store_true", dest="multiprocess",
     #                    required=False,
     #                    help="Whether to use multiprocessing in vis sim calls.")
 
     # Point source sim params
-    parser.add_argument("--ra-bounds", type=float, action="store", default=(0, 1),
-                        nargs=2, required=False, dest="ra_bounds",
-                        help="Bounds for the Right Ascension of the randomly simulated sources")
-    parser.add_argument("--dec-bounds", type=float, action="store", default=(-0.6, 0.4),
-                        nargs=2, required=False, dest="dec_bounds",
-                        help="Bounds for the Declination of the randomly simulated sources")
-    parser.add_argument("--lst-bounds", type=float, action="store", default=(0.2, 0.5),
-                        nargs=2, required=False, dest="lst_bounds",
-                        help="Bounds for the LST range of the simulation, in radians.")
-    parser.add_argument("--freq-bounds", type=float, action="store", default=(100., 120.),
-                        nargs=2, required=False, dest="freq_bounds",
-                        help="Bounds for the frequency range of the simulation, in MHz.")
-    
+    parser.add_argument(
+        "--ra-bounds",
+        type=float,
+        action="store",
+        default=(0, 1),
+        nargs=2,
+        required=False,
+        dest="ra_bounds",
+        help="Bounds for the Right Ascension of the randomly simulated sources",
+    )
+    parser.add_argument(
+        "--dec-bounds",
+        type=float,
+        action="store",
+        default=(-0.6, 0.4),
+        nargs=2,
+        required=False,
+        dest="dec_bounds",
+        help="Bounds for the Declination of the randomly simulated sources",
+    )
+    parser.add_argument(
+        "--lst-bounds",
+        type=float,
+        action="store",
+        default=(0.2, 0.5),
+        nargs=2,
+        required=False,
+        dest="lst_bounds",
+        help="Bounds for the LST range of the simulation, in radians.",
+    )
+    parser.add_argument(
+        "--freq-bounds",
+        type=float,
+        action="store",
+        default=(100.0, 120.0),
+        nargs=2,
+        required=False,
+        dest="freq_bounds",
+        help="Bounds for the frequency range of the simulation, in MHz.",
+    )
 
     # Point source sampler parameters
-    parser.add_argument("--ptsrc-amp-prior-level", type=float, action="store", default=0.1,
-                        required=False, dest="ptsrc_amp_prior_level",
-                        help="Fractional prior on point source amplitudes")
-    #parser.add_argument("--vis-prior-level", type=float, action="store", default=0.1,
+    parser.add_argument(
+        "--ptsrc-amp-prior-level",
+        type=float,
+        action="store",
+        default=0.1,
+        required=False,
+        dest="ptsrc_amp_prior_level",
+        help="Fractional prior on point source amplitudes",
+    )
+    # parser.add_argument("--vis-prior-level", type=float, action="store", default=0.1,
     #                    required=False, dest="vis_prior_level",
     #                    help="Prior on visibility values")
-    parser.add_argument("--calsrc-std", type=float, action="store", default=-1.,
-                        required=False, dest="calsrc_std",
-                        help="Define a different std. dev. for the amplitude prior of a calibration source. If -1, do not use a calibration source.")
-    parser.add_argument("--calsrc-radius", type=float, action="store", default=10.,
-                        required=False, dest="calsrc_radius",
-                        help="Radius around declination of the zenith in which to search for brightest source, which is then identified as the calibration source.")
-    
-    # Diffuse model simulation parameters
-    parser.add_argument("--sim-diffuse-sky-model", type=str, action="store",
-                        default="none", required=False, dest="sim_diffuse_sky_model",
-                        help="Which global sky model to use to define diffuse model. "
-                             "The options all come from pyGDSM: 'gsm2008', 'gsm2016', 'haslam', "
-                             "'lfss', or 'none'.")
-    parser.add_argument("--sim-diffuse-nside", type=int, action="store",
-                        default=16, required=False, dest="sim_diffuse_nside",
-                        help="Healpix nside to use for the region maps.")
+    parser.add_argument(
+        "--calsrc-std",
+        type=float,
+        action="store",
+        default=-1.0,
+        required=False,
+        dest="calsrc_std",
+        help="Define a different std. dev. for the amplitude prior of a calibration source. If -1, do not use a calibration source.",
+    )
+    parser.add_argument(
+        "--calsrc-radius",
+        type=float,
+        action="store",
+        default=10.0,
+        required=False,
+        dest="calsrc_radius",
+        help="Radius around declination of the zenith in which to search for brightest source, which is then identified as the calibration source.",
+    )
 
+    # Diffuse model simulation parameters
+    parser.add_argument(
+        "--sim-diffuse-sky-model",
+        type=str,
+        action="store",
+        default="none",
+        required=False,
+        dest="sim_diffuse_sky_model",
+        help="Which global sky model to use to define diffuse model. "
+        "The options all come from pyGDSM: 'gsm2008', 'gsm2016', 'haslam', "
+        "'lfss', or 'none'.",
+    )
+    parser.add_argument(
+        "--sim-diffuse-nside",
+        type=int,
+        action="store",
+        default=16,
+        required=False,
+        dest="sim_diffuse_nside",
+        help="Healpix nside to use for the region maps.",
+    )
 
     # Cosmo field parameters
-    parser.add_argument("--cosmo-ra-bounds", type=float, action="store", default=(0, 60),
-                        nargs=2, required=False, dest="cosmo_field_ra_bounds",
-                        help="Bounds for the RA of the cosmo field sample points (in degrees).")
-    parser.add_argument("--cosmo-dec-bounds", type=float, action="store", default=(-40., -20.),
-                        nargs=2, required=False, dest="cosmo_field_dec_bounds",
-                        help="Bounds for the Dec of the cosmo field sample points (in degrees).")
-    parser.add_argument("--cosmo-ra-ngrid", type=int, action="store", default=10,
-                        required=False, dest="cosmo_field_ra_ngrid",
-                        help="Number of cosmo field sample points in the RA direction.")
-    parser.add_argument("--cosmo-dec-ngrid", type=int, action="store", default=10,
-                        required=False, dest="cosmo_field_dec_ngrid",
-                        help="Number of cosmo field sample points in the Dec direction.")
+    parser.add_argument(
+        "--cosmo-ra-bounds",
+        type=float,
+        action="store",
+        default=(0, 60),
+        nargs=2,
+        required=False,
+        dest="cosmo_field_ra_bounds",
+        help="Bounds for the RA of the cosmo field sample points (in degrees).",
+    )
+    parser.add_argument(
+        "--cosmo-dec-bounds",
+        type=float,
+        action="store",
+        default=(-40.0, -20.0),
+        nargs=2,
+        required=False,
+        dest="cosmo_field_dec_bounds",
+        help="Bounds for the Dec of the cosmo field sample points (in degrees).",
+    )
+    parser.add_argument(
+        "--cosmo-ra-ngrid",
+        type=int,
+        action="store",
+        default=10,
+        required=False,
+        dest="cosmo_field_ra_ngrid",
+        help="Number of cosmo field sample points in the RA direction.",
+    )
+    parser.add_argument(
+        "--cosmo-dec-ngrid",
+        type=int,
+        action="store",
+        default=10,
+        required=False,
+        dest="cosmo_field_dec_ngrid",
+        help="Number of cosmo field sample points in the Dec direction.",
+    )
 
     # Gain prior
-    parser.add_argument("--gain-prior-amp", type=float, action="store", default=0.1,
-                        required=False, dest="gain_prior_amp",
-                        help="Overall amplitude of gain prior.")
-    parser.add_argument("--gain-nmax-freq", type=int, action="store", default=2,
-                        required=False, dest="gain_nmaxfreq",
-                        help="Max. Fourier mode index for gain perturbations (freq. direction).")
-    parser.add_argument("--gain-nmax-time", type=int, action="store", default=2,
-                        required=False, dest="gain_nmaxtime",
-                        help="Max. Fourier mode index for gain perturbations (time direction).")
-    parser.add_argument("--gain-prior-zero-mode-std", type=float, action="store", default=None,
-                        required=False, dest="gain_prior_zero_mode_std",
-                        help="Separately specify the gain prior standard deviaiton for the zero mode.")
-    parser.add_argument("--gain-only-positive-modes", type=bool, action="store", default=True,
-                        required=False, dest="gain_only_positive_modes",
-                        help="Whether to only permit positive wavenumber gain modes.")
+    parser.add_argument(
+        "--gain-prior-amp",
+        type=float,
+        action="store",
+        default=0.1,
+        required=False,
+        dest="gain_prior_amp",
+        help="Overall amplitude of gain prior.",
+    )
+    parser.add_argument(
+        "--gain-nmax-freq",
+        type=int,
+        action="store",
+        default=2,
+        required=False,
+        dest="gain_nmaxfreq",
+        help="Max. Fourier mode index for gain perturbations (freq. direction).",
+    )
+    parser.add_argument(
+        "--gain-nmax-time",
+        type=int,
+        action="store",
+        default=2,
+        required=False,
+        dest="gain_nmaxtime",
+        help="Max. Fourier mode index for gain perturbations (time direction).",
+    )
+    parser.add_argument(
+        "--gain-prior-zero-mode-std",
+        type=float,
+        action="store",
+        default=None,
+        required=False,
+        dest="gain_prior_zero_mode_std",
+        help="Separately specify the gain prior standard deviaiton for the zero mode.",
+    )
+    parser.add_argument(
+        "--gain-only-positive-modes",
+        type=bool,
+        action="store",
+        default=True,
+        required=False,
+        dest="gain_only_positive_modes",
+        help="Whether to only permit positive wavenumber gain modes.",
+    )
 
     # Gain simulation
-    parser.add_argument("--sim-gain-amp-std", type=float, action="store", default=0.05,
-                        required=False, dest="sim_gain_amp_std",
-                        help="Std. dev. of amplitude of simulated gain.")
+    parser.add_argument(
+        "--sim-gain-amp-std",
+        type=float,
+        action="store",
+        default=0.05,
+        required=False,
+        dest="sim_gain_amp_std",
+        help="Std. dev. of amplitude of simulated gain.",
+    )
     # parser.add_argument("--sim-gain-sigma-frate", type=float, action="store", default=None,
     #                     required=False, dest="sim_gain_sigma_frate",
     #                     help="Width of a Gaussian in fringe rate, in units of mHz.")
@@ -190,65 +424,166 @@ def get_config():
     #                     help="The central delay of the Gaussian taper (ns).")
 
     # Beam parameters
-    parser.add_argument("--beam-sim-type", type=str, action="store", default="gaussian",
-                        required=False, dest="beam_sim_type",
-                        help="Which type of beam to use for the simulation. ['gaussian', 'polybeam']")
-    parser.add_argument("--beam-prior-std", type=float, action="store", default=1,
-                        required=False, dest="beam_prior_std",
-                        help="Std. dev. of beam coefficient prior, in units of Zernike coefficient")
-    parser.add_argument("--beam-nmax", type=int, action="store",
-                        default=16, required=False, dest="beam_nmax",
-                        help="Maximum radial degree of the Fourier-Bessel basis for the beams.")
-    parser.add_argument("--beam-mmax", type=int, action="store",
-                        default=0, required=False, dest="beam_mmax",
-                        help="Maximum azimuthal degree of the Fourier-Bessel basis for the beams.")
-    parser.add_argument("--rho-const", type=float, action="store", 
-                        default=np.sqrt(1-np.cos(np.pi * 23 / 45)),
-                        required=False, dest="rho_const",
-                        help="A constant to define the radial projection for the beam spatial basis")
+    parser.add_argument(
+        "--beam-sim-type",
+        type=str,
+        action="store",
+        default="gaussian",
+        required=False,
+        dest="beam_sim_type",
+        help="Which type of beam to use for the simulation. ['gaussian', 'polybeam']",
+    )
+    parser.add_argument(
+        "--beam-prior-std",
+        type=float,
+        action="store",
+        default=1,
+        required=False,
+        dest="beam_prior_std",
+        help="Std. dev. of beam coefficient prior, in units of Zernike coefficient",
+    )
+    parser.add_argument(
+        "--beam-nmax",
+        type=int,
+        action="store",
+        default=16,
+        required=False,
+        dest="beam_nmax",
+        help="Maximum radial degree of the Fourier-Bessel basis for the beams.",
+    )
+    parser.add_argument(
+        "--beam-mmax",
+        type=int,
+        action="store",
+        default=0,
+        required=False,
+        dest="beam_mmax",
+        help="Maximum azimuthal degree of the Fourier-Bessel basis for the beams.",
+    )
+    parser.add_argument(
+        "--rho-const",
+        type=float,
+        action="store",
+        default=np.sqrt(1 - np.cos(np.pi * 23 / 45)),
+        required=False,
+        dest="rho_const",
+        help="A constant to define the radial projection for the beam spatial basis",
+    )
 
     # Spherical harmonic parameters
-    parser.add_argument("--sim-sh-lmax", type=int, action="store", default=8,
-                        required=False, dest="sim_sh_lmax",
-                        help="Maximum ell value to include for spherical harmonic simulation.")
-    parser.add_argument("--sim-sh-nside", type=int, action="store",
-                        default=16, required=False, dest="sim_sh_nside",
-                        help="Healpix nside used to construct simulated spherical harmonic response.")
-    parser.add_argument("--sh-lmax", type=int, action="store", default=8,
-                        required=False, dest="sh_lmax",
-                        help="Maximum ell value to include for spherical harmonic sampler.")
-    parser.add_argument("--sh-nside", type=int, action="store",
-                        default=16, required=False, dest="sh_nside",
-                        help="Healpix nside used to construct spherical harmonic response function.")
-    parser.add_argument("--sh-prior-std", type=float, action="store",
-                        default=0.1, required=False, dest="sh_prior_std",
-                        help="Prior standard deviation for spherical harmonic modes.")
-    parser.add_argument("--sh-ref-freq", type=float, action="store",
-                        default=100., required=False, dest="sh_ref_freq",
-                        help="Reference frequency for the SH spectral dependence, in MHz.")
-    parser.add_argument("--sh-spectral-idx", type=float, action="store",
-                        default=0., required=False, dest="sh_spectral_idx",
-                        help="Spectral index for the SH power law spectral dependence.")
+    parser.add_argument(
+        "--sim-sh-lmax",
+        type=int,
+        action="store",
+        default=8,
+        required=False,
+        dest="sim_sh_lmax",
+        help="Maximum ell value to include for spherical harmonic simulation.",
+    )
+    parser.add_argument(
+        "--sim-sh-nside",
+        type=int,
+        action="store",
+        default=16,
+        required=False,
+        dest="sim_sh_nside",
+        help="Healpix nside used to construct simulated spherical harmonic response.",
+    )
+    parser.add_argument(
+        "--sh-lmax",
+        type=int,
+        action="store",
+        default=8,
+        required=False,
+        dest="sh_lmax",
+        help="Maximum ell value to include for spherical harmonic sampler.",
+    )
+    parser.add_argument(
+        "--sh-nside",
+        type=int,
+        action="store",
+        default=16,
+        required=False,
+        dest="sh_nside",
+        help="Healpix nside used to construct spherical harmonic response function.",
+    )
+    parser.add_argument(
+        "--sh-prior-std",
+        type=float,
+        action="store",
+        default=0.1,
+        required=False,
+        dest="sh_prior_std",
+        help="Prior standard deviation for spherical harmonic modes.",
+    )
+    parser.add_argument(
+        "--sh-ref-freq",
+        type=float,
+        action="store",
+        default=100.0,
+        required=False,
+        dest="sh_ref_freq",
+        help="Reference frequency for the SH spectral dependence, in MHz.",
+    )
+    parser.add_argument(
+        "--sh-spectral-idx",
+        type=float,
+        action="store",
+        default=0.0,
+        required=False,
+        dest="sh_spectral_idx",
+        help="Spectral index for the SH power law spectral dependence.",
+    )
 
     # Region parameters
-    parser.add_argument("--region-nregions", type=int, action="store", default=10,
-                        required=False, dest="region_nregions",
-                        help="No. of regions to break up a diffuse map into.")
-    parser.add_argument("--region-smoothing-fwhm", type=float, action="store",
-                        default=None, required=False, dest="region_smoothing_fwhm",
-                        help="Smoothing FWHM to apply to segmented diffuse map, to smooth "
-                        "sharp edges.")
-    parser.add_argument("--region-sky-model", type=str, action="store",
-                        default="gsm2016", required=False, dest="region_sky_model",
-                        help="Which global sky model to use to define diffuse regions. "
-                             "The options all come from pyGDSM: 'gsm2008', 'gsm2016', 'haslam', "
-                             "'lfss'.")
-    parser.add_argument("--region-nside", type=int, action="store",
-                        default=16, required=False, dest="region_nside",
-                        help="Healpix nside to use for the region maps.")
-    parser.add_argument("--region-amp-prior-level", type=float, action="store", default=0.1,
-                        required=False, dest="region_amp_prior_level",
-                        help="Fractional prior on diffuse region amplitudes")
+    parser.add_argument(
+        "--region-nregions",
+        type=int,
+        action="store",
+        default=10,
+        required=False,
+        dest="region_nregions",
+        help="No. of regions to break up a diffuse map into.",
+    )
+    parser.add_argument(
+        "--region-smoothing-fwhm",
+        type=float,
+        action="store",
+        default=None,
+        required=False,
+        dest="region_smoothing_fwhm",
+        help="Smoothing FWHM to apply to segmented diffuse map, to smooth "
+        "sharp edges.",
+    )
+    parser.add_argument(
+        "--region-sky-model",
+        type=str,
+        action="store",
+        default="gsm2016",
+        required=False,
+        dest="region_sky_model",
+        help="Which global sky model to use to define diffuse regions. "
+        "The options all come from pyGDSM: 'gsm2008', 'gsm2016', 'haslam', "
+        "'lfss'.",
+    )
+    parser.add_argument(
+        "--region-nside",
+        type=int,
+        action="store",
+        default=16,
+        required=False,
+        dest="region_nside",
+        help="Healpix nside to use for the region maps.",
+    )
+    parser.add_argument(
+        "--region-amp-prior-level",
+        type=float,
+        action="store",
+        default=0.1,
+        required=False,
+        dest="region_amp_prior_level",
+        help="Fractional prior on diffuse region amplitudes",
+    )
 
     args = parser.parse_args()
-    return args
\ No newline at end of file
+    return args
diff --git a/hydra/cosmo_sampler.py b/hydra/cosmo_sampler.py
index 3e64c2d..7781d48 100644
--- a/hydra/cosmo_sampler.py
+++ b/hydra/cosmo_sampler.py
@@ -1,16 +1,15 @@
-
 import numpy as np
 
 
 def make_cosmo_field_grid(args):
     """
-    Make a regular Cartesian grid of points in RA and Dec that sample a 
+    Make a regular Cartesian grid of points in RA and Dec that sample a
     cosmological 21cm field.
 
     Parameters:
         args (argparse object):
-            An argparse object containing the following settings: 
-            `cosmo_field_ra_bounds`, `cosmo_field_dec_bounds`, 
+            An argparse object containing the following settings:
+            `cosmo_field_ra_bounds`, `cosmo_field_dec_bounds`,
             `cosmo_field_ra_ngrid`, `cosmo_field_dec_ngrid`.
 
     Returns:
@@ -18,36 +17,42 @@ def make_cosmo_field_grid(args):
             RA and Dec values of the sample points, in radians.
     """
     # Define sample points
-    ra = np.linspace(min(args.cosmo_field_ra_bounds), 
-                     max(args.cosmo_field_ra_bounds), 
-                     args.cosmo_field_ra_ngrid)
-    dec = np.linspace(min(args.cosmo_field_dec_bounds), 
-                      max(args.cosmo_field_dec_bounds), 
-                      args.cosmo_field_dec_ngrid)
+    ra = np.linspace(
+        min(args.cosmo_field_ra_bounds),
+        max(args.cosmo_field_ra_bounds),
+        args.cosmo_field_ra_ngrid,
+    )
+    dec = np.linspace(
+        min(args.cosmo_field_dec_bounds),
+        max(args.cosmo_field_dec_bounds),
+        args.cosmo_field_dec_ngrid,
+    )
 
     # Define 2D grid
     ra_grid, dec_grid = np.meshgrid(ra, dec)
-    return ra_grid.flatten(), dec_grid.flatten() 
-
-
-def precompute_mpi(comm,
-                   ants, 
-                   antpairs,
-                   freqs,
-                   freq_chunk, 
-                   time_chunk,
-                   proj_chunk,
-                   data_chunk,
-                   inv_noise_var_chunk,
-                   current_data_model_chunk,
-                   gain_chunk, 
-                   amp_prior_std, 
-                   realisation=True):
+    return ra_grid.flatten(), dec_grid.flatten()
+
+
+def precompute_mpi(
+    comm,
+    ants,
+    antpairs,
+    freqs,
+    freq_chunk,
+    time_chunk,
+    proj_chunk,
+    data_chunk,
+    inv_noise_var_chunk,
+    current_data_model_chunk,
+    gain_chunk,
+    amp_prior_std,
+    realisation=True,
+):
     """
-    Precompute the projection operator and matrix operator in parallel. 
+    Precompute the projection operator and matrix operator in parallel.
 
-    The projection operator is computed in chunks in time and frequency. 
-    The overall matrix operator can be computed by summing the matrix 
+    The projection operator is computed in chunks in time and frequency.
+    The overall matrix operator can be computed by summing the matrix
     operator for the time and frequency chunks.
     """
     myid = comm.Get_rank()
@@ -56,7 +61,7 @@ def precompute_mpi(comm,
     assert data_chunk.shape == (len(antpairs), freq_chunk.size, time_chunk.size)
     assert data_chunk.shape == inv_noise_var_chunk.shape
     assert data_chunk.shape == current_data_model_chunk.shape
-    proj = proj_chunk.copy() # make a copy so we don't alter the original proj!
+    proj = proj_chunk.copy()  # make a copy so we don't alter the original proj!
 
     # FIXME: Check for unused args!
 
@@ -65,83 +70,97 @@ def precompute_mpi(comm,
         ant1, ant2 = bl
         i1 = np.where(ants == ant1)[0][0]
         i2 = np.where(ants == ant2)[0][0]
-        proj[k,:,:,:] *= gain_chunk[i1,:,:,np.newaxis] \
-                       * gain_chunk[i2,:,:,np.newaxis].conj()
+        proj[k, :, :, :] *= (
+            gain_chunk[i1, :, :, np.newaxis] * gain_chunk[i2, :, :, np.newaxis].conj()
+        )
 
-    # (2) Precompute linear system operator for each frequency (for the 
+    # (2) Precompute linear system operator for each frequency (for the
     # likelihood part of the operator, the freqs. don't talk to each other)
     Npix = proj.shape[-1]
     my_linear_op = np.zeros((freqs.size, Npix, Npix), dtype=proj.real.dtype)
 
-    # inv_noise_var has shape (Nbls, Nfreqs, Ntimes); proj has shape (Nbls, Nfreqs, Ntimes, Npix) 
-    v_re = (proj.real * np.sqrt(inv_noise_var_chunk[...,np.newaxis])).reshape((-1, Npix))
-    v_im = (proj.imag * np.sqrt(inv_noise_var_chunk[...,np.newaxis])).reshape(((-1, Npix)))
+    # inv_noise_var has shape (Nbls, Nfreqs, Ntimes); proj has shape (Nbls, Nfreqs, Ntimes, Npix)
+    v_re = (proj.real * np.sqrt(inv_noise_var_chunk[..., np.newaxis])).reshape(
+        (-1, Npix)
+    )
+    v_im = (proj.imag * np.sqrt(inv_noise_var_chunk[..., np.newaxis])).reshape(
+        ((-1, Npix))
+    )
 
     # Treat real and imaginary separately; treat frequencies separately
     # FIXME: Is this neglecting real/imag cross-terms?
     for j in range(freq_chunk.size):
         i = np.where(freqs == freq_chunk[j])[0]
-        my_linear_op[i,:,:] = v_re[:,j,:,:].T @ v_re[:,j,:,:] \
-                            + v_im[:,j,:,:].T @ v_im[:,j,:,:]
+        my_linear_op[i, :, :] = (
+            v_re[:, j, :, :].T @ v_re[:, j, :, :]
+            + v_im[:, j, :, :].T @ v_im[:, j, :, :]
+        )
     del v_re, v_im
 
     # Do Reduce (sum) operation to get total operator on root node
-    linear_op = np.zeros((1,1,1), dtype=my_linear_op.dtype) # dummy data for non-root workers
+    linear_op = np.zeros(
+        (1, 1, 1), dtype=my_linear_op.dtype
+    )  # dummy data for non-root workers
     if myid == 0:
         linear_op = np.zeros_like(my_linear_op)
-    
-    comm.Reduce(my_linear_op,
-                linear_op,
-                op=MPI_SUM,
-                root=0)
+
+    comm.Reduce(my_linear_op, linear_op, op=MPI_SUM, root=0)
 
     #######################
     ### FIXME: Got to here
 
     # Include prior and identity terms to finish constructing LHS operator on root worker
     if myid == 0:
-        linear_op = np.eye(linear_op.shape[0]) \
-                  + np.diag(amp_prior_std) @ linear_op @ np.diag(amp_prior_std)
-    
+        linear_op = np.eye(linear_op.shape[0]) + np.diag(
+            amp_prior_std
+        ) @ linear_op @ np.diag(amp_prior_std)
+
     # (3) Calculate linear system RHS
     proj = proj.reshape((-1, nsrcs))
-    realisation_switch = 1.0 if realisation else 0.0 # Turn random realisations on or off
+    realisation_switch = (
+        1.0 if realisation else 0.0
+    )  # Turn random realisations on or off
 
     # Construct current state of model (residual from amplitudes = 1)
     # (proj now includes gains)
-    resid_chunk = data_chunk.copy() \
-          - (  proj.reshape((-1, nsrcs)) 
-             @ np.ones_like(amp_prior_std) ).reshape(current_data_model_chunk.shape)
+    resid_chunk = data_chunk.copy() - (
+        proj.reshape((-1, nsrcs)) @ np.ones_like(amp_prior_std)
+    ).reshape(current_data_model_chunk.shape)
 
     # (Terms 1+3): S^1/2 A^\dagger [ N^{-1} r + N^{-1/2} \omega_r ]
     omega_n = (
         realisation_switch
-        * (1.0 * np.random.randn(*resid_chunk.shape) + 1.0j * np.random.randn(*resid_chunk.shape))
+        * (
+            1.0 * np.random.randn(*resid_chunk.shape)
+            + 1.0j * np.random.randn(*resid_chunk.shape)
+        )
         / np.sqrt(2.0)
     )
 
     # Separate complex part of RHS into real and imaginary parts, and apply
     # the real and imaginary parts of the projection operator separately.
     # This is necessary to get a real RHS vector
-    y = ((resid_chunk * inv_noise_var_chunk) + (omega_n * np.sqrt(inv_noise_var_chunk))).flatten()
+    y = (
+        (resid_chunk * inv_noise_var_chunk) + (omega_n * np.sqrt(inv_noise_var_chunk))
+    ).flatten()
     b = amp_prior_std * (proj.T.real @ y.real + proj.T.imag @ y.imag)
 
     # Reduce (sum) operation on b
-    linear_rhs = np.zeros((1,), dtype=b.dtype) # dummy data for non-root workers
+    linear_rhs = np.zeros((1,), dtype=b.dtype)  # dummy data for non-root workers
     if myid == 0:
         linear_rhs = np.zeros_like(b)
     comm.Reduce(b, linear_rhs, op=MPI_SUM, root=0)
 
     # (Term 2): \omega_a
     if myid == 0:
-        linear_rhs += realisation_switch * np.random.randn(nsrcs) # real vector
+        linear_rhs += realisation_switch * np.random.randn(nsrcs)  # real vector
 
     return linear_op, linear_rhs
 
 
 def apply_operator(x, freqs, ra_pix, dec_pix, linear_op_term, pspec):
     """
-    
+
     Parameters:
         x (array_like):
             1D array of cosmo field values that can be reshaped to `(Nfreqs, Npix)`.
@@ -160,4 +179,4 @@ def apply_operator(x, freqs, ra_pix, dec_pix, linear_op_term, pspec):
 
     # Apply prior term to x vector
     x_arr = x.reshape((Nfreqs, Nx, Ny))
-    y_vecnp.fft.ifftn(pspec * np.fft.fftn(x_arr)).reshape(y_vec.shape)
\ No newline at end of file
+    y_vecnp.fft.ifftn(pspec * np.fft.fftn(x_arr)).reshape(y_vec.shape)
diff --git a/hydra/diffuse_sampler.py b/hydra/diffuse_sampler.py
index 546bfba..db3927b 100644
--- a/hydra/diffuse_sampler.py
+++ b/hydra/diffuse_sampler.py
@@ -1,15 +1,16 @@
-
 from .vis_simulator import simulate_vis
 import numpy as np
 import healpy as hp
 
-def spectral_index_segments_basic(maps, freqs, boundaries, region_npix=None, 
-                                  fill_masked=True, avg_fn=np.mean):
+
+def spectral_index_segments_basic(
+    maps, freqs, boundaries, region_npix=None, fill_masked=True, avg_fn=np.mean
+):
     """
-    Split map into regions based on the spectral index calculated between 
-    two frequencies. This can result in region with quite disjoint morphologies 
+    Split map into regions based on the spectral index calculated between
+    two frequencies. This can result in region with quite disjoint morphologies
     and complicated boundaries.
-    
+
     Parameters:
         maps (array_like):
             Array of healpix maps, with shape (2, Npix).
@@ -18,45 +19,44 @@ def spectral_index_segments_basic(maps, freqs, boundaries, region_npix=None,
         boundaries (array_like):
             Boundaries of bins in spectral index.
         region_npix (int):
-            If specified, downgrade the resolution of the beta map to 
-            simplify the region boundaries, and then upgrade to the 
+            If specified, downgrade the resolution of the beta map to
+            simplify the region boundaries, and then upgrade to the
             original resolution.
         fill_masked (bool):
-            If True, fill masked pixels with the mean of the neighbouring 
+            If True, fill masked pixels with the mean of the neighbouring
             pixels of the spectral index map.
         avg_fn (function):
-            Which function to use to calculate the representative 
-            spectral index of each region. By default, it uses the 
+            Which function to use to calculate the representative
+            spectral index of each region. By default, it uses the
             mean.
-    
+
     Returns:
         beta (array_like):
-            Healpix map with the same resolution as the input map, 
+            Healpix map with the same resolution as the input map,
             but with integer region IDs as the map values.
     """
     assert maps.shape[0] == 2, "Requires a pair of Healpix maps"
     assert len(freqs) == 2, "Requires a pair of frequencies"
-    
+
     # Make sure boundaries are ordered
     if not np.all(boundaries[:-1] < boundaries[1:]):
         raise ValueError("boundaries array must be sorted in ascending order")
-    
+
     # Calculate per-pixel spectral index map
-    beta = np.log(maps[1] / maps[0]) \
-         / np.log(freqs[1] / freqs[0])
-    
+    beta = np.log(maps[1] / maps[0]) / np.log(freqs[1] / freqs[0])
+
     # Fill masked regions if requested
     if fill_masked:
         idxs = np.where(np.isnan(beta))
         print(idxs)
-    
+
     # Loop over bins
     regions = np.zeros_like(maps[0])
     regions[:] = np.nan
     for i in range(len(boundaries) - 1):
-        regions[(beta >= boundaries[i]) & (beta < boundaries[i+1])] = i
-    
-    # If region_npix is set, downgrade and then upgrade the 
+        regions[(beta >= boundaries[i]) & (beta < boundaries[i + 1])] = i
+
+    # If region_npix is set, downgrade and then upgrade the
     # resolution to simplify the region shapes
     if region_npix is not None:
         npix_in = hp.npix2nside(maps[0].size)
@@ -67,18 +67,25 @@ def spectral_index_segments_basic(maps, freqs, boundaries, region_npix=None,
 
 
 def calc_proj_operator_per_region(
-    fluxes, region_idxs, ant_pos, antpairs, freqs, times, beams,
-    latitude=-0.5361913261514378, multiprocess=True
+    fluxes,
+    region_idxs,
+    ant_pos,
+    antpairs,
+    freqs,
+    times,
+    beams,
+    latitude=-0.5361913261514378,
+    multiprocess=True,
 ):
     """
-    Calculate a visibility vector for each region of a healpix/healpy map, 
-    as a function of frequency, time, and baseline. This is the projection 
+    Calculate a visibility vector for each region of a healpix/healpy map,
+    as a function of frequency, time, and baseline. This is the projection
     operator from region amplitude to visibilities. Gains are not included.
-    
+
     Parameters:
         fluxes (array_like):
-            Flux for each pixel source as a function of frequency. Assumed to 
-            be a healpy/healpix array (in Galactic coords) per frequency, with 
+            Flux for each pixel source as a function of frequency. Assumed to
+            be a healpy/healpix array (in Galactic coords) per frequency, with
             shape `(Npix, Nfreq)`.
         region_idxs (array_like):
             A healpix/healpy map with the integer region ID of each pixel.
@@ -100,8 +107,8 @@ def calc_proj_operator_per_region(
             calculation
     Returns:
         vis_proj_operator (array_like):
-            The projection operator from region amplitudes to visibilities. 
-            This is an array of the visibility value contributed by each 
+            The projection operator from region amplitudes to visibilities.
+            This is an array of the visibility value contributed by each
             region if its amplitude were 1.
         unique_regions (array_like):
             An ordered list of integer region IDs.
@@ -111,27 +118,31 @@ def calc_proj_operator_per_region(
     Nants = len(ant_pos)
     Nvis = len(antpairs)
     assert fluxes.shape[1] == len(freqs), "`fluxes` must have shape (Npix, Nfreqs)"
-    
+
     # Get pixel Galactic and then equatorial coords of each map pixel (radians)
     nside = hp.npix2nside(re.size)
-    theta_gal, phi_gal = hp.pix2ang(nside=nside, ipix=np.arange(m[0].size), lonlat=False)
-    theta_eq, phi_eq = hp.Rotator(coord='ge', deg=False)(theta_gal, phi_gal)
+    theta_gal, phi_gal = hp.pix2ang(
+        nside=nside, ipix=np.arange(m[0].size), lonlat=False
+    )
+    theta_eq, phi_eq = hp.Rotator(coord="ge", deg=False)(theta_gal, phi_gal)
 
     # Empty array of per-point source visibilities
-    vis_regions = np.zeros((Nvis, freqs.size, times.size, Nregions), dtype=np.complex128)
+    vis_regions = np.zeros(
+        (Nvis, freqs.size, times.size, Nregions), dtype=np.complex128
+    )
 
     # Get visibility for each region
     for i in range(Nregions):
-        
+
         # Get indices of pixels that are in this region
         idxs = np.where(region_idxs == unique_regions[i])
-        
+
         # Simulate visibility for this region
         # Returns shape (Nfreqs, Ntimes, Nants, Nants)
         vis = simulate_vis(
             ants=ant_pos,
-            fluxes=fluxes[idxs,:],
-            ra=phi_eq[idxs], # FIXME: Make sure these occupy the correct range!
+            fluxes=fluxes[idxs, :],
+            ra=phi_eq[idxs],  # FIXME: Make sure these occupy the correct range!
             dec=theta_eq[idxs],
             freqs=freqs * 1e6,
             lsts=times,
@@ -140,7 +151,7 @@ def calc_proj_operator_per_region(
             precision=2,
             latitude=latitude,
             use_feed="x",
-            multiprocess=multiprocess
+            multiprocess=multiprocess,
         )
 
         # Allocate computed visibilities to only available baselines (saves memory)
diff --git a/hydra/example.py b/hydra/example.py
index 42618af..ea55a36 100644
--- a/hydra/example.py
+++ b/hydra/example.py
@@ -1,4 +1,3 @@
-
 import numpy as np
 
 from scipy.signal.windows import blackmanharris
@@ -6,29 +5,37 @@
 from hera_sim.beams import PolyBeam
 import time, os, resource
 from .vis_simulator import simulate_vis
-from .utils import flatten_vector, reconstruct_vector, timing_info, \
-                   build_hex_array, get_flux_from_ptsrc_amp, \
-                   convert_to_tops, gain_prior_pspec_sqrt, extract_vis_from_sim
-
-
-
-def generate_random_ptsrc_catalogue(Nptsrc, ra_bounds, dec_bounds, logflux_bounds=(-1., 2.)):
+from .utils import (
+    flatten_vector,
+    reconstruct_vector,
+    timing_info,
+    build_hex_array,
+    get_flux_from_ptsrc_amp,
+    convert_to_tops,
+    gain_prior_pspec_sqrt,
+    extract_vis_from_sim,
+)
+
+
+def generate_random_ptsrc_catalogue(
+    Nptsrc, ra_bounds, dec_bounds, logflux_bounds=(-1.0, 2.0)
+):
     """
-    Generate a catalogue of point sources with random positions and log(flux) 
+    Generate a catalogue of point sources with random positions and log(flux)
     values.
 
     Parameters:
         Nptsrc (int):
             Number of point sources in catalogue.
         ra_bounds, dec_bounds (tuple of float):
-            Tuples with the lower and upper bounds of the RA and Dec ranges. 
-            Sources will be placed according to a uniform random distribution 
+            Tuples with the lower and upper bounds of the RA and Dec ranges.
+            Sources will be placed according to a uniform random distribution
             within the specified interval.
         logflux_bounds (tuple of float):
-            The lower and upper bounds of log10(flux) at some reference 
-            frequency. The reference frequency is not specified here, so these 
-            values can be treated purely as amplitude/scaling factors. Sources 
-            will be randomly assigned fluxes/amplitudes in this interval in 
+            The lower and upper bounds of log10(flux) at some reference
+            frequency. The reference frequency is not specified here, so these
+            values can be treated purely as amplitude/scaling factors. Sources
+            will be randomly assigned fluxes/amplitudes in this interval in
             log10(flux).
 
     Returns:
@@ -41,26 +48,31 @@ def generate_random_ptsrc_catalogue(Nptsrc, ra_bounds, dec_bounds, logflux_bound
     ra_low, ra_high = (min(ra_bounds), max(ra_bounds))
     dec_low, dec_high = (min(dec_bounds), max(dec_bounds))
     logflux_low, logflux_high = (min(logflux_bounds), max(logflux_bounds))
-    
+
     # Generate random point source locations
     # RA goes from [0, 2 pi] and Dec from [-pi / 2, +pi / 2].
     ra = np.random.uniform(low=ra_low, high=ra_high, size=Nptsrc)
-    
+
     # Inversion sample to get them uniform on the sphere, in case wide bounds are used
     U = np.random.uniform(low=0, high=1, size=Nptsrc)
     dsin = np.sin(dec_high) - np.sin(dec_low)
-    dec = np.arcsin(U * dsin + np.sin(dec_low)) # np.arcsin returns on [-pi / 2, +pi / 2]
+    dec = np.arcsin(
+        U * dsin + np.sin(dec_low)
+    )  # np.arcsin returns on [-pi / 2, +pi / 2]
 
     # Generate fluxes
-    ptsrc_amps = 10.**np.random.uniform(low=logflux_low, high=logflux_high, size=Nptsrc)
+    ptsrc_amps = 10.0 ** np.random.uniform(
+        low=logflux_low, high=logflux_high, size=Nptsrc
+    )
     return ra, dec, ptsrc_amps
 
 
-def run_example_simulation(args, times, freqs, output_dir, ra, dec, ptsrc_amps, 
-                           array_latitude, verbose=False):
+def run_example_simulation(
+    args, times, freqs, output_dir, ra, dec, ptsrc_amps, array_latitude, verbose=False
+):
     """
-    Run an example visibility simulation for testing purposes, using 
-    semi-realistic beams, randomly placed point sources, and a hexagonal 
+    Run an example visibility simulation for testing purposes, using
+    semi-realistic beams, randomly placed point sources, and a hexagonal
     array layout.
 
     Parameters:
@@ -81,14 +93,14 @@ def run_example_simulation(args, times, freqs, output_dir, ra, dec, ptsrc_amps,
 
     Returns:
         model0 (array_like):
-            Model visibility computed by the simulator. The shape is given 
+            Model visibility computed by the simulator. The shape is given
             by `extract_vis_from_sim()`.
         fluxes (array_like):
             Point source fluxes, per source and per frequency.
         beams (list):
             List of `UVBeam` or `AnalyticBeam`-compatible beam objects.
         ant_info (tuple):
-            Tuple of arrays and dictionaries containing antenna information, 
+            Tuple of arrays and dictionaries containing antenna information,
             in the order: `(ants, ant_pos, antpairs, ants1, ants2)`.
     """
     # Dimensions of simulation
@@ -106,7 +118,7 @@ def run_example_simulation(args, times, freqs, output_dir, ra, dec, ptsrc_amps,
         for j in range(i, len(ants)):
             if i != j:
                 # Exclude autos
-                antpairs.append((i,j))
+                antpairs.append((i, j))
 
     ants1, ants2 = list(zip(*antpairs))
 
@@ -118,42 +130,57 @@ def run_example_simulation(args, times, freqs, output_dir, ra, dec, ptsrc_amps,
     # Beams
     if "polybeam" in args.beam_sim_type.lower():
         # PolyBeam fitted to HERA Fagnoni beam
-        beam_coeffs=[  0.29778665, -0.44821433, 0.27338272, 
-                      -0.10030698, -0.01195859, 0.06063853, 
-                      -0.04593295,  0.0107879,  0.01390283, 
-                      -0.01881641, -0.00177106, 0.01265177, 
-                      -0.00568299, -0.00333975, 0.00452368,
-                       0.00151808, -0.00593812, 0.00351559
-                     ]
-        beams = [PolyBeam(beam_coeffs, spectral_index=-0.6975, ref_freq=1.e8)
-                 for ant in ants]
+        beam_coeffs = [
+            0.29778665,
+            -0.44821433,
+            0.27338272,
+            -0.10030698,
+            -0.01195859,
+            0.06063853,
+            -0.04593295,
+            0.0107879,
+            0.01390283,
+            -0.01881641,
+            -0.00177106,
+            0.01265177,
+            -0.00568299,
+            -0.00333975,
+            0.00452368,
+            0.00151808,
+            -0.00593812,
+            0.00351559,
+        ]
+        beams = [
+            PolyBeam(beam_coeffs, spectral_index=-0.6975, ref_freq=1.0e8)
+            for ant in ants
+        ]
     else:
-        beams = [pyuvsim.analyticbeam.AnalyticBeam('gaussian', diameter=14.)
-                 for ant in ants]
+        beams = [
+            pyuvsim.analyticbeam.AnalyticBeam("gaussian", diameter=14.0) for ant in ants
+        ]
 
     # Run a simulation
     t0 = time.time()
     _sim_vis = simulate_vis(
-            ants=ant_pos,
-            fluxes=fluxes,
-            ra=ra,
-            dec=dec,
-            freqs=freqs*1e6, # MHz -> Hz
-            lsts=times,
-            beams=beams,
-            polarized=False,
-            precision=2,
-            latitude=array_latitude,
-            use_feed="x"
-        )
-    #timing_info(ftime, 0, "(0) Simulation", time.time() - t0)
-    #print("(0) Simulation", time.time() - t0)
+        ants=ant_pos,
+        fluxes=fluxes,
+        ra=ra,
+        dec=dec,
+        freqs=freqs * 1e6,  # MHz -> Hz
+        lsts=times,
+        beams=beams,
+        polarized=False,
+        precision=2,
+        latitude=array_latitude,
+        use_feed="x",
+    )
+    # timing_info(ftime, 0, "(0) Simulation", time.time() - t0)
+    # print("(0) Simulation", time.time() - t0)
 
     # Allocate computed visibilities to only the requested baselines (saves memory)
     model0 = extract_vis_from_sim(ants, antpairs, _sim_vis)
-    del _sim_vis # save some memory
+    del _sim_vis  # save some memory
 
     # Return
     ant_info = (ants, ant_pos, antpairs, ants1, ants2)
     return model0, fluxes, beams, ant_info
-
diff --git a/hydra/gain_sampler.py b/hydra/gain_sampler.py
index ed6585c..2a7ac6c 100644
--- a/hydra/gain_sampler.py
+++ b/hydra/gain_sampler.py
@@ -3,7 +3,8 @@
 import scipy.sparse
 import numpy.fft as fft
 from mpi4py.MPI import SUM as MPI_SUM
-#import pyfftw.interfaces.numpy_fft as fft
+
+# import pyfftw.interfaces.numpy_fft as fft
 
 from scipy.sparse import dok_matrix
 from .utils import flatten_vector, reconstruct_vector, freqs_times_for_worker
@@ -140,8 +141,16 @@ def apply_proj_conj(v, A_real, A_imag, model_vis, gain_shape):
 
 
 def construct_rhs_mpi(
-    comm, resid, inv_noise_var, pspec_sqrt, A_real, A_imag, model_vis, Fbasis, 
-    realisation=True, seed=None
+    comm,
+    resid,
+    inv_noise_var,
+    pspec_sqrt,
+    A_real,
+    A_imag,
+    model_vis,
+    Fbasis,
+    realisation=True,
+    seed=None,
 ):
     """
     MPI version of RHS constructor.
@@ -176,8 +185,8 @@ def construct_rhs_mpi(
     # Broadcast this random realisation
     comm.Bcast(b, root=0)
 
-    # The following quantities are calculated on each worker; each worker holds 
-    # a chunk of the data, and the calculated quantities are reduced/summed across 
+    # The following quantities are calculated on each worker; each worker holds
+    # a chunk of the data, and the calculated quantities are reduced/summed across
     # the full set of workers at the end
 
     # (Terms 1+3): S^1/2 F^dagger A^\dagger [ N^{-1} r + N^{-1/2} \omega_r ]
@@ -197,19 +206,19 @@ def construct_rhs_mpi(
         model_vis,
         gain_shape,
     )
-    
+
     # Conjugate basis
     FFc = Fbasis.conj().reshape((Fbasis.shape[0], -1))
-    
+
     # Do FT to go into Fourier space again; also apply sqrtS operator
     my_y = np.zeros_like(b)
     for k in range(Nants):
-        my_y[k,:] = pspec_sqrt * np.tensordot(FFc, 
-                                              yy[k, :, :].flatten(), 
-                                              axes=((1,), (0,)))
+        my_y[k, :] = pspec_sqrt * np.tensordot(
+            FFc, yy[k, :, :].flatten(), axes=((1,), (0,))
+        )
 
-    # Do reduce (sum) operation for all workers' my_y arrays, which contain the 
-    # projection of that worker's chunk of the result onto a chunk of the Fourier 
+    # Do reduce (sum) operation for all workers' my_y arrays, which contain the
+    # projection of that worker's chunk of the result onto a chunk of the Fourier
     # basis. The result of the reduce/sum gives the full FT over *all* data points
     total_y = np.zeros_like(my_y)
     comm.Allreduce(my_y.flatten(), total_y, op=MPI_SUM)
@@ -220,66 +229,65 @@ def construct_rhs_mpi(
     return bvec
 
 
-def apply_operator_mpi(comm, x, inv_noise_var, pspec_sqrt, A_real, A_imag, model_vis, Fbasis):
+def apply_operator_mpi(
+    comm, x, inv_noise_var, pspec_sqrt, A_real, A_imag, model_vis, Fbasis
+):
     """
     MPI version of gain linear operator.
     """
     myid = comm.Get_rank()
 
     assert inv_noise_var.shape == model_vis.shape
-    
+
     # Reshape input vector
     Nmodes, Nfreqs, Ntimes = Fbasis.shape
     Nants = A_real.shape[1]
-    #vec = x.reshape((Nants, Nmodes))
+    # vec = x.reshape((Nants, Nmodes))
     assert pspec_sqrt.shape == (Nmodes,)
 
     # Broadcast this input vector to all workers, to make sure it's synced
-    vec = np.zeros(2*Nants*Nmodes, dtype=x.dtype)
+    vec = np.zeros(2 * Nants * Nmodes, dtype=x.dtype)
     if myid == 0:
         vec[:] = x
     comm.Bcast(vec, root=0)
-    
+
     # Extract real and imaginary parts, reshape, and multiply by sqrt of prior var
-    xre = pspec_sqrt[np.newaxis,:] * vec[:vec.size//2].reshape((Nants, Nmodes))
-    xim = pspec_sqrt[np.newaxis,:] * vec[vec.size//2:].reshape((Nants, Nmodes))
-    
+    xre = pspec_sqrt[np.newaxis, :] * vec[: vec.size // 2].reshape((Nants, Nmodes))
+    xim = pspec_sqrt[np.newaxis, :] * vec[vec.size // 2 :].reshape((Nants, Nmodes))
+
     # Conjugate basis
     FFc = Fbasis.conj().reshape((Fbasis.shape[0], -1))
 
-    # The following quantities are calculated on each worker; each worker holds 
-    # a chunk of the data, and the calculated quantities are reduced/summed across 
+    # The following quantities are calculated on each worker; each worker holds
+    # a chunk of the data, and the calculated quantities are reduced/summed across
     # the full set of workers at the end
-    
+
     # Multiply Fourier x values by S^1/2 and FT
     sqrtSx = np.zeros((Nants, Fbasis.shape[1], Fbasis.shape[2]), dtype=np.complex128)
     for k in range(sqrtSx.shape[0]):
-        sqrtSx[k, :, :] = np.tensordot(Fbasis, 
-                                       xre[k, :] + 1.j*xim[k, :], 
-                                       axes=((0,), (0,)))
+        sqrtSx[k, :, :] = np.tensordot(
+            Fbasis, xre[k, :] + 1.0j * xim[k, :], axes=((0,), (0,))
+        )
     gain_shape = (Nants, Nfreqs, Ntimes)
 
     # Apply projection operator to real-space sqrt(S)-weighted x values,
     # weight by inverse noise variance, then apply (conjugated) projection
     # operator
     y = apply_proj_conj(
-                        apply_proj(sqrtSx,
-                                   A_real,
-                                   A_imag,
-                                   model_vis) \
-                        * inv_noise_var,
-                        A_real,
-                        A_imag,
-                        model_vis,
-                        gain_shape)
+        apply_proj(sqrtSx, A_real, A_imag, model_vis) * inv_noise_var,
+        A_real,
+        A_imag,
+        model_vis,
+        gain_shape,
+    )
 
     # Do inverse FT and multiply by S^1/2 again
     yy = np.zeros((Nants, Nmodes), dtype=np.complex128)
     for k in range(y.shape[0]):
-        yy[k, :] = pspec_sqrt * np.tensordot(FFc, 
-                                             y[k, :, :].flatten(), 
-                                             axes=((1,), (0,)))
-    
+        yy[k, :] = pspec_sqrt * np.tensordot(
+            FFc, y[k, :, :].flatten(), axes=((1,), (0,))
+        )
+
     # Expand into 1D vector of real and imag. parts
     my_y = np.concatenate((yy.real.flatten(), yy.imag.flatten()))
 
@@ -290,17 +298,20 @@ def apply_operator_mpi(comm, x, inv_noise_var, pspec_sqrt, A_real, A_imag, model
     # Add identity and return
     return x + total_y
 
-#--------------------------
 
-def nonfunctioning_apply_linear_op_mpi(comm, ants, antpairs, Fbasis, vec, vec_shape, 
-                                       inv_noise_var, gain_pspec_sqrt, ggV):
+# --------------------------
+
+
+def nonfunctioning_apply_linear_op_mpi(
+    comm, ants, antpairs, Fbasis, vec, vec_shape, inv_noise_var, gain_pspec_sqrt, ggV
+):
     """
-    This was an attempted MPI rewrite of the linear operator function using a 
+    This was an attempted MPI rewrite of the linear operator function using a
     different mathematical approach, but it seems to have a bug.
 
     Apply the linear system operator on each chunk.
 
-    Assume that only the root worker has the correct x vector and distribute 
+    Assume that only the root worker has the correct x vector and distribute
     it.
     """
     myid = comm.Get_rank()
@@ -309,41 +320,40 @@ def nonfunctioning_apply_linear_op_mpi(comm, ants, antpairs, Fbasis, vec, vec_sh
     # FIXME: Need to test for correctness
 
     Nants = len(ants)
-    Nmodes = Fbasis.shape[0] # Fbasis: (Ngain_modes, Nfreqs, Ntimes)
+    Nmodes = Fbasis.shape[0]  # Fbasis: (Ngain_modes, Nfreqs, Ntimes)
 
-    # coeffs of Fbasis should be coeff of real part and coeff of imag part 
+    # coeffs of Fbasis should be coeff of real part and coeff of imag part
     # (*not* real and imag parts of the Fourier coeffs)
-    Fb = Fbasis.reshape((Nmodes, -1)).T # (Nblock_freqstimes, Ngain_modes)
+    Fb = Fbasis.reshape((Nmodes, -1)).T  # (Nblock_freqstimes, Ngain_modes)
 
     # Broadcast x vector
-    #if myid != 0:
+    # if myid != 0:
     #    vec = np.zeros(vec_shape, dtype=np.float64)
-    #comm.Bcast(vec, root=0)
+    # comm.Bcast(vec, root=0)
     x = vec
 
     # Extract real and imaginary parts, reshape, and multiply by sqrt of prior var
-    xre = gain_pspec_sqrt[np.newaxis,:] * x[:x.size//2].reshape((Nants, Nmodes))
-    xim = gain_pspec_sqrt[np.newaxis,:] * x[x.size//2:].reshape((Nants, Nmodes))
+    xre = gain_pspec_sqrt[np.newaxis, :] * x[: x.size // 2].reshape((Nants, Nmodes))
+    xim = gain_pspec_sqrt[np.newaxis, :] * x[x.size // 2 :].reshape((Nants, Nmodes))
 
     # Calculate delta gain for each antenna (in freq/time space)
     # dg_block ~ (Nants, Nblockfreqs, Nblocktimes)
-    dg_block = np.tensordot(xre + 1.j*xim, Fbasis, axes=((1,), (0,)))
+    dg_block = np.tensordot(xre + 1.0j * xim, Fbasis, axes=((1,), (0,)))
 
-
-    # Maths note: We can expand the (F^T P^T N^-1 P F) operator acting on x, and 
-    # turn it into a simple sum over terms for each row of the result vector. 
-    # Here, F is the partial Fourier basis operator (from per-antenna coefficients 
-    # to per-antenna gains vs freq. and time), and P is the sparse operator that 
-    # mixes the gains for each visibility. The inverse noise term is actually the 
+    # Maths note: We can expand the (F^T P^T N^-1 P F) operator acting on x, and
+    # turn it into a simple sum over terms for each row of the result vector.
+    # Here, F is the partial Fourier basis operator (from per-antenna coefficients
+    # to per-antenna gains vs freq. and time), and P is the sparse operator that
+    # mixes the gains for each visibility. The inverse noise term is actually the
     # product with the corresponding visibility times the mean gain, i.e.
     # N^-1 = (g_i g_i^* V_ij)^T (N_ij)^-1 (g_i g_i^* V_ij)^T.
 
-    # If we evaluate this operator acting on x, we get the following expression 
+    # If we evaluate this operator acting on x, we get the following expression
     # for each (block) row:
     # y_i = F^T (Sum_j N^-1_{ij}) (F x_i) \
     #     + F^T (Sum_j N^-1_{ij} (F x_j))
-    # where i is the antenna element of this (block) row, for visibility V_ij, 
-    # and j is the second antenna in the pair, for all pairs where i is the first 
+    # where i is the antenna element of this (block) row, for visibility V_ij,
+    # and j is the second antenna in the pair, for all pairs where i is the first
     # antenna. To simplify, we can write F x_i = (delta g)_i = gg_i, to obtain
     # y_i = F^T (Sum_j N^-1_{ij}) gg_i \
     #     + F^T (Sum_j N^-1_{ij} gg_j)
@@ -353,7 +363,7 @@ def nonfunctioning_apply_linear_op_mpi(comm, ants, antpairs, Fbasis, vec, vec_sh
     ant2 = np.array([j for i, j in antpairs])
 
     # Initialise result vector for single block (i.e. for gain coeffs for single ant)
-    my_y = np.zeros((Nants, 2*Nmodes), dtype=xre.dtype) # y is real
+    my_y = np.zeros((Nants, 2 * Nmodes), dtype=xre.dtype)  # y is real
 
     # Loop over antennas in order
     for k, ant in enumerate(ants):
@@ -377,18 +387,22 @@ def nonfunctioning_apply_linear_op_mpi(comm, ants, antpairs, Fbasis, vec, vec_sh
             # Sum (V^T N^-1 V g) over baselines where i = ant
             # (Should have shape (Nants_in_group, Nfreqs, Ntimes) before sum)
             # All terms with x_ant
-            term1 = np.sum(   ggV[idxs1,:,:].conj() \
-                            * inv_noise_var[idxs1,:,:] \
-                            * ggV[idxs1,:,:] \
-                            * dg_block[k][np.newaxis,:,:],
-                            axis=0) # x_ant
+            term1 = np.sum(
+                ggV[idxs1, :, :].conj()
+                * inv_noise_var[idxs1, :, :]
+                * ggV[idxs1, :, :]
+                * dg_block[k][np.newaxis, :, :],
+                axis=0,
+            )  # x_ant
 
             # All terms with x_j (where i=ant)
-            term2c = np.sum(   ggV[idxs1,:,:].conj() \
-                             * inv_noise_var[idxs1,:,:] \
-                             * ggV[idxs1,:,:] \
-                             * dg_block[gain_idxs1].conj(),
-                             axis=0) # x_j^*
+            term2c = np.sum(
+                ggV[idxs1, :, :].conj()
+                * inv_noise_var[idxs1, :, :]
+                * ggV[idxs1, :, :]
+                * dg_block[gain_idxs1].conj(),
+                axis=0,
+            )  # x_j^*
 
         # Gains where this antenna is ant j
         if len(idxs2) > 0:
@@ -398,19 +412,22 @@ def nonfunctioning_apply_linear_op_mpi(comm, ants, antpairs, Fbasis, vec, vec_sh
 
             # Sum (V^T N^-1 V g) over baselines where j = ant
             # All terms with x_ant^*
-            term1c = np.sum(   ggV[idxs2,:,:].conj() \
-                             * inv_noise_var[idxs2,:,:] \
-                             * ggV[idxs2,:,:] \
-                             * dg_block[k][np.newaxis,:,:].conj(),
-                             axis=0) # x_ant^*
+            term1c = np.sum(
+                ggV[idxs2, :, :].conj()
+                * inv_noise_var[idxs2, :, :]
+                * ggV[idxs2, :, :]
+                * dg_block[k][np.newaxis, :, :].conj(),
+                axis=0,
+            )  # x_ant^*
 
             # All terms with x_i (where j=ant)
-            term2 = np.sum(   ggV[idxs2,:,:].conj() \
-                            * inv_noise_var[idxs2,:,:] \
-                            * ggV[idxs2,:,:] \
-                            * dg_block[gain_idxs2],
-                            axis=0) # x_i
-
+            term2 = np.sum(
+                ggV[idxs2, :, :].conj()
+                * inv_noise_var[idxs2, :, :]
+                * ggV[idxs2, :, :]
+                * dg_block[gain_idxs2],
+                axis=0,
+            )  # x_i
 
         # Apply final factor of conjugate transpose F matrix and store as result
         FFc = Fbasis.conj().reshape((Fbasis.shape[0], -1))
@@ -418,16 +435,15 @@ def nonfunctioning_apply_linear_op_mpi(comm, ants, antpairs, Fbasis, vec, vec_sh
         y1c = np.tensordot(FFc, term1c.flatten(), axes=((1,), (0,)))
         y2 = np.tensordot(FFc, term2.flatten(), axes=((1,), (0,)))
         y2c = np.tensordot(FFc, term2c.flatten(), axes=((1,), (0,)))
-        my_y[k,:Nmodes] = gain_pspec_sqrt * (y1 + y1c + y2 + y2c).real
+        my_y[k, :Nmodes] = gain_pspec_sqrt * (y1 + y1c + y2 + y2c).real
 
         # FIXME: Conjugates?
-        my_y[k,Nmodes:] = gain_pspec_sqrt * (y1 + y1c + y2 + y2c).imag
-        
+        my_y[k, Nmodes:] = gain_pspec_sqrt * (y1 + y1c + y2 + y2c).imag
 
     # Reduce all y values onto all workers
-    #total_y = np.zeros_like(my_y)
-    #comm.Allreduce(my_y.flatten(), total_y, op=MPI_SUM)
-    return x + my_y.flatten() #total_y.flatten() # add identity times input too
+    # total_y = np.zeros_like(my_y)
+    # comm.Allreduce(my_y.flatten(), total_y, op=MPI_SUM)
+    return x + my_y.flatten()  # total_y.flatten() # add identity times input too
 
 
 def legacy_apply_sqrt_pspec(sqrt_pspec, x):
@@ -457,8 +473,9 @@ def legacy_apply_sqrt_pspec(sqrt_pspec, x):
         return sqrt_pspec[np.newaxis, :, :] * x
 
 
-def legacy_apply_operator(x, inv_noise_var, pspec_sqrt, A_real, A_imag, model_vis, 
-                   reduced_idxs=None):
+def legacy_apply_operator(
+    x, inv_noise_var, pspec_sqrt, A_real, A_imag, model_vis, reduced_idxs=None
+):
     r"""
     Apply LHS operator to a vector of Fourier coefficients.
 
@@ -498,22 +515,19 @@ def legacy_apply_operator(x, inv_noise_var, pspec_sqrt, A_real, A_imag, model_vi
     # Apply projection operator to real-space sqrt(S)-weighted x values,
     # weight by inverse noise variance, then apply (conjugated) projection
     # operator
-    y = apply_proj_conj(apply_proj(sqrtSx,
-                                   A_real,
-                                   A_imag,
-                                   model_vis) \
-                          * inv_noise_var,
-                        A_real,
-                        A_imag,
-                        model_vis,
-                        gain_shape)
+    y = apply_proj_conj(
+        apply_proj(sqrtSx, A_real, A_imag, model_vis) * inv_noise_var,
+        A_real,
+        A_imag,
+        model_vis,
+        gain_shape,
+    )
 
     # Do inverse FT and multiply by S^1/2 again
     for k in range(y.shape[0]):
         y[k, :, :] = fft.fft2(y[k, :, :])
 
     return x + apply_sqrt_pspec(pspec_sqrt, y)
-        
 
 
 def legacy_construct_rhs(
@@ -546,8 +560,8 @@ def legacy_construct_rhs(
         realisation (bool):
             Whether to include Gaussian random realisation terms (True) or just
             the Wiener filter terms (False).
-        
-        
+
+
     """
     # fft: data -> fourier
     # ifft: fourier -> data
@@ -589,7 +603,7 @@ def legacy_construct_rhs(
 
     # Do FT to go into Fourier space again
     for k in range(Nants):
-        yy[k,:,:] = fft.fft2(yy[k,:,:])
+        yy[k, :, :] = fft.fft2(yy[k, :, :])
 
     # Apply sqrt(S) operator
     yy = apply_sqrt_pspec(pspec_sqrt, yy)
diff --git a/hydra/linear_solver.py b/hydra/linear_solver.py
index c975572..0d62cba 100644
--- a/hydra/linear_solver.py
+++ b/hydra/linear_solver.py
@@ -1,5 +1,4 @@
-
-from mpi4py.MPI import SUM as MPI_SUM 
+from mpi4py.MPI import SUM as MPI_SUM
 from mpi4py.MPI import LAND as MPI_LAND
 import numpy as np
 
@@ -7,23 +6,23 @@
 def matvec_mpi(comm_row, mat_block, vec_block):
     """
     Do matrix-vector product for a row of a block matrix.
-    
-    Each block in the matrix row is multiplied by the corresponding row block 
-    of the vector. The result on each worker is then summed together to give 
+
+    Each block in the matrix row is multiplied by the corresponding row block
+    of the vector. The result on each worker is then summed together to give
     the result for the corresponding row of the result vector.
-    
-    All workers in the row will posses the result for the same row of the 
+
+    All workers in the row will posses the result for the same row of the
     result vector.
-    
+
     For example, for the first row of this (block) linear system:
     ( A B C )     ( x )     ( r0 )
     ( D E F )  .  ( y )  =  ( r1 )
     ( G H I )     ( z )     ( r2 )
-    
-    workers 0, 1, and 2 will compute Ax, By, and Cz respectively. They will 
-    then collectively sum over their results to obtain `r0 = Ax + By + Cz`. The 
+
+    workers 0, 1, and 2 will compute Ax, By, and Cz respectively. They will
+    then collectively sum over their results to obtain `r0 = Ax + By + Cz`. The
     three workers will all possess copies of `r0`.
-    
+
     Parameters:
         comm_row (MPI.Intracomm):
             MPI group communicator for a row of the block matrix.
@@ -31,14 +30,14 @@ def matvec_mpi(comm_row, mat_block, vec_block):
             Block of the matrix belonging to this worker.
         vec_block (array_like):
             Block of the vector belonging to this worker.
-    
+
     Returns:
         res_block (array_like):
             Block of the result vector corresponding to this row.
     """
     # Do matrix-vector product for the available blocks
     y = mat_block @ vec_block
-    
+
     # Do reduce to all members of this column group
     ytot = np.zeros_like(y)
     comm_row.Allreduce(y, ytot, op=MPI_SUM)
@@ -47,105 +46,107 @@ def matvec_mpi(comm_row, mat_block, vec_block):
 
 def setup_mpi_blocks(comm, matrix_shape, split=1):
     """
-    Set up a scheme for dividing the linear system into blocks. This function 
-    determines the number and size of the blocks, creates a map between MPI 
-    workers and blocks, and sets up some MPI communicator groups that are 
+    Set up a scheme for dividing the linear system into blocks. This function
+    determines the number and size of the blocks, creates a map between MPI
+    workers and blocks, and sets up some MPI communicator groups that are
     needed by the CG solver to communicate intermediate results.
-    
-    The linear system matrix operator is assumed to be square, and the blocks 
-    must also be square. The blocks will be zero-padded at the edges if the 
+
+    The linear system matrix operator is assumed to be square, and the blocks
+    must also be square. The blocks will be zero-padded at the edges if the
     operator cannot be evenly divided into the blocks.
-    
+
     Parameters:
         comm (MPI.Communicator):
             MPI communicator object for all active workers.
         matrix_shape (tuple of int):
-            The shape of the linear operator matrix that is to be divided into 
+            The shape of the linear operator matrix that is to be divided into
             blocks.
         split (int):
-            How many rows and columns to split the matrix into. For instance, 
-            `split = 2` will split the matrix into 2 rows and 2 columns, for a 
+            How many rows and columns to split the matrix into. For instance,
+            `split = 2` will split the matrix into 2 rows and 2 columns, for a
             total of 4 blocks.
-    
+
     Returns:
         comm_groups (tuple of MPI.Intracomm):
-            Group communicators for the blocks (active, row, col, diag). 
-            
-            These correspond to the MPI workers that are active, and the ones 
-            for each row, each column, and along the diagonal of the block 
+            Group communicators for the blocks (active, row, col, diag).
+
+            These correspond to the MPI workers that are active, and the ones
+            for each row, each column, and along the diagonal of the block
             structure, respectively.
-            
-            Each worker will return its own set of communicators (e.g. for the 
-            row or column it belongs to). Where it is not a member of a 
+
+            Each worker will return its own set of communicators (e.g. for the
+            row or column it belongs to). Where it is not a member of a
             relevant group, `None` will be returned instead.
         block_map (dict):
-            Dictionary of tuples, one for each worker in the `active` 
-            communicator group, with the row and column ID of the block that it 
+            Dictionary of tuples, one for each worker in the `active`
+            communicator group, with the row and column ID of the block that it
             is managing.
         block_shape (tuple of int):
-            Shape of the square blocks that the full matrix operator (and RHS 
+            Shape of the square blocks that the full matrix operator (and RHS
             vector) should be split up into. These will be square.
     """
-    assert len(matrix_shape) == 2, \
-        "'matrix_shape' must be a tuple with 2 entries"
-    assert matrix_shape[0] == matrix_shape[1], \
-        "Only square matrices are currently supported"
+    assert len(matrix_shape) == 2, "'matrix_shape' must be a tuple with 2 entries"
+    assert (
+        matrix_shape[0] == matrix_shape[1]
+    ), "Only square matrices are currently supported"
     myid = comm.Get_rank()
     nworkers = comm.Get_size()
-    
+
     # Check that enough workers are available
     nblocks = split * split
     assert nworkers >= nblocks, "Specified more blocks than workers"
     workers = np.arange(nblocks).reshape((split, split))
-    
+
     # Handle workers that do not have an assigned block
     if myid >= nblocks:
         return None, {}, {}
-    
+
     # Construct map of block row/column IDs vs worker IDs
     block_map = {}
     for w in workers.flatten():
         _row, _col = np.where(workers == w)
         block_map[w] = (_row[0], _col[0])
     myrow, mycol = block_map[myid]
-    
+
     # Setup communicator groups for each row, columns, and the diagonals
     grp_active = workers.flatten()
-    grp_row = workers[myrow,:]
-    grp_col = workers[:,mycol]
+    grp_row = workers[myrow, :]
+    grp_col = workers[:, mycol]
     grp_diag = np.diag(workers)
-    comm_active = comm.Create( comm.group.Incl(grp_active) )
-    comm_row = comm.Create( comm.group.Incl(grp_row) )
-    comm_col = comm.Create( comm.group.Incl(grp_col) )
-    comm_diag = comm.Create( comm.group.Incl(grp_diag) )
-    
+    comm_active = comm.Create(comm.group.Incl(grp_active))
+    comm_row = comm.Create(comm.group.Incl(grp_row))
+    comm_col = comm.Create(comm.group.Incl(grp_col))
+    comm_diag = comm.Create(comm.group.Incl(grp_diag))
+
     # Calculate block size (all blocks must have the same shape, so some zero-
-    # padding will be done if the matrix operator shape is not exactly 
+    # padding will be done if the matrix operator shape is not exactly
     # divisible by 'split')
-    block_rows = int(np.ceil(matrix_shape[0] / split)) # rows per block
-    block_cols = int(np.ceil(matrix_shape[1] / split)) # cols per block
+    block_rows = int(np.ceil(matrix_shape[0] / split))  # rows per block
+    block_cols = int(np.ceil(matrix_shape[1] / split))  # cols per block
     block_shape = (block_rows, block_cols)
-    assert block_rows == block_cols, \
-        "Current implementation assumes that blocks are square"
-    
+    assert (
+        block_rows == block_cols
+    ), "Current implementation assumes that blocks are square"
+
     comms = (comm_active, comm_row, comm_col, comm_diag)
     return comms, block_map, block_shape
 
 
-def collect_linear_sys_blocks(comm, block_map, block_shape, Amat=None, 
-                              bvec=None, verbose=False):
+def collect_linear_sys_blocks(
+    comm, block_map, block_shape, Amat=None, bvec=None, verbose=False
+):
     """
     Send LHS operator matrix and RHS vector blocks to assigned workers.
-    
+
     Parameters:
         comm (MPI.Communicator):
             MPI communicator object for all active workers.
         block_map (dict):
-            Dictionary of tuples, one for each worker in the `active` 
-            communicator group, with the row and column ID of the block that it 
+            Dictionary of tuples, one for each worker in the `active`
+            communicator group, with the row and column ID of the block that it
             is managing.
         block_shape (tuple of int):
-            Shape of the square blocks that the full matrix operator (and RHS 
+            Shape of the square blocks that the full matrix operator (and RHS
             vector) should be split up into. These must be square.
         Amat (array_like):
             The full LHS matrix operator, which will be split into blocks.
@@ -153,246 +154,263 @@ def collect_linear_sys_blocks(comm, block_map, block_shape, Amat=None,
             The full right-hand side vector, which will be split into blocks.
         verbose (bool):
             If `True`, print status messages when MPI communication is complete.
-    
+
     Returns:
         my_Amat (array_like):
-            The single block of the matrix operator belonging to this worker. 
-            It will have shape `block_shape`. If the matrix operator cannot be 
-            exactly divided into same-sized blocks, the blocks at the far edges 
+            The single block of the matrix operator belonging to this worker.
+            It will have shape `block_shape`. If the matrix operator cannot be
+            exactly divided into same-sized blocks, the blocks at the far edges
             will be zero-padded. Returns `None` if worker is not active.
         my_bvec (array_like):
-            The single block of the RHS vector belonging to this worker. Note 
-            that workers in the same column have the same block. Returns `None` 
+            The single block of the RHS vector belonging to this worker. Note
+            that workers in the same column have the same block. Returns `None`
             if worker is not active.
     """
     myid = comm.Get_rank()
     dtype = bvec.dtype
-    
+
     # Determine whether this worker is participating
     workers_used = np.array(list(block_map.keys()))
     workers_used.sort()
     if myid not in workers_used:
         return None, None
-    
+
     # Initialise blocks of A matrix and b vector
     block_rows, block_cols = block_shape
     my_Amat = np.zeros((block_rows, block_cols), dtype=dtype)
     my_bvec = np.zeros((block_rows,), dtype=dtype)
-    
+
     # Send blocks from root worker
     if myid == 0:
         reqs = []
         for w in workers_used:
-            
+
             # Get row and column indices for this worker
             wrow, wcol = block_map[w]
-            
+
             # Start and end indices of block (handles edges)
-            ii = wrow*block_rows
-            jj = wcol*block_cols
+            ii = wrow * block_rows
+            jj = wcol * block_cols
             iip = ii + block_rows
             jjp = jj + block_cols
-            if iip  > Amat.shape[0]: iip = Amat.shape[0]
-            if jjp  > Amat.shape[1]: jjp = Amat.shape[1]
-            
+            if iip > Amat.shape[0]:
+                iip = Amat.shape[0]
+            if jjp > Amat.shape[1]:
+                jjp = Amat.shape[1]
+
             if w == 0:
                 # Block belongs to root worker
-                my_Amat[:iip-ii,:jjp-jj] = Amat[ii:iip, jj:jjp]
-                my_bvec[:jjp-jj] = bvec[jj:jjp]
-                
+                my_Amat[: iip - ii, : jjp - jj] = Amat[ii:iip, jj:jjp]
+                my_bvec[: jjp - jj] = bvec[jj:jjp]
+
             else:
                 # Send blocks to other worker
                 # Handles zero-padding of blocks at edge of matrix
                 # FIXME: Do we have to copy here, to get contiguous memory?
                 Amat_buf = np.zeros_like(my_Amat)
                 bvec_buf = np.zeros_like(my_bvec)
-                
-                Amat_buf[:iip-ii,:jjp-jj] = Amat[ii:iip, jj:jjp]
-                bvec_buf[:jjp-jj] = bvec[jj:jjp]
-                
+
+                Amat_buf[: iip - ii, : jjp - jj] = Amat[ii:iip, jj:jjp]
+                bvec_buf[: jjp - jj] = bvec[jj:jjp]
+
                 # The flattened Amat_buf array is reshaped into 2D when received
                 comm.Send(Amat_buf.flatten().copy(), dest=w)
                 comm.Send(bvec_buf, dest=w)
-                
+
         if verbose:
             print("All send operations completed.")
     else:
         # Receive this worker's assigned blocks of Amat and bvec
         comm.Recv(my_Amat, source=0)
         comm.Recv(my_bvec, source=0)
-        
+
         if verbose:
             print("Worker %d finished receive" % myid)
-    
+
     return my_Amat, my_bvec
-    
 
-def cg_mpi(comm_groups, Amat_block, bvec_block, vec_size, block_map, 
-           maxiters=1000, abs_tol=1e-8):
+
+def cg_mpi(
+    comm_groups,
+    Amat_block,
+    bvec_block,
+    vec_size,
+    block_map,
+    maxiters=1000,
+    abs_tol=1e-8,
+):
     """
-    A distributed CG solver for linear systems with square matrix operators. 
-    The linear operator matrix is split into square blocks, each of which is 
+    A distributed CG solver for linear systems with square matrix operators.
+    The linear operator matrix is split into square blocks, each of which is
     handled by a single worker.
-    
+
     Parameters:
         comm_groups (tuple of MPI.Intracomm):
-            Group communicators for the blocks (active, row, col, diag). 
-            
-            These are set up by `setup_mpi_blocks`, and correspond to the MPI 
-            workers that are active, the ones for each row, each column, and 
+            Group communicators for the blocks (active, row, col, diag).
+
+            These are set up by `setup_mpi_blocks`, and correspond to the MPI
+            workers that are active, the ones for each row, each column, and
             along the diagonal of the block structure, respectively.
-            
-            If `None`, this is assumed to be an inactive worker and nothing is 
+
+            If `None`, this is assumed to be an inactive worker and nothing is
             done.
         Amat_block (array_like):
             The block of the matrix operator belonging to this worker.
         bvec_block (array_like):
-            The block of the right-hand side vector corresponding to this 
+            The block of the right-hand side vector corresponding to this
             worker's matrix operator block.
         vec_size (int):
             The size of the total result vector, across all blocks.
         block_map (dict):
-            Dictionary of tuples, one for each worker in the `active` 
-            communicator group, with the row and column ID of the block that it 
+            Dictionary of tuples, one for each worker in the `active`
+            communicator group, with the row and column ID of the block that it
             is managing.
         maxiters (int):
-            Maximum number of iterations of the solver to perform before 
+            Maximum number of iterations of the solver to perform before
             returning.
         abs_tol (float):
-            Absolute tolerance on each element of the residual. Once this 
-            tolerance has been reached for all entries of the residual vector, 
+            Absolute tolerance on each element of the residual. Once this
+            tolerance has been reached for all entries of the residual vector,
             the solution is considered to have converged.
-    
+
     Returns:
         x (array_like):
-            Solution vector for the full system. Only workers on the diagonal 
+            Solution vector for the full system. Only workers on the diagonal
             have the correct solution vector; other workers will return `None`.
     """
     if comm_groups is None:
         # FIXME: Need to fix this so non-active workers are ignored without hanging
         return None
-        
+
     comm_active, comm_row, comm_col, comm_diag = comm_groups
-    #grp_active, grp_row, grp_col, grp_diag = groups
+    # grp_active, grp_row, grp_col, grp_diag = groups
     myid = comm_active.Get_rank()
     myrow, mycol = block_map[myid]
-    
+
     # Initialise solution vector
     x_block = np.zeros_like(bvec_block)
-    
+
     # Calculate initial residual (if x0 is not 0, need to rewrite this, i.e.
     # be careful with x_block ordering)
-    r_block = bvec_block[:] #- matvec_mpi(comm_row, Amat_block, x_block)
+    r_block = bvec_block[:]  # - matvec_mpi(comm_row, Amat_block, x_block)
     pvec_block = r_block[:]
-    
+
     # Iterate
     niter = 0
     finished = False
     while niter < maxiters and not finished:
-        
+
         # Check convergence criterion from all workers
         converged = np.all(np.abs(r_block) < abs_tol)
-        
+
         # Check if convergence is reached
         # (reduce with logical-AND operation)
         converged = comm_active.allreduce(converged, op=MPI_LAND)
         if converged:
             finished = True
             break
-        
+
         # Distribute pvec_block to all workers (identical in each column)
-        # NOTE: Assumes that the rank IDs in comm_col are in the same order as 
-        # for the original comm_world communicator, in which case the rank with 
-        # ID = mycol will be the one on the diagonal that we want to broadcast 
+        # NOTE: Assumes that the rank IDs in comm_col are in the same order as
+        # for the original comm_world communicator, in which case the rank with
+        # ID = mycol will be the one on the diagonal that we want to broadcast
         # the up-to-date value of pvec_block from
         if myrow != mycol:
-            pvec_block *= 0.
+            pvec_block *= 0.0
         comm_col.Bcast(pvec_block, root=mycol)
-        
-        # Calculate matrix operator product with p-vector (returns result for 
+
+        # Calculate matrix operator product with p-vector (returns result for
         # this row)
         A_dot_p = matvec_mpi(comm_row, Amat_block, pvec_block)
-        
-        # Only workers with mycol == myrow will give correct updates, so only 
+
+        # Only workers with mycol == myrow will give correct updates, so only
         # calculate using those
         if mycol == myrow:
             # Calculate residual norm, summed across all (diagonal) workers
             r_dot_r = comm_diag.allreduce(np.dot(r_block.T, r_block), op=MPI_SUM)
-            
+
             # Calculate quadratic, summed across all (diagonal) workers
             pAp = comm_diag.allreduce(np.dot(pvec_block.T, A_dot_p), op=MPI_SUM)
-            
+
             # Calculate alpha (valid on all diagonal workers)
             alpha = r_dot_r / pAp
-            
+
             # Update solution vector and residual for this worker
             x_block = x_block + alpha * pvec_block
             r_block = r_block - alpha * A_dot_p
-            
+
             # Calculate updated residual norm
-            rnew_dot_rnew = comm_diag.allreduce(np.dot(r_block.T, r_block), 
-                                                op=MPI_SUM)
-            
+            rnew_dot_rnew = comm_diag.allreduce(np.dot(r_block.T, r_block), op=MPI_SUM)
+
             # Calculate beta (valid on all diagonal workers)
             beta = rnew_dot_rnew / r_dot_r
-            
+
             # Update pvec_block (valid on all diagonal workers)
             pvec_block = r_block + beta * pvec_block
-        
+
         comm_active.barrier()
-        
+
         # Increment iteration
         niter += 1
-    
+
     # Gather all the blocks into a single array (on diagonal workers only)
     if myrow == mycol:
         x_all = np.zeros((vec_size), dtype=x_block.dtype)
-        x_all_blocks = np.zeros((x_block.size * comm_diag.Get_size()), 
-                                dtype=x_block.dtype) # needed for zero-padding
+        x_all_blocks = np.zeros(
+            (x_block.size * comm_diag.Get_size()), dtype=x_block.dtype
+        )  # needed for zero-padding
         comm_diag.Allgather(x_block, x_all_blocks)
-        
+
         # Remove zero padding if necessary
         x_all[:] = x_all_blocks[:vec_size]
     else:
         x_all = None
-    
+
     comm_active.barrier()
     return x_all
 
 
-def cg(Amat, bvec, maxiters=1000, abs_tol=1e-8, use_norm_tol=False, 
-       x0=None, linear_op=None, comm=None):
+def cg(
+    Amat,
+    bvec,
+    maxiters=1000,
+    abs_tol=1e-8,
+    use_norm_tol=False,
+    x0=None,
+    linear_op=None,
+    comm=None,
+):
     """
-    Simple Conjugate Gradient solver that operates in serial. This uses the 
-    same algorithm as `cg_mpi()` and so can be used for testing/comparison of 
+    Simple Conjugate Gradient solver that operates in serial. This uses the
+    same algorithm as `cg_mpi()` and so can be used for testing/comparison of
     results.
-    
+
     Note that this function will still permit threading used within numpy.
-    
+
     Parameters:
         Amat (array_like):
             Linear operator matrix.
         bvec (array_like):
             Right-hand side vector.
         maxiters (int):
-            Maximum number of iterations of the solver to perform before 
+            Maximum number of iterations of the solver to perform before
             returning.
         abs_tol (float):
-            Absolute tolerance on each element of the residual. Once this 
-            tolerance has been reached for all entries of the residual vector, 
-            the solution is considered to have converged.    
+            Absolute tolerance on each element of the residual. Once this
+            tolerance has been reached for all entries of the residual vector,
+            the solution is considered to have converged.
         use_norm_tol (bool):
-            Whether to use the tolerance on each element (as above), or an 
+            Whether to use the tolerance on each element (as above), or an
             overall tolerance on the norm of the residual.
         x0 (array_like):
-            Initial guess for the solution vector. Will be set to zero 
+            Initial guess for the solution vector. Will be set to zero
             otherwise.
         linear_op (func):
-            If specified, this function will be used to operate on vectors, 
+            If specified, this function will be used to operate on vectors,
             instead of the Amat matrix. Must have call signature `func(x)`.
         comm (MPI communicator):
-            If specified, the CG solver will be run only on the root worker, 
-            but the 
+            If specified, the CG solver will be run only on the root worker,
+            but the
 
     Returns:
         x (array_like):
@@ -406,7 +424,7 @@ def cg(Amat, bvec, maxiters=1000, abs_tol=1e-8, use_norm_tol=False,
     # Use Amat as the linear operator if function not specified
     if linear_op is None:
         linear_op = lambda v: Amat @ v
-    
+
     # Initialise solution vector
     if x0 is None:
         x = np.zeros_like(bvec)
@@ -416,18 +434,18 @@ def cg(Amat, bvec, maxiters=1000, abs_tol=1e-8, use_norm_tol=False,
         x = x0.copy()
 
     # Calculate initial residual
-    # NOTE: linear_op may have internal MPI calls; we assume that it 
-    # handles synchronising its input itself, but that only the root 
+    # NOTE: linear_op may have internal MPI calls; we assume that it
+    # handles synchronising its input itself, but that only the root
     # worker receives a correct return value.
 
     r = bvec - linear_op(x)
     pvec = r[:]
-    
+
     # Blocks indexed by i,j: y = A . x = Sum_j A_ij b_j
     niter = 0
     finished = False
     while niter < maxiters and not finished:
-        
+
         try:
             if myid == 0:
                 # Root worker checks for convergence
@@ -441,7 +459,11 @@ def cg(Amat, bvec, maxiters=1000, abs_tol=1e-8, use_norm_tol=False,
                         finished = True
 
             # Broadcast finished flag to all workers (need to use a non-immutable type)
-            finished_arr = np.array([finished,])
+            finished_arr = np.array(
+                [
+                    finished,
+                ]
+            )
             if comm is not None:
                 comm.Bcast(finished_arr, root=0)
                 finished = bool(finished_arr[0])
@@ -450,8 +472,8 @@ def cg(Amat, bvec, maxiters=1000, abs_tol=1e-8, use_norm_tol=False,
 
             # Do CG iteration
             r_dot_r = np.dot(r.T, r)
-            A_dot_p = linear_op(pvec) # root worker will broadcast correct pvec
-            
+            A_dot_p = linear_op(pvec)  # root worker will broadcast correct pvec
+
             # Only root worker needs to do these updates; other workers idle
             if myid == 0:
                 pAp = pvec.T @ A_dot_p
@@ -466,12 +488,12 @@ def cg(Amat, bvec, maxiters=1000, abs_tol=1e-8, use_norm_tol=False,
 
             # Update pvec on all workers
             comm.Bcast(pvec, root=0)
-            
+
             # Increment iteration
             niter += 1
         except:
             raise
-    
+
     if comm is not None:
         comm.barrier()
 
diff --git a/hydra/plot.py b/hydra/plot.py
index c7d59c7..789ddbe 100644
--- a/hydra/plot.py
+++ b/hydra/plot.py
@@ -1,38 +1,41 @@
-
-
-def plot_beam_cross(beam_coeffs, ant_ind, iter, output_dir, tag='', type='cross'):
+def plot_beam_cross(beam_coeffs, ant_ind, iter, output_dir, tag="", type="cross"):
     # Shape ncoeffs, Nfreqs, Nant -- just use a ref freq
     coeff_use = beam_coeffs[:, 0, :]
     Nants = coeff_use.shape[1]
-    if type == 'cross':
-        fig, ax = plt.subplots(figsize=(16, 9), nrows=Nants, ncols=Nants,
-                               subplot_kw={'projection': 'polar'})
+    if type == "cross":
+        fig, ax = plt.subplots(
+            figsize=(16, 9),
+            nrows=Nants,
+            ncols=Nants,
+            subplot_kw={"projection": "polar"},
+        )
         for ant_ind1 in range(Nants):
 
             beam_use1 = bess_matr_fit @ coeff_use[:, ant_ind1, 0, 0]
             for ant_ind2 in range(Nants):
                 beam_use2 = bess_matr_fit @ (coeff_use[:, ant_ind2, 0, 0])
-                beam_cross = (beam_use1 * beam_use2.conj())
+                beam_cross = beam_use1 * beam_use2.conj()
                 if ant_ind1 >= ant_ind2:
-                    ax[ant_ind1, ant_ind2].pcolormesh(PHI,
-                                                      RHO,
-                                                      np.abs(beam_cross),
-                                                      vmin=0, vmax=1)
+                    ax[ant_ind1, ant_ind2].pcolormesh(
+                        PHI, RHO, np.abs(beam_cross), vmin=0, vmax=1
+                    )
                 else:
-                    ax[ant_ind1, ant_ind2].pcolormesh(PHI,
-                                                      RHO,
-                                                      np.angle(beam_cross),
-                                                      vmin=-np.pi,
-                                                      vmax=np.pi,
-                                                      cmap='twilight')
+                    ax[ant_ind1, ant_ind2].pcolormesh(
+                        PHI,
+                        RHO,
+                        np.angle(beam_cross),
+                        vmin=-np.pi,
+                        vmax=np.pi,
+                        cmap="twilight",
+                    )
     else:
-        fig, ax = plt.subplots(ncols=2, subplot_kw={'projection': 'polar'})
-        beam_use = (bess_matr_fit@(coeff_use[:, ant_ind, 0, 0]))
-        ax[0].pcolormesh(PHI, RHO, np.abs(beam_use),
-                         vmin=0, vmax=1)
-        ax[1].pcolormesh(PHI, RHO, np.angle(beam_use),
-                         vmin=-np.pi, vmax=np.pi, cmap='twilight')
+        fig, ax = plt.subplots(ncols=2, subplot_kw={"projection": "polar"})
+        beam_use = bess_matr_fit @ (coeff_use[:, ant_ind, 0, 0])
+        ax[0].pcolormesh(PHI, RHO, np.abs(beam_use), vmin=0, vmax=1)
+        ax[1].pcolormesh(
+            PHI, RHO, np.angle(beam_use), vmin=-np.pi, vmax=np.pi, cmap="twilight"
+        )
 
     fig.savefig(f"{output_dir}/beam_plot_ant_{ant_ind}_iter_{iter}_{type}_{tag}.png")
     plt.close(fig)
-    return
\ No newline at end of file
+    return
diff --git a/hydra/pspec_sampler.py b/hydra/pspec_sampler.py
index 81b7292..eaa3258 100644
--- a/hydra/pspec_sampler.py
+++ b/hydra/pspec_sampler.py
@@ -5,7 +5,7 @@
 from scipy.stats import invgamma
 from scipy.optimize import minimize, Bounds
 
-#from multiprocess import Pool
+# from multiprocess import Pool
 from . import utils
 import os, time
 
@@ -151,8 +151,7 @@ def gcr_fgmodes_1d(
 
 
 def gcr_fgmodes(
-    vis, w, matrices, fgmodes, f0=None, nproc=1, map_estimate=False,
-    verbose=False
+    vis, w, matrices, fgmodes, f0=None, nproc=1, map_estimate=False, verbose=False
 ):
     """
     Perform the GCR step on all time samples, using parallelisation if
@@ -199,18 +198,20 @@ def gcr_fgmodes(
     if verbose:
         st = time.time()
     with Pool(nproc) as pool:
-        samples, residuals, info = zip(*pool.map(
-            lambda idx: gcr_fgmodes_1d(
-                vis=vis[idx],
-                w=w,
-                matrices=matrices,
-                fgmodes=fgmodes,
-                f0=f0,
-                map_estimate=map_estimate,
-                verbose=verbose
-            ),
-            idxs,
-        ))
+        samples, residuals, info = zip(
+            *pool.map(
+                lambda idx: gcr_fgmodes_1d(
+                    vis=vis[idx],
+                    w=w,
+                    matrices=matrices,
+                    fgmodes=fgmodes,
+                    f0=f0,
+                    map_estimate=map_estimate,
+                    verbose=verbose,
+                ),
+                idxs,
+            )
+        )
     samples = np.array(samples).reshape((vis.shape[0], -1))
     residuals = np.array(residuals)
     info = np.array(info)
@@ -231,19 +232,19 @@ def covariance_from_pspec(ps, fourier_op):
     Nfreqs = ps.size
     Csigfft = np.zeros((Nfreqs, Nfreqs), dtype=complex)
     Csigfft[np.diag_indices(Nfreqs)] = ps
-    C = (fourier_op.T.conj() @ Csigfft @ fourier_op)
+    C = fourier_op.T.conj() @ Csigfft @ fourier_op
     return C
 
 
 def build_matrices(Nparams, flags, signal_S, Ninv, fgmodes):
     """
     Calculate matrices and build A in Ax=b for the GCR step.
-    
+
     Parameters:
         Nparams (int):
             Number of model parameters.
         flags (array_like):
-            Array of flags (1 for unflagged, 0 for flagged), with shape 
+            Array of flags (1 for unflagged, 0 for flagged), with shape
             `(Nfreqs,)`.
         signal_S (array_like):
             Current value of the EoR signal frequency-frequency covariance.
@@ -255,14 +256,14 @@ def build_matrices(Nparams, flags, signal_S, Ninv, fgmodes):
             Foreground mode array, of shape (Nfreqs, Nmodes). This should be
             derived from a PCA decomposition of a model foreground covariance
             matrix or similar.
-    
+
     Returns:
         matrices (list of array_like):
             List containing necessary GCR operators (`matrices[0]`) and the
             linear operator A in the GCR Ax=b solve step.
     """
     Nfreqs = signal_S.shape[0]
-    
+
     # Construct matrix structure
     matrices = [0, 0]
     matrices[0] = np.zeros((4, Nfreqs, Nfreqs), dtype=complex)
@@ -283,13 +284,21 @@ def build_matrices(Nparams, flags, signal_S, Ninv, fgmodes):
 
     matrices[1][0] = A
     matrices[1][1] = np.linalg.pinv(A)  # pseudo-inverse, to be used as a preconditioner
-    
+
     return matrices
 
 
 def gibbs_step_fgmodes(
-    vis, flags, signal_S, fgmodes, Ninv, ps_prior=None, f0=None, nproc=1,
-    map_estimate=False, verbose=False
+    vis,
+    flags,
+    signal_S,
+    fgmodes,
+    Ninv,
+    ps_prior=None,
+    f0=None,
+    nproc=1,
+    map_estimate=False,
+    verbose=False,
 ):
     """
     Perform a single Gibbs iteration for a Gibbs sampling scheme using a foreground model
@@ -300,7 +309,7 @@ def gibbs_step_fgmodes(
             Array of complex visibilities for a single baseline, of shape
             `(Ntimes, Nfreqs)`.
         flags (array_like):
-            Array of flags (1 for unflagged, 0 for flagged), with shape 
+            Array of flags (1 for unflagged, 0 for flagged), with shape
             `(Nfreqs,)`.
         signal_S (array_like):
             Current value of the EoR signal frequency-frequency covariance.
@@ -348,8 +357,14 @@ def gibbs_step_fgmodes(
 
     # (1) Solve GCR equation to get EoR signal and foreground amplitude realisations
     cr = gcr_fgmodes(
-        vis=vis, w=flags, matrices=matrices, fgmodes=fgmodes, f0=f0, nproc=nproc,
-        map_estimate=map_estimate, verbose=verbose
+        vis=vis,
+        w=flags,
+        matrices=matrices,
+        fgmodes=fgmodes,
+        f0=f0,
+        nproc=nproc,
+        map_estimate=map_estimate,
+        verbose=verbose,
     )
 
     # Extract separate signal and FG parts from the solution
@@ -357,12 +372,14 @@ def gibbs_step_fgmodes(
     fg_amps = cr[:, -fgmodes.shape[1] :]
 
     # Full model of data is sum of EoR (GCR) + FG model
-    model = signal_cr + fg_amps @ fgmodes.T  # np.einsum('ijk,lk->ijl', fg_amps, fgmodes)
+    model = (
+        signal_cr + fg_amps @ fgmodes.T
+    )  # np.einsum('ijk,lk->ijl', fg_amps, fgmodes)
     # Chi-squared is computed as the sum of ( |data - model| / noise )^2,
     # i.e. as a sum of standard normal random variables.
     # FIXME: this will need to be changed to account for time-dependent
     # flags (i.e. when we have a different N per time).
-    chisq = np.abs(vis - model)**2 * Ninv.diagonal()[None, :]
+    chisq = np.abs(vis - model) ** 2 * Ninv.diagonal()[None, :]
     if verbose:
         chisq_mean = chisq[:, flags].mean()
         if chisq_mean > 10:
@@ -383,21 +400,23 @@ def gibbs_step_fgmodes(
     # WARNING: np.linalg.inv should be avoided for general, dense matrices.
     # S_sample should be diagonally dominant and thus this should be okay.
     Sinv = np.linalg.inv(S_sample)
-    ln_post = np.sum(np.diagonal(
-        -(
-            (vis - model)[:, flags].conj()
-            @ Ninv[flags][:, flags]
-            @ (vis - model)[:, flags].T
-        )
-        - (
-            signal_cr[:, flags].conj()
-            @ Sinv[flags][:, flags]
-            @ signal_cr[:, flags].T
+    ln_post = np.sum(
+        np.diagonal(
+            -(
+                (vis - model)[:, flags].conj()
+                @ Ninv[flags][:, flags]
+                @ (vis - model)[:, flags].T
+            )
+            - (
+                signal_cr[:, flags].conj()
+                @ Sinv[flags][:, flags]
+                @ signal_cr[:, flags].T
+            )
         )
-    ))
+    )
     ln_post = ln_post.real
     if verbose:
         print(f"{ln_post:<12.1f}")
 
     # Return samples
-    return signal_cr, S_sample, ps_sample, fg_amps, chisq, ln_post
\ No newline at end of file
+    return signal_cr, S_sample, ps_sample, fg_amps, chisq, ln_post
diff --git a/hydra/ptsrc_sampler.py b/hydra/ptsrc_sampler.py
index 4cb7bf4..560ba3c 100644
--- a/hydra/ptsrc_sampler.py
+++ b/hydra/ptsrc_sampler.py
@@ -1,4 +1,3 @@
-
 from mpi4py.MPI import SUM as MPI_SUM
 import numpy as np
 import numpy.fft as fft
@@ -10,23 +9,24 @@
 from .vis_simulator import simulate_vis_per_source
 
 
-
-def precompute_mpi(comm,
-                   ants, 
-                   antpairs, 
-                   freq_chunk, 
-                   time_chunk,
-                   proj_chunk,
-                   data_chunk,
-                   inv_noise_var_chunk,
-                   gain_chunk, 
-                   amp_prior_std, 
-                   realisation=True):
+def precompute_mpi(
+    comm,
+    ants,
+    antpairs,
+    freq_chunk,
+    time_chunk,
+    proj_chunk,
+    data_chunk,
+    inv_noise_var_chunk,
+    gain_chunk,
+    amp_prior_std,
+    realisation=True,
+):
     """
-    Precompute the projection operator and matrix operator in parallel. 
+    Precompute the projection operator and matrix operator in parallel.
 
-    The projection operator is computed in chunks in time and frequency. 
-    The overall matrix operator can be computed by summing the matrix 
+    The projection operator is computed in chunks in time and frequency.
+    The overall matrix operator can be computed by summing the matrix
     operator for the time and frequency chunks.
     """
     myid = comm.Get_rank()
@@ -34,7 +34,7 @@ def precompute_mpi(comm,
     # Check input dimensions
     assert data_chunk.shape == (len(antpairs), freq_chunk.size, time_chunk.size)
     assert data_chunk.shape == inv_noise_var_chunk.shape
-    proj = proj_chunk.copy() # make a copy so we don't alter the original proj!
+    proj = proj_chunk.copy()  # make a copy so we don't alter the original proj!
 
     # FIXME: Check for unused args!
 
@@ -43,77 +43,95 @@ def precompute_mpi(comm,
         ant1, ant2 = bl
         i1 = np.where(ants == ant1)[0][0]
         i2 = np.where(ants == ant2)[0][0]
-        proj[k,:,:,:] *= gain_chunk[i1,:,:,np.newaxis] \
-                       * gain_chunk[i2,:,:,np.newaxis].conj()
+        proj[k, :, :, :] *= (
+            gain_chunk[i1, :, :, np.newaxis] * gain_chunk[i2, :, :, np.newaxis].conj()
+        )
 
     # (2) Precompute linear system operator
     nsrcs = proj.shape[-1]
     my_linear_op = np.zeros((nsrcs, nsrcs), dtype=proj.real.dtype)
 
     # inv_noise_var has shape (Nbls, Nfreqs, Ntimes)
-    v_re = (proj.real * np.sqrt(inv_noise_var_chunk[...,np.newaxis])).reshape((-1, nsrcs))
-    v_im = (proj.imag * np.sqrt(inv_noise_var_chunk[...,np.newaxis])).reshape(((-1, nsrcs)))
+    v_re = (proj.real * np.sqrt(inv_noise_var_chunk[..., np.newaxis])).reshape(
+        (-1, nsrcs)
+    )
+    v_im = (proj.imag * np.sqrt(inv_noise_var_chunk[..., np.newaxis])).reshape(
+        ((-1, nsrcs))
+    )
 
     # Treat real and imaginary separately, and get copies, to massively
     # speed-up the matrix multiplication!
-    my_linear_op[:,:] = v_re.T @ v_re + v_im.T @ v_im
+    my_linear_op[:, :] = v_re.T @ v_re + v_im.T @ v_im
     del v_re, v_im
 
     # Do Reduce (sum) operation to get total operator on root node
-    linear_op = np.zeros((1,1), dtype=my_linear_op.dtype) # dummy data for non-root workers
+    linear_op = np.zeros(
+        (1, 1), dtype=my_linear_op.dtype
+    )  # dummy data for non-root workers
     if myid == 0:
         linear_op = np.zeros_like(my_linear_op)
-    
-    comm.Reduce(my_linear_op,
-                linear_op,
-                op=MPI_SUM,
-                root=0)
+
+    comm.Reduce(my_linear_op, linear_op, op=MPI_SUM, root=0)
 
     # Include prior and identity terms to finish constructing LHS operator on root worker
     if myid == 0:
-        linear_op = np.eye(linear_op.shape[0]) \
-                  + np.diag(amp_prior_std) @ linear_op @ np.diag(amp_prior_std)
-    
+        linear_op = np.eye(linear_op.shape[0]) + np.diag(
+            amp_prior_std
+        ) @ linear_op @ np.diag(amp_prior_std)
+
     # (3) Calculate linear system RHS
     proj = proj.reshape((-1, nsrcs))
-    realisation_switch = 1.0 if realisation else 0.0 # Turn random realisations on or off
+    realisation_switch = (
+        1.0 if realisation else 0.0
+    )  # Turn random realisations on or off
 
     # Calculate residual of data vs fiducial model (residual from amplitudes = 1)
     # (proj now includes gains)
-    resid_chunk = data_chunk.copy() \
-          - (  proj.reshape((-1, nsrcs)) 
-             @ np.ones_like(amp_prior_std) ).reshape(data_chunk.shape)
+    resid_chunk = data_chunk.copy() - (
+        proj.reshape((-1, nsrcs)) @ np.ones_like(amp_prior_std)
+    ).reshape(data_chunk.shape)
 
     # (Terms 1+3): S^1/2 A^\dagger [ N^{-1} r + N^{-1/2} \omega_r ]
     omega_n = (
         realisation_switch
-        * (   1.0 * np.random.randn(*resid_chunk.shape) 
-            + 1.0j * np.random.randn(*resid_chunk.shape)  )
+        * (
+            1.0 * np.random.randn(*resid_chunk.shape)
+            + 1.0j * np.random.randn(*resid_chunk.shape)
+        )
         / np.sqrt(2.0)
     )
 
     # Separate complex part of RHS into real and imaginary parts, and apply
     # the real and imaginary parts of the projection operator separately.
     # This is necessary to get a real RHS vector
-    y = ((resid_chunk * inv_noise_var_chunk) + (omega_n * np.sqrt(inv_noise_var_chunk))).flatten()
+    y = (
+        (resid_chunk * inv_noise_var_chunk) + (omega_n * np.sqrt(inv_noise_var_chunk))
+    ).flatten()
     b = amp_prior_std * (proj.T.real @ y.real + proj.T.imag @ y.imag)
 
     # Reduce (sum) operation on b
-    linear_rhs = np.zeros((1,), dtype=b.dtype) # dummy data for non-root workers
+    linear_rhs = np.zeros((1,), dtype=b.dtype)  # dummy data for non-root workers
     if myid == 0:
         linear_rhs = np.zeros_like(b)
     comm.Reduce(b, linear_rhs, op=MPI_SUM, root=0)
 
     # (Term 2): \omega_a
     if myid == 0:
-        linear_rhs += realisation_switch * np.random.randn(nsrcs) # real vector
+        linear_rhs += realisation_switch * np.random.randn(nsrcs)  # real vector
 
     return linear_op, linear_rhs
 
 
 def calc_proj_operator(
-    ra, dec, fluxes, ant_pos, antpairs, freqs, times, beams,
-    latitude=-0.5361913261514378
+    ra,
+    dec,
+    fluxes,
+    ant_pos,
+    antpairs,
+    freqs,
+    times,
+    beams,
+    latitude=-0.5361913261514378,
 ):
     """
     Calculate a visibility vector for each point source, as a function of
@@ -166,7 +184,7 @@ def calc_proj_operator(
         polarized=False,
         precision=2,
         latitude=latitude,
-        use_feed="x"
+        use_feed="x",
     )
 
     # Allocate computed visibilities to only available baselines (saves memory)
@@ -179,7 +197,6 @@ def calc_proj_operator(
     return vis_ptsrc
 
 
-
 def legacy_precompute_op(vis_proj_operator, inv_noise_var):
     """
     Precompute the real and imaginary blocks of the matrix operator
@@ -207,21 +224,21 @@ def legacy_precompute_op(vis_proj_operator, inv_noise_var):
     """
     nsrcs = vis_proj_operator.shape[-1]
 
-    #v = vis_proj_operator * np.sqrt(inv_noise_var)[:, :, :, np.newaxis]
+    # v = vis_proj_operator * np.sqrt(inv_noise_var)[:, :, :, np.newaxis]
     # Treat real and imaginary separately, and get copies, to massively
     # speed-up the matrix multiplication!
-    #v_re = v.reshape((-1, nsrcs)).real.copy()
-    #v_im = v.reshape((-1, nsrcs)).imag.copy()
-    #return v_re.T @ v_re + v_im.T @ v_im
+    # v_re = v.reshape((-1, nsrcs)).real.copy()
+    # v_im = v.reshape((-1, nsrcs)).imag.copy()
+    # return v_re.T @ v_re + v_im.T @ v_im
 
     # Treat real and imaginary separately, and get copies, to massively
     # speed-up the matrix multiplication!
     v_re = vis_proj_operator.real * np.sqrt(inv_noise_var)[:, :, :, np.newaxis]
-    y = np.einsum('ji,ik->ik', v_re, v_re)
+    y = np.einsum("ji,ik->ik", v_re, v_re)
     del v_re
 
     v_im = vis_proj_operator.imag * np.sqrt(inv_noise_var)[:, :, :, np.newaxis]
-    y += np.einsum('ji,ik->ik', v_im, v_im)
+    y += np.einsum("ji,ik->ik", v_im, v_im)
     del v_im
     return y
 
@@ -285,7 +302,7 @@ def legacy_construct_rhs(
     realisation_switch = 1.0 if realisation else 0.0
 
     # (Term 2): \omega_a
-    b = realisation_switch * np.random.randn(Nptsrc) # real vector
+    b = realisation_switch * np.random.randn(Nptsrc)  # real vector
 
     # (Terms 1+3): S^1/2 A^\dagger [ N^{-1} r + N^{-1/2} \omega_r ]
     omega_n = (
@@ -299,4 +316,4 @@ def legacy_construct_rhs(
     # This is necessary to get a real RHS vector
     y = ((resid * inv_noise_var) + (omega_n * np.sqrt(inv_noise_var))).flatten()
     b += amp_prior_std * (proj.T.real @ y.real + proj.T.imag @ y.imag)
-    return b
\ No newline at end of file
+    return b
diff --git a/hydra/region_sampler.py b/hydra/region_sampler.py
index d882474..56f597f 100644
--- a/hydra/region_sampler.py
+++ b/hydra/region_sampler.py
@@ -1,4 +1,3 @@
-
 import numpy as np
 from .vis_simulator import simulate_vis
 
@@ -8,12 +7,12 @@
 import healpy as hp
 
 
-def get_diffuse_sky_model_pixels(freqs, nside=32, sky_model='gsm2016'):
+def get_diffuse_sky_model_pixels(freqs, nside=32, sky_model="gsm2016"):
     """
-    Returns arrays of the pixel RA, Dec locations and per-pixel frequency 
-    spectra from a given sky model. By default this is GSM, as implemented 
+    Returns arrays of the pixel RA, Dec locations and per-pixel frequency
+    spectra from a given sky model. By default this is GSM, as implemented
     in `pygdsm`.
-    
+
     Parameters:
         freqs (array_like):
             Frequencies, in MHz.
@@ -22,7 +21,7 @@ def get_diffuse_sky_model_pixels(freqs, nside=32, sky_model='gsm2016'):
             Healpix nside to use when constructing the sky model.
 
         sky_model (str):
-            Which sky modle to use, from pyGDSM. One of: 
+            Which sky modle to use, from pyGDSM. One of:
             'gsm2008', 'gsm2016', 'haslam', 'lfss'.
 
     Returns:
@@ -34,22 +33,26 @@ def get_diffuse_sky_model_pixels(freqs, nside=32, sky_model='gsm2016'):
     """
     # Get expected frequency units
     freqs_MHz = freqs
-    
+
     # Initialise sky model and extract data cube
-    assert sky_model in ['gsm2008', 'gsm2016', 'haslam', 'lfsm'], \
-        "Available sky models: 'gsm2008', 'gsm2016', 'haslam', 'lfsm'"
-    if sky_model == 'gsm2008':
-        model = pygdsm.GlobalSkyModel(freq_unit='MHz', include_cmb=False)
-    if sky_model == 'gsm2016':
-        model = pygdsm.GlobalSkyModel16(freq_unit='MHz', include_cmb=False)
-    if sky_model == 'haslam':
-        model = pygdsm.HaslamSkyModel(freq_unit='MHz', include_cmb=False)
-    if sky_model == 'lfsm':
-        model = pygdsm.LowFrequencySkyModel(freq_unit='MHz', include_cmb=False)
-    
+    assert sky_model in [
+        "gsm2008",
+        "gsm2016",
+        "haslam",
+        "lfsm",
+    ], "Available sky models: 'gsm2008', 'gsm2016', 'haslam', 'lfsm'"
+    if sky_model == "gsm2008":
+        model = pygdsm.GlobalSkyModel(freq_unit="MHz", include_cmb=False)
+    if sky_model == "gsm2016":
+        model = pygdsm.GlobalSkyModel16(freq_unit="MHz", include_cmb=False)
+    if sky_model == "haslam":
+        model = pygdsm.HaslamSkyModel(freq_unit="MHz", include_cmb=False)
+    if sky_model == "lfsm":
+        model = pygdsm.LowFrequencySkyModel(freq_unit="MHz", include_cmb=False)
+
     model.generate(freqs_MHz)
-    sky_maps = model.generated_map_data # (Nfreqs, Npix), should be in Kelvin
-    
+    sky_maps = model.generated_map_data  # (Nfreqs, Npix), should be in Kelvin
+
     # Must change nside to make compute practical
     nside_gsm = hp.npix2nside(sky_maps[0].size)
     sky_maps = hp.ud_grade(sky_maps, nside_out=nside)
@@ -58,31 +61,32 @@ def get_diffuse_sky_model_pixels(freqs, nside=32, sky_model='gsm2016'):
     equiv = u.brightness_temperature(freqs_MHz * u.MHz, beam_area=1 * u.sr)
     Ksr_per_Jy = ((1 * u.Jy).to(u.K, equivalencies=equiv) * u.sr / u.Jy).value
     for i in range(Ksr_per_Jy.size):
-        sky_maps[i] /= Ksr_per_Jy[i] # converts K to Jy/sr
-    
+        sky_maps[i] /= Ksr_per_Jy[i]  # converts K to Jy/sr
+
     # Get pixel RA/Dec coords (assumes Galactic coords for sky map data)
     idxs = np.arange(sky_maps[0].size)
     pix_lon, pix_lat = hp.pix2ang(nside, idxs, lonlat=True)
-    gal_coords = Galactic(l=pix_lon*u.deg, b=pix_lat*u.deg)
-    
+    gal_coords = Galactic(l=pix_lon * u.deg, b=pix_lat * u.deg)
+
     icrs_frame = ICRS()
     eq_coords = gal_coords.transform_to(icrs_frame)
     ra = eq_coords.ra.rad
     dec = eq_coords.dec.rad
-    
+
     # Returns list of pixels with spectra as effective point sources
     return ra, dec, sky_maps.T
 
 
-def segmented_diffuse_sky_model_pixels(ra, dec, sky_maps, freqs, nregions, 
-                                       smoothing_fwhm=None):
+def segmented_diffuse_sky_model_pixels(
+    ra, dec, sky_maps, freqs, nregions, smoothing_fwhm=None
+):
     """
-    Returns a list of pixel indices for each region in a diffuse sky map. 
-    This function currently uses a crude spectral index measurement to 
-    segment the map into regions with roughly equal numbers of pixels. 
-    Smoothing can be used to reduce sharp edges. The regions can be 
+    Returns a list of pixel indices for each region in a diffuse sky map.
+    This function currently uses a crude spectral index measurement to
+    segment the map into regions with roughly equal numbers of pixels.
+    Smoothing can be used to reduce sharp edges. The regions can be
     disconnected.
-    
+
     Parameters:
         ra, dec (array_like):
             ICRS RA and Dec locations of each pixel.
@@ -94,33 +98,33 @@ def segmented_diffuse_sky_model_pixels(ra, dec, sky_maps, freqs, nregions,
             Frequencies, in MHz.
 
         nregions (int):
-            The number of regions of roughly equal numbers of pixels to 
+            The number of regions of roughly equal numbers of pixels to
             segment the sky map into.
 
         smoothing_fwhm (float):
-            Smoothing FWHM (in degrees) to apply to the segmented map in 
-            order to reduce sharp edges. The smoothing is applied to a map 
-            of region indices. It is then re-segmented, which can result 
-            in a slight reduction in the number of segments due to 
+            Smoothing FWHM (in degrees) to apply to the segmented map in
+            order to reduce sharp edges. The smoothing is applied to a map
+            of region indices. It is then re-segmented, which can result
+            in a slight reduction in the number of segments due to
             round-off of region indices.
 
     Returns:
         idxs (list of array_like):
-            List of arrays, with each array containing the array indices 
+            List of arrays, with each array containing the array indices
             of the pixels that belong to each region.
     """
     # Crude spectral index map
-    beta = np.log(sky_maps[:,0] / sky_maps[:,1]) / np.log(freqs[0] / freqs[1])
+    beta = np.log(sky_maps[:, 0] / sky_maps[:, 1]) / np.log(freqs[0] / freqs[1])
 
     # Sort spectral index map and break up into ~equal-sized segments
     beta_sorted = np.sort(beta)
-    bounds = beta_sorted[::beta_sorted.size // nregions]
+    bounds = beta_sorted[:: beta_sorted.size // nregions]
 
     # Loop over regions and select pixels belonging to each region
     regions = np.zeros(beta.size, dtype=int)
     for i in range(bounds.size - 1):
         # These have >= and <= to ensure that all pixels belong somewhere
-        idxs = np.where(np.logical_and(beta >= bounds[i], beta <= bounds[i+1]))
+        idxs = np.where(np.logical_and(beta >= bounds[i], beta <= bounds[i + 1]))
         regions[idxs] = i
 
     # Apply smoothing and then re-segment
@@ -137,8 +141,16 @@ def segmented_diffuse_sky_model_pixels(ra, dec, sky_maps, freqs, nregions,
 
 
 def calc_proj_operator(
-    region_pixel_ra, region_pixel_dec, region_fluxes, region_idxs, ant_pos, 
-    antpairs, freqs, times, beams, latitude=-0.5361913261514378
+    region_pixel_ra,
+    region_pixel_dec,
+    region_fluxes,
+    region_idxs,
+    ant_pos,
+    antpairs,
+    freqs,
+    times,
+    beams,
+    latitude=-0.5361913261514378,
 ):
     """
     Calculate a visibility vector for each point source, as a function of
@@ -169,7 +181,7 @@ def calc_proj_operator(
 
     Returns:
         proj_operator (array_like):
-            The projection operator from region amplitudes to visibilities. This 
+            The projection operator from region amplitudes to visibilities. This
             is an array of the visibility values for each region.
     """
     Nregions = len(region_idxs)
@@ -185,7 +197,7 @@ def calc_proj_operator(
         # Returns shape (NFREQS, NTIMES, NANTS, NANTS)
         vis = simulate_vis(
             ants=ant_pos,
-            fluxes=region_fluxes[region_idxs[j],:],
+            fluxes=region_fluxes[region_idxs[j], :],
             ra=region_pixel_ra[region_idxs[j]],
             dec=region_pixel_dec[region_idxs[j]],
             freqs=freqs * 1e6,
@@ -194,7 +206,7 @@ def calc_proj_operator(
             polarized=False,
             precision=2,
             latitude=latitude,
-            use_feed="x"
+            use_feed="x",
         )
 
         # Allocate computed visibilities to only available baselines (saves memory)
@@ -203,4 +215,4 @@ def calc_proj_operator(
             idx2 = ants.index(bl[1])
             vis_region[i, :, :, j] = vis[:, :, idx1, idx2]
 
-    return vis_region
\ No newline at end of file
+    return vis_region
diff --git a/hydra/sh_sampler.py b/hydra/sh_sampler.py
index 7234ab4..62d2b18 100644
--- a/hydra/sh_sampler.py
+++ b/hydra/sh_sampler.py
@@ -24,11 +24,11 @@ def get_em_ell_idx(lmax):
     """
     With (m,l) ordering!
     """
-    ells_list = np.arange(0,lmax+1)
-    em_real = np.arange(0,lmax+1)
-    em_imag = np.arange(1,lmax+1)
+    ells_list = np.arange(0, lmax + 1)
+    em_real = np.arange(0, lmax + 1)
+    em_imag = np.arange(1, lmax + 1)
     # ylabel = []
-    
+
     # First append all real (l,m) values
     Nreal = 0
     i = 0
@@ -43,12 +43,12 @@ def get_em_ell_idx(lmax):
                 ells.append(ell)
                 Nreal += 1
                 i += 1
-    
+
     # Then all imaginary -- note: no m=0 modes!
     Nimag = 0
     for em in em_imag:
         for ell in ells_list:
-            if ell >= em :
+            if ell >= em:
                 idx.append(i)
                 ems.append(em)
                 ells.append(ell)
@@ -56,22 +56,25 @@ def get_em_ell_idx(lmax):
                 i += 1
     return ems, ells, idx
 
-def vis_proj_operator_no_rot(freqs, 
-                             lsts, 
-                             beams, 
-                             ant_pos, 
-                             lmax, 
-                             nside, 
-                             latitude=-0.5361913261514378, 
-                             include_autos=False, 
-                             autos_only=False,
-                             ref_freq=100.,
-                             spectral_idx=0.):
+
+def vis_proj_operator_no_rot(
+    freqs,
+    lsts,
+    beams,
+    ant_pos,
+    lmax,
+    nside,
+    latitude=-0.5361913261514378,
+    include_autos=False,
+    autos_only=False,
+    ref_freq=100.0,
+    spectral_idx=0.0,
+):
     """
-    Precompute the real and imaginary blocks of the visibility response 
+    Precompute the real and imaginary blocks of the visibility response
     operator. This should only be done once and then "apply_vis_response()"
     is used to get the actual visibilities.
-    
+
     Parameters:
         freqs (array_like):
             Frequencies, in MHz.
@@ -81,44 +84,47 @@ def vis_proj_operator_no_rot(freqs,
             List of pyuveam objects, one for each antenna
         ant_pos (dict):
             Dictionary of antenna positions, [x, y, z], in m. The keys should
-            be the numerical antenna IDs.    
+            be the numerical antenna IDs.
         lmax (int):
             Maximum ell value. Determines the number of modes used.
         nside (int):
-            Healpix nside to use for the calculation (longer baselines should 
+            Healpix nside to use for the calculation (longer baselines should
             use higher nside).
         latitude (float):
-            Latitude in decimal format of the simulated array/visibilities. 
+            Latitude in decimal format of the simulated array/visibilities.
         include_autos (bool):
             If `True`, the auto baselines are included.
         ref_freq (float):
             Reference frequency for the spectral dependence, in MHz.
         spectral_idx (float):
-            Spectral index, `beta`, for the spectral dependence, 
+            Spectral index, `beta`, for the spectral dependence,
             `~(freqs / ref_freq)^beta`.
-    
+
     Returns:
         vis_response_2D (array_like):
-            Visibility operator (δV_ij) for each (l,m) mode, frequency, 
+            Visibility operator (δV_ij) for each (l,m) mode, frequency,
             baseline and lst. Shape (Nvis, Nalms) where Nvis is Nbl x Ntimes x Nfreqs.
         ell (array of int):
             Array of ell-values for the visiblity simulation
         m  (array of int):
             Array of ell-values for the visiblity simulation
     """
-    ell, m, vis_alm = simulate_vis_per_alm(lmax=lmax, 
-                                           nside=nside, 
-                                           ants=ant_pos, 
-                                           freqs=freqs*1e6, # MHz -> Hz 
-                                           lsts=lsts, 
-                                           beams=beams,
-                                           latitude=latitude)
-    
-    # Removing visibility responses corresponding to the m=0 imaginary parts 
-    vis_alm = np.concatenate((vis_alm[:,:,:,:,:len(ell)],
-                              vis_alm[:,:,:,:,len(ell)+(lmax+1):]), 
-                             axis=4)
-    
+    ell, m, vis_alm = simulate_vis_per_alm(
+        lmax=lmax,
+        nside=nside,
+        ants=ant_pos,
+        freqs=freqs * 1e6,  # MHz -> Hz
+        lsts=lsts,
+        beams=beams,
+        latitude=latitude,
+    )
+
+    # Removing visibility responses corresponding to the m=0 imaginary parts
+    vis_alm = np.concatenate(
+        (vis_alm[:, :, :, :, : len(ell)], vis_alm[:, :, :, :, len(ell) + (lmax + 1) :]),
+        axis=4,
+    )
+
     ants = list(ant_pos.keys())
     antpairs = []
     if autos_only == False and include_autos == False:
@@ -137,41 +143,48 @@ def vis_proj_operator_no_rot(freqs,
                     auto_ants.append((ants[i], ants[j]))
                 if j > i:
                     antpairs.append((ants[i], ants[j]))
-                
-    vis_response = np.zeros((len(antpairs), len(freqs), len(lsts), 2*len(ell)-(lmax+1)), 
-                            dtype=np.complex128)
-    
+
+    vis_response = np.zeros(
+        (len(antpairs), len(freqs), len(lsts), 2 * len(ell) - (lmax + 1)),
+        dtype=np.complex128,
+    )
+
     ## Collapse the two antenna dimensions into one baseline dimension
-    # Nfreqs, Ntimes, Nant1, Nant2, Nalms --> Nbl, Nfreqs, Ntimes, Nalms 
+    # Nfreqs, Ntimes, Nant1, Nant2, Nalms --> Nbl, Nfreqs, Ntimes, Nalms
     for i, bl in enumerate(antpairs):
         idx1 = ants.index(bl[0])
         idx2 = ants.index(bl[1])
-        vis_response[i, :] = vis_alm[:, :, idx1, idx2, :]  
-    
+        vis_response[i, :] = vis_alm[:, :, idx1, idx2, :]
+
     # Multiply by spectral dependence model (a powerlaw)
-    # Shape: Nbl, Nfreqs, Ntimes, Nalms 
-    vis_response *= ((freqs / ref_freq)**spectral_idx)[np.newaxis,:,np.newaxis,np.newaxis]
+    # Shape: Nbl, Nfreqs, Ntimes, Nalms
+    vis_response *= ((freqs / ref_freq) ** spectral_idx)[
+        np.newaxis, :, np.newaxis, np.newaxis
+    ]
 
     # Reshape to 2D
     # TODO: Make this into a "pack" and "unpack" function
     # Nbl, Nfreqs, Ntimes, Nalms --> Nvis, Nalms
     Nvis = len(antpairs) * len(freqs) * len(lsts)
-    vis_response_2D = vis_response.reshape(Nvis, 2*len(ell)-(lmax+1))
-    
+    vis_response_2D = vis_response.reshape(Nvis, 2 * len(ell) - (lmax + 1))
+
     if autos_only == False and include_autos == False:
-        autos = np.zeros((len(auto_ants),len(freqs),len(lsts),2*len(ell)-(lmax+1)), dtype=np.complex128)
+        autos = np.zeros(
+            (len(auto_ants), len(freqs), len(lsts), 2 * len(ell) - (lmax + 1)),
+            dtype=np.complex128,
+        )
         ## Collapse the two antenna dimensions into one baseline dimension
-        # Nfreqs, Ntimes, Nant1, Nant2, Nalms --> Nbl, Nfreqs, Ntimes, Nalms 
+        # Nfreqs, Ntimes, Nant1, Nant2, Nalms --> Nbl, Nfreqs, Ntimes, Nalms
         for i, bl in enumerate(auto_ants):
             idx1 = ants.index(bl[0])
             idx2 = ants.index(bl[1])
-            autos[i, :] = vis_alm[:, :, idx1, idx2, :]   
+            autos[i, :] = vis_alm[:, :, idx1, idx2, :]
 
         ## Reshape to 2D
         ## TODO: Make this into a "pack" and "unpack" function
         # Nbl, Nfreqs, Ntimes, Nalms --> Nvis, Nalms
         Nautos = len(auto_ants) * len(freqs) * len(lsts)
-        autos_2D = autos.reshape(Nautos, 2*len(ell)-(lmax+1))
+        autos_2D = autos.reshape(Nautos, 2 * len(ell) - (lmax + 1))
 
     if autos_only == False and include_autos == False:
         return vis_response_2D, autos_2D, ell, m
@@ -181,65 +194,66 @@ def vis_proj_operator_no_rot(freqs,
 
 def alms2healpy(alms, lmax):
     """
-    Takes a real array split as [real, imag] (without the m=0 modes 
-    imag-part) and turns it into a complex array of alms (positive 
+    Takes a real array split as [real, imag] (without the m=0 modes
+    imag-part) and turns it into a complex array of alms (positive
     modes only) ordered as in HEALpy.
-      
+
     Parameters:
         alms (array_like):
-            Array of zeros except for the specified mode. 
-            The array represents all positive (+m) modes including zero 
-            and has double length, as real and imaginary values are split. 
+            Array of zeros except for the specified mode.
+            The array represents all positive (+m) modes including zero
+            and has double length, as real and imaginary values are split.
             The first half is the real values.
-    
+
     Returns:
         healpy_modes (array_like):
-            Array of zeros except for the specified mode. 
+            Array of zeros except for the specified mode.
             The array represents all positive (+m) modes including zeroth modes.
     """
-    
-    real_imag_split_index = int((np.size(alms)+(lmax+1))/2)
+
+    real_imag_split_index = int((np.size(alms) + (lmax + 1)) / 2)
     real = alms[:real_imag_split_index]
-    
-    add_imag_m0_modes = np.zeros(lmax+1)
+
+    add_imag_m0_modes = np.zeros(lmax + 1)
     imag = np.concatenate((add_imag_m0_modes, alms[real_imag_split_index:]))
-    
-    healpy_modes = real + 1.j*imag
-    
+
+    healpy_modes = real + 1.0j * imag
+
     return healpy_modes
-    
-    
+
+
 def healpy2alms(healpy_modes):
     """
     Takes a complex array of alms (positive modes only) and turns into
-    a real array split as [real, imag] making sure to remove the 
+    a real array split as [real, imag] making sure to remove the
     m=0 modes from the imag-part.
-      
+
     Parameters:
         healpy_modes (array_like, complex):
-            Array of zeros except for the specified mode. 
+            Array of zeros except for the specified mode.
             The array represents all positive (+m) modes including zeroth modes.
-    
+
     Returns:
         alms (array_like):
-            Array of zeros except for the specified mode. 
-            The array represents all positive (+m) modes including zero 
+            Array of zeros except for the specified mode.
+            The array represents all positive (+m) modes including zero
             and is split into a real (first) and imag (second) part. The
-            Imag part is smaller as the m=0 modes shouldn't contain and 
-            imaginary part. 
+            Imag part is smaller as the m=0 modes shouldn't contain and
+            imaginary part.
     """
-    lmax = hp.sphtfunc.Alm.getlmax(healpy_modes.size) # to remove the m=0 imag modes
-    alms = np.concatenate((healpy_modes.real,healpy_modes.imag[(lmax+1):]))
-        
-    return alms   
+    lmax = hp.sphtfunc.Alm.getlmax(healpy_modes.size)  # to remove the m=0 imag modes
+    alms = np.concatenate((healpy_modes.real, healpy_modes.imag[(lmax + 1) :]))
 
+    return alms
 
-def get_healpy_from_gsm(freq, lmax, nside=64, resolution="low", output_model=False, 
-                        output_map=False):
+
+def get_healpy_from_gsm(
+    freq, lmax, nside=64, resolution="low", output_model=False, output_map=False
+):
     """
-    Generate an array of alms (HEALpy ordered) from gsm 2016 
+    Generate an array of alms (HEALpy ordered) from gsm 2016
     (https://github.com/telegraphic/pygdsm)
-    
+
     Parameters:
         freqs (array_like):
             Frequency (in MHz) for which to return GSM model.
@@ -249,51 +263,52 @@ def get_healpy_from_gsm(freq, lmax, nside=64, resolution="low", output_model=Fal
             The nside to upgrade/downgrade the map to. Default is nside=64.
         resolution (str):
             if "low/lo/l":  The GSM nside = 64  (default)
-            if "hi/high/h": The GSM nside = 1024 
+            if "hi/high/h": The GSM nside = 1024
         output_model (bool):
-            If output_model=True: Outputs model generated from the GSM data. 
+            If output_model=True: Outputs model generated from the GSM data.
             If output_model=False (default): no model output.
         output_map (bool):
-            If output_map=True: Outputs map generated from the GSM data. 
+            If output_map=True: Outputs map generated from the GSM data.
             If output_map=False (default): no map output.
 
     Returns:
         healpy_modes (array_like):
             Complex array of alms with same size and ordering as in healpy (m,l)
         gsm_2016 (PyGDSM 2016 model):
-            If output_model=True: Outputs model generated from the GSM data. 
+            If output_model=True: Outputs model generated from the GSM data.
             If output_model=False (default): no model output.
         gsm_map (healpy map):
-            If output_map=True: Outputs map generated from the GSM data. 
+            If output_map=True: Outputs map generated from the GSM data.
             If output_map=False (default): no map output.
-    
+
     """
     # Instantiate GSM model and extract alms
-    gsm_2016 = GlobalSkyModel2016(freq_unit='MHz', resolution=resolution) 
+    gsm_2016 = GlobalSkyModel2016(freq_unit="MHz", resolution=resolution)
     gsm_map = gsm_2016.generate(freqs=freq)
     gsm_upgrade = hp.ud_grade(gsm_map, nside)
-    healpy_modes_gal = hp.map2alm(maps=gsm_upgrade,lmax=lmax)
+    healpy_modes_gal = hp.map2alm(maps=gsm_upgrade, lmax=lmax)
 
     # By default it is in gal-coordinates, convert to equatorial
     rot_gal2eq = hp.Rotator(coord="GC")
     healpy_modes_eq = rot_gal2eq.rotate_alm(healpy_modes_gal)
 
-    if output_model == False and output_map == False: # default
+    if output_model == False and output_map == False:  # default
         return healpy_modes_eq
     elif output_model == False and output_map == True:
-        return healpy_modes_eq, gsm_map 
+        return healpy_modes_eq, gsm_map
     elif output_model == True and output_map == False:
-        return healpy_modes_eq, gsm_2016 
+        return healpy_modes_eq, gsm_2016
     else:
         return healpy_modes_eq, gsm_2016, gsm_map
 
 
-def get_alms_from_gsm(freq, lmax, nside=64, resolution='low', output_model=False, 
-                      output_map=False):
+def get_alms_from_gsm(
+    freq, lmax, nside=64, resolution="low", output_model=False, output_map=False
+):
     """
-    Generate a real array split as [real, imag] (without the m=0 modes 
+    Generate a real array split as [real, imag] (without the m=0 modes
     imag-part) from gsm 2016 (https://github.com/telegraphic/pygdsm)
-    
+
     Parameters:
     freqs (float or array_like):
         Frequency (in MHz) for which to return GSM model
@@ -303,42 +318,48 @@ def get_alms_from_gsm(freq, lmax, nside=64, resolution='low', output_model=False
         The nside to upgrade/downgrade the map to. Default is nside=64.
     resolution (str):
         if "low/lo/l":  nside = 64  (default)
-        if "hi/high/h": nside = 1024 
+        if "hi/high/h": nside = 1024
     output_model (bool):
-        If output_model=True: Outputs model generated from the GSM data. 
+        If output_model=True: Outputs model generated from the GSM data.
         If output_model=False (default): no model output.
     output_map (bool):
-        If output_map=True: Outputs map generated from the GSM data. 
+        If output_map=True: Outputs map generated from the GSM data.
         If output_map=False (default): no map output.
 
     Returns:
         alms (array_like):
-            Array of zeros except for the specified mode. 
-            The array represents all positive (+m) modes including zero 
-            and has double length, as real and imaginary values are split. 
+            Array of zeros except for the specified mode.
+            The array represents all positive (+m) modes including zero
+            and has double length, as real and imaginary values are split.
             The first half is the real values.
         gsm_2016 (PyGDSM 2016 model):
-            If output_model=True: Outputs model generated from the GSM data. 
+            If output_model=True: Outputs model generated from the GSM data.
             If output_model=False (default): no model output.
         gsm_map (healpy map):
-            If output_map=True: Outputs map generated from the GSM data. 
+            If output_map=True: Outputs map generated from the GSM data.
             If output_map=False (default): no map output.
     """
-    return healpy2alms(get_healpy_from_gsm(freq, lmax, nside, resolution, output_model, output_map))
+    return healpy2alms(
+        get_healpy_from_gsm(freq, lmax, nside, resolution, output_model, output_map)
+    )
 
 
-def construct_rhs_no_rot(data, inv_noise_var, inv_prior_var, omega_0, omega_1, a_0, vis_response):
+def construct_rhs_no_rot(
+    data, inv_noise_var, inv_prior_var, omega_0, omega_1, a_0, vis_response
+):
     """
     Construct RHS of linear system.
     """
-    real_data_term = vis_response.real.T @ (inv_noise_var*data.real 
-                                            + np.sqrt(inv_noise_var)*omega_1.real)
-    imag_data_term = vis_response.imag.T @ (inv_noise_var*data.imag 
-                                            + np.sqrt(inv_noise_var)*omega_1.imag)
-    prior_term = inv_prior_var*a_0 + np.sqrt(inv_prior_var)*omega_0
-
-    right_hand_side = real_data_term + imag_data_term + prior_term 
-    
+    real_data_term = vis_response.real.T @ (
+        inv_noise_var * data.real + np.sqrt(inv_noise_var) * omega_1.real
+    )
+    imag_data_term = vis_response.imag.T @ (
+        inv_noise_var * data.imag + np.sqrt(inv_noise_var) * omega_1.imag
+    )
+    prior_term = inv_prior_var * a_0 + np.sqrt(inv_prior_var) * omega_0
+
+    right_hand_side = real_data_term + imag_data_term + prior_term
+
     return right_hand_side
 
 
@@ -346,20 +367,21 @@ def apply_lhs_no_rot(a_cr, inv_noise_var, inv_prior_var, vis_response):
     """
     Apply LHS operator of linear system to an input vector.
     """
-    real_noise_term = vis_response.real.T \
-                    @ ( inv_noise_var[:,np.newaxis] * vis_response.real ) \
-                    @ a_cr
-    imag_noise_term = vis_response.imag.T \
-                    @ ( inv_noise_var[:,np.newaxis]* vis_response.imag ) \
-                    @ a_cr
+    real_noise_term = (
+        vis_response.real.T @ (inv_noise_var[:, np.newaxis] * vis_response.real) @ a_cr
+    )
+    imag_noise_term = (
+        vis_response.imag.T @ (inv_noise_var[:, np.newaxis] * vis_response.imag) @ a_cr
+    )
     signal_term = inv_prior_var * a_cr
-    
-    left_hand_side = (real_noise_term + imag_noise_term + signal_term) 
+
+    left_hand_side = real_noise_term + imag_noise_term + signal_term
     return left_hand_side
 
 
-def construct_rhs_no_rot_mpi(comm, data, inv_noise_var, inv_prior_var, 
-                             omega_a, omega_n, a_0, vis_response):
+def construct_rhs_no_rot_mpi(
+    comm, data, inv_noise_var, inv_prior_var, omega_a, omega_n, a_0, vis_response
+):
     """
     Construct RHS of linear system from data split across multiple MPI workers.
     """
@@ -367,63 +389,65 @@ def construct_rhs_no_rot_mpi(comm, data, inv_noise_var, inv_prior_var,
 
     # Synchronise omega_a across all workers
     if myid != 0:
-        omega_a *= 0.
+        omega_a *= 0.0
     comm.Bcast(omega_a, root=0)
 
     # Calculate data terms
-    my_data_term = vis_response.real.T @ ((inv_noise_var * data.real).flatten()
-                                          + np.sqrt(inv_noise_var).flatten()
-                                            * omega_n.real.flatten()) \
-                 + vis_response.imag.T @ ((inv_noise_var * data.imag).flatten()
-                                          + np.sqrt(inv_noise_var).flatten() 
-                                            * omega_n.imag.flatten())
-    
+    my_data_term = vis_response.real.T @ (
+        (inv_noise_var * data.real).flatten()
+        + np.sqrt(inv_noise_var).flatten() * omega_n.real.flatten()
+    ) + vis_response.imag.T @ (
+        (inv_noise_var * data.imag).flatten()
+        + np.sqrt(inv_noise_var).flatten() * omega_n.imag.flatten()
+    )
+
     # Do Reduce (sum) operation to get total operator on root node
-    data_term = np.zeros((1,), dtype=my_data_term.dtype) # dummy data for non-root workers
+    data_term = np.zeros(
+        (1,), dtype=my_data_term.dtype
+    )  # dummy data for non-root workers
     if myid == 0:
         data_term = np.zeros_like(my_data_term)
-    
+
     comm.Reduce(my_data_term, data_term, op=MPI_SUM, root=0)
     comm.barrier()
 
     # Return result (only root worker has correct result)
     if myid == 0:
-        return data_term \
-             + inv_prior_var * a_0 \
-             + np.sqrt(inv_prior_var) * omega_a
+        return data_term + inv_prior_var * a_0 + np.sqrt(inv_prior_var) * omega_a
     else:
         return np.zeros_like(a_0)
 
 
 def apply_lhs_no_rot_mpi(comm, a_cr, inv_noise_var, inv_prior_var, vis_response):
     """
-    Apply LHS operator of linear system to an input vector that has been 
+    Apply LHS operator of linear system to an input vector that has been
     split into chunks between MPI workers.
     """
     myid = comm.Get_rank()
 
     # Synchronise a_cr across all workers
     if myid != 0:
-        a_cr *= 0.
+        a_cr *= 0.0
     comm.Bcast(a_cr, root=0)
 
     # Calculate noise terms for this rank
-    my_tot_noise_term = vis_response.real.T \
-                        @ ( inv_noise_var.flatten()[:,np.newaxis] * vis_response.real ) \
-                        @ a_cr \
-                      + vis_response.imag.T \
-                        @ ( inv_noise_var.flatten()[:,np.newaxis] * vis_response.imag ) \
-                        @ a_cr
+    my_tot_noise_term = (
+        vis_response.real.T
+        @ (inv_noise_var.flatten()[:, np.newaxis] * vis_response.real)
+        @ a_cr
+        + vis_response.imag.T
+        @ (inv_noise_var.flatten()[:, np.newaxis] * vis_response.imag)
+        @ a_cr
+    )
 
     # Do Reduce (sum) operation to get total operator on root node
-    tot_noise_term = np.zeros((1,), dtype=my_tot_noise_term.dtype) # dummy data for non-root workers
+    tot_noise_term = np.zeros(
+        (1,), dtype=my_tot_noise_term.dtype
+    )  # dummy data for non-root workers
     if myid == 0:
         tot_noise_term = np.zeros_like(my_tot_noise_term)
-    
-    comm.Reduce(my_tot_noise_term,
-                tot_noise_term,
-                op=MPI_SUM,
-                root=0)
+
+    comm.Reduce(my_tot_noise_term, tot_noise_term, op=MPI_SUM, root=0)
 
     # Return result (only root worker has correct result)
     if myid == 0:
@@ -433,101 +457,114 @@ def apply_lhs_no_rot_mpi(comm, a_cr, inv_noise_var, inv_prior_var, vis_response)
         return np.zeros_like(a_cr)
 
 
-def radiometer_eq(auto_visibilities, ants, delta_time, delta_freq, Nnights = 1, include_autos=False):
+def radiometer_eq(
+    auto_visibilities, ants, delta_time, delta_freq, Nnights=1, include_autos=False
+):
     nbls = len(ants)
-    indx = auto_visibilities.shape[0]//nbls
-    
-    sigma_full = np.empty((0))#, autos.shape[-1]))
+    indx = auto_visibilities.shape[0] // nbls
+
+    sigma_full = np.empty((0))  # , autos.shape[-1]))
 
     for i in ants:
-        vis_ii = auto_visibilities[i*indx:(i+1)*indx]#,:]
+        vis_ii = auto_visibilities[i * indx : (i + 1) * indx]  # ,:]
 
         for j in ants:
             if include_autos == True:
                 if j >= i:
-                    vis_jj = auto_visibilities[j*indx:(j+1)*indx]#,:]
-                    sigma_ij = ( vis_ii*vis_jj ) / ( Nnights*delta_time*delta_freq )
-                    sigma_full = np.concatenate((sigma_full,sigma_ij))
+                    vis_jj = auto_visibilities[j * indx : (j + 1) * indx]  # ,:]
+                    sigma_ij = (vis_ii * vis_jj) / (Nnights * delta_time * delta_freq)
+                    sigma_full = np.concatenate((sigma_full, sigma_ij))
             else:
-                if j > i:  # only keep this line if you don't want the auto baseline sigmas
-                    vis_jj = auto_visibilities[j*indx:(j+1)*indx]#,:]
-                    sigma_ij = ( vis_ii*vis_jj ) / ( Nnights*delta_time*delta_freq )
-                    sigma_full = np.concatenate((sigma_full,sigma_ij))
-                    
-    return sigma_full
+                if (
+                    j > i
+                ):  # only keep this line if you don't want the auto baseline sigmas
+                    vis_jj = auto_visibilities[j * indx : (j + 1) * indx]  # ,:]
+                    sigma_ij = (vis_ii * vis_jj) / (Nnights * delta_time * delta_freq)
+                    sigma_full = np.concatenate((sigma_full, sigma_ij))
 
+    return sigma_full
 
 
-# MAIN    
+# MAIN
 if __name__ == "__main__":
     start_time = time.time()
-    
+
     # Creating directory for output
-    if ARGS['directory']: 
-        directory = str(ARGS['directory'])
+    if ARGS["directory"]:
+        directory = str(ARGS["directory"])
     else:
         directory = "output"
 
-    path = f'/cosma8/data/dp270/dc-bull2/{directory}/'
-    try: 
+    path = f"/cosma8/data/dp270/dc-bull2/{directory}/"
+    try:
         os.makedirs(path)
     except FileExistsError:
-        print('folder already exists')
-    
+        print("folder already exists")
+
     # Defining the data_seed for the precomputation random seed
-    if ARGS['data_seed']:
-        data_seed = int(ARGS['data_seed'])
+    if ARGS["data_seed"]:
+        data_seed = int(ARGS["data_seed"])
     else:
         # if none is passed go back to 10 as before
         data_seed = 10
 
     # Defining the jobid to distinguish multiple runs in one go
-    if ARGS['jobid']: 
-        jobid = int(ARGS['jobid'])
+    if ARGS["jobid"]:
+        jobid = int(ARGS["jobid"])
     else:
         # if none is passed then don't change the keys
         jobid = 0
 
-    ant_pos = build_hex_array(hex_spec=(3,4), d=14.6)  #builds array with (3,4,3) ants = 10 total
+    ant_pos = build_hex_array(
+        hex_spec=(3, 4), d=14.6
+    )  # builds array with (3,4,3) ants = 10 total
     ants = list(ant_pos.keys())
     lmax = 20
     nside = 128
-    beam_diameter = 14.
-    beams = [pyuvsim.AnalyticBeam('gaussian', diameter=beam_diameter) for ant in ants]
+    beam_diameter = 14.0
+    beams = [pyuvsim.AnalyticBeam("gaussian", diameter=beam_diameter) for ant in ants]
     freqs = np.linspace(100e6, 102e6, 2)
-    lsts_hours = np.linspace(0.,8.,10)      # in hours for easy setting
-    lsts = np.deg2rad((lsts_hours/24)*360) # in radian, used by HYDRA (and this code)
-    delta_time = 60 # s
-    delta_freq = 1e+06 # (M)Hz
-    latitude = 31.7215 * np.pi / 180  # HERA loc in decimal numbers ## There's some sign error in the code, so this missing sign is a quick fix
+    lsts_hours = np.linspace(0.0, 8.0, 10)  # in hours for easy setting
+    lsts = np.deg2rad(
+        (lsts_hours / 24) * 360
+    )  # in radian, used by HYDRA (and this code)
+    delta_time = 60  # s
+    delta_freq = 1e06  # (M)Hz
+    latitude = (
+        31.7215 * np.pi / 180
+    )  # HERA loc in decimal numbers ## There's some sign error in the code, so this missing sign is a quick fix
     solver = cg
 
-    vis_response, autos, ell, m = vis_proj_operator_no_rot(freqs=freqs, 
-                                                        lsts=lsts, 
-                                                        beams=beams, 
-                                                        ant_pos=ant_pos, 
-                                                        lmax=lmax, 
-                                                        nside=nside,
-                                                        latitude=latitude)
+    vis_response, autos, ell, m = vis_proj_operator_no_rot(
+        freqs=freqs,
+        lsts=lsts,
+        beams=beams,
+        ant_pos=ant_pos,
+        lmax=lmax,
+        nside=nside,
+        latitude=latitude,
+    )
 
     np.random.seed(data_seed)
-    x_true = get_alms_from_gsm(freq=100,lmax=lmax, nside=nside)
+    x_true = get_alms_from_gsm(freq=100, lmax=lmax, nside=nside)
     model_true = vis_response @ x_true
 
     # Inverse noise covariance and noise on data
-    noise_cov = radiometer_eq(autos@x_true, ants, delta_time, delta_freq)
-    inv_noise_var = 1/noise_cov
-    data_noise = np.random.randn(noise_cov.size)*np.sqrt(noise_cov) 
+    noise_cov = radiometer_eq(autos @ x_true, ants, delta_time, delta_freq)
+    inv_noise_var = 1 / noise_cov
+    data_noise = np.random.randn(noise_cov.size) * np.sqrt(noise_cov)
     data_vec = model_true + data_noise
 
     # Inverse signal covariance
     zero_value = 0.001
-    prior_cov = (x_true*0.1)**2     # if 0.1 = 10% prior
+    prior_cov = (x_true * 0.1) ** 2  # if 0.1 = 10% prior
     prior_cov[prior_cov == 0] = zero_value
-    inv_prior_var = 1./prior_cov
-    a_0 = np.random.randn(x_true.size)*np.sqrt(prior_cov) + x_true # gaussian centered on alms with S variance 
+    inv_prior_var = 1.0 / prior_cov
+    a_0 = (
+        np.random.randn(x_true.size) * np.sqrt(prior_cov) + x_true
+    )  # gaussian centered on alms with S variance
 
-    # Define left hand side operator 
+    # Define left hand side operator
     def lhs_operator(x):
         y = apply_lhs_no_rot(x, inv_noise_var, inv_prior_var, vis_response)
 
@@ -536,133 +573,136 @@ def lhs_operator(x):
     # Wiener filter solution to provide initial guess:
     omega_0_wf = np.zeros_like(a_0)
     omega_1_wf = np.zeros_like(model_true, dtype=np.complex128)
-    rhs_wf = construct_rhs_no_rot(data_vec,
-                                  inv_noise_var, 
-                                  inv_prior_var, 
-                                  omega_0_wf, 
-                                  omega_1_wf, 
-                                  a_0, 
-                                  vis_response)
-    
-    # Build linear operator object 
+    rhs_wf = construct_rhs_no_rot(
+        data_vec,
+        inv_noise_var,
+        inv_prior_var,
+        omega_0_wf,
+        omega_1_wf,
+        a_0,
+        vis_response,
+    )
+
+    # Build linear operator object
     lhs_shape = (rhs_wf.size, rhs_wf.size)
-    lhs_linear_op = LinearOperator(matvec = lhs_operator,
-                                       shape = lhs_shape)
+    lhs_linear_op = LinearOperator(matvec=lhs_operator, shape=lhs_shape)
 
     # Get the Wiener Filter solution for initial guess
-    wf_soln, wf_convergence_info = solver(A = lhs_linear_op,
-                                          b = rhs_wf,
-                                          # tol = 1e-07,
-                                          maxiter = 15000)
-    
+    wf_soln, wf_convergence_info = solver(
+        A=lhs_linear_op,
+        b=rhs_wf,
+        # tol = 1e-07,
+        maxiter=15000,
+    )
+
     def samples(key):
         t_iter = time.time()
 
         # Set a random seed defined by the key
-        random_seed = 100*jobid + key
+        random_seed = 100 * jobid + key
         np.random.seed(random_seed)
-        #random_seed = np.random.get_state()[1][0] #for test/output purposes
+        # random_seed = np.random.get_state()[1][0] #for test/output purposes
 
         # Generate random maps for the realisations
         omega_0 = np.random.randn(a_0.size)
-        omega_1 = (np.random.randn(model_true.size) + 1.j*np.random.randn(model_true.size))/np.sqrt(2)
+        omega_1 = (
+            np.random.randn(model_true.size) + 1.0j * np.random.randn(model_true.size)
+        ) / np.sqrt(2)
 
         # Construct the right hand side
-        rhs = construct_rhs_no_rot(data_vec,
-                                   inv_noise_var, 
-                                   inv_prior_var,
-                                   omega_0,
-                                   omega_1,
-                                   a_0,
-                                   vis_response)
+        rhs = construct_rhs_no_rot(
+            data_vec, inv_noise_var, inv_prior_var, omega_0, omega_1, a_0, vis_response
+        )
 
         # Run and time solver
         time_start_solver = time.time()
-        x_soln, convergence_info = solver(A = lhs_linear_op,
-                                          b = rhs,
-                                          # tol = 1e-07,
-                                          maxiter = 15000,
-                                          x0 = wf_soln) #initial guess
+        x_soln, convergence_info = solver(
+            A=lhs_linear_op,
+            b=rhs,
+            # tol = 1e-07,
+            maxiter=15000,
+            x0=wf_soln,
+        )  # initial guess
         solver_time = time.time() - time_start_solver
-        iteration_time = time.time()-t_iter
-        
+        iteration_time = time.time() - t_iter
+
         # Save output
-        np.savez(path+'results_'+f'{data_seed}_'+f'{random_seed}',
-                 omega_0=omega_0,
-                 omega_1=omega_1,
-                 key=key,
-                 x_soln=x_soln,
-                 rhs=rhs,
-                 convergence_info=convergence_info,
-                 solver_time=solver_time,
-                 iteration_time=iteration_time
+        np.savez(
+            path + "results_" + f"{data_seed}_" + f"{random_seed}",
+            omega_0=omega_0,
+            omega_1=omega_1,
+            key=key,
+            x_soln=x_soln,
+            rhs=rhs,
+            convergence_info=convergence_info,
+            solver_time=solver_time,
+            iteration_time=iteration_time,
         )
-        
+
         return key, iteration_time
-    
+
     # Time for all precomputations
-    precomp_time = time.time()-start_time
-    print(f'\nprecomputation took:\n{precomp_time}\n')
-    
+    precomp_time = time.time() - start_time
+    print(f"\nprecomputation took:\n{precomp_time}\n")
+
     avg_iter_time = 0
 
-    # Multiprocessing, getting the samples    
-    number_of_cores = int(os.environ['SLURM_CPUS_PER_TASK'])
-    print(f'\nSLURM_CPUS_PER_TASK = {number_of_cores}')
+    # Multiprocessing, getting the samples
+    number_of_cores = int(os.environ["SLURM_CPUS_PER_TASK"])
+    print(f"\nSLURM_CPUS_PER_TASK = {number_of_cores}")
 
     with Pool(number_of_cores) as pool:
         # issue tasks and process results
         for result in pool.map(samples, range(100)):
             key, iteration_time = result
             avg_iter_time += iteration_time
-            #print(f'Iteration {key} completed in {iteration_time:.2f} seconds')
+            # print(f'Iteration {key} completed in {iteration_time:.2f} seconds')
+
+    avg_iter_time /= key + 1
+    print(f"average_iter_time:\n{avg_iter_time}\n")
 
-    avg_iter_time /= (key+1)
-    print(f'average_iter_time:\n{avg_iter_time}\n')
+    total_time = time.time() - start_time
+    print(f"total_time:\n{total_time}\n")
+    print(f"All output saved in folder {path}\n")
+    print(f"Note, ant_pos (dict) is saved in own file in {path}\n")
 
-    total_time = time.time()-start_time
-    print(f'total_time:\n{total_time}\n')
-    print(f'All output saved in folder {path}\n')
-    print(f'Note, ant_pos (dict) is saved in own file in {path}\n')
-    
     # Saving all globally calculated data
-    np.savez(path+'precomputed_data_'+f'{data_seed}_'+f'{jobid}',
-             vis_response=vis_response,
-             x_true=x_true,
-             inv_noise_var=inv_noise_var,
-             zero_value=zero_value,
-             inv_prior_var=inv_prior_var,
-             wf_soln=wf_soln,
-             nside=nside,
-             lmax=lmax,
-             ants=ants,
-             beam_diameter=beam_diameter,
-             freqs=freqs,
-             lsts_hours=lsts_hours,
-             precomp_time=precomp_time,
-             total_time=total_time
-            )
+    np.savez(
+        path + "precomputed_data_" + f"{data_seed}_" + f"{jobid}",
+        vis_response=vis_response,
+        x_true=x_true,
+        inv_noise_var=inv_noise_var,
+        zero_value=zero_value,
+        inv_prior_var=inv_prior_var,
+        wf_soln=wf_soln,
+        nside=nside,
+        lmax=lmax,
+        ants=ants,
+        beam_diameter=beam_diameter,
+        freqs=freqs,
+        lsts_hours=lsts_hours,
+        precomp_time=precomp_time,
+        total_time=total_time,
+    )
     # creating a dictionary with string-keys as required by .npz files
     ant_dict = dict((str(ant), ant_pos[ant]) for ant in ant_pos)
-    np.savez(path+'ant_pos',**ant_dict)
-    
-    
+    np.savez(path + "ant_pos", **ant_dict)
+
+
 def sample_cl(alms, ell, m):
     """
-    Sample C_ell from an inverse gamma distribution, given a set of 
+    Sample C_ell from an inverse gamma distribution, given a set of
     SH coefficients. See Eq. 7 of Eriksen et al. (arXiv:0709.1058).
     """
     # Get m, ell ordering
     m_vals, ell_vals, lm_idxs = get_em_ell_idx(lmax)
-    
+
     # Calculate sigma_ell = 1/(2 l + 1) sum_m |a_lm|^2
     for ell in np.unique(ell_vals):
-        
-        idxs = np.where(ell_vals == ell)
-
 
-    sigma_ell = 
-    x = invgamma.rvs(loc=1, scale=1)
-    C_l = x ((2l+1)/2) sigma_l
-    #a = (2l-1)/2
+        idxs = np.where(ell_vals == ell)
 
+    # sigma_ell =
+    # x = invgamma.rvs(loc=1, scale=1)
+    # C_l = x ((2l+1)/2) sigma_l
+    # a = (2l-1)/2
diff --git a/hydra/sparse_beam.py b/hydra/sparse_beam.py
index 84ced3b..6fbe866 100644
--- a/hydra/sparse_beam.py
+++ b/hydra/sparse_beam.py
@@ -7,35 +7,57 @@
 
 
 class sparse_beam(UVBeam):
-    
-    def __init__(self, filename, nmax, mmodes, za_range=(0, 90), 
-                 save_fn='', load=False, bound="Dirichlet", Nfeeds=None, 
-                 do_fit=True, alpha=np.sqrt(1 - np.cos(46 * np.pi / 90)), 
-                 num_modes_comp=64, nmodes_comp=None, mmodes_comp=None, 
-                 sparse_fit_coeffs=None, perturb=False, za_ml=np.deg2rad(18.),
-                 dza=np.deg2rad(3.), Nsin_pert=8, sin_pert_coeffs=None,
-                 cSL=0.2, gam=None, sqrt=False,
-                 rot=0.,stretch_x=1.,stretch_y=1., trans_x=0., trans_y=0.,
-                 **kwargs):
+
+    def __init__(
+        self,
+        filename,
+        nmax,
+        mmodes,
+        za_range=(0, 90),
+        save_fn="",
+        load=False,
+        bound="Dirichlet",
+        Nfeeds=None,
+        do_fit=True,
+        alpha=np.sqrt(1 - np.cos(46 * np.pi / 90)),
+        num_modes_comp=64,
+        nmodes_comp=None,
+        mmodes_comp=None,
+        sparse_fit_coeffs=None,
+        perturb=False,
+        za_ml=np.deg2rad(18.0),
+        dza=np.deg2rad(3.0),
+        Nsin_pert=8,
+        sin_pert_coeffs=None,
+        cSL=0.2,
+        gam=None,
+        sqrt=False,
+        rot=0.0,
+        stretch_x=1.0,
+        stretch_y=1.0,
+        trans_x=0.0,
+        trans_y=0.0,
+        **kwargs,
+    ):
         """
         Construct the sparse_beam instance, which is a subclass of UVBeam
 
         Parameters:
-            filename (str): 
+            filename (str):
                 The filename to for the UVBeam compatible file.
-            nmax (int): 
+            nmax (int):
                 Maximum number of radial (Bessel) modes.
-            mmodes (array of int): 
+            mmodes (array of int):
                 Which azimuthal (Fourier) modes to include.
-            za_range (tuple): 
+            za_range (tuple):
                 Minimum and maximum zenith angle to read in.
-            save_fn (str): 
+            save_fn (str):
                 filepath to save a numpy array of coefficients once fitting is
                 complete.
-            load (bool): 
-                Whether to load coefficients from save_fn rather than fitting 
+            load (bool):
+                Whether to load coefficients from save_fn rather than fitting
                 anew.
-            bound (str): 
+            bound (str):
                 Options are 'Dirichlet' or 'Neumann'. Refers to the boundary
                 conditions for the Laplace equation in cylindrical coordinates
                 i.e. it determines whether to use 0th order or 1st order
@@ -61,22 +83,20 @@ def __init__(self, filename, nmax, mmodes, za_range=(0, 90),
             self.gam = gam
             self.sin_pert_coeffs = sin_pert_coeffs
             self.cSL = cSL
-            
 
-        if Nfeeds is not None:         # power beam may not have the Nfeeds set
+        if Nfeeds is not None:  # power beam may not have the Nfeeds set
             assert self.Nfeeds is None, "Nfeeds already set on the beam"
             self.Nfeeds = Nfeeds
 
         if sqrt:
             self.data_array = np.sqrt(self.data_array)
-        
-        
+
         self.alpha = alpha
-        
+
         self.save_fn = save_fn
         self.nmax = nmax
         self.mmodes = mmodes
-        
+
         self.az_array = self.axis1_array
         self.rad_array = self.get_rad_array()
 
@@ -90,26 +110,29 @@ def __init__(self, filename, nmax, mmodes, za_range=(0, 90),
             self.ncoeff_bess = self.nmax * len(self.mmodes)
             self.bess_matr, self.trig_matr = self.get_dmatr()
             self.bess_fits, self.bess_beam = self.get_fits(load=load)
-            self.bess_ps = np.abs(self.bess_fits)**2
+            self.bess_ps = np.abs(self.bess_fits) ** 2
 
             self.num_modes_comp = num_modes_comp
             self.nmodes_comp, self.mmodes_comp = self.get_comp_inds()
             self.comp_fits, self.comp_beam = self.sparse_fit_loop()
 
         elif sparse_fit_coeffs is None:
-            raise ValueError("Must either set full_fit=True or supply "
-                             "sparse_fit_coeffs")
+            raise ValueError(
+                "Must either set full_fit=True or supply " "sparse_fit_coeffs"
+            )
         else:
             self.comp_fits = sparse_fit_coeffs
             self.num_modes_comp = self.comp_fits.shape[-1]
             if (nmodes_comp is None) or (mmodes_comp is None):
-                raise ValueError("Sparse fit coeffs supplied without "
-                                 "corresponding nmodes or mmodes. Check "
-                                 "sparse_nomdes and sparse_mmodes kwargs.")
+                raise ValueError(
+                    "Sparse fit coeffs supplied without "
+                    "corresponding nmodes or mmodes. Check "
+                    "sparse_nomdes and sparse_mmodes kwargs."
+                )
             else:
                 self.nmodes_comp = nmodes_comp
                 self.mmodes_comp = mmodes_comp
-        
+
         # Cache dicts for repeated interpolation
         self.az_array_dict = {}
         self.za_array_dict = {}
@@ -117,11 +140,9 @@ def __init__(self, filename, nmax, mmodes, za_range=(0, 90),
         self.bess_matr_interp_dict = {}
         self.bt_matr_interp_dict = {}
 
-
-
     def get_rad_array(self, za_array=None):
         """
-        Get the radial coordinates corresponding to the zenith angles in 
+        Get the radial coordinates corresponding to the zenith angles in
         za_array, calculated according to the formula in Hydra Beam Paper I.
 
         Parameters:
@@ -135,20 +156,19 @@ def get_rad_array(self, za_array=None):
         if za_array is None:
             za_array = self.axis2_array
         rad_array = np.sqrt(1 - np.cos(za_array)) / self.alpha
-        
-        return rad_array
 
+        return rad_array
 
     def get_bzeros(self):
         """
         Get the zeros of the appropriate Bessel function based on the
-        desired basis specified by the 'bound' attribute, along with the 
+        desired basis specified by the 'bound' attribute, along with the
         associated normalization.
 
         Returns:
-            zeros (array): 
+            zeros (array):
                 The zeros of the appropriate Bessel function
-            norm (array): 
+            norm (array):
                 The normalization for the Bessel functions so that their L2 norm
                 on the unit disc is 1.
 
@@ -159,21 +179,20 @@ def get_bzeros(self):
         else:
             zeros = jn_zeros(1, self.nmax - 1)
             norm = jn(2, zeros)
-            
+
             zeros = np.append(0, zeros)
             norm = np.append(1, norm)
         norm = norm / np.sqrt(2)
 
-        return zeros, norm 
-        
-        
+        return zeros, norm
+
     def get_dmatr(self):
         """
-        Compute the factored design matrix that maps from Fourier-Bessel 
+        Compute the factored design matrix that maps from Fourier-Bessel
         coefficients to pixel centers on the sky. Assumes az/za coordinates,
-        AND uniform sampling in azimuth. Full design matrix is the tensor 
+        AND uniform sampling in azimuth. Full design matrix is the tensor
         product of these two factors.
-                
+
         Returns:
             bess_matr (array, complex):
                 Has shape (Nza, Nn). Contains the radial information of the
@@ -181,28 +200,30 @@ def get_dmatr(self):
             trig_matr (array, complex):
                 Has shape (Naz, Nm). Contains the azimuthal information of the
                 design matrix.
-        """  
+        """
         zeros, norm = self.get_bzeros()
-        
+
         Naz = len(self.az_array)
-                
+
         bess_matr = jn(0, zeros[np.newaxis] * self.rad_array[:, np.newaxis]) / norm
         # Assume a regular az spacing and just make a unitary DFT matrix; better for fitting later
-        trig_matr = np.exp(1.0j * np.array(self.mmodes)[np.newaxis] * self.az_array[:, np.newaxis]) / np.sqrt(Naz)
-        
+        trig_matr = np.exp(
+            1.0j * np.array(self.mmodes)[np.newaxis] * self.az_array[:, np.newaxis]
+        ) / np.sqrt(Naz)
+
         return bess_matr, trig_matr
-    
+
     def get_dmatr_interp(self, az_array, za_array):
         """
         Get a design matrix specialized for interpolation rather than fitting.
 
         Parameters:
             az_array (array):
-                Azimuth angles to evaluate bassis functions at. Does not have to 
+                Azimuth angles to evaluate bassis functions at. Does not have to
                 be on a uniform grid, unlike the fitting design matrix. Should
                 be 1-dimensional (i.e. flattened).
             za_array (array):
-                Zenith angles to evaluate basis functions at. Should be 
+                Zenith angles to evaluate basis functions at. Should be
                 1-dimensional (i.e. flattened)
 
         Returns:
@@ -211,7 +232,7 @@ def get_dmatr_interp(self, az_array, za_array):
             trig_matr (array, complex):
                 The Fourier part of the design matrix.
         """
-        
+
         rad_array = self.get_rad_array(za_array)
         if self.perturb:
             if self.rot > 0:
@@ -220,7 +241,10 @@ def get_dmatr_interp(self, az_array, za_array):
                 if self.stretch_x == self.stretch_y:
                     rad_array /= self.stretch_x
                 else:
-                    rad_array *= np.sqrt((np.cos(az_array) / self.stretch_x)**2 + (np.sin(az_array) / self.stretch_y)**2)
+                    rad_array *= np.sqrt(
+                        (np.cos(az_array) / self.stretch_x) ** 2
+                        + (np.sin(az_array) / self.stretch_y) ** 2
+                    )
             if self.trans_x != 0 or self.trans_y != 0:
                 xtrans = rad_array * np.cos(az_array) - self.trans_x
                 ytrans = rad_array * np.sin(az_array) - self.trans_y
@@ -234,17 +258,18 @@ def get_dmatr_interp(self, az_array, za_array):
         if self.perturb:
             bess_matr *= self.SL_pert(rad_array=rad_array)[:, None]
         # Need to use the same normalization as in the dmatr used for fitting
-        trig_matr = np.exp(1.j *  np.array(self.mmodes)[np.newaxis] * az_array[:, np.newaxis]) / np.sqrt(Naz)
+        trig_matr = np.exp(
+            1.0j * np.array(self.mmodes)[np.newaxis] * az_array[:, np.newaxis]
+        ) / np.sqrt(Naz)
 
         return bess_matr, trig_matr
-    
-    
+
     def get_fits(self, load=False, data_array=None):
         """
         Compute Fourier-Bessel fits up to nmax and for all m-modes.
 
         Parameters:
-            load (bool): 
+            load (bool):
                 Whether to load precomputed solutions
 
         Returns:
@@ -252,10 +277,10 @@ def get_fits(self, load=False, data_array=None):
                 The coefficients for the Fourier-Bessel fit. Has shape
                 (nmax, len(mmodes), Naxes_vec, 1, Nfeeds, Nfreqs)
             fit_beam (array, complex):
-                The fit beam in sky coordinates. Has shape 
+                The fit beam in sky coordinates. Has shape
                 (Naxes_vec, 1, Nfeeds, Nfreqs, Nza, Naz)
         """
-        
+
         if load:
             fit_coeffs = np.load(f"{self.save_fn}_bess_fit_coeffs.npy")
             fit_beam = np.load(f"{self.save_fn}_bess_fit_beam.npy")
@@ -264,25 +289,32 @@ def get_fits(self, load=False, data_array=None):
                 data_array = self.data_array
             # az_modes are discretely orthonormal so just project onto the basis
             # Saves loads of memory and time
-            az_fit = data_array @ self.trig_matr.conj() # Naxes_vec, 1, Nfeeds, Nfreq, Nza, Nm
+            az_fit = (
+                data_array @ self.trig_matr.conj()
+            )  # Naxes_vec, 1, Nfeeds, Nfreq, Nza, Nm
 
             BtB = self.bess_matr.T @ self.bess_matr
-            Baz = self.bess_matr.T @ az_fit # Naxes_vec, 1, Nfeeds, Nfreq, Nn, Nm
-            Baz = Baz.transpose(4, 5, 0, 1, 2, 3) # Nn, Nm, Naxes_vec, 1, Nfeeds, Nfreq
+            Baz = self.bess_matr.T @ az_fit  # Naxes_vec, 1, Nfeeds, Nfreq, Nn, Nm
+            Baz = Baz.transpose(4, 5, 0, 1, 2, 3)  # Nn, Nm, Naxes_vec, 1, Nfeeds, Nfreq
 
-
-            fit_coeffs = solve(BtB, Baz, assume_a="sym") # Nn, Nm, Naxes_vec, 1, Nfeeds, Nfreq
+            fit_coeffs = solve(
+                BtB, Baz, assume_a="sym"
+            )  # Nn, Nm, Naxes_vec, 1, Nfeeds, Nfreq
 
             # Apply design matrices to get fit beams
-            fit_beam_az = np.tensordot(self.trig_matr, fit_coeffs, axes=((1,), (1,))) # Naz, Nn, Naxes_vec, 1, Nfeeds, Nfreq
-            fit_beam = np.tensordot(self.bess_matr, fit_beam_az, axes=((1,), (1,))) # Nza, Naz, Naxes_vec, 1, Nfeeds, Nfreq
+            fit_beam_az = np.tensordot(
+                self.trig_matr, fit_coeffs, axes=((1,), (1,))
+            )  # Naz, Nn, Naxes_vec, 1, Nfeeds, Nfreq
+            fit_beam = np.tensordot(
+                self.bess_matr, fit_beam_az, axes=((1,), (1,))
+            )  # Nza, Naz, Naxes_vec, 1, Nfeeds, Nfreq
             fit_beam = fit_beam.transpose(2, 3, 4, 5, 0, 1)
 
             np.save(f"{self.save_fn}_bess_fit_coeffs.npy", fit_coeffs)
-            np.save(f"{self.save_fn}_bess_fit_beam.npy", fit_beam)        
+            np.save(f"{self.save_fn}_bess_fit_beam.npy", fit_beam)
 
         return fit_coeffs, fit_beam
-    
+
     def get_comp_inds(self, make_const_in_freq=True):
         """
         Get the indices for the self.num_modes most significant modes for each
@@ -296,36 +328,50 @@ def get_comp_inds(self, make_const_in_freq=True):
 
         Returns:
             nmodes_comp (array, int):
-                The radial mode numbers corresponding to the top num_modes 
-                Fourier-Bessl modes, in descending order of significance. Has 
+                The radial mode numbers corresponding to the top num_modes
+                Fourier-Bessl modes, in descending order of significance. Has
                 shape (num_modes, Naxes_vec, 1, Nfeeds, Nfreqs).
-            mmodes_comp (array, int):    
-                The azimuthal modes numbers corresponding to the top num_modes 
+            mmodes_comp (array, int):
+                The azimuthal modes numbers corresponding to the top num_modes
                 Fourier-Bessel modes, in descending order of significance. Has
                 shape (num_modes, Naxes_vec, 1, Nfeeds, Nfreqs).
         """
 
-        ps_sort_inds = np.argsort(self.bess_ps.reshape((self.ncoeff_bess, 
-                                                        self.Naxes_vec, 1, 
-                                                        self.Nfeeds, 
-                                                        self.Nfreqs)),
-                                  axis=0)
+        ps_sort_inds = np.argsort(
+            self.bess_ps.reshape(
+                (self.ncoeff_bess, self.Naxes_vec, 1, self.Nfeeds, self.Nfreqs)
+            ),
+            axis=0,
+        )
         # Highest modes start from the end
-        sort_inds_flip = np.flip(ps_sort_inds, axis=0)[:self.num_modes_comp]
-        nmodes_comp, mmodes_comp = np.unravel_index(sort_inds_flip, 
-                                                    (self.nmax, len(self.mmodes)))
+        sort_inds_flip = np.flip(ps_sort_inds, axis=0)[: self.num_modes_comp]
+        nmodes_comp, mmodes_comp = np.unravel_index(
+            sort_inds_flip, (self.nmax, len(self.mmodes))
+        )
         if make_const_in_freq:
             mid_freq_ind = self.Nfreqs // 2
-            nmodes_comp = np.repeat(nmodes_comp[:, :, :, :, mid_freq_ind:mid_freq_ind + 1],
-                                    self.Nfreqs, axis=4)
-            mmodes_comp = np.repeat(mmodes_comp[:, :, :, :, mid_freq_ind:mid_freq_ind + 1],
-                                    self.Nfreqs, axis=4)
+            nmodes_comp = np.repeat(
+                nmodes_comp[:, :, :, :, mid_freq_ind : mid_freq_ind + 1],
+                self.Nfreqs,
+                axis=4,
+            )
+            mmodes_comp = np.repeat(
+                mmodes_comp[:, :, :, :, mid_freq_ind : mid_freq_ind + 1],
+                self.Nfreqs,
+                axis=4,
+            )
 
-        
         return nmodes_comp, mmodes_comp
-    
-    def sparse_fit_loop(self, do_fit=True, bess_matr=None, trig_matr=None,
-                        fit_coeffs=None, freq_array=None, data_array=None):
+
+    def sparse_fit_loop(
+        self,
+        do_fit=True,
+        bess_matr=None,
+        trig_matr=None,
+        fit_coeffs=None,
+        freq_array=None,
+        data_array=None,
+    ):
         """
         Do a loop over all the axes and fit/evaluate fit in position space.
 
@@ -338,7 +384,7 @@ def sparse_fit_loop(self, do_fit=True, bess_matr=None, trig_matr=None,
                 Bessel part of design matrix.
             trig_matr (array, complex):
                 Fourier part of design matrix.
-        
+
         Returns:
             fit_coeffs (array, complex; if do_fit is True):
                 The newly calculated fit coefficients in the sparse basis.
@@ -348,25 +394,28 @@ def sparse_fit_loop(self, do_fit=True, bess_matr=None, trig_matr=None,
         # nmodes might vary from pol to pol, freq to freq. The fit is fast, just do a big for loop.
         interp_kwargs = [bess_matr, trig_matr, fit_coeffs]
         if do_fit:
-            fit_coeffs = np.zeros([self.Naxes_vec, 1, self.Nfeeds, self.Nfreqs, 
-                                   self.num_modes_comp], dtype=complex)
+            fit_coeffs = np.zeros(
+                [self.Naxes_vec, 1, self.Nfeeds, self.Nfreqs, self.num_modes_comp],
+                dtype=complex,
+            )
             if data_array is None:
                 data_array = self.data_array
             beam_shape = data_array.shape
             bess_matr = self.bess_matr
             trig_matr = self.trig_matr
         elif any([item is None for item in interp_kwargs[:2]]):
-            raise ValueError("Must supply bess_matr, and trig_matr "
-                             "if not doing fit.")
+            raise ValueError(
+                "Must supply bess_matr, and trig_matr " "if not doing fit."
+            )
         else:
             Npos = bess_matr.shape[0]
             beam_shape = (self.Naxes_vec, 1, self.Nfeeds, self.Nfreqs, Npos)
-            if fit_coeffs is None: # already have some fits
+            if fit_coeffs is None:  # already have some fits
                 fit_coeffs = self.comp_fits
         fit_beam = np.zeros(beam_shape, dtype=complex)
-        
+
         Nfreqs = self.Nfreqs if freq_array is None else len(freq_array)
-        
+
         for vec_ind in range(self.Naxes_vec):
             for feed_ind in range(self.Nfeeds):
                 for freq_ind in range(Nfreqs):
@@ -375,52 +424,67 @@ def sparse_fit_loop(self, do_fit=True, bess_matr=None, trig_matr=None,
                     nmodes_iter = self.nmodes_comp[:, vec_ind, 0, feed_ind, freq_ind]
                     mmodes_iter = self.mmodes_comp[:, vec_ind, 0, feed_ind, freq_ind]
                     unique_mmodes_iter = np.unique(mmodes_iter)
-       
+
                     for mmode in unique_mmodes_iter:
                         mmode_inds = mmodes_iter == mmode
                         # Get the nmodes that this mmode is used for
-                        nmodes_mmode = nmodes_iter[mmode_inds] 
+                        nmodes_mmode = nmodes_iter[mmode_inds]
 
                         bess_matr_mmode = bess_matr[:, nmodes_mmode]
                         trig_mode = trig_matr[:, mmode]
 
                         if do_fit:
-                            az_fit_mmode = dat_iter @ trig_mode.conj() # Nza
+                            az_fit_mmode = dat_iter @ trig_mode.conj()  # Nza
 
                             fit_coeffs_mmode = lstsq(bess_matr_mmode, az_fit_mmode)[0]
-                            fit_coeffs[vec_ind, 0, feed_ind, freq_ind, mmode_inds] = fit_coeffs_mmode
-                            fit_beam[vec_ind, 0, feed_ind, freq_ind] += np.outer(bess_matr_mmode @ fit_coeffs_mmode, trig_mode)
+                            fit_coeffs[vec_ind, 0, feed_ind, freq_ind, mmode_inds] = (
+                                fit_coeffs_mmode
+                            )
+                            fit_beam[vec_ind, 0, feed_ind, freq_ind] += np.outer(
+                                bess_matr_mmode @ fit_coeffs_mmode, trig_mode
+                            )
                         else:
-                            fit_coeffs_mmode = fit_coeffs[vec_ind, 0, feed_ind, freq_ind, mmode_inds]
-                            fit_beam[vec_ind, 0, feed_ind, freq_ind] += (bess_matr_mmode @ fit_coeffs_mmode) * trig_mode
-                        
+                            fit_coeffs_mmode = fit_coeffs[
+                                vec_ind, 0, feed_ind, freq_ind, mmode_inds
+                            ]
+                            fit_beam[vec_ind, 0, feed_ind, freq_ind] += (
+                                bess_matr_mmode @ fit_coeffs_mmode
+                            ) * trig_mode
+
         if do_fit:
-            return fit_coeffs,fit_beam
+            return fit_coeffs, fit_beam
         else:
             return fit_beam
-    
-    def interp(self, sparse_fit=False, fit_coeffs=None, az_array=None, 
-               za_array=None, reuse_spline=False, freq_array=None,
-               freq_interp_kind="cubic",
-               **kwargs):
+
+    def interp(
+        self,
+        sparse_fit=False,
+        fit_coeffs=None,
+        az_array=None,
+        za_array=None,
+        reuse_spline=False,
+        freq_array=None,
+        freq_interp_kind="cubic",
+        **kwargs,
+    ):
         """
         A very paired down override of UVBeam.interp that more resembles
         pyuvsim.AnalyticBeam.interp. Any kwarg for UVBeam.interp that is not
         explicitly listed in this version of interp will do nothing.
 
         Parameters:
-            sparse_fit (bool): 
+            sparse_fit (bool):
                 Whether a sparse fit is being supplied. If False (default), just
                 uses the full fit specified at instantiation.
             fit_coeffs (bool):
-                The sparse fit coefficients being supplied if sparse_fit is 
+                The sparse fit coefficients being supplied if sparse_fit is
                 True.
-            az_array (array): 
+            az_array (array):
                 Flattened azimuth angles to interpolate to.
             za_array (array):
                 Flattened zenith angles to interpolate to.
             reuse_spline (bool):
-                Whether to reuse the spatial design matrix for a particular 
+                Whether to reuse the spatial design matrix for a particular
                 az_array and za_array (named to keep consistency with UVBeam).
             freq_array (array):
                 Frequencies to interpolate to. If None (default), just computes
@@ -432,20 +496,20 @@ def interp(self, sparse_fit=False, fit_coeffs=None, az_array=None,
 
         Returns:
             beam_vals (array, complex):
-                The values of the beam at the interpolated 
-                frequencies/spatial positions. Has shape 
+                The values of the beam at the interpolated
+                frequencies/spatial positions. Has shape
                 (Naxes_vec, 1, Npols, Nfreqs, Npos).
         """
         if az_array is None and za_array is None and freq_array is not None:
             # vis_cpu wants to get a new object with frequency freq_array. No can do. Return the whole thing.
             # FIXME: Can make a new_sparse_beam_from_self method to accomplish this
             return self
-    
+
         if az_array is None:
             raise ValueError("Must specify an azimuth array.")
         if za_array is None:
             raise ValueError("Must specify a zenith-angle array.")
-        
+
         if reuse_spline:
             az_hash = hashlib.sha1(az_array).hexdigest()
             za_hash = hashlib.sha1(za_array).hexdigest()
@@ -456,10 +520,10 @@ def interp(self, sparse_fit=False, fit_coeffs=None, az_array=None,
             else:
                 self.az_array_dict[az_hash] = az_array
                 self.za_array_dict[za_hash] = za_array
-    
+
                 bess_matr, trig_matr = self.get_dmatr_interp(az_array, za_array)
                 bt_matr = trig_matr[:, np.newaxis] * bess_matr[:, :, np.newaxis]
-                
+
                 self.trig_matr_interp_dict[az_hash] = trig_matr
                 self.bess_matr_interp_dict[za_hash] = bess_matr
                 self.bt_matr_interp_dict[(az_hash, za_hash)] = bt_matr
@@ -467,52 +531,62 @@ def interp(self, sparse_fit=False, fit_coeffs=None, az_array=None,
             bess_matr, trig_matr = self.get_dmatr_interp(az_array, za_array)
             bt_matr = trig_matr[:, np.newaxis] * bess_matr[:, :, np.newaxis]
 
-        
         if sparse_fit:
             if freq_array is None:
                 fit_coeffs = self.comp_fits
             else:
                 for ind_ob in [self.nmodes_comp, self.mmodes_comp]:
                     if not np.all(ind_ob == ind_ob[:, :, :, :, :1]):
-                        raise NotImplementedError("Basis is not constant in "
-                                                  "frequency. Cannot do "
-                                                  "frequency interpolation for "
-                                                  "sparse_fit=True")
-                freq_array, freq_array_knots = self.prep_freq_array_for_interp(freq_array)
-                    
+                        raise NotImplementedError(
+                            "Basis is not constant in "
+                            "frequency. Cannot do "
+                            "frequency interpolation for "
+                            "sparse_fit=True"
+                        )
+                freq_array, freq_array_knots = self.prep_freq_array_for_interp(
+                    freq_array
+                )
+
                 fit_coeffs_interp = interp1d(freq_array_knots, self.comp_fits, axis=3)
                 fit_coeffs = fit_coeffs_interp(freq_array)
-            beam_vals = self.sparse_fit_loop(do_fit=False, 
-                                             fit_coeffs=fit_coeffs, 
-                                             bess_matr=bess_matr,
-                                             trig_matr=trig_matr,
-                                             freq_array=freq_array)
+            beam_vals = self.sparse_fit_loop(
+                do_fit=False,
+                fit_coeffs=fit_coeffs,
+                bess_matr=bess_matr,
+                trig_matr=trig_matr,
+                freq_array=freq_array,
+            )
         else:
             if freq_array is None:
                 bess_fits = self.bess_fits
             else:
-                freq_array, freq_array_knots = self.prep_freq_array_for_interp(freq_array)
-                bess_fits_interp = interp1d(freq_array_knots, self.bess_fits, axis=5,
-                                            kind=freq_interp_kind)
+                freq_array, freq_array_knots = self.prep_freq_array_for_interp(
+                    freq_array
+                )
+                bess_fits_interp = interp1d(
+                    freq_array_knots, self.bess_fits, axis=5, kind=freq_interp_kind
+                )
                 bess_fits = bess_fits_interp(freq_array)
-            beam_vals = np.tensordot(bt_matr, bess_fits, axes=2).transpose(1, 2, 3, 4, 0)
+            beam_vals = np.tensordot(bt_matr, bess_fits, axes=2).transpose(
+                1, 2, 3, 4, 0
+            )
         if self.beam_type == "power":
             # FIXME: This assumes you are reading in a power beam and is just to get rid of the imaginary component
             beam_vals = np.abs(beam_vals)
-            
+
         return beam_vals, None
 
     def prep_freq_array_for_interp(self, freq_array):
         freq_array = np.atleast_1d(freq_array)
         assert freq_array.ndim == 1, "Freq array for interp must be exactly 1d"
-                
-                # FIXME: More explicit and complete future_array_shapes compatibility throughout code base desired
+
+        # FIXME: More explicit and complete future_array_shapes compatibility throughout code base desired
         if self.freq_array.ndim > 1:
             freq_array_knots = self.freq_array[0]
         else:
             freq_array_knots = self.freq_array
-        return freq_array,freq_array_knots
-    
+        return freq_array, freq_array_knots
+
     def clear_cache(self):
         self.az_array_dict.clear()
         self.za_array_dict.clear()
@@ -521,24 +595,28 @@ def clear_cache(self):
         self.bt_matr_interp_dict.clear()
 
         return
-    
+
     def efield_to_power(*args, **kwargs):
         raise NotImplementedError("efield_to_power is not implemented yet.")
-    
+
     def efield_to_pstokes(*args, **kwargs):
         raise NotImplementedError("efield_to_pstokes is not implemented yet.")
-    
+
     def sigmoid_mod(self, rad_array=None):
         if rad_array is None:
             rad_array = self.rad_array
-        za_array = np.arccos(1 - (self.alpha * rad_array)**2)
+        za_array = np.arccos(1 - (self.alpha * rad_array) ** 2)
         return 0.5 * (1 + np.tanh((za_array - self.za_ml) / self.dza))
 
     def sin_perts(self, rad_array=None):
         if rad_array is None:
             rad_array = self.rad_array
-        L = self.rad_array[-1] # Always make this zero-out at the horizon in unstretched coordinates
-        dmatr = np.array([np.sin(2 * np.pi * m * rad_array / L) for m in range(self.Nsin_pert)]).T
+        L = self.rad_array[
+            -1
+        ]  # Always make this zero-out at the horizon in unstretched coordinates
+        dmatr = np.array(
+            [np.sin(2 * np.pi * m * rad_array / L) for m in range(self.Nsin_pert)]
+        ).T
         sin_pert_unnorm = dmatr @ self.sin_pert_coeffs
         sp_range = np.amax(sin_pert_unnorm) - np.amin(sin_pert_unnorm)
         return sin_pert_unnorm / sp_range
@@ -547,23 +625,29 @@ def SL_pert(self, rad_array=None):
         if self.sin_pert_coeffs is None:
             return np.ones_like(rad_array)
         else:
-            return 1 + self.cSL * self.sin_perts(rad_array=rad_array) * self.sigmoid_mod(rad_array=rad_array)
+            return 1 + self.cSL * self.sin_perts(
+                rad_array=rad_array
+            ) * self.sigmoid_mod(rad_array=rad_array)
 
     def ML_gauss_term(self, gam):
-        return np.exp(-0.5 * self.axis2_array**2/(gam * self.za_ml)**2)
+        return np.exp(-0.5 * self.axis2_array**2 / (gam * self.za_ml) ** 2)
 
     def ML_pert(self):
-        sig_factor = (1 - self.sigmoid_mod())
-        gauss_diff = self.ML_gauss_term(gam=self.gam) -self. ML_gauss_term(gam=1)
+        sig_factor = 1 - self.sigmoid_mod()
+        gauss_diff = self.ML_gauss_term(gam=self.gam) - self.ML_gauss_term(gam=1)
         return sig_factor * gauss_diff
-    
-    def az_pert(self):
-        dmatr_cos = np.array([np.cos(m * self.axis1_array) for m in range(1, 1 + self.Naz_pert)]).T
-        dmatr_sin = np.array([np.sin(m * self.axis1_array) for m in range(1, 1 + self.Naz_pert)]).T
 
-        cos_modes = dmatr_cos @ self.az_cos_pert_coeffs 
+    def az_pert(self):
+        dmatr_cos = np.array(
+            [np.cos(m * self.axis1_array) for m in range(1, 1 + self.Naz_pert)]
+        ).T
+        dmatr_sin = np.array(
+            [np.sin(m * self.axis1_array) for m in range(1, 1 + self.Naz_pert)]
+        ).T
+
+        cos_modes = dmatr_cos @ self.az_cos_pert_coeffs
         sin_modes = dmatr_sin @ self.az_sin_pert_coeffs
 
         az_pert_unnorm = cos_modes + sin_modes
         az_pert_range = np.amax(az_pert_unnorm) - np.amin(az_pert_unnorm)
-        return az_pert_unnorm / az_pert_range
\ No newline at end of file
+        return az_pert_unnorm / az_pert_range
diff --git a/hydra/utils.py b/hydra/utils.py
index 97d47c1..982450f 100644
--- a/hydra/utils.py
+++ b/hydra/utils.py
@@ -1,4 +1,3 @@
-
 import numpy as np
 from matvis import conversions
 import pyuvdata
@@ -23,10 +22,10 @@ def flatten_vector(v, reduced_idxs=None):
     Flatten a complex vector with shape (N, Ntimes, Nfreq) into a block
     vector of shape (N x Ntimes x Nfreqs x 2), i.e. the real and imaginary
     blocks stuck together.
-    
+
     Parameters:
         reduced_idxs (array_like):
-            If specified, this is an array of indices that maps a reduced x 
+            If specified, this is an array of indices that maps a reduced x
             vector to the full x vector. The unspecified modes are set to zero.
     """
     # If only certain indices were kept, remove others
@@ -34,26 +33,27 @@ def flatten_vector(v, reduced_idxs=None):
         return np.concatenate((v.real.flatten(), v.imag.flatten()))[reduced_idxs]
     else:
         return np.concatenate((v.real.flatten(), v.imag.flatten()))
-    
 
 
 def reconstruct_vector(v, shape, reduced_idxs=None):
     """
     Undo the flattening of a complex vector.
-    
+
     Parameters:
         reduced_idxs (array_like):
-            If specified, this is an array of indices that maps a reduced x 
-            vector to the full x vector expected by the linear operator. The 
+            If specified, this is an array of indices that maps a reduced x
+            vector to the full x vector expected by the linear operator. The
             unspecified modes are set to zero.
     """
     # If only certain indices were kept, populate those into a vector of zeros
     if reduced_idxs is not None:
-        vv = np.zeros(2*np.prod(shape), dtype=v.dtype) # required size = product of shape tuple
+        vv = np.zeros(
+            2 * np.prod(shape), dtype=v.dtype
+        )  # required size = product of shape tuple
         vv[reduced_idxs] = v[:]
     else:
         vv = v
-        
+
     # Now unpack the full-sized array
     y = vv[: vv.size // 2] + 1.0j * vv[vv.size // 2 :]
     return y.reshape(shape)
@@ -77,8 +77,8 @@ def apply_gains(v, gains, ants, antpairs, perturbation=None, inline=False):
         antpairs (list of tuples):
             List of antenna pair tuples.
         perturbation (array_like):
-            Linear perturbations to the gains. If set, this will apply these 
-            perturbations under the linear approximation, i.e. 
+            Linear perturbations to the gains. If set, this will apply these
+            perturbations under the linear approximation, i.e.
             `g_i g_j* \approx \bar{g}_i \bar{g}_j^* (1 + x_i + x_j^*)`.
             Expected shape is (Nants, Nfreqs, Ntimes).
         inline (bool):
@@ -93,22 +93,23 @@ def apply_gains(v, gains, ants, antpairs, perturbation=None, inline=False):
         ggv = v
     else:
         ggv = v.copy()
-    assert v.shape[0] == len(antpairs), \
-        "Input array `v` has shape that is incompatible with `antpairs`"
+    assert v.shape[0] == len(
+        antpairs
+    ), "Input array `v` has shape that is incompatible with `antpairs`"
 
     # Apply gains
     for k, bl in enumerate(antpairs):
         ant1, ant2 = bl
         i1 = np.where(ants == ant1)[0][0]
         i2 = np.where(ants == ant2)[0][0]
-        fac = 1.
+        fac = 1.0
         if perturbation is not None:
-            fac = 1. + perturbation[i1] + perturbation[i2].conj()
-        ggv[k,:,:] *= gains[i1] * gains[i2].conj() * fac
+            fac = 1.0 + perturbation[i1] + perturbation[i2].conj()
+        ggv[k, :, :] *= gains[i1] * gains[i2].conj() * fac
     return ggv
 
 
-def load_gain_model(gain_model_file, lst_pad=[0,0], freq_pad=[0,0], pad_value=1.):
+def load_gain_model(gain_model_file, lst_pad=[0, 0], freq_pad=[0, 0], pad_value=1.0):
     """
     Load complex gain model for each antenna into a single array, and
     zero-pad the edges of the array in the time and frequency dimensions
@@ -138,15 +139,19 @@ def load_gain_model(gain_model_file, lst_pad=[0,0], freq_pad=[0,0], pad_value=1.
     orig_model = np.load(gain_model_file)
 
     # Pad the array as requested
-    padded_shape = (orig_model.shape[0],
-                    orig_model.shape[1] + freq_pad[0] + freq_pad[1],
-                    orig_model.shape[2] + lst_pad[0] + lst_pad[1],)
+    padded_shape = (
+        orig_model.shape[0],
+        orig_model.shape[1] + freq_pad[0] + freq_pad[1],
+        orig_model.shape[2] + lst_pad[0] + lst_pad[1],
+    )
     gain_model = np.zeros(padded_shape, dtype=orig_model.dtype) + pad_value
 
     # Put gain model into padded array
-    gain_model[:,
-               freq_pad[0]:orig_model.shape[1]+freq_pad[0],
-               lst_pad[0]:orig_model.shape[2]+lst_pad[0]] = orig_model[:,:,:]
+    gain_model[
+        :,
+        freq_pad[0] : orig_model.shape[1] + freq_pad[0],
+        lst_pad[0] : orig_model.shape[2] + lst_pad[0],
+    ] = orig_model[:, :, :]
     return gain_model
 
 
@@ -188,11 +193,11 @@ def extract_vis_from_sim(ants, antpairs, sim_vis):
         # Extract data for this antenna pair
         idx1 = np.where(ants == ant1)[0][0]
         idx2 = np.where(ants == ant2)[0][0]
-        vis[i,:,:] = sim_vis[:,:,idx1,idx2]
+        vis[i, :, :] = sim_vis[:, :, idx1, idx2]
     return vis
 
 
-def extract_vis_from_uvdata(uvd, exclude_autos=True, lst_pad=[0,0], freq_pad=[0,0]):
+def extract_vis_from_uvdata(uvd, exclude_autos=True, lst_pad=[0, 0], freq_pad=[0, 0]):
     """
     Extract only the desired set of visibilities from a UVData object, in
     the desired order.
@@ -221,8 +226,8 @@ def extract_vis_from_uvdata(uvd, exclude_autos=True, lst_pad=[0,0], freq_pad=[0,
     ants = []
     antpairs = []
     vis = []
-    uvd.conjugate_bls(convention='ant1<ant2') # conjugate baselines
-    for antpair, bl_data in uvd.antpairpol_iter(squeeze='full'):
+    uvd.conjugate_bls(convention="ant1<ant2")  # conjugate baselines
+    for antpair, bl_data in uvd.antpairpol_iter(squeeze="full"):
         ant1 = antpair[0]
         ant2 = antpair[1]
         Nfreqs = bl_data.shape[1]
@@ -230,15 +235,17 @@ def extract_vis_from_uvdata(uvd, exclude_autos=True, lst_pad=[0,0], freq_pad=[0,
 
         # Add antpair to the list
         if ant1 == ant2 and exclude_autos:
-            continue # exclude autos
+            continue  # exclude autos
         antpairs.append((ant1, ant2))
 
         # Add data for this bl to array (with zero-padding if necessary)
-        dat = np.zeros((freq_pad[0] + Nfreqs + freq_pad[1],
-                        lst_pad[0] + Ntimes + lst_pad[1]),
-                       dtype=bl_data.dtype)
-        dat[freq_pad[0]:Nfreqs+freq_pad[0],
-            lst_pad[0]:Ntimes+lst_pad[0]] = bl_data[:,:].T
+        dat = np.zeros(
+            (freq_pad[0] + Nfreqs + freq_pad[1], lst_pad[0] + Ntimes + lst_pad[1]),
+            dtype=bl_data.dtype,
+        )
+        dat[freq_pad[0] : Nfreqs + freq_pad[0], lst_pad[0] : Ntimes + lst_pad[0]] = (
+            bl_data[:, :].T
+        )
         vis.append(dat)
 
         # Add antenna to list if not there already
@@ -250,7 +257,7 @@ def extract_vis_from_uvdata(uvd, exclude_autos=True, lst_pad=[0,0], freq_pad=[0,
     return np.array(vis), antpairs, np.array(ants)
 
 
-def extend_coords_with_padding(arr, pad=[0,0]):
+def extend_coords_with_padding(arr, pad=[0, 0]):
     """
     Take an array with (e.g.) frequencies or times and extend it by a number
     of elements of padding on either side. The values are extrapolated into
@@ -271,12 +278,12 @@ def extend_coords_with_padding(arr, pad=[0,0]):
     """
     # Create new copy of the array with zero padding around it
     arr_new = np.zeros(pad[0] + arr.size + pad[1], dtype=arr.dtype)
-    arr_new[pad[0]:arr_new.size-pad[1]] = arr[:]
+    arr_new[pad[0] : arr_new.size - pad[1]] = arr[:]
 
     # Extrapolate into the padded regions
     diff = arr[1] - arr[0]
-    arr_new[arr_new.size-pad[1]:] = arr[-1] + diff * (1.+np.arange(pad[1]))
-    arr_new[:pad[0]] = arr[0] - diff * (np.arange(pad[0]) + 1.)[::-1]
+    arr_new[arr_new.size - pad[1] :] = arr[-1] + diff * (1.0 + np.arange(pad[1]))
+    arr_new[: pad[0]] = arr[0] - diff * (np.arange(pad[0]) + 1.0)[::-1]
     return arr_new
 
 
@@ -304,8 +311,8 @@ def timing_info(fname, iter, task, duration, verbose=True):
 
 def freqs_times_for_worker(comm, freqs, times, fchunks, tchunks=1):
     """
-    Get the unique frequency and time chunk for a worker. The workers 
-    are arranged onto an `fchunks * tchunks` grid, and each worker is 
+    Get the unique frequency and time chunk for a worker. The workers
+    are arranged onto an `fchunks * tchunks` grid, and each worker is
     given the corresponding chunk of the 2D `freqs * times` grid.
 
     Parameters:
@@ -326,25 +333,27 @@ def freqs_times_for_worker(comm, freqs, times, fchunks, tchunks=1):
         time_idx_chunk (array_like):
             The indices of the time array belonging to this worker.
         worker_map (dict):
-            A dictionary containing the chunk indices and frequency and 
-            time indices for each worker. This is useful for cases where 
-            individual workers have to communicate with each other. The 
+            A dictionary containing the chunk indices and frequency and
+            time indices for each worker. This is useful for cases where
+            individual workers have to communicate with each other. The
             format of the dictionary entries is:
             `worker_id: (freq_chunk_idx, time_chunk_idx, freq_idxs, time_idxs)`
-            where `freq_chunk_idx` and `time_chunk_idx` are indices in the 
-            grid of chunks for this worker, and `freq_idxs` and `time_idxs` 
-            are the actual indices in the `freqs` and `times` arrays belonging 
+            where `freq_chunk_idx` and `time_chunk_idx` are indices in the
+            grid of chunks for this worker, and `freq_idxs` and `time_idxs`
+            are the actual indices in the `freqs` and `times` arrays belonging
             to this worker.
     """
     myid = comm.Get_rank()
     nworkers = comm.Get_size()
-    assert nworkers <= fchunks * tchunks, "There are more workers than time and frequency chunks"
+    assert (
+        nworkers <= fchunks * tchunks
+    ), "There are more workers than time and frequency chunks"
 
     # Get chunk ID for this worker
     allidxs = np.arange(fchunks * tchunks).reshape((fchunks, tchunks))
     fidx, tidx = np.where(allidxs == myid)
     fidx, tidx = int(fidx[0]), int(tidx[0])
-    
+
     # Get chunk of freq/time idxs for each worker
     freq_idxs = np.arange(freqs.size)
     time_idxs = np.arange(times.size)
@@ -361,7 +370,7 @@ def freqs_times_for_worker(comm, freqs, times, fchunks, tchunks=1):
     return freq_idx_chunks[fidx], time_idx_chunks[tidx], worker_map
 
 
-def build_hex_array(hex_spec=(3,4), ants_per_row=None, d=14.6):
+def build_hex_array(hex_spec=(3, 4), ants_per_row=None, d=14.6):
     """
     Build an antenna position dict for a hexagonally close-packed array.
 
@@ -390,21 +399,21 @@ def build_hex_array(hex_spec=(3,4), ants_per_row=None, d=14.6):
 
     # If ants_per_row isn't given, build it from hex_spec
     if ants_per_row is None:
-        r = np.arange(hex_spec[0], hex_spec[1]+1).tolist()
+        r = np.arange(hex_spec[0], hex_spec[1] + 1).tolist()
         ants_per_row = r[:-1] + r[::-1]
 
     # Assign antennas
     k = -1
-    y = 0.
-    dy = d * np.sqrt(3) / 2. # delta y = d sin(60 deg)
+    y = 0.0
+    dy = d * np.sqrt(3) / 2.0  # delta y = d sin(60 deg)
     for j, r in enumerate(ants_per_row):
 
         # Calculate y coord and x offset
-        y = -0.5 * dy * (len(ants_per_row)-1) + dy * j
-        x = np.linspace(-d*(r-1)/2., d*(r-1)/2., r)
+        y = -0.5 * dy * (len(ants_per_row) - 1) + dy * j
+        x = np.linspace(-d * (r - 1) / 2.0, d * (r - 1) / 2.0, r)
         for i in range(r):
             k += 1
-            ants[k] = (x[i], y, 0.)
+            ants[k] = (x[i], y, 0.0)
 
     return ants
 
@@ -445,10 +454,12 @@ def convert_to_tops(ra, dec, lsts, latitude, precision=1):
         tys.append(ty)
         tzs.append(tz)
 
-    return(np.array(txs), np.array(tys), np.array(tzs))
+    return (np.array(txs), np.array(tys), np.array(tzs))
 
 
-def get_flux_from_ptsrc_amp(ptsrc_amps, freqs, beta_ptsrc, curv_ptsrc=None, ref_freq=100.):
+def get_flux_from_ptsrc_amp(
+    ptsrc_amps, freqs, beta_ptsrc, curv_ptsrc=None, ref_freq=100.0
+):
     """
     Calculate flux as a function of frequency for each point source, assuming
     a power-law spectrum.
@@ -481,9 +492,11 @@ def get_flux_from_ptsrc_amp(ptsrc_amps, freqs, beta_ptsrc, curv_ptsrc=None, ref_
     assert ptsrc_amps.size == curv_ptsrc.size
 
     # Make SEDs
-    spec_idx = beta_ptsrc[:,np.newaxis] + curv_ptsrc[:,np.newaxis] \
-                                          * np.log(freqs / ref_freq)[np.newaxis,:]
-    fluxes = ptsrc_amps[:,np.newaxis] * (freqs / ref_freq)[np.newaxis,:]**spec_idx
+    spec_idx = (
+        beta_ptsrc[:, np.newaxis]
+        + curv_ptsrc[:, np.newaxis] * np.log(freqs / ref_freq)[np.newaxis, :]
+    )
+    fluxes = ptsrc_amps[:, np.newaxis] * (freqs / ref_freq)[np.newaxis, :] ** spec_idx
     return fluxes
 
 
@@ -503,14 +516,17 @@ def antenna_dict_from_uvd(uvd):
             in ENU coordinates local to the array (values).
     """
     # Get ENU positions of antennas
-    enu = pyuvdata.utils.ENU_from_ECEF(uvd.antenna_positions + uvd.telescope_location,
-                                       *uvd.telescope_location_lat_lon_alt)
+    enu = pyuvdata.utils.ENU_from_ECEF(
+        uvd.antenna_positions + uvd.telescope_location,
+        *uvd.telescope_location_lat_lon_alt
+    )
 
     # Get antenna IDs
     ant_ids = uvd.antenna_numbers
 
-    assert ant_ids.size == enu.shape[0], \
-        "antenna_number and antenna_positions are mis-matched"
+    assert (
+        ant_ids.size == enu.shape[0]
+    ), "antenna_number and antenna_positions are mis-matched"
 
     # Output dict
     ants = {}
@@ -539,18 +555,22 @@ def get_beam_interp_shape(beam):
         raise ValueError("beam object is not AnalyticBeam or UVBeam object.")
 
     return spw_axis_present
-    
 
-def gain_prior_pspec_sqrt(lsts, freqs, 
-                          gain_prior_amp, 
-                          gain_prior_sigma_frate=None, 
-                          gain_prior_sigma_delay=None, 
-                          gain_prior_zeropoint_std=None,
-                          frate0=0., delay0=0.):
+
+def gain_prior_pspec_sqrt(
+    lsts,
+    freqs,
+    gain_prior_amp,
+    gain_prior_sigma_frate=None,
+    gain_prior_sigma_delay=None,
+    gain_prior_zeropoint_std=None,
+    frate0=0.0,
+    delay0=0.0,
+):
     """
-    Construct a gain prior power spectrum, which can be uniform or have a 
+    Construct a gain prior power spectrum, which can be uniform or have a
     Gaussian taper in the delay and/or fringe rate directions.
-    
+
     Parameters:
         lsts (array_like):
             LST grid (same as for the data), in radians.
@@ -559,10 +579,10 @@ def gain_prior_pspec_sqrt(lsts, freqs,
         gain_prior_amp (float):
             Amplitude of the prior power spectrum.
         gain_prior_sigma_frate (float):
-            Width of a Gaussian prior in fringe rate, in units of mHz. If 
+            Width of a Gaussian prior in fringe rate, in units of mHz. If
             `None`, the prior will be uniform in fringe rate.
         gain_prior_sigma_delay (float):
-            Width of a Gaussian prior in delay, in units of ns. If `None`, the 
+            Width of a Gaussian prior in delay, in units of ns. If `None`, the
             prior will be uniform in delay.
         gain_prior_zeropoint_std (float):
             If not `None`, fix the std. dev. of the (0,0) mode to some value.
@@ -570,37 +590,38 @@ def gain_prior_pspec_sqrt(lsts, freqs,
             The central fringe rate of the Gaussian taper (mHz).
         delay0 (float):
             The central delay of the Gaussian taper (ns).
-    
+
     Returns:
         gain_pspec_sqrt (array_like):
-            Square-root of the gain prior power spectrum, in delay-fringe rate 
+            Square-root of the gain prior power spectrum, in delay-fringe rate
             space (FFT ordering). Shape (Nfreqs, Ntimes).
     """
-    times = 24.*60.*60. * lsts / (2.*np.pi) # seconds
-    frate = 1e3 * np.fft.fftfreq(times.size, d=times[1] - times[0]) # mHz
-    delay = 1e3 * np.fft.fftfreq(freqs.size, d=freqs[1] - freqs[0]) # ns
-    
-    # Start with flat prior and then multiply by 
+    times = 24.0 * 60.0 * 60.0 * lsts / (2.0 * np.pi)  # seconds
+    frate = 1e3 * np.fft.fftfreq(times.size, d=times[1] - times[0])  # mHz
+    delay = 1e3 * np.fft.fftfreq(freqs.size, d=freqs[1] - freqs[0])  # ns
+
+    # Start with flat prior and then multiply by
     # Ordering is usual FFT ordering
     gain_pspec_sqrt = gain_prior_amp * np.ones((freqs.size, times.size))
     if gain_prior_sigma_frate is not None:
-        xt = (frate[np.newaxis,:] - frate0) / gain_prior_sigma_frate
-        gain_pspec_sqrt *= np.exp(-0.5 * xt**2.)
+        xt = (frate[np.newaxis, :] - frate0) / gain_prior_sigma_frate
+        gain_pspec_sqrt *= np.exp(-0.5 * xt**2.0)
     if gain_prior_sigma_delay is not None:
-        xf = (delay[:,np.newaxis] - delay0) / gain_prior_sigma_delay
-        gain_pspec_sqrt *= np.exp(-0.5 * xf**2.)
-    
+        xf = (delay[:, np.newaxis] - delay0) / gain_prior_sigma_delay
+        gain_pspec_sqrt *= np.exp(-0.5 * xf**2.0)
+
     # Zero-point prior
     if gain_prior_zeropoint_std is not None:
-        gain_pspec_sqrt[delay==0., frate == 0.] = gain_prior_zeropoint_std
-    
+        gain_pspec_sqrt[delay == 0.0, frate == 0.0] = gain_prior_zeropoint_std
+
     return gain_pspec_sqrt
 
 
-def partial_fourier_basis_2d(freqs, times, nfreq, ntime, Lfreq, Ltime, freq0=None, time0=None, 
-                             shape0=None):
+def partial_fourier_basis_2d(
+    freqs, times, nfreq, ntime, Lfreq, Ltime, freq0=None, time0=None, shape0=None
+):
     """
-    Construct a set of 2D Fourier modes from a list of wavenumber integers, 
+    Construct a set of 2D Fourier modes from a list of wavenumber integers,
     to form an incomplete set of 2D Fourier modes.
     """
     # Decide on origin of frequency axis for FT
@@ -609,9 +630,9 @@ def partial_fourier_basis_2d(freqs, times, nfreq, ntime, Lfreq, Ltime, freq0=Non
     if freq0 is None:
         freq0 = freqs[0]
 
-    # Determine normalising factors. If being used as a standalone Fourier operator, 
-    # these are just the lengths of the freq and time arrays. If being used as a 
-    # chunk of a Fourier operator across multiple workers, use the overall shape 
+    # Determine normalising factors. If being used as a standalone Fourier operator,
+    # these are just the lengths of the freq and time arrays. If being used as a
+    # chunk of a Fourier operator across multiple workers, use the overall shape
     # from 'shape0'
     if shape0 is None:
         Nfreqs = freqs.size
@@ -629,24 +650,37 @@ def partial_fourier_basis_2d(freqs, times, nfreq, ntime, Lfreq, Ltime, freq0=Non
     t2d, f2d = np.meshgrid(times - time0, freqs - freq0)
 
     # Calculate wavenumbers for each mode
-    kfreq = 2. * np.pi * nfreq / Lfreq # inverse freq. units
-    ktime = 2. * np.pi * ntime / Ltime # inverse time units
+    kfreq = 2.0 * np.pi * nfreq / Lfreq  # inverse freq. units
+    ktime = 2.0 * np.pi * ntime / Ltime  # inverse time units
 
     # Shape: (Nmodes, Nfreqs, Ntimes)
-    basis_fns = np.exp(1.j \
-                        * (  (kfreq[:,np.newaxis,np.newaxis] * f2d[np.newaxis,:,:]) \
-                           + (ktime[:,np.newaxis,np.newaxis] * t2d[np.newaxis,:,:])) ) \
-              / np.sqrt(Nfreqs * Ntimes)
+    basis_fns = np.exp(
+        1.0j
+        * (
+            (kfreq[:, np.newaxis, np.newaxis] * f2d[np.newaxis, :, :])
+            + (ktime[:, np.newaxis, np.newaxis] * t2d[np.newaxis, :, :])
+        )
+    ) / np.sqrt(Nfreqs * Ntimes)
     return basis_fns, kfreq, ktime
 
 
-def partial_fourier_basis_2d_from_nmax(freqs, times, nmaxfreq, nmaxtime, Lfreq, Ltime,
-                                       freq0=None, time0=None, shape0=None, positive_only=False):
-    """
-    Convenience function to construct a set of 2D Fourier modes with wavenumbers 
+def partial_fourier_basis_2d_from_nmax(
+    freqs,
+    times,
+    nmaxfreq,
+    nmaxtime,
+    Lfreq,
+    Ltime,
+    freq0=None,
+    time0=None,
+    shape0=None,
+    positive_only=False,
+):
+    """
+    Convenience function to construct a set of 2D Fourier modes with wavenumbers
     between -nmax <= 0 < nmax.
     """
-    # Make grid of wavenumber values for both frequency and time, in the range 
+    # Make grid of wavenumber values for both frequency and time, in the range
     # -nmax <= 0 < nmax (matches the modes you get from an FFT for 2*nmax points)
     if positive_only:
         _nfreq = np.arange(0, nmaxfreq)
@@ -658,15 +692,16 @@ def partial_fourier_basis_2d_from_nmax(freqs, times, nmaxfreq, nmaxtime, Lfreq,
 
     # Use generic function to construct Fourier basis functions
     basis_fns, kfreq, ktime = partial_fourier_basis_2d(
-                                         freqs=freqs, 
-                                         times=times, 
-                                         nfreq=nfreq.flatten(), 
-                                         ntime=ntime.flatten(), 
-                                         Lfreq=Lfreq, 
-                                         Ltime=Ltime,
-                                         freq0=freq0, 
-                                         time0=time0,
-                                         shape0=shape0)
+        freqs=freqs,
+        times=times,
+        nfreq=nfreq.flatten(),
+        ntime=ntime.flatten(),
+        Lfreq=Lfreq,
+        Ltime=Ltime,
+        freq0=freq0,
+        time0=time0,
+        shape0=shape0,
+    )
     return basis_fns, kfreq, ktime
 
 
@@ -674,18 +709,18 @@ def status(myid, message, colour=None):
     """
     Print a status message.
     """
-    #terminal_ = '\033[95m'
-    endchar = '\033[0m'
+    # terminal_ = '\033[95m'
+    endchar = "\033[0m"
     colours = {
-        'r':    '\033[91m',
-        'g':    '\033[92m',
-        'y':    '\033[93m',
-        'b':    '\033[94m',
-        'm':    '\033[95m',
-        'c':    '\033[96m',
-        'bold': '\033[1m',
-        'ul':   '\033[4m',
-        }
+        "r": "\033[91m",
+        "g": "\033[92m",
+        "y": "\033[93m",
+        "b": "\033[94m",
+        "m": "\033[95m",
+        "c": "\033[96m",
+        "bold": "\033[1m",
+        "ul": "\033[4m",
+    }
 
     # Worker ID, if present
     if myid is None:
diff --git a/hydra/vis_sampler.py b/hydra/vis_sampler.py
index 10cd3f4..912211b 100644
--- a/hydra/vis_sampler.py
+++ b/hydra/vis_sampler.py
@@ -85,10 +85,10 @@ def apply_operator(v, inv_noise_var, sqrt_pspec, group_id, gains, ants, antpairs
     assert len(antpairs) == v.shape[0]
 
     # Apply sqrt power spectrum, inverse FFT, and divide by noise variance
-    y = apply_sqrt_pspec(sqrt_pspec=sqrt_pspec,
-                         v=v,
-                         group_id=group_id,
-                         ifft=True) * inv_noise_var
+    y = (
+        apply_sqrt_pspec(sqrt_pspec=sqrt_pspec, v=v, group_id=group_id, ifft=True)
+        * inv_noise_var
+    )
 
     # Multiply by gains, then FFT back
     for k, bl in enumerate(antpairs):
@@ -100,16 +100,14 @@ def apply_operator(v, inv_noise_var, sqrt_pspec, group_id, gains, ants, antpairs
 
         # Apply S^1/2 to input vector, transform to data space, multiply by
         # gain product, and divide by noise variance
-        y[k, :, :] = fft.fft2(  y[k,:,:]
-                              * gains[i1] * gains[i2].conj()
-                              * gains[i2] * gains[i1].conj()
-                              )
+        y[k, :, :] = fft.fft2(
+            y[k, :, :] * gains[i1] * gains[i2].conj() * gains[i2] * gains[i1].conj()
+        )
 
     # Return full application of the LHS operator
-    return v + apply_sqrt_pspec(sqrt_pspec=sqrt_pspec,
-                                v=y,
-                                group_id=group_id,
-                                ifft=False)
+    return v + apply_sqrt_pspec(
+        sqrt_pspec=sqrt_pspec, v=y, group_id=group_id, ifft=False
+    )
 
 
 def construct_rhs(
@@ -189,7 +187,7 @@ def construct_rhs(
         i1 = np.where(ants == ant1)[0][0]
         i2 = np.where(ants == ant2)[0][0]
 
-        y[k,:,:] = fft.fft2(y[k,:,:] * gains[i1].conj() * gains[i2])
+        y[k, :, :] = fft.fft2(y[k, :, :] * gains[i1].conj() * gains[i2])
 
     # Apply sqrt(S) operator
     y = apply_sqrt_pspec(sqrt_pspec, y, group_id, ifft=False)
diff --git a/hydra/vis_simulator.py b/hydra/vis_simulator.py
index b9f3bac..fb76912 100644
--- a/hydra/vis_simulator.py
+++ b/hydra/vis_simulator.py
@@ -2,6 +2,7 @@
 Trimmed-down version of matvis that can be modified to return fragments of the
 visibilities.
 """
+
 import numpy as np
 import warnings
 from astropy.constants import c
@@ -14,6 +15,7 @@
 from datetime import datetime
 from . import utils
 
+
 def run_checks(
     antpos: np.ndarray,
     freq: float,
@@ -75,9 +77,11 @@ def run_checks(
     if beam_list is not None:
         # make sure we interpolate to the right frequency first.
         beam_list = [
-            bm.interp(freq_array=np.array([freq]), new_object=True, run_check=False)
-            if isinstance(bm, UVBeam)
-            else bm
+            (
+                bm.interp(freq_array=np.array([freq]), new_object=True, run_check=False)
+                if isinstance(bm, UVBeam)
+                else bm
+            )
             for bm in beam_list
         ]
 
@@ -108,7 +112,7 @@ def vis_sim_per_source(
     polarized: bool = False,
     beam_idx: Optional[np.ndarray] = None,
     subarr_ant=None,
-    force_no_beam_sqrt=False
+    force_no_beam_sqrt=False,
 ):
     """
     Calculate visibility from an input intensity map and beam model. This is
@@ -157,7 +161,7 @@ def vis_sim_per_source(
             Optional length-NANT array specifying a beam index for each antenna.
             By default, either a single beam is assumed to apply to all
             antennas or each antenna gets its own beam.
-        subarr_ant (int): 
+        subarr_ant (int):
             Used to calculate only those visibilities associated
             with a particular antenna.
         force_no_beam_sqrt (bool):
@@ -227,8 +231,8 @@ def vis_sim_per_source(
         # Simulate even if sources are below the horizon, since we need a
         # visibility per source regardless
         above_horizon = tz > 0
-        #tx = tx[above_horizon]
-        #ty = ty[above_horizon]
+        # tx = tx[above_horizon]
+        # ty = ty[above_horizon]
         nsrcs_up = len(tx)
 
         A_s = np.zeros((nax, nfeed, nbeam, nsrcs_up), dtype=complex_dtype)
@@ -275,8 +279,8 @@ def vis_sim_per_source(
             raise ValueError("Beam interpolation resulted in an invalid value")
 
         # Calculate delays, where tau = (b * s) / c
-        #np.dot(antpos, crd_top[:, above_horizon], out=tau)
-        np.dot(antpos, crd_top[:,:], out=tau)
+        # np.dot(antpos, crd_top[:, above_horizon], out=tau)
+        np.dot(antpos, crd_top[:, :], out=tau)
         tau /= c.value
 
         # Component of complex phase factor for one antenna
@@ -285,15 +289,15 @@ def vis_sim_per_source(
         np.exp(1.0j * (ang_freq * tau), out=v)
 
         # Complex voltages.
-        #v *= Isqrt[above_horizon]
+        # v *= Isqrt[above_horizon]
         v *= Isqrt[:]
-        v[:,~above_horizon] *= 0. # zero-out sources below the horizon
+        v[:, ~above_horizon] *= 0.0  # zero-out sources below the horizon
 
         # Compute visibilities using product of complex voltages (upper triangle).
         # Input arrays have shape (Nax, Nfeed, [Nants], Nsrcs
         v = A_s[:, :, beam_idx] * v[np.newaxis, np.newaxis, :]
 
-        # If a subarray is requested, only compute the visibilities that involve 
+        # If a subarray is requested, only compute the visibilities that involve
         # a specified antenna (useful for beam computations etc.)
         if subarr_ant is None:
             for i in range(len(antpos)):
@@ -308,9 +312,9 @@ def vis_sim_per_source(
         else:
             # Get the ones where the antenna in question is not conjugated
             vis[:, :, t] = np.einsum(
-                "jiln,jkmn->ikln", # summing over m just sums over one antenna (squeezes an axis)
+                "jiln,jkmn->ikln",  # summing over m just sums over one antenna (squeezes an axis)
                 v[:, :, :, :].conj(),
-                v[:, :, subarr_ant: subarr_ant + 1, :],
+                v[:, :, subarr_ant : subarr_ant + 1, :],
                 optimize=True,
             )
 
@@ -331,7 +335,7 @@ def simulate_vis_per_source(
     latitude=-30.7215 * np.pi / 180.0,
     use_feed="x",
     subarr_ant=None,
-    force_no_beam_sqrt=False
+    force_no_beam_sqrt=False,
 ):
     """
     Run a basic simulation, returning the visibility for each source
@@ -391,10 +395,16 @@ def simulate_vis_per_source(
     ), "The `fluxes` array must have shape (NSRCS, NFREQS)."
 
     # Check RA and Dec ranges
-    if np.any(ra < 0.) or np.any(ra > 2.*np.pi):
-        warnings.warn("One or more ra values is outside the allowed range (0, 2 pi)", RuntimeWarning)
+    if np.any(ra < 0.0) or np.any(ra > 2.0 * np.pi):
+        warnings.warn(
+            "One or more ra values is outside the allowed range (0, 2 pi)",
+            RuntimeWarning,
+        )
     if np.any(dec < -np.pi) or np.any(dec > np.pi):
-        warnings.warn("One or more dec values is outside the allowed range (-pi, +pi)", RuntimeWarning)
+        warnings.warn(
+            "One or more dec values is outside the allowed range (-pi, +pi)",
+            RuntimeWarning,
+        )
 
     # Determine precision
     if precision == 1:
@@ -430,7 +440,7 @@ def simulate_vis_per_source(
     if polarized:
         vis_shape = (naxes, nfeeds, freqs.size, lsts.size, nants, nants, nsrcs)
     elif subarr_ant is not None:
-    # When polarized beams implemented, need to have similar block in polarized case above
+        # When polarized beams implemented, need to have similar block in polarized case above
         vis_shape = (freqs.size, lsts.size, nants, nsrcs)
     else:
         vis_shape = (freqs.size, lsts.size, nants, nants, nsrcs)
@@ -440,32 +450,31 @@ def simulate_vis_per_source(
     vv = np.zeros_like(vis)
     for i in range(len(freqs)):
         vv[i] = vis_sim_per_source(
-                              antpos,
-                              freqs[i],
-                              eq2tops,
-                              crd_eq,
-                              fluxes[:, i],
-                              beam_list=beams,
-                              precision=precision,
-                              polarized=polarized,
-                              subarr_ant=subarr_ant,
-                              force_no_beam_sqrt=force_no_beam_sqrt
-                             )
+            antpos,
+            freqs[i],
+            eq2tops,
+            crd_eq,
+            fluxes[:, i],
+            beam_list=beams,
+            precision=precision,
+            polarized=polarized,
+            subarr_ant=subarr_ant,
+            force_no_beam_sqrt=force_no_beam_sqrt,
+        )
 
     # Assign returned values to array
     for i in range(freqs.size):
         if polarized:
             vis[:, :, i] = vv[i]  # v.shape: (nax, nfeed, ntimes, nant, nant, nsrcs)
         else:
-            vis[i] = vv[i]  # v.shape: (ntimes, nant, nant, nsrcs) (unless subarr_ant is not None)
+            vis[i] = vv[
+                i
+            ]  # v.shape: (ntimes, nant, nant, nsrcs) (unless subarr_ant is not None)
 
     return vis
 
 
-
-def simulate_vis(
-    *args, **kwargs
-):
+def simulate_vis(*args, **kwargs):
     """
     Run a basic simulation, based on ``matvis``.
 
@@ -530,8 +539,8 @@ def simulate_vis_per_alm(
     latitude=-30.7215 * np.pi / 180.0,
     use_feed="x",
     multiprocess=True,
-    amplitude=1.,
-    logfile=None
+    amplitude=1.0,
+    logfile=None,
 ):
     """
     Run a basic simulation, returning the visibility for each spherical harmonic mode
@@ -599,16 +608,16 @@ def simulate_vis_per_alm(
     # Make sure these are array_like
     freqs = np.atleast_1d(freqs)
     lsts = np.atleast_1d(lsts)
-    
+
     # Array of ell, m values in healpy ordering
     ell, m = hp.Alm().getlm(lmax=lmax)
 
     # Get Healpix pixel coords
     npix = hp.nside2npix(nside)
-    pix_area = 4.*np.pi / npix # steradians per pixel
+    pix_area = 4.0 * np.pi / npix  # steradians per pixel
     dec, ra = hp.pix2ang(nside=nside, ipix=np.arange(npix), lonlat=False)
     # RA must be in range [0, 2 pi] and Dec in range [-pi, +pi]
-    dec = dec - 0.5*np.pi # shift Dec coords
+    dec = dec - 0.5 * np.pi  # shift Dec coords
 
     # Dummy fluxes (one everywhere)
     fluxes = np.ones((npix, freqs.size))
@@ -617,29 +626,33 @@ def simulate_vis_per_alm(
     # visibility contrib. from each pixel
     if logfile is not None:
         t0 = time.time()
-        with open(logfile, 'a') as f:
+        with open(logfile, "a") as f:
             f.write("%s Starting simulate_vis_per_source\n" % (datetime.now()))
 
-    vis_pix = simulate_vis_per_source(ants=ants,
-                                      fluxes=fluxes,
-                                      ra=ra,
-                                      dec=dec,
-                                      freqs=freqs,
-                                      lsts=lsts,
-                                      beams=beams,
-                                      polarized=polarized,
-                                      precision=precision,
-                                      latitude=latitude,
-                                      use_feed=use_feed)
+    vis_pix = simulate_vis_per_source(
+        ants=ants,
+        fluxes=fluxes,
+        ra=ra,
+        dec=dec,
+        freqs=freqs,
+        lsts=lsts,
+        beams=beams,
+        polarized=polarized,
+        precision=precision,
+        latitude=latitude,
+        use_feed=use_feed,
+    )
 
     if logfile is not None:
-        with open(logfile, 'a') as f:
-            f.write("%s Finished simulate_vis_per_source in %5.2f sec\n" \
-                     % (datetime.now(), time.time() - t0))
+        with open(logfile, "a") as f:
+            f.write(
+                "%s Finished simulate_vis_per_source in %5.2f sec\n"
+                % (datetime.now(), time.time() - t0)
+            )
 
     # Empty array with the right shape (no. visibilities times no. l,m modes)
     shape = list(vis_pix.shape)
-    shape[-1] = 2*ell.size # replace last dim. with Nmodes (real + imag.)
+    shape[-1] = 2 * ell.size  # replace last dim. with Nmodes (real + imag.)
     vis = np.zeros(shape, dtype=np.complex128)
 
     # Loop over (ell, m) modes, weighting the precomputed visibility sim
@@ -648,14 +661,14 @@ def simulate_vis_per_alm(
     for n in range(ell.size):
 
         if logfile is not None:
-            with open(logfile, 'a') as f:
+            with open(logfile, "a") as f:
                 f.write("%s ell %d / %d\n" % (datetime.now(), n, ell.size))
 
         # Start with zero vector for all modes
         alm *= 0
 
         # Loop over real, imaginary values for this mode only
-        for j, val in enumerate([1., 1.j]):
+        for j, val in enumerate([1.0, 1.0j]):
 
             # Make healpix map for this mode only
             alm[n] = val
@@ -666,21 +679,22 @@ def simulate_vis_per_alm(
             # Multiply visibility for each pixel by the pixel value for this mode
             if polarized:
                 # vis_pix: (NAXES, NFEED, NFREQS, NTIMES, NANTS, NANTS, NSRCS)
-                vis[:,:,:,:,:,:,n + j*ell.size] = np.sum(vis_pix * skymap, axis=-1)
+                vis[:, :, :, :, :, :, n + j * ell.size] = np.sum(
+                    vis_pix * skymap, axis=-1
+                )
                 # Last dim. of vis is in blocks of real (first ell.size modes) and
                 # imaginary (last ell.size modes)
             else:
                 # vis_pix: (NFREQS, NTIMES, NANTS, NANTS, NSRCS)
-                vis[:,:,:,:,n + j*ell.size] = np.sum(vis_pix * skymap, axis=-1)
+                vis[:, :, :, :, n + j * ell.size] = np.sum(vis_pix * skymap, axis=-1)
 
     if logfile is not None:
-        with open(logfile, 'a') as f:
+        with open(logfile, "a") as f:
             f.write("%s Finished all.\n" % (datetime.now()))
 
     return ell, m, vis
 
 
-
 def simulate_vis_per_region(
     lmax,
     nside,
@@ -693,8 +707,8 @@ def simulate_vis_per_region(
     latitude=-30.7215 * np.pi / 180.0,
     use_feed="x",
     multiprocess=True,
-    amplitude=1.,
-    logfile=None
+    amplitude=1.0,
+    logfile=None,
 ):
     """
     Run a basic simulation, returning the visibility for each spherical harmonic mode
@@ -768,10 +782,10 @@ def simulate_vis_per_region(
 
     # Get Healpix pixel coords
     npix = hp.nside2npix(nside)
-    pix_area = 4.*np.pi / npix # steradians per pixel
+    pix_area = 4.0 * np.pi / npix  # steradians per pixel
     dec, ra = hp.pix2ang(nside=nside, ipix=np.arange(npix), lonlat=False)
     # RA must be in range [0, 2 pi] and Dec in range [-pi, +pi]
-    dec = dec - 0.5*np.pi # shift Dec coords
+    dec = dec - 0.5 * np.pi  # shift Dec coords
 
     # Dummy fluxes (one everywhere)
     fluxes = np.ones((npix, freqs.size))
@@ -780,30 +794,34 @@ def simulate_vis_per_region(
     # visibility contrib. from each pixel
     if logfile is not None:
         t0 = time.time()
-        with open(logfile, 'a') as f:
+        with open(logfile, "a") as f:
             f.write("%s Starting simulate_vis_per_source\n" % (datetime.now()))
 
-    vis_pix = simulate_vis_per_source(ants=ants,
-                                      fluxes=fluxes,
-                                      ra=ra,
-                                      dec=dec,
-                                      freqs=freqs,
-                                      lsts=lsts,
-                                      beams=beams,
-                                      polarized=polarized,
-                                      precision=precision,
-                                      latitude=latitude,
-                                      use_feed=use_feed,
-                                      multiprocess=multiprocess)
+    vis_pix = simulate_vis_per_source(
+        ants=ants,
+        fluxes=fluxes,
+        ra=ra,
+        dec=dec,
+        freqs=freqs,
+        lsts=lsts,
+        beams=beams,
+        polarized=polarized,
+        precision=precision,
+        latitude=latitude,
+        use_feed=use_feed,
+        multiprocess=multiprocess,
+    )
 
     if logfile is not None:
-        with open(logfile, 'a') as f:
-            f.write("%s Finished simulate_vis_per_source in %5.2f sec\n" \
-                     % (datetime.now(), time.time() - t0))
+        with open(logfile, "a") as f:
+            f.write(
+                "%s Finished simulate_vis_per_source in %5.2f sec\n"
+                % (datetime.now(), time.time() - t0)
+            )
 
     # Empty array with the right shape (no. visibilities times no. l,m modes)
     shape = list(vis_pix.shape)
-    shape[-1] = 2*ell.size # replace last dim. with Nmodes (real + imag.)
+    shape[-1] = 2 * ell.size  # replace last dim. with Nmodes (real + imag.)
     vis = np.zeros(shape, dtype=np.complex128)
 
     # Loop over (ell, m) modes, weighting the precomputed visibility sim
@@ -812,14 +830,14 @@ def simulate_vis_per_region(
     for n in range(ell.size):
 
         if logfile is not None:
-            with open(logfile, 'a') as f:
+            with open(logfile, "a") as f:
                 f.write("%s ell %d / %d\n" % (datetime.now(), n, ell.size))
 
         # Start with zero vector for all modes
         alm *= 0
 
         # Loop over real, imaginary values for this mode only
-        for j, val in enumerate([1., 1.j]):
+        for j, val in enumerate([1.0, 1.0j]):
 
             # Make healpix map for this mode only
             alm[n] = val
@@ -830,33 +848,33 @@ def simulate_vis_per_region(
             # Multiply visibility for each pixel by the pixel value for this mode
             if polarized:
                 # vis_pix: (NAXES, NFEED, NFREQS, NTIMES, NANTS, NANTS, NSRCS)
-                vis[:,:,:,:,:,:,n + j*ell.size] = np.sum(vis_pix * skymap, axis=-1)
+                vis[:, :, :, :, :, :, n + j * ell.size] = np.sum(
+                    vis_pix * skymap, axis=-1
+                )
                 # Last dim. of vis is in blocks of real (first ell.size modes) and
                 # imaginary (last ell.size modes)
             else:
                 # vis_pix: (NFREQS, NTIMES, NANTS, NANTS, NSRCS)
-                vis[:,:,:,:,n + j*ell.size] = np.sum(vis_pix * skymap, axis=-1)
+                vis[:, :, :, :, n + j * ell.size] = np.sum(vis_pix * skymap, axis=-1)
 
     if logfile is not None:
-        with open(logfile, 'a') as f:
+        with open(logfile, "a") as f:
             f.write("%s Finished all.\n" % (datetime.now()))
 
     return ell, m, vis
 
 
-
-
 def vis_sim_per_source_new(
     antpos: np.ndarray,
     freq: float,
     lsts: np.ndarray,
     alt: np.ndarray,
     az: np.ndarray,
-    I_sky: np.ndarray,  
+    I_sky: np.ndarray,
     beam_list: Sequence[UVBeam],
     precision: int = 2,
     polarized: bool = False,
-    subarr_ant=None
+    subarr_ant=None,
 ):
     """
     Calculate visibility from an input intensity map and beam model. This is
@@ -944,7 +962,7 @@ def vis_sim_per_source_new(
 
         # Simulate even if sources are below the horizon, since we need a
         # visibility per source regardless
-        above_horizon = alt[tidx] > 0.
+        above_horizon = alt[tidx] > 0.0
         nsrcs_up = nsrcs
 
         A_s = np.zeros((nax, nfeed, nbeam, nsrcs_up), dtype=complex_dtype)
@@ -952,7 +970,7 @@ def vis_sim_per_source_new(
         v = np.zeros((nant, nsrcs_up), dtype=complex_dtype)
 
         # Primary beam pattern using direct interpolation of UVBeam object
-        za = 0.5*np.pi - alt
+        za = 0.5 * np.pi - alt
         for i, bm in enumerate(beam_list):
             spw_axis_present = utils.get_beam_interp_shape(bm)
             kw = (
@@ -962,7 +980,10 @@ def vis_sim_per_source_new(
             )
 
             interp_beam = bm.interp(
-                az_array=az[tidx], za_array=za[tidx], freq_array=np.atleast_1d(freq), **kw
+                az_array=az[tidx],
+                za_array=za[tidx],
+                freq_array=np.atleast_1d(freq),
+                **kw
             )[0]
 
             if polarized:
@@ -985,11 +1006,15 @@ def vis_sim_per_source_new(
             raise ValueError("Beam interpolation resulted in an invalid value")
 
         # Calculate delays, where tau = (b * s) / c
-        #enu_e, enu_n, enu_u = crd_top
-        crd_top = np.array([np.sin(za[tidx])*np.cos(az[tidx]), 
-                            np.sin(za[tidx])*np.sin(az[tidx]), 
-                            np.cos(za[tidx])])
-        np.dot(antpos, crd_top[:,:], out=tau)
+        # enu_e, enu_n, enu_u = crd_top
+        crd_top = np.array(
+            [
+                np.sin(za[tidx]) * np.cos(az[tidx]),
+                np.sin(za[tidx]) * np.sin(az[tidx]),
+                np.cos(za[tidx]),
+            ]
+        )
+        np.dot(antpos, crd_top[:, :], out=tau)
         tau /= c.value
 
         # Component of complex phase factor for one antenna
@@ -998,17 +1023,17 @@ def vis_sim_per_source_new(
         np.exp(1.0j * (ang_freq * tau), out=v)
 
         # Complex voltages.
-        #v *= Isqrt[above_horizon]
+        # v *= Isqrt[above_horizon]
         v *= Isqrt[:]
-        v[:,~above_horizon] *= 0. # zero-out sources below the horizon
+        v[:, ~above_horizon] *= 0.0  # zero-out sources below the horizon
 
         # Compute visibilities using product of complex voltages (upper triangle).
         # Input arrays have shape (Nax, Nfeed, [Nants], Nsrcs
         v = A_s[:, :, :] * v[np.newaxis, np.newaxis, :]
 
-        #print(">>>", A_s)
-        #print("\n\n\n")
-        #print("***", v)
+        # print(">>>", A_s)
+        # print("\n\n\n")
+        # print("***", v)
 
         if subarr_ant is None:
             for i in range(len(antpos)):
@@ -1023,12 +1048,11 @@ def vis_sim_per_source_new(
         else:
             # Get the ones where the antenna in question is not conjugated
             vis[:, :, tidx] = np.einsum(
-                "jiln,jkmn->ikln", # summing over m just sums over one antenna (squeezes an axis)
+                "jiln,jkmn->ikln",  # summing over m just sums over one antenna (squeezes an axis)
                 v[:, :, :, :].conj(),
-                v[:, :, subarr_ant: subarr_ant + 1, :],
+                v[:, :, subarr_ant : subarr_ant + 1, :],
                 optimize=True,
             )
 
     # Return visibilities with or without multiple polarization channels
     return vis if polarized else vis[0, 0]
-