tmc2130: add SET_TMC_CURRENT and SET_TMC_FIELD commands (#1419)

Signed-off-by: Frank Kang <amblidex@outlook.com>

1 files changed, 54 insertions, 1 deletions

diff --git a/klippy/extras/tmc2208.py b/klippy/extras/tmc2208.py
index a6f954c5..4cb67e3e 100644
--- a/klippy/extras/tmc2208.py
+++ b/klippy/extras/tmc2208.py
@@ -270,9 +270,15 @@ class TMC2208:
         # Add DUMP_TMC, INIT_TMC command
         gcode = self.printer.lookup_object("gcode")
         gcode.register_mux_command(
+            "SET_TMC_CURRENT", "STEPPER", self.name,
+            self.cmd_SET_TMC_CURRENT, desc=self.cmd_SET_TMC_CURRENT_help)
+        gcode.register_mux_command(
             "DUMP_TMC", "STEPPER", self.name,
             self.cmd_DUMP_TMC, desc=self.cmd_DUMP_TMC_help)
         gcode.register_mux_command(
+            "SET_TMC_FIELD", "STEPPER", self.name,
+            self.cmd_SET_TMC_FIELD, desc=self.cmd_SET_TMC_FIELD_help)
+        gcode.register_mux_command(
             "INIT_TMC", "STEPPER", self.name,
             self.cmd_INIT_TMC, desc=self.cmd_INIT_TMC_help)
         # Setup basic register values
@@ -282,7 +288,8 @@ class TMC2208:
         self.fields.set_field("pdn_disable", True)
         self.fields.set_field("mstep_reg_select", True)
         self.fields.set_field("multistep_filt", True)
-        vsense, irun, ihold = tmc2130.get_config_current(config)
+        vsense, irun, ihold, self.sense_resistor = \
+            tmc2130.get_config_current(config)
         self.fields.set_field("vsense", vsense)
         self.fields.set_field("IHOLD", ihold)
         self.fields.set_field("IRUN", irun)
@@ -351,6 +358,40 @@ class TMC2208:
     def get_phase(self):
         mscnt = self.fields.get_field("MSCNT", self.get_register("MSCNT"))
         return mscnt >> self.fields.get_field("MRES")
+    cmd_SET_TMC_CURRENT_help = "Set the current of a TMC2208 driver"
+    def cmd_SET_TMC_CURRENT(self, params):
+        gcode = self.printer.lookup_object('gcode')
+        vsense = bool(self.fields.get_field("vsense"))
+        if 'HOLDCURRENT' in params:
+            hold_current = gcode.get_float(
+                'HOLDCURRENT', params, above=0., maxval=2.)
+        else:
+            hold_current = tmc2130.bits_to_current(
+                    self.fields.get_field("IHOLD"),
+                    self.sense_resistor,
+                    vsense)
+        if 'CURRENT' in params:
+            run_current = gcode.get_float(
+                'CURRENT', params, minval=hold_current, maxval=2.)
+        else:
+            run_current = tmc2130.bits_to_current(
+                    self.fields.get_field("IRUN"),
+                    self.sense_resistor,
+                    vsense)
+        if 'HOLDCURRENT' in params or 'CURRENT' in params:
+            vsense_calc, irun, ihold = tmc2130.calc_current_config(run_current,
+                                              hold_current, self.sense_resistor)
+            self.printer.lookup_object('toolhead').wait_moves()
+            if (vsense_calc != vsense):
+                self.fields.set_field("vsense", vsense_calc)
+                self.set_register("CHOPCONF", self.regs["CHOPCONF"])
+            self.fields.set_field("IHOLD", ihold)
+            self.fields.set_field("IRUN", irun)
+            self.set_register("IHOLD_IRUN", self.regs["IHOLD_IRUN"])
+        else:
+            gcode.respond_info(
+                "Run Current: %0.2fA Hold Current: %0.2fA"
+                % (run_current, hold_current))
     cmd_DUMP_TMC_help = "Read and display TMC stepper driver registers"
     def cmd_DUMP_TMC(self, params):
         self.printer.lookup_object('toolhead').get_last_move_time()
@@ -377,6 +418,18 @@ class TMC2208:
         logging.info("INIT_TMC 2208 %s", self.name)
         self.printer.lookup_object('toolhead').wait_moves()
         self._init_registers()
+    cmd_SET_TMC_FIELD_help = "Set a register field of a TMC2208 driver"
+    def cmd_SET_TMC_FIELD(self, params):
+        gcode = self.printer.lookup_object('gcode')
+        if ('FIELD' not in params or
+            'VALUE' not in params):
+            raise gcode.error("Invalid command format")
+        field = gcode.get_str('FIELD', params)
+        reg = self.fields.field_to_register[field]
+        value = gcode.get_int('VALUE', params)
+        self.fields.set_field(field, value)
+        self.printer.lookup_object('toolhead').wait_moves()
+        self.set_register(reg, self.regs[reg])
 
 def load_config_prefix(config):
     return TMC2208(config)