externalized common func

K-ShakeTune/scripts/common_func.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python3
+
+# Common functions for the Shake&Tune package
+# Written by Frix_x#0161 #
+
+
+import numpy as np
+from scipy.signal import spectrogram
+
+
+# This is Klipper's spectrogram generation function adapted to use Scipy
+def compute_spectrogram(data):
+    N = data.shape[0]
+    Fs = N / (data[-1, 0] - data[0, 0])
+    # Round up to a power of 2 for faster FFT
+    M = 1 << int(.5 * Fs - 1).bit_length()
+    window = np.kaiser(M, 6.)
+
+    def _specgram(x):
+        x_detrended = x - np.mean(x)  # Detrending by subtracting the mean value
+        return spectrogram(
+            x_detrended, fs=Fs, window=window, nperseg=M, noverlap=M//2,
+            detrend='constant', scaling='density', mode='psd')
+
+    d = {'x': data[:, 1], 'y': data[:, 2], 'z': data[:, 3]}
+    f, t, pdata = _specgram(d['x'])
+    for axis in 'yz':
+        pdata += _specgram(d[axis])[2]
+    return pdata, t, f
+
+
+# This find all the peaks in a curve by looking at when the derivative term goes from positive to negative
+# Then only the peaks found above a threshold are kept to avoid capturing peaks in the low amplitude noise of a signal
+def detect_peaks(data, indices, detection_threshold, relative_height_threshold=None, window_size=5, vicinity=3):
+    # Smooth the curve using a moving average to avoid catching peaks everywhere in noisy signals
+    kernel = np.ones(window_size) / window_size
+    smoothed_data = np.convolve(data, kernel, mode='valid')
+    mean_pad = [np.mean(data[:window_size])] * (window_size // 2)
+    smoothed_data = np.concatenate((mean_pad, smoothed_data))
+    
+    # Find peaks on the smoothed curve
+    smoothed_peaks = np.where((smoothed_data[:-2] < smoothed_data[1:-1]) & (smoothed_data[1:-1] > smoothed_data[2:]))[0] + 1
+    smoothed_peaks = smoothed_peaks[smoothed_data[smoothed_peaks] > detection_threshold]
+    
+    # Additional validation for peaks based on relative height
+    valid_peaks = smoothed_peaks
+    if relative_height_threshold is not None:
+        valid_peaks = []
+        for peak in smoothed_peaks:
+            peak_height = smoothed_data[peak] - np.min(smoothed_data[max(0, peak-vicinity):min(len(smoothed_data), peak+vicinity+1)])
+            if peak_height > relative_height_threshold * smoothed_data[peak]:
+                valid_peaks.append(peak)
+
+    # Refine peak positions on the original curve
+    refined_peaks = []
+    for peak in valid_peaks:
+        local_max = peak + np.argmax(data[max(0, peak-vicinity):min(len(data), peak+vicinity+1)]) - vicinity
+        refined_peaks.append(local_max)
+    
+    num_peaks = len(refined_peaks)
+    
+    return num_peaks, np.array(refined_peaks), indices[refined_peaks]

K-ShakeTune/scripts/graph_belts.py

@@ -14,15 +14,17 @@
 import optparse, matplotlib, sys, importlib, os
 from collections import namedtuple
 import numpy as np
-import scipy
+from scipy.interpolate import griddata
 import matplotlib.pyplot, matplotlib.dates, matplotlib.font_manager
 import matplotlib.ticker, matplotlib.gridspec, matplotlib.colors
 import matplotlib.patches
-from locale_utils import set_locale, print_with_c_locale
 from datetime import datetime
 
 matplotlib.use('Agg')
 
+from locale_utils import set_locale, print_with_c_locale
+from common_func import compute_spectrogram, detect_peaks
+
 
 ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" # For paired peaks names
 
@@ -48,12 +50,6 @@ KLIPPAIN_COLORS = {
 # Computation of the PSD graph
 ######################################################################
 
-# Calculate estimated "power spectral density" using existing Klipper tools
-def calc_freq_response(data):
-    helper = shaper_calibrate.ShaperCalibrate(printer=None)
-    return helper.process_accelerometer_data(data)
-
-
 # Calculate or estimate a "similarity" factor between two PSD curves and scale it to a percentage. This is
 # used here to quantify how close the two belts path behavior and responses are close together.
 def compute_curve_similarity_factor(signal1, signal2):
@@ -76,29 +72,6 @@ def compute_curve_similarity_factor(signal1, signal2):
     return scaled_similarity
 
 
-# This find all the peaks in a curve by looking at when the derivative term goes from positive to negative
-# Then only the peaks found above a threshold are kept to avoid capturing peaks in the low amplitude noise of a signal
-def detect_peaks(psd, freqs, window_size=5, vicinity=3):
-    # Smooth the curve using a moving average to avoid catching peaks everywhere in noisy signals
-    kernel = np.ones(window_size) / window_size
-    smoothed_psd = np.convolve(psd, kernel, mode='valid')
-    mean_pad = [np.mean(psd[:window_size])] * (window_size // 2)
-    smoothed_psd = np.concatenate((mean_pad, smoothed_psd))
-    
-    # Find peaks on the smoothed curve
-    smoothed_peaks = np.where((smoothed_psd[:-2] < smoothed_psd[1:-1]) & (smoothed_psd[1:-1] > smoothed_psd[2:]))[0] + 1
-    detection_threshold = PEAKS_DETECTION_THRESHOLD * psd.max()
-    smoothed_peaks = smoothed_peaks[smoothed_psd[smoothed_peaks] > detection_threshold]
-    
-    # Refine peak positions on the original curve
-    refined_peaks = []
-    for peak in smoothed_peaks:
-        local_max = peak + np.argmax(psd[max(0, peak-vicinity):min(len(psd), peak+vicinity+1)]) - vicinity
-        refined_peaks.append(local_max)
-    
-    return np.array(refined_peaks), freqs[refined_peaks]
-
-
 # This function create pairs of peaks that are close in frequency on two curves (that are known
 # to be resonances points and must be similar on both belts on a CoreXY kinematic)
 def pair_peaks(peaks1, freqs1, psd1, peaks2, freqs2, psd2):
@@ -143,30 +116,6 @@ def pair_peaks(peaks1, freqs1, psd1, peaks2, freqs2, psd2):
     return paired_peaks, unpaired_peaks1, unpaired_peaks2
 
 
-######################################################################
-# Computation of a basic signal spectrogram
-######################################################################
-
-def compute_spectrogram(data):
-    N = data.shape[0]
-    Fs = N / (data[-1, 0] - data[0, 0])
-    # Round up to a power of 2 for faster FFT
-    M = 1 << int(.5 * Fs - 1).bit_length()
-    window = np.kaiser(M, 6.)
-
-    def _specgram(x):
-        x_detrended = x - np.mean(x)  # Detrending by subtracting the mean value
-        return scipy.signal.spectrogram(
-            x_detrended, fs=Fs, window=window, nperseg=M, noverlap=M//2,
-            detrend='constant', scaling='density', mode='psd')
-
-    d = {'x': data[:, 1], 'y': data[:, 2], 'z': data[:, 3]}
-    f, t, pdata = _specgram(d['x'])
-    for axis in 'yz':
-        pdata += _specgram(d[axis])[2]
-    return pdata, t, f
-
-
 ######################################################################
 # Computation of the differential spectrogram
 ######################################################################
@@ -182,7 +131,7 @@ def interpolate_2d(target_x, target_y, source_x, source_y, source_data):
     source_values = source_data.flatten()
 
     # Interpolate and reshape the interpolated data to match the target grid shape and replace NaN with zeros
-    interpolated_data = scipy.interpolate.griddata(source_points, source_values, target_points, method='nearest')
+    interpolated_data = griddata(source_points, source_values, target_points, method='nearest')
     interpolated_data = interpolated_data.reshape((len(target_y), len(target_x)))
     interpolated_data = np.nan_to_num(interpolated_data)
 
@@ -425,13 +374,17 @@ def sigmoid_scale(x, k=1):
 
 # Original Klipper function to get the PSD data of a raw accelerometer signal
 def compute_signal_data(data, max_freq):
-    calibration_data = calc_freq_response(data)
+    helper = shaper_calibrate.ShaperCalibrate(printer=None)
+    calibration_data = helper.process_accelerometer_data(data)
+
     freqs = calibration_data.freq_bins[calibration_data.freq_bins <= max_freq]
     psd = calibration_data.get_psd('all')[calibration_data.freq_bins <= max_freq]
-    peaks, _ = detect_peaks(psd, freqs)
+
+    _, peaks, _ = detect_peaks(psd, freqs, PEAKS_DETECTION_THRESHOLD * psd.max())
+
     return SignalData(freqs=freqs, psd=psd, peaks=peaks, paired_peaks=None, unpaired_peaks=None)
  
-
+ 
 ######################################################################
 # Startup and main routines
 ######################################################################

K-ShakeTune/scripts/graph_shaper.py

@@ -15,16 +15,16 @@
 #####################################################################
 
 import optparse, matplotlib, sys, importlib, os, math
-from textwrap import wrap
 import numpy as np
-import scipy
 import matplotlib.pyplot, matplotlib.dates, matplotlib.font_manager
 import matplotlib.ticker, matplotlib.gridspec
-from locale_utils import set_locale, print_with_c_locale
 from datetime import datetime
 
 matplotlib.use('Agg')
 
+from locale_utils import set_locale, print_with_c_locale
+from common_func import compute_spectrogram, detect_peaks
+
 
 PEAKS_DETECTION_THRESHOLD = 0.05
 PEAKS_EFFECT_THRESHOLD = 0.12
@@ -82,59 +82,6 @@ def compute_damping_ratio(psd, freqs):
     return fr, zeta
 
 
-def compute_spectrogram(data):
-    N = data.shape[0]
-    Fs = N / (data[-1, 0] - data[0, 0])
-    # Round up to a power of 2 for faster FFT
-    M = 1 << int(.5 * Fs - 1).bit_length()
-    window = np.kaiser(M, 6.)
-
-    def _specgram(x):
-        x_detrended = x - np.mean(x)  # Detrending by subtracting the mean value
-        return scipy.signal.spectrogram(
-            x_detrended, fs=Fs, window=window, nperseg=M, noverlap=M//2,
-            detrend='constant', scaling='density', mode='psd')
-
-    d = {'x': data[:, 1], 'y': data[:, 2], 'z': data[:, 3]}
-    f, t, pdata = _specgram(d['x'])
-    for axis in 'yz':
-        pdata += _specgram(d[axis])[2]
-    return pdata, t, f
-
-
-# This find all the peaks in a curve by looking at when the derivative term goes from positive to negative
-# Then only the peaks found above a threshold are kept to avoid capturing peaks in the low amplitude noise of a signal
-# An added "virtual" threshold allow me to quantify in an opiniated way the peaks that "could have" effect on the printer
-# behavior and are likely known to produce or contribute to the ringing/ghosting in printed parts
-def detect_peaks(psd, freqs, window_size=5, vicinity=3):
-    # Smooth the curve using a moving average to avoid catching peaks everywhere in noisy signals
-    kernel = np.ones(window_size) / window_size
-    smoothed_psd = np.convolve(psd, kernel, mode='valid')
-    mean_pad = [np.mean(psd[:window_size])] * (window_size // 2)
-    smoothed_psd = np.concatenate((mean_pad, smoothed_psd))
-
-    # Find peaks on the smoothed curve
-    smoothed_peaks = np.where((smoothed_psd[:-2] < smoothed_psd[1:-1]) & (smoothed_psd[1:-1] > smoothed_psd[2:]))[0] + 1
-    detection_threshold = PEAKS_DETECTION_THRESHOLD * psd.max()
-    effect_threshold = PEAKS_EFFECT_THRESHOLD * psd.max()
-    smoothed_peaks = smoothed_peaks[smoothed_psd[smoothed_peaks] > detection_threshold]
- 
-    # Refine peak positions on the original curve
-    refined_peaks = []
-    for peak in smoothed_peaks:
-        local_max = peak + np.argmax(psd[max(0, peak-vicinity):min(len(psd), peak+vicinity+1)]) - vicinity
-        refined_peaks.append(local_max)
-
-    peak_freqs = ["{:.1f}".format(f) for f in freqs[refined_peaks]]
-
-    num_peaks = len(refined_peaks)
-    num_peaks_above_effect_threshold = np.sum(psd[refined_peaks] > effect_threshold)
-    
-    print_with_c_locale("Peaks detected on the graph: %d @ %s Hz (%d above effect threshold)" % (num_peaks, ", ".join(map(str, peak_freqs)), num_peaks_above_effect_threshold))
-
-    return np.array(refined_peaks), num_peaks, num_peaks_above_effect_threshold
-
-
 ######################################################################
 # Graphing
 ######################################################################
@@ -216,11 +163,15 @@ def plot_freq_response_with_damping(ax, calibration_data, shapers, performance_s
 
     # Draw the detected peaks and name them
     # This also draw the detection threshold and warning threshold (aka "effect zone")
-    peaks, _, _ = detect_peaks(psd, freqs)
     peaks_warning_threshold = PEAKS_DETECTION_THRESHOLD * psd.max()
     peaks_effect_threshold = PEAKS_EFFECT_THRESHOLD * psd.max()
+    num_peaks, peaks, peaks_freqs = detect_peaks(psd, freqs, peaks_warning_threshold)
     
-    ax.plot(freqs[peaks], psd[peaks], "x", color='black', markersize=8)
+    peak_freqs_formated = ["{:.1f}".format(f) for f in peaks_freqs]
+    num_peaks_above_effect_threshold = np.sum(psd[peaks] > peaks_effect_threshold)
+    print_with_c_locale("Peaks detected on the graph: %d @ %s Hz (%d above effect threshold)" % (num_peaks, ", ".join(map(str, peak_freqs_formated)), num_peaks_above_effect_threshold))
+
+    ax.plot(peaks_freqs, psd[peaks], "x", color='black', markersize=8)
     for idx, peak in enumerate(peaks):
         if psd[peak] > peaks_effect_threshold:
             fontcolor = 'red'
@@ -242,7 +193,7 @@ def plot_freq_response_with_damping(ax, calibration_data, shapers, performance_s
     ax.legend(loc='upper left', prop=fontP)
     ax2.legend(loc='upper right', prop=fontP)
 
-    return freqs[peaks]
+    return peaks_freqs
 
 
 # Plot a time-frequency spectrogram to see how the system respond over time during the

K-ShakeTune/scripts/graph_vibrations.py

@@ -17,11 +17,13 @@ from collections import OrderedDict
 import numpy as np
 import matplotlib.pyplot, matplotlib.dates, matplotlib.font_manager
 import matplotlib.ticker, matplotlib.gridspec
-from locale_utils import set_locale, print_with_c_locale
 from datetime import datetime
 
 matplotlib.use('Agg')
 
+from locale_utils import set_locale, print_with_c_locale
+from common_func import detect_peaks
+
 
 PEAKS_DETECTION_THRESHOLD = 0.05
 PEAKS_RELATIVE_HEIGHT_THRESHOLD = 0.04
@@ -127,38 +129,6 @@ def calc_vibration_profile(power_spectral_densities):
     return total_vibration
 
 
-# This find all the peaks in a curve by looking at when the derivative term goes from positive to negative
-# Then only the peaks found above a threshold are kept to avoid capturing peaks in the low amplitude noise of a signal
-# Additionaly, we validate that a peak is a real peak based of its neighbors as we can have pretty flat zones in vibration
-# graphs with a lot of false positive due to small "noise" in these flat zones
-def detect_peaks(power_total, speeds, window_size=10, vicinity=10):
-    # Smooth the curve using a moving average to avoid catching peaks everywhere in noisy signals
-    kernel = np.ones(window_size) / window_size
-    smoothed_psd = np.convolve(power_total, kernel, mode='valid')
-    mean_pad = [np.mean(power_total[:window_size])] * (window_size // 2)
-    smoothed_psd = np.concatenate((mean_pad, smoothed_psd))
-
-    # Find peaks on the smoothed curve (and excluding the last value of the serie often detected when in a flat zone)
-    smoothed_peaks = np.where((smoothed_psd[:-3] < smoothed_psd[1:-2]) & (smoothed_psd[1:-2] > smoothed_psd[2:-1]))[0] + 1
-    detection_threshold = PEAKS_DETECTION_THRESHOLD * power_total.max()
-    
-    valid_peaks = []
-    for peak in smoothed_peaks:
-        peak_height = smoothed_psd[peak] - np.min(smoothed_psd[max(0, peak-vicinity):min(len(smoothed_psd), peak+vicinity+1)])
-        if peak_height > PEAKS_RELATIVE_HEIGHT_THRESHOLD * smoothed_psd[peak] and smoothed_psd[peak] > detection_threshold:
-            valid_peaks.append(peak)
- 
-    # Refine peak positions on the original curve
-    refined_peaks = []
-    for peak in valid_peaks:
-        local_max = peak + np.argmax(power_total[max(0, peak-vicinity):min(len(power_total), peak+vicinity+1)]) - vicinity
-        refined_peaks.append(local_max)
-
-    num_peaks = len(refined_peaks)
-
-    return np.array(refined_peaks), num_peaks
-
-
 # The goal is to find zone outside of peaks (flat low energy zones) to advise them as good speeds range to use in the slicer
 def identify_low_energy_zones(power_total):
     valleys = []
@@ -233,9 +203,10 @@ def plot_speed_profile(ax, speeds, power_total):
     ax.plot(speeds, power_total[2], label="Y", color='green')
     ax.plot(speeds, power_total[3], label="Z", color='blue')
 
-    peaks, num_peaks = detect_peaks(resampled_power_total, resampled_speeds)
+    detection_threshold = PEAKS_DETECTION_THRESHOLD * resampled_power_total.max()
+    num_peaks, peaks, _ = detect_peaks(resampled_power_total, resampled_speeds, detection_threshold, PEAKS_RELATIVE_HEIGHT_THRESHOLD, 10, 10)
     low_energy_zones = identify_low_energy_zones(resampled_power_total)
-    
+
     peak_speeds = ["{:.1f}".format(resampled_speeds[i]) for i in peaks]
     print_with_c_locale("Vibrations peaks detected: %d @ %s mm/s (avoid setting a speed near these values in your slicer print profile)" % (num_peaks, ", ".join(map(str, peak_speeds))))