ENH: Enable mne.compute_covariance() to handle nan values

Describe the new feature or enhancement

With the advent of "big data", I feel like the current support for handling and concatenating multiple recordings could be improved. The main bottleneck here is the handling of bad channels. Currently, either a segment is bad or a channel is bad. There is no way to mark a channel as bad for a part of a recording. This means, when concatenating many recordings into one raw file, one either loses important data because 1) an entire recording is marked as a bad segment due to a single bad channel or because 2) the given channel is marked as bad for all recordings. This is already discussed in #11481.

In that issue, I propose to handle partial bad channels by marking them as "good" and setting their bad data segments as nan. I currently do this manually and works perfectly with raw.compute_psd() using the maximum amount of data. nan data are simply automatically dropped by that function, however, only for the given channel.

However, this does not work with mne.compute_covariance():

def _create_toy_data(n_channels=3, sfreq=250, duration=100, seed=None):
    rng = np.random.default_rng(seed)
    data = rng.standard_normal(size=(n_channels, duration * sfreq)) * 5e-6
    info = mne.create_info(n_channels, sfreq, "eeg")
    return mne.io.RawArray(data, info)

raw = _create_toy_data(n_channels=n_ch, sfreq=s_freq, duration=duration)
raw._data[1, :2*s_freq] = np.nan

epochs = mne.make_fixed_length_epochs(raw, duration=1, preload=True)
mne.compute_covariance(epochs)

ValueError                                Traceback (most recent call last)
Cell In[37], line 1
----> 1 mne.compute_covariance(epochs)

File :12, in compute_covariance(epochs, keep_sample_mean, tmin, tmax, projs, method, method_params, cv, scalings, n_jobs, return_estimators, on_mismatch, rank, verbose)

File [/opt/anaconda3/envs/bids_neu/lib/python3.11/site-packages/mne/cov.py:961](https://file+.vscode-resource.vscode-cdn.net/opt/anaconda3/envs/bids_neu/lib/python3.11/site-packages/mne/cov.py:961), in compute_covariance(***failed resolving arguments***)
    958 _check_n_samples(n_samples_tot, len(picks_meeg))
    960 epochs = epochs.T  # sklearn | C-order
--> 961 cov_data = _compute_covariance_auto(
    962     epochs, method=method, method_params=_method_params, info=info,
    963     cv=cv, n_jobs=n_jobs, stop_early=True, picks_list=picks_list,
    964     scalings=scalings, rank=rank)
    966 if keep_sample_mean is False:
    967     cov = cov_data['empirical']['data']

File [/opt/anaconda3/envs/bids_neu/lib/python3.11/site-packages/mne/cov.py:1028](https://file+.vscode-resource.vscode-cdn.net/opt/anaconda3/envs/bids_neu/lib/python3.11/site-packages/mne/cov.py:1028), in _compute_covariance_auto(data, method, info, method_params, cv, scalings, n_jobs, stop_early, picks_list, rank)
   1026 # rescale to improve numerical stability
   1027 orig_rank = rank
-> 1028 rank = compute_rank(
   1029     RawArray(data.T, info, copy=None, verbose=_verbose_safe_false()),
   1030     rank, scalings, info)
   1031 with _scaled_array(data.T, picks_list, scalings):
   1032     C = np.dot(data.T, data)
...
--> 627     raise ValueError(
    628         "array must not contain infs or NaNs")
    629 return a

ValueError: array must not contain infs or NaNs

Describe your proposed implementation

I already looked at the function mne.compute_covariance(), however, it appears really complicated to me. Given that it relies on scipy, I don't know how easy it would be to support nan values.

Naively speaking, shouldn't it be possible to calculate the covariance matrices for all epochs separately and average them, in the end, using np.nanmean?

Describe possible alternatives

Maybe supporting nan values is not possible mathematically. I'm happy to learn.

Additional context

No response

Reproduce the provided RawArray and Epochs example, then inspect mne/cov.py around compute_covariance() and _compute_covariance_auto(), where the traceback reaches rank computation. Determine the intended covariance behavior when epochs contain NaN values and add coverage for the example; done means compute_covariance() handles the stated partial-channel case without the shown ValueError.