def tensorize_fg_model_comps_dict(fg_model_comps_dict, ants_map, nfreqs, use_redundancy=False, dtype=np.float32, notebook_progressbar=False, verbose=False, grp_size_threshold=5):
    'Convert per-baseline model components into a Ndata x Ncomponent tensor\n\n    Parameters\n    ----------\n    fg_model_comps_dict: dict\n        dictionary where each key is a 2-tuple (nbl, nvecs) referring to the number\n        of baselines in each vector and the number of vectors. Each 2-tuple points to\n        a dictionary where each key is the fitting group in fg_comps_dict that includes\n        nbl baselines. Each key in the referenced dict points to an (nred_grps * nfreqs x nvecs)\n        numpy.ndarray describing the modeling components for each fitting group in the chunk.\n    ants_map: dict mapping integers to integers\n        map between each antenna number to a unique index between 0 and Nants_data\n        (typically the index of each antenna in ants_map)\n    nfreqs: int, optional\n        number of frequency channels\n    dtype: numpy.dtype\n        tensor data types\n        default is np.float32\n\n    Returns\n    -------\n    fg_model_comps: list\n        list of tf.Tensor objects where each tensor has shape (nvecs, ngrps, nbls, nfreqs)\n        where nbls varies from tensor to tensor. Fitting groups with vectors that span nbls are lumped into the same\n        modeling tensor along the ngrps axis. nvecs is chosen in chunk_fg_comp_dict_by_nbls\n        to be the maximum number of vectors representing any of the ngrps baseline grps\n        which means that many rows in nvecs will be zero. For example, if we are modeling with\n        vectors that all span nbls=1 baseline and using delay-modes to model our data\n        then nvecs will equal the largest number of delay modes necessary to model the wedge\n        on all baselines even though the short baselines are described by far fewer modes\n        on short baselines, most of the rows along the vector dimension will therefor be zero.\n        This is wasteful of memory but it allows us to take advantage of the fast\n        dense matrix operations on a GPU.\n\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\n    '
    echo(f'''{datetime.datetime.now()} Computing foreground components matrices...
''', verbose=verbose)
    fg_model_comps_dict = chunk_fg_comp_dict_by_nbls(fg_model_comps_dict, use_redundancy=use_redundancy, grp_size_threshold=grp_size_threshold)
    fg_model_comps = []
    corr_inds = []
    for (nbls, nvecs) in fg_model_comps_dict:
        ngrps = len(fg_model_comps_dict[(nbls, nvecs)])
        modeling_matrix = np.zeros((nvecs, ngrps, nbls, nfreqs))
        corr_inds_chunk = []
        for (grpnum, modeling_grp) in enumerate(fg_model_comps_dict[(nbls, nvecs)]):
            corr_inds_grp = []
            nbl = 0
            for (rgrpnum, red_grp) in enumerate(modeling_grp):
                nred = len(red_grp)
                for ap in red_grp:
                    (i, j) = (ants_map[ap[0]], ants_map[ap[1]])
                    corr_inds_grp.append((i, j))
                    vecslice = slice(0, fg_model_comps_dict[(nbls, nvecs)][modeling_grp].shape[1])
                    compslice = slice((rgrpnum * nfreqs), ((rgrpnum + 1) * nfreqs))
                    dslice = slice((nbl * nfreqs), ((nbl + 1) * nfreqs))
                    modeling_matrix[(vecslice, grpnum, nbl)] = fg_model_comps_dict[(nbls, nvecs)][modeling_grp][compslice].T
                    nbl += 1
            corr_inds_chunk.append(corr_inds_grp)
        fg_model_comps.append(tf.convert_to_tensor(modeling_matrix, dtype=dtype))
        corr_inds.append(corr_inds_chunk)
    return (fg_model_comps, corr_inds)
Convert per-baseline model components into a Ndata x Ncomponent tensor

Parameters
----------
fg_model_comps_dict: dict
    dictionary where each key is a 2-tuple (nbl, nvecs) referring to the number
    of baselines in each vector and the number of vectors. Each 2-tuple points to
    a dictionary where each key is the fitting group in fg_comps_dict that includes
    nbl baselines. Each key in the referenced dict points to an (nred_grps * nfreqs x nvecs)
    numpy.ndarray describing the modeling components for each fitting group in the chunk.
ants_map: dict mapping integers to integers
    map between each antenna number to a unique index between 0 and Nants_data
    (typically the index of each antenna in ants_map)
nfreqs: int, optional
    number of frequency channels
dtype: numpy.dtype
    tensor data types
    default is np.float32

Returns
-------
fg_model_comps: list
    list of tf.Tensor objects where each tensor has shape (nvecs, ngrps, nbls, nfreqs)
    where nbls varies from tensor to tensor. Fitting groups with vectors that span nbls are lumped into the same
    modeling tensor along the ngrps axis. nvecs is chosen in chunk_fg_comp_dict_by_nbls
    to be the maximum number of vectors representing any of the ngrps baseline grps
    which means that many rows in nvecs will be zero. For example, if we are modeling with
    vectors that all span nbls=1 baseline and using delay-modes to model our data
    then nvecs will equal the largest number of delay modes necessary to model the wedge
    on all baselines even though the short baselines are described by far fewer modes
    on short baselines, most of the rows along the vector dimension will therefor be zero.
    This is wasteful of memory but it allows us to take advantage of the fast
    dense matrix operations on a GPU.

corr_inds: list
    list of list of lists of 2-tuples. Hierarchy of lists is
    chunk
        group
            baseline - (int 2-tuple)
def tensorize_fg_model_comps_dict(fg_model_comps_dict, ants_map, nfreqs, use_redundancy=False, dtype=np.float32, notebook_progressbar=False, verbose=False, grp_size_threshold=5):
    'Convert per-baseline model components into a Ndata x Ncomponent tensor\n\n    Parameters\n    ----------\n    fg_model_comps_dict: dict\n        dictionary where each key is a 2-tuple (nbl, nvecs) referring to the number\n        of baselines in each vector and the number of vectors. Each 2-tuple points to\n        a dictionary where each key is the fitting group in fg_comps_dict that includes\n        nbl baselines. Each key in the referenced dict points to an (nred_grps * nfreqs x nvecs)\n        numpy.ndarray describing the modeling components for each fitting group in the chunk.\n    ants_map: dict mapping integers to integers\n        map between each antenna number to a unique index between 0 and Nants_data\n        (typically the index of each antenna in ants_map)\n    nfreqs: int, optional\n        number of frequency channels\n    dtype: numpy.dtype\n        tensor data types\n        default is np.float32\n\n    Returns\n    -------\n    fg_model_comps: list\n        list of tf.Tensor objects where each tensor has shape (nvecs, ngrps, nbls, nfreqs)\n        where nbls varies from tensor to tensor. Fitting groups with vectors that span nbls are lumped into the same\n        modeling tensor along the ngrps axis. nvecs is chosen in chunk_fg_comp_dict_by_nbls\n        to be the maximum number of vectors representing any of the ngrps baseline grps\n        which means that many rows in nvecs will be zero. For example, if we are modeling with\n        vectors that all span nbls=1 baseline and using delay-modes to model our data\n        then nvecs will equal the largest number of delay modes necessary to model the wedge\n        on all baselines even though the short baselines are described by far fewer modes\n        on short baselines, most of the rows along the vector dimension will therefor be zero.\n        This is wasteful of memory but it allows us to take advantage of the fast\n        dense matrix operations on a GPU.\n\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\n    '
    echo(f'{datetime.datetime.now()} Computing foreground components matrices...
', verbose=verbose)
    fg_model_comps_dict = chunk_fg_comp_dict_by_nbls(fg_model_comps_dict, use_redundancy=use_redundancy, grp_size_threshold=grp_size_threshold)
    fg_model_comps = []
    corr_inds = []
    for (nbls, nvecs) in fg_model_comps_dict:
        ngrps = len(fg_model_comps_dict[(nbls, nvecs)])
        modeling_matrix = np.zeros((nvecs, ngrps, nbls, nfreqs))
        corr_inds_chunk = []
        for (grpnum, modeling_grp) in enumerate(fg_model_comps_dict[(nbls, nvecs)]):
            corr_inds_grp = []
            nbl = 0
            for (rgrpnum, red_grp) in enumerate(modeling_grp):
                nred = len(red_grp)
                for ap in red_grp:
                    (i, j) = (ants_map[ap[0]], ants_map[ap[1]])
                    corr_inds_grp.append((i, j))
                    vecslice = slice(0, fg_model_comps_dict[(nbls, nvecs)][modeling_grp].shape[1])
                    compslice = slice((rgrpnum * nfreqs), ((rgrpnum + 1) * nfreqs))
                    dslice = slice((nbl * nfreqs), ((nbl + 1) * nfreqs))
                    modeling_matrix[(vecslice, grpnum, nbl)] = fg_model_comps_dict[(nbls, nvecs)][modeling_grp][compslice].T
                    nbl += 1
            corr_inds_chunk.append(corr_inds_grp)
        fg_model_comps.append(tf.convert_to_tensor(modeling_matrix, dtype=dtype))
        corr_inds.append(corr_inds_chunk)
    return (fg_model_comps, corr_inds)