from operator import inv


def retrieval_gain_matrix(K, S_a, S_y):
    """Calculate the averaging kernel matrix.

    Parameters:
        K (np.array): Simulated Jacobians.
        S_a (np.array): A priori error covariance matrix.
        S_y (np.array): Measurement covariance matrix.

    Returns:
        np.array: Retrieval gain matrix.
    """
    return inv(inv(S_a) + K.T @ inv(S_y) @ K) @ K.T @ inv(S_y)