# AGPL-3 Karl Semich 2022 import numpy as np # At the moment, wavelets must be: # - Zero at zero f(0) == 0 # - Even functions f(x) == f(-x) # - Odd functions when slid by 90 degrees f(0.25 - x) == f(0.25 + x) # Presently x is passed in the modular range [0,1) but this is arbitrary, just remove it if helpful COSINE = lambda x: np.cos(2 * np.pi * x) STEP = lambda x: COSINE(x).round() def complex_wavelet(wavelet, x): return wavelet(x % 1) + wavelet((x - 0.25) % 1) * 1j def fftfreq(freq_count = None, sample_rate = None, min_freq = None, max_freq = None, dc_offset = True, complex = True, sample_time = None, repetition_time = None, repetition_samples = None, freq_sample_rate = None, freq_sample_time = None): ''' Calculates and returns the frequency bins used to convert between the time and frequency domains with a discrete Fourier transform. With no optional arguments, this function should be equivalent to numpy.fft.fftfreq . To specify frequencies in terms of repetitions / sample_count, set sample_rate to sample_count. If frequencies were in Hz (cycles/sec), this implies that sample_count should be considered to last 1 second, so the frequency becomes equal to the cycle count. Parameters: - freq_count: the number of frequency bins to generate - sample_rate: the time-domain sample rate (default: 1 sample/time_unit) - min_freq: the minimum frequency the signal contains - max_fraq: the maximum frequency the signal contains - dc_offset: whether to include a DC offset component (0 Hz) - complex: whether to generate negative frequencies - sample_time: sample_rate as the duration of a single sample - repetition_time: min_freq as a period time, possibly of a subsignal - repetition_samples: min_freq as a period size in samples - freq_sample_rate: convert to or from a different frequency-domain sample rate - freq_sample_time: freq_sample_rate as the duration of a single sample Returns a vector of sinusoid time scalings that can be used to perform or analyse a discrete Fourier transform. Multiply this vector by the time-domain sample rate to get the frequencies. ''' assert not sample_time or not sample_rate assert not freq_sample_time or not freq_sample_rate assert (min_freq, repetition_time, repetition_samples).count(None) >= 2 assert freq_count or min_freq or repetition_time or repetition_samples if sample_time is not None: sample_rate = 1 / sample_time if freq_sample_time is not None: freq_sample_rate = 1 / freq_sample_time sample_rate = sample_rate or freq_sample_rate or 1 freq_sample_rate = freq_sample_rate or sample_rate or 1 if repetition_time: min_freq = 1 / repetition_time elif repetition_samples: min_freq = freq_sample_rate / repetition_samples # here freq_count adopts the complex + DC offset value, for calculations # it might be clearer to use e.g. sample_count, at least consolidate the cases with lower if freq_count is None: freq_count = int(np.ceil(freq_sample_rate / min_freq)) else: if not dc_offset: freq_count += 1 if not complex: freq_count = int((freq_count - 1) * 2) if not min_freq: min_freq = freq_sample_rate / freq_count if not max_freq: #max_freq = freq_sample_rate / 2 max_freq = freq_count * min_freq / 2 if freq_count % 2 != 0: #max_freq -= freq_sample_rate / (2 * freq_count) max_freq -= min_freq / 2 min_freq /= sample_rate max_freq /= sample_rate if freq_count % 2 == 0: if complex: neg_freqs = np.linspace(-max_freq, -min_freq, num=freq_count // 2, endpoint=True) pos_freqs = -neg_freqs[:0:-1] else: pos_freqs = np.linspace(min_freq, max_freq, num=freq_count // 2, endpoint=True) neg_freqs = pos_freqs[:0] else: pos_freqs = np.linspace(min_freq, max_freq, num=freq_count // 2, endpoint=True) neg_freqs = -pos_freqs[::-1] if complex else pos_freqs[:0] return np.concatenate([ np.array([0] if dc_offset else []), pos_freqs, neg_freqs ]) def create_freq2time(time_count = None, freqs = None, wavelet = COSINE): ''' Creates a matrix that will perform an inverse discrete Fourier transform when it post-multiplies a vector of complex frequency magnitudes. Example time_data = spectrum @ create_freq2time(len(spectrum)) Parameters: - time_count: size of the output vector, defaults to the frequency bincount - freqs: frequency bins to convert, defaults to traditional FFT frequencies Returns: - an inverse discrete Fourier matrix of shape (len(freqs), time_count) ''' assert (time_count is not None) or (freqs is not None) if freqs is None: freqs = fftfreq(time_count) is_complex = (freqs < 0).any() has_dc_offset = bool(freqs[0] == 0) freq_count = len(freqs) if not has_dc_offset: freq_count += 1 if not is_complex: freq_count = (freq_count - 1) * 2 if time_count is None: time_count = freq_count offsets = np.arange(time_count) if not is_complex: has_nyquist = bool(freqs[-1] == 0.5) head_idx = 1 if has_dc_offset else 0 tail_idx = len(freqs) - (1 if has_nyquist else 0) mat = complex_wavelet(wavelet, np.outer(np.concatenate([freqs[:tail_idx], -freqs[tail_idx:]]), offsets)) # todo: remove test data full_mat = complex_wavelet(wavelet, np.outer(np.concatenate([freqs[:tail_idx], -freqs[tail_idx:], -freqs[head_idx:tail_idx][::-1]]), offsets)) time_data = np.random.random(full_mat.shape[1]) extended_freq_data = time_data @ np.linalg.inv(full_mat) freq_data = extended_freq_data[:len(freqs)] import pdb; pdb.set_trace() mat[head_idx:tail_idx,:] += complex_wavelet(wavelet, np.outer(-freqs[head_idx:tail_idx], offsets)) mat = mat.real else: mat = complex_wavelet(wavelet, np.outer(freqs, offsets)) return mat / freq_count # scaled to match numpy convention def create_time2freq(time_count = None, freqs = None, wavelet = COSINE): ''' Creates a matrix that will perform a forward discrete Fourier transform when it post-multiplies a vector of time series data. If time_count is too small or large, the minimal least squares solution over all the data passed will be produced. This function is equivalent to calling .pinv() on the return value of create_freq2time. If the return value is single-use, it is more efficient and accurate to use numpy.linalg.lstsq . Example spectrum = time_data @ create_time2freq(len(time_data)) Parameters: - time_count: size of the input vector, defaults to the frequency bincount - freqs: frequency bins to produce, defaults to traditional FFT frequencies Returns: - a discrete Fourier matrix of shape (time_count, len(freqs)) ''' assert (time_count is not None) or (freqs is not None) if freqs is None: freqs = fftfreq(time_count) is_complex = (freqs < 0).any() # TODO: provide for DC offsets and nyquists not at index 0 by excluding items in (0,0.5) has_dc_offset = bool(freqs[0] == 0) has_nyquist = bool(freqs[-1] == 0.5) head_idx = 1 if has_dc_offset else 0 tail_idx = len(freqs)-1 if has_nyquist else len(freqs) neg_start_idx = len(freqs) if not is_complex: freqs = np.concatenate([ freqs[:head_idx], freqs[head_idx:tail_idx], -freqs[tail_idx:neg_start_idx], -freqs[head_idx:tail_idx][::-1] ]) inverse_mat = create_freq2time(time_count, freqs, wavelet) forward_mat = np.linalg.pinv(inverse_mat) if not is_complex: forward_mat = np.concatenate([ forward_mat[:,:tail_idx], forward_mat[:,tail_idx:neg_start_idx].conj() ], axis=1) return forward_mat def peak_pair_idcs(freq_data, dc_offset=True, sum=True): dc_offset = int(dc_offset) freq_heights = abs(freq_data) # squares and sums the components if sum: while len(freq_heights).shape > 1: freq_heights = freq_height.sum(axis=0) paired_heights = freq_heights[...,dc_offset:-1] + freq_heights[...,dc_offset+1:] peak_idx = paired_heights.argmax(axis=-1, keepdims=True) + dc_offset return np.concatenate(peak_idx, peak_idx + 1, axis=-1) def improve_peak(time_data, min_freq, max_freq): freqs = fftfreq(time_data.shape[-1], min_freq = min_freq, max_freq = max_freq) freq_data = time_data @ np.linalg.inv(create_freq2time(freqs)) left, right = peak_pair_idcs(freq_data) return freq_data[left], freq_data[right] def test(): np.random.seed(0) randvec = np.random.random(16) freqs16 = fftfreq(16) ift16 = create_freq2time(16, freqs=freqs16) ft16 = create_time2freq(16) randvec2time = randvec @ ift16 randvec2freq = randvec @ ft16 randvec2ifft = np.fft.ifft(randvec) randvec2fft = np.fft.fft(randvec) assert np.allclose(freqs16, np.fft.fftfreq(16)) assert np.allclose(randvec2ifft, randvec2time) assert np.allclose(randvec2fft, randvec2freq) assert np.allclose(randvec2ifft, np.linalg.solve(ft16.T, randvec)) assert np.allclose(randvec2fft, np.linalg.solve(ift16.T, randvec)) rfreqs15t = fftfreq(repetition_samples=15, complex=False) rfreqs15f = fftfreq(15, complex=False) irft15 = create_freq2time(freqs=rfreqs15f) rft15 = create_time2freq(15, freqs=rfreqs15t) randvec2rtime15 = randvec[:15] @ irft15 randvec2rfreq15 = randvec[:15] @ rft15 randvec2irfft = np.fft.irfft(randvec[:15]) randvec2rfft = np.fft.rfft(randvec[:15]) assert np.allclose(rfreqs15t, np.fft.rfftfreq(15)) assert np.allclose(rfreqs15f, np.fft.rfftfreq(28)) assert np.allclose(randvec2rtime15, randvec2irfft) assert np.allclose(randvec2rfreq15, randvec2rfft) rfreqs16t = fftfreq(repetition_samples=16, complex=False) rfreqs16f = fftfreq(16, complex=False) irft16 = create_freq2time(freqs=rfreqs16f) rft16 = create_time2freq(16, freqs=rfreqs16t) randvec2rtime16 = randvec[:16] @ irft16 randvec2rfreq16 = randvec[:16] @ rft16 randvec2irfft = np.fft.irfft(randvec[:16]) randvec2rfft = np.fft.rfft(randvec[:16]) assert np.allclose(rfreqs16t, np.fft.rfftfreq(16)) assert np.allclose(rfreqs16f, np.fft.rfftfreq(30)) assert np.allclose(randvec2rtime16, randvec2irfft) assert np.allclose(randvec2rfreq16, randvec2rfft) # [ 0, 1, 2, 3, 4, 3, 2, 1] #irft9_16 = create_freq2time(16, fftfreq(9, complex = False)) #rft16_30 = create_time2freq(16, fftfreq(30, complex = False)) #randvec2 #assert np.allclose((randvec @ rft16) @ irft16, randvec) # sample data at a differing rate time_rate = np.random.random() * 2 freq_rate = 1.0 freqs = np.fft.fftfreq(len(randvec)) rescaling_freqs = fftfreq(len(randvec), freq_sample_rate = freq_rate, sample_rate = time_rate) rescaling_ift = create_freq2time(freqs = rescaling_freqs) rescaling_ft = create_time2freq(freqs = rescaling_freqs) rescaled_time_data = np.array([ np.mean([ randvec[freqidx] * np.exp(2j * np.pi * freqs[freqidx] * sampleidx / time_rate) for freqidx in range(len(randvec)) ]) for sampleidx in range(len(randvec)) ]) assert np.allclose(rescaled_time_data, randvec @ rescaling_ift) assert np.allclose(rescaled_time_data, np.linalg.solve(rescaling_ft.T, randvec)) unscaled_freq_data = rescaled_time_data @ rescaling_ft unscaled_time_data = unscaled_freq_data @ ift16 assert np.allclose(unscaled_freq_data, randvec) assert np.allclose(unscaled_time_data, randvec2time) assert np.allclose(np.linalg.solve(rescaling_ift.T, rescaled_time_data), randvec) # extract a repeating wave longvec = np.empty(len(randvec)*100) short_count = np.random.random() * 4 + 1 short_duration = len(longvec) / short_count for idx in range(len(longvec)): longvec[idx] = randvec[int((idx / len(longvec) * short_duration) % len(randvec))] shortspace_freqs = fftfreq(None, complex = False, dc_offset = True, repetition_samples = len(randvec)) longspace_freqs = fftfreq(len(shortspace_freqs), complex = False, dc_offset = True, repetition_samples = short_duration) assert np.allclose(longspace_freqs, shortspace_freqs / (short_duration / len(randvec))) inserting_ift = create_freq2time(len(longvec), longspace_freqs, wavelet = STEP) extracting_ft = create_time2freq(len(longvec), longspace_freqs, wavelet = STEP) extracting_ift = create_freq2time(freqs = shortspace_freqs, wavelet = STEP) unextracting_ft = create_time2freq(freqs = shortspace_freqs, wavelet = STEP) inserting_spectrum = randvec @ unextracting_ft assert np.allclose(inserting_spectrum @ extracting_ift, randvec) assert np.allclose(longvec, inserting_spectrum @ inserting_ift) extracted_freqs = longvec @ extracting_ft extracted_randvec = extracted_freqs @ extracting_ift assert np.allclose(extracted_randvec, randvec) if __name__ == '__main__': test()