Merge branch 'develop' into cross-belts

shaketune/__init__.py

@@ -0,0 +1,18 @@
+#!/usr/bin/env python3
+
+############################################
+###### INPUT SHAPER KLIPPAIN WORKFLOW ######
+############################################
+# Written by Frix_x#0161 #
+
+# This module functions as a plugin within Klipper, aimed at enhancing printer diagnostics. It serves multiple purposes:
+# 1. Diagnosing and pinpointing vibration sources in the printer.
+# 2. Conducting standard axis input shaper tests on the XY axes to determine the optimal input shaper filter.
+# 3. Executing a specialized half-axis test for CoreXY printers to analyze and compare the frequency profiles of individual belts.
+
+
+from .shaketune import ShakeTune as ShakeTune
+
+
+def load_config(config) -> ShakeTune:
+    return ShakeTune(config)

shaketune/helpers/common_func.py

@@ -0,0 +1,245 @@
+#!/usr/bin/env python3
+
+# Common functions for the Shake&Tune package
+# Written by Frix_x#0161 #
+
+import math
+import os
+import sys
+from importlib import import_module
+from pathlib import Path
+
+import numpy as np
+from scipy.signal import spectrogram
+
+from .console_output import ConsoleOutput
+
+# Constant used to define the standard axis direction and names
+AXIS_CONFIG = [
+    {'axis': 'x', 'direction': (1, 0, 0), 'label': 'axis_X'},
+    {'axis': 'y', 'direction': (0, 1, 0), 'label': 'axis_Y'},
+    {'axis': 'a', 'direction': (1, -1, 0), 'label': 'belt_A'},
+    {'axis': 'b', 'direction': (1, 1, 0), 'label': 'belt_B'},
+]
+
+
+def parse_log(logname):
+    try:
+        with open(logname) as f:
+            header = None
+            for line in f:
+                cleaned_line = line.strip()
+
+                # Check for a PSD file generated by Klipper and raise a warning
+                if cleaned_line.startswith('#freq,psd_x,psd_y,psd_z,psd_xyz'):
+                    ConsoleOutput.print(
+                        'Warning: %s does not contain raw accelerometer data. '
+                        'Please use the official Klipper script to process it instead. '
+                        'It will be ignored by Shake&Tune!' % (logname,)
+                    )
+                    return None
+
+                # Check for the expected header for Shake&Tune (raw accelerometer data from Klipper)
+                elif cleaned_line.startswith('#time,accel_x,accel_y,accel_z'):
+                    header = cleaned_line
+                    break
+
+            if not header:
+                ConsoleOutput.print(
+                    'Warning: file %s has an incorrect header and will be ignored by Shake&Tune!\n'
+                    "Expected '#time,accel_x,accel_y,accel_z', but got '%s'." % (logname, header.strip())
+                )
+                return None
+
+            # If we have the correct raw data header, proceed to load the data
+            data = np.loadtxt(logname, comments='#', delimiter=',', skiprows=1)
+            if data.ndim == 1 or data.shape[1] != 4:
+                ConsoleOutput.print(
+                    'Warning: %s does not have the correct data format; expected 4 columns. '
+                    'It will be ignored by Shake&Tune!' % (logname,)
+                )
+                return None
+
+            return data
+
+    except Exception as err:
+        ConsoleOutput.print(f'Error while reading {logname}: {err}. It will be ignored by Shake&Tune!')
+        return None
+
+
+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
+        from git import GitCommandError, Repo
+
+        script_path = Path(__file__).resolve()
+        repo_path = script_path.parents[1]
+        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:
+        return None
+
+
+# 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(0.5 * Fs - 1).bit_length()
+    window = np.kaiser(M, 6.0)
+
+    def _specgram(x):
+        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'])
+    for axis in 'yz':
+        pdata += _specgram(d[axis])[2]
+    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, min_freq=None):
+    max_under_min_freq = False
+
+    if min_freq is not None:
+        min_freq_index = np.searchsorted(freqs, min_freq, side='left')
+        if min_freq_index >= len(freqs):
+            return None, None, None, max_under_min_freq
+        if np.argmax(psd) < min_freq_index:
+            max_under_min_freq = True
+    else:
+        min_freq_index = 0
+
+    # Consider only the part of the signal above min_freq
+    psd_above_min_freq = psd[min_freq_index:]
+    if len(psd_above_min_freq) == 0:
+        return None, None, None, max_under_min_freq
+
+    max_power_index_above_min_freq = np.argmax(psd_above_min_freq)
+    max_power_index = max_power_index_above_min_freq + min_freq_index
+    fr = freqs[max_power_index]
+    max_power = psd[max_power_index]
+
+    half_power = max_power / math.sqrt(2)
+    indices_below = np.where(psd[:max_power_index] <= half_power)[0]
+    indices_above = np.where(psd[max_power_index:] <= half_power)[0]
+
+    # If we are not able to find points around the half power, we can't compute the damping ratio and return None instead
+    if len(indices_below) == 0 or len(indices_above) == 0:
+        return fr, None, max_power_index, max_under_min_freq
+
+    idx_below = indices_below[-1]
+    idx_above = indices_above[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
+    bw1 = math.pow(bandwidth / fr, 2)
+    bw2 = math.pow(bandwidth / fr, 4)
+
+    try:
+        zeta = math.sqrt(0.5 - math.sqrt(1 / (4 + 4 * bw1 - bw2)))
+    except ValueError:
+        # If a math problem arise such as a negative sqrt term, we also return None instead for damping ratio
+        return fr, None, max_power_index, max_under_min_freq
+
+    return fr, zeta, max_power_index, max_under_min_freq
+
+
+# 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]
+
+
+# The goal is to find zone outside of peaks (flat low energy zones) in a signal
+def identify_low_energy_zones(power_total, detection_threshold=0.1):
+    valleys = []
+
+    # Calculate the a "mean + 1/4" and standard deviation of the entire power_total
+    mean_energy = np.mean(power_total) + (np.max(power_total) - np.min(power_total)) / 4
+    std_energy = np.std(power_total)
+
+    # Define a threshold value as "mean + 1/4" minus a certain number of standard deviations
+    threshold_value = mean_energy - detection_threshold * std_energy
+
+    # Find valleys in power_total based on the threshold
+    in_valley = False
+    start_idx = 0
+    for i, value in enumerate(power_total):
+        if not in_valley and value < threshold_value:
+            in_valley = True
+            start_idx = i
+        elif in_valley and value >= threshold_value:
+            in_valley = False
+            valleys.append((start_idx, i))
+
+    # If the last point is still in a valley, close the valley
+    if in_valley:
+        valleys.append((start_idx, len(power_total) - 1))
+
+    max_signal = np.max(power_total)
+
+    # Calculate mean energy for each valley as a percentage of the maximum of the signal
+    valley_means_percentage = []
+    for start, end in valleys:
+        if not np.isnan(np.mean(power_total[start:end])):
+            valley_means_percentage.append((start, end, (np.mean(power_total[start:end]) / max_signal) * 100))
+
+    # Sort valleys based on mean percentage values
+    sorted_valleys = sorted(valley_means_percentage, key=lambda x: x[2])
+
+    return sorted_valleys

shaketune/helpers/console_output.py

@@ -0,0 +1,24 @@
+import io
+from typing import Callable, Optional
+
+
+class ConsoleOutput:
+    """
+    Print output to stdout or to an alternative like the Klipper console through a callback
+    """
+
+    _output_func: Optional[Callable[[str], None]] = None
+
+    @classmethod
+    def register_output_callback(cls, output_func: Optional[Callable[[str], None]]):
+        cls._output_func = output_func
+
+    @classmethod
+    def print(cls, *args, **kwargs):
+        if not cls._output_func:
+            print(*args, **kwargs)
+            return
+
+        with io.StringIO() as mem_output:
+            print(*args, file=mem_output, **kwargs)
+            cls._output_func(mem_output.getvalue())

shaketune/helpers/filemanager.py

@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+
+# Common file management functions for the Shake&Tune package
+# Written by Frix_x#0161 #
+
+import os
+import time
+from pathlib import Path
+
+
+def wait_file_ready(filepath: Path, timeout: int = 60) -> None:
+    file_busy = True
+    loop_count = 0
+
+    while file_busy:
+        if loop_count >= timeout:
+            raise TimeoutError(f'Klipper is taking too long to release the CSV file ({filepath})!')
+
+        # Try to open the file in write-only mode to check if it is in use
+        # If we successfully open and close the file, it is not in use
+        try:
+            fd = os.open(filepath, os.O_WRONLY)
+            os.close(fd)
+            file_busy = False
+        except OSError:
+            # If OSError is caught, it indicates the file is still being used
+            pass
+        except Exception:
+            # If another exception is raised, it's not a problem, we just loop again
+            pass
+
+        loop_count += 1
+        time.sleep(1)
+
+
+def ensure_folders_exist(folders: list[Path]) -> None:
+    for folder in folders:
+        folder.mkdir(parents=True, exist_ok=True)

shaketune/measurement/__init__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+
+from .axes_input_shaper import axes_shaper_calibration as axes_shaper_calibration
+from .axes_map import axes_map_calibration as axes_map_calibration
+from .belts_comparison import compare_belts_responses as compare_belts_responses
+from .static_freq import excitate_axis_at_freq as excitate_axis_at_freq
+from .vibrations_profile import create_vibrations_profile as create_vibrations_profile

shaketune/measurement/accelerometer.py

@@ -0,0 +1,60 @@
+#!/usr/bin/env python3
+
+# This file provides a custom and internal Shake&Tune Accelerometer helper that is
+# an interface to Klipper's own accelerometer classes. It is used to start and
+# stop accelerometer measurements and write the data to a file in a blocking manner.
+
+import time
+
+# from ..helpers.console_output import ConsoleOutput
+
+
+class Accelerometer:
+    def __init__(self, klipper_accelerometer):
+        self._k_accelerometer = klipper_accelerometer
+        self._bg_client = None
+
+    @staticmethod
+    def find_axis_accelerometer(printer, axis: str = 'xy'):
+        accel_chip_names = printer.lookup_object('resonance_tester').accel_chip_names
+        for chip_axis, chip_name in accel_chip_names:
+            if axis in ['x', 'y'] and chip_axis == 'xy':
+                return chip_name
+            elif chip_axis == axis:
+                return chip_name
+        return None
+
+    def start_measurement(self):
+        if self._bg_client is None:
+            self._bg_client = self._k_accelerometer.start_internal_client()
+            # ConsoleOutput.print('Accelerometer measurements started')
+        else:
+            raise ValueError('measurements already started!')
+
+    def stop_measurement(self, name: str = None, append_time: bool = True):
+        if self._bg_client is None:
+            raise ValueError('measurements need to be started first!')
+
+        timestamp = time.strftime('%Y%m%d_%H%M%S')
+        if name is None:
+            name = timestamp
+        elif append_time:
+            name += f'_{timestamp}'
+
+        if not name.replace('-', '').replace('_', '').isalnum():
+            raise ValueError('invalid file name!')
+
+        bg_client = self._bg_client
+        self._bg_client = None
+        bg_client.finish_measurements()
+
+        filename = f'/tmp/shaketune-{name}.csv'
+        self._write_to_file(bg_client, filename)
+        # ConsoleOutput.print(f'Accelerometer measurements stopped. Data written to {filename}')
+
+    def _write_to_file(self, bg_client, filename):
+        with open(filename, 'w') as f:
+            f.write('#time,accel_x,accel_y,accel_z\n')
+            samples = bg_client.samples or bg_client.get_samples()
+            for t, accel_x, accel_y, accel_z in samples:
+                f.write('%.6f,%.6f,%.6f,%.6f\n' % (t, accel_x, accel_y, accel_z))

shaketune/measurement/axes_input_shaper.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+
+
+from ..helpers.common_func import AXIS_CONFIG
+from ..helpers.console_output import ConsoleOutput
+from ..shaketune_thread import ShakeTuneThread
+from .accelerometer import Accelerometer
+from .resonance_test import vibrate_axis
+
+
+def axes_shaper_calibration(gcmd, config, st_thread: ShakeTuneThread) -> None:
+    min_freq = gcmd.get_float('FREQ_START', default=5, minval=1)
+    max_freq = gcmd.get_float('FREQ_END', default=133.33, minval=1)
+    hz_per_sec = gcmd.get_float('HZ_PER_SEC', default=1, minval=1)
+    accel_per_hz = gcmd.get_float('ACCEL_PER_HZ', default=None)
+    axis_input = gcmd.get('AXIS', default='all').lower()
+    if axis_input not in ['x', 'y', 'all']:
+        gcmd.error('AXIS selection invalid. Should be either x, y, or all!')
+    scv = gcmd.get_float('SCV', default=None, minval=0)
+    max_sm = gcmd.get_float('MAX_SMOOTHING', default=None, minval=0)
+    feedrate_travel = gcmd.get_float('TRAVEL_SPEED', default=120.0, minval=20.0)
+    z_height = gcmd.get_float('Z_HEIGHT', default=None, minval=1)
+
+    printer = config.get_printer()
+    gcode = printer.lookup_object('gcode')
+    toolhead = printer.lookup_object('toolhead')
+    res_tester = printer.lookup_object('resonance_tester')
+    systime = printer.get_reactor().monotonic()
+
+    if scv is None:
+        toolhead_info = toolhead.get_status(systime)
+        scv = toolhead_info['square_corner_velocity']
+
+    if accel_per_hz is None:
+        accel_per_hz = res_tester.test.accel_per_hz
+    max_accel = max_freq * accel_per_hz
+
+    # Move to the starting point
+    test_points = res_tester.test.get_start_test_points()
+    if len(test_points) > 1:
+        gcmd.error('Only one test point in the [resonance_tester] section is supported by Shake&Tune.')
+    if test_points[0] == (-1, -1, -1):
+        if z_height is None:
+            gcmd.error(
+                'Z_HEIGHT parameter is required if the test_point in [resonance_tester] section is set to -1,-1,-1'
+            )
+        # Use center of bed in case the test point in [resonance_tester] is set to -1,-1,-1
+        # This is usefull to get something automatic and is also used in the Klippain modular config
+        kin_info = toolhead.kin.get_status(systime)
+        mid_x = (kin_info['axis_minimum'].x + kin_info['axis_maximum'].x) / 2
+        mid_y = (kin_info['axis_minimum'].y + kin_info['axis_maximum'].y) / 2
+        point = (mid_x, mid_y, z_height)
+    else:
+        x, y, z = test_points[0]
+        if z_height is not None:
+            z = z_height
+        point = (x, y, z)
+
+    toolhead.manual_move(point, feedrate_travel)
+
+    # Configure the graph creator
+    creator = st_thread.get_graph_creator()
+    creator.configure(scv, max_sm, accel_per_hz)
+
+    # set the needed acceleration values for the test
+    toolhead_info = toolhead.get_status(systime)
+    old_accel = toolhead_info['max_accel']
+    old_mcr = toolhead_info['minimum_cruise_ratio']
+    gcode.run_script_from_command(f'SET_VELOCITY_LIMIT ACCEL={max_accel} MINIMUM_CRUISE_RATIO=0')
+
+    # Deactivate input shaper if it is active to get raw movements
+    input_shaper = printer.lookup_object('input_shaper', None)
+    if input_shaper is not None:
+        input_shaper.disable_shaping()
+    else:
+        input_shaper = None
+
+    # Filter axis configurations based on user input, assuming 'axis_input' can be 'x', 'y', 'all' (that means 'x' and 'y')
+    filtered_config = [
+        a for a in AXIS_CONFIG if a['axis'] == axis_input or (axis_input == 'all' and a['axis'] in ('x', 'y'))
+    ]
+    for config in filtered_config:
+        # First we need to find the accelerometer chip suited for the axis
+        accel_chip = Accelerometer.find_axis_accelerometer(printer, config['axis'])
+        if accel_chip is None:
+            gcmd.error(
+                'No suitable accelerometer found for measurement! Multi-accelerometer configurations are not supported for this macro.'
+            )
+        accelerometer = Accelerometer(printer.lookup_object(accel_chip))
+
+        # Then do the actual measurements
+        accelerometer.start_measurement()
+        vibrate_axis(toolhead, gcode, config['direction'], min_freq, max_freq, hz_per_sec, accel_per_hz)
+        accelerometer.stop_measurement(config['label'], append_time=True)
+
+        # And finally generate the graph for each measured axis
+        ConsoleOutput.print(f'{config["axis"].upper()} axis frequency profile generation...')
+        ConsoleOutput.print('This may take some time (1-3min)')
+        st_thread.run()
+
+    # Re-enable the input shaper if it was active
+    if input_shaper is not None:
+        input_shaper.enable_shaping()
+
+    # Restore the previous acceleration values
+    gcode.run_script_from_command(f'SET_VELOCITY_LIMIT ACCEL={old_accel} MINIMUM_CRUISE_RATIO={old_mcr}')

shaketune/measurement/axes_map.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python3
+
+
+from ..helpers.console_output import ConsoleOutput
+from ..shaketune_thread import ShakeTuneThread
+from .accelerometer import Accelerometer
+
+
+def axes_map_calibration(gcmd, config, st_thread: ShakeTuneThread) -> None:
+    z_height = gcmd.get_float('Z_HEIGHT', default=20.0)
+    speed = gcmd.get_float('SPEED', default=80.0, minval=20.0)
+    accel = gcmd.get_int('ACCEL', default=1500, minval=100)
+    feedrate_travel = gcmd.get_float('TRAVEL_SPEED', default=120.0, minval=20.0)
+    accel_chip = gcmd.get('ACCEL_CHIP', default=None)
+
+    printer = config.get_printer()
+    gcode = printer.lookup_object('gcode')
+    toolhead = printer.lookup_object('toolhead')
+    systime = printer.get_reactor().monotonic()
+
+    if accel_chip is None:
+        accel_chip = Accelerometer.find_axis_accelerometer(printer, 'xy')
+        if accel_chip is None:
+            gcmd.error(
+                'No accelerometer specified for measurement! Multi-accelerometer configurations are not supported for this macro.'
+            )
+    accelerometer = Accelerometer(printer.lookup_object(accel_chip))
+
+    toolhead_info = toolhead.get_status(systime)
+    old_accel = toolhead_info['max_accel']
+    old_mcr = toolhead_info['minimum_cruise_ratio']
+    old_sqv = toolhead_info['square_corner_velocity']
+
+    # set the wanted acceleration values
+    gcode.run_script_from_command(f'SET_VELOCITY_LIMIT ACCEL={accel} MINIMUM_CRUISE_RATIO=0 SQUARE_CORNER_VELOCITY=5.0')
+
+    # Deactivate input shaper if it is active to get raw movements
+    input_shaper = printer.lookup_object('input_shaper', None)
+    if input_shaper is not None:
+        input_shaper.disable_shaping()
+    else:
+        input_shaper = None
+
+    kin_info = toolhead.kin.get_status(systime)
+    mid_x = (kin_info['axis_minimum'].x + kin_info['axis_maximum'].x) / 2
+    mid_y = (kin_info['axis_minimum'].y + kin_info['axis_maximum'].y) / 2
+    _, _, _, E = toolhead.get_position()
+
+    # Going to the start position
+    toolhead.move([mid_x - 15, mid_y - 15, z_height, E], feedrate_travel)
+    toolhead.dwell(0.5)
+
+    # Start the measurements and do the movements (+X, +Y and then +Z)
+    accelerometer.start_measurement()
+    toolhead.dwell(1)
+    toolhead.move([mid_x + 15, mid_y - 15, z_height, E], speed)
+    toolhead.dwell(1)
+    toolhead.move([mid_x + 15, mid_y + 15, z_height, E], speed)
+    toolhead.dwell(1)
+    toolhead.move([mid_x + 15, mid_y + 15, z_height + 15, E], speed)
+    toolhead.dwell(1)
+    accelerometer.stop_measurement('axemap')
+
+    # Re-enable the input shaper if it was active
+    if input_shaper is not None:
+        input_shaper.enable_shaping()
+
+    # Restore the previous acceleration values
+    gcode.run_script_from_command(
+        f'SET_VELOCITY_LIMIT ACCEL={old_accel} MINIMUM_CRUISE_RATIO={old_mcr} SQUARE_CORNER_VELOCITY={old_sqv}'
+    )
+    toolhead.wait_moves()
+
+    # Run post-processing
+    ConsoleOutput.print('Analysis of the movements...')
+    creator = st_thread.get_graph_creator()
+    creator.configure(accel)
+    st_thread.run()

shaketune/measurement/belts_comparison.py

@@ -0,0 +1,95 @@
+#!/usr/bin/env python3
+
+
+from ..helpers.common_func import AXIS_CONFIG
+from ..helpers.console_output import ConsoleOutput
+from ..shaketune_thread import ShakeTuneThread
+from .accelerometer import Accelerometer
+from .motorsconfigparser import MotorsConfigParser
+from .resonance_test import vibrate_axis
+
+
+def compare_belts_responses(gcmd, config, st_thread: ShakeTuneThread) -> None:
+    min_freq = gcmd.get_float('FREQ_START', default=5.0, minval=1)
+    max_freq = gcmd.get_float('FREQ_END', default=133.33, minval=1)
+    hz_per_sec = gcmd.get_float('HZ_PER_SEC', default=1.0, minval=1)
+    accel_per_hz = gcmd.get_float('ACCEL_PER_HZ', default=None)
+    feedrate_travel = gcmd.get_float('TRAVEL_SPEED', default=120.0, minval=20.0)
+    z_height = gcmd.get_float('Z_HEIGHT', default=None, minval=1)
+
+    printer = config.get_printer()
+    gcode = printer.lookup_object('gcode')
+    toolhead = printer.lookup_object('toolhead')
+    res_tester = printer.lookup_object('resonance_tester')
+    systime = printer.get_reactor().monotonic()
+
+    accel_chip = Accelerometer.find_axis_accelerometer(printer, 'xy')
+    if accel_chip is None:
+        gcmd.error(
+            'No suitable accelerometer found for measurement! Multi-accelerometer configurations are not supported for this macro.'
+        )
+    accelerometer = Accelerometer(printer.lookup_object(accel_chip))
+
+    if accel_per_hz is None:
+        accel_per_hz = res_tester.test.accel_per_hz
+    max_accel = max_freq * accel_per_hz
+
+    # Move to the starting point
+    test_points = res_tester.test.get_start_test_points()
+    if len(test_points) > 1:
+        gcmd.error('Only one test point in the [resonance_tester] section is supported by Shake&Tune.')
+    if test_points[0] == (-1, -1, -1):
+        if z_height is None:
+            gcmd.error(
+                'Z_HEIGHT parameter is required if the test_point in [resonance_tester] section is set to -1,-1,-1'
+            )
+        # Use center of bed in case the test point in [resonance_tester] is set to -1,-1,-1
+        # This is usefull to get something automatic and is also used in the Klippain modular config
+        kin_info = toolhead.kin.get_status(systime)
+        mid_x = (kin_info['axis_minimum'].x + kin_info['axis_maximum'].x) / 2
+        mid_y = (kin_info['axis_minimum'].y + kin_info['axis_maximum'].y) / 2
+        point = (mid_x, mid_y, z_height)
+    else:
+        x, y, z = test_points[0]
+        if z_height is not None:
+            z = z_height
+        point = (x, y, z)
+
+    toolhead.manual_move(point, feedrate_travel)
+
+    # Configure the graph creator
+    motors_config_parser = MotorsConfigParser(config, motors=None)
+    creator = st_thread.get_graph_creator()
+    creator.configure(motors_config_parser.kinematics, accel_per_hz)
+
+    # set the needed acceleration values for the test
+    toolhead_info = toolhead.get_status(systime)
+    old_accel = toolhead_info['max_accel']
+    old_mcr = toolhead_info['minimum_cruise_ratio']
+    gcode.run_script_from_command(f'SET_VELOCITY_LIMIT ACCEL={max_accel} MINIMUM_CRUISE_RATIO=0')
+
+    # Deactivate input shaper if it is active to get raw movements
+    input_shaper = printer.lookup_object('input_shaper', None)
+    if input_shaper is not None:
+        input_shaper.disable_shaping()
+    else:
+        input_shaper = None
+
+    # Filter axis configurations to get the A and B axis only
+    filtered_config = [a for a in AXIS_CONFIG if a['axis'] in ('a', 'b')]
+    for config in filtered_config:
+        accelerometer.start_measurement()
+        vibrate_axis(toolhead, gcode, config['direction'], min_freq, max_freq, hz_per_sec, accel_per_hz)
+        accelerometer.stop_measurement(config['label'], append_time=True)
+
+    # Re-enable the input shaper if it was active
+    if input_shaper is not None:
+        input_shaper.enable_shaping()
+
+    # Restore the previous acceleration values
+    gcode.run_script_from_command(f'SET_VELOCITY_LIMIT ACCEL={old_accel} MINIMUM_CRUISE_RATIO={old_mcr}')
+
+    # Run post-processing
+    ConsoleOutput.print('Belts comparative frequency profile generation...')
+    ConsoleOutput.print('This may take some time (3-5min)')
+    st_thread.run()

shaketune/measurement/macros.cfg

@@ -0,0 +1,8 @@
+
+# [gcode_macro AXES_MAP_CALIBRATION]
+# gcode:
+#     {% set z_height = params.Z_HEIGHT|default(20)|int %} # z height to put the toolhead before starting the movements
+#     {% set speed = params.SPEED|default(80)|float * 60 %} # feedrate for the movements
+#     {% set accel = params.ACCEL|default(1500)|int %} # accel value used to move on the pattern
+#     {% set feedrate_travel = params.TRAVEL_SPEED|default(120)|int * 60 %} # travel feedrate between moves
+#     {% set accel_chip = params.ACCEL_CHIP|default("adxl345") %} # ADXL chip name in the config

shaketune/measurement/motorsconfigparser.py

@@ -0,0 +1,189 @@
+#!/usr/bin/env python3
+
+# Classes to retrieve a couple of motors infos and extract the relevant information
+# from the Klipper configuration and the TMC registers
+# Written by Frix_x#0161 #
+
+from typing import Any, Dict, List, Optional
+
+TRINAMIC_DRIVERS = ['tmc2130', 'tmc2208', 'tmc2209', 'tmc2240', 'tmc2660', 'tmc5160']
+MOTORS = ['stepper_x', 'stepper_y', 'stepper_x1', 'stepper_y1', 'stepper_z', 'stepper_z1', 'stepper_z2', 'stepper_z3']
+RELEVANT_TMC_REGISTERS = ['CHOPCONF', 'PWMCONF', 'COOLCONF', 'TPWMTHRS', 'TCOOLTHRS']
+
+
+class Motor:
+    def __init__(self, name: str):
+        self.name: str = name
+        self._registers: Dict[str, Dict[str, Any]] = {}
+        self._config: Dict[str, Any] = {}
+
+    def set_register(self, register: str, value_dict: dict) -> None:
+        # First we filter out entries with a value of 0 to avoid having too much uneeded data
+        value_dict = {k: v for k, v in value_dict.items() if v != 0}
+
+        # Special parsing for CHOPCONF to extract meaningful values
+        if register == 'CHOPCONF':
+            # Add intpol=0 if missing from the register dump to force printing it as it's important
+            if 'intpol' not in value_dict:
+                value_dict['intpol'] = '0'
+            # Remove the microsteps entry as the format here is not easy to read and
+            # it's already read in the correct format directly from the Klipper config
+            if 'mres' in value_dict:
+                del value_dict['mres']
+
+        # Special parsing for CHOPCONF to avoid pwm_ before each values
+        if register == 'PWMCONF':
+            new_value_dict = {}
+            for key, val in value_dict.items():
+                if key.startswith('pwm_'):
+                    key = key[4:]
+                new_value_dict[key] = val
+            value_dict = new_value_dict
+
+        # Then gets merged all the thresholds into the same THRS virtual register
+        if register in ['TPWMTHRS', 'TCOOLTHRS']:
+            existing_thrs = self._registers.get('THRS', {})
+            merged_values = {**existing_thrs, **value_dict}
+            self._registers['THRS'] = merged_values
+        else:
+            self._registers[register] = value_dict
+
+    def get_register(self, register: str) -> Optional[Dict[str, Any]]:
+        return self._registers.get(register)
+
+    def get_registers(self) -> Dict[str, Dict[str, Any]]:
+        return self._registers
+
+    def set_config(self, field: str, value: Any) -> None:
+        self._config[field] = value
+
+    def get_config(self, field: str) -> Optional[Any]:
+        return self._config.get(field)
+
+    def __str__(self):
+        return f'Stepper: {self.name}\nKlipper config: {self._config}\nTMC Registers: {self._registers}'
+
+    # Return the other motor config and registers that are different from the current motor
+    def compare_to(self, other: 'Motor') -> Optional[Dict[str, Dict[str, Any]]]:
+        differences = {'config': {}, 'registers': {}}
+
+        # Compare Klipper config
+        all_keys = self._config.keys() | other._config.keys()
+        for key in all_keys:
+            val1 = self._config.get(key)
+            val2 = other._config.get(key)
+            if val1 != val2:
+                differences['config'][key] = val2
+
+        # Compare TMC registers
+        all_keys = self._registers.keys() | other._registers.keys()
+        for key in all_keys:
+            reg1 = self._registers.get(key, {})
+            reg2 = other._registers.get(key, {})
+            if reg1 != reg2:
+                reg_diffs = {}
+                sub_keys = reg1.keys() | reg2.keys()
+                for sub_key in sub_keys:
+                    reg_val1 = reg1.get(sub_key)
+                    reg_val2 = reg2.get(sub_key)
+                    if reg_val1 != reg_val2:
+                        reg_diffs[sub_key] = reg_val2
+                if reg_diffs:
+                    differences['registers'][key] = reg_diffs
+
+        # Clean up: remove empty sections if there are no differences
+        if not differences['config']:
+            del differences['config']
+        if not differences['registers']:
+            del differences['registers']
+
+        if not differences:
+            return None
+
+        return differences
+
+
+class MotorsConfigParser:
+    def __init__(self, config, motors: List[str] = MOTORS, drivers: List[str] = TRINAMIC_DRIVERS):
+        self._printer = config.get_printer()
+
+        self._motors: List[Motor] = []
+
+        if motors is not None:
+            for motor_name in motors:
+                for driver in drivers:
+                    tmc_object = self._printer.lookup_object(f'{driver} {motor_name}', None)
+                    if tmc_object is None:
+                        continue
+                    motor = self._create_motor(motor_name, driver, tmc_object)
+                    self._motors.append(motor)
+
+        pconfig = self._printer.lookup_object('configfile')
+        self.kinematics = pconfig.status_raw_config['printer']['kinematics']
+
+    # Create a Motor object with the given name, driver and TMC object
+    # and fill it with the relevant configuration and registers
+    def _create_motor(self, motor_name: str, driver: str, tmc_object: Any) -> Motor:
+        motor = Motor(motor_name)
+        motor.set_config('tmc', driver)
+        self._parse_klipper_config(motor, tmc_object)
+        self._parse_tmc_registers(motor, tmc_object)
+        return motor
+
+    def _parse_klipper_config(self, motor: Motor, tmc_object: Any) -> None:
+        # The TMCCommandHelper isn't a direct member of the TMC object... but we can still get it this way
+        tmc_cmdhelper = tmc_object.get_status.__self__
+
+        motor_currents = tmc_cmdhelper.current_helper.get_current()
+        motor.set_config('run_current', motor_currents[0])
+        motor.set_config('hold_current', motor_currents[1])
+
+        pconfig = self._printer.lookup_object('configfile')
+        motor.set_config('microsteps', int(pconfig.status_raw_config[motor.name]['microsteps']))
+
+        autotune_object = self._printer.lookup_object(f'autotune_tmc {motor.name}', None)
+        if autotune_object is not None:
+            motor.set_config('autotune_enabled', True)
+            motor.set_config('motor', autotune_object.motor)
+            motor.set_config('voltage', autotune_object.voltage)
+        else:
+            motor.set_config('autotune_enabled', False)
+
+    def _parse_tmc_registers(self, motor: Motor, tmc_object: Any) -> None:
+        # The TMCCommandHelper isn't a direct member of the TMC object... but we can still get it this way
+        tmc_cmdhelper = tmc_object.get_status.__self__
+
+        for register in RELEVANT_TMC_REGISTERS:
+            val = tmc_cmdhelper.fields.registers.get(register)
+            if (val is not None) and (register not in tmc_cmdhelper.read_registers):
+                # write-only register
+                fields_string = self._extract_register_values(tmc_cmdhelper, register, val)
+            elif register in tmc_cmdhelper.read_registers:
+                # readable register
+                val = tmc_cmdhelper.mcu_tmc.get_register(register)
+                if tmc_cmdhelper.read_translate is not None:
+                    register, val = tmc_cmdhelper.read_translate(register, val)
+                fields_string = self._extract_register_values(tmc_cmdhelper, register, val)
+
+            motor.set_register(register, fields_string)
+
+    def _extract_register_values(self, tmc_cmdhelper, register, val):
+        # Provide a dictionary of register values
+        reg_fields = tmc_cmdhelper.fields.all_fields.get(register, {})
+        reg_fields = sorted([(mask, name) for name, mask in reg_fields.items()])
+        fields = {}
+        for _, field_name in reg_fields:
+            field_value = tmc_cmdhelper.fields.get_field(field_name, val, register)
+            fields[field_name] = field_value
+        return fields
+
+    # Find and return the motor by its name
+    def get_motor(self, motor_name: str) -> Optional[Motor]:
+        for motor in self._motors:
+            if motor.name == motor_name:
+                return motor
+        return None
+
+    # Get all the motor list at once
+    def get_motors(self) -> List[Motor]:
+        return self._motors

shaketune/measurement/resonance_test.py

@@ -0,0 +1,50 @@
+#!/usr/bin/env python3
+
+# The logic in this file was "extracted" from Klipper's orignal resonance_tester.py file
+# Courtesy of Dmitry Butyugin <dmbutyugin@google.com> for the original implementation
+
+# This derive a bit from Klipper's implementation as there are two main changes:
+# 1. Original code doesn't use euclidean distance for the moves calculation with projection. The new approach implemented here
+#    ensures that the vector's total length remains constant (= L), regardless of the direction components. It's especially
+#    important when the direction vector involves combinations of movements along multiple axes like for the diagonal belt tests.
+# 2. Original code doesn't allow Z axis movement that was added here for later use
+
+import math
+
+from ..helpers.console_output import ConsoleOutput
+
+
+# This function is used to vibrate the toolhead in a specific axis direction
+# to test the resonance frequency of the printer and its components
+def vibrate_axis(toolhead, gcode, axis_direction, min_freq, max_freq, hz_per_sec, accel_per_hz):
+    freq = min_freq
+    X, Y, Z, E = toolhead.get_position()  # Get current position
+    sign = 1.0
+
+    while freq <= max_freq + 0.000001:
+        t_seg = 0.25 / freq  # Time segment for one vibration cycle
+        accel = accel_per_hz * freq  # Acceleration for each half-cycle
+        max_v = accel * t_seg  # Max velocity for each half-cycle
+        toolhead.cmd_M204(gcode.create_gcode_command('M204', 'M204', {'S': accel}))
+        L = 0.5 * accel * t_seg**2  # Distance for each half-cycle
+
+        # Calculate move points based on axis direction (X, Y and Z)
+        magnitude = math.sqrt(sum([component**2 for component in axis_direction]))
+        normalized_direction = tuple(component / magnitude for component in axis_direction)
+        dX, dY, dZ = normalized_direction[0] * L, normalized_direction[1] * L, normalized_direction[2] * L
+        nX = X + sign * dX
+        nY = Y + sign * dY
+        nZ = Z + sign * dZ
+
+        # Execute movement
+        toolhead.move([nX, nY, nZ, E], max_v)
+        toolhead.move([X, Y, Z, E], max_v)
+        sign *= -1
+
+        # Increase frequency for next cycle
+        old_freq = freq
+        freq += 2 * t_seg * hz_per_sec
+        if int(freq) > int(old_freq):
+            ConsoleOutput.print(f'Testing frequency: {freq:.0f} Hz')
+
+    toolhead.wait_moves()

shaketune/measurement/static_freq.py

@@ -0,0 +1,57 @@
+#!/usr/bin/env python3
+
+from ..helpers.common_func import AXIS_CONFIG
+from ..helpers.console_output import ConsoleOutput
+from .resonance_test import vibrate_axis
+
+
+def excitate_axis_at_freq(gcmd, config) -> None:
+    freq = gcmd.get_int('FREQUENCY', default=25, minval=1)
+    duration = gcmd.get_int('DURATION', default=10, minval=1)
+    accel_per_hz = gcmd.get_float('ACCEL_PER_HZ', default=None)
+    axis = gcmd.get('AXIS', default='x').lower()
+    feedrate_travel = gcmd.get_float('TRAVEL_SPEED', default=120.0, minval=20.0)
+    z_height = gcmd.get_float('Z_HEIGHT', default=None, minval=1)
+
+    axis_config = next((item for item in AXIS_CONFIG if item['axis'] == axis), None)
+    if axis_config is None:
+        gcmd.error('AXIS selection invalid. Should be either x, y, a or b!')
+
+    ConsoleOutput.print(f'Excitating {axis.upper()} axis at {freq}Hz for {duration} seconds')
+
+    printer = config.get_printer()
+    gcode = printer.lookup_object('gcode')
+    toolhead = printer.lookup_object('toolhead')
+    res_tester = printer.lookup_object('resonance_tester')
+    systime = printer.get_reactor().monotonic()
+
+    if accel_per_hz is None:
+        accel_per_hz = res_tester.test.accel_per_hz
+
+    # Move to the starting point
+    test_points = res_tester.test.get_start_test_points()
+    if len(test_points) > 1:
+        gcmd.error('Only one test point in the [resonance_tester] section is supported by Shake&Tune.')
+    if test_points[0] == (-1, -1, -1):
+        if z_height is None:
+            gcmd.error(
+                'Z_HEIGHT parameter is required if the test_point in [resonance_tester] section is set to -1,-1,-1'
+            )
+        # Use center of bed in case the test point in [resonance_tester] is set to -1,-1,-1
+        # This is usefull to get something automatic and is also used in the Klippain modular config
+        kin_info = toolhead.kin.get_status(systime)
+        mid_x = (kin_info['axis_minimum'].x + kin_info['axis_maximum'].x) / 2
+        mid_y = (kin_info['axis_minimum'].y + kin_info['axis_maximum'].y) / 2
+        point = (mid_x, mid_y, z_height)
+    else:
+        x, y, z = test_points[0]
+        if z_height is not None:
+            z = z_height
+        point = (x, y, z)
+
+    toolhead.manual_move(point, feedrate_travel)
+
+    min_freq = freq - 1
+    max_freq = freq + 1
+    hz_per_sec = 1 / (duration / 3)
+    vibrate_axis(toolhead, gcode, axis_config['direction'], min_freq, max_freq, hz_per_sec, accel_per_hz)

shaketune/measurement/vibrations_profile.py

@@ -0,0 +1,137 @@
+#!/usr/bin/env python3
+
+
+import math
+
+from ..helpers.console_output import ConsoleOutput
+from ..shaketune_thread import ShakeTuneThread
+from .accelerometer import Accelerometer
+from .motorsconfigparser import MotorsConfigParser
+
+MIN_SPEED = 2  # mm/s
+
+
+def create_vibrations_profile(gcmd, config, st_thread: ShakeTuneThread) -> None:
+    size = gcmd.get_float('SIZE', default=100.0, minval=50.0)
+    z_height = gcmd.get_float('Z_HEIGHT', default=20.0)
+    max_speed = gcmd.get_float('MAX_SPEED', default=200.0, minval=10.0)
+    speed_increment = gcmd.get_float('SPEED_INCREMENT', default=2.0, minval=1.0)
+    accel = gcmd.get_int('ACCEL', default=3000, minval=100)
+    feedrate_travel = gcmd.get_float('TRAVEL_SPEED', default=120.0, minval=20.0)
+    accel_chip = gcmd.get('ACCEL_CHIP', default=None)
+
+    if (size / (max_speed / 60)) < 0.25:
+        gcmd.error('The size of the movement is too small for the given speed! Increase SIZE or decrease MAX_SPEED!')
+
+    printer = config.get_printer()
+    gcode = printer.lookup_object('gcode')
+    toolhead = printer.lookup_object('toolhead')
+    input_shaper = printer.lookup_object('input_shaper', None)
+    systime = printer.get_reactor().monotonic()
+
+    # Check that input shaper is already configured
+    if input_shaper is None:
+        gcmd.error('Input shaper is not configured! Please run the shaper calibration macro first.')
+
+    motors_config_parser = MotorsConfigParser(config, motors=['stepper_x', 'stepper_y'])
+
+    if motors_config_parser.kinematics == 'cartesian' or motors_config_parser.kinematics == 'corexz':
+        # Cartesian motors are on X and Y axis directly, same for CoreXZ
+        main_angles = [0, 90]
+    elif motors_config_parser.kinematics == 'corexy':
+        # CoreXY motors are on A and B axis (45 and 135 degrees)
+        main_angles = [45, 135]
+    else:
+        gcmd.error(
+            'Only Cartesian and CoreXY kinematics are supported at the moment for the vibrations measurement tool!'
+        )
+    ConsoleOutput.print(f'{motors_config_parser.kinematics.upper()} kinematics mode')
+
+    toolhead_info = toolhead.get_status(systime)
+    old_accel = toolhead_info['max_accel']
+    old_mcr = toolhead_info['minimum_cruise_ratio']
+    old_sqv = toolhead_info['square_corner_velocity']
+
+    # set the wanted acceleration values
+    gcode.run_script_from_command(f'SET_VELOCITY_LIMIT ACCEL={accel} MINIMUM_CRUISE_RATIO=0 SQUARE_CORNER_VELOCITY=5.0')
+
+    kin_info = toolhead.kin.get_status(systime)
+    mid_x = (kin_info['axis_minimum'].x + kin_info['axis_maximum'].x) / 2
+    mid_y = (kin_info['axis_minimum'].y + kin_info['axis_maximum'].y) / 2
+    X, Y, _, E = toolhead.get_position()
+
+    # Going to the start position
+    toolhead.move([X, Y, z_height, E], feedrate_travel / 10)
+    toolhead.move([mid_x - 15, mid_y - 15, z_height, E], feedrate_travel)
+    toolhead.dwell(0.5)
+
+    nb_speed_samples = int((max_speed - MIN_SPEED) / speed_increment + 1)
+    for curr_angle in main_angles:
+        ConsoleOutput.print(f'-> Measuring angle: {curr_angle} degrees...')
+        radian_angle = math.radians(curr_angle)
+
+        # Find the best accelerometer chip for the current angle if not specified
+        if curr_angle == 0:
+            accel_axis = 'x'
+        elif curr_angle == 90:
+            accel_axis = 'y'
+        else:
+            accel_axis = 'xy'
+        if accel_chip is None:
+            accel_chip = Accelerometer.find_axis_accelerometer(printer, accel_axis)
+            if accel_chip is None:
+                gcmd.error(
+                    'No accelerometer specified for measurement! Multi-accelerometer configurations are not supported for this macro.'
+                )
+        accelerometer = Accelerometer(printer.lookup_object(accel_chip))
+
+        # Sweep the speed range to record the vibrations at different speeds
+        for curr_speed_sample in range(nb_speed_samples):
+            curr_speed = MIN_SPEED + curr_speed_sample * speed_increment
+            ConsoleOutput.print(f'Current speed: {curr_speed} mm/s')
+
+            # Reduce the segments length for the lower speed range (0-100mm/s). The minimum length is 1/3 of the SIZE and is gradually increased
+            # to the nominal SIZE at 100mm/s. No further size changes are made above this speed. The goal is to ensure that the print head moves
+            # enough to collect enough data for vibration analysis, without doing unnecessary distance to save time. At higher speeds, the full
+            # segments lengths are used because the head moves faster and travels more distance in the same amount of time and we want enough data
+            if curr_speed < 100:
+                segment_length_multiplier = 1 / 5 + 4 / 5 * curr_speed / 100
+            else:
+                segment_length_multiplier = 1
+
+            # Calculate angle coordinates using trigonometry and length multiplier and move to start point
+            dX = (size / 2) * math.cos(radian_angle) * segment_length_multiplier
+            dY = (size / 2) * math.sin(radian_angle) * segment_length_multiplier
+            toolhead.move([mid_x - dX, mid_y - dY, z_height, E], feedrate_travel)
+
+            # Adjust the number of back and forth movements based on speed to also save time on lower speed range
+            # 3 movements are done by default, reduced to 2 between 150-250mm/s and to 1 under 150mm/s.
+            movements = 3
+            if curr_speed < 150:
+                movements = 1
+            elif curr_speed < 250:
+                movements = 2
+
+            # Back and forth movements to record the vibrations at constant speed in both direction
+            accelerometer.start_measurement()
+            for _ in range(movements):
+                toolhead.move([mid_x + dX, mid_y + dY, z_height, E], curr_speed)
+                toolhead.move([mid_x - dX, mid_y - dY, z_height, E], curr_speed)
+            name = f'vib_an{curr_angle:.2f}sp{curr_speed:.2f}'.replace('.', '_')
+            accelerometer.stop_measurement(name)
+
+            toolhead.dwell(0.3)
+            toolhead.wait_moves()
+
+    # Restore the previous acceleration values
+    gcode.run_script_from_command(
+        f'SET_VELOCITY_LIMIT ACCEL={old_accel} MINIMUM_CRUISE_RATIO={old_mcr} SQUARE_CORNER_VELOCITY={old_sqv}'
+    )
+    toolhead.wait_moves()
+
+    # Run post-processing
+    ConsoleOutput.print('Machine vibrations profile generation...')
+    ConsoleOutput.print('This may take some time (5-8min)')
+    creator = st_thread.get_graph_creator()
+    creator.configure(motors_config_parser.kinematics, accel, motors_config_parser)
+    st_thread.run()

shaketune/post_processing/__init__.py

@@ -0,0 +1,7 @@
+#!/usr/bin/env python3
+
+from .graph_creator import AxesMapFinder as AxesMapFinder
+from .graph_creator import BeltsGraphCreator as BeltsGraphCreator
+from .graph_creator import GraphCreator as GraphCreator
+from .graph_creator import ShaperGraphCreator as ShaperGraphCreator
+from .graph_creator import VibrationsGraphCreator as VibrationsGraphCreator

shaketune/post_processing/analyze_axesmap.py

@@ -0,0 +1,155 @@
+#!/usr/bin/env python3
+
+######################################
+###### AXE_MAP DETECTION SCRIPT ######
+######################################
+# Written by Frix_x#0161 #
+
+import optparse
+
+import numpy as np
+from scipy.signal import butter, filtfilt
+
+from ..helpers.console_output import ConsoleOutput
+
+NUM_POINTS = 500
+
+
+######################################################################
+# Computation
+######################################################################
+
+
+def accel_signal_filter(data, cutoff=2, fs=100, order=5):
+    nyq = 0.5 * fs
+    normal_cutoff = cutoff / nyq
+    b, a = butter(order, normal_cutoff, btype='low', analog=False)
+    filtered_data = filtfilt(b, a, data)
+    filtered_data -= np.mean(filtered_data)
+    return filtered_data
+
+
+def find_first_spike(data):
+    min_index, max_index = np.argmin(data), np.argmax(data)
+    return ('-', min_index) if min_index < max_index else ('', max_index)
+
+
+def get_movement_vector(data, start_idx, end_idx):
+    if start_idx < end_idx:
+        vector = []
+        for i in range(3):
+            vector.append(np.mean(data[i][start_idx:end_idx], axis=0))
+        return vector
+    else:
+        return np.zeros(3)
+
+
+def angle_between(v1, v2):
+    v1_u = v1 / np.linalg.norm(v1)
+    v2_u = v2 / np.linalg.norm(v2)
+    return np.arccos(np.clip(np.dot(v1_u, v2_u), -1.0, 1.0))
+
+
+def compute_errors(filtered_data, spikes_sorted, accel_value, num_points):
+    # Get the movement start points in the correct order from the sorted bag of spikes
+    movement_starts = [spike[0][1] for spike in spikes_sorted]
+
+    # Theoretical unit vectors for X, Y, Z printer axes
+    printer_axes = {'x': np.array([1, 0, 0]), 'y': np.array([0, 1, 0]), 'z': np.array([0, 0, 1])}
+
+    alignment_errors = {}
+    sensitivity_errors = {}
+    for i, axis in enumerate(['x', 'y', 'z']):
+        movement_start = movement_starts[i]
+        movement_end = movement_start + num_points
+        movement_vector = get_movement_vector(filtered_data, movement_start, movement_end)
+        alignment_errors[axis] = angle_between(movement_vector, printer_axes[axis])
+
+        measured_accel_magnitude = np.linalg.norm(movement_vector)
+        if accel_value != 0:
+            sensitivity_errors[axis] = abs(measured_accel_magnitude - accel_value) / accel_value * 100
+        else:
+            sensitivity_errors[axis] = None
+
+    return alignment_errors, sensitivity_errors
+
+
+######################################################################
+# 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 axesmap_calibration(lognames, accel=None):
+    # Parse the raw data and get them ready for analysis
+    raw_datas = [parse_log(filename) for filename in lognames]
+    if len(raw_datas) > 1:
+        raise ValueError('Analysis of multiple CSV files at once is not possible with this script')
+
+    filtered_data = [accel_signal_filter(raw_datas[0][:, i + 1]) for i in range(3)]
+    spikes = [find_first_spike(filtered_data[i]) for i in range(3)]
+    spikes_sorted = sorted([(spikes[0], 'x'), (spikes[1], 'y'), (spikes[2], 'z')], key=lambda x: x[0][1])
+
+    # Using the previous variables to get the axes_map and errors
+    axes_map = ','.join([f'{spike[0][0]}{spike[1]}' for spike in spikes_sorted])
+    # alignment_error, sensitivity_error = compute_errors(filtered_data, spikes_sorted, accel, NUM_POINTS)
+
+    results = f'Be aware that this macro is experimental and has been known to sometimes produce incorrect results. Use it with caution and always check the results!\n'
+    results += f'Detected axes_map:\n  {axes_map}\n'
+
+    # TODO: work on this function that is currently not giving good results...
+    # results += "Accelerometer angle deviation:\n"
+    # for axis, angle in alignment_error.items():
+    #     angle_degrees = np.degrees(angle) # Convert radians to degrees
+    #     results += f"  {axis.upper()} axis: {angle_degrees:.2f} degrees\n"
+
+    # results += "Accelerometer sensitivity error:\n"
+    # for axis, error in sensitivity_error.items():
+    #     results += f"  {axis.upper()} axis: {error:.2f}%\n"
+
+    return results
+
+
+def main():
+    # Parse command-line arguments
+    usage = '%prog [options] <raw logs>'
+    opts = optparse.OptionParser(usage)
+    opts.add_option('-o', '--output', type='string', dest='output', default=None, help='filename of output graph')
+    opts.add_option(
+        '-a', '--accel', type='string', dest='accel', default=None, help='acceleration value used to do the movements'
+    )
+    options, args = opts.parse_args()
+    if len(args) < 1:
+        opts.error('No CSV file(s) to analyse')
+    if options.accel is None:
+        opts.error('You must specify the acceleration value used when generating the CSV file (option -a)')
+    try:
+        accel_value = float(options.accel)
+    except ValueError:
+        opts.error('Invalid acceleration value. It should be a numeric value.')
+
+    results = axesmap_calibration(args, accel_value)
+    ConsoleOutput.print(results)
+
+    if options.output is not None:
+        with open(options.output, 'w') as f:
+            f.write(results)
+
+
+if __name__ == '__main__':
+    main()

shaketune/post_processing/graph_belts.py

@@ -0,0 +1,562 @@
+#!/usr/bin/env python3
+
+#################################################
+######## CoreXY BELTS CALIBRATION SCRIPT ########
+#################################################
+# Written by Frix_x#0161 #
+
+import optparse
+import os
+from collections import namedtuple
+from datetime import datetime
+
+import matplotlib
+import matplotlib.colors
+import matplotlib.font_manager
+import matplotlib.pyplot as plt
+import matplotlib.ticker
+import numpy as np
+
+matplotlib.use('Agg')
+
+from ..helpers.common_func import detect_peaks, parse_log, setup_klipper_import
+from ..helpers.console_output import ConsoleOutput
+
+ALPHABET = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ'  # For paired peaks names
+
+PEAKS_DETECTION_THRESHOLD = 0.1  # Threshold to detect peaks in the PSD signal (10% of max)
+DC_MAX_PEAKS = 2  # Maximum ideal number of peaks
+DC_MAX_UNPAIRED_PEAKS_ALLOWED = 0  # No unpaired peaks are tolerated
+
+# Define the SignalData namedtuple
+SignalData = namedtuple('CalibrationData', ['freqs', 'psd', 'peaks', 'paired_peaks', 'unpaired_peaks'])
+
+KLIPPAIN_COLORS = {
+    'purple': '#70088C',
+    'orange': '#FF8D32',
+    'dark_purple': '#150140',
+    'dark_orange': '#F24130',
+    'red_pink': '#F2055C',
+}
+
+
+######################################################################
+# Computation of the PSD graph
+######################################################################
+
+
+# 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):
+    # Compute a dynamic detection threshold to filter and pair peaks efficiently
+    # even if the signal is very noisy (this get clipped to a maximum of 10Hz diff)
+    distances = []
+    for p1 in peaks1:
+        for p2 in peaks2:
+            distances.append(abs(freqs1[p1] - freqs2[p2]))
+    distances = np.array(distances)
+
+    median_distance = np.median(distances)
+    iqr = np.percentile(distances, 75) - np.percentile(distances, 25)
+
+    threshold = median_distance + 1.5 * iqr
+    threshold = min(threshold, 10)
+
+    # Pair the peaks using the dynamic thresold
+    paired_peaks = []
+    unpaired_peaks1 = list(peaks1)
+    unpaired_peaks2 = list(peaks2)
+
+    while unpaired_peaks1 and unpaired_peaks2:
+        min_distance = threshold + 1
+        pair = None
+
+        for p1 in unpaired_peaks1:
+            for p2 in unpaired_peaks2:
+                distance = abs(freqs1[p1] - freqs2[p2])
+                if distance < min_distance:
+                    min_distance = distance
+                    pair = (p1, p2)
+
+        if pair is None:  # No more pairs below the threshold
+            break
+
+        p1, p2 = pair
+        paired_peaks.append(((p1, freqs1[p1], psd1[p1]), (p2, freqs2[p2], psd2[p2])))
+        unpaired_peaks1.remove(p1)
+        unpaired_peaks2.remove(p2)
+
+    return paired_peaks, unpaired_peaks1, unpaired_peaks2
+
+
+######################################################################
+# Computation of the differential spectrogram
+######################################################################
+
+
+def compute_mhi(similarity_factor, signal1, signal2):
+    num_unpaired_peaks = len(signal1.unpaired_peaks) + len(signal2.unpaired_peaks)
+    num_paired_peaks = len(signal1.paired_peaks)
+    # Combine unpaired peaks from both signals, tagging each peak with its respective signal
+    combined_unpaired_peaks = [(peak, signal1) for peak in signal1.unpaired_peaks] + [
+        (peak, signal2) for peak in signal2.unpaired_peaks
+    ]
+    psd_highest_max = max(signal1.psd.max(), signal2.psd.max())
+
+    # Start with the similarity factor directly scaled to a percentage
+    mhi = similarity_factor
+
+    # Bonus for ideal number of total peaks (1 or 2)
+    if num_paired_peaks >= DC_MAX_PEAKS:
+        mhi *= DC_MAX_PEAKS / num_paired_peaks  # Reduce MHI if more than ideal number of peaks
+
+    # Penalty from unpaired peaks weighted by their amplitude relative to the maximum PSD amplitude
+    unpaired_peak_penalty = 0
+    if num_unpaired_peaks > DC_MAX_UNPAIRED_PEAKS_ALLOWED:
+        for peak, signal in combined_unpaired_peaks:
+            unpaired_peak_penalty += (signal.psd[peak] / psd_highest_max) * 30
+        mhi -= unpaired_peak_penalty
+
+    # Ensure the result lies between 0 and 100 by clipping the computed value
+    mhi = np.clip(mhi, 0, 100)
+
+    return mhi_lut(mhi)
+
+
+# LUT to transform the MHI into a textual value easy to understand for the users of the script
+def mhi_lut(mhi):
+    ranges = [
+        (70, 100, 'Excellent mechanical health'),
+        (55, 70, 'Good mechanical health'),
+        (45, 55, 'Acceptable mechanical health'),
+        (30, 45, 'Potential signs of a mechanical issue'),
+        (15, 30, 'Likely a mechanical issue'),
+        (0, 15, 'Mechanical issue detected'),
+    ]
+    mhi = np.clip(mhi, 1, 100)
+    for lower, upper, message in ranges:
+        if lower < mhi <= upper:
+            return message
+
+    return 'Unknown mechanical health'  # Should never happen
+
+
+######################################################################
+# Graphing
+######################################################################
+
+
+def plot_compare_frequency(ax, signal1, signal2, signal1_belt, signal2_belt, max_freq):
+    # Plot the two belts PSD signals
+    ax.plot(signal1.freqs, signal1.psd, label='Belt ' + signal1_belt, color=KLIPPAIN_COLORS['purple'])
+    ax.plot(signal2.freqs, signal2.psd, label='Belt ' + signal2_belt, color=KLIPPAIN_COLORS['orange'])
+
+    # Trace and annotate the peaks on the graph
+    paired_peak_count = 0
+    unpaired_peak_count = 0
+    offsets_table_data = []
+
+    for _, (peak1, peak2) in enumerate(signal1.paired_peaks):
+        label = ALPHABET[paired_peak_count]
+        amplitude_offset = abs(
+            ((signal2.psd[peak2[0]] - signal1.psd[peak1[0]]) / max(signal1.psd[peak1[0]], signal2.psd[peak2[0]])) * 100
+        )
+        frequency_offset = abs(signal2.freqs[peak2[0]] - signal1.freqs[peak1[0]])
+        offsets_table_data.append([f'Peaks {label}', f'{frequency_offset:.1f} Hz', f'{amplitude_offset:.1f} %'])
+
+        ax.plot(signal1.freqs[peak1[0]], signal1.psd[peak1[0]], 'x', color='black')
+        ax.plot(signal2.freqs[peak2[0]], signal2.psd[peak2[0]], 'x', color='black')
+        ax.plot(
+            [signal1.freqs[peak1[0]], signal2.freqs[peak2[0]]],
+            [signal1.psd[peak1[0]], signal2.psd[peak2[0]]],
+            ':',
+            color='gray',
+        )
+
+        ax.annotate(
+            label + '1',
+            (signal1.freqs[peak1[0]], signal1.psd[peak1[0]]),
+            textcoords='offset points',
+            xytext=(8, 5),
+            ha='left',
+            fontsize=13,
+            color='black',
+        )
+        ax.annotate(
+            label + '2',
+            (signal2.freqs[peak2[0]], signal2.psd[peak2[0]]),
+            textcoords='offset points',
+            xytext=(8, 5),
+            ha='left',
+            fontsize=13,
+            color='black',
+        )
+        paired_peak_count += 1
+
+    for peak in signal1.unpaired_peaks:
+        ax.plot(signal1.freqs[peak], signal1.psd[peak], 'x', color='black')
+        ax.annotate(
+            str(unpaired_peak_count + 1),
+            (signal1.freqs[peak], signal1.psd[peak]),
+            textcoords='offset points',
+            xytext=(8, 5),
+            ha='left',
+            fontsize=13,
+            color='red',
+            weight='bold',
+        )
+        unpaired_peak_count += 1
+
+    for peak in signal2.unpaired_peaks:
+        ax.plot(signal2.freqs[peak], signal2.psd[peak], 'x', color='black')
+        ax.annotate(
+            str(unpaired_peak_count + 1),
+            (signal2.freqs[peak], signal2.psd[peak]),
+            textcoords='offset points',
+            xytext=(8, 5),
+            ha='left',
+            fontsize=13,
+            color='red',
+            weight='bold',
+        )
+        unpaired_peak_count += 1
+
+    # Add estimated similarity to the graph
+    ax2 = ax.twinx()  # To split the legends in two box
+    ax2.yaxis.set_visible(False)
+    ax2.plot([], [], ' ', label=f'Number of unpaired peaks: {unpaired_peak_count}')
+
+    # Setting axis parameters, grid and graph title
+    ax.set_xlabel('Frequency (Hz)')
+    ax.set_xlim([0, max_freq])
+    ax.set_ylabel('Power spectral density')
+    psd_highest_max = max(signal1.psd.max(), signal2.psd.max())
+    ax.set_ylim([0, psd_highest_max * 1.1])
+
+    ax.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
+    ax.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
+    ax.ticklabel_format(axis='y', style='scientific', scilimits=(0, 0))
+    ax.grid(which='major', color='grey')
+    ax.grid(which='minor', color='lightgrey')
+    fontP = matplotlib.font_manager.FontProperties()
+    fontP.set_size('small')
+    ax.set_title(
+        'Belts frequency profiles',
+        fontsize=14,
+        color=KLIPPAIN_COLORS['dark_orange'],
+        weight='bold',
+    )
+
+    # Print the table of offsets ontop of the graph below the original legend (upper right)
+    if len(offsets_table_data) > 0:
+        columns = [
+            '',
+            'Frequency delta',
+            'Amplitude delta',
+        ]
+        offset_table = ax.table(
+            cellText=offsets_table_data,
+            colLabels=columns,
+            bbox=[0.66, 0.79, 0.33, 0.15],
+            loc='upper right',
+            cellLoc='center',
+        )
+        offset_table.auto_set_font_size(False)
+        offset_table.set_fontsize(8)
+        offset_table.auto_set_column_width([0, 1, 2])
+        offset_table.set_zorder(100)
+        cells = [key for key in offset_table.get_celld().keys()]
+        for cell in cells:
+            offset_table[cell].set_facecolor('white')
+            offset_table[cell].set_alpha(0.6)
+
+    ax.legend(loc='upper left', prop=fontP)
+    ax2.legend(loc='upper right', prop=fontP)
+
+    return
+
+
+# Compute quantile-quantile plot to compare the two belts
+def plot_versus_belts(ax, common_freqs, signal1, signal2, interp_psd1, interp_psd2, signal1_belt, signal2_belt):
+    ax.set_title('Cross-belts comparison plot', fontsize=14, color=KLIPPAIN_COLORS['dark_orange'], weight='bold')
+
+    max_psd = max(np.max(interp_psd1), np.max(interp_psd2))
+    ideal_line = np.linspace(0, max_psd * 1.1, 500)
+    green_boundary = ideal_line + (0.35 * max_psd * np.exp(-ideal_line / (0.6 * max_psd)))
+    ax.fill_betweenx(ideal_line, ideal_line, green_boundary, color='green', alpha=0.15)
+    ax.fill_between(ideal_line, ideal_line, green_boundary, color='green', alpha=0.15, label='Good zone')
+    ax.plot(
+        ideal_line,
+        ideal_line,
+        '--',
+        label='Ideal line',
+        color='red',
+        linewidth=2,
+    )
+
+    ax.plot(interp_psd1, interp_psd2, color='dimgrey', marker='o', markersize=1.5)
+    ax.fill_betweenx(interp_psd2, interp_psd1, color=KLIPPAIN_COLORS['red_pink'], alpha=0.1)
+
+    paired_peak_count = 0
+    unpaired_peak_count = 0
+
+    for _, (peak1, peak2) in enumerate(signal1.paired_peaks):
+        label = ALPHABET[paired_peak_count]
+        freq1 = signal1.freqs[peak1[0]]
+        freq2 = signal2.freqs[peak2[0]]
+        nearest_idx1 = np.argmin(np.abs(common_freqs - freq1))
+        nearest_idx2 = np.argmin(np.abs(common_freqs - freq2))
+
+        if nearest_idx1 == nearest_idx2:
+            psd1_peak_value = interp_psd1[nearest_idx1]
+            psd2_peak_value = interp_psd2[nearest_idx1]
+            ax.plot(psd1_peak_value, psd2_peak_value, marker='o', color='black', markersize=7)
+            ax.annotate(
+                f'{label}1/{label}2',
+                (psd1_peak_value, psd2_peak_value),
+                textcoords='offset points',
+                xytext=(-7, 7),
+                fontsize=13,
+                color='black',
+            )
+        else:
+            psd1_peak_value = interp_psd1[nearest_idx1]
+            psd1_on_peak = interp_psd1[nearest_idx2]
+            psd2_peak_value = interp_psd2[nearest_idx2]
+            psd2_on_peak = interp_psd2[nearest_idx1]
+            ax.plot(psd1_on_peak, psd2_peak_value, marker='o', color=KLIPPAIN_COLORS['orange'], markersize=7)
+            ax.plot(psd1_peak_value, psd2_on_peak, marker='o', color=KLIPPAIN_COLORS['purple'], markersize=7)
+            ax.annotate(
+                f'{label}1',
+                (psd1_peak_value, psd2_on_peak),
+                textcoords='offset points',
+                xytext=(0, 7),
+                fontsize=13,
+                color='black',
+            )
+            ax.annotate(
+                f'{label}2',
+                (psd1_on_peak, psd2_peak_value),
+                textcoords='offset points',
+                xytext=(0, 7),
+                fontsize=13,
+                color='black',
+            )
+        paired_peak_count += 1
+
+    for _, peak_index in enumerate(signal1.unpaired_peaks):
+        freq1 = signal1.freqs[peak_index]
+        freq2 = signal2.freqs[peak_index]
+        nearest_idx1 = np.argmin(np.abs(common_freqs - freq1))
+        nearest_idx2 = np.argmin(np.abs(common_freqs - freq2))
+        psd1_peak_value = interp_psd1[nearest_idx1]
+        psd2_peak_value = interp_psd2[nearest_idx1]
+        ax.plot(psd1_peak_value, psd2_peak_value, marker='o', color=KLIPPAIN_COLORS['purple'], markersize=7)
+        ax.annotate(
+            str(unpaired_peak_count + 1),
+            (psd1_peak_value, psd2_peak_value),
+            textcoords='offset points',
+            fontsize=13,
+            weight='bold',
+            color=KLIPPAIN_COLORS['red_pink'],
+            xytext=(0, 7),
+        )
+        unpaired_peak_count += 1
+
+    for _, peak_index in enumerate(signal2.unpaired_peaks):
+        freq1 = signal1.freqs[peak_index]
+        freq2 = signal2.freqs[peak_index]
+        nearest_idx1 = np.argmin(np.abs(common_freqs - freq1))
+        nearest_idx2 = np.argmin(np.abs(common_freqs - freq2))
+        psd1_peak_value = interp_psd1[nearest_idx1]
+        psd2_peak_value = interp_psd2[nearest_idx1]
+        ax.plot(psd1_peak_value, psd2_peak_value, marker='o', color=KLIPPAIN_COLORS['orange'], markersize=7)
+        ax.annotate(
+            str(unpaired_peak_count + 1),
+            (psd1_peak_value, psd2_peak_value),
+            textcoords='offset points',
+            fontsize=13,
+            weight='bold',
+            color=KLIPPAIN_COLORS['red_pink'],
+            xytext=(0, 7),
+        )
+        unpaired_peak_count += 1
+
+    ax.set_xlabel(f'Belt {signal1_belt}')
+    ax.set_ylabel(f'Belt {signal2_belt}')
+    ax.set_xlim([0, max_psd * 1.1])
+    ax.set_ylim([0, max_psd * 1.1])
+
+    ax.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
+    ax.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
+    ax.ticklabel_format(axis='y', style='scientific', scilimits=(0, 0))
+    ax.grid(which='major', color='grey')
+    ax.grid(which='minor', color='lightgrey')
+
+    fontP = matplotlib.font_manager.FontProperties()
+    fontP.set_size('medium')
+    ax.legend(loc='upper left', prop=fontP)
+
+    return
+
+
+######################################################################
+# Custom tools
+######################################################################
+
+
+# Original Klipper function to get the PSD data of a raw accelerometer signal
+def compute_signal_data(data, max_freq):
+    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_DETECTION_THRESHOLD * psd.max())
+
+    return SignalData(freqs=freqs, psd=psd, peaks=peaks, paired_peaks=None, unpaired_peaks=None)
+
+
+######################################################################
+# Startup and main routines
+######################################################################
+
+
+def belts_calibration(
+    lognames, kinematics, klipperdir='~/klipper', max_freq=200.0, accel_per_hz=None, st_version='unknown'
+):
+    global shaper_calibrate
+    shaper_calibrate = setup_klipper_import(klipperdir)
+
+    # Parse data from the log files while ignoring CSV in the wrong format
+    datas = [data for data in (parse_log(fn) for fn in lognames) if data is not None]
+    if len(datas) > 2:
+        raise ValueError('Incorrect number of .csv files used (this function needs exactly two files to compare them)!')
+
+    # Get the belts name for the legend to avoid putting the full file name
+    belt_info = {'A': ' (axis 1,-1)', 'B': ' (axis 1, 1)'}
+    signal1_belt = (lognames[0].split('/')[-1]).split('_')[-1][0]
+    signal2_belt = (lognames[1].split('/')[-1]).split('_')[-1][0]
+    signal1_belt += belt_info.get(signal1_belt, '')
+    signal2_belt += belt_info.get(signal2_belt, '')
+
+    # 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)
+    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)
+
+    # Re-interpolate the PSD signals to a common frequency range to be able to plot them one against the other point by point
+    common_freqs = np.linspace(0, max_freq, 500)
+    interp_psd1 = np.interp(common_freqs, signal1.freqs, signal1.psd)
+    interp_psd2 = np.interp(common_freqs, signal2.freqs, signal2.psd)
+
+    # Calculating R^2 to y=x line to compute the similarity between the two belts
+    ss_res = np.sum((interp_psd2 - interp_psd1) ** 2)
+    ss_tot = np.sum((interp_psd2 - np.mean(interp_psd2)) ** 2)
+    similarity_factor = (1 - (ss_res / ss_tot)) * 100
+    ConsoleOutput.print(f'Belts estimated similarity: {similarity_factor:.1f}%')
+
+    # mhi = compute_mhi(similarity_factor, num_peaks, num_unpaired_peaks)
+    mhi = compute_mhi(similarity_factor, signal1, signal2)
+    ConsoleOutput.print(f'[experimental] Mechanical health: {mhi}')
+
+    fig, ((ax1, ax3)) = plt.subplots(
+        1,
+        2,
+        gridspec_kw={
+            'width_ratios': [5, 3],
+            'bottom': 0.080,
+            'top': 0.840,
+            'left': 0.050,
+            'right': 0.985,
+            'hspace': 0.166,
+            'wspace': 0.138,
+        },
+    )
+    fig.set_size_inches(15, 7)
+
+    # Add title
+    title_line1 = 'RELATIVE BELTS CALIBRATION TOOL'
+    fig.text(
+        0.060, 0.947, title_line1, ha='left', va='bottom', fontsize=20, color=KLIPPAIN_COLORS['purple'], weight='bold'
+    )
+    try:
+        filename = lognames[0].split('/')[-1]
+        dt = datetime.strptime(f"{filename.split('_')[1]} {filename.split('_')[2]}", '%Y%m%d %H%M%S')
+        title_line2 = dt.strftime('%x %X')
+        if kinematics is not None:
+            title_line2 += ' -- ' + kinematics.upper() + ' kinematics'
+    except Exception:
+        ConsoleOutput.print(
+            'Warning: CSV filenames look to be different than expected (%s , %s)' % (lognames[0], lognames[1])
+        )
+        title_line2 = lognames[0].split('/')[-1] + ' / ' + lognames[1].split('/')[-1]
+    fig.text(0.060, 0.939, title_line2, ha='left', va='top', fontsize=16, color=KLIPPAIN_COLORS['dark_purple'])
+
+    # We add the estimated similarity and the MHI value to the title only if the kinematics is CoreXY
+    # as it make no sense to compute these values for other kinematics that doesn't have paired belts
+    if kinematics == 'corexy':
+        title_line3 = f'| Estimated similarity: {similarity_factor:.1f}%'
+        title_line4 = f'| {mhi} (experimental)'
+        fig.text(0.55, 0.985, title_line3, ha='left', va='top', fontsize=14, color=KLIPPAIN_COLORS['dark_purple'])
+        fig.text(0.55, 0.950, title_line4, ha='left', va='top', fontsize=14, color=KLIPPAIN_COLORS['dark_purple'])
+
+    # Add the accel_per_hz value to the title
+    title_line5 = f'| Accel per Hz used: {accel_per_hz} mm/s²/Hz'
+    fig.text(0.55, 0.915, title_line5, ha='left', va='top', fontsize=14, color=KLIPPAIN_COLORS['dark_purple'])
+
+    # Plot the graphs
+    plot_compare_frequency(ax1, signal1, signal2, signal1_belt, signal2_belt, max_freq)
+    plot_versus_belts(ax3, common_freqs, signal1, signal2, interp_psd1, interp_psd2, signal1_belt, signal2_belt)
+
+    # Adding a small Klippain logo to the top left corner of the figure
+    ax_logo = fig.add_axes([0.001, 0.894, 0.105, 0.105], 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
+    if st_version != 'unknown':
+        fig.text(0.995, 0.985, st_version, ha='right', va='bottom', fontsize=8, color=KLIPPAIN_COLORS['purple'])
+
+    return fig
+
+
+def main():
+    # Parse command-line arguments
+    usage = '%prog [options] <raw logs>'
+    opts = optparse.OptionParser(usage)
+    opts.add_option('-o', '--output', type='string', dest='output', default=None, help='filename of output graph')
+    opts.add_option('-f', '--max_freq', type='float', default=200.0, help='maximum frequency to graph')
+    opts.add_option('--accel_per_hz', type='float', default=None, help='accel_per_hz used during the measurement')
+    opts.add_option(
+        '-k', '--klipper_dir', type='string', dest='klipperdir', default='~/klipper', help='main klipper directory'
+    )
+    opts.add_option(
+        '-m',
+        '--kinematics',
+        type='string',
+        dest='kinematics',
+        help='machine kinematics configuration',
+    )
+    options, args = opts.parse_args()
+    if len(args) < 1:
+        opts.error('Incorrect number of arguments')
+    if options.output is None:
+        opts.error('You must specify an output file.png to use the script (option -o)')
+
+    fig = belts_calibration(
+        args, options.kinematics, options.klipperdir, options.max_freq, options.accel_per_hz, 'unknown'
+    )
+    fig.savefig(options.output, dpi=150)
+
+
+if __name__ == '__main__':
+    main()

shaketune/post_processing/graph_creator.py

@@ -0,0 +1,279 @@
+#!/usr/bin/env python3
+
+import abc
+import re
+import shutil
+import tarfile
+from datetime import datetime
+from pathlib import Path
+from typing import Callable, Optional
+
+from matplotlib.figure import Figure
+
+from ..helpers import filemanager as fm
+from ..helpers.console_output import ConsoleOutput
+from ..measurement.motorsconfigparser import MotorsConfigParser
+from ..shaketune_config import ShakeTuneConfig
+from .analyze_axesmap import axesmap_calibration
+from .graph_belts import belts_calibration
+from .graph_shaper import shaper_calibration
+from .graph_vibrations import vibrations_profile
+
+
+class GraphCreator(abc.ABC):
+    def __init__(self, config: ShakeTuneConfig):
+        self._config = config
+
+        self._graph_date = datetime.now().strftime('%Y%m%d_%H%M%S')
+        self._version = ShakeTuneConfig.get_git_version()
+
+        self._type = None
+        self._folder = None
+
+    def _setup_folder(self, graph_type: str) -> None:
+        self._type = graph_type
+        self._folder = self._config.get_results_folder(graph_type)
+
+    def _move_and_prepare_files(
+        self,
+        glob_pattern: str,
+        min_files_required: Optional[int] = None,
+        custom_name_func: Optional[Callable[[Path], str]] = None,
+    ) -> list[Path]:
+        tmp_path = Path('/tmp')
+        globbed_files = list(tmp_path.glob(glob_pattern))
+
+        # If min_files_required is not set, use the number of globbed files as the minimum
+        min_files_required = min_files_required or len(globbed_files)
+
+        if not globbed_files:
+            raise FileNotFoundError(f'no CSV files found in the /tmp folder to create the {self._type} graphs!')
+        if len(globbed_files) < min_files_required:
+            raise FileNotFoundError(f'{min_files_required} CSV files are needed to create the {self._type} graphs!')
+
+        lognames = []
+        for filename in sorted(globbed_files, key=lambda f: f.stat().st_mtime, reverse=True)[:min_files_required]:
+            fm.wait_file_ready(filename)
+            custom_name = custom_name_func(filename) if custom_name_func else filename.name
+            new_file = self._folder / f'{self._type}_{self._graph_date}_{custom_name}.csv'
+            # shutil.move() is needed to move the file across filesystems (mainly for BTT CB1 Pi default OS image)
+            shutil.move(filename, new_file)
+            lognames.append(new_file)
+        return lognames
+
+    def _save_figure_and_cleanup(self, fig: Figure, lognames: list[Path], axis_label: Optional[str] = None) -> None:
+        axis_suffix = f'_{axis_label}' if axis_label else ''
+        png_filename = self._folder / f'{self._type}_{self._graph_date}{axis_suffix}.png'
+        fig.savefig(png_filename, dpi=self._config.dpi)
+
+        if self._config.keep_csv:
+            self._archive_files(lognames)
+        else:
+            self._remove_files(lognames)
+
+    def _archive_files(self, _: list[Path]) -> None:
+        return
+
+    def _remove_files(self, lognames: list[Path]) -> None:
+        for csv in lognames:
+            csv.unlink(missing_ok=True)
+
+    def get_type(self) -> str:
+        return self._type
+
+    @abc.abstractmethod
+    def create_graph(self) -> None:
+        pass
+
+    @abc.abstractmethod
+    def clean_old_files(self, keep_results: int) -> None:
+        pass
+
+
+class BeltsGraphCreator(GraphCreator):
+    def __init__(self, config: ShakeTuneConfig):
+        super().__init__(config)
+
+        self._kinematics = None
+        self._accel_per_hz = None
+
+        self._setup_folder('belts')
+
+    def configure(self, kinematics: str = None, accel_per_hz: float = None) -> None:
+        self._kinematics = kinematics
+        self._accel_per_hz = accel_per_hz
+
+    def create_graph(self) -> None:
+        lognames = self._move_and_prepare_files(
+            glob_pattern='shaketune-belt_*.csv',
+            min_files_required=2,
+            custom_name_func=lambda f: f.stem.split('_')[1].upper(),
+        )
+        fig = belts_calibration(
+            lognames=[str(path) for path in lognames],
+            klipperdir=str(self._config.klipper_folder),
+            accel_per_hz=self._accel_per_hz,
+            st_version=self._version,
+        )
+        self._save_figure_and_cleanup(fig, lognames)
+
+    def clean_old_files(self, keep_results: int = 3) -> None:
+        # Get all PNG files in the directory as a list of Path objects
+        files = sorted(self._folder.glob('*.png'), key=lambda f: f.stat().st_mtime, reverse=True)
+
+        if len(files) <= keep_results:
+            return  # No need to delete any files
+
+        # Delete the older files
+        for old_file in files[keep_results:]:
+            file_date = '_'.join(old_file.stem.split('_')[1:3])
+            for suffix in ['A', 'B']:
+                csv_file = self._folder / f'belts_{file_date}_{suffix}.csv'
+                csv_file.unlink(missing_ok=True)
+            old_file.unlink()
+
+
+class ShaperGraphCreator(GraphCreator):
+    def __init__(self, config: ShakeTuneConfig):
+        super().__init__(config)
+
+        self._max_smoothing = None
+        self._scv = None
+
+        self._setup_folder('shaper')
+
+    def configure(self, scv: float, max_smoothing: float = None, accel_per_hz: float = None) -> None:
+        self._scv = scv
+        self._max_smoothing = max_smoothing
+        self._accel_per_hz = accel_per_hz
+
+    def create_graph(self) -> None:
+        if not self._scv:
+            raise ValueError('scv must be set to create the input shaper graph!')
+
+        lognames = self._move_and_prepare_files(
+            glob_pattern='shaketune-axis_*.csv',
+            min_files_required=1,
+            custom_name_func=lambda f: f.stem.split('_')[1].upper(),
+        )
+        fig = shaper_calibration(
+            lognames=[str(path) for path in lognames],
+            klipperdir=str(self._config.klipper_folder),
+            max_smoothing=self._max_smoothing,
+            scv=self._scv,
+            accel_per_hz=self._accel_per_hz,
+            st_version=self._version,
+        )
+        self._save_figure_and_cleanup(fig, lognames, lognames[0].stem.split('_')[-1])
+
+    def clean_old_files(self, keep_results: int = 3) -> None:
+        # Get all PNG files in the directory as a list of Path objects
+        files = sorted(self._folder.glob('*.png'), key=lambda f: f.stat().st_mtime, reverse=True)
+
+        if len(files) <= 2 * keep_results:
+            return  # No need to delete any files
+
+        # Delete the older files
+        for old_file in files[2 * keep_results :]:
+            csv_file = old_file.with_suffix('.csv')
+            csv_file.unlink(missing_ok=True)
+            old_file.unlink()
+
+
+class VibrationsGraphCreator(GraphCreator):
+    def __init__(self, config: ShakeTuneConfig):
+        super().__init__(config)
+
+        self._kinematics = None
+        self._accel = None
+        self._motors = None
+
+        self._setup_folder('vibrations')
+
+    def configure(self, kinematics: str, accel: float, motor_config_parser: MotorsConfigParser) -> None:
+        self._kinematics = kinematics
+        self._accel = accel
+        self._motors = motor_config_parser.get_motors()
+
+    def _archive_files(self, lognames: list[Path]) -> None:
+        tar_path = self._folder / f'{self._type}_{self._graph_date}.tar.gz'
+        with tarfile.open(tar_path, 'w:gz') as tar:
+            for csv_file in lognames:
+                tar.add(csv_file, arcname=csv_file.name, recursive=False)
+                csv_file.unlink()
+
+    def create_graph(self) -> None:
+        if not self._accel or not self._kinematics:
+            raise ValueError('accel, chip_name and kinematics must be set to create the vibrations profile graph!')
+
+        lognames = self._move_and_prepare_files(
+            glob_pattern='shaketune-vib_*.csv',
+            min_files_required=None,
+            custom_name_func=lambda f: re.search(r'shaketune-vib_(.*?)_\d{8}_\d{6}', f.name).group(1),
+        )
+        fig = vibrations_profile(
+            lognames=[str(path) for path in lognames],
+            klipperdir=str(self._config.klipper_folder),
+            kinematics=self._kinematics,
+            accel=self._accel,
+            st_version=self._version,
+            motors=self._motors,
+        )
+        self._save_figure_and_cleanup(fig, lognames)
+
+    def clean_old_files(self, keep_results: int = 3) -> None:
+        # Get all PNG files in the directory as a list of Path objects
+        files = sorted(self._folder.glob('*.png'), key=lambda f: f.stat().st_mtime, reverse=True)
+
+        if len(files) <= keep_results:
+            return  # No need to delete any files
+
+        # Delete the older files
+        for old_file in files[keep_results:]:
+            old_file.unlink()
+            tar_file = old_file.with_suffix('.tar.gz')
+            tar_file.unlink(missing_ok=True)
+
+
+class AxesMapFinder(GraphCreator):
+    def __init__(self, config: ShakeTuneConfig):
+        super().__init__(config)
+
+        self._graph_date = datetime.now().strftime('%Y%m%d_%H%M%S')
+        self._type = 'axesmap'
+        self._folder = config.get_results_folder()
+
+        self._accel = None
+
+    def configure(self, accel: int) -> None:
+        self._accel = accel
+
+    def find_axesmap(self) -> None:
+        tmp_folder = Path('/tmp')
+        globbed_files = list(tmp_folder.glob('shaketune-axemap_*.csv'))
+
+        if not globbed_files:
+            raise FileNotFoundError('no CSV files found in the /tmp folder to find the axes map!')
+
+        # Find the CSV files with the latest timestamp and process it
+        logname = sorted(globbed_files, key=lambda f: f.stat().st_mtime, reverse=True)[0]
+        results = axesmap_calibration(
+            lognames=[str(logname)],
+            accel=self._accel,
+        )
+        ConsoleOutput.print(results)
+
+        result_filename = self._folder / f'{self._type}_{self._graph_date}.txt'
+        with result_filename.open('w') as f:
+            f.write(results)
+
+    # While the AxesMapFinder doesn't directly create a graph, we need to implement this
+    # method to allow using it seemlessly like all the other GraphCreator objects
+    def create_graph(self) -> None:
+        self.find_axesmap()
+
+    def clean_old_files(self, keep_results: int) -> None:
+        tmp_folder = Path('/tmp')
+        globbed_files = list(tmp_folder.glob('shaketune-axemap_*.csv'))
+        for csv_file in globbed_files:
+            csv_file.unlink()

shaketune/post_processing/graph_shaper.py

@@ -0,0 +1,435 @@
+#!/usr/bin/env python3
+
+#################################################
+######## INPUT SHAPER CALIBRATION SCRIPT ########
+#################################################
+# Derived from the calibrate_shaper.py official Klipper script
+# Copyright (C) 2020  Dmitry Butyugin <dmbutyugin@google.com>
+# Copyright (C) 2020  Kevin O'Connor <kevin@koconnor.net>
+# Highly modified and improved by Frix_x#0161 #
+
+import optparse
+import os
+from datetime import datetime
+
+import matplotlib
+import matplotlib.font_manager
+import matplotlib.pyplot as plt
+import matplotlib.ticker
+import numpy as np
+
+matplotlib.use('Agg')
+
+from ..helpers.common_func import (
+    compute_mechanical_parameters,
+    compute_spectrogram,
+    detect_peaks,
+    parse_log,
+    setup_klipper_import,
+)
+from ..helpers.console_output import ConsoleOutput
+
+PEAKS_DETECTION_THRESHOLD = 0.05
+PEAKS_EFFECT_THRESHOLD = 0.12
+SPECTROGRAM_LOW_PERCENTILE_FILTER = 5
+MAX_SMOOTHING = 0.1
+
+KLIPPAIN_COLORS = {
+    'purple': '#70088C',
+    'orange': '#FF8D32',
+    'dark_purple': '#150140',
+    'dark_orange': '#F24130',
+    'red_pink': '#F2055C',
+}
+
+
+######################################################################
+# Computation
+######################################################################
+
+
+# Find the best shaper parameters using Klipper's official algorithm selection with
+# a proper precomputed damping ratio (zeta) and using the configured printer SQV value
+def calibrate_shaper(datas, max_smoothing, scv, max_freq):
+    helper = shaper_calibrate.ShaperCalibrate(printer=None)
+    calibration_data = helper.process_accelerometer_data(datas)
+    calibration_data.normalize_to_frequencies()
+
+    fr, zeta, _, _ = compute_mechanical_parameters(calibration_data.psd_sum, calibration_data.freq_bins)
+
+    # If the damping ratio computation fail, we use Klipper default value instead
+    if zeta is None:
+        zeta = 0.1
+
+    compat = False
+    try:
+        shaper, all_shapers = helper.find_best_shaper(
+            calibration_data,
+            shapers=None,
+            damping_ratio=zeta,
+            scv=scv,
+            shaper_freqs=None,
+            max_smoothing=max_smoothing,
+            test_damping_ratios=None,
+            max_freq=max_freq,
+            logger=ConsoleOutput.print,
+        )
+    except TypeError:
+        ConsoleOutput.print(
+            '[WARNING] You seem to be using an older version of Klipper that is not compatible with all the latest Shake&Tune features!'
+        )
+        ConsoleOutput.print(
+            'Shake&Tune now runs in compatibility mode: be aware that the results may be slightly off, since the real damping ratio cannot be used to create the filter recommendations'
+        )
+        compat = True
+        shaper, all_shapers = helper.find_best_shaper(calibration_data, max_smoothing, ConsoleOutput.print)
+
+    ConsoleOutput.print(
+        '\n-> Recommended shaper is %s @ %.1f Hz (when using a square corner velocity of %.1f and a damping ratio of %.3f)'
+        % (shaper.name.upper(), shaper.freq, scv, zeta)
+    )
+
+    return shaper.name, all_shapers, calibration_data, fr, zeta, compat
+
+
+######################################################################
+# Graphing
+######################################################################
+
+
+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')
+
+    ax.set_xlabel('Frequency (Hz)')
+    ax.set_xlim([0, max_freq])
+    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', 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')
+
+    ax.xaxis.set_minor_locator(matplotlib.ticker.MultipleLocator(5))
+    ax.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
+    ax.ticklabel_format(axis='y', style='scientific', scilimits=(0, 0))
+    ax.grid(which='major', color='grey')
+    ax.grid(which='minor', color='lightgrey')
+
+    ax2 = ax.twinx()
+    ax2.yaxis.set_visible(False)
+
+    lowvib_shaper_vibrs = float('inf')
+    lowvib_shaper = None
+    lowvib_shaper_freq = None
+    lowvib_shaper_accel = 0
+
+    # Draw the shappers curves and add their specific parameters in the legend
+    # This adds also a way to find the best shaper with a low level of vibrations (with a resonable level of smoothing)
+    for shaper in shapers:
+        shaper_max_accel = round(shaper.max_accel / 100.0) * 100.0
+        label = '%s (%.1f Hz, vibr=%.1f%%, sm~=%.2f, accel<=%.f)' % (
+            shaper.name.upper(),
+            shaper.freq,
+            shaper.vibrs * 100.0,
+            shaper.smoothing,
+            shaper_max_accel,
+        )
+        ax2.plot(freqs, shaper.vals, label=label, linestyle='dotted')
+
+        # Get the performance shaper
+        if shaper.name == performance_shaper:
+            performance_shaper_freq = shaper.freq
+            performance_shaper_vibr = shaper.vibrs * 100.0
+            performance_shaper_vals = shaper.vals
+
+        # Get the low vibration shaper
+        if (
+            shaper.vibrs * 100 < lowvib_shaper_vibrs
+            or (shaper.vibrs * 100 == lowvib_shaper_vibrs and shaper_max_accel > lowvib_shaper_accel)
+        ) and shaper.smoothing < MAX_SMOOTHING:
+            lowvib_shaper_accel = shaper_max_accel
+            lowvib_shaper = shaper.name
+            lowvib_shaper_freq = shaper.freq
+            lowvib_shaper_vibrs = shaper.vibrs * 100
+            lowvib_shaper_vals = shaper.vals
+
+    # User recommendations are added to the legend: one is Klipper's original suggestion that is usually good for performances
+    # and the other one is the custom "low vibration" recommendation that looks for a suitable shaper that doesn't have excessive
+    # smoothing and that have a lower vibration level. If both recommendation are the same shaper, or if no suitable "low
+    # vibration" shaper is found, then only a single line as the "best shaper" recommendation is added to the legend
+    if (
+        lowvib_shaper is not None
+        and lowvib_shaper != performance_shaper
+        and lowvib_shaper_vibrs <= performance_shaper_vibr
+    ):
+        ax2.plot(
+            [],
+            [],
+            ' ',
+            label='Recommended performance shaper: %s @ %.1f Hz'
+            % (performance_shaper.upper(), performance_shaper_freq),
+        )
+        ax.plot(
+            freqs, psd * performance_shaper_vals, label='With %s applied' % (performance_shaper.upper()), color='cyan'
+        )
+        ax2.plot(
+            [],
+            [],
+            ' ',
+            label='Recommended low vibrations shaper: %s @ %.1f Hz' % (lowvib_shaper.upper(), lowvib_shaper_freq),
+        )
+        ax.plot(freqs, psd * lowvib_shaper_vals, label='With %s applied' % (lowvib_shaper.upper()), color='lime')
+    else:
+        ax2.plot(
+            [],
+            [],
+            ' ',
+            label='Recommended best shaper: %s @ %.1f Hz' % (performance_shaper.upper(), performance_shaper_freq),
+        )
+        ax.plot(
+            freqs, psd * performance_shaper_vals, label='With %s applied' % (performance_shaper.upper()), color='cyan'
+        )
+
+    # And the estimated damping ratio is finally added at the end of the legend
+    ax2.plot([], [], ' ', label='Estimated damping ratio (ζ): %.3f' % (zeta))
+
+    # Draw the detected peaks and name them
+    # This also draw the detection threshold and warning threshold (aka "effect zone")
+    ax.plot(peaks_freqs, psd[peaks], 'x', color='black', markersize=8)
+    for idx, peak in enumerate(peaks):
+        if psd[peak] > peaks_threshold[1]:
+            fontcolor = 'red'
+            fontweight = 'bold'
+        else:
+            fontcolor = 'black'
+            fontweight = 'normal'
+        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_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
+
+
+# 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, 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)
+
+    # 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.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
+
+
+######################################################################
+# Startup and main routines
+######################################################################
+
+
+def shaper_calibration(
+    lognames,
+    klipperdir='~/klipper',
+    max_smoothing=None,
+    scv=5.0,
+    max_freq=200.0,
+    accel_per_hz=None,
+    st_version='unknown',
+):
+    global shaper_calibrate
+    shaper_calibrate = setup_klipper_import(klipperdir)
+
+    # Parse data from the log files while ignoring CSV in the wrong format
+    datas = [data for data in (parse_log(fn) for fn in lognames) if data is not None]
+    if len(datas) > 1:
+        ConsoleOutput.print('Warning: incorrect number of .csv files detected. Only the first one will be used!')
+
+    # Compute shapers, PSD outputs and spectrogram
+    performance_shaper, shapers, calibration_data, fr, zeta, compat = calibrate_shaper(
+        datas[0], max_smoothing, scv, max_freq
+    )
+    pdata, bins, t = compute_spectrogram(datas[0])
+    del datas
+
+    # 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])
+    ConsoleOutput.print(
+        '\nPeaks 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 a title with some test info
+    title_line1 = 'INPUT SHAPER CALIBRATION TOOL'
+    fig.text(
+        0.12, 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]}', '%Y%m%d %H%M%S')
+        title_line2 = dt.strftime('%x %X') + ' -- ' + filename_parts[3].upper().split('.')[0] + ' axis'
+        if compat:
+            title_line3 = '| Older Klipper version detected, damping ratio'
+            title_line4 = '| and SCV are not used for filter recommendations!'
+            title_line5 = f'| Accel per Hz used: {accel_per_hz} mm/s²/Hz' if accel_per_hz is not None else ''
+        else:
+            title_line3 = f'| Square corner velocity: {scv} mm/s'
+            title_line4 = f'| Max allowed smoothing: {max_smoothing}'
+            title_line5 = f'| Accel per Hz used: {accel_per_hz} mm/s²/Hz' if accel_per_hz is not None else ''
+    except Exception:
+        ConsoleOutput.print('Warning: CSV filename look to be different than expected (%s)' % (lognames[0]))
+        title_line2 = lognames[0].split('/')[-1]
+        title_line3 = ''
+        title_line4 = ''
+        title_line5 = ''
+    fig.text(0.12, 0.957, title_line2, ha='left', va='top', fontsize=16, color=KLIPPAIN_COLORS['dark_purple'])
+    fig.text(0.58, 0.965, title_line3, ha='left', va='top', fontsize=10, color=KLIPPAIN_COLORS['dark_purple'])
+    fig.text(0.58, 0.951, title_line4, ha='left', va='top', fontsize=10, color=KLIPPAIN_COLORS['dark_purple'])
+    fig.text(0.58, 0.919, title_line5, ha='left', va='top', fontsize=10, color=KLIPPAIN_COLORS['dark_purple'])
+
+    # Plot the graphs
+    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.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
+    if st_version != 'unknown':
+        fig.text(0.995, 0.985, st_version, ha='right', va='bottom', fontsize=8, color=KLIPPAIN_COLORS['purple'])
+
+    return fig
+
+
+def main():
+    # Parse command-line arguments
+    usage = '%prog [options] <logs>'
+    opts = optparse.OptionParser(usage)
+    opts.add_option('-o', '--output', type='string', dest='output', default=None, help='filename of output graph')
+    opts.add_option('-f', '--max_freq', type='float', default=200.0, help='maximum frequency to graph')
+    opts.add_option('-s', '--max_smoothing', type='float', default=None, help='maximum shaper smoothing to allow')
+    opts.add_option(
+        '--scv', '--square_corner_velocity', type='float', dest='scv', default=5.0, help='square corner velocity'
+    )
+    opts.add_option('--accel_per_hz', type='float', default=None, help='accel_per_hz used during the measurement')
+    opts.add_option(
+        '-k', '--klipper_dir', type='string', dest='klipperdir', default='~/klipper', help='main klipper directory'
+    )
+    options, args = opts.parse_args()
+    if len(args) < 1:
+        opts.error('Incorrect number of arguments')
+    if options.output is None:
+        opts.error('You must specify an output file.png to use the script (option -o)')
+    if options.max_smoothing is not None and options.max_smoothing < 0.05:
+        opts.error('Too small max_smoothing specified (must be at least 0.05)')
+
+    fig = shaper_calibration(
+        args, options.klipperdir, options.max_smoothing, options.scv, options.max_freq, options.accel_per_hz, 'unknown'
+    )
+    fig.savefig(options.output, dpi=150)
+
+
+if __name__ == '__main__':
+    main()

shaketune/post_processing/graph_vibrations.py

@@ -0,0 +1,841 @@
+#!/usr/bin/env python3
+
+##################################################
+#### DIRECTIONAL VIBRATIONS PLOTTING SCRIPT ######
+##################################################
+# Written by Frix_x#0161 #
+
+import math
+import optparse
+import os
+import re
+from collections import defaultdict
+from datetime import datetime
+
+import matplotlib
+import matplotlib.font_manager
+import matplotlib.gridspec
+import matplotlib.pyplot as plt
+import matplotlib.ticker
+import numpy as np
+
+matplotlib.use('Agg')
+
+from ..helpers.common_func import (
+    compute_mechanical_parameters,
+    detect_peaks,
+    identify_low_energy_zones,
+    parse_log,
+    setup_klipper_import,
+)
+from ..helpers.console_output import ConsoleOutput
+
+PEAKS_DETECTION_THRESHOLD = 0.05
+PEAKS_RELATIVE_HEIGHT_THRESHOLD = 0.04
+CURVE_SIMILARITY_SIGMOID_K = 0.5
+SPEEDS_VALLEY_DETECTION_THRESHOLD = 0.7  # Lower is more sensitive
+SPEEDS_AROUND_PEAK_DELETION = 3  # to delete +-3mm/s around a peak
+ANGLES_VALLEY_DETECTION_THRESHOLD = 1.1  # Lower is more sensitive
+
+KLIPPAIN_COLORS = {
+    'purple': '#70088C',
+    'orange': '#FF8D32',
+    'dark_purple': '#150140',
+    'dark_orange': '#F24130',
+    'red_pink': '#F2055C',
+}
+
+
+######################################################################
+# Computation
+######################################################################
+
+
+# Call to the official Klipper input shaper object to do the PSD computation
+def calc_freq_response(data):
+    helper = shaper_calibrate.ShaperCalibrate(printer=None)
+    return helper.process_accelerometer_data(data)
+
+
+# Calculate motor frequency profiles based on the measured Power Spectral Density (PSD) measurements for the machine kinematics
+# main angles and then create a global motor profile as a weighted average (from their own vibrations) of all calculated profiles
+def compute_motor_profiles(freqs, psds, all_angles_energy, measured_angles=None, energy_amplification_factor=2):
+    if measured_angles is None:
+        measured_angles = [0, 90]
+
+    motor_profiles = {}
+    weighted_sum_profiles = np.zeros_like(freqs)
+    total_weight = 0
+    conv_filter = np.ones(20) / 20
+
+    # Creating the PSD motor profiles for each angles
+    for angle in measured_angles:
+        # Calculate the sum of PSDs for the current angle and then convolve
+        sum_curve = np.sum(np.array([psds[angle][speed] for speed in psds[angle]]), axis=0)
+        motor_profiles[angle] = np.convolve(sum_curve / len(psds[angle]), conv_filter, mode='same')
+
+        # Calculate weights
+        angle_energy = (
+            all_angles_energy[angle] ** energy_amplification_factor
+        )  # First weighting factor is based on the total vibrations of the machine at the specified angle
+        curve_area = (
+            np.trapz(motor_profiles[angle], freqs) ** energy_amplification_factor
+        )  # Additional weighting factor is based on the area under the current motor profile at this specified angle
+        total_angle_weight = angle_energy * curve_area
+
+        # Update weighted sum profiles to get the global motor profile
+        weighted_sum_profiles += motor_profiles[angle] * total_angle_weight
+        total_weight += total_angle_weight
+
+    # Creating a global average motor profile that is the weighted average of all the PSD motor profiles
+    global_motor_profile = weighted_sum_profiles / total_weight if total_weight != 0 else weighted_sum_profiles
+
+    return motor_profiles, global_motor_profile
+
+
+# Since it was discovered that there is no non-linear mixing in the stepper "steps" vibrations, instead of measuring
+# the effects of each speeds at each angles, this function simplify it by using only the main motors axes (X/Y for Cartesian
+# printers and A/B for CoreXY) measurements and project each points on the [0,360] degrees range using trigonometry
+# to "sum" the vibration impact of each axis at every points of the generated spectrogram. The result is very similar at the end.
+def compute_dir_speed_spectrogram(measured_speeds, data, kinematics='cartesian', measured_angles=None):
+    if measured_angles is None:
+        measured_angles = [0, 90]
+
+    # We want to project the motor vibrations measured on their own axes on the [0, 360] range
+    spectrum_angles = np.linspace(0, 360, 720)  # One point every 0.5 degrees
+    spectrum_speeds = np.linspace(min(measured_speeds), max(measured_speeds), len(measured_speeds) * 6)
+    spectrum_vibrations = np.zeros((len(spectrum_angles), len(spectrum_speeds)))
+
+    def get_interpolated_vibrations(data, speed, speeds):
+        idx = np.clip(np.searchsorted(speeds, speed, side='left'), 1, len(speeds) - 1)
+        lower_speed = speeds[idx - 1]
+        upper_speed = speeds[idx]
+        lower_vibrations = data.get(lower_speed, 0)
+        upper_vibrations = data.get(upper_speed, 0)
+        return lower_vibrations + (speed - lower_speed) * (upper_vibrations - lower_vibrations) / (
+            upper_speed - lower_speed
+        )
+
+    # Precompute trigonometric values and constant before the loop
+    angle_radians = np.deg2rad(spectrum_angles)
+    cos_vals = np.cos(angle_radians)
+    sin_vals = np.sin(angle_radians)
+    sqrt_2_inv = 1 / math.sqrt(2)
+
+    # Compute the spectrum vibrations for each angle and speed combination
+    for target_angle_idx, (cos_val, sin_val) in enumerate(zip(cos_vals, sin_vals)):
+        for target_speed_idx, target_speed in enumerate(spectrum_speeds):
+            if kinematics == 'cartesian':
+                speed_1 = np.abs(target_speed * cos_val)
+                speed_2 = np.abs(target_speed * sin_val)
+            elif kinematics == 'corexy':
+                speed_1 = np.abs(target_speed * (cos_val + sin_val) * sqrt_2_inv)
+                speed_2 = np.abs(target_speed * (cos_val - sin_val) * sqrt_2_inv)
+
+            vibrations_1 = get_interpolated_vibrations(data[measured_angles[0]], speed_1, measured_speeds)
+            vibrations_2 = get_interpolated_vibrations(data[measured_angles[1]], speed_2, measured_speeds)
+            spectrum_vibrations[target_angle_idx, target_speed_idx] = vibrations_1 + vibrations_2
+
+    return spectrum_angles, spectrum_speeds, spectrum_vibrations
+
+
+def compute_angle_powers(spectrogram_data):
+    angles_powers = np.trapz(spectrogram_data, axis=1)
+
+    # Since we want to plot it on a continuous polar plot later on, we need to append parts of
+    # the array to start and end of it to smooth transitions when doing the convolution
+    # and get the same value at modulo 360. Then we return the array without the extras
+    extended_angles_powers = np.concatenate([angles_powers[-9:], angles_powers, angles_powers[:9]])
+    convolved_extended = np.convolve(extended_angles_powers, np.ones(15) / 15, mode='same')
+
+    return convolved_extended[9:-9]
+
+
+def compute_speed_powers(spectrogram_data, smoothing_window=15):
+    min_values = np.amin(spectrogram_data, axis=0)
+    max_values = np.amax(spectrogram_data, axis=0)
+    var_values = np.var(spectrogram_data, axis=0)
+
+    # rescale the variance to the same range as max_values to plot it on the same graph
+    var_values = var_values / var_values.max() * max_values.max()
+
+    # Create a vibration metric that is the product of the max values and the variance to quantify the best
+    # speeds that have at the same time a low global energy level that is also consistent at every angles
+    vibration_metric = max_values * var_values
+
+    # utility function to pad and smooth the data avoiding edge effects
+    conv_filter = np.ones(smoothing_window) / smoothing_window
+    window = int(smoothing_window / 2)
+
+    def pad_and_smooth(data):
+        data_padded = np.pad(data, (window,), mode='edge')
+        smoothed_data = np.convolve(data_padded, conv_filter, mode='valid')
+        return smoothed_data
+
+    # Stack the arrays and apply padding and smoothing in batch
+    data_arrays = np.stack([min_values, max_values, var_values, vibration_metric])
+    smoothed_arrays = np.array([pad_and_smooth(data) for data in data_arrays])
+
+    return smoothed_arrays
+
+
+# Function that filter and split the good_speed ranges. The goal is to remove some zones around
+# additional detected small peaks in order to suppress them if there is a peak, even if it's low,
+# that's probably due to a crossing in the motor resonance pattern that still need to be removed
+def filter_and_split_ranges(all_speeds, good_speeds, peak_speed_indices, deletion_range):
+    # Process each range to filter out and split based on peak indices
+    filtered_good_speeds = []
+    for start, end, energy in good_speeds:
+        start_speed, end_speed = all_speeds[start], all_speeds[end]
+        # Identify peaks that intersect with the current speed range
+        intersecting_peaks_indices = [
+            idx for speed, idx in peak_speed_indices.items() if start_speed <= speed <= end_speed
+        ]
+
+        if not intersecting_peaks_indices:
+            filtered_good_speeds.append((start, end, energy))
+        else:
+            intersecting_peaks_indices.sort()
+            current_start = start
+
+            for peak_index in intersecting_peaks_indices:
+                before_peak_end = max(current_start, peak_index - deletion_range)
+                if current_start < before_peak_end:
+                    filtered_good_speeds.append((current_start, before_peak_end, energy))
+                current_start = peak_index + deletion_range + 1
+
+            if current_start < end:
+                filtered_good_speeds.append((current_start, end, energy))
+
+    # Sorting by start point once and then merge overlapping ranges
+    sorted_ranges = sorted(filtered_good_speeds, key=lambda x: x[0])
+    merged_ranges = [sorted_ranges[0]]
+
+    for current in sorted_ranges[1:]:
+        last_merged_start, last_merged_end, last_merged_energy = merged_ranges[-1]
+        if current[0] <= last_merged_end:
+            new_end = max(last_merged_end, current[1])
+            new_energy = min(last_merged_energy, current[2])
+            merged_ranges[-1] = (last_merged_start, new_end, new_energy)
+        else:
+            merged_ranges.append(current)
+
+    return merged_ranges
+
+
+# This function allow the computation of a symmetry score that reflect the spectrogram apparent symmetry between
+# measured axes on both the shape of the signal and the energy level consistency across both side of the signal
+def compute_symmetry_analysis(all_angles, spectrogram_data, measured_angles=None):
+    if measured_angles is None:
+        measured_angles = [0, 90]
+
+    total_spectrogram_angles = len(all_angles)
+    half_spectrogram_angles = total_spectrogram_angles // 2
+
+    # Extend the spectrogram by adding half to the beginning (in order to not get an out of bounds error later)
+    extended_spectrogram = np.concatenate((spectrogram_data[-half_spectrogram_angles:], spectrogram_data), axis=0)
+
+    # Calculate the split index directly within the slicing
+    midpoint_angle = np.mean(measured_angles)
+    split_index = int(midpoint_angle * (total_spectrogram_angles / 360) + half_spectrogram_angles)
+    half_segment_length = half_spectrogram_angles // 2
+
+    # Slice out the two segments of the spectrogram and flatten them for comparison
+    segment_1_flattened = extended_spectrogram[split_index - half_segment_length : split_index].flatten()
+    segment_2_flattened = extended_spectrogram[split_index : split_index + half_segment_length].flatten()
+
+    # Compute the correlation coefficient between the two segments of spectrogram
+    correlation = np.corrcoef(segment_1_flattened, segment_2_flattened)[0, 1]
+    percentage_correlation_biased = (100 * np.power(correlation, 0.75)) + 10
+
+    return np.clip(0, 100, percentage_correlation_biased)
+
+
+######################################################################
+# Graphing
+######################################################################
+
+
+def plot_angle_profile_polar(ax, angles, angles_powers, low_energy_zones, symmetry_factor):
+    angles_radians = np.deg2rad(angles)
+
+    ax.set_title('Polar angle energy profile', fontsize=14, color=KLIPPAIN_COLORS['dark_orange'], weight='bold')
+    ax.set_theta_zero_location('E')
+    ax.set_theta_direction(1)
+
+    ax.plot(angles_radians, angles_powers, color=KLIPPAIN_COLORS['purple'], zorder=5)
+    ax.fill(angles_radians, angles_powers, color=KLIPPAIN_COLORS['purple'], alpha=0.3)
+    ax.set_xlim([0, np.deg2rad(360)])
+    ymax = angles_powers.max() * 1.05
+    ax.set_ylim([0, ymax])
+    ax.set_thetagrids([theta * 15 for theta in range(360 // 15)])
+
+    ax.text(
+        0,
+        0,
+        f'Symmetry: {symmetry_factor:.1f}%',
+        ha='center',
+        va='center',
+        color=KLIPPAIN_COLORS['red_pink'],
+        fontsize=12,
+        fontweight='bold',
+        zorder=6,
+    )
+
+    for _, (start, end, _) in enumerate(low_energy_zones):
+        ax.axvline(
+            angles_radians[start],
+            angles_powers[start] / ymax,
+            color=KLIPPAIN_COLORS['red_pink'],
+            linestyle='dotted',
+            linewidth=1.5,
+        )
+        ax.axvline(
+            angles_radians[end],
+            angles_powers[end] / ymax,
+            color=KLIPPAIN_COLORS['red_pink'],
+            linestyle='dotted',
+            linewidth=1.5,
+        )
+        ax.fill_between(
+            angles_radians[start:end], angles_powers[start:end], angles_powers.max() * 1.05, color='green', alpha=0.2
+        )
+
+    ax.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
+    ax.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
+    ax.grid(which='major', color='grey')
+    ax.grid(which='minor', color='lightgrey')
+
+    # Polar plot doesn't follow the gridspec margin, so we adjust it manually here
+    pos = ax.get_position()
+    new_pos = [pos.x0 - 0.01, pos.y0 - 0.01, pos.width, pos.height]
+    ax.set_position(new_pos)
+
+    return
+
+
+def plot_global_speed_profile(
+    ax,
+    all_speeds,
+    sp_min_energy,
+    sp_max_energy,
+    sp_variance_energy,
+    vibration_metric,
+    num_peaks,
+    peaks,
+    low_energy_zones,
+):
+    ax.set_title('Global speed energy profile', fontsize=14, color=KLIPPAIN_COLORS['dark_orange'], weight='bold')
+    ax.set_xlabel('Speed (mm/s)')
+    ax.set_ylabel('Energy')
+    ax2 = ax.twinx()
+    ax2.yaxis.set_visible(False)
+
+    ax.plot(all_speeds, sp_min_energy, label='Minimum', color=KLIPPAIN_COLORS['dark_purple'], zorder=5)
+    ax.plot(all_speeds, sp_max_energy, label='Maximum', color=KLIPPAIN_COLORS['purple'], zorder=5)
+    ax.plot(all_speeds, sp_variance_energy, label='Variance', color=KLIPPAIN_COLORS['orange'], zorder=5, linestyle='--')
+    ax2.plot(
+        all_speeds,
+        vibration_metric,
+        label=f'Vibration metric ({num_peaks} bad peaks)',
+        color=KLIPPAIN_COLORS['red_pink'],
+        zorder=5,
+    )
+
+    ax.set_xlim([all_speeds.min(), all_speeds.max()])
+    ax.set_ylim([0, sp_max_energy.max() * 1.15])
+
+    y2min = -(vibration_metric.max() * 0.025)
+    y2max = vibration_metric.max() * 1.07
+    ax2.set_ylim([y2min, y2max])
+
+    if peaks is not None and len(peaks) > 0:
+        ax2.plot(all_speeds[peaks], vibration_metric[peaks], 'x', color='black', markersize=8, zorder=10)
+        for idx, peak in enumerate(peaks):
+            ax2.annotate(
+                f'{idx+1}',
+                (all_speeds[peak], vibration_metric[peak]),
+                textcoords='offset points',
+                xytext=(5, 5),
+                fontweight='bold',
+                ha='left',
+                fontsize=13,
+                color=KLIPPAIN_COLORS['red_pink'],
+                zorder=10,
+            )
+
+    for idx, (start, end, _) in enumerate(low_energy_zones):
+        # ax2.axvline(all_speeds[start], color=KLIPPAIN_COLORS['red_pink'], linestyle='dotted', linewidth=1.5, zorder=8)
+        # ax2.axvline(all_speeds[end], color=KLIPPAIN_COLORS['red_pink'], linestyle='dotted', linewidth=1.5, zorder=8)
+        ax2.fill_between(
+            all_speeds[start:end],
+            y2min,
+            vibration_metric[start:end],
+            color='green',
+            alpha=0.2,
+            label=f'Zone {idx+1}: {all_speeds[start]:.1f} to {all_speeds[end]:.1f} mm/s',
+        )
+
+    ax.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
+    ax.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
+    ax.grid(which='major', color='grey')
+    ax.grid(which='minor', color='lightgrey')
+
+    fontP = matplotlib.font_manager.FontProperties()
+    fontP.set_size('small')
+    ax.legend(loc='upper left', prop=fontP)
+    ax2.legend(loc='upper right', prop=fontP)
+
+    return
+
+
+def plot_angular_speed_profiles(ax, speeds, angles, spectrogram_data, kinematics='cartesian'):
+    ax.set_title('Angular speed energy profiles', fontsize=14, color=KLIPPAIN_COLORS['dark_orange'], weight='bold')
+    ax.set_xlabel('Speed (mm/s)')
+    ax.set_ylabel('Energy')
+
+    # Define mappings for labels and colors to simplify plotting commands
+    angle_settings = {
+        0: ('X (0 deg)', 'purple', 10),
+        90: ('Y (90 deg)', 'dark_purple', 5),
+        45: ('A (45 deg)' if kinematics == 'corexy' else '45 deg', 'orange', 10),
+        135: ('B (135 deg)' if kinematics == 'corexy' else '135 deg', 'dark_orange', 5),
+    }
+
+    # Plot each angle using settings from the dictionary
+    for angle, (label, color, zorder) in angle_settings.items():
+        idx = np.searchsorted(angles, angle, side='left')
+        ax.plot(speeds, spectrogram_data[idx], label=label, color=KLIPPAIN_COLORS[color], zorder=zorder)
+
+    ax.set_xlim([speeds.min(), speeds.max()])
+    max_value = max(spectrogram_data[angle].max() for angle in [0, 45, 90, 135])
+    ax.set_ylim([0, max_value * 1.1])
+
+    ax.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
+    ax.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
+    ax.grid(which='major', color='grey')
+    ax.grid(which='minor', color='lightgrey')
+
+    fontP = matplotlib.font_manager.FontProperties()
+    fontP.set_size('small')
+    ax.legend(loc='upper right', prop=fontP)
+
+    return
+
+
+def plot_motor_profiles(ax, freqs, main_angles, motor_profiles, global_motor_profile, max_freq):
+    ax.set_title('Motor frequency profile', fontsize=14, color=KLIPPAIN_COLORS['dark_orange'], weight='bold')
+    ax.set_ylabel('Energy')
+    ax.set_xlabel('Frequency (Hz)')
+
+    ax2 = ax.twinx()
+    ax2.yaxis.set_visible(False)
+
+    # Global weighted average motor profile
+    ax.plot(freqs, global_motor_profile, label='Combined', color=KLIPPAIN_COLORS['purple'], zorder=5)
+    max_value = global_motor_profile.max()
+
+    # Mapping of angles to axis names
+    angle_settings = {0: 'X', 90: 'Y', 45: 'A', 135: 'B'}
+
+    # And then plot the motor profiles at each measured angles
+    for angle in main_angles:
+        profile_max = motor_profiles[angle].max()
+        if profile_max > max_value:
+            max_value = profile_max
+        label = f'{angle_settings[angle]} ({angle} deg)' if angle in angle_settings else f'{angle} deg'
+        ax.plot(freqs, motor_profiles[angle], linestyle='--', label=label, zorder=2)
+
+    ax.set_xlim([0, max_freq])
+    ax.set_ylim([0, max_value * 1.1])
+    ax.ticklabel_format(axis='y', style='scientific', scilimits=(0, 0))
+
+    # Then add the motor resonance peak to the graph and print some infos about it
+    motor_fr, motor_zeta, motor_res_idx, lowfreq_max = compute_mechanical_parameters(global_motor_profile, freqs, 30)
+    if lowfreq_max:
+        ConsoleOutput.print(
+            '[WARNING] There are a lot of low frequency vibrations that can alter the readings. This is probably due to the test being performed at too high an acceleration!'
+        )
+        ConsoleOutput.print(
+            'Try lowering the ACCEL value and/or increasing the SIZE value before restarting the macro to ensure that only constant speeds are being recorded and that the dynamic behavior of the machine is not affecting the measurements'
+        )
+    if motor_zeta is not None:
+        ConsoleOutput.print(
+            'Motors have a main resonant frequency at %.1fHz with an estimated damping ratio of %.3f'
+            % (motor_fr, motor_zeta)
+        )
+    else:
+        ConsoleOutput.print(
+            'Motors have a main resonant frequency at %.1fHz but it was impossible to estimate a damping ratio.'
+            % (motor_fr)
+        )
+
+    ax.plot(freqs[motor_res_idx], global_motor_profile[motor_res_idx], 'x', color='black', markersize=10)
+    ax.annotate(
+        'R',
+        (freqs[motor_res_idx], global_motor_profile[motor_res_idx]),
+        textcoords='offset points',
+        xytext=(15, 5),
+        ha='right',
+        fontsize=14,
+        color=KLIPPAIN_COLORS['red_pink'],
+        weight='bold',
+    )
+
+    ax2.plot([], [], ' ', label='Motor resonant frequency (ω0): %.1fHz' % (motor_fr))
+    if motor_zeta is not None:
+        ax2.plot([], [], ' ', label='Motor damping ratio (ζ): %.3f' % (motor_zeta))
+    else:
+        ax2.plot([], [], ' ', label='No damping ratio computed')
+
+    ax.xaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
+    ax.yaxis.set_minor_locator(matplotlib.ticker.AutoMinorLocator())
+    ax.grid(which='major', color='grey')
+    ax.grid(which='minor', color='lightgrey')
+
+    fontP = matplotlib.font_manager.FontProperties()
+    fontP.set_size('small')
+    ax.legend(loc='upper left', prop=fontP)
+    ax2.legend(loc='upper right', prop=fontP)
+
+    return
+
+
+def plot_vibration_spectrogram_polar(ax, angles, speeds, spectrogram_data):
+    angles_radians = np.radians(angles)
+
+    # Assuming speeds defines the radial distance from the center, we need to create a meshgrid
+    # for both angles and speeds to map the spectrogram data onto a polar plot correctly
+    radius, theta = np.meshgrid(speeds, angles_radians)
+
+    ax.set_title(
+        'Polar vibrations heatmap', fontsize=14, color=KLIPPAIN_COLORS['dark_orange'], weight='bold', va='bottom'
+    )
+    ax.set_theta_zero_location('E')
+    ax.set_theta_direction(1)
+
+    ax.pcolormesh(theta, radius, spectrogram_data, norm=matplotlib.colors.LogNorm(), cmap='inferno', shading='auto')
+    ax.set_thetagrids([theta * 15 for theta in range(360 // 15)])
+    ax.tick_params(axis='y', which='both', colors='white', labelsize='medium')
+    ax.set_ylim([0, max(speeds)])
+
+    # Polar plot doesn't follow the gridspec margin, so we adjust it manually here
+    pos = ax.get_position()
+    new_pos = [pos.x0 - 0.01, pos.y0 - 0.01, pos.width, pos.height]
+    ax.set_position(new_pos)
+
+    return
+
+
+def plot_vibration_spectrogram(ax, angles, speeds, spectrogram_data, peaks):
+    ax.set_title('Vibrations heatmap', fontsize=14, color=KLIPPAIN_COLORS['dark_orange'], weight='bold')
+    ax.set_xlabel('Speed (mm/s)')
+    ax.set_ylabel('Angle (deg)')
+
+    ax.imshow(
+        spectrogram_data,
+        norm=matplotlib.colors.LogNorm(),
+        cmap='inferno',
+        aspect='auto',
+        extent=[speeds[0], speeds[-1], angles[0], angles[-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 and len(peaks) > 0:
+        for idx, peak in enumerate(peaks):
+            ax.axvline(speeds[peak], color='cyan', linewidth=0.75)
+            ax.annotate(
+                f'Peak {idx+1}',
+                (speeds[peak], angles[-1] * 0.9),
+                textcoords='data',
+                color='cyan',
+                rotation=90,
+                fontsize=10,
+                verticalalignment='top',
+                horizontalalignment='right',
+            )
+
+    return
+
+
+def plot_motor_config_txt(fig, motors, differences):
+    motor_details = [(motors[0], 'X motor'), (motors[1], 'Y motor')]
+
+    distance = 0.12
+    if motors[0].get_config('autotune_enabled'):
+        distance = 0.27
+        config_blocks = [
+            f"| {lbl}: {mot.get_config('motor').upper()} on {mot.get_config('tmc').upper()} @ {mot.get_config('voltage'):0.1f}V {mot.get_config('run_current'):0.2f}A - {mot.get_config('microsteps')}usteps"
+            for mot, lbl in motor_details
+        ]
+        config_blocks.append('| TMC Autotune enabled')
+    else:
+        config_blocks = [
+            f"| {lbl}: {mot.get_config('tmc').upper()} @ {mot.get_config('run_current'):0.2f}A - {mot.get_config('microsteps')}usteps"
+            for mot, lbl in motor_details
+        ]
+        config_blocks.append('| TMC Autotune not detected')
+
+    for idx, block in enumerate(config_blocks):
+        fig.text(
+            0.41, 0.990 - 0.015 * idx, block, ha='left', va='top', fontsize=10, color=KLIPPAIN_COLORS['dark_purple']
+        )
+
+    tmc_registers = motors[0].get_registers()
+    idx = -1
+    for idx, (register, settings) in enumerate(tmc_registers.items()):
+        settings_str = ' '.join(f'{k}={v}' for k, v in settings.items())
+        tmc_block = f'| {register.upper()}: {settings_str}'
+        fig.text(
+            0.41 + distance,
+            0.990 - 0.015 * idx,
+            tmc_block,
+            ha='left',
+            va='top',
+            fontsize=10,
+            color=KLIPPAIN_COLORS['dark_purple'],
+        )
+
+    if differences is not None:
+        differences_text = f'| Y motor diff: {differences}'
+        fig.text(
+            0.41 + distance,
+            0.990 - 0.015 * (idx + 1),
+            differences_text,
+            ha='left',
+            va='top',
+            fontsize=10,
+            color=KLIPPAIN_COLORS['dark_purple'],
+        )
+
+
+######################################################################
+# Startup and main routines
+######################################################################
+
+
+def extract_angle_and_speed(logname):
+    try:
+        match = re.search(r'an(\d+)_\d+sp(\d+)_\d+', os.path.basename(logname))
+        if match:
+            angle = match.group(1)
+            speed = match.group(2)
+        else:
+            raise ValueError(f'File {logname} does not match expected format. Clean your /tmp folder and start again!')
+    except AttributeError as err:
+        raise ValueError(
+            f'File {logname} does not match expected format. Clean your /tmp folder and start again!'
+        ) from err
+    return float(angle), float(speed)
+
+
+def vibrations_profile(
+    lognames, klipperdir='~/klipper', kinematics='cartesian', accel=None, max_freq=1000.0, st_version=None, motors=None
+):
+    global shaper_calibrate
+    shaper_calibrate = setup_klipper_import(klipperdir)
+
+    if kinematics == 'cartesian':
+        main_angles = [0, 90]
+    elif kinematics == 'corexy':
+        main_angles = [45, 135]
+    else:
+        raise ValueError('Only Cartesian and CoreXY kinematics are supported by this tool at the moment!')
+
+    psds = defaultdict(lambda: defaultdict(list))
+    psds_sum = defaultdict(lambda: defaultdict(list))
+    target_freqs_initialized = False
+
+    for logname in lognames:
+        data = parse_log(logname)
+        if data is None:
+            continue  # File is not in the expected format, skip it
+        angle, speed = extract_angle_and_speed(logname)
+        freq_response = calc_freq_response(data)
+        first_freqs = freq_response.freq_bins
+        psd_sum = freq_response.psd_sum
+
+        if not target_freqs_initialized:
+            target_freqs = first_freqs[first_freqs <= max_freq]
+            target_freqs_initialized = True
+
+        psd_sum = psd_sum[first_freqs <= max_freq]
+        first_freqs = first_freqs[first_freqs <= max_freq]
+
+        # Store the interpolated PSD and integral values
+        psds[angle][speed] = np.interp(target_freqs, first_freqs, psd_sum)
+        psds_sum[angle][speed] = np.trapz(psd_sum, first_freqs)
+
+    measured_angles = sorted(psds_sum.keys())
+    measured_speeds = sorted({speed for angle_speeds in psds_sum.values() for speed in angle_speeds.keys()})
+
+    for main_angle in main_angles:
+        if main_angle not in measured_angles:
+            raise ValueError('Measurements not taken at the correct angles for the specified kinematics!')
+
+    # Precompute the variables used in plot functions
+    all_angles, all_speeds, spectrogram_data = compute_dir_speed_spectrogram(
+        measured_speeds, psds_sum, kinematics, main_angles
+    )
+    all_angles_energy = compute_angle_powers(spectrogram_data)
+    sp_min_energy, sp_max_energy, sp_variance_energy, vibration_metric = compute_speed_powers(spectrogram_data)
+    motor_profiles, global_motor_profile = compute_motor_profiles(target_freqs, psds, all_angles_energy, main_angles)
+
+    # symmetry_factor = compute_symmetry_analysis(all_angles, all_angles_energy)
+    symmetry_factor = compute_symmetry_analysis(all_angles, spectrogram_data, main_angles)
+    ConsoleOutput.print(f'Machine estimated vibration symmetry: {symmetry_factor:.1f}%')
+
+    # Analyze low variance ranges of vibration energy across all angles for each speed to identify clean speeds
+    # and highlight them. Also find the peaks to identify speeds to avoid due to high resonances
+    num_peaks, vibration_peaks, peaks_speeds = detect_peaks(
+        vibration_metric,
+        all_speeds,
+        PEAKS_DETECTION_THRESHOLD * vibration_metric.max(),
+        PEAKS_RELATIVE_HEIGHT_THRESHOLD,
+        10,
+        10,
+    )
+    formated_peaks_speeds = ['{:.1f}'.format(pspeed) for pspeed in peaks_speeds]
+    ConsoleOutput.print(
+        '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)))
+    )
+
+    good_speeds = identify_low_energy_zones(vibration_metric, SPEEDS_VALLEY_DETECTION_THRESHOLD)
+    if good_speeds is not None:
+        deletion_range = int(SPEEDS_AROUND_PEAK_DELETION / (all_speeds[1] - all_speeds[0]))
+        peak_speed_indices = {pspeed: np.where(all_speeds == pspeed)[0][0] for pspeed in set(peaks_speeds)}
+
+        # Filter and split ranges based on peak indices, avoiding overlaps
+        good_speeds = filter_and_split_ranges(all_speeds, good_speeds, peak_speed_indices, deletion_range)
+
+        # Add some logging about the good speeds found
+        ConsoleOutput.print(f'Lowest vibrations speeds ({len(good_speeds)} ranges sorted from best to worse):')
+        for idx, (start, end, _) in enumerate(good_speeds):
+            ConsoleOutput.print(f'{idx+1}: {all_speeds[start]:.1f} to {all_speeds[end]:.1f} mm/s')
+
+    # Angle low energy valleys identification (good angles ranges) and print them to the console
+    good_angles = identify_low_energy_zones(all_angles_energy, ANGLES_VALLEY_DETECTION_THRESHOLD)
+    if good_angles is not None:
+        ConsoleOutput.print(f'Lowest vibrations angles ({len(good_angles)} ranges sorted from best to worse):')
+        for idx, (start, end, energy) in enumerate(good_angles):
+            ConsoleOutput.print(
+                f'{idx+1}: {all_angles[start]:.1f}° to {all_angles[end]:.1f}° (mean vibrations energy: {energy:.2f}% of max)'
+            )
+
+    # Create graph layout
+    fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(
+        2,
+        3,
+        gridspec_kw={
+            'height_ratios': [1, 1],
+            'width_ratios': [4, 8, 6],
+            'bottom': 0.050,
+            'top': 0.890,
+            'left': 0.040,
+            'right': 0.985,
+            'hspace': 0.166,
+            'wspace': 0.138,
+        },
+    )
+
+    # Transform ax3 and ax4 to polar plots
+    ax1.remove()
+    ax1 = fig.add_subplot(2, 3, 1, projection='polar')
+    ax4.remove()
+    ax4 = fig.add_subplot(2, 3, 4, projection='polar')
+
+    # Set the global .png figure size
+    fig.set_size_inches(20, 11.5)
+
+    # Add title
+    title_line1 = 'MACHINE VIBRATIONS ANALYSIS TOOL'
+    fig.text(
+        0.060, 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')
+        if accel is not None:
+            title_line2 += ' at ' + str(accel) + ' mm/s² -- ' + kinematics.upper() + ' kinematics'
+    except Exception:
+        ConsoleOutput.print('Warning: CSV filenames appear to be different than expected (%s)' % (lognames[0]))
+        title_line2 = lognames[0].split('/')[-1]
+    fig.text(0.060, 0.957, title_line2, ha='left', va='top', fontsize=16, color=KLIPPAIN_COLORS['dark_purple'])
+
+    # Add the motors infos to the top of the graph
+    if motors is not None and len(motors) == 2:
+        differences = motors[0].compare_to(motors[1])
+        plot_motor_config_txt(fig, motors, differences)
+        if differences is not None and kinematics == 'corexy':
+            ConsoleOutput.print(f'Warning: motors have different TMC configurations!\n{differences}')
+
+    # Plot the graphs
+    plot_angle_profile_polar(ax1, all_angles, all_angles_energy, good_angles, symmetry_factor)
+    plot_vibration_spectrogram_polar(ax4, all_angles, all_speeds, spectrogram_data)
+
+    plot_global_speed_profile(
+        ax2,
+        all_speeds,
+        sp_min_energy,
+        sp_max_energy,
+        sp_variance_energy,
+        vibration_metric,
+        num_peaks,
+        vibration_peaks,
+        good_speeds,
+    )
+    plot_angular_speed_profiles(ax3, all_speeds, all_angles, spectrogram_data, kinematics)
+    plot_vibration_spectrogram(ax5, all_angles, all_speeds, spectrogram_data, vibration_peaks)
+
+    plot_motor_profiles(ax6, target_freqs, main_angles, motor_profiles, global_motor_profile, 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')
+    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
+    if st_version != 'unknown':
+        fig.text(0.995, 0.985, st_version, ha='right', va='bottom', fontsize=8, color=KLIPPAIN_COLORS['purple'])
+
+    return fig
+
+
+def main():
+    # Parse command-line arguments
+    usage = '%prog [options] <raw logs>'
+    opts = optparse.OptionParser(usage)
+    opts.add_option('-o', '--output', type='string', dest='output', default=None, help='filename of output 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.0, help='maximum frequency to graph')
+    opts.add_option(
+        '-k', '--klipper_dir', type='string', dest='klipperdir', default='~/klipper', help='main klipper directory'
+    )
+    opts.add_option(
+        '-m',
+        '--kinematics',
+        type='string',
+        dest='kinematics',
+        default='cartesian',
+        help='machine kinematics configuration',
+    )
+    options, args = opts.parse_args()
+    if len(args) < 1:
+        opts.error('No CSV file(s) to analyse')
+    if options.output is None:
+        opts.error('You must specify an output file.png to use the script (option -o)')
+    if options.kinematics not in ['cartesian', 'corexy']:
+        opts.error('Only cartesian and corexy kinematics are supported by this tool at the moment!')
+
+    fig = vibrations_profile(args, options.klipperdir, options.kinematics, options.accel, options.max_freq)
+    fig.savefig(options.output, dpi=150)
+
+
+if __name__ == '__main__':
+    main()

shaketune/shaketune.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python3
+
+
+from pathlib import Path
+
+from .helpers.console_output import ConsoleOutput
+from .measurement import (
+    axes_map_calibration,
+    axes_shaper_calibration,
+    compare_belts_responses,
+    create_vibrations_profile,
+    excitate_axis_at_freq,
+)
+from .post_processing import AxesMapFinder, BeltsGraphCreator, ShaperGraphCreator, VibrationsGraphCreator
+from .shaketune_config import ShakeTuneConfig
+from .shaketune_thread import ShakeTuneThread
+
+
+class ShakeTune:
+    def __init__(self, config) -> None:
+        self._pconfig = config
+        self._printer = config.get_printer()
+        gcode = self._printer.lookup_object('gcode')
+
+        res_tester = self._printer.lookup_object('resonance_tester')
+        if res_tester is None:
+            config.error('No [resonance_tester] config section found in printer.cfg! Please add one to use Shake&Tune')
+
+        self.timeout = config.getfloat('timeout', 2.0, above=0.0)
+
+        result_folder = config.get('result_folder', default='~/printer_data/config/K-ShakeTune_results')
+        result_folder_path = Path(result_folder).expanduser() if result_folder else None
+        keep_n_results = config.getint('number_of_results_to_keep', default=3, minval=0)
+        keep_csv = config.getboolean('keep_raw_csv', default=False)
+        dpi = config.getint('dpi', default=150, minval=100, maxval=500)
+
+        self._config = ShakeTuneConfig(result_folder_path, keep_n_results, keep_csv, dpi)
+        ConsoleOutput.register_output_callback(gcode.respond_info)
+
+        gcode.register_command(
+            'EXCITATE_AXIS_AT_FREQ',
+            self.cmd_EXCITATE_AXIS_AT_FREQ,
+            desc=self.cmd_EXCITATE_AXIS_AT_FREQ_help,
+        )
+        gcode.register_command(
+            'AXES_MAP_CALIBRATION',
+            self.cmd_AXES_MAP_CALIBRATION,
+            desc=self.cmd_AXES_MAP_CALIBRATION_help,
+        )
+        gcode.register_command(
+            'COMPARE_BELTS_RESPONSES',
+            self.cmd_COMPARE_BELTS_RESPONSES,
+            desc=self.cmd_COMPARE_BELTS_RESPONSES_help,
+        )
+        gcode.register_command(
+            'AXES_SHAPER_CALIBRATION',
+            self.cmd_AXES_SHAPER_CALIBRATION,
+            desc=self.cmd_AXES_SHAPER_CALIBRATION_help,
+        )
+        gcode.register_command(
+            'CREATE_VIBRATIONS_PROFILE',
+            self.cmd_CREATE_VIBRATIONS_PROFILE,
+            desc=self.cmd_CREATE_VIBRATIONS_PROFILE_help,
+        )
+
+    cmd_EXCITATE_AXIS_AT_FREQ_help = (
+        'Maintain a specified excitation frequency for a period of time to diagnose and locate a source of vibration'
+    )
+
+    def cmd_EXCITATE_AXIS_AT_FREQ(self, gcmd) -> None:
+        ConsoleOutput.print(f'Shake&Tune version: {ShakeTuneConfig.get_git_version()}')
+        excitate_axis_at_freq(gcmd, self._pconfig)
+
+    cmd_AXES_MAP_CALIBRATION_help = 'Perform a set of movements to measure the orientation of the accelerometer and help you set the best axes_map configuration for your printer'
+
+    def cmd_AXES_MAP_CALIBRATION(self, gcmd) -> None:
+        ConsoleOutput.print(f'Shake&Tune version: {ShakeTuneConfig.get_git_version()}')
+        axes_map_finder = AxesMapFinder(self._config)
+        st_thread = ShakeTuneThread(self._config, axes_map_finder, self._printer.get_reactor(), self.timeout)
+        axes_map_calibration(gcmd, self._pconfig, st_thread)
+
+    cmd_COMPARE_BELTS_RESPONSES_help = 'Perform a custom half-axis test to analyze and compare the frequency profiles of individual belts on CoreXY printers'
+
+    def cmd_COMPARE_BELTS_RESPONSES(self, gcmd) -> None:
+        ConsoleOutput.print(f'Shake&Tune version: {ShakeTuneConfig.get_git_version()}')
+        belt_graph_creator = BeltsGraphCreator(self._config)
+        st_thread = ShakeTuneThread(self._config, belt_graph_creator, self._printer.get_reactor(), self.timeout)
+        compare_belts_responses(gcmd, self._pconfig, st_thread)
+
+    cmd_AXES_SHAPER_CALIBRATION_help = (
+        'Perform standard axis input shaper tests on one or both XY axes to select the best input shaper filter'
+    )
+
+    def cmd_AXES_SHAPER_CALIBRATION(self, gcmd) -> None:
+        ConsoleOutput.print(f'Shake&Tune version: {ShakeTuneConfig.get_git_version()}')
+        shaper_graph_creator = ShaperGraphCreator(self._config)
+        st_thread = ShakeTuneThread(self._config, shaper_graph_creator, self._printer.get_reactor(), self.timeout)
+        axes_shaper_calibration(gcmd, self._pconfig, st_thread)
+
+    cmd_CREATE_VIBRATIONS_PROFILE_help = 'Perform a set of movements to measure the orientation of the accelerometer and help you set the best axes_map configuration for your printer'
+
+    def cmd_CREATE_VIBRATIONS_PROFILE(self, gcmd) -> None:
+        ConsoleOutput.print(f'Shake&Tune version: {ShakeTuneConfig.get_git_version()}')
+        vibration_profile_creator = VibrationsGraphCreator(self._config)
+        st_thread = ShakeTuneThread(self._config, vibration_profile_creator, self._printer.get_reactor(), self.timeout)
+        create_vibrations_profile(gcmd, self._pconfig, st_thread)

shaketune/shaketune_config.py

@@ -0,0 +1,53 @@
+#!/usr/bin/env python3
+
+from pathlib import Path
+
+from .helpers.console_output import ConsoleOutput
+
+KLIPPER_FOLDER = Path.home() / 'klipper'
+KLIPPER_LOG_FOLDER = Path.home() / 'printer_data/logs'
+RESULTS_BASE_FOLDER = Path.home() / 'printer_data/config/K-ShakeTune_results'
+RESULTS_SUBFOLDERS = {'belts': 'belts', 'shaper': 'inputshaper', 'vibrations': 'vibrations'}
+
+
+class ShakeTuneConfig:
+    def __init__(
+        self, result_folder: Path = RESULTS_BASE_FOLDER, keep_n_results: int = 3, keep_csv: bool = False, dpi: int = 150
+    ) -> None:
+        self._result_folder = result_folder
+
+        self.keep_n_results = keep_n_results
+        self.keep_csv = keep_csv
+        self.dpi = dpi
+
+        self.klipper_folder = KLIPPER_FOLDER
+        self.klipper_log_folder = KLIPPER_LOG_FOLDER
+
+    def get_results_folder(self, type: str = None) -> Path:
+        if type is None:
+            return self._result_folder
+        else:
+            return self._result_folder / RESULTS_SUBFOLDERS[type]
+
+    def get_results_subfolders(self) -> Path:
+        subfolders = [self._result_folder / subfolder for subfolder in RESULTS_SUBFOLDERS.values()]
+        return subfolders
+
+    @staticmethod
+    def get_git_version() -> str:
+        try:
+            from git import GitCommandError, Repo
+
+            # Get the absolute path of the script, resolving any symlinks
+            # Then get 1 times to parent dir to be at the git root folder
+            script_path = Path(__file__).resolve()
+            repo_path = script_path.parents[1]
+            repo = Repo(repo_path)
+            try:
+                version = repo.git.describe('--tags')
+            except GitCommandError:
+                version = repo.head.commit.hexsha[:7]  # If no tag is found, use the simplified commit SHA instead
+            return version
+        except Exception as e:
+            ConsoleOutput.print(f'Warning: unable to retrieve Shake&Tune version number: {e}')
+            return 'unknown'

shaketune/shaketune_thread.py

@@ -0,0 +1,66 @@
+#!/usr/bin/env python3
+
+
+import os
+import threading
+import traceback
+
+from .helpers import filemanager as fm
+from .helpers.console_output import ConsoleOutput
+from .shaketune_config import ShakeTuneConfig
+
+
+class ShakeTuneThread(threading.Thread):
+    def __init__(self, config: ShakeTuneConfig, graph_creator, reactor, timeout: float):
+        super(ShakeTuneThread, self).__init__()
+        self._config = config
+        self.graph_creator = graph_creator
+        self._reactor = reactor
+        self._timeout = timeout
+
+    def get_graph_creator(self):
+        return self.graph_creator
+
+    def run(self) -> None:
+        # Start the target function in a new thread
+        internal_thread = threading.Thread(target=self._shaketune_thread, args=(self.graph_creator,))
+        internal_thread.start()
+
+        # Monitor the thread execution and stop it if it takes too long
+        event_time = self._reactor.monotonic()
+        end_time = event_time + self._timeout
+        while event_time < end_time:
+            event_time = self._reactor.pause(event_time + 0.05)
+            if not internal_thread.is_alive():
+                break
+
+    # This function run in its own thread is used to do the CSV analysis and create the graphs
+    def _shaketune_thread(self, graph_creator) -> None:
+        # Trying to reduce the Shake&Tune prost-processing thread priority to avoid slowing down the main Klipper process
+        # as this could lead to random "Timer" errors when already running CANbus, etc...
+        try:
+            os.nice(20)
+        except Exception:
+            ConsoleOutput.print('Warning: failed reducing Shake&Tune thread priority, continuing...')
+
+        fm.ensure_folders_exist(self._config.get_results_subfolders())
+
+        try:
+            graph_creator.create_graph()
+        except FileNotFoundError as e:
+            ConsoleOutput.print(f'FileNotFound error: {e}')
+            return
+        except TimeoutError as e:
+            ConsoleOutput.print(f'Timeout error: {e}')
+            return
+        except Exception as e:
+            ConsoleOutput.print(f'Error while generating the graphs: {e}\n{traceback.print_exc()}')
+            return
+
+        graph_creator.clean_old_files(self._config.keep_n_results)
+
+        if graph_creator.get_type() != 'axesmap':
+            ConsoleOutput.print(f'{graph_creator.get_type()} graphs created successfully!')
+            ConsoleOutput.print(
+                f'Cleaned up the output folder (only the last {self._config.keep_n_results} results were kept)!'
+            )