Handling missing data for mne_connectivity

MNE version: 1.6.1
Operating system: openSUSE Leap 15.3

Hello,

I would like to ask about how MNE-Connectivity (or MNE in general) handle bad data segments / missing data? For my research, I need to keep track of the timing and bad trials while calculating one functional connectivity metric per trial. I do this via a rather tedious and admittedly error prone procedure. Specifically, I initialize a new empty np array con_epoch below with the shape n_chan, n_chan, ntrial, nfreq. Then I loop across trials, if the trial is in epochs.selection, then calculate the FC and add it to the con_epoch, else, put nans to con_epoch. In code, looks like this:

    ntrial = len(epochs.drop_log)
    nfreq = len(fmin)
    nchan = epochs.get_data().shape[1]
    con_epoch = np.zeros((nchan, nchan, ntrial, nfreq))
    c = 0
    for i in range(ntrial):
        if i in epochs.selection:
            con = spectral_connectivity_epochs(
                epochs[c],
                method=method,
                mode="multitaper",
                sfreq=epochs.info["sfreq"],
            )
            con_epoch[:, :, i, :] = con.get_data(output="dense")
            c += 1
        else:
            con_epoch[:, :, i] = np.nan

Note that I have to use two iteration variables, i and c since epochs do not contain the same amount of trials as epochs.drop_log which keeps track of bad data.

Is there a cleaner way to handle the procedure?

Thanks for the amazing work with the package,
Yasir

edit: added more clarification