From 78dcce412fc608d18471459481e8393a847175e2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?F=C3=A9lix=20Boisselier?= <felix@fboisselier.fr>
Date: Fri, 3 May 2024 20:22:43 +0200
Subject: [PATCH] added kinematics info to belt comparison graph

---
 K-ShakeTune/K-SnT_belts.cfg       |  4 +++-
 src/graph_creators/graph_belts.py | 13 +++++++++++--
 src/is_workflow.py                |  8 +++++++-
 3 files changed, 21 insertions(+), 4 deletions(-)

diff --git a/K-ShakeTune/K-SnT_belts.cfg b/K-ShakeTune/K-SnT_belts.cfg
index cd4987c..12f5d67 100644
--- a/K-ShakeTune/K-SnT_belts.cfg
+++ b/K-ShakeTune/K-SnT_belts.cfg
@@ -12,6 +12,8 @@ gcode:
     {% set keep_results = params.KEEP_N_RESULTS|default(3)|int %}
     {% set keep_csv = params.KEEP_CSV|default(0)|int %}
 
+    {% set kinematics = printer.configfile.settings.printer.kinematics %}
+
     TEST_RESONANCES AXIS=1,1 OUTPUT=raw_data NAME=b FREQ_START={min_freq} FREQ_END={max_freq} HZ_PER_SEC={hz_per_sec}
     M400
 
@@ -20,4 +22,4 @@ gcode:
 
     RESPOND MSG="Belts comparative frequency profile generation..."
     RESPOND MSG="This may take some time (3-5min)"
-    RUN_SHELL_COMMAND CMD=shaketune PARAMS="--type belts {% if keep_csv %}--keep_csv{% endif %} --keep_results {keep_results}"
+    RUN_SHELL_COMMAND CMD=shaketune PARAMS="--type belts --kinematics {kinematics} {% if keep_csv %}--keep_csv{% endif %} --keep_results {keep_results}"
diff --git a/src/graph_creators/graph_belts.py b/src/graph_creators/graph_belts.py
index a3104c5..d966302 100644
--- a/src/graph_creators/graph_belts.py
+++ b/src/graph_creators/graph_belts.py
@@ -424,7 +424,7 @@ def compute_signal_data(data, max_freq):
 ######################################################################
 
 
-def belts_calibration(lognames, klipperdir='~/klipper', max_freq=200.0, st_version=None):
+def belts_calibration(lognames, kinematics, klipperdir='~/klipper', max_freq=200.0, st_version=None):
     set_locale()
     global shaper_calibrate
     shaper_calibrate = setup_klipper_import(klipperdir)
@@ -493,6 +493,8 @@ def belts_calibration(lognames, klipperdir='~/klipper', max_freq=200.0, st_versi
         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:
         print_with_c_locale(
             'Warning: CSV filenames look to be different than expected (%s , %s)' % (lognames[0], lognames[1])
@@ -530,13 +532,20 @@ def main():
     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.klipperdir, options.max_freq)
+    fig = belts_calibration(args, options.kinematics, options.klipperdir, options.max_freq)
     fig.savefig(options.output, dpi=150)
 
 
diff --git a/src/is_workflow.py b/src/is_workflow.py
index 5ae1370..c188413 100755
--- a/src/is_workflow.py
+++ b/src/is_workflow.py
@@ -207,8 +207,13 @@ class BeltsGraphCreator(GraphCreator):
     def __init__(self, keep_csv: bool = False, dpi: int = 150):
         super().__init__(keep_csv, dpi)
 
+        self._kinematics = None
+
         self._setup_folder('belts')
 
+    def configure(self, kinematics: str) -> None:
+        self._kinematics = kinematics
+
     def create_graph(self) -> None:
         lognames = self._move_and_prepare_files(
             glob_pattern='raw_data_axis*.csv',
@@ -217,6 +222,7 @@ class BeltsGraphCreator(GraphCreator):
         )
         fig = belts_calibration(
             lognames=[str(path) for path in lognames],
+            kinematics=self._kinematics,
             klipperdir=str(Config.KLIPPER_FOLDER),
             st_version=self._version,
         )
@@ -381,7 +387,7 @@ def main():
     print_with_c_locale(f'Shake&Tune version: {Config.get_git_version()}')
 
     graph_creators = {
-        'belts': (BeltsGraphCreator, None),
+        'belts': (BeltsGraphCreator, lambda gc: gc.configure(options.kinematics)),
         'shaper': (ShaperGraphCreator, lambda gc: gc.configure(options.scv, options.max_smoothing)),
         'vibrations': (
             VibrationsGraphCreator,