input_shaper: Factored out AxisInputShaper class

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

3 files changed, 129 insertions, 113 deletions

diff --git a/klippy/chelper/__init__.py b/klippy/chelper/__init__.py
index b5eb12dd..66d701d3 100644
--- a/klippy/chelper/__init__.py
+++ b/klippy/chelper/__init__.py
@@ -140,9 +140,8 @@ defs_kin_extruder = """
 defs_kin_shaper = """
     double input_shaper_get_step_generation_window(int n, double a[]
         , double t[]);
-    int input_shaper_set_shaper_params(struct stepper_kinematics *sk
-        , int n_x, double a_x[], double t_x[]
-        , int n_y, double a_y[], double t_y[]);
+    int input_shaper_set_shaper_params(struct stepper_kinematics *sk, char axis
+        , int n, double a[], double t[]);
     int input_shaper_set_sk(struct stepper_kinematics *sk
         , struct stepper_kinematics *orig_sk);
     struct stepper_kinematics * input_shaper_alloc(void);
diff --git a/klippy/chelper/kin_shaper.c b/klippy/chelper/kin_shaper.c
index efba013d..133f46a9 100644
--- a/klippy/chelper/kin_shaper.c
+++ b/klippy/chelper/kin_shaper.c
@@ -38,12 +38,12 @@ shift_pulses(struct shaper_pulses *sp)
         sp->pulses[i].t -= ts;
 }
 
-static void
+static int
 init_shaper(int n, double a[], double t[], struct shaper_pulses *sp)
 {
     if (n < 0 || n > ARRAY_SIZE(sp->pulses)) {
         sp->num_pulses = 0;
-        return;
+        return -1;
     }
     int i;
     double sum_a = 0.;
@@ -57,6 +57,7 @@ init_shaper(int n, double a[], double t[], struct shaper_pulses *sp)
     }
     sp->num_pulses = n;
     shift_pulses(sp);
+    return 0;
 }
 
 
@@ -192,21 +193,20 @@ shaper_note_generation_time(struct input_shaper *is)
 }
 
 int __visible
-input_shaper_set_shaper_params(struct stepper_kinematics *sk
-                               , int n_x, double a_x[], double t_x[]
-                               , int n_y, double a_y[], double t_y[])
+input_shaper_set_shaper_params(struct stepper_kinematics *sk, char axis
+                               , int n, double a[], double t[])
 {
+    if (axis != 'x' && axis != 'y')
+        return -1;
     struct input_shaper *is = container_of(sk, struct input_shaper, sk);
-    if (is->orig_sk->active_flags & AF_X)
-        init_shaper(n_x, a_x, t_x, &is->sx);
+    struct shaper_pulses *sp = axis == 'x' ? &is->sx : &is->sy;
+    int status = 0;
+    if (is->orig_sk->active_flags & (axis == 'x' ? AF_X : AF_Y))
+        status = init_shaper(n, a, t, sp);
     else
-        is->sx.num_pulses = 0;
-    if (is->orig_sk->active_flags & AF_Y)
-        init_shaper(n_y, a_y, t_y, &is->sy);
-    else
-        is->sy.num_pulses = 0;
+        sp->num_pulses = 0;
     shaper_note_generation_time(is);
-    return 0;
+    return status;
 }
 
 double __visible
diff --git a/klippy/extras/input_shaper.py b/klippy/extras/input_shaper.py
index 1d8fe0c7..c4963c8d 100644
--- a/klippy/extras/input_shaper.py
+++ b/klippy/extras/input_shaper.py
@@ -4,32 +4,99 @@
 # Copyright (C) 2020  Dmitry Butyugin <dmbutyugin@google.com>
 #
 # This file may be distributed under the terms of the GNU GPLv3 license.
+import collections
 import chelper
 from . import shaper_defs
 
+class InputShaperParams:
+    def __init__(self, axis, config):
+        self.axis = axis
+        self.shapers = {s.name : s.init_func for s in shaper_defs.INPUT_SHAPERS}
+        self.shaper_type = config.get('shaper_type_' + axis, 'mzv').lower()
+        if self.shaper_type not in self.shapers:
+            raise config.error(
+                    'Unsupported shaper type: %s' % (self.shaper_type,))
+        self.damping_ratio = config.getfloat('damping_ratio_' + axis,
+                                             shaper_defs.DEFAULT_DAMPING_RATIO,
+                                             minval=0., maxval=1.)
+        self.shaper_freq = config.getfloat('shaper_freq_' + axis, 0., minval=0.)
+    def update(self, gcmd):
+        axis = self.axis.upper()
+        self.damping_ratio = gcmd.get_float('DAMPING_RATIO_' + axis,
+                                            self.damping_ratio,
+                                            minval=0., maxval=1.)
+        self.shaper_freq = gcmd.get_float('SHAPER_FREQ_' + axis,
+                                          self.shaper_freq, minval=0.)
+        shaper_type = gcmd.get('SHAPER_TYPE', None)
+        if shaper_type is None:
+            shaper_type = gcmd.get('SHAPER_TYPE_' + axis, self.shaper_type)
+        if shaper_type.lower() not in self.shapers:
+            raise gcmd.error('Unsupported shaper type: %s' % (shaper_type,))
+        self.shaper_type = shaper_type.lower()
+    def get_shaper(self):
+        if not self.shaper_freq:
+            A, T = shaper_defs.get_none_shaper()
+        else:
+            A, T = self.shapers[self.shaper_type](
+                    self.shaper_freq, self.damping_ratio)
+        return len(A), A, T
+    def get_status(self):
+        return collections.OrderedDict([
+            ('shaper_type', self.shaper_type),
+            ('shaper_freq', '%.3f' % (self.shaper_freq,)),
+            ('damping_ratio', '%.6f' % (self.damping_ratio,))])
+
+class AxisInputShaper:
+    def __init__(self, axis, config):
+        self.axis = axis
+        self.params = InputShaperParams(axis, config)
+        self.n, self.A, self.T = self.params.get_shaper()
+        self.saved = None
+    def get_name(self):
+        return 'shaper_' + self.axis
+    def get_shaper(self):
+        return self.n, self.A, self.T
+    def update(self, gcmd):
+        self.params.update(gcmd)
+        old_n, old_A, old_T = self.n, self.A, self.T
+        self.n, self.A, self.T = self.params.get_shaper()
+        return (old_n, old_A, old_T) != (self.n, self.A, self.T)
+    def set_shaper_kinematics(self, sk):
+        ffi_main, ffi_lib = chelper.get_ffi()
+        success = ffi_lib.input_shaper_set_shaper_params(
+                sk, bytes(self.axis), self.n, self.A, self.T) == 0
+        if not success:
+            self.disable_shaping()
+            ffi_lib.input_shaper_set_shaper_params(
+                    sk, bytes(self.axis), self.n, self.A, self.T)
+        return success
+    def get_step_generation_window(self):
+        ffi_main, ffi_lib = chelper.get_ffi()
+        return ffi_lib.input_shaper_get_step_generation_window(self.n,
+                                                               self.A, self.T)
+    def disable_shaping(self):
+        if self.saved is None and self.n:
+            self.saved = (self.n, self.A, self.T)
+        A, T = shaper_defs.get_none_shaper()
+        self.n, self.A, self.T = len(A), A, T
+    def enable_shaping(self):
+        if self.saved is None:
+            # Input shaper was not disabled
+            return
+        self.n, self.A, self.T = self.saved
+        self.saved = None
+    def report(self, gcmd):
+        info = ' '.join(["%s_%s:%s" % (key, self.axis, value)
+                         for (key, value) in self.params.get_status().items()])
+        gcmd.respond_info(info)
+
 class InputShaper:
     def __init__(self, config):
         self.printer = config.get_printer()
         self.printer.register_event_handler("klippy:connect", self.connect)
         self.toolhead = None
-        self.shapers = {s.name : s.init_func for s in shaper_defs.INPUT_SHAPERS}
-        self.damping_ratio_x = config.getfloat(
-                'damping_ratio_x', shaper_defs.DEFAULT_DAMPING_RATIO,
-                minval=0., maxval=1.)
-        self.damping_ratio_y = config.getfloat(
-                'damping_ratio_y', shaper_defs.DEFAULT_DAMPING_RATIO,
-                minval=0., maxval=1.)
-        self.shaper_freq_x = config.getfloat('shaper_freq_x', 0., minval=0.)
-        self.shaper_freq_y = config.getfloat('shaper_freq_y', 0., minval=0.)
-        self.shaper_type_x = config.get('shaper_type_x', 'mzv').lower()
-        if self.shaper_type_x not in self.shapers:
-            raise config.error(
-                    'Unsupported shaper type: %s' % (self.shaper_type_x,))
-        self.shaper_type_y = config.get('shaper_type_y', 'mzv').lower()
-        if self.shaper_type_y not in self.shapers:
-            raise config.error(
-                    'Unsupported shaper type: %s' % (self.shaper_type_y,))
-        self.saved_shaper_freq_x = self.saved_shaper_freq_y = 0.
+        self.shapers = [AxisInputShaper('x', config),
+                        AxisInputShaper('y', config)]
         self.stepper_kinematics = []
         self.orig_stepper_kinematics = []
         # Register gcode commands
@@ -37,6 +104,8 @@ class InputShaper:
         gcode.register_command("SET_INPUT_SHAPER",
                                self.cmd_SET_INPUT_SHAPER,
                                desc=self.cmd_SET_INPUT_SHAPER_help)
+    def get_shapers(self):
+        return self.shapers
     def connect(self):
         self.toolhead = self.printer.lookup_object("toolhead")
         kin = self.toolhead.get_kinematics()
@@ -54,92 +123,40 @@ class InputShaper:
             self.orig_stepper_kinematics.append(orig_sk)
         # Configure initial values
         self.old_delay = 0.
-        self._set_input_shaper(self.shaper_type_x, self.shaper_type_y,
-                               self.shaper_freq_x, self.shaper_freq_y,
-                               self.damping_ratio_x, self.damping_ratio_y)
-    def _get_shaper(self, shaper_type, shaper_freq, damping_ratio):
-        if not shaper_freq:
-            return shaper_defs.get_none_shaper()
-        A, T = self.shapers[shaper_type](shaper_freq, damping_ratio)
-        return len(A), A, T
-    def _set_input_shaper(self, shaper_type_x, shaper_type_y
-                          , shaper_freq_x, shaper_freq_y
-                          , damping_ratio_x, damping_ratio_y):
-        if (shaper_type_x != self.shaper_type_x
-                or shaper_type_y != self.shaper_type_y):
-            self.toolhead.flush_step_generation()
-        n_x, A_x, T_x = self._get_shaper(
-                shaper_type_x, shaper_freq_x, damping_ratio_x)
-        n_y, A_y, T_y = self._get_shaper(
-                shaper_type_y, shaper_freq_y, damping_ratio_y)
-        ffi_main, ffi_lib = chelper.get_ffi()
-        new_delay = max(
-                ffi_lib.input_shaper_get_step_generation_window(n_x, A_x, T_x),
-                ffi_lib.input_shaper_get_step_generation_window(n_y, A_y, T_y))
+        self._update_input_shaping(error=self.printer.config_error)
+    def _update_input_shaping(self, error=None):
+        self.toolhead.flush_step_generation()
+        new_delay = max([s.get_step_generation_window() for s in self.shapers])
         self.toolhead.note_step_generation_scan_time(new_delay,
                                                      old_delay=self.old_delay)
-        self.old_delay = new_delay
-        self.shaper_type_x = shaper_type_x
-        self.shaper_type_y = shaper_type_y
-        self.shaper_freq_x = shaper_freq_x
-        self.shaper_freq_y = shaper_freq_y
-        self.damping_ratio_x = damping_ratio_x
-        self.damping_ratio_y = damping_ratio_y
+        failed = []
         for sk in self.stepper_kinematics:
-            ffi_lib.input_shaper_set_shaper_params(
-                    sk, len(A_x), A_x, T_x, len(A_y), A_y, T_y)
+            for shaper in self.shapers:
+                if shaper in failed:
+                    continue
+                if not shaper.set_shaper_kinematics(sk):
+                    failed.append(shaper)
+        if failed:
+            error = error or self.printer.command_error
+            raise error("Failed to configure shaper(s) %s with given parameters"
+                        % (', '.join([s.get_name() for s in failed])))
     def disable_shaping(self):
-        if (self.saved_shaper_freq_x or self.saved_shaper_freq_y) and not (
-                self.shaper_freq_x or self.shaper_freq_y):
-            # Input shaper is already disabled
-            return
-        self.saved_shaper_freq_x = self.shaper_freq_x
-        self.saved_shaper_freq_y = self.shaper_freq_y
-        self._set_input_shaper(self.shaper_type_x, self.shaper_type_y, 0., 0.,
-                               self.damping_ratio_x, self.damping_ratio_y)
+        for shaper in self.shapers:
+            shaper.disable_shaping()
+        self._update_input_shaping()
     def enable_shaping(self):
-        saved = self.saved_shaper_freq_x or self.saved_shaper_freq_y
-        if saved:
-            self._set_input_shaper(self.shaper_type_x, self.shaper_type_y,
-                                   self.saved_shaper_freq_x,
-                                   self.saved_shaper_freq_y,
-                                   self.damping_ratio_x, self.damping_ratio_y)
-        self.saved_shaper_freq_x = self.saved_shaper_freq_y = 0.
-        return saved
+        for shaper in self.shapers:
+            shaper.enable_shaping()
+        self._update_input_shaping()
     cmd_SET_INPUT_SHAPER_help = "Set cartesian parameters for input shaper"
     def cmd_SET_INPUT_SHAPER(self, gcmd):
-        damping_ratio_x = gcmd.get_float(
-                'DAMPING_RATIO_X', self.damping_ratio_x, minval=0., maxval=1.)
-        damping_ratio_y = gcmd.get_float(
-                'DAMPING_RATIO_Y', self.damping_ratio_y, minval=0., maxval=1.)
-        shaper_freq_x = gcmd.get_float(
-                'SHAPER_FREQ_X', self.shaper_freq_x, minval=0.)
-        shaper_freq_y = gcmd.get_float(
-                'SHAPER_FREQ_Y', self.shaper_freq_y, minval=0.)
-
-        shaper_type = gcmd.get('SHAPER_TYPE', None)
-        if shaper_type is None:
-            shaper_type_x = gcmd.get(
-                    'SHAPER_TYPE_X', self.shaper_type_x).lower()
-            shaper_type_y = gcmd.get(
-                    'SHAPER_TYPE_Y', self.shaper_type_y).lower()
-        else:
-            shaper_type_x = shaper_type_y = shaper_type.lower()
-        if shaper_type_x not in self.shapers:
-            raise gcmd.error('Unsupported shaper type: %s' % (shaper_type_x,))
-        if shaper_type_y not in self.shapers:
-            raise gcmd.error('Unsupported shaper type: %s' % (shaper_type_y,))
-
-        self._set_input_shaper(shaper_type_x, shaper_type_y,
-                               shaper_freq_x, shaper_freq_y,
-                               damping_ratio_x, damping_ratio_y)
-
-        gcmd.respond_info("shaper_type_x:%s shaper_type_y:%s "
-                          "shaper_freq_x:%.3f shaper_freq_y:%.3f "
-                          "damping_ratio_x:%.6f damping_ratio_y:%.6f"
-                          % (self.shaper_type_x, self.shaper_type_y,
-                             self.shaper_freq_x, self.shaper_freq_y,
-                             self.damping_ratio_x, self.damping_ratio_y))
+        updated = False
+        for shaper in self.shapers:
+            updated |= shaper.update(gcmd)
+        if updated:
+            self._update_input_shaping()
+        for shaper in self.shapers:
+            shaper.report(gcmd)
 
 def load_config(config):
     return InputShaper(config)