Lot of refactoring, memory and speed optimizations for all the graphs scripts

K-ShakeTune/scripts/common_func.py

@@ -3,9 +3,53 @@
 # Common functions for the Shake&Tune package
 # Written by Frix_x#0161 #
 
-
+import math
+import os, sys
+from importlib import import_module
+from pathlib import Path
 import numpy as np
 from scipy.signal import spectrogram
+from git import GitCommandError, Repo
+
+
+def parse_log(logname):
+    with open(logname) as f:
+        for header in f:
+            if not header.startswith('#'):
+                break
+        if not header.startswith('freq,psd_x,psd_y,psd_z,psd_xyz'):
+            # Raw accelerometer data
+            return np.loadtxt(logname, comments='#', delimiter=',')
+    # Power spectral density data or shaper calibration data
+    raise ValueError("File %s does not contain raw accelerometer data and therefore "
+               "is not supported by Shake&Tune. Please use the official Klipper "
+               "script to process it instead." % (logname,))
+
+
+def setup_klipper_import(kdir):
+    kdir = os.path.expanduser(kdir)
+    sys.path.append(os.path.join(kdir, 'klippy'))
+    return import_module('.shaper_calibrate', 'extras')
+
+
+# This is used to print the current S&T version on top of the png graph file
+def get_git_version():
+    try:
+        # Get the absolute path of the script, resolving any symlinks
+        # Then get 2 times to parent dir to be at the git root folder
+        script_path = Path(__file__).resolve()
+        repo_path = script_path.parents[2]
+        repo = Repo(repo_path)
+
+        try:
+            version = repo.git.describe('--tags')
+        except GitCommandError:
+            # If no tag is found, use the simplified commit SHA instead
+            version = repo.head.commit.hexsha[:7]
+        return version
+
+    except Exception as e:
+        return None
 
 
 # This is Klipper's spectrogram generation function adapted to use Scipy
@@ -17,10 +61,9 @@ def compute_spectrogram(data):
     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')
+        x -= np.mean(x) # Detrending by subtracting the mean value
+        return spectrogram(x, 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'])
@@ -29,6 +72,23 @@ def compute_spectrogram(data):
     return pdata, t, f
 
 
+# Compute natural resonant frequency and damping ratio by using the half power bandwidth method with interpolated frequencies
+def compute_mechanical_parameters(psd, freqs):
+    max_power_index = np.argmax(psd)
+    fr = freqs[max_power_index]
+    max_power = psd[max_power_index]
+
+    half_power = max_power / math.sqrt(2)
+    idx_below = np.where(psd[:max_power_index] <= half_power)[0][-1]
+    idx_above = np.where(psd[max_power_index:] <= half_power)[0][0] + max_power_index
+    freq_below_half_power = freqs[idx_below] + (half_power - psd[idx_below]) * (freqs[idx_below + 1] - freqs[idx_below]) / (psd[idx_below + 1] - psd[idx_below])
+    freq_above_half_power = freqs[idx_above - 1] + (half_power - psd[idx_above - 1]) * (freqs[idx_above] - freqs[idx_above - 1]) / (psd[idx_above] - psd[idx_above - 1])
+
+    bandwidth = freq_above_half_power - freq_below_half_power
+    zeta = bandwidth / (2 * fr)
+
+    return fr, zeta, max_power_index
+
 # 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):

K-ShakeTune/scripts/graph_belts.py

@@ -12,18 +12,17 @@
 #####################################################################
 
 import optparse, matplotlib, sys, importlib, os
+from datetime import datetime
 from collections import namedtuple
 import numpy as np
+import matplotlib.pyplot as plt 
+import matplotlib.font_manager, matplotlib.ticker, matplotlib.colors
 from scipy.interpolate import griddata
-import matplotlib.pyplot, matplotlib.dates, matplotlib.font_manager
-import matplotlib.ticker, matplotlib.gridspec, matplotlib.colors
-import matplotlib.patches
-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
+from common_func import compute_spectrogram, detect_peaks, get_git_version, parse_log, setup_klipper_import
 
 
 ALPHABET = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" # For paired peaks names
@@ -141,14 +140,14 @@ def interpolate_2d(target_x, target_y, source_x, source_y, source_data):
 # Main logic function to combine two similar spectrogram - ie. from both belts paths - by substracting signals in order to create
 # a new composite spectrogram. This result of a divergent but mostly centered new spectrogram (center will be white) with some colored zones
 # highlighting differences in the belts paths. The summative spectrogram is used for the MHI calculation.
-def combined_spectrogram(data1, data2):
+def compute_combined_spectrogram(data1, data2):
     pdata1, bins1, t1 = compute_spectrogram(data1)
     pdata2, bins2, t2 = compute_spectrogram(data2)
 
     # Interpolate the spectrograms
     pdata2_interpolated = interpolate_2d(bins1, t1, bins2, t2, pdata2)
     
-    # Cobine them in two form: a summed diff for the MHI computation and a diverging diff for the spectrogram colors
+    # Combine them in two form: a summed diff for the MHI computation and a diverging diff for the spectrogram colors
     combined_sum = np.abs(pdata1 - pdata2_interpolated)
     combined_divergent = pdata1 - pdata2_interpolated
 
@@ -184,26 +183,27 @@ def compute_mhi(combined_data, similarity_coefficient, num_unpaired_peaks):
 
 
 # LUT to transform the MHI into a textual value easy to understand for the users of the script
-def mhi_lut(mhi):
-    if 0 <= mhi <= 30:
-        return "Excellent mechanical health"
-    elif 30 < mhi <= 45:
-        return "Good mechanical health"
-    elif 45 < mhi <= 55:
-        return "Acceptable mechanical health"
-    elif 55 < mhi <= 70:
-        return "Potential signs of a mechanical issue"
-    elif 70 < mhi <= 85:
-        return "Likely a mechanical issue"
-    elif 85 < mhi <= 100:
-        return "Mechanical issue detected"
+def mhi_lut(mhi):    
+    ranges = [
+        (0, 30, "Excellent mechanical health"),
+        (30, 45, "Good mechanical health"),
+        (45, 55, "Acceptable mechanical health"),
+        (55, 70, "Potential signs of a mechanical issue"),
+        (70, 85, "Likely a mechanical issue"),
+        (85, 100, "Mechanical issue detected")
+    ]
+    for lower, upper, message in ranges:
+        if lower < mhi <= upper:
+            return message
+
+    return "Error computing MHI value"
 
 
 ######################################################################
 # Graphing
 ######################################################################
 
-def plot_compare_frequency(ax, lognames, signal1, signal2, max_freq):
+def plot_compare_frequency(ax, lognames, signal1, signal2, similarity_factor, max_freq):
     # Get the belt name for the legend to avoid putting the full file name
     signal1_belt = (lognames[0].split('/')[-1]).split('_')[-1][0]
     signal2_belt = (lognames[1].split('/')[-1]).split('_')[-1][0]
@@ -264,13 +264,11 @@ def plot_compare_frequency(ax, lognames, signal1, signal2, max_freq):
                     ha='left', fontsize=13, color='red', weight='bold')
         unpaired_peak_count += 1
 
-    # Compute the similarity (using cross-correlation of the PSD signals)
+    # Add estimated similarity to the graph
     ax2 = ax.twinx() # To split the legends in two box
     ax2.yaxis.set_visible(False)
-    similarity_factor = compute_curve_similarity_factor(signal1, signal2)
     ax2.plot([], [], ' ', label=f'Estimated similarity: {similarity_factor:.1f}%')
     ax2.plot([], [], ' ', label=f'Number of unpaired peaks: {unpaired_peak_count}')
-    print_with_c_locale(f"Belts estimated similarity: {similarity_factor:.1f}%")
 
     # Setting axis parameters, grid and graph title
     ax.set_xlabel('Frequency (Hz)')
@@ -304,25 +302,20 @@ def plot_compare_frequency(ax, lognames, signal1, signal2, max_freq):
     ax.legend(loc='upper left', prop=fontP)
     ax2.legend(loc='upper right', prop=fontP)
 
-    return similarity_factor, unpaired_peak_count
+    return
 
 
-def plot_difference_spectrogram(ax, data1, data2, signal1, signal2, similarity_factor, max_freq):
-    combined_sum, combined_divergent, bins, t = combined_spectrogram(data1, data2)
-
-    # Compute the MHI value from the differential spectrogram sum of gradient, salted with
-    # the similarity factor and the number or unpaired peaks from the belts frequency profile
-    # Be careful, this value is highly opinionated and is pretty experimental!
-    mhi, textual_mhi = compute_mhi(combined_sum, similarity_factor, len(signal1.unpaired_peaks) + len(signal2.unpaired_peaks))
-    print_with_c_locale(f"[experimental] Mechanical Health Indicator: {textual_mhi.lower()} ({mhi:.1f}%)")
+def plot_difference_spectrogram(ax, signal1, signal2, t, bins, combined_divergent, textual_mhi, max_freq):
     ax.set_title(f"Differential Spectrogram", fontsize=14, color=KLIPPAIN_COLORS['dark_orange'], weight='bold')
     ax.plot([], [], ' ', label=f'{textual_mhi} (experimental)')
     
     # Draw the differential spectrogram with a specific custom norm to get orange or purple values where there is signal or white near zeros
+    # imgshow is better suited here than pcolormesh since its result is already rasterized and we doesn't need to keep vector graphics
+    # when saving to a final .png file. Using it also allow to save ~150-200MB of RAM during the "fig.savefig" operation.
     colors = [KLIPPAIN_COLORS['dark_orange'], KLIPPAIN_COLORS['orange'], 'white', KLIPPAIN_COLORS['purple'], KLIPPAIN_COLORS['dark_purple']]  
     cm = matplotlib.colors.LinearSegmentedColormap.from_list('klippain_divergent', list(zip([0, 0.25, 0.5, 0.75, 1], colors)))
     norm = matplotlib.colors.TwoSlopeNorm(vmin=np.min(combined_divergent), vcenter=0, vmax=np.max(combined_divergent))
-    ax.pcolormesh(t, bins, combined_divergent.T, cmap=cm, norm=norm, shading='gouraud')
+    ax.imshow(combined_divergent.T, cmap=cm, norm=norm, aspect='auto', extent=[t[0], t[-1], bins[0], bins[-1]], interpolation='bilinear', origin='lower')
 
     ax.set_xlabel('Frequency (hz)')
     ax.set_xlim([0., max_freq])
@@ -389,50 +382,47 @@ def compute_signal_data(data, max_freq):
 # Startup and main routines
 ######################################################################
 
-def parse_log(logname):
-    with open(logname) as f:
-        for header in f:
-            if not header.startswith('#'):
-                break
-        if not header.startswith('freq,psd_x,psd_y,psd_z,psd_xyz'):
-            # Raw accelerometer data
-            return np.loadtxt(logname, comments='#', delimiter=',')
-    # Power spectral density data or shaper calibration data
-    raise ValueError("File %s does not contain raw accelerometer data and therefore "
-               "is not supported by this script. Please use the official Klipper "
-               "graph_accelerometer.py script to process it instead." % (logname,))
-
-
-def setup_klipper_import(kdir):
-    global shaper_calibrate
-    kdir = os.path.expanduser(kdir)
-    sys.path.append(os.path.join(kdir, 'klippy'))
-    shaper_calibrate = importlib.import_module('.shaper_calibrate', 'extras')
-
-
 def belts_calibration(lognames, klipperdir="~/klipper", max_freq=200.):
     set_locale()
-    setup_klipper_import(klipperdir)
+    global shaper_calibrate
+    shaper_calibrate = setup_klipper_import(klipperdir)
 
     # Parse data
     datas = [parse_log(fn) for fn in lognames]
     if len(datas) > 2:
-        raise ValueError("Incorrect number of .csv files used (this function needs two files to compare them)")
+        raise ValueError("Incorrect number of .csv files used (this function needs exactly two files to compare them)!")
 
     # Compute calibration data for the two datasets with automatic peaks detection
     signal1 = compute_signal_data(datas[0], max_freq)
     signal2 = compute_signal_data(datas[1], max_freq)
+    combined_sum, combined_divergent, bins, t = compute_combined_spectrogram(datas[0], datas[1])
+    del datas
 
     # Pair the peaks across the two datasets
     paired_peaks, unpaired_peaks1, unpaired_peaks2 = pair_peaks(signal1.peaks, signal1.freqs, signal1.psd,
                                                                signal2.peaks, signal2.freqs, signal2.psd)
-    signal1 = signal1._replace(paired_peaks=paired_peaks, unpaired_peaks=unpaired_peaks1)
-    signal2 = signal2._replace(paired_peaks=paired_peaks, unpaired_peaks=unpaired_peaks2)
+    signal1 = signal1._replace(paired_peaks = paired_peaks, unpaired_peaks = unpaired_peaks1)
+    signal2 = signal2._replace(paired_peaks = paired_peaks, unpaired_peaks = unpaired_peaks2)
 
-    fig = matplotlib.pyplot.figure()
-    gs = matplotlib.gridspec.GridSpec(2, 1, height_ratios=[4, 3])
-    ax1 = fig.add_subplot(gs[0])
-    ax2 = fig.add_subplot(gs[1])
+    # Compute the similarity (using cross-correlation of the PSD signals)
+    similarity_factor = compute_curve_similarity_factor(signal1, signal2)
+    print_with_c_locale(f"Belts estimated similarity: {similarity_factor:.1f}%")
+    # Compute the MHI value from the differential spectrogram sum of gradient, salted with the similarity factor and the number of
+    # unpaired peaks from the belts frequency profile. Be careful, this value is highly opinionated and is pretty experimental!
+    mhi, textual_mhi = compute_mhi(combined_sum, similarity_factor, len(signal1.unpaired_peaks) + len(signal2.unpaired_peaks))
+    print_with_c_locale(f"[experimental] Mechanical Health Indicator: {textual_mhi.lower()} ({mhi:.1f}%)")
+
+    # Create graph layout
+    fig, (ax1, ax2) = plt.subplots(2, 1, gridspec_kw={
+            'height_ratios':[4, 3],
+            'bottom':0.050,
+            'top':0.890,
+            'left':0.085,
+            'right':0.966,
+            'hspace':0.169,
+            'wspace':0.200
+            })
+    fig.set_size_inches(8.3, 11.6)
 
     # Add title
     title_line1 = "RELATIVE BELT CALIBRATION TOOL"
@@ -447,18 +437,19 @@ def belts_calibration(lognames, klipperdir="~/klipper", max_freq=200.):
     fig.text(0.12, 0.957, title_line2, ha='left', va='top', fontsize=16, color=KLIPPAIN_COLORS['dark_purple'])
 
     # Plot the graphs
-    similarity_factor, _ = plot_compare_frequency(ax1, lognames, signal1, signal2, max_freq)
-    plot_difference_spectrogram(ax2, datas[0], datas[1], signal1, signal2, similarity_factor, max_freq)
-
-    fig.set_size_inches(8.3, 11.6)
-    fig.tight_layout()
-    fig.subplots_adjust(top=0.89)
+    plot_compare_frequency(ax1, lognames, signal1, signal2, similarity_factor, max_freq)
+    plot_difference_spectrogram(ax2, signal1, signal2, t, bins, combined_divergent, textual_mhi, max_freq)
     
     # Adding a small Klippain logo to the top left corner of the figure
-    ax_logo = fig.add_axes([0.001, 0.899, 0.1, 0.1], anchor='NW', zorder=-1)
-    ax_logo.imshow(matplotlib.pyplot.imread(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'klippain.png')))
+    ax_logo = fig.add_axes([0.001, 0.8995, 0.1, 0.1], anchor='NW')
+    ax_logo.imshow(plt.imread(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'klippain.png')))
     ax_logo.axis('off')
-    
+
+    # Adding Shake&Tune version in the top right corner
+    st_version = get_git_version()
+    if st_version is not None:
+        fig.text(0.995, 0.985, st_version, ha='right', va='bottom', fontsize=8, color=KLIPPAIN_COLORS['purple'])
+
     return fig
 
 
@@ -479,7 +470,7 @@ def main():
         opts.error("You must specify an output file.png to use the script (option -o)")
 
     fig = belts_calibration(args, options.klipperdir, options.max_freq)
-    fig.savefig(options.output)
+    fig.savefig(options.output, dpi=150)
 
 
 if __name__ == '__main__':

K-ShakeTune/scripts/graph_shaper.py

@@ -14,16 +14,16 @@
 ################ !!! DO NOT EDIT BELOW THIS LINE !!! ################
 #####################################################################
 
-import optparse, matplotlib, sys, importlib, os, math
-import numpy as np
-import matplotlib.pyplot, matplotlib.dates, matplotlib.font_manager
-import matplotlib.ticker, matplotlib.gridspec
+import optparse, matplotlib, os
 from datetime import datetime
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.font_manager, matplotlib.ticker
 
 matplotlib.use('Agg')
 
 from locale_utils import set_locale, print_with_c_locale
-from common_func import compute_spectrogram, detect_peaks
+from common_func import compute_mechanical_parameters, compute_spectrogram, detect_peaks, get_git_version, parse_log, setup_klipper_import
 
 
 PEAKS_DETECTION_THRESHOLD = 0.05
@@ -43,20 +43,13 @@ KLIPPAIN_COLORS = {
 ######################################################################
 
 # Find the best shaper parameters using Klipper's official algorithm selection
-def calibrate_shaper_with_damping(datas, max_smoothing):
+def calibrate_shaper(datas, max_smoothing):
     helper = shaper_calibrate.ShaperCalibrate(printer=None)
-
-    calibration_data = helper.process_accelerometer_data(datas[0])
-    for data in datas[1:]:
-        calibration_data.add_data(helper.process_accelerometer_data(data))
-
+    calibration_data = helper.process_accelerometer_data(datas)
     calibration_data.normalize_to_frequencies()
 
     shaper, all_shapers = helper.find_best_shaper(calibration_data, max_smoothing, print_with_c_locale)
-
-    freqs = calibration_data.freq_bins
-    psd = calibration_data.psd_sum
-    fr, zeta = compute_damping_ratio(psd, freqs)
+    fr, zeta, _ = compute_mechanical_parameters(calibration_data.psd_sum, calibration_data.freq_bins)
 
     print_with_c_locale("Recommended shaper is %s @ %.1f Hz" % (shaper.name, shaper.freq))
     print_with_c_locale("Axis has a main resonant frequency at %.1fHz with an estimated damping ratio of %.3f" % (fr, zeta))
@@ -64,36 +57,17 @@ def calibrate_shaper_with_damping(datas, max_smoothing):
     return shaper.name, all_shapers, calibration_data, fr, zeta
 
 
-# Compute damping ratio by using the half power bandwidth method with interpolated frequencies
-def compute_damping_ratio(psd, freqs):
-    max_power_index = np.argmax(psd)
-    fr = freqs[max_power_index]
-    max_power = psd[max_power_index]
-
-    half_power = max_power / math.sqrt(2)
-    idx_below = np.where(psd[:max_power_index] <= half_power)[0][-1]
-    idx_above = np.where(psd[max_power_index:] <= half_power)[0][0] + max_power_index
-    freq_below_half_power = freqs[idx_below] + (half_power - psd[idx_below]) * (freqs[idx_below + 1] - freqs[idx_below]) / (psd[idx_below + 1] - psd[idx_below])
-    freq_above_half_power = freqs[idx_above - 1] + (half_power - psd[idx_above - 1]) * (freqs[idx_above] - freqs[idx_above - 1]) / (psd[idx_above] - psd[idx_above - 1])
-
-    bandwidth = freq_above_half_power - freq_below_half_power
-    zeta = bandwidth / (2 * fr)
-
-    return fr, zeta
-
-
 ######################################################################
 # Graphing
 ######################################################################
 
-def plot_freq_response_with_damping(ax, calibration_data, shapers, performance_shaper, fr, zeta, max_freq):
-    freqs = calibration_data.freq_bins
-    psd = calibration_data.psd_sum[freqs <= max_freq]
-    px = calibration_data.psd_x[freqs <= max_freq]
-    py = calibration_data.psd_y[freqs <= max_freq]
-    pz = calibration_data.psd_z[freqs <= max_freq]
-    freqs = freqs[freqs <= max_freq]
-
+def plot_freq_response(ax, calibration_data, shapers, performance_shaper, peaks, peaks_freqs, peaks_threshold, fr, zeta, max_freq):
+    freqs = calibration_data.freqs
+    psd = calibration_data.psd_sum
+    px = calibration_data.psd_x
+    py = calibration_data.psd_y
+    pz = calibration_data.psd_z
+    
     fontP = matplotlib.font_manager.FontProperties()
     fontP.set_size('x-small')
 
@@ -102,7 +76,7 @@ def plot_freq_response_with_damping(ax, calibration_data, shapers, performance_s
     ax.set_ylabel('Power spectral density')
     ax.set_ylim([0, psd.max() + psd.max() * 0.05])
 
-    ax.plot(freqs, psd, label='X+Y+Z', color='purple')
+    ax.plot(freqs, psd, label='X+Y+Z', color='purple', zorder=5)
     ax.plot(freqs, px, label='X', color='red')
     ax.plot(freqs, py, label='Y', color='green')
     ax.plot(freqs, pz, label='Z', color='blue')
@@ -163,17 +137,9 @@ 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_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)
-    
-    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:
+        if psd[peak] > peaks_threshold[1]:
             fontcolor = 'red'
             fontweight = 'bold'
         else:
@@ -182,46 +148,48 @@ def plot_freq_response_with_damping(ax, calibration_data, shapers, performance_s
         ax.annotate(f"{idx+1}", (freqs[peak], psd[peak]), 
                     textcoords="offset points", xytext=(8, 5), 
                     ha='left', fontsize=13, color=fontcolor, weight=fontweight)
-    ax.axhline(y=peaks_warning_threshold, color='black', linestyle='--', linewidth=0.5)
-    ax.axhline(y=peaks_effect_threshold, color='black', linestyle='--', linewidth=0.5)
-    ax.fill_between(freqs, 0, peaks_warning_threshold, color='green', alpha=0.15, label='Relax Region')
-    ax.fill_between(freqs, peaks_warning_threshold, peaks_effect_threshold, color='orange', alpha=0.2, label='Warning Region')
-
+    ax.axhline(y=peaks_threshold[0], color='black', linestyle='--', linewidth=0.5)
+    ax.axhline(y=peaks_threshold[1], color='black', linestyle='--', linewidth=0.5)
+    ax.fill_between(freqs, 0, peaks_threshold[0], color='green', alpha=0.15, label='Relax Region')
+    ax.fill_between(freqs, peaks_threshold[0], peaks_threshold[1], color='orange', alpha=0.2, label='Warning Region')
 
     # Add the main resonant frequency and damping ratio of the axis to the graph title
     ax.set_title("Axis Frequency Profile (ω0=%.1fHz, ζ=%.3f)" % (fr, zeta), fontsize=14, color=KLIPPAIN_COLORS['dark_orange'], weight='bold')
     ax.legend(loc='upper left', prop=fontP)
     ax2.legend(loc='upper right', prop=fontP)
 
-    return peaks_freqs
+    return
 
 
 # Plot a time-frequency spectrogram to see how the system respond over time during the
 # resonnance test. This can highlight hidden spots from the standard PSD graph from other harmonics
-def plot_spectrogram(ax, data, peaks, max_freq):
-    pdata, bins, t = compute_spectrogram(data)
-
+def plot_spectrogram(ax, t, bins, pdata, peaks, max_freq):
+    ax.set_title("Time-Frequency Spectrogram", fontsize=14, color=KLIPPAIN_COLORS['dark_orange'], weight='bold')
+    
     # We need to normalize the data to get a proper signal on the spectrogram
     # However, while using "LogNorm" provide too much background noise, using
     # "Normalize" make only the resonnance appearing and hide interesting elements
     # So we need to filter out the lower part of the data (ie. find the proper vmin for LogNorm)
     vmin_value = np.percentile(pdata, SPECTROGRAM_LOW_PERCENTILE_FILTER)
 
-    ax.set_title("Time-Frequency Spectrogram", fontsize=14, color=KLIPPAIN_COLORS['dark_orange'], weight='bold')
-    ax.pcolormesh(t, bins, pdata.T, norm=matplotlib.colors.LogNorm(vmin=vmin_value),
-                  cmap='inferno', shading='gouraud')
-
-    # Add peaks lines in the spectrogram to get hint from peaks found in the first graph
-    if peaks is not None:
-        for idx, peak in enumerate(peaks):
-            ax.axvline(peak, color='cyan', linestyle='dotted', linewidth=0.75)
-            ax.annotate(f"Peak {idx+1}", (peak, t[-1]*0.9), 
-                        textcoords="data", color='cyan', rotation=90, fontsize=10,
-                        verticalalignment='top', horizontalalignment='right')
+    # Draw the spectrogram using imgshow that is better suited here than pcolormesh since its result is already rasterized and
+    # we doesn't need to keep vector graphics when saving to a final .png file. Using it also allow to
+    # save ~150-200MB of RAM during the "fig.savefig" operation.
+    cm = 'inferno'
+    norm = matplotlib.colors.LogNorm(vmin=vmin_value)
+    ax.imshow(pdata.T, norm=norm, cmap=cm, aspect='auto', extent=[t[0], t[-1], bins[0], bins[-1]], origin='lower', interpolation='antialiased')
 
     ax.set_xlim([0., max_freq])
     ax.set_ylabel('Time (s)')
     ax.set_xlabel('Frequency (Hz)')
+    
+    # Add peaks lines in the spectrogram to get hint from peaks found in the first graph
+    if peaks is not None:
+        for idx, peak in enumerate(peaks):
+            ax.axvline(peak, color='cyan', linestyle='dotted', linewidth=1)
+            ax.annotate(f"Peak {idx+1}", (peak, bins[-1]*0.9), 
+                        textcoords="data", color='cyan', rotation=90, fontsize=10,
+                        verticalalignment='top', horizontalalignment='right')
 
     return
 
@@ -230,41 +198,52 @@ def plot_spectrogram(ax, data, peaks, max_freq):
 # Startup and main routines
 ######################################################################
 
-def parse_log(logname):
-    with open(logname) as f:
-        for header in f:
-            if not header.startswith('#'):
-                break
-        if not header.startswith('freq,psd_x,psd_y,psd_z,psd_xyz'):
-            # Raw accelerometer data
-            return np.loadtxt(logname, comments='#', delimiter=',')
-    # Power spectral density data or shaper calibration data
-    raise ValueError("File %s does not contain raw accelerometer data and therefore "
-               "is not supported by this script. Please use the official Klipper "
-               "calibrate_shaper.py script to process it instead." % (logname,))
-
-
-def setup_klipper_import(kdir):
-    global shaper_calibrate
-    kdir = os.path.expanduser(kdir)
-    sys.path.append(os.path.join(kdir, 'klippy'))
-    shaper_calibrate = importlib.import_module('.shaper_calibrate', 'extras')
-
-
 def shaper_calibration(lognames, klipperdir="~/klipper", max_smoothing=None, max_freq=200.):
     set_locale()
-    setup_klipper_import(klipperdir)
+    global shaper_calibrate
+    shaper_calibrate = setup_klipper_import(klipperdir)
 
     # Parse data
     datas = [parse_log(fn) for fn in lognames]
+    if len(datas) > 1:
+        print_with_c_locale("Warning: incorrect number of .csv files detected. Only the first one will be used!")
 
-    # Calibrate shaper and generate outputs
-    performance_shaper, shapers, calibration_data, fr, zeta = calibrate_shaper_with_damping(datas, max_smoothing)
+    # Compute shapers, PSD outputs and spectrogram
+    performance_shaper, shapers, calibration_data, fr, zeta = calibrate_shaper(datas[0], max_smoothing)
+    pdata, bins, t = compute_spectrogram(datas[0])
+    del datas
 
-    fig = matplotlib.pyplot.figure()
-    gs = matplotlib.gridspec.GridSpec(2, 1, height_ratios=[4, 3])
-    ax1 = fig.add_subplot(gs[0])
-    ax2 = fig.add_subplot(gs[1])
+    # Select only the relevant part of the PSD data
+    freqs = calibration_data.freq_bins
+    calibration_data.psd_sum = calibration_data.psd_sum[freqs <= max_freq]
+    calibration_data.psd_x = calibration_data.psd_x[freqs <= max_freq]
+    calibration_data.psd_y = calibration_data.psd_y[freqs <= max_freq]
+    calibration_data.psd_z = calibration_data.psd_z[freqs <= max_freq]
+    calibration_data.freqs = freqs[freqs <= max_freq]
+
+    # Peak detection algorithm
+    peaks_threshold = [
+        PEAKS_DETECTION_THRESHOLD * calibration_data.psd_sum.max(),
+        PEAKS_EFFECT_THRESHOLD * calibration_data.psd_sum.max()
+    ]
+    num_peaks, peaks, peaks_freqs = detect_peaks(calibration_data.psd_sum, calibration_data.freqs, peaks_threshold[0])
+    
+    # Print the peaks info in the console
+    peak_freqs_formated = ["{:.1f}".format(f) for f in peaks_freqs]
+    num_peaks_above_effect_threshold = np.sum(calibration_data.psd_sum[peaks] > peaks_threshold[1])
+    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))
+
+    # Create graph layout
+    fig, (ax1, ax2) = plt.subplots(2, 1, gridspec_kw={
+            'height_ratios':[4, 3],
+            'bottom':0.050,
+            'top':0.890,
+            'left':0.085,
+            'right':0.966,
+            'hspace':0.169,
+            'wspace':0.200
+            })
+    fig.set_size_inches(8.3, 11.6)
     
     # Add title
     title_line1 = "INPUT SHAPER CALIBRATION TOOL"
@@ -279,18 +258,19 @@ def shaper_calibration(lognames, klipperdir="~/klipper", max_smoothing=None, max
     fig.text(0.12, 0.957, title_line2, ha='left', va='top', fontsize=16, color=KLIPPAIN_COLORS['dark_purple'])
 
     # Plot the graphs
-    peaks = plot_freq_response_with_damping(ax1, calibration_data, shapers, performance_shaper, fr, zeta, max_freq)
-    plot_spectrogram(ax2, datas[0], peaks, max_freq)
-
-    fig.set_size_inches(8.3, 11.6)
-    fig.tight_layout()
-    fig.subplots_adjust(top=0.89)
+    plot_freq_response(ax1, calibration_data, shapers, performance_shaper, peaks, peaks_freqs, peaks_threshold, fr, zeta, max_freq)
+    plot_spectrogram(ax2, t, bins, pdata, peaks_freqs, max_freq)
 
     # Adding a small Klippain logo to the top left corner of the figure
-    ax_logo = fig.add_axes([0.001, 0.899, 0.1, 0.1], anchor='NW', zorder=-1)
-    ax_logo.imshow(matplotlib.pyplot.imread(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'klippain.png')))
+    ax_logo = fig.add_axes([0.001, 0.8995, 0.1, 0.1], anchor='NW')
+    ax_logo.imshow(plt.imread(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'klippain.png')))
     ax_logo.axis('off')
 
+    # Adding Shake&Tune version in the top right corner
+    st_version = get_git_version()
+    if st_version is not None:
+        fig.text(0.995, 0.985, st_version, ha='right', va='bottom', fontsize=8, color=KLIPPAIN_COLORS['purple'])
+
     return fig
 
 
@@ -315,7 +295,7 @@ def main():
         opts.error("Too small max_smoothing specified (must be at least 0.05)")
 
     fig = shaper_calibration(args, options.klipperdir, options.max_smoothing, options.max_freq)
-    fig.savefig(options.output)
+    fig.savefig(options.output, dpi=150)
 
 
 if __name__ == '__main__':

K-ShakeTune/scripts/graph_vibrations.py

@@ -11,18 +11,17 @@
 ################ !!! DO NOT EDIT BELOW THIS LINE !!! ################
 #####################################################################
 
-import math
-import optparse, matplotlib, re, sys, importlib, os, operator
+import optparse, matplotlib, re, os, operator
+from datetime import datetime
 from collections import OrderedDict
 import numpy as np
-import matplotlib.pyplot, matplotlib.dates, matplotlib.font_manager
-import matplotlib.ticker, matplotlib.gridspec
-from datetime import datetime
+import matplotlib.pyplot as plt
+import matplotlib.font_manager, matplotlib.ticker, matplotlib.gridspec
 
 matplotlib.use('Agg')
 
 from locale_utils import set_locale, print_with_c_locale
-from common_func import detect_peaks
+from common_func import compute_mechanical_parameters, detect_peaks, get_git_version, parse_log, setup_klipper_import
 
 
 PEAKS_DETECTION_THRESHOLD = 0.05
@@ -42,13 +41,13 @@ KLIPPAIN_COLORS = {
 # Computation
 ######################################################################
 
+# Call to the official Klipper input shaper object to do the PSD computation
 def calc_freq_response(data):
-    # Use Klipper standard input shaper objects to do the computation
     helper = shaper_calibrate.ShaperCalibrate(printer=None)
     return helper.process_accelerometer_data(data)
 
 
-def calc_psd(datas, group, max_freq):
+def compute_vibration_spectrogram(datas, group, max_freq):
     psd_list = []
     first_freqs = None
     signal_axes = ['x', 'y', 'z', 'all']
@@ -104,10 +103,10 @@ def calc_psd(datas, group, max_freq):
             pz = signal_normalized['z'][first_freqs <= max_freq]
             psd_list.append([psd, px, py, pz])
 
-    return first_freqs[first_freqs <= max_freq], psd_list
+    return np.array(first_freqs[first_freqs <= max_freq]), np.array(psd_list)
 
 
-def calc_speed_profile(psd_list, freqs):
+def compute_speed_profile(speeds, freqs, psd_list):
     # Preallocate arrays as psd_list is known and consistent
     pwrtot_sum = np.zeros(len(psd_list))
     pwrtot_x = np.zeros(len(psd_list))
@@ -119,14 +118,27 @@ def calc_speed_profile(psd_list, freqs):
         pwrtot_x[i] = np.trapz(psd[1], freqs)
         pwrtot_y[i] = np.trapz(psd[2], freqs)
         pwrtot_z[i] = np.trapz(psd[3], freqs)
+    
+    # Resample the signals to get a better detection of the valleys of low energy
+    # and avoid getting limited by the speed increment defined by the user
+    resampled_speeds, resampled_power_sum = resample_signal(speeds, pwrtot_sum)
+    _, resampled_pwrtot_x = resample_signal(speeds, pwrtot_x)
+    _, resampled_pwrtot_y = resample_signal(speeds, pwrtot_y)
+    _, resampled_pwrtot_z = resample_signal(speeds, pwrtot_z)
 
-    return [pwrtot_sum, pwrtot_x, pwrtot_y, pwrtot_z]
+    return resampled_speeds, [resampled_power_sum, resampled_pwrtot_x, resampled_pwrtot_y, resampled_pwrtot_z]
 
 
-def calc_vibration_profile(power_spectral_densities):
-    # Sum the PSD across all speeds for each frequency
-    total_vibration = np.sum([psd[0] for psd in power_spectral_densities], axis=0)
-    return total_vibration
+def compute_motor_profile(power_spectral_densities):
+    # Sum the PSD across all speeds for each frequency of the spectrogram. Basically this
+    # is equivalent to sum up all the spectrogram column by column to plot the total on the right
+    motor_total_vibration = np.sum([psd[0] for psd in power_spectral_densities], axis=0)
+    
+    # Then a very little smoothing of the signal is applied to avoid too much noise and sharp peaks on it and simplify
+    # the resonance frequency and damping ratio estimation later on. Also, too much smoothing is bad and would alter the results
+    smoothed_motor_total_vibration = np.convolve(motor_total_vibration, np.ones(10)/10, mode='same')
+    
+    return smoothed_motor_total_vibration
 
 
 # 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
@@ -174,16 +186,16 @@ def identify_low_energy_zones(power_total):
 def resample_signal(speeds, power_total, new_spacing=0.1):
     new_speeds = np.arange(speeds[0], speeds[-1] + new_spacing, new_spacing)
     new_power_total = np.interp(new_speeds, speeds, power_total)
-    return new_speeds, new_power_total
+    return np.array(new_speeds), np.array(new_power_total)
 
 
 ######################################################################
 # Graphing
 ######################################################################
 
-def plot_speed_profile(ax, speeds, power_total):
-    resampled_speeds, resampled_power_total = resample_signal(speeds, power_total[0])
-
+def plot_speed_profile(ax, speeds, power_total, num_peaks, peaks, low_energy_zones):
+    # For this function, we have two array for the speeds. Indeed, since the power total sum was resampled to better detect
+    # the valleys of low energy later on, we also need the resampled speed array to plot it. For the rest
     ax.set_title("Machine speed profile", fontsize=14, color=KLIPPAIN_COLORS['dark_orange'], weight='bold')
     ax.set_xlabel('Speed (mm/s)')
     ax.set_ylabel('Energy')
@@ -191,31 +203,22 @@ def plot_speed_profile(ax, speeds, power_total):
     ax2 = ax.twinx()
     ax2.yaxis.set_visible(False)
     
-    power_total_sum = np.array(resampled_power_total)
-    speed_array = np.array(resampled_speeds)
-    max_y = power_total_sum.max() + power_total_sum.max() * 0.05
-    ax.set_xlim([speed_array.min(), speed_array.max()])
+    max_y = power_total[0].max() + power_total[0].max() * 0.05
+    ax.set_xlim([speeds.min(), speeds.max()])
     ax.set_ylim([0, max_y])
     ax2.set_ylim([0, max_y])
 
-    ax.plot(resampled_speeds, resampled_power_total, label="X+Y+Z", color='purple')
+    ax.plot(speeds, power_total[0], label="X+Y+Z", color='purple', zorder=5)
     ax.plot(speeds, power_total[1], label="X", color='red')
     ax.plot(speeds, power_total[2], label="Y", color='green')
     ax.plot(speeds, power_total[3], label="Z", color='blue')
 
-    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))))
-
     if peaks.size:
-        ax.plot(speed_array[peaks], power_total_sum[peaks], "x", color='black', markersize=8)
+        ax.plot(speeds[peaks], power_total[0][peaks], "x", color='black', markersize=8)
         for idx, peak in enumerate(peaks):
             fontcolor = 'red'
             fontweight = 'bold'
-            ax.annotate(f"{idx+1}", (speed_array[peak], power_total_sum[peak]), 
+            ax.annotate(f"{idx+1}", (speeds[peak], power_total[0][peak]), 
                         textcoords="offset points", xytext=(8, 5), 
                         ha='left', fontsize=13, color=fontcolor, weight=fontweight)
         ax2.plot([], [], ' ', label=f'Number of peaks: {num_peaks}')
@@ -223,9 +226,9 @@ def plot_speed_profile(ax, speeds, power_total):
         ax2.plot([], [], ' ', label=f'No peaks detected')
 
     for idx, (start, end, energy) in enumerate(low_energy_zones):
-        ax.axvline(speed_array[start], color='red', linestyle='dotted', linewidth=1.5)
-        ax.axvline(speed_array[end], color='red', linestyle='dotted', linewidth=1.5)
-        ax2.fill_between(speed_array[start:end], 0, power_total_sum[start:end], color='green', alpha=0.2, label=f'Zone {idx+1}: {speed_array[start]:.1f} to {speed_array[end]:.1f} mm/s (mean energy: {energy:.2f}%)')
+        ax.axvline(speeds[start], color='red', linestyle='dotted', linewidth=1.5)
+        ax.axvline(speeds[end], color='red', linestyle='dotted', linewidth=1.5)
+        ax2.fill_between(speeds[start:end], 0, power_total[0][start:end], color='green', alpha=0.2, label=f'Zone {idx+1}: {speeds[start]:.1f} to {speeds[end]:.1f} mm/s (mean energy: {energy:.2f}%)')
 
     ax.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
     ax.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
@@ -236,22 +239,23 @@ def plot_speed_profile(ax, speeds, power_total):
     ax.legend(loc='upper left', prop=fontP)
     ax2.legend(loc='upper right', prop=fontP)
 
-    if peaks.size:
-        return speed_array[peaks]
-    else:
-        return None
+    return
 
 
-def plot_spectrogram(ax, speeds, freqs, power_spectral_densities, peaks, fr, max_freq):
+def plot_vibration_spectrogram(ax, speeds, freqs, power_spectral_densities, peaks, fr, max_freq):
+    # Prepare the spectrum data
     spectrum = np.empty([len(freqs), len(speeds)])
-
     for i in range(len(speeds)):
         for j in range(len(freqs)):
             spectrum[j, i] = power_spectral_densities[i][0][j]
 
     ax.set_title("Vibrations spectrogram", fontsize=14, color=KLIPPAIN_COLORS['dark_orange'], weight='bold')
-    ax.pcolormesh(speeds, freqs, spectrum, norm=matplotlib.colors.LogNorm(),
-            cmap='inferno', shading='gouraud')
+    # ax.pcolormesh(speeds, freqs, spectrum, norm=matplotlib.colors.LogNorm(),
+    #         cmap='inferno', shading='gouraud')
+    
+    ax.imshow(spectrum, norm=matplotlib.colors.LogNorm(), cmap='inferno',
+              aspect='auto', extent=[speeds[0], speeds[-1], freqs[0], freqs[-1]],
+              origin='lower', interpolation='antialiased')
 
     # Add peaks lines in the spectrogram to get hint from peaks found in the first graph
     if peaks is not None:
@@ -262,7 +266,7 @@ def plot_spectrogram(ax, speeds, freqs, power_spectral_densities, peaks, fr, max
                         verticalalignment='top', horizontalalignment='right')
     
     # Add motor resonance line
-    if fr is not None:
+    if fr is not None and fr > 25:
         ax.axhline(fr, color='cyan', linestyle='dotted', linewidth=1)
         ax.annotate(f"Motor resonance", (speeds[-1]*0.95, fr+2), 
                     textcoords="data", color='cyan', fontsize=10,
@@ -275,10 +279,7 @@ def plot_spectrogram(ax, speeds, freqs, power_spectral_densities, peaks, fr, max
     return
 
 
-def plot_vibration_profile(ax, freqs, vibration_power):
-    kernel = np.ones(10)/10
-    smoothed_vibration_power = np.convolve(vibration_power, kernel, mode='same')
-
+def plot_motor_profile(ax, freqs, motor_vibration_power, motor_fr, motor_zeta, motor_max_power_index):
     ax.set_title("Motors frequency profile", fontsize=14, color=KLIPPAIN_COLORS['dark_orange'], weight='bold')
     ax.set_xlabel('Energy')
     ax.set_ylabel('Frequency (hz)')
@@ -286,40 +287,23 @@ def plot_vibration_profile(ax, freqs, vibration_power):
     ax2 = ax.twinx()
     ax2.yaxis.set_visible(False)
     
-    vibr_power_array = np.array(smoothed_vibration_power)
-    freq_array = np.array(freqs)
-    max_x = vibr_power_array.max() + vibr_power_array.max() * 0.1
-    ax.set_ylim([freq_array.min(), freq_array.max()])
-    ax.set_xlim([0, max_x])
-    ax2.set_xlim([0, max_x])
+    ax.set_ylim([freqs.min(), freqs.max()])
+    ax.set_xlim([0, motor_vibration_power.max() + motor_vibration_power.max() * 0.1])
 
-    ax.plot(smoothed_vibration_power, freqs, color=KLIPPAIN_COLORS['orange'])
+    # Plot the profile curve
+    ax.plot(motor_vibration_power, freqs, color=KLIPPAIN_COLORS['orange'])
 
-    max_power_index = np.argmax(vibr_power_array)
-    fr = freq_array[max_power_index]
-    max_power = vibr_power_array[max_power_index]
-    half_power = max_power / math.sqrt(2)
-    idx_below = np.where(vibr_power_array[:max_power_index] <= half_power)[0][-1]
-    idx_above = np.where(vibr_power_array[max_power_index:] <= half_power)[0][0] + max_power_index
-    freq_below_half_power = freqs[idx_below] + (half_power - vibr_power_array[idx_below]) * (freqs[idx_below + 1] - freqs[idx_below]) / (vibr_power_array[idx_below + 1] - vibr_power_array[idx_below])
-    freq_above_half_power = freqs[idx_above - 1] + (half_power - vibr_power_array[idx_above - 1]) * (freqs[idx_above] - freqs[idx_above - 1]) / (vibr_power_array[idx_above] - vibr_power_array[idx_above - 1])
-    bandwidth = freq_above_half_power - freq_below_half_power
-    zeta = bandwidth / (2 * fr)
-
-    if fr > 20:
-        print_with_c_locale("Motors have a main resonant frequency at %.1fHz with an estimated damping ratio of %.3f" % (fr, zeta))
-    else:
-        print_with_c_locale("The resonance frequency of the motors is too low (%.1fHz). This is probably due to the test run with too high acceleration!" % fr)
-        print_with_c_locale("Try lowering the ACCEL value before restarting the macro to ensure that only constant speeds are recorded and that the dynamic behavior in the corners is not impacting the measurements.")
-
-    ax.plot(vibr_power_array[max_power_index], freq_array[max_power_index], "x", color='black', markersize=8)
+    # Tag the resonance peak
+    ax.plot(motor_vibration_power[motor_max_power_index], freqs[motor_max_power_index], "x", color='black', markersize=8)
     fontcolor = KLIPPAIN_COLORS['purple']
     fontweight = 'bold'
-    ax.annotate(f"R", (vibr_power_array[max_power_index], freq_array[max_power_index]), 
+    ax.annotate(f"R", (motor_vibration_power[motor_max_power_index], freqs[motor_max_power_index]), 
                 textcoords="offset points", xytext=(8, 8), 
                 ha='right', fontsize=13, color=fontcolor, weight=fontweight)
-    ax2.plot([], [], ' ', label="Motor resonant frequency (ω0): %.1fHz" % (fr))
-    ax2.plot([], [], ' ', label="Motor damping ratio (ζ): %.3f" % (zeta))
+    
+    # Add the legend
+    ax2.plot([], [], ' ', label="Motor resonant frequency (ω0): %.1fHz" % (motor_fr))
+    ax2.plot([], [], ' ', label="Motor damping ratio (ζ): %.3f" % (motor_zeta))
 
     ax.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
     ax.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
@@ -329,27 +313,13 @@ def plot_vibration_profile(ax, freqs, vibration_power):
     fontP.set_size('small')
     ax2.legend(loc='upper right', prop=fontP)
 
-    return fr if fr > 20 else None
+    return
 
 
 ######################################################################
 # Startup and main routines
 ######################################################################
 
-def parse_log(logname):
-    with open(logname) as f:
-        for header in f:
-            if not header.startswith('#'):
-                break
-        if not header.startswith('freq,psd_x,psd_y,psd_z,psd_xyz'):
-            # Raw accelerometer data
-            return np.loadtxt(logname, comments='#', delimiter=',')
-    # Power spectral density data or shaper calibration data
-    raise ValueError("File %s does not contain raw accelerometer data and therefore "
-               "is not supported by this script. Please use the official Klipper"
-               "calibrate_shaper.py script to process it instead." % (logname,))
-
-
 def extract_speed(logname):
     try:
         speed = re.search('sp(.+?)n', os.path.basename(logname)).group(1).replace('_','.')
@@ -363,73 +333,104 @@ def sort_and_slice(raw_speeds, raw_datas, remove):
     # Sort to get the speeds and their datas aligned and in ascending order
     raw_speeds, raw_datas = zip(*sorted(zip(raw_speeds, raw_datas), key=operator.itemgetter(0)))
 
-    # Remove beginning and end of the datas for each file to get only
-    # constant speed data and remove the start/stop phase of the movements
-    datas = []
+    # Optionally remove the beginning and end of each data file to get only
+    # the constant speed part of the segments and remove the start/stop phase
+    sliced_datas = []
     for data in raw_datas:
         sliced = round((len(data) * remove / 100) / 2)
-        datas.append(data[sliced:len(data)-sliced])
+        sliced_datas.append(data[sliced:len(data)-sliced])
 
-    return raw_speeds, datas
+    return raw_speeds, sliced_datas
 
 
-def setup_klipper_import(kdir):
-    global shaper_calibrate
-    kdir = os.path.expanduser(kdir)
-    sys.path.append(os.path.join(kdir, 'klippy'))
-    shaper_calibrate = importlib.import_module('.shaper_calibrate', 'extras')
-
-
-def vibrations_calibration(lognames, klipperdir="~/klipper", axisname=None, max_freq=1000., remove=0):
+def vibrations_calibration(lognames, klipperdir="~/klipper", axisname=None, accel=None, max_freq=1000., remove=0):
     set_locale()
-    setup_klipper_import(klipperdir)
+    global shaper_calibrate
+    shaper_calibrate = setup_klipper_import(klipperdir)
 
     # Parse the raw data and get them ready for analysis
     raw_datas = [parse_log(filename) for filename in lognames]
     raw_speeds = [extract_speed(filename) for filename in lognames]
     speeds, datas = sort_and_slice(raw_speeds, raw_datas, remove)
+    del raw_datas, raw_speeds
 
-    # As we assume that we have the same number of file for each speeds. We can group
-    # the PSD results by this number (to combine vibrations at given speed on all movements)
+    # As we assume that we have the same number of file for each speed increment, we can group
+    # the PSD results by this number (to combine all the segments of the pattern at a constant speed)
     group_by = speeds.count(speeds[0])
-    # Compute psd and total power of the signal
-    freqs, power_spectral_densities = calc_psd(datas, group_by, max_freq)
-    speed_power = calc_speed_profile(power_spectral_densities, freqs)
-    vibration_power = calc_vibration_profile(power_spectral_densities)
 
-    fig = matplotlib.pyplot.figure()
-    gs = matplotlib.gridspec.GridSpec(2, 2, height_ratios=[4, 3], width_ratios=[5, 3])
-    ax1 = fig.add_subplot(gs[0])
-    ax2 = fig.add_subplot(gs[2])
-    ax4 = fig.add_subplot(gs[3])
+    # Remove speeds duplicates and graph the processed datas
+    speeds = list(OrderedDict((x, True) for x in speeds).keys())
 
+    # Compute speed profile, vibration spectrogram and motor resonance profile
+    freqs, psd = compute_vibration_spectrogram(datas, group_by, max_freq)
+    upsampled_speeds, speeds_powers = compute_speed_profile(speeds, freqs, psd)
+    motor_vibration_power = compute_motor_profile(psd)
+
+    # Peak detection and low energy valleys (good speeds) identification between the peaks
+    num_peaks, vibration_peaks, peaks_speeds = detect_peaks(
+        speeds_powers[0], upsampled_speeds,
+        PEAKS_DETECTION_THRESHOLD * speeds_powers[0].max(),
+        PEAKS_RELATIVE_HEIGHT_THRESHOLD, 10, 10
+        )
+    low_energy_zones = identify_low_energy_zones(speeds_powers[0])
+
+    # Print the vibration peaks info in the console
+    formated_peaks_speeds = ["{:.1f}".format(pspeed) for pspeed in peaks_speeds]
+    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, formated_peaks_speeds))))
+
+    # Motor resonance estimation
+    motor_fr, motor_zeta, motor_max_power_index = compute_mechanical_parameters(motor_vibration_power, freqs)
+    if motor_fr > 25:
+        print_with_c_locale("Motors have a main resonant frequency at %.1fHz with an estimated damping ratio of %.3f" % (motor_fr, motor_zeta))
+    else:
+        print_with_c_locale("The detected resonance frequency of the motors is too low (%.1fHz). This is probably due to the test run with too high acceleration!" % motor_fr)
+        print_with_c_locale("Try lowering the ACCEL value before restarting the macro to ensure that only constant speeds are recorded and that the dynamic behavior in the corners is not impacting the measurements.")
+
+    # Create graph layout
+    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, gridspec_kw={
+            'height_ratios':[4, 3],
+            'width_ratios':[5, 3],
+            'bottom':0.050,
+            'top':0.890,
+            'left':0.057,
+            'right':0.985,
+            'hspace':0.166,
+            'wspace':0.138
+            })
+    ax2.remove() # top right graph is not used and left blank for now...
+    fig.set_size_inches(14, 11.6)
+
+    # Add title
     title_line1 = "VIBRATIONS MEASUREMENT TOOL"
     fig.text(0.075, 0.965, title_line1, ha='left', va='bottom', fontsize=20, color=KLIPPAIN_COLORS['purple'], weight='bold')
     try:
         filename_parts = (lognames[0].split('/')[-1]).split('_')
         dt = datetime.strptime(f"{filename_parts[1]} {filename_parts[2].split('-')[0]}", "%Y%m%d %H%M%S")
-        title_line2 = dt.strftime('%x %X') + ' -- ' + axisname.upper() + ' axis'
+        title_line2 = dt.strftime('%x %X')
+        if axisname is not None:
+            title_line2 += ' -- ' + str(axisname).upper() + ' axis'
+        if accel is not None:
+            title_line2 += ' at ' + str(accel) + ' mm/s²'
     except:
         print_with_c_locale("Warning: CSV filename look to be different than expected (%s)" % (lognames[0]))
         title_line2 = lognames[0].split('/')[-1]
     fig.text(0.075, 0.957, title_line2, ha='left', va='top', fontsize=16, color=KLIPPAIN_COLORS['dark_purple'])
 
-    # Remove speeds duplicates and graph the processed datas
-    speeds = list(OrderedDict((x, True) for x in speeds).keys())
-
-    speed_peaks = plot_speed_profile(ax1, speeds, speed_power)
-    fr = plot_vibration_profile(ax4, freqs, vibration_power)
-    plot_spectrogram(ax2, speeds, freqs, power_spectral_densities, speed_peaks, fr, max_freq)
-
-    fig.set_size_inches(14, 11.6)
-    fig.tight_layout()
-    fig.subplots_adjust(top=0.89)
+    # Plot the graphs
+    plot_speed_profile(ax1, upsampled_speeds, speeds_powers, num_peaks, vibration_peaks, low_energy_zones)
+    plot_motor_profile(ax4, freqs, motor_vibration_power, motor_fr, motor_zeta, motor_max_power_index)
+    plot_vibration_spectrogram(ax3, speeds, freqs, psd, peaks_speeds, motor_fr, max_freq)
 
     # Adding a small Klippain logo to the top left corner of the figure
-    ax_logo = fig.add_axes([0.001, 0.924, 0.075, 0.075], anchor='NW', zorder=-1)
-    ax_logo.imshow(matplotlib.pyplot.imread(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'klippain.png')))
+    ax_logo = fig.add_axes([0.001, 0.924, 0.075, 0.075], anchor='NW')
+    ax_logo.imshow(plt.imread(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'klippain.png')))
     ax_logo.axis('off')
 
+    # Adding Shake&Tune version in the top right corner
+    st_version = get_git_version()
+    if st_version is not None:
+        fig.text(0.995, 0.985, st_version, ha='right', va='bottom', fontsize=8, color=KLIPPAIN_COLORS['purple'])
+
     return fig
 
 
@@ -440,11 +441,13 @@ def main():
     opts.add_option("-o", "--output", type="string", dest="output",
                     default=None, help="filename of output graph")
     opts.add_option("-a", "--axis", type="string", dest="axisname",
-                    default=None, help="axis name to be shown on the side of the graph")
+                    default=None, help="axis name to be printed on the graph")
+    opts.add_option("-c", "--accel", type="int", dest="accel",
+                    default=None, help="accel value to be printed on the graph")
     opts.add_option("-f", "--max_freq", type="float", default=1000.,
                     help="maximum frequency to graph")
     opts.add_option("-r", "--remove", type="int", default=0,
-                    help="percentage of data removed at start/end of each files")
+                    help="percentage of data removed at start/end of each CSV files")
     opts.add_option("-k", "--klipper_dir", type="string", dest="klipperdir",
                     default="~/klipper", help="main klipper directory")
     options, args = opts.parse_args()
@@ -455,8 +458,8 @@ def main():
     if options.remove > 50 or options.remove < 0:
         opts.error("You must specify a correct percentage (option -r) in the 0-50 range")
 
-    fig = vibrations_calibration(args, options.klipperdir, options.axisname, options.max_freq, options.remove)
-    fig.savefig(options.output)
+    fig = vibrations_calibration(args, options.klipperdir, options.axisname, options.accel, options.max_freq, options.remove)
+    fig.savefig(options.output, dpi=150)
 
 
 if __name__ == '__main__':

K-ShakeTune/scripts/is_workflow.py

@@ -125,7 +125,7 @@ def create_shaper_graph():
     return
 
 
-def create_vibrations_graph(axis_name):
+def create_vibrations_graph(axis_name, accel):
     current_date = datetime.now().strftime('%Y%m%d_%H%M%S')
     lognames = []
 
@@ -155,7 +155,7 @@ def create_vibrations_graph(axis_name):
     time.sleep(5)
 
     # Generate the vibration graph and its name
-    fig = vibrations_calibration(lognames, KLIPPER_FOLDER, axis_name)
+    fig = vibrations_calibration(lognames, KLIPPER_FOLDER, axis_name, accel)
     png_filename = os.path.join(RESULTS_FOLDER, RESULTS_SUBFOLDERS[2], f'vibrations_{current_date}_{axis_name}.png')
     
     # Archive all the csv files in a tarball and remove them to clean up the results folder
@@ -251,7 +251,7 @@ def main():
     elif sys.argv[1].lower() == 'shaper':
         create_shaper_graph()
     elif sys.argv[1].lower() == 'vibrations':
-        create_vibrations_graph(axis_name=sys.argv[2])
+        create_vibrations_graph(axis_name=sys.argv[2], accel=sys.argv[3])
     elif sys.argv[1].lower() == 'axesmap':
         find_axesmap(accel=sys.argv[2])
     else: