Static freq optional graphs (#112)

shaketune/post_processing/graph_creator.py

@@ -15,6 +15,7 @@ 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_static import static_frequency_tool
 from .graph_vibrations import vibrations_profile
 
 
@@ -275,3 +276,52 @@ class AxesMapFinder(GraphCreator):
                 csv_file = self._folder / f'axesmap_{file_date}_{suffix}.csv'
                 csv_file.unlink(missing_ok=True)
             old_file.unlink()
+
+
+class StaticGraphCreator(GraphCreator):
+    def __init__(self, config: ShakeTuneConfig):
+        super().__init__(config)
+
+        self._freq = None
+        self._duration = None
+        self._accel_per_hz = None
+
+        self._setup_folder('staticfreq')
+
+    def configure(self, freq: float, duration: float, accel_per_hz: float = None) -> None:
+        self._freq = freq
+        self._duration = duration
+        self._accel_per_hz = accel_per_hz
+
+    def create_graph(self) -> None:
+        if not self._freq or not self._duration or not self._accel_per_hz:
+            raise ValueError('freq, duration and accel_per_hz must be set to create the static frequency graph!')
+
+        lognames = self._move_and_prepare_files(
+            glob_pattern='shaketune-staticfreq_*.csv',
+            min_files_required=1,
+            custom_name_func=lambda f: f.stem.split('_')[1].upper(),
+        )
+        fig = static_frequency_tool(
+            lognames=[str(path) for path in lognames],
+            klipperdir=str(self._config.klipper_folder),
+            freq=self._freq,
+            duration=self._duration,
+            max_freq=200.0,
+            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) <= keep_results:
+            return  # No need to delete any files
+
+        # Delete the older files
+        for old_file in files[keep_results:]:
+            csv_file = old_file.with_suffix('.csv')
+            csv_file.unlink(missing_ok=True)
+            old_file.unlink()