wave_mode_clustering cleanup fails on localized waves

[dhasegan]

Part of wave_mode_clustering script, is running a PCA and clustering the clusters. For such a clustering, the current algorithm requires that there are no NaN (not a number) so there is a step of cleanup. There are multiple steps of cleanup, including removing all rows and all columns that have all NaNs. Then within the remaining matrix, there are NaNs that are for some waves but not for others: showing localization to a specific part of the grid. From what I see, the remaining cleanup finds the closest (5, but parameterized) neighbors for each wave and fills the NaN based on the average value from those neighbors for the missing channels.
https://github.com/NeuralEnsemble/cobrawap/blob/master/cobrawap/pipeline/stage04_wave_detection/scripts/wave_mode_clustering.py

The bug: if the closest neighbors do not cover the rest of the channels, this method does not fill up the channels data.

Possible Solutions:

1. change the PCA to be able to use the missing channels, NaNs
2. fill the missing channels with the average channel data from all electrodes. This might result in distorted waves

Heres the error:

Traceback (most recent call last):
  File "/Users/angelinazuo/workspace/TravelingWaves/cobrawap/stage04_wave_detection/scripts/wave_mode_clustering.py", line 336, in <module>
    timelag_matrix_transformed = pca_transform(timelag_df, dims=args.pca_dims)
                                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/angelinazuo/workspace/TravelingWaves/cobrawap/stage04_wave_detection/scripts/wave_mode_clustering.py", line 144, in pca_transform
    pca_out = PCA(n_components=dims).fit(x_scaled)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/angelinazuo/anaconda3/envs/test_env/lib/python3.11/site-packages/sklearn/base.py", line 1389, in wrapper
    return fit_method(estimator, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/angelinazuo/anaconda3/envs/test_env/lib/python3.11/site-packages/sklearn/decomposition/_pca.py", line 442, in fit
    self._fit(X)
  File "/Users/angelinazuo/anaconda3/envs/test_env/lib/python3.11/site-packages/sklearn/decomposition/_pca.py", line 505, in _fit
    X = validate_data(
        ^^^^^^^^^^^^^^
  File "/Users/angelinazuo/anaconda3/envs/test_env/lib/python3.11/site-packages/sklearn/utils/validation.py", line 2944, in validate_data
    out = check_array(X, input_name="X", **check_params)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/angelinazuo/anaconda3/envs/test_env/lib/python3.11/site-packages/sklearn/utils/validation.py", line 1107, in check_array
    _assert_all_finite(
  File "/Users/angelinazuo/anaconda3/envs/test_env/lib/python3.11/site-packages/sklearn/utils/validation.py", line 120, in _assert_all_finite
    _assert_all_finite_element_wise(
  File "/Users/angelinazuo/anaconda3/envs/test_env/lib/python3.11/site-packages/sklearn/utils/validation.py", line 169, in _assert_all_finite_element_wise
    raise ValueError(msg_err)
ValueError: Input X contains NaN.