def test_notch_filters(): """Test notch filters.""" # let's use an ugly, prime sfreq for fun sfreq = 487.0 sig_len_secs = 20 t = np.arange(0, int(sig_len_secs * sfreq)) / sfreq freqs = np.arange(60, 241, 60) # make a "signal" a = rng.randn(int(sig_len_secs * sfreq)) orig_power = np.sqrt(np.mean(a ** 2)) # make line noise a += np.sum([np.sin(2 * np.pi * f * t) for f in freqs], axis=0) # only allow None line_freqs with 'spectrum_fit' mode assert_raises(ValueError, notch_filter, a, sfreq, None, 'fft') assert_raises(ValueError, notch_filter, a, sfreq, None, 'iir') methods = ['spectrum_fit', 'spectrum_fit', 'fft', 'fft', 'iir'] filter_lengths = ['auto', 'auto', 'auto', 8192, 'auto'] line_freqs = [None, freqs, freqs, freqs, freqs] tols = [2, 1, 1, 1] for meth, lf, fl, tol in zip(methods, line_freqs, filter_lengths, tols): with catch_logging() as log_file: with warnings.catch_warnings(record=True): b = notch_filter(a, sfreq, lf, fl, method=meth, verbose=True) if lf is None: out = log_file.getvalue().split('\n')[:-1] if len(out) != 2 and len(out) != 3: # force_serial: len(out) == 3 raise ValueError('Detected frequencies not logged properly') out = np.fromstring(out[-1], sep=', ') assert_array_almost_equal(out, freqs) new_power = np.sqrt(sum_squared(b) / b.size) assert_almost_equal(new_power, orig_power, tol)
def test_notch_filters(): """Test notch filters """ # let's use an ugly, prime sfreq for fun sfreq = 487.0 sig_len_secs = 20 t = np.arange(0, int(sig_len_secs * sfreq)) / sfreq freqs = np.arange(60, 241, 60) # make a "signal" rng = np.random.RandomState(0) a = rng.randn(int(sig_len_secs * sfreq)) orig_power = np.sqrt(np.mean(a ** 2)) # make line noise a += np.sum([np.sin(2 * np.pi * f * t) for f in freqs], axis=0) # only allow None line_freqs with 'spectrum_fit' mode assert_raises(ValueError, notch_filter, a, sfreq, None, "fft") assert_raises(ValueError, notch_filter, a, sfreq, None, "iir") methods = ["spectrum_fit", "spectrum_fit", "fft", "fft", "iir"] filter_lengths = [None, None, None, 8192, None] line_freqs = [None, freqs, freqs, freqs, freqs] tols = [2, 1, 1, 1] for meth, lf, fl, tol in zip(methods, line_freqs, filter_lengths, tols): with catch_logging() as log_file: b = notch_filter(a, sfreq, lf, filter_length=fl, method=meth, verbose="INFO") if lf is None: out = log_file.getvalue().split("\n")[:-1] if len(out) != 2: raise ValueError("Detected frequencies not logged properly") out = np.fromstring(out[1], sep=", ") assert_array_almost_equal(out, freqs) new_power = np.sqrt(sum_squared(b) / b.size) assert_almost_equal(new_power, orig_power, tol)
def test_notch_filters(method, filter_length, line_freq, tol): """Test notch filters.""" # let's use an ugly, prime sfreq for fun rng = np.random.RandomState(0) sfreq = 487 sig_len_secs = 21 t = np.arange(0, int(round(sig_len_secs * sfreq))) / sfreq # make a "signal" a = rng.randn(int(sig_len_secs * sfreq)) orig_power = np.sqrt(np.mean(a**2)) # make line noise a += np.sum([np.sin(2 * np.pi * f * t) for f in line_freqs], axis=0) # only allow None line_freqs with 'spectrum_fit' mode for kind in ('fir', 'iir'): with pytest.raises(ValueError, match='freqs=None can only be used wi'): notch_filter(a, sfreq, None, kind) with catch_logging() as log_file: b = notch_filter(a, sfreq, line_freq, filter_length, method=method, verbose=True) if line_freq is None: out = [ line.strip().split(':')[0] for line in log_file.getvalue().split('\n') if line.startswith(' ') ] assert len(out) == 4, 'Detected frequencies not logged properly' out = np.array(out, float) assert_array_almost_equal(out, line_freqs) new_power = np.sqrt(sum_squared(b) / b.size) assert_almost_equal(new_power, orig_power, tol)
def test_csd_on_artificial_data(): """Test computing CSD on artificial data. """ epochs = _get_data(mode='sin') sfreq = epochs.info['sfreq'] # Computing signal power in the time domain signal_power = sum_squared(epochs._data) signal_power_per_sample = signal_power / len(epochs.times) # Computing signal power in the frequency domain data_csd_mt, freqs_mt = csd_array(epochs._data, sfreq, mode='multitaper') data_csd_fourier, freqs_fft = csd_array(epochs._data, sfreq, mode='fourier') fourier_power = np.abs(data_csd_fourier[0, 0]) * sfreq mt_power = np.abs(data_csd_mt[0, 0]) * sfreq assert_true(abs(fourier_power - signal_power) <= 0.5) assert_true(abs(mt_power - signal_power) <= 1) # Power per sample should not depend on time window length for tmax in [0.2, 0.8]: tslice = np.where(epochs.times <= tmax)[0] for add_n_fft in [0, 30]: t_mask = (epochs.times >= 0) & (epochs.times <= tmax) n_samples = sum(t_mask) n_fft = n_samples + add_n_fft data_csd_fourier, _ = csd_array(epochs._data[:, :, tslice], sfreq, mode='fourier', fmin=0, fmax=np.inf, n_fft=n_fft) first_samp = data_csd_fourier[0, 0] fourier_power_per_sample = np.abs(first_samp) * sfreq / n_fft assert_true( abs(signal_power_per_sample - fourier_power_per_sample) < 0.003) # Power per sample should not depend on number of tapers for n_tapers in [1, 2, 5]: for add_n_fft in [0, 30]: mt_bandwidth = sfreq / float(n_samples) * (n_tapers + 1) data_csd_mt, _ = csd_array(epochs._data[:, :, tslice], sfreq, mt_bandwidth=mt_bandwidth, n_fft=n_fft) mt_power_per_sample = np.abs(data_csd_mt[0, 0]) *\ sfreq / n_fft # The estimate of power gets worse for small time windows when # more tapers are used if n_tapers == 5 and tmax == 0.2: delta = 0.05 else: delta = 0.004 assert_true( abs(signal_power_per_sample - mt_power_per_sample) < delta)
def test_compute_epochs_csd_on_artificial_data(): """Test computing CSD on artificial data """ epochs, epochs_sin = _get_data() sfreq = epochs_sin.info['sfreq'] # Computing signal power in the time domain signal_power = sum_squared(epochs_sin._data) signal_power_per_sample = signal_power / len(epochs_sin.times) # Computing signal power in the frequency domain data_csd_fourier = compute_epochs_csd(epochs_sin, mode='fourier') data_csd_mt = compute_epochs_csd(epochs_sin, mode='multitaper') fourier_power = np.abs(data_csd_fourier.data[0, 0]) * sfreq mt_power = np.abs(data_csd_mt.data[0, 0]) * sfreq assert_true(abs(fourier_power - signal_power) <= 0.5) assert_true(abs(mt_power - signal_power) <= 1) # Power per sample should not depend on time window length for tmax in [0.2, 0.4, 0.6, 0.8]: for add_n_fft in [30, 0, 30]: t_mask = (epochs_sin.times >= 0) & (epochs_sin.times <= tmax) n_samples = sum(t_mask) n_fft = n_samples + add_n_fft data_csd_fourier = compute_epochs_csd(epochs_sin, mode='fourier', tmin=None, tmax=tmax, fmin=0, fmax=np.inf, n_fft=n_fft) fourier_power_per_sample = np.abs(data_csd_fourier.data[0, 0]) *\ sfreq / data_csd_fourier.n_fft assert_true( abs(signal_power_per_sample - fourier_power_per_sample) < 0.003) # Power per sample should not depend on number of tapers for n_tapers in [1, 2, 3, 5]: for add_n_fft in [30, 0, 30]: mt_bandwidth = sfreq / float(n_samples) * (n_tapers + 1) data_csd_mt = compute_epochs_csd(epochs_sin, mode='multitaper', tmin=None, tmax=tmax, fmin=0, fmax=np.inf, mt_bandwidth=mt_bandwidth, n_fft=n_fft) mt_power_per_sample = np.abs(data_csd_mt.data[0, 0]) *\ sfreq / data_csd_mt.n_fft # The estimate of power gets worse for small time windows when # more tapers are used if n_tapers == 5 and tmax == 0.2: delta = 0.05 else: delta = 0.004 assert_true( abs(signal_power_per_sample - mt_power_per_sample) < delta)
def test_csd_multitaper(): """Test computing cross-spectral density using multitapers.""" epochs = _generate_coherence_data() sfreq = epochs.info['sfreq'] _test_fourier_multitaper_parameters(epochs, csd_multitaper, csd_array_multitaper) # Compute CSDs using various parameters times = [(None, None), (1, 9)] as_arrays = [False, True] adaptives = [False, True] parameters = product(times, as_arrays, adaptives) for (tmin, tmax), as_array, adaptive in parameters: if as_array: csd = csd_array_multitaper(epochs.get_data(), sfreq, epochs.tmin, adaptive=adaptive, fmin=9, fmax=23, tmin=tmin, tmax=tmax, ch_names=epochs.ch_names) else: csd = csd_multitaper(epochs, adaptive=adaptive, fmin=9, fmax=23, tmin=tmin, tmax=tmax) if tmin is None and tmax is None: assert csd.tmin == 0 and csd.tmax == 9.98 else: assert csd.tmin == tmin and csd.tmax == tmax csd = csd.mean([9.9, 14.9, 21.9], [10.1, 15.1, 22.1]) _test_csd_matrix(csd) # Test equivalence with PSD psd, psd_freqs = psd_multitaper(epochs, fmin=1e-3, normalization='full') # omit DC csd = csd_multitaper(epochs) assert_allclose(psd_freqs, csd.frequencies) csd = np.array([np.diag(csd.get_data(index=ii)) for ii in range(len(csd))]).T assert_allclose(psd[0], csd) # For the next test, generate a simple sine wave with a known power times = np.arange(20 * sfreq) / sfreq # 20 seconds of signal signal = np.sin(2 * np.pi * 10 * times)[None, None, :] # 10 Hz wave signal_power_per_sample = sum_squared(signal) / len(times) # Power per sample should not depend on time window length for tmax in [12, 18]: t_mask = (times <= tmax) n_samples = sum(t_mask) n_fft = len(times) # Power per sample should not depend on number of tapers for n_tapers in [1, 2, 5]: bandwidth = sfreq / float(n_samples) * (n_tapers + 1) csd_mt = csd_array_multitaper(signal, sfreq, tmax=tmax, bandwidth=bandwidth, n_fft=n_fft).sum().get_data() mt_power_per_sample = np.abs(csd_mt[0, 0]) * sfreq / n_fft assert abs(signal_power_per_sample - mt_power_per_sample) < 0.001
def test_csd_on_artificial_data(): """Test computing CSD on artificial data. """ # Ignore deprecation warnings for this test epochs = _generate_simple_data() sfreq = epochs.info['sfreq'] # Computing signal power in the time domain signal_power = sum_squared(epochs._data) signal_power_per_sample = signal_power / len(epochs.times) # Computing signal power in the frequency domain with warnings.catch_warnings(record=True): # deprecation csd_mt = csd_array(epochs._data, sfreq, mode='multitaper').get_data() csd_fourier = csd_array(epochs._data, sfreq, mode='fourier').get_data() fourier_power = np.abs(csd_fourier[0, 0]) * sfreq mt_power = np.abs(csd_mt[0, 0]) * sfreq assert abs(fourier_power - signal_power) <= 0.5 assert abs(mt_power - signal_power) <= 1 # Power per sample should not depend on time window length for tmax in [0.2, 0.8]: tslice = np.where(epochs.times <= tmax)[0] for add_n_fft in [0, 30]: t_mask = (epochs.times >= 0) & (epochs.times <= tmax) n_samples = sum(t_mask) n_fft = n_samples + add_n_fft with warnings.catch_warnings(record=True): # deprecation csd_fourier = csd_array(epochs._data[:, :, tslice], sfreq, mode='fourier', fmin=0, fmax=np.inf, n_fft=n_fft).get_data() first_samp = csd_fourier[0, 0] fourier_power_per_sample = np.abs(first_samp) * sfreq / n_fft assert abs(signal_power_per_sample - fourier_power_per_sample) < 0.003 # Power per sample should not depend on number of tapers for n_tapers in [1, 2, 5]: mt_bandwidth = sfreq / float(n_samples) * (n_tapers + 1) with warnings.catch_warnings(record=True): # deprecation csd_mt = csd_array( epochs._data[:, :, tslice], sfreq, mt_bandwidth=mt_bandwidth, n_fft=n_fft ).get_data() mt_power_per_sample = np.abs(csd_mt[0, 0]) * sfreq / n_fft # The estimate of power gets worse for small time windows when more # tapers are used if n_tapers == 5 and tmax == 0.2: delta = 0.05 else: delta = 0.004 assert abs(signal_power_per_sample - mt_power_per_sample) < delta
def test_csd_on_artificial_data(): """Test computing CSD on artificial data. """ epochs = _get_data(mode='sin') sfreq = epochs.info['sfreq'] # Computing signal power in the time domain signal_power = sum_squared(epochs._data) signal_power_per_sample = signal_power / len(epochs.times) # Computing signal power in the frequency domain data_csd_mt, freqs_mt = csd_array(epochs._data, sfreq, mode='multitaper') data_csd_fourier, freqs_fft = csd_array(epochs._data, sfreq, mode='fourier') fourier_power = np.abs(data_csd_fourier[0, 0]) * sfreq mt_power = np.abs(data_csd_mt[0, 0]) * sfreq assert_true(abs(fourier_power - signal_power) <= 0.5) assert_true(abs(mt_power - signal_power) <= 1) # Power per sample should not depend on time window length for tmax in [0.2, 0.8]: tslice = np.where(epochs.times <= tmax)[0] for add_n_fft in [0, 30]: t_mask = (epochs.times >= 0) & (epochs.times <= tmax) n_samples = sum(t_mask) n_fft = n_samples + add_n_fft data_csd_fourier, _ = csd_array(epochs._data[:, :, tslice], sfreq, mode='fourier', fmin=0, fmax=np.inf, n_fft=n_fft) first_samp = data_csd_fourier[0, 0] fourier_power_per_sample = np.abs(first_samp) * sfreq / n_fft assert_true(abs(signal_power_per_sample - fourier_power_per_sample) < 0.003) # Power per sample should not depend on number of tapers for n_tapers in [1, 2, 5]: for add_n_fft in [0, 30]: mt_bandwidth = sfreq / float(n_samples) * (n_tapers + 1) data_csd_mt, _ = csd_array(epochs._data[:, :, tslice], sfreq, mt_bandwidth=mt_bandwidth, n_fft=n_fft) mt_power_per_sample = np.abs(data_csd_mt[0, 0]) *\ sfreq / n_fft # The estimate of power gets worse for small time windows when # more tapers are used if n_tapers == 5 and tmax == 0.2: delta = 0.05 else: delta = 0.004 assert_true(abs(signal_power_per_sample - mt_power_per_sample) < delta)
def test_csd_multitaper(): """Test computing cross-spectral density using multitapers.""" epochs = _generate_coherence_data() sfreq = epochs.info['sfreq'] _test_fourier_multitaper_parameters(epochs, csd_multitaper, csd_array_multitaper) # Compute CSDs using various parameters times = [(None, None), (1, 9)] as_arrays = [False, True] adaptives = [False, True] parameters = product(times, as_arrays, adaptives) for (tmin, tmax), as_array, adaptive in parameters: if as_array: csd = csd_array_multitaper(epochs.get_data(), sfreq, epochs.tmin, adaptive=adaptive, fmin=9, fmax=23, tmin=tmin, tmax=tmax) else: csd = csd_multitaper(epochs, adaptive=adaptive, fmin=9, fmax=23, tmin=tmin, tmax=tmax) csd = csd.mean([9.9, 14.9, 21.9], [10.1, 15.1, 22.1]) _test_csd_matrix(csd) # For the next test, generate a simple sine wave with a known power times = np.arange(20 * sfreq) / sfreq # 20 seconds of signal signal = np.sin(2 * np.pi * 10 * times)[None, None, :] # 10 Hz wave signal_power_per_sample = sum_squared(signal) / len(times) # Power per sample should not depend on time window length for tmax in [12, 18]: t_mask = (times <= tmax) n_samples = sum(t_mask) n_fft = len(times) # Power per sample should not depend on number of tapers for n_tapers in [1, 2, 5]: bandwidth = sfreq / float(n_samples) * (n_tapers + 1) csd_mt = csd_array_multitaper(signal, sfreq, tmax=tmax, bandwidth=bandwidth, n_fft=n_fft).sum().get_data() mt_power_per_sample = np.abs(csd_mt[0, 0]) * sfreq / n_fft assert abs(signal_power_per_sample - mt_power_per_sample) < 0.001
def test_compute_epochs_csd_on_artificial_data(): """Test computing CSD on artificial data """ epochs, epochs_sin = _get_data() sfreq = epochs_sin.info['sfreq'] # Computing signal power in the time domain signal_power = sum_squared(epochs_sin._data) signal_power_per_sample = signal_power / len(epochs_sin.times) # Computing signal power in the frequency domain data_csd_fourier = compute_epochs_csd(epochs_sin, mode='fourier') data_csd_mt = compute_epochs_csd(epochs_sin, mode='multitaper') fourier_power = np.abs(data_csd_fourier.data[0, 0]) * sfreq mt_power = np.abs(data_csd_mt.data[0, 0]) * sfreq assert_almost_equal(fourier_power, signal_power, delta=0.5) assert_almost_equal(mt_power, signal_power, delta=1) # Power per sample should not depend on time window length for tmax in [0.2, 0.4, 0.6, 0.8]: for add_n_fft in [30, 0, 30]: t_mask = (epochs_sin.times >= 0) & (epochs_sin.times <= tmax) n_samples = sum(t_mask) n_fft = n_samples + add_n_fft data_csd_fourier = compute_epochs_csd(epochs_sin, mode='fourier', tmin=None, tmax=tmax, fmin=0, fmax=np.inf, n_fft=n_fft) fourier_power_per_sample = np.abs(data_csd_fourier.data[0, 0]) *\ sfreq / data_csd_fourier.n_fft assert_almost_equal(signal_power_per_sample, fourier_power_per_sample, delta=0.003) # Power per sample should not depend on number of tapers for n_tapers in [1, 2, 3, 5]: for add_n_fft in [30, 0, 30]: mt_bandwidth = sfreq / float(n_samples) * (n_tapers + 1) data_csd_mt = compute_epochs_csd(epochs_sin, mode='multitaper', tmin=None, tmax=tmax, fmin=0, fmax=np.inf, mt_bandwidth=mt_bandwidth, n_fft=n_fft) mt_power_per_sample = np.abs(data_csd_mt.data[0, 0]) *\ sfreq / data_csd_mt.n_fft # The estimate of power gets worse for small time windows when # more tapers are used if n_tapers == 5 and tmax == 0.2: delta = 0.05 else: delta = 0.004 assert_almost_equal(signal_power_per_sample, mt_power_per_sample, delta=delta)
def test_csd_fourier(): """Test computing cross-spectral density using short-term Fourier.""" epochs = _generate_coherence_data() sfreq = epochs.info['sfreq'] _test_fourier_multitaper_parameters(epochs, csd_fourier, csd_array_fourier) # Compute CSDs using various parameters times = [(None, None), (1, 9)] as_arrays = [False, True] parameters = product(times, as_arrays) for (tmin, tmax), as_array in parameters: if as_array: csd = csd_array_fourier(epochs.get_data(), sfreq, epochs.tmin, fmin=9, fmax=23, tmin=tmin, tmax=tmax, ch_names=epochs.ch_names) else: csd = csd_fourier(epochs, fmin=9, fmax=23, tmin=tmin, tmax=tmax) if tmin is None and tmax is None: assert csd.tmin == 0 and csd.tmax == 9.98 else: assert csd.tmin == tmin and csd.tmax == tmax csd = csd.mean([9.9, 14.9, 21.9], [10.1, 15.1, 22.1]) _test_csd_matrix(csd) # For the next test, generate a simple sine wave with a known power times = np.arange(20 * sfreq) / sfreq # 20 seconds of signal signal = np.sin(2 * np.pi * 10 * times)[None, None, :] # 10 Hz wave signal_power_per_sample = sum_squared(signal) / len(times) # Power per sample should not depend on time window length for tmax in [12, 18]: t_mask = (times <= tmax) n_samples = sum(t_mask) # Power per sample should not depend on number of FFT points for add_n_fft in [0, 30]: n_fft = n_samples + add_n_fft csd = csd_array_fourier(signal, sfreq, tmax=tmax, n_fft=n_fft).sum().get_data() first_samp = csd[0, 0] fourier_power_per_sample = np.abs(first_samp) * sfreq / n_fft assert abs(signal_power_per_sample - fourier_power_per_sample) < 0.001
def test_notch_filters(): """Test notch filters """ tempdir = _TempDir() log_file = op.join(tempdir, 'temp_log.txt') # let's use an ugly, prime sfreq for fun sfreq = 487.0 sig_len_secs = 20 t = np.arange(0, int(sig_len_secs * sfreq)) / sfreq freqs = np.arange(60, 241, 60) # make a "signal" rng = np.random.RandomState(0) a = rng.randn(int(sig_len_secs * sfreq)) orig_power = np.sqrt(np.mean(a**2)) # make line noise a += np.sum([np.sin(2 * np.pi * f * t) for f in freqs], axis=0) # only allow None line_freqs with 'spectrum_fit' mode assert_raises(ValueError, notch_filter, a, sfreq, None, 'fft') assert_raises(ValueError, notch_filter, a, sfreq, None, 'iir') methods = ['spectrum_fit', 'spectrum_fit', 'fft', 'fft', 'iir'] filter_lengths = [None, None, None, 8192, None] line_freqs = [None, freqs, freqs, freqs, freqs] tols = [2, 1, 1, 1] for meth, lf, fl, tol in zip(methods, line_freqs, filter_lengths, tols): if lf is None: set_log_file(log_file, overwrite=True) b = notch_filter(a, sfreq, lf, filter_length=fl, method=meth, verbose='INFO') if lf is None: set_log_file() with open(log_file) as fid: out = fid.readlines() if len(out) != 2: raise ValueError('Detected frequencies not logged properly') out = np.fromstring(out[1], sep=', ') assert_array_almost_equal(out, freqs) new_power = np.sqrt(sum_squared(b) / b.size) assert_almost_equal(new_power, orig_power, tol)
def test_notch_filters(): """Test notch filters """ tempdir = _TempDir() log_file = op.join(tempdir, 'temp_log.txt') # let's use an ugly, prime sfreq for fun sfreq = 487.0 sig_len_secs = 20 t = np.arange(0, int(sig_len_secs * sfreq)) / sfreq freqs = np.arange(60, 241, 60) # make a "signal" rng = np.random.RandomState(0) a = rng.randn(int(sig_len_secs * sfreq)) orig_power = np.sqrt(np.mean(a ** 2)) # make line noise a += np.sum([np.sin(2 * np.pi * f * t) for f in freqs], axis=0) # only allow None line_freqs with 'spectrum_fit' mode assert_raises(ValueError, notch_filter, a, sfreq, None, 'fft') assert_raises(ValueError, notch_filter, a, sfreq, None, 'iir') methods = ['spectrum_fit', 'spectrum_fit', 'fft', 'fft', 'iir'] filter_lengths = [None, None, None, 8192, None] line_freqs = [None, freqs, freqs, freqs, freqs] tols = [2, 1, 1, 1] for meth, lf, fl, tol in zip(methods, line_freqs, filter_lengths, tols): if lf is None: set_log_file(log_file, overwrite=True) b = notch_filter(a, sfreq, lf, filter_length=fl, method=meth, verbose='INFO') if lf is None: set_log_file() with open(log_file) as fid: out = fid.readlines() if len(out) != 2: raise ValueError('Detected frequencies not logged properly') out = np.fromstring(out[1], sep=', ') assert_array_almost_equal(out, freqs) new_power = np.sqrt(sum_squared(b) / b.size) assert_almost_equal(new_power, orig_power, tol)
def test_sum_squared(): """Test optimized sum of squares """ X = np.random.randint(0, 50, (3, 3)) assert_equal(np.sum(X**2), sum_squared(X))
def test_sum_squared(): """Test optimized sum of squares.""" X = np.random.RandomState(0).randint(0, 50, (3, 3)) assert np.sum(X**2) == sum_squared(X)
def test_sum_squared(): """Test optimized sum of squares.""" X = np.random.RandomState(0).randint(0, 50, (3, 3)) assert_equal(np.sum(X ** 2), sum_squared(X))