shaper_calibrate: Fixed crashes in SHAPER_CALIBRATE and TEST_RESONANCES

Fixed crashes due to wrong parameter passed to the shaper selection function
and when the custom FREQ_END is specified.

Signed-off-by: Dmitry Butyugin <dmbutyugin@google.com>

2 files changed, 16 insertions, 12 deletions

diff --git a/klippy/extras/resonance_tester.py b/klippy/extras/resonance_tester.py
index 4b29d6b8..d249412c 100644
--- a/klippy/extras/resonance_tester.py
+++ b/klippy/extras/resonance_tester.py
@@ -52,7 +52,7 @@ class VibrationPulseTest:
         self.min_freq = config.getfloat('min_freq', 5., minval=1.)
         # Defaults are such that max_freq * accel_per_hz == 10000 (max_accel)
         self.max_freq = config.getfloat('max_freq', 10000. / 75.,
-                                        minval=self.min_freq, maxval=200.)
+                                        minval=self.min_freq, maxval=300.)
         self.accel_per_hz = config.getfloat('accel_per_hz', 75., above=0.)
         self.hz_per_sec = config.getfloat('hz_per_sec', 1.,
                                           minval=0.1, maxval=2.)
@@ -64,7 +64,7 @@ class VibrationPulseTest:
     def prepare_test(self, gcmd):
         self.freq_start = gcmd.get_float("FREQ_START", self.min_freq, minval=1.)
         self.freq_end = gcmd.get_float("FREQ_END", self.max_freq,
-                                       minval=self.freq_start, maxval=200.)
+                                       minval=self.freq_start, maxval=300.)
         self.hz_per_sec = gcmd.get_float("HZ_PER_SEC", self.hz_per_sec,
                                          above=0., maxval=2.)
     def run_test(self, axis, gcmd):
@@ -219,6 +219,8 @@ class ResonanceTester:
             chip = self.printer.lookup_object(chip_lookup_name)
             parsed_chips.append(chip)
         return parsed_chips
+    def _get_max_calibration_freq(self):
+        return 1.5 * self.test.get_max_freq()
     cmd_TEST_RESONANCES_help = ("Runs the resonance test for a specifed axis")
     def cmd_TEST_RESONANCES(self, gcmd):
         # Parse parameters
@@ -263,9 +265,9 @@ class ResonanceTester:
                 raw_name_suffix=name_suffix if raw_output else None,
                 accel_chips=accel_chips, test_point=test_point)[axis]
         if csv_output:
-            csv_name = self.save_calibration_data('resonances', name_suffix,
-                                                  helper, axis, data,
-                                                  point=test_point)
+            csv_name = self.save_calibration_data(
+                    'resonances', name_suffix, helper, axis, data,
+                    point=test_point, max_freq=self._get_max_calibration_freq())
             gcmd.respond_info(
                     "Resonances data written to %s file" % (csv_name,))
     cmd_SHAPER_CALIBRATE_help = (
@@ -308,10 +310,10 @@ class ResonanceTester:
             toolhead = self.printer.lookup_object('toolhead')
             toolhead_info = toolhead.get_status(systime)
             scv = toolhead_info['square_corner_velocity']
+            max_freq = self._get_max_calibration_freq()
             best_shaper, all_shapers = helper.find_best_shaper(
                     calibration_data[axis], max_smoothing=max_smoothing,
-                    scv=scv, max_freq=1.5*self.test.get_max_freq(),
-                    logging=gcmd.respond_info)
+                    scv=scv, max_freq=max_freq, logger=gcmd.respond_info)
             gcmd.respond_info(
                     "Recommended shaper_type_%s = %s, shaper_freq_%s = %.1f Hz"
                     % (axis_name, best_shaper.name,
@@ -323,7 +325,7 @@ class ResonanceTester:
                                best_shaper.name, best_shaper.freq)
             csv_name = self.save_calibration_data(
                     'calibration_data', name_suffix, helper, axis,
-                    calibration_data[axis], all_shapers)
+                    calibration_data[axis], all_shapers, max_freq=max_freq)
             gcmd.respond_info(
                     "Shaper calibration data written to %s file" % (csv_name,))
         gcmd.respond_info(
@@ -369,10 +371,10 @@ class ResonanceTester:
 
     def save_calibration_data(self, base_name, name_suffix, shaper_calibrate,
                               axis, calibration_data,
-                              all_shapers=None, point=None):
+                              all_shapers=None, point=None, max_freq=None):
         output = self.get_filename(base_name, name_suffix, axis, point)
         shaper_calibrate.save_calibration_data(output, calibration_data,
-                                               all_shapers)
+                                               all_shapers, max_freq)
         return output
 
 def load_config(config):
diff --git a/klippy/extras/shaper_calibrate.py b/klippy/extras/shaper_calibrate.py
index f3bfd8d2..6891fefb 100644
--- a/klippy/extras/shaper_calibrate.py
+++ b/klippy/extras/shaper_calibrate.py
@@ -368,8 +368,10 @@ class ShaperCalibrate:
                     "SHAPER_TYPE_" + axis: shaper_name,
                     "SHAPER_FREQ_" + axis: shaper_freq}))
 
-    def save_calibration_data(self, output, calibration_data, shapers=None):
+    def save_calibration_data(self, output, calibration_data, shapers=None,
+                              max_freq=None):
         try:
+            max_freq = max_freq or MAX_FREQ
             with open(output, "w") as csvfile:
                 csvfile.write("freq,psd_x,psd_y,psd_z,psd_xyz")
                 if shapers:
@@ -378,7 +380,7 @@ class ShaperCalibrate:
                 csvfile.write("\n")
                 num_freqs = calibration_data.freq_bins.shape[0]
                 for i in range(num_freqs):
-                    if calibration_data.freq_bins[i] >= MAX_FREQ:
+                    if calibration_data.freq_bins[i] >= max_freq:
                         break
                     csvfile.write("%.1f,%.3e,%.3e,%.3e,%.3e" % (
                         calibration_data.freq_bins[i],