resonance tester: disable limits (not optimal)

klippy/extras/resonance_tester.py

@@ -6,6 +6,7 @@
 import math
 import os
 import time
+from contextlib import contextmanager
 from . import shaper_calibrate
 
 
@@ -52,6 +53,73 @@ def _parse_axis(gcmd, raw_axis):
     return TestAxis(vib_dir=(dir_x, dir_y))
 
 
+@contextmanager
+def suspend_limits(printer, max_accel, max_velocity, input_shaping):
+    # Override maximum acceleration and acceleration to
+    # deceleration based on the maximum test frequency
+    gcode = printer.lookup_object("gcode")
+    input_shaper = printer.lookup_object("input_shaper", None)
+    if input_shaper is not None and not input_shaping:
+        input_shaper.disable_shaping()
+        gcode.respond_info("Disabled [input_shaper] for resonance testing")
+    else:
+        input_shaper = None
+    toolhead = printer.lookup_object("toolhead")
+    systime = printer.get_reactor().monotonic()
+    toolhead_info = toolhead.get_status(systime)
+    old_max_accel = toolhead_info["max_accel"]
+    old_max_accel_to_decel = toolhead_info["max_accel_to_decel"]
+    old_max_velocity = toolhead_info["max_velocity"]
+    gcode.run_script_from_command(
+        "SET_VELOCITY_LIMIT ACCEL=%.3f ACCEL_TO_DECEL=%.3f VELOCITY=%.3f"
+        % (max_accel, max_accel, max_velocity)
+    )
+    kin = toolhead.get_kinematics()
+    old_max_velocities = getattr(kin, "max_velocities", None)
+    if old_max_velocities is not None:
+        kin.max_velocities = [
+            max_velocity,
+            max_velocity,
+            old_max_velocities[-1],
+        ]
+    old_max_accels = getattr(kin, "max_accels", None)
+    if old_max_accels is not None:
+        kin.max_accels = [max_accel, max_accel, old_max_accels[-1]]
+    # FIXME: could be cleaner if limited_corexy were using the same format than
+    # limited_cartesian
+    old_max_x_accel = getattr(kin, "max_x_accel", None)
+    if old_max_x_accel is not None:
+        kin.max_x_accel = max_accel
+    old_max_y_accel = getattr(kin, "max_y_accel", None)
+    if old_max_y_accel is not None:
+        kin.max_y_accel = max_accel
+    old_scale_per_axis = getattr(kin, "scale_per_axis", None)
+    if old_scale_per_axis is not None:
+        kin.scale_per_axis = False
+    try:
+        yield
+    finally:
+        # Restore input shaper if it was disabled for resonance testing
+        if input_shaper is not None:
+            input_shaper.enable_shaping()
+            gcode.respond_info("Re-enabled [input_shaper]")
+        # Restore the original acceleration values
+        gcode.run_script_from_command(
+            "SET_VELOCITY_LIMIT ACCEL=%.3f ACCEL_TO_DECEL=%.3f VELOCITY=%.3f"
+            % (old_max_accel, old_max_accel_to_decel, old_max_velocity)
+        )
+        if old_max_velocities is not None:
+            kin.max_velocities = old_max_velocities
+        if old_max_accels is not None:
+            kin.max_accels = old_max_accels
+        if old_max_x_accel is not None:
+            kin.max_x_accel = old_max_x_accel
+        if old_max_y_accel is not None:
+            kin.max_y_accel = old_max_y_accel
+        if old_scale_per_axis is not None:
+            kin.scale_per_axis = old_scale_per_axis
+
+
 class VibrationPulseTest:
     def __init__(self, config):
         self.printer = config.get_printer()
@@ -85,27 +153,19 @@ def prepare_test(self, gcmd):
         )
 
     def run_test(self, axis, gcmd):
+        with suspend_limits(
+            self.printer,
+            self.freq_end * self.accel_per_hz + 10.0,
+            self.accel_per_hz * 0.25 + 1.0,
+            gcmd.get_int("INPUT_SHAPING", 0),
+        ):
+            self._run_test(axis, gcmd)
+
+    def _run_test(self, axis, gcmd):
         toolhead = self.printer.lookup_object("toolhead")
         X, Y, Z, E = toolhead.get_position()
         sign = 1.0
         freq = self.freq_start
-        # Override maximum acceleration and acceleration to
-        # deceleration based on the maximum test frequency
-        systime = self.printer.get_reactor().monotonic()
-        toolhead_info = toolhead.get_status(systime)
-        old_max_accel = toolhead_info["max_accel"]
-        old_max_accel_to_decel = toolhead_info["max_accel_to_decel"]
-        max_accel = self.freq_end * self.accel_per_hz
-        self.gcode.run_script_from_command(
-            "SET_VELOCITY_LIMIT ACCEL=%.3f ACCEL_TO_DECEL=%.3f"
-            % (max_accel, max_accel)
-        )
-        input_shaper = self.printer.lookup_object("input_shaper", None)
-        if input_shaper is not None and not gcmd.get_int("INPUT_SHAPING", 0):
-            input_shaper.disable_shaping()
-            gcmd.respond_info("Disabled [input_shaper] for resonance testing")
-        else:
-            input_shaper = None
         gcmd.respond_info("Testing frequency %.0f Hz" % (freq,))
         while freq <= self.freq_end + 0.000001:
             t_seg = 0.25 / freq
@@ -125,15 +185,6 @@ def run_test(self, axis, gcmd):
             freq += 2.0 * t_seg * self.hz_per_sec
             if math.floor(freq) > math.floor(old_freq):
                 gcmd.respond_info("Testing frequency %.0f Hz" % (freq,))
-        # Restore the original acceleration values
-        self.gcode.run_script_from_command(
-            "SET_VELOCITY_LIMIT ACCEL=%.3f ACCEL_TO_DECEL=%.3f"
-            % (old_max_accel, old_max_accel_to_decel)
-        )
-        # Restore input shaper if it was disabled for resonance testing
-        if input_shaper is not None:
-            input_shaper.enable_shaping()
-            gcmd.respond_info("Re-enabled [input_shaper]")
 
 
 class ResonanceTester: