Refactoring of the graph creators

README.md

@@ -47,6 +47,10 @@ Follow these steps to install Shake&Tune on your printer:
      #    RAM, and should work for everyone. However, if you are using a powerful computer, you may
      #    wish to increase this value to keep more measurements in memory (e.g., 15-20) before writing
      #    the chunk and avoid stressing the filesystem too much.
+     # max_freq: 200
+     #    This setting defines the maximum frequency at which the calculation of the power spectral density
+     #    is cutoff. The default value should be fine for most machines and accelerometer combinations and
+     #    avoid touching it unless you know what you're doing.
      # dpi: 300
      #    Controls the resolution of the generated graphs. The default value of 300 dpi was optimized
      #    and strikes a balance between performance and readability, ensuring that graphs are clear

shaketune/cli.py

@@ -0,0 +1,140 @@
+import argparse
+import os
+import sys
+from importlib import import_module
+from pathlib import Path
+
+from .graph_creators.graph_creator_factory import GraphCreatorFactory
+from .helpers.accelerometer import MeasurementsManager
+from .shaketune_config import ShakeTuneConfig
+
+
+def add_common_arguments(parser):
+    """Helper function to add common arguments to all subparsers."""
+    parser.add_argument('-o', '--output', required=True, help='Output filename')
+    parser.add_argument('files', nargs='+', help='Input data files (.csv or .stdata)')
+    parser.add_argument('--max_freq', type=float, help='Maximum frequency to graph')
+    parser.add_argument('--dpi', type=int, help='DPI value to use for the graph')
+
+
+def configure_graph_creator(graph_type, args, dummy_config):
+    """Helper function to get and configure a graph creator based on graph type and args."""
+    graph_creator = GraphCreatorFactory.create_graph_creator(graph_type, dummy_config)
+    config_kwargs = {}
+
+    # Dynamically configure the graph creator based on graph type
+    if graph_type == 'axes map':
+        config_kwargs |= {'accel': args.accel, 'segment_length': args.length}
+    elif graph_type == 'static frequency':
+        config_kwargs |= {'accel_per_hz': args.accel_per_hz, 'freq': args.frequency, 'duration': args.duration}
+    elif graph_type == 'belts comparison':
+        config_kwargs |= {'kinematics': args.kinematics, 'accel_per_hz': args.accel_per_hz}
+    elif graph_type == 'input shaper':
+        config_kwargs |= {'scv': args.scv, 'max_smoothing': args.max_smoothing, 'accel_per_hz': args.accel_per_hz}
+    elif graph_type == 'vibrations profile':
+        config_kwargs |= {'kinematics': args.kinematics, 'accel': args.accel}
+
+    graph_creator.configure(**config_kwargs)
+    return graph_creator
+
+
+def load_klipper_module(args):
+    """Helper function to load the shaper_calibrate module from the specified Klipper folder."""
+    if hasattr(args, 'klipper_dir') and args.klipper_dir:
+        kdir = os.path.expanduser(args.klipper_dir)
+        sys.path.append(os.path.join(kdir, 'klippy'))
+        sys.modules['shaper_calibrate'] = import_module('.shaper_calibrate', 'extras')
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Shake&Tune command line interface')
+    subparsers = parser.add_subparsers(dest='graph_type', help='Type of graph to create')
+
+    # Static frequency graph parser
+    static_freq_parser = subparsers.add_parser('static_freq', help='Create static frequency graph')
+    add_common_arguments(static_freq_parser)
+    static_freq_parser.add_argument('--accel_per_hz', type=float, help='Accel per Hz used during the measurement')
+    static_freq_parser.add_argument('--frequency', type=float, help='Maintained frequency of the measurement')
+    static_freq_parser.add_argument('--duration', type=float, help='Duration of the measurement')
+
+    # Axes map detection graph parser
+    axes_map_parser = subparsers.add_parser('axes_map', help='Create axes map detection graph')
+    add_common_arguments(axes_map_parser)
+    axes_map_parser.add_argument('--accel', required=True, type=float, help='Accel value used for the measurement')
+    axes_map_parser.add_argument('--length', required=True, type=float, help='Recorded length for each segment')
+
+    # Belts graph parser
+    belts_parser = subparsers.add_parser('belts', help='Create belts comparison graph')
+    add_common_arguments(belts_parser)
+    belts_parser.add_argument('-k', '--klipper_dir', default='~/klipper', help='Main klipper directory')
+    belts_parser.add_argument('--kinematics', help='Machine kinematics configuration')
+    belts_parser.add_argument('--accel_per_hz', type=float, help='Accel per Hz used during the measurement')
+
+    # Input Shaper graph parser
+    shaper_parser = subparsers.add_parser('input_shaper', help='Create input shaper graph')
+    add_common_arguments(shaper_parser)
+    shaper_parser.add_argument('-k', '--klipper_dir', default='~/klipper', help='Main klipper directory')
+    shaper_parser.add_argument('--scv', type=float, default=5.0, help='Square corner velocity')
+    shaper_parser.add_argument('--max_smoothing', type=float, help='Maximum shaper smoothing to allow')
+    shaper_parser.add_argument('--accel_per_hz', type=float, help='Accel per Hz used during the measurement')
+
+    # Vibrations graph parser
+    vibrations_parser = subparsers.add_parser('vibrations', help='Create vibrations profile graph')
+    add_common_arguments(vibrations_parser)
+    vibrations_parser.add_argument('-k', '--klipper_dir', default='~/klipper', help='Main klipper directory')
+    vibrations_parser.add_argument('--kinematics', required=True, default='cartesian', help='Used kinematics')
+    vibrations_parser.add_argument('--accel', type=int, help='Accel value to be printed on the graph')
+
+    args = parser.parse_args()
+
+    if args.graph_type is None:
+        parser.print_help()
+        exit(1)
+
+    graph_type_map = {
+        'static_freq': 'static frequency',
+        'axes_map': 'axes map',
+        'belts': 'belts comparison',
+        'input_shaper': 'input shaper',
+        'vibrations': 'vibrations profile',
+    }
+    graph_type = graph_type_map[args.graph_type]
+
+    # Load configuration
+    dummy_config = ShakeTuneConfig()
+    if args.dpi is not None:
+        dummy_config.dpi = args.dpi
+    if args.max_freq is not None:
+        if graph_type == 'vibrations profile':
+            dummy_config.max_freq_vibrations = args.max_freq
+        else:
+            dummy_config.max_freq = args.max_freq
+
+    # Load shaper_calibrate module if needed
+    load_klipper_module(args)
+
+    # Create the graph creator and configure it
+    graph_creator = configure_graph_creator(graph_type, args, dummy_config)
+    graph_creator.override_output_target(args.output)
+
+    print(f'Creating {graph_type} graph...')
+
+    # Load measurements
+    measurements_manager = MeasurementsManager(10)
+    args.files = [Path(f) for f in args.files]
+    if args.files[0].suffix == '.csv':
+        measurements_manager.load_from_csvs(args.files)
+    elif args.files[0].suffix == '.stdata':
+        measurements_manager.load_from_stdata(args.files[0])
+    else:
+        raise ValueError('Only .stdata or legacy Klipper raw accelerometer CSV files are supported!')
+
+    # Create graph
+    graph_creator.create_graph(measurements_manager)
+
+    print('...done!')
+
+
+if __name__ == '__main__':
+    os.environ['SHAKETUNE_IN_CLI'] = '1'
+    main()

shaketune/graph_creators/__init__.py

@@ -6,10 +6,22 @@
 # File: __init__.py
 # Description: Imports various graph creator classes for the Shake&Tune package.
 
+import os
+import sys
 
-from .axes_map_graph_creator import AxesMapGraphCreator as AxesMapGraphCreator
-from .belts_graph_creator import BeltsGraphCreator as BeltsGraphCreator
-from .graph_creator import GraphCreator as GraphCreator
-from .shaper_graph_creator import ShaperGraphCreator as ShaperGraphCreator
-from .static_graph_creator import StaticGraphCreator as StaticGraphCreator
-from .vibrations_graph_creator import VibrationsGraphCreator as VibrationsGraphCreator
+
+def get_shaper_calibrate_module():
+    if os.environ.get('SHAKETUNE_IN_CLI') != '1':
+        from ... import shaper_calibrate
+    else:
+        shaper_calibrate = sys.modules['shaper_calibrate']
+    return shaper_calibrate
+
+
+from .axes_map_graph_creator import AxesMapGraphCreator as AxesMapGraphCreator  # noqa: E402
+from .belts_graph_creator import BeltsGraphCreator as BeltsGraphCreator  # noqa: E402
+from .graph_creator import GraphCreator as GraphCreator  # noqa: E402
+from .graph_creator_factory import GraphCreatorFactory as GraphCreatorFactory  # noqa: E402
+from .shaper_graph_creator import ShaperGraphCreator as ShaperGraphCreator  # noqa: E402
+from .static_graph_creator import StaticGraphCreator as StaticGraphCreator  # noqa: E402
+from .vibrations_graph_creator import VibrationsGraphCreator as VibrationsGraphCreator  # noqa: E402