Implementing the coherence and amplitude spectra and fixing bugs in the cross spectra.

[sortega87]

This modifies the "get_cross_spectrum" function to return not only the cross spectra but also the coherence and phase. It also replaces multiple function that where not correctly implemented. It also includes an example of the functionality.

[sortega87]

Hi @HF7weatherman! Let me know what you think. Thanks!

[HF7weatherman]

Hi @sortega87,
Thanks for implementing these changes! I have a few comments:

• I responded to the normalization issue (Normalization of the Power Spectra #6). Given on how we decide there on how to handle the normalization, we should include this decision into the present PR. Therefore, I wouldn't merge at the moment.
• You removed the functions for calculating the cospectrum and quadrature spectrum, I guess because of your changes in the function for calculating the cross spectrum. Although I haven't had a real use case for them yet, I would be keen on keeping them since they are only a thin wrapper on top of the function that calculates the cross spectrum.
• You basically added the calculation of the coherence spectrum and the angle to the function get_cross_spectrum(). I think it would be nice to split those two, i.e. keep a separate function to calculate the cross spectrum and create a new function get_coherence_spectrum(). The latter could still return the cross spectrum as it does currently, but I think that a separate function would be cleaner and offers more flexibility. This might require some changes on the hidden functions as well, in particular _get_cross_spectrum(), but I think that shouldn't be too much work...

If this is unclear, let's discuss it in person! :)

[sortega87]

Thanks for the review @HF7weatherman!

Regarding the first two points you raise:

• I think the first point is mostly covered as the function computes recomputes the fft on each of the fields, instead of relying on the get_spectra funciton. Thus, the pull request will only modify the get_cross_spectrum function. However, I modified the definition also there to match what we have for the spectra.
• I have added back the functions for calculating the cospectrum and quadrature spectrum.

Regarding the last point, I separated the functionality to compute the coherence into a function:

def coherence_squared(cross_spectrum: xr.DataArray) -> xr.DataArray:
    """Compute the squared coherence and phase of a cross spectrum."""
    sxy2 = np.abs(cross_spectrum.cross)**2
    sxx = np.abs(cross_spectrum.spectra1)
    syy = np.abs(cross_spectrum.spectra2)
    coh2 =  sxy2 / (sxx * syy)
    coh2.name = "coh2"
    return coh2

Thus one would use this as follows:

symmetric_spectrum = get_cross_spectrum(
    variables_normalized,
    "symmetric",
    data_frequency,
    window_length,
    overlap_length,
    taper_alpha=taper_alpha,
)
coh2 = coherence_squared(symmetric_spectrum)

I think it might be better this way as all the information to compute the coherence is already in the cross spectrum.

Let me know what you think

Thanks!
Sebastian

[HF7weatherman]

Hi @sortega87,
thanks for implementing the changes based on my suggestions. I have a few more comments:

• If in a normal use case for the coherence spectrum one is only interested in the coherence spectrum in the end, and not in the cross spectrum, one could even wrap the calculation of the cross spectrum into the function for calculating the squared coherence spectrum:

def coherence_squared(*args, **kwargs) -> xr.DataArray:
    """Compute the squared coherence and phase of a cross spectrum."""
    cross_spectrum = get_cross_spectrum(*args, **kwargs)
    sxy2 = np.abs(cross_spectrum.cross)**2
    sxx = np.abs(cross_spectrum.spectra1)
    syy = np.abs(cross_spectrum.spectra2)
    coh2 =  sxy2 / (sxx * syy)
    coh2.name = "coh2"
    return coh2

• I wondered how you deal with the normalization of the coherence spectrum now. Currently it is normalized with (n_time * n_lon)**2, which makes sense for the other spectra, but as you said may be not so meaningful for the coherence spectrum. If we would wrap the calculation of the cross spectrum into the function for calculating the coherence spectrum as suggested in my previous point, one could easily multiply the coherence spectrum by (n_time * n_lon) as n_lon and n_time could be extracted from *args, thereby achieving the optimal normalization of the coherence spectrum without changing the one for the other spectra.
• You changed the output of the function _compute_hayashi_cross_spectrum() from the xr.DataArray cross_spectrum to the xr.Dataset xr.merge([spectrum_1, spectrum_2, cross_1_2]). That's fine as it is needed for the calculation of the coherence spectrum, but the functions get_co_spectrum() and get_quadrature_spectrum() need to be adjusted in that case. I have added my suggestions as review comments.

[sortega87]

Hi @HF7weatherman! I think this is ready for merging. Let me know what you think. Thanks!

[HF7weatherman]

Hi @sortega87! Apart from one small comment, I have no more things to add, so it's ready for merging. Thanks a lot for implementing this! :)

[sortega87]

Hi @sortega87! Apart from one small comment, I have no more things to add, so it's ready for merging. Thanks a lot for implementing this! :)

Thanks @HF7weatherman for the review!! :D I will merge now the branch into main.