<commit_before>from __future__ import print_function
import numpy as np
import scipy.stats, scipy.optimize
import acq4.pyqtgraph as pg
from acq4.Manager import getManager
from acq4.util import Qt
from acq4.util.target import Target


class CalibrationWindow(Qt.QWidget):
    def __init__(self, device):
        self.dev = device
        self._cammod = None
        self._camdev = None
        self.transform = None

        Qt.QWidget.__init__(self)
        self.resize(600, 300)
        self.setWindowTitle("Calibration: %s" % device.name())

        self.layout = Qt.QGridLayout()
        self.setLayout(self.layout)

        # tree columns:
        #   stage x, y, z   global x, y, z   error
        self.pointTree = Qt.QTreeWidget()
        self.pointTree.setHeaderLabels(['stage pos', 'parent pos', 'error'])
        self.pointTree.setColumnCount(3)
        self.layout.addWidget(self.pointTree, 0, 0)
        self.pointTree.setColumnWidth(0, 200)
        self.pointTree.setColumnWidth(1, 200)

        self.btnPanel = Qt.QWidget()
        self.btnPanelLayout = Qt.QHBoxLayout()
        self.btnPanelLayout.setContentsMargins(0, 0, 0, 0)
        self.btnPanel.setLayout(self.btnPanelLayout)
        self.layout.addWidget(self.btnPanel, 1, 0)

        self.addPointBtn = Qt.QPushButton("add point")
        self.addPointBtn.setCheckable(True)
        self.btnPanelLayout.addWidget(self.addPointBtn)

        self.removePointBtn = Qt.QPushButton("remove point")
        self.btnPanelLayout.addWidget(self.removePointBtn)
        
        self.saveBtn = Qt.QPushButton("save calibration")
        self.btnPanelLayout.addWidget(self.saveBtn)

        self.addPointBtn.toggled.connect(self.addPointToggled)
        self.removePointBtn.clicked.connect(self.removePointClicked)
        self.saveBtn.clicked.connect(self.saveClicked)

        # more controls:
        #    Show calibration points (in camera module)
        #    Force orthogonal axes: xy, xz, yz

        self.loadCalibrationFromDevice()

        cam = self.getCameraDevice()
        cam.sigGlobalTransformChanged.connect(self.cameraTransformChanged)

    def addPointToggled(self):
        cammod = self.getCameraModule()
        if self.addPointBtn.isChecked():
            cammod.window().getView().scene().sigMouseClicked.connect(self.cameraModuleClicked)
            self.addPointBtn.setText("click new point..")
        else:
            pg.disconnect(cammod.window().getView().scene().sigMouseClicked, self.cameraModuleClicked)
            self.addPointBtn.setText("add point")

    def cameraModuleClicked(self, ev):
        if ev.button() != Qt.Qt.LeftButton:
            return

        camera = self.getCameraDevice()
        cameraPos = camera.mapToGlobal([0, 0, 0])

        globalPos = self._cammod.window().getView().mapSceneToView(ev.scenePos())
        globalPos = [globalPos.x(), globalPos.y(), cameraPos[2]]
        parentDev = self.dev.parentDevice()
        if parentDev is None:
            parentPos = globalPos
        else:
            parentPos = parentDev.mapFromGlobal(globalPos)

        stagePos = self.dev.getPosition()

        self.calibration['points'].append((stagePos, parentPos))
        item = self._addCalibrationPoint(stagePos, parentPos)

        target = Target(movable=False)
        self._cammod.window().addItem(target)
        target.setPos(pg.Point(globalPos[:2]))
        target.setDepth(globalPos[2])
        target.setFocusDepth(globalPos[2])
        item.target = target

        self.addPointBtn.setChecked(False)
        self.recalculate()
        self.saveBtn.setText("*save calibration*")

    def cameraTransformChanged(self):
        cam = self.getCameraDevice()
        fdepth = cam.mapToGlobal([0, 0, 0])[2]

        items = [self.pointTree.topLevelItem(i) for i in range(self.pointTree.topLevelItemCount())]
        for item in items:
            if item.target is None:
                continue
            item.target.setFocusDepth(fdepth)

    def removePointClicked(self):
        sel = self.pointTree.selectedItems()[0]
        index = self.pointTree.indexOfTopLevelItem(sel)
        self.pointTree.takeTopLevelItem(index)
        if sel.target is not None:
            sel.target.scene().removeItem(sel.target)
        items = [self.pointTree.topLevelItem(i) for i in range(self.pointTree.topLevelItemCount())]
        self.calibration['points'] = [(item.stagePos, item.parentPos) for item in items]
        self.recalculate()
        self.saveBtn.setText("*save calibration*")

    def saveClicked(self):
        self.saveCalibrationToDevice()

    def loadCalibrationFromDevice(self):
        self.calibration = self.dev.readConfigFile('calibration')
        self.calibration.setdefault('points', [])
        for stagePos, parentPos in self.calibration['points']:
            self._addCalibrationPoint(stagePos, parentPos)
        self.recalculate()

    def saveCalibrationToDevice(self):
        self.recalculate()
        self.calibration['transform'] = None if self.transform is None else [list(row) for row in self.transform.matrix()]
        self.dev.writeConfigFile(self.calibration, 'calibration')
        self.saveBtn.setText("save calibration")

    def _addCalibrationPoint(self, stagePos, parentPos):
        item = Qt.QTreeWidgetItem(["%0.3g, %0.3g, %0.3g" % tuple(stagePos), "%0.3g, %0.3g, %0.3g" % tuple(parentPos), ""])
        self.pointTree.addTopLevelItem(item)
        item.stagePos = stagePos
        item.parentPos = parentPos
        item.target = None
        return item

    def recalculate(self):
        # identity affine axis transform matrix

        npts = len(self.calibration['points'])

        # Need at least 4 points to generate a calibration
        if npts < 4:
            for i in range(npts):
                item = self.pointTree.topLevelItem(i)
                item.setText(2, "")
            self.transform = None
            return

        stagePos = np.empty((npts, 3))
        parentPos = np.empty((npts, 3))
        for i, pt in enumerate(self.calibration['points']):
            stagePos[i] = pt[0]
            parentPos[i] = pt[1]

        def mapPoint(axisTr, stagePos, localPos):
            # given a stage position and axis transform, map from localPos to parent coordinate system
            if isinstance(axisTr, np.ndarray):
                m = np.eye(4)
                m[:3] = axisTr.reshape(3, 4)
                axisTr = pg.Transform3D(m)
            st = self.dev._makeStageTransform(stagePos, axisTr)[0]
            tr = pg.Transform3D(self.dev.baseTransform() * st)
            return tr.map(localPos)

        def mapError(axisTr, stagePos, parentPos):
            # Goal is to map origin to parent position correctly
            return [mapPoint(axisTr, sp, [0, 0, 0]) - pp for sp, pp in zip(stagePos, parentPos)]<commit_msg>initial code for independent axis calibration<commit_after>from __future__ import print_function
import numpy as np
import scipy.stats, scipy.optimize
import acq4.pyqtgraph as pg
from acq4.Manager import getManager
from acq4.util import Qt
from acq4.util.target import Target


class CalibrationWindow(Qt.QWidget):
    def __init__(self, device):
        self.dev = device
        self._cammod = None
        self._camdev = None
        self.transform = None

        Qt.QWidget.__init__(self)
        self.resize(600, 300)
        self.setWindowTitle("Calibration: %s" % device.name())

        self.layout = Qt.QGridLayout()
        self.setLayout(self.layout)

        # tree columns:
        #   stage x, y, z   global x, y, z   error
        self.pointTree = Qt.QTreeWidget()
        self.pointTree.setHeaderLabels(['stage pos', 'parent pos', 'error'])
        self.pointTree.setColumnCount(3)
        self.layout.addWidget(self.pointTree, 0, 0)
        self.pointTree.setColumnWidth(0, 200)
        self.pointTree.setColumnWidth(1, 200)

        self.btnPanel = Qt.QWidget()
        self.btnPanelLayout = Qt.QHBoxLayout()
        self.btnPanelLayout.setContentsMargins(0, 0, 0, 0)
        self.btnPanel.setLayout(self.btnPanelLayout)
        self.layout.addWidget(self.btnPanel, 1, 0)

        self.addPointBtn = Qt.QPushButton("add point")
        self.addPointBtn.setCheckable(True)
        self.btnPanelLayout.addWidget(self.addPointBtn)

        self.removePointBtn = Qt.QPushButton("remove point")
        self.btnPanelLayout.addWidget(self.removePointBtn)
        
        self.saveBtn = Qt.QPushButton("save calibration")
        self.btnPanelLayout.addWidget(self.saveBtn)

        self.addPointBtn.toggled.connect(self.addPointToggled)
        self.removePointBtn.clicked.connect(self.removePointClicked)
        self.saveBtn.clicked.connect(self.saveClicked)

        # more controls:
        #    Show calibration points (in camera module)
        #    Force orthogonal axes: xy, xz, yz

        self.loadCalibrationFromDevice()

        cam = self.getCameraDevice()
        cam.sigGlobalTransformChanged.connect(self.cameraTransformChanged)

    def addPointToggled(self):
        cammod = self.getCameraModule()
        if self.addPointBtn.isChecked():
            cammod.window().getView().scene().sigMouseClicked.connect(self.cameraModuleClicked)
            self.addPointBtn.setText("click new point..")
        else:
            pg.disconnect(cammod.window().getView().scene().sigMouseClicked, self.cameraModuleClicked)
            self.addPointBtn.setText("add point")

    def cameraModuleClicked(self, ev):
        if ev.button() != Qt.Qt.LeftButton:
            return

        camera = self.getCameraDevice()
        cameraPos = camera.mapToGlobal([0, 0, 0])

        globalPos = self._cammod.window().getView().mapSceneToView(ev.scenePos())
        globalPos = [globalPos.x(), globalPos.y(), cameraPos[2]]
        parentDev = self.dev.parentDevice()
        if parentDev is None:
            parentPos = globalPos
        else:
            parentPos = parentDev.mapFromGlobal(globalPos)

        stagePos = self.dev.getPosition()

        self.calibration['points'].append((stagePos, parentPos))
        item = self._addCalibrationPoint(stagePos, parentPos)

        target = Target(movable=False)
        self._cammod.window().addItem(target)
        target.setPos(pg.Point(globalPos[:2]))
        target.setDepth(globalPos[2])
        target.setFocusDepth(globalPos[2])
        item.target = target

        self.addPointBtn.setChecked(False)
        self.recalculate()
        self.saveBtn.setText("*save calibration*")

    def cameraTransformChanged(self):
        cam = self.getCameraDevice()
        fdepth = cam.mapToGlobal([0, 0, 0])[2]

        items = [self.pointTree.topLevelItem(i) for i in range(self.pointTree.topLevelItemCount())]
        for item in items:
            if item.target is None:
                continue
            item.target.setFocusDepth(fdepth)

    def removePointClicked(self):
        sel = self.pointTree.selectedItems()[0]
        index = self.pointTree.indexOfTopLevelItem(sel)
        self.pointTree.takeTopLevelItem(index)
        if sel.target is not None:
            sel.target.scene().removeItem(sel.target)
        items = [self.pointTree.topLevelItem(i) for i in range(self.pointTree.topLevelItemCount())]
        self.calibration['points'] = [(item.stagePos, item.parentPos) for item in items]
        self.recalculate()
        self.saveBtn.setText("*save calibration*")

    def saveClicked(self):
        self.saveCalibrationToDevice()

    def loadCalibrationFromDevice(self):
        self.calibration = self.dev.readConfigFile('calibration')
        self.calibration.setdefault('points', [])
        for stagePos, parentPos in self.calibration['points']:
            self._addCalibrationPoint(stagePos, parentPos)
        self.recalculate()

    def saveCalibrationToDevice(self):
        self.recalculate()
        self.calibration['transform'] = None if self.transform is None else [list(row) for row in self.transform.matrix()]
        self.dev.writeConfigFile(self.calibration, 'calibration')
        self.saveBtn.setText("save calibration")

    def _addCalibrationPoint(self, stagePos, parentPos):
        item = Qt.QTreeWidgetItem(["%0.3g, %0.3g, %0.3g" % tuple(stagePos), "%0.3g, %0.3g, %0.3g" % tuple(parentPos), ""])
        self.pointTree.addTopLevelItem(item)
        item.stagePos = stagePos
        item.parentPos = parentPos
        item.target = None
        return item

    def recalculate(self):
        # identity affine axis transform matrix

        # method: user generates many calibration points that are all colinear along each of the
        # stage axes. In this way, we can independently determine the orientation of each stage axis,
        # and combine these into a full transformation matrix.

        npts = len(self.calibration['points'])

        # build arrays of calibration points
        stagePos = np.empty((npts, 3))
        parentPos = np.empty((npts, 3))
        for i, pt in enumerate(self.calibration['points']):
            stagePos[i] = pt[0]
            parentPos[i] = pt[1]

        def changeAxis(p1, p2):
            # Which single axis has changed between 2 points?
            diff = np.abs(p2-p1)
            axis = np.argmax(diff)
            if diff[axis] > 10e-6 and diff[(axis+1)%3] < 1e-6 and diff[(axis+2)%3] < 1e-6:
                return axis
            else:
                return None

        # find the point groupings for each axis
        axisPoints = [set(), set(), set()]
        currentAxis = None
        for i in range(1, npts):
            currentAxis = changeAxis(stagePos[i-1], stagePos[i])
            if currentAxis is None:
                continue
            axisPoints[currentAxis].add(i-1)
            axisPoints[currentAxis].add(i)
                
        stagePos = [stagePos[list(axisPoints[ax]), ax] for ax in (0,1,2)]
        parentPos = [parentPos[list(axisPoints[ax])] for ax in (0,1,2)]

        # find optimal linear mapping for each axis
        m = np.eye(4)
        for i in (0,1,2):
            for j in (0,1,2):
                line = scipy.stats.linregress(stagePos[j], parentPos[j][:,i])
                m[i, j] = line.slope

        self.transform = pg.Transform3D(m)

        # measure and display errors for each point
        def mapPoint(axisTr, stagePos, localPos):
            # given a stage position and axis transform, map from localPos to parent coordinate system
            if isinstance(axisTr, np.ndarray):
                m = np.eye(4)
                m[:3] = axisTr.reshape(3, 4)
                axisTr = pg.Transform3D(m)
            st = self.dev._makeStageTransform(stagePos, axisTr)[0]
            tr = pg.Transform3D(self.dev.baseTransform() * st)
            return tr.map(localPos)

        def mapError(axisTr, stagePos, parentPos):
            # Goal is to map origin to parent position correctly
            return [mapPoint(axisTr, sp, [0, 0, 0]) - pp for sp, pp in zip(stagePos, parentPos)]

        error = mapError(self.transform, stagePos, parentPos)
        for i in range(npts):
            item = self.pointTree.topLevelItem(i)
            dist = np.linalg.norm(error[i])
            item.setText(2, "%0.2f um  (%0.3g, %0.3g, %0.3g)" % (1e6*dist, error[i][0], error[i][1], error[i][2]))

        # send new transform to device
        self.dev._axisTransform = self.transform
        self.dev._inverseAxisTransform = None
        self.dev._updateTransform()

    def _recalculate(self):
        # identity affine axis transform matrix

        # method: given >= 4 arbitrarily-located calibration points, attempt to find the 
        # affine transform that best matches all points.

        # For any 4 point-pairs, we can get an exact solution that may not work as well for other points.
        # With more points, we can try to pick an optimal transform that minimizes errors, but this
        # turns out to be a tricky minimization problem.

        npts = len(self.calibration['points'])

        # Need at least 4 points to generate a calibration
        if npts < 4:
            for i in range(npts):
                item = self.pointTree.topLevelItem(i)
                item.setText(2, "")
            self.transform = None
            return

        stagePos = np.empty((npts, 3))
        parentPos = np.empty((npts, 3))
        for i, pt in enumerate(self.calibration['points']):
            stagePos[i] = pt[0]
            parentPos[i] = pt[1]

        def mapPoint(axisTr, stagePos, localPos):
            # given a stage position and axis transform, map from localPos to parent coordinate system
            if isinstance(axisTr, np.ndarray):
                m = np.eye(4)
                m[:3] = axisTr.reshape(3, 4)
                axisTr = pg.Transform3D(m)
            st = self.dev._makeStageTransform(stagePos, axisTr)[0]
            tr = pg.Transform3D(self.dev.baseTransform() * st)
            return tr.map(localPos)

        def mapError(axisTr, stagePos, parentPos):
            # Goal is to map origin to parent position correctly
            return [mapPoint(axisTr, sp, [0, 0, 0]) - pp for sp, pp in zip(stagePos, parentPos)]
{'fraction_non_alphanumeric': 0.10933115060804491, 'fraction_numerical': 0.011634705332086062, 'mean_word_length': 2.4306056959486564, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6180264', 'n_tokens_mistral': 5184, 'n_tokens_neox': 4970, 'n_words': 1132}