# Path: cubical/kernels/generics.py
# def compute_2x2_inverse(x, xinv, flags, eps, flagbit):
# def compute_diag_inverse(x, xinv, flags, eps, flagbit):
# def compute_chisq(r, chisq):
# def compute_chisq_diag(r, chisq):
# def compute_chisq_offdiag(r, chisq):
# 
# Path: cubical/kernels/full_complex.py
# def allocate_vis_array(shape, dtype, zeros=False):
# def allocate_gain_array(shape, dtype, zeros=False):
# def allocate_flag_array(shape, dtype, zeros=False):
# def compute_residual(m, g, gh, r, t_int, f_int):
# def compute_jh(m, g, jh, t_int, f_int):
# def compute_jhr(jh, r, jhr, t_int, f_int):
# def compute_jhj(jh, jhj, t_int, f_int):
# def compute_update(jhr, jhjinv, upd):
# def compute_corrected(o, g, gh, corr, t_int, f_int):
# def apply_gains(m, g, gh, t_int, f_int):
# def right_multiply_gains(g, g_next, t_int, f_int):
# 
# Path: cubical/kernels/diag_complex.py
# def compute_residual(m, g, gh, r, t_int, f_int):
# def compute_jh(m, g, jh, t_int, f_int):
# def compute_update(jhr, jhjinv, upd):
# def compute_corrected(o, g, gh, corr, t_int, f_int):
# def apply_gains(m, g, gh, t_int, f_int):
# def right_multiply_gains(g, g_next, t_int, f_int):
# 
# Path: cubical/kernels/diagdiag_complex.py
# def compute_residual(m, g, gh, r, t_int, f_int):
# def compute_jh(m, g, jh, t_int, f_int):
# def compute_jhr(jh, r, jhr, t_int, f_int):
# def compute_jhj(jh, jhj, t_int, f_int):
# def compute_update(jhr, jhjinv, upd):
# def compute_corrected(o, g, gh, corr, t_int, f_int):
# def apply_gains(m, g, gh, t_int, f_int):
# def right_multiply_gains(g, g_next, t_int, f_int):
# 
# Path: cubical/kernels/phase_only.py
# def compute_jhj(m, jhj, t_int=1, f_int=1):
# def compute_jhr(gh, jh, r, jhr, t_int=1, f_int=1):
# def compute_update(jhr, jhjinv, upd):


from builtins import range
from numba import jit, prange
from cubical.kernels import generics
from cubical.kernels import full_complex
from cubical.kernels import diag_complex
from cubical.kernels import diagdiag_complex
from cubical.kernels import phase_only
import numpy as np
import cubical.kernels
    broadcast_dirs = np.array([d%g_dir for d in range(n_dir)])

    for ibl in prange(n_bl):
        aa, ab = all_bls[ibl,0], all_bls[ibl,1]
        for i in range(n_mod):
            for t in range(n_tim):
                bt = broadcast_times[t]
                for f in range(n_fre):
                    bf = broadcast_freqs[f]
                    for d in range(n_dir):
                        bd = broadcast_dirs[d]

                        r00 = r[i,t,f,aa,ab,0,0]
                        r11 = r[i,t,f,aa,ab,1,1]

                        jhh00 = jh[d,i,t,f,ab,aa,0,0]
                        jhh11 = jh[d,i,t,f,ab,aa,1,1]

                        jhr[d,bt,bf,aa,0,0] += gh[bd,bt,bf,aa,0,0]*r00*jhh00
                        jhr[d,bt,bf,aa,1,1] += gh[bd,bt,bf,aa,1,1]*r11*jhh11

# Remaining kernel functions are reused.

# J^H is computed assuming diagonal gains.
compute_jh = diagdiag_complex.compute_jh

# J^H.J inverse is computed assuming diagonal blocks.
compute_jhjinv = generics.compute_diag_inverse

# Residuals computed assuming diagonal gains.

# Path: cubical/kernels/generics.py
# def compute_2x2_inverse(x, xinv, flags, eps, flagbit):
# def compute_diag_inverse(x, xinv, flags, eps, flagbit):
# def compute_chisq(r, chisq):
# def compute_chisq_diag(r, chisq):
# def compute_chisq_offdiag(r, chisq):
# 
# Path: cubical/kernels/full_complex.py
# def allocate_vis_array(shape, dtype, zeros=False):
# def allocate_gain_array(shape, dtype, zeros=False):
# def allocate_flag_array(shape, dtype, zeros=False):
# def compute_residual(m, g, gh, r, t_int, f_int):
# def compute_jh(m, g, jh, t_int, f_int):
# def compute_jhr(jh, r, jhr, t_int, f_int):
# def compute_jhj(jh, jhj, t_int, f_int):
# def compute_update(jhr, jhjinv, upd):
# def compute_corrected(o, g, gh, corr, t_int, f_int):
# def apply_gains(m, g, gh, t_int, f_int):
# def right_multiply_gains(g, g_next, t_int, f_int):
# 
# Path: cubical/kernels/diag_complex.py
# def compute_residual(m, g, gh, r, t_int, f_int):
# def compute_jh(m, g, jh, t_int, f_int):
# def compute_update(jhr, jhjinv, upd):
# def compute_corrected(o, g, gh, corr, t_int, f_int):
# def apply_gains(m, g, gh, t_int, f_int):
# def right_multiply_gains(g, g_next, t_int, f_int):
# 
# Path: cubical/kernels/diagdiag_complex.py
# def compute_residual(m, g, gh, r, t_int, f_int):
# def compute_jh(m, g, jh, t_int, f_int):
# def compute_jhr(jh, r, jhr, t_int, f_int):
# def compute_jhj(jh, jhj, t_int, f_int):
# def compute_update(jhr, jhjinv, upd):
# def compute_corrected(o, g, gh, corr, t_int, f_int):
# def apply_gains(m, g, gh, t_int, f_int):
# def right_multiply_gains(g, g_next, t_int, f_int):
# 
# Path: cubical/kernels/phase_only.py
# def compute_jhj(m, jhj, t_int=1, f_int=1):
# def compute_jhr(gh, jh, r, jhr, t_int=1, f_int=1):
# def compute_update(jhr, jhjinv, upd):


# Path: cubical/kernels/diag_phase_only.py
from builtins import range
from numba import jit, prange
from cubical.kernels import generics
from cubical.kernels import full_complex
from cubical.kernels import diag_complex
from cubical.kernels import diagdiag_complex
from cubical.kernels import phase_only
import numpy as np
import cubical.kernels

    broadcast_dirs = np.array([d%g_dir for d in range(n_dir)])

    for ibl in prange(n_bl):
        aa, ab = all_bls[ibl,0], all_bls[ibl,1]
        for i in range(n_mod):
            for t in range(n_tim):
                bt = broadcast_times[t]
                for f in range(n_fre):
                    bf = broadcast_freqs[f]
                    for d in range(n_dir):
                        bd = broadcast_dirs[d]

                        r00 = r[i,t,f,aa,ab,0,0]
                        r11 = r[i,t,f,aa,ab,1,1]

                        jhh00 = jh[d,i,t,f,ab,aa,0,0]
                        jhh11 = jh[d,i,t,f,ab,aa,1,1]

                        jhr[d,bt,bf,aa,0,0] += gh[bd,bt,bf,aa,0,0]*r00*jhh00
                        jhr[d,bt,bf,aa,1,1] += gh[bd,bt,bf,aa,1,1]*r11*jhh11

# Remaining kernel functions are reused.

# J^H is computed assuming diagonal gains.
compute_jh = diagdiag_complex.compute_jh

# J^H.J inverse is computed assuming diagonal blocks.
compute_jhjinv = generics.compute_diag_inverse

# Residuals computed assuming diagonal gains.

compute_residual = diagdiag_complex.compute_residual