def fit_gains_and_foregrounds(g_r, g_i, fg_r, fg_i, data_r, data_i, wgts, fg_comps, corr_inds, use_min=False, tol=1e-14, maxsteps=10000, optimizer='Adamax', freeze_model=False, verbose=False, notebook_progressbar=False, dtype=np.float32, graph_mode=False, n_profile_steps=0, profile_log_dir='./logdir', sky_model_r=None, sky_model_i=None, model_regularization=None, graph_args_dict=None, **opt_kwargs):
    'Run optimization loop to fit gains and foreground components.\n\n    Parameters\n    ----------\n    g_r: tf.Tensor object.\n        tf.Tensor object holding real parts of gains.\n    g_i: tf.Tensor object.\n        tf.Tensor object holding imag parts of gains.\n    fg_r: list\n        list of tf.Tensor objects. Each has shape (nvecs, ngrps, 1, 1)\n        tf.Tensor object holding foreground coeffs.\n    fg_i: list\n        list of tf.Tensor objects. Each has shape (nvecs, ngrps, 1, 1)\n        tf.Tensor object holding imag coeffs.\n    data_r: list\n        list of tf.Tensor objects. Each has shape (ngrps, nbls, nfreqs)\n        real part of data to fit.\n    data_i: list\n        list of tf.Tensor objects. Each has shape (ngrps, nbls, nfreqs)\n        imag part of data to fit.\n    wgts: list\n        list of tf.Tensor objects. Each has shape (ngrps, nbls, nfreqs)\n    fg_comps: list:\n        list of tf.Tensor objects. Each has shape (nvecs, ngrps, nbls, nfreqs)\n        represents vectors to be used in modeling visibilities.\n    corr_inds: list\n        list of list of lists of 2-tuples. Hierarchy of lists is\n        chunk\n            group\n                baseline - (int 2-tuple)\n    use_min: bool, optional\n        if True, use the value that minimizes the loss function\n        regardless of where optimization loop ended up\n        (prevents overshooting due to excess momentum)\n    tol: float, optional\n        halt optimization loop once the loss changes by less then this value.\n        default is 1e-14\n    maxsteps: int, optional\n        maximum number of opt.minimize calls before halting.\n        default is 10000\n    optimizer: string\n        Name of optimizer. See OPTIMIZERS dictionary which contains optimizers described in\n        https://www.tensorflow.org/api_docs/python/tf/keras/optimizers\n        default is \'Adamax\'\n    freeze_model: bool, optional\n        Only optimize loss function wrt gain variables. This is effectively traditional model-based calibration\n        with sky_model as the model (but projected onto the foreground basis vectors).\n        default is False.\n    verbose: bool, optional\n        lots of text output\n        default is False.\n    notebook_progressbar: bool, optional\n        use progress bar optimized for notebook output.\n        default is False.\n    graph_mode: bool, optional\n        if True, compile gradient update step in graph mode to speed up\n        runtime by ~2-3x. I\'ve found that this helps on CPUs but on GPUs\n        it actually increases runtime by a similar factor.\n    n_profile_steps: bool, optional\n        number of steps to run profiling on\n        default is 0.\n    profile_log_dir: str, optional\n        directory to save profile logs to\n        default is \'./logdir\'\n    sky_model_r: list of tf.Tensor objects, optional\n        chunked tensors containing model in same format as data_r\n    sky_model_i: list of tf.Tensor objects, optional\n        chunked tensors containing model in the same format as data_i\n    model_regularization: str, optional\n        type of model regularization to perform. Currently support "sum"\n        where the sums of real and imaginary parts (across all bls and freqs)\n        are constrained to be the same as the sum of real and imag parts\n        of data.\n    opt_kwargs: kwarg dict\n        additional kwargs for tf.opt.Optimizer(). See tensorflow docs.\n\n    Returns\n    -------\n    g_r_opt: tf.Tensor object\n        real part of optimized gains.\n    g_i_opt: tf.Tensor object\n        imag part of optimized gains.\n    fg_r_opt: tf.Tensor object\n        real part of foreground coeffs.\n    fg_i_opt: tf.Tensor object.\n        imag part of optimized foreground coeffs.\n    fit_history: dict\n        dictionary containing fit history for each time-step and polarization in the data with fields:\n        \'loss_history\': list of values of the loss function in each minimization iteration.\n    '
    if (graph_args_dict is None):
        graph_args_dict = {}
    echo(f'Using {str(dtype)} precision.')
    echo(f'{datetime.datetime.now()} Provided the following opt_kwargs')
    for k in opt_kwargs:
        echo(f'{k}: {opt_kwargs[k]}')
    opt = OPTIMIZERS[optimizer](**opt_kwargs)
    fit_history = {'loss': []}
    min_loss = 9e+99
    nants = g_r.shape[0]
    nfreqs = g_r.shape[1]
    ant0_inds = []
    ant1_inds = []
    nchunks = len(fg_comps)
    for cnum in range(nchunks):
        ant0_chunk = []
        ant1_chunk = []
        ngrps = len(corr_inds[cnum])
        for gnum in range(ngrps):
            ant0_grp = []
            ant1_grp = []
            for cpair in corr_inds[cnum][gnum]:
                ant0_grp.append(cpair[0])
                ant1_grp.append(cpair[1])
            ant0_chunk.append(ant0_grp)
            ant1_chunk.append(ant1_grp)
        ant0_inds.append(ant0_chunk)
        ant1_inds.append(ant1_chunk)
    g_r = tf.Variable(g_r)
    g_i = tf.Variable(g_i)
    if (not freeze_model):
        fg_r = [tf.Variable(fgr) for fgr in fg_r]
        fg_i = [tf.Variable(fgi) for fgi in fg_i]
        vars = (([g_r, g_i] + fg_r) + fg_i)
    else:
        vars = [g_r, g_i]
    echo(f'{datetime.datetime.now()} Performing gradient descent on {np.prod(g_r.shape)} complex gain parameters...', verbose=verbose)
    if (not freeze_model):
        echo(f'Performing gradient descent on total of {int(np.sum([(fgr.shape[0] * fgr.shape[1]) for fgr in fg_r]))} complex foreground parameters', verbose=verbose)
        echo(f"Foreground Parameters grouped into chunks of shape ((nvecs, ngrps): nbls) {[((str(fgr.shape[:2]) + ':') + str(dc.shape[1])) for (fgr, dc) in zip(fg_r, data_r)]}", verbose=verbose)
    if (model_regularization == 'sum'):
        prior_r_sum = tf.reduce_sum(tf.stack([tf.reduce_sum((sky_model_r[cnum] * wgts[cnum])) for cnum in range(nchunks)]))
        prior_i_sum = tf.reduce_sum(tf.stack([tf.reduce_sum((sky_model_i[cnum] * wgts[cnum])) for cnum in range(nchunks)]))

        def loss_function():
            return mse_chunked_sum_regularized(g_r=g_r, g_i=g_i, fg_r=fg_r, fg_i=fg_i, fg_comps=fg_comps, nchunks=nchunks, data_r=data_r, data_i=data_i, wgts=wgts, ant0_inds=ant0_inds, ant1_inds=ant1_inds, dtype=dtype, prior_r_sum=prior_r_sum, prior_i_sum=prior_i_sum)
    else:

        def loss_function():
            return mse_chunked(g_r=g_r, g_i=g_i, fg_r=fg_r, fg_i=fg_i, fg_comps=fg_comps, nchunks=nchunks, data_r=data_r, data_i=data_i, wgts=wgts, ant0_inds=ant0_inds, ant1_inds=ant1_inds, dtype=dtype)

    def train_step_code():
        with tf.GradientTape() as tape:
            loss = loss_function()
        grads = tape.gradient(loss, vars)
        opt.apply_gradients(zip(grads, vars))
        return loss
    if graph_mode:

        @tf.function(**graph_args_dict)
        def train_step():
            return train_step_code()
    else:

        def train_step():
            return train_step_code()
    if (n_profile_steps > 0):
        echo(f'{datetime.datetime.now()} Profiling with {n_profile_steps}. And writing output to {profile_log_dir}...')
        tf.profiler.experimental.start(profile_log_dir)
        for step in PBARS[notebook_progressbar](range(n_profile_steps)):
            with tf.profiler.experimental.Trace('train', step_num=step):
                train_step()
        tf.profiler.experimental.stop()
    echo(f'''{datetime.datetime.now()} Building Computational Graph...
''', verbose=verbose)
    loss = train_step()
    echo(f'''{datetime.datetime.now()} Performing Gradient Descent. Initial MSE of {loss:.2e}...
''', verbose=verbose)
    for step in PBARS[notebook_progressbar](range(maxsteps)):
        loss = train_step()
        fit_history['loss'].append(loss.numpy())
        if (use_min and (fit_history['loss'][(- 1)] < min_loss)):
            min_loss = fit_history['loss'][(- 1)]
            g_r_opt = g_r.value()
            g_i_opt = g_i.value()
            if (not freeze_model):
                fg_r_opt = [fgr.value() for fgr in fg_r]
                fg_i_opt = [fgi.value() for fgi in fg_i]
        if ((step >= 1) and (np.abs((fit_history['loss'][(- 1)] - fit_history['loss'][(- 2)])) < tol)):
            echo(f'''Tolerance thresshold met with delta of {np.abs((fit_history['loss'][(- 1)] - fit_history['loss'][(- 2)])):.2e}. Terminating...
 ''', verbose=verbose)
            break
    if (not use_min):
        min_loss = fit_history['loss'][(- 1)]
        g_r_opt = g_r.value()
        g_i_opt = g_i.value()
        if (not freeze_model):
            fg_r_opt = [fgr.value() for fgr in fg_r]
            fg_i_opt = [fgi.value() for fgi in fg_i]
        else:
            fg_r_opt = fg_r
            fg_i_opt = fg_i
    echo(f'''{datetime.datetime.now()} Finished Gradient Descent. MSE of {min_loss:.2e}...
''', verbose=verbose)
    return (g_r_opt, g_i_opt, fg_r_opt, fg_i_opt, fit_history)
Run optimization loop to fit gains and foreground components.

Parameters
----------
g_r: tf.Tensor object.
    tf.Tensor object holding real parts of gains.
g_i: tf.Tensor object.
    tf.Tensor object holding imag parts of gains.
fg_r: list
    list of tf.Tensor objects. Each has shape (nvecs, ngrps, 1, 1)
    tf.Tensor object holding foreground coeffs.
fg_i: list
    list of tf.Tensor objects. Each has shape (nvecs, ngrps, 1, 1)
    tf.Tensor object holding imag coeffs.
data_r: list
    list of tf.Tensor objects. Each has shape (ngrps, nbls, nfreqs)
    real part of data to fit.
data_i: list
    list of tf.Tensor objects. Each has shape (ngrps, nbls, nfreqs)
    imag part of data to fit.
wgts: list
    list of tf.Tensor objects. Each has shape (ngrps, nbls, nfreqs)
fg_comps: list:
    list of tf.Tensor objects. Each has shape (nvecs, ngrps, nbls, nfreqs)
    represents vectors to be used in modeling visibilities.
corr_inds: list
    list of list of lists of 2-tuples. Hierarchy of lists is
    chunk
        group
            baseline - (int 2-tuple)
use_min: bool, optional
    if True, use the value that minimizes the loss function
    regardless of where optimization loop ended up
    (prevents overshooting due to excess momentum)
tol: float, optional
    halt optimization loop once the loss changes by less then this value.
    default is 1e-14
maxsteps: int, optional
    maximum number of opt.minimize calls before halting.
    default is 10000
optimizer: string
    Name of optimizer. See OPTIMIZERS dictionary which contains optimizers described in
    https://www.tensorflow.org/api_docs/python/tf/keras/optimizers
    default is 'Adamax'
freeze_model: bool, optional
    Only optimize loss function wrt gain variables. This is effectively traditional model-based calibration
    with sky_model as the model (but projected onto the foreground basis vectors).
    default is False.
verbose: bool, optional
    lots of text output
    default is False.
notebook_progressbar: bool, optional
    use progress bar optimized for notebook output.
    default is False.
graph_mode: bool, optional
    if True, compile gradient update step in graph mode to speed up
    runtime by ~2-3x. I've found that this helps on CPUs but on GPUs
    it actually increases runtime by a similar factor.
n_profile_steps: bool, optional
    number of steps to run profiling on
    default is 0.
profile_log_dir: str, optional
    directory to save profile logs to
    default is './logdir'
sky_model_r: list of tf.Tensor objects, optional
    chunked tensors containing model in same format as data_r
sky_model_i: list of tf.Tensor objects, optional
    chunked tensors containing model in the same format as data_i
model_regularization: str, optional
    type of model regularization to perform. Currently support "sum"
    where the sums of real and imaginary parts (across all bls and freqs)
    are constrained to be the same as the sum of real and imag parts
    of data.
opt_kwargs: kwarg dict
    additional kwargs for tf.opt.Optimizer(). See tensorflow docs.

Returns
-------
g_r_opt: tf.Tensor object
    real part of optimized gains.
g_i_opt: tf.Tensor object
    imag part of optimized gains.
fg_r_opt: tf.Tensor object
    real part of foreground coeffs.
fg_i_opt: tf.Tensor object.
    imag part of optimized foreground coeffs.
fit_history: dict
    dictionary containing fit history for each time-step and polarization in the data with fields:
    'loss_history': list of values of the loss function in each minimization iteration.
def fit_gains_and_foregrounds(g_r, g_i, fg_r, fg_i, data_r, data_i, wgts, fg_comps, corr_inds, use_min=False, tol=1e-14, maxsteps=10000, optimizer='Adamax', freeze_model=False, verbose=False, notebook_progressbar=False, dtype=np.float32, graph_mode=False, n_profile_steps=0, profile_log_dir='./logdir', sky_model_r=None, sky_model_i=None, model_regularization=None, graph_args_dict=None, **opt_kwargs):
    'Run optimization loop to fit gains and foreground components.\n\n    Parameters\n    ----------\n    g_r: tf.Tensor object.\n        tf.Tensor object holding real parts of gains.\n    g_i: tf.Tensor object.\n        tf.Tensor object holding imag parts of gains.\n    fg_r: list\n        list of tf.Tensor objects. Each has shape (nvecs, ngrps, 1, 1)\n        tf.Tensor object holding foreground coeffs.\n    fg_i: list\n        list of tf.Tensor objects. Each has shape (nvecs, ngrps, 1, 1)\n        tf.Tensor object holding imag coeffs.\n    data_r: list\n        list of tf.Tensor objects. Each has shape (ngrps, nbls, nfreqs)\n        real part of data to fit.\n    data_i: list\n        list of tf.Tensor objects. Each has shape (ngrps, nbls, nfreqs)\n        imag part of data to fit.\n    wgts: list\n        list of tf.Tensor objects. Each has shape (ngrps, nbls, nfreqs)\n    fg_comps: list:\n        list of tf.Tensor objects. Each has shape (nvecs, ngrps, nbls, nfreqs)\n        represents vectors to be used in modeling visibilities.\n    corr_inds: list\n        list of list of lists of 2-tuples. Hierarchy of lists is\n        chunk\n            group\n                baseline - (int 2-tuple)\n    use_min: bool, optional\n        if True, use the value that minimizes the loss function\n        regardless of where optimization loop ended up\n        (prevents overshooting due to excess momentum)\n    tol: float, optional\n        halt optimization loop once the loss changes by less then this value.\n        default is 1e-14\n    maxsteps: int, optional\n        maximum number of opt.minimize calls before halting.\n        default is 10000\n    optimizer: string\n        Name of optimizer. See OPTIMIZERS dictionary which contains optimizers described in\n        https://www.tensorflow.org/api_docs/python/tf/keras/optimizers\n        default is \'Adamax\'\n    freeze_model: bool, optional\n        Only optimize loss function wrt gain variables. This is effectively traditional model-based calibration\n        with sky_model as the model (but projected onto the foreground basis vectors).\n        default is False.\n    verbose: bool, optional\n        lots of text output\n        default is False.\n    notebook_progressbar: bool, optional\n        use progress bar optimized for notebook output.\n        default is False.\n    graph_mode: bool, optional\n        if True, compile gradient update step in graph mode to speed up\n        runtime by ~2-3x. I\'ve found that this helps on CPUs but on GPUs\n        it actually increases runtime by a similar factor.\n    n_profile_steps: bool, optional\n        number of steps to run profiling on\n        default is 0.\n    profile_log_dir: str, optional\n        directory to save profile logs to\n        default is \'./logdir\'\n    sky_model_r: list of tf.Tensor objects, optional\n        chunked tensors containing model in same format as data_r\n    sky_model_i: list of tf.Tensor objects, optional\n        chunked tensors containing model in the same format as data_i\n    model_regularization: str, optional\n        type of model regularization to perform. Currently support "sum"\n        where the sums of real and imaginary parts (across all bls and freqs)\n        are constrained to be the same as the sum of real and imag parts\n        of data.\n    opt_kwargs: kwarg dict\n        additional kwargs for tf.opt.Optimizer(). See tensorflow docs.\n\n    Returns\n    -------\n    g_r_opt: tf.Tensor object\n        real part of optimized gains.\n    g_i_opt: tf.Tensor object\n        imag part of optimized gains.\n    fg_r_opt: tf.Tensor object\n        real part of foreground coeffs.\n    fg_i_opt: tf.Tensor object.\n        imag part of optimized foreground coeffs.\n    fit_history: dict\n        dictionary containing fit history for each time-step and polarization in the data with fields:\n        \'loss_history\': list of values of the loss function in each minimization iteration.\n    '
    if (graph_args_dict is None):
        graph_args_dict = {}
    echo(f'Using {str(dtype)} precision.')
    echo(f'{datetime.datetime.now()} Provided the following opt_kwargs')
    for k in opt_kwargs:
        echo(f'{k}: {opt_kwargs[k]}')
    opt = OPTIMIZERS[optimizer](**opt_kwargs)
    fit_history = {'loss': []}
    min_loss = 9e+99
    nants = g_r.shape[0]
    nfreqs = g_r.shape[1]
    ant0_inds = []
    ant1_inds = []
    nchunks = len(fg_comps)
    for cnum in range(nchunks):
        ant0_chunk = []
        ant1_chunk = []
        ngrps = len(corr_inds[cnum])
        for gnum in range(ngrps):
            ant0_grp = []
            ant1_grp = []
            for cpair in corr_inds[cnum][gnum]:
                ant0_grp.append(cpair[0])
                ant1_grp.append(cpair[1])
            ant0_chunk.append(ant0_grp)
            ant1_chunk.append(ant1_grp)
        ant0_inds.append(ant0_chunk)
        ant1_inds.append(ant1_chunk)
    g_r = tf.Variable(g_r)
    g_i = tf.Variable(g_i)
    if (not freeze_model):
        fg_r = [tf.Variable(fgr) for fgr in fg_r]
        fg_i = [tf.Variable(fgi) for fgi in fg_i]
        vars = (([g_r, g_i] + fg_r) + fg_i)
    else:
        vars = [g_r, g_i]
    echo(f'{datetime.datetime.now()} Performing gradient descent on {np.prod(g_r.shape)} complex gain parameters...', verbose=verbose)
    if (not freeze_model):
        echo(f'Performing gradient descent on total of {int(np.sum([(fgr.shape[0] * fgr.shape[1]) for fgr in fg_r]))} complex foreground parameters', verbose=verbose)
        echo(f"Foreground Parameters grouped into chunks of shape ((nvecs, ngrps): nbls) {[((str(fgr.shape[:2]) + ':') + str(dc.shape[1])) for (fgr, dc) in zip(fg_r, data_r)]}", verbose=verbose)
    if (model_regularization == 'sum'):
        prior_r_sum = tf.reduce_sum(tf.stack([tf.reduce_sum((sky_model_r[cnum] * wgts[cnum])) for cnum in range(nchunks)]))
        prior_i_sum = tf.reduce_sum(tf.stack([tf.reduce_sum((sky_model_i[cnum] * wgts[cnum])) for cnum in range(nchunks)]))

        def loss_function():
            return mse_chunked_sum_regularized(g_r=g_r, g_i=g_i, fg_r=fg_r, fg_i=fg_i, fg_comps=fg_comps, nchunks=nchunks, data_r=data_r, data_i=data_i, wgts=wgts, ant0_inds=ant0_inds, ant1_inds=ant1_inds, dtype=dtype, prior_r_sum=prior_r_sum, prior_i_sum=prior_i_sum)
    else:

        def loss_function():
            return mse_chunked(g_r=g_r, g_i=g_i, fg_r=fg_r, fg_i=fg_i, fg_comps=fg_comps, nchunks=nchunks, data_r=data_r, data_i=data_i, wgts=wgts, ant0_inds=ant0_inds, ant1_inds=ant1_inds, dtype=dtype)

    def train_step_code():
        with tf.GradientTape() as tape:
            loss = loss_function()
        grads = tape.gradient(loss, vars)
        opt.apply_gradients(zip(grads, vars))
        return loss
    if graph_mode:

        @tf.function(**graph_args_dict)
        def train_step():
            return train_step_code()
    else:

        def train_step():
            return train_step_code()
    if (n_profile_steps > 0):
        echo(f'{datetime.datetime.now()} Profiling with {n_profile_steps}. And writing output to {profile_log_dir}...')
        tf.profiler.experimental.start(profile_log_dir)
        for step in PBARS[notebook_progressbar](range(n_profile_steps)):
            with tf.profiler.experimental.Trace('train', step_num=step):
                train_step()
        tf.profiler.experimental.stop()
    echo(f'{datetime.datetime.now()} Building Computational Graph...
', verbose=verbose)
    loss = train_step()
    echo(f'{datetime.datetime.now()} Performing Gradient Descent. Initial MSE of {loss:.2e}...
', verbose=verbose)
    for step in PBARS[notebook_progressbar](range(maxsteps)):
        loss = train_step()
        fit_history['loss'].append(loss.numpy())
        if (use_min and (fit_history['loss'][(- 1)] < min_loss)):
            min_loss = fit_history['loss'][(- 1)]
            g_r_opt = g_r.value()
            g_i_opt = g_i.value()
            if (not freeze_model):
                fg_r_opt = [fgr.value() for fgr in fg_r]
                fg_i_opt = [fgi.value() for fgi in fg_i]
        if ((step >= 1) and (np.abs((fit_history['loss'][(- 1)] - fit_history['loss'][(- 2)])) < tol)):
            echo(f'Tolerance thresshold met with delta of {np.abs((fit_history['loss'][(- 1)] - fit_history['loss'][(- 2)])):.2e}. Terminating...
 ', verbose=verbose)
            break
    if (not use_min):
        min_loss = fit_history['loss'][(- 1)]
        g_r_opt = g_r.value()
        g_i_opt = g_i.value()
        if (not freeze_model):
            fg_r_opt = [fgr.value() for fgr in fg_r]
            fg_i_opt = [fgi.value() for fgi in fg_i]
        else:
            fg_r_opt = fg_r
            fg_i_opt = fg_i
    echo(f'{datetime.datetime.now()} Finished Gradient Descent. MSE of {min_loss:.2e}...
', verbose=verbose)
    return (g_r_opt, g_i_opt, fg_r_opt, fg_i_opt, fit_history)