axis_twist_compensation: Implement Y-axis support

This commit implements support for the Y-axis in the axis_twist_compensation
module. This update enables the module to handle corrections for printers
with a twisted Y rail.

Signed-off-by: Jorge Apaza Merma <yochiwarez@gmail.com>

1 files changed, 140 insertions, 45 deletions

diff --git a/klippy/extras/axis_twist_compensation.py b/klippy/extras/axis_twist_compensation.py
index e7aad52c..17437a97 100644
--- a/klippy/extras/axis_twist_compensation.py
+++ b/klippy/extras/axis_twist_compensation.py
@@ -23,18 +23,27 @@ class AxisTwistCompensation:
         self.horizontal_move_z = config.getfloat('horizontal_move_z',
                                                  DEFAULT_HORIZONTAL_MOVE_Z)
         self.speed = config.getfloat('speed', DEFAULT_SPEED)
-        self.calibrate_start_x = config.getfloat('calibrate_start_x')
-        self.calibrate_end_x = config.getfloat('calibrate_end_x')
-        self.calibrate_y = config.getfloat('calibrate_y')
+        self.calibrate_start_x = config.getfloat('calibrate_start_x',
+                                                default=None)
+        self.calibrate_end_x = config.getfloat('calibrate_end_x', default=None)
+        self.calibrate_y = config.getfloat('calibrate_y', default=None)
         self.z_compensations = config.getlists('z_compensations',
                                                default=[], parser=float)
         self.compensation_start_x = config.getfloat('compensation_start_x',
                                                     default=None)
-        self.compensation_end_x = config.getfloat('compensation_start_y',
+        self.compensation_end_x = config.getfloat('compensation_end_x',
                                                   default=None)
 
-        self.m = None
-        self.b = None
+        self.calibrate_start_y = config.getfloat('calibrate_start_y',
+                                                  default=None)
+        self.calibrate_end_y = config.getfloat('calibrate_end_y', default=None)
+        self.calibrate_x = config.getfloat('calibrate_x', default=None)
+        self.compensation_start_y = config.getfloat('compensation_start_y',
+                                                    default=None)
+        self.compensation_end_y = config.getfloat('compensation_end_y',
+                                                  default=None)
+        self.zy_compensations = config.getlists('zy_compensations',
+                                                default=[], parser=float)
 
         # setup calibrater
         self.calibrater = Calibrater(self, config)
@@ -43,28 +52,46 @@ class AxisTwistCompensation:
                                             self._update_z_compensation_value)
 
     def _update_z_compensation_value(self, pos):
-        if not self.z_compensations:
-            return
+        if self.z_compensations:
+            pos[2] += self._get_interpolated_z_compensation(
+                pos[0], self.z_compensations,
+                self.compensation_start_x,
+                self.compensation_end_x
+                )
+
+        if self.zy_compensations:
+            pos[2] += self._get_interpolated_z_compensation(
+                pos[1], self.zy_compensations,
+                self.compensation_start_y,
+                self.compensation_end_y
+                )
+
+    def _get_interpolated_z_compensation(
+            self, coord, z_compensations,
+            comp_start,
+            comp_end
+            ):
 
-        x_coord = pos[0]
-        z_compensations = self.z_compensations
         sample_count = len(z_compensations)
-        spacing = ((self.calibrate_end_x - self.calibrate_start_x)
+        spacing = ((comp_end - comp_start)
                    / (sample_count - 1))
-        interpolate_t = (x_coord - self.calibrate_start_x) / spacing
+        interpolate_t = (coord - comp_start) / spacing
         interpolate_i = int(math.floor(interpolate_t))
         interpolate_i = bed_mesh.constrain(interpolate_i, 0, sample_count - 2)
         interpolate_t -= interpolate_i
         interpolated_z_compensation = bed_mesh.lerp(
             interpolate_t, z_compensations[interpolate_i],
             z_compensations[interpolate_i + 1])
-        pos[2] += interpolated_z_compensation
-
-    def clear_compensations(self):
-        self.z_compensations = []
-        self.m = None
-        self.b = None
+        return interpolated_z_compensation
 
+    def clear_compensations(self, axis=None):
+        if axis is None:
+            self.z_compensations = []
+            self.zy_compensations = []
+        elif axis == 'X':
+            self.z_compensations = []
+        elif axis == 'Y':
+            self.zy_compensations = []
 
 class Calibrater:
     def __init__(self, compensation, config):
@@ -80,10 +107,14 @@ class Calibrater:
                                             self._handle_connect)
         self.speed = compensation.speed
         self.horizontal_move_z = compensation.horizontal_move_z
-        self.start_point = (compensation.calibrate_start_x,
+        self.x_start_point = (compensation.calibrate_start_x,
                             compensation.calibrate_y)
-        self.end_point = (compensation.calibrate_end_x,
+        self.x_end_point = (compensation.calibrate_end_x,
                           compensation.calibrate_y)
+        self.y_start_point = (compensation.calibrate_x,
+                            compensation.calibrate_start_y)
+        self.y_end_point = (compensation.calibrate_x,
+                            compensation.calibrate_end_y)
         self.results = None
         self.current_point_index = None
         self.gcmd = None
@@ -119,20 +150,75 @@ class Calibrater:
     def cmd_AXIS_TWIST_COMPENSATION_CALIBRATE(self, gcmd):
         self.gcmd = gcmd
         sample_count = gcmd.get_int('SAMPLE_COUNT', DEFAULT_SAMPLE_COUNT)
+        axis = gcmd.get('AXIS', 'X')
 
         # check for valid sample_count
-        if sample_count is None or sample_count < 2:
+        if sample_count < 2:
             raise self.gcmd.error(
                 "SAMPLE_COUNT to probe must be at least 2")
 
-        # clear the current config
-        self.compensation.clear_compensations()
+        # calculate the points to put the probe at, returned as a list of tuples
+        nozzle_points = []
+
+        if axis == 'X':
+
+            self.compensation.clear_compensations('X')
+
+            if not all([
+                self.x_start_point[0],
+                self.x_end_point[0],
+                self.x_start_point[1]
+                ]):
+                raise self.gcmd.error(
+                    """AXIS_TWIST_COMPENSATION for X axis requires
+                    calibrate_start_x, calibrate_end_x and calibrate_y
+                    to be defined
+                    """
+                    )
+
+            start_point = self.x_start_point
+            end_point = self.x_end_point
+
+            x_axis_range = end_point[0] - start_point[0]
+            interval_dist = x_axis_range / (sample_count - 1)
+
+            for i in range(sample_count):
+                x = start_point[0] + i * interval_dist
+                y = start_point[1]
+                nozzle_points.append((x, y))
+
+        elif axis == 'Y':
+
+            self.compensation.clear_compensations('Y')
+
+            if not all([
+                self.y_start_point[0],
+                self.y_end_point[0],
+                self.y_start_point[1]
+                ]):
+                raise self.gcmd.error(
+                    """AXIS_TWIST_COMPENSATION for Y axis requires
+                    calibrate_start_y, calibrate_end_y and calibrate_x
+                    to be defined
+                    """
+                    )
+
+            start_point = self.y_start_point
+            end_point = self.y_end_point
+
+            y_axis_range = end_point[1] - start_point[1]
+            interval_dist = y_axis_range / (sample_count - 1)
+
+            for i in range(sample_count):
+                x = start_point[0]
+                y = start_point[1] + i * interval_dist
+                nozzle_points.append((x, y))
+
+        else:
+            raise self.gcmd.error(
+                "AXIS_TWIST_COMPENSATION_CALIBRATE: "
+                "Invalid axis.")
 
-        # calculate some values
-        x_range = self.end_point[0] - self.start_point[0]
-        interval_dist = x_range / (sample_count - 1)
-        nozzle_points = self._calculate_nozzle_points(sample_count,
-                                                      interval_dist)
         probe_points = self._calculate_probe_points(
             nozzle_points, self.probe_x_offset, self.probe_y_offset)
 
@@ -142,17 +228,9 @@ class Calibrater:
         # begin calibration
         self.current_point_index = 0
         self.results = []
+        self.current_axis = axis
         self._calibration(probe_points, nozzle_points, interval_dist)
 
-    def _calculate_nozzle_points(self, sample_count, interval_dist):
-        # calculate the points to put the probe at, returned as a list of tuples
-        nozzle_points = []
-        for i in range(sample_count):
-            x = self.start_point[0] + i * interval_dist
-            y = self.start_point[1]
-            nozzle_points.append((x, y))
-        return nozzle_points
-
     def _calculate_probe_points(self, nozzle_points,
         probe_x_offset, probe_y_offset):
         # calculate the points to put the nozzle at
@@ -238,14 +316,31 @@ class Calibrater:
         configfile = self.printer.lookup_object('configfile')
         values_as_str = ', '.join(["{:.6f}".format(x)
                                    for x in self.results])
-        configfile.set(self.configname, 'z_compensations', values_as_str)
-        configfile.set(self.configname, 'compensation_start_x',
-                       self.start_point[0])
-        configfile.set(self.configname, 'compensation_end_x',
-                       self.end_point[0])
-        self.compensation.z_compensations = self.results
-        self.compensation.compensation_start_x = self.start_point[0]
-        self.compensation.compensation_end_x = self.end_point[0]
+
+        if(self.current_axis == 'X'):
+
+            configfile.set(self.configname, 'z_compensations', values_as_str)
+            configfile.set(self.configname, 'compensation_start_x',
+                        self.x_start_point[0])
+            configfile.set(self.configname, 'compensation_end_x',
+                        self.x_end_point[0])
+
+            self.compensation.z_compensations = self.results
+            self.compensation.compensation_start_x = self.x_start_point[0]
+            self.compensation.compensation_end_x = self.x_end_point[0]
+
+        elif(self.current_axis == 'Y'):
+
+            configfile.set(self.configname, 'zy_compensations', values_as_str)
+            configfile.set(self.configname, 'compensation_start_y',
+                        self.y_start_point[1])
+            configfile.set(self.configname, 'compensation_end_y',
+                        self.y_end_point[1])
+
+            self.compensation.zy_compensations = self.results
+            self.compensation.compensation_start_y = self.y_start_point[1]
+            self.compensation.compensation_end_y = self.y_end_point[1]
+
         self.gcode.respond_info(
             "AXIS_TWIST_COMPENSATION state has been saved "
             "for the current session.  The SAVE_CONFIG command will "