Marius Isken 11 months ago
parent
commit
e83cdb1bb8
  1. 73
      pyrocko_das/goldstein.py
  2. 10
      src/lib.rs

73
pyrocko_das/goldstein.py

@ -7,10 +7,19 @@ import matplotlib.pyplot as plt
import numpy as num
from matplotlib.widgets import Slider
from scipy import signal
from matplotlib.colors import Normalize
from scipy.signal import butter, lfilter
from .utils import timeit
data = num.load(op.join(op.dirname(__file__), "data", "data-DAS-gfz2020wswf.npy"))
data = num.load("/home/marius/Downloads/sample_data.npy").astype(num.float32)
def butter_bandpass_filter(data, lowcut, highcut, fs, order=4):
b, a = butter(order, (lowcut, highcut), btype="bandpass", fs=fs)
y = lfilter(b, a, data, axis=0)
return y
@lru_cache
@ -42,7 +51,7 @@ def goldstein_filter(data, window_size: int = 32, overlap: int = 14, exponent=0.
if num.log2(window_size) % 1.0 or window_size < 4:
raise ValueError("window_size has to be pow(2) and > 4.")
if overlap > window_size / 2 - 1:
raise ValueError("Overlap is large small. Minimum window_size)/ 2 - 1.")
raise ValueError("Overlap is too large. Maximum overlap: window_size / 2 - 1.")
window_stride = window_size - overlap
window_non_overlap = window_size - 2 * overlap
@ -65,11 +74,11 @@ def goldstein_filter(data, window_size: int = 32, overlap: int = 14, exponent=0.
window_data = data[slice_x, slice_y]
window_fft = num.fft.rfft2(window_data)
window_fft_abs = num.abs(window_fft)
weights = num.abs(window_fft) ** exponent
# Optional
# window_fft_abs = signal.medfilt2d(window_fft_abs, kernel_size=3)
# weights = signal.medfilt2d(weights, kernel_size=3)
# window_coherency = data_coherency(window_data)
window_fft *= window_fft_abs ** exponent
window_fft *= weights
window_filtered = num.fft.irfft2(window_fft)
taper_this = taper[: window_filtered.shape[0], : window_filtered.shape[1]]
@ -84,36 +93,36 @@ def goldstein_filter(data, window_size: int = 32, overlap: int = 14, exponent=0.
def normalize_diff(data, filtered_data):
norm = num.abs(data).max()
filtered_data_norm = filtered_data / num.abs(filtered_data).max() * (1.0 / norm)
filtered_data_norm -= data / norm
return filtered_data_norm
return Normalize()(data) - Normalize()(filtered_data)
def plot_goldstein():
def r(data):
v = num.std(data)
return -v, v
from .rust import goldstein_filter as goldstein_filter_rust
goldstein_filter_rust = timeit(goldstein_filter_rust)
window_size = 64
overlap = 30
window_size = 32
overlap = 14
exponent = 0.5
data_filtered = goldstein_filter(data, window_size, overlap, exponent)
data_filtered_rust = goldstein_filter_rust(data, window_size, overlap, exponent)
imshow_kwargs = dict(aspect=5, cmap="viridis", vmin=-1.0, vmax=1.0)
fig, (ax1, ax2, ax3, ax4) = plt.subplots(1, 4, sharex=True, sharey=True)
ax1.imshow(data, aspect=5, cmap="viridis")
image_python = ax2.imshow(data_filtered, aspect=5, cmap="viridis")
image_rust = ax3.imshow(data_filtered_rust, aspect=5, cmap="viridis")
image_diff = ax4.imshow(
normalize_diff(data, data_filtered), aspect=5, cmap="viridis"
)
data_bp = butter_bandpass_filter(data, 1.0, 10.0, 50.0).astype(num.float32)
fig, (ax1, ax3, ax4) = plt.subplots(1, 3, sharex=True, sharey=True)
ax1.imshow(data, **imshow_kwargs)
# image_python = ax2.imshow(data, **imshow_kwargs)
image_rust = ax3.imshow(data, **imshow_kwargs)
image_diff = ax4.imshow(data, **imshow_kwargs)
ax_slider = plt.axes((0.25, 0.1, 0.65, 0.03))
ax1.set_title("Data Input")
ax2.set_title("Data Filtered (Python)")
# ax2.set_title("Data Filtered (Python)")
ax3.set_title("Data Filtered (Rust)")
ax4.set_title("Data Residual")
@ -128,25 +137,29 @@ def plot_goldstein():
def update(val):
global data
# data = data_bp
data_filtered = goldstein_filter(
data, window_size=window_size, exponent=val, overlap=overlap
)
# data_filtered = goldstein_filter(
# data, window_size=window_size, exponent=val, overlap=overlap
# )
data_filtered_rust = goldstein_filter_rust(
data_filtered = goldstein_filter_rust(
data, window_size=window_size, exponent=val, overlap=overlap
)
image_python.set_data(data_filtered)
image_python.set_norm(None)
# image_python.set_data(data_filtered)
# image_python.set_norm(Normalize(*r(data_filtered)))
image_rust.set_data(data_filtered_rust)
image_rust.set_norm(None)
image_rust.set_data(data_filtered)
image_rust.set_norm(Normalize(*r(data_filtered)))
image_diff.set_data(normalize_diff(data, data_filtered))
image_diff.set_norm(None)
norm_diff = normalize_diff(data, data_filtered)
image_diff.set_data(norm_diff)
image_diff.set_norm(Normalize(*r(norm_diff)))
fig.canvas.draw_idle()
update(exponent)
exp_slider.on_changed(update)
plt.show()

10
src/lib.rs

@ -36,6 +36,16 @@ fn goldstein_filter(
overlap: usize,
exponent: f32,
) -> Array2<f32> {
assert!(
((window_size as f64).log2() % 1.0 == 0.) && window_size > 4,
"window_size has to be pow(2) and > 4."
);
assert!(
overlap < (window_size / 2 - 1),
"Overlap is too large. Maximum overlap: window_size / 2 - 1."
);
let window_stride = window_size - overlap;
let window_non_overlap = window_size - 2 * overlap;
let npx_x = data.nrows();

Loading…
Cancel
Save