Klipper plugin refactoring with embedded macros

2024-05-09 16:08:47 +02:00
parent d9060fed3b
commit 30a1910513
21 changed files with 762 additions and 644 deletions
--- a/shaketune/graph_creators/graph_creator.py
+++ b/shaketune/graph_creators/graph_creator.py
@@ -0,0 +1,276 @@
+#!/usr/bin/env python3
+
+import abc
+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 ..helpers.motorlogparser import MotorLogParser
+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)
+            fm.wait_file_ready(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._setup_folder('belts')
+
+    def create_graph(self) -> None:
+        lognames = self._move_and_prepare_files(
+            glob_pattern='raw_data_axis*.csv',
+            min_files_required=2,
+            custom_name_func=lambda f: f.stem.split('_')[3].upper(),
+        )
+        fig = belts_calibration(
+            lognames=[str(path) for path in lognames],
+            klipperdir=str(self._config.klipper_folder),
+            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) -> None:
+        self._scv = scv
+        self._max_smoothing = max_smoothing
+
+    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='raw_data*.csv',
+            min_files_required=1,
+            custom_name_func=lambda f: f.stem.split('_')[3].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,
+            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._chip_name = None
+        self._motors = None
+
+        self._setup_folder('vibrations')
+
+    def configure(self, kinematics: str, accel: float, chip_name: str, metadata: str) -> None:
+        self._kinematics = kinematics
+        self._accel = accel
+        self._chip_name = chip_name
+
+        parser = MotorLogParser(self._config.klipper_log_folder / 'klippy.log', metadata)
+        self._motors = 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)
+
+    def create_graph(self) -> None:
+        if not self._accel or not self._chip_name 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=f'{self._chip_name}-*.csv',
+            min_files_required=None,
+            custom_name_func=lambda f: f.name.replace(self._chip_name, self._type),
+        )
+        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
+        self._chip_name = None
+
+    def configure(self, accel: int, chip_name: str) -> None:
+        self._accel = accel
+        self._chip_name = chip_name
+
+    def find_axesmap(self) -> None:
+        tmp_folder = Path('/tmp')
+        globbed_files = list(tmp_folder.glob(f'{self._chip_name}-*.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 wait for it to be released by Klipper
+        logname = sorted(globbed_files, key=lambda f: f.stat().st_mtime, reverse=True)[0]
+        fm.wait_file_ready(logname)
+
+        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(f'{self._chip_name}-*.csv'))
+        for csv_file in globbed_files:
+            csv_file.unlink()