nanonisTCPIP.py

import socket
import struct
import numpy as np

class nanonisTCP:    
    def __init__(self, ip = '127.0.0.1', port = 6501, max_buf_size = 1024):
        """Initialize the NanonisTCPIP class with the IP and port."""
        self.ip = ip
        self.port = port
        self.sock = None
        self.max_buf_size = max_buf_size  # Default buffer size; you can adjust it as needed
        
    def connect(self):
        try:
            self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            self.sock.settimeout(5.0)
            self.sock.connect((self.ip, self.port))
            print(f"Connected to {self.ip}:{self.port}")
            return True
        except socket.timeout:
            raise TimeoutError(f"Connection to {self.ip}:{self.port} timed out")
            self.sock = None
            return False
        except socket.error as e:
            raise ConnectionError(f"Failed to connect to {self.ip}:{self.port}: {e}")
            self.sock = None
            return False

    def send_command(self, message):
        try:
            self.sock.settimeout(2.0)
            self.sock.sendall(bytes.fromhex(message))
        except socket.timeout:
            print("Client 1: Send operation timed out.")
        except socket.error as e:
            print(f"Client 1: Socket error during send: {e}")

    def receive_response(self, error_index=-1, keep_header = False):
        """
        Parameters
        error_index : index of 'error status' within the body. -1 skip check
        keep_header : if true: return entire response. if false: return body
        
        Returns
        response    : either header + body or body only (keep_header)
        
        """
        try:
            self.sock.settimeout(180.0)
            response = self.sock.recv(self.max_buf_size)                               # Read the response
        except socket.timeout:
            print("Client 1: Receive operation timed out.")
        except socket.error as e:
            print(f"Client 1: Socket error during receive: {e}")
        body_size = self.hex_to_int32(response[32:36])
        while(True): 
            if(len(response) == body_size + 40): break                          # body_size + header size (40)
            response += self.sock.recv(self.max_buf_size)
        
        if(error_index > -1): self.check_error(response[40:],error_index)       # error_index < 0 skips error check
        
        if(not keep_header):
            return response[40:]                                                # Header is fixed to 40 bytes - drop it
        
        return response

    def check_error(self,response,error_index):
        """
        Checks the response from nanonis for error messages

        Parameters
        response : response body (not inc. header) from nanonis (bytes)
        error_index : index of error status within the body

        Raises
        Exception   : error message returned from Nanonis

        """
        i = error_index                                                         # error_index points to start-byte in the body, which is after the 40-byte header
        error_status = self.hex_to_uint16(response[i:i+4])                      # error_status is 4 bytes long
        
        if(error_status):
            i += 8                                                              # index of error description is 8 bytes after error status
            error_description = response[i:].decode()                           # just grab from start index to the end of the message
            raise Exception(error_description)                                  # raise the exception
                
    def close_socket(self):
        """ Close the socket """
        if self.sock:
            self.sock.close()

    def hex_to_int32(self,h32):
        return struct.unpack("<i",struct.pack("I",int("0x"+h32.hex(),16)))[0]

    def to_hex(self,conv,num_bytes):
        if(conv >= 0): return hex(conv)[2:].zfill(2*num_bytes)
        if(conv < 0):  return hex((conv + (1 << 8*num_bytes)) % (1 << 8*num_bytes))[2:]

    def float64_to_hex(self,f64):
        # see https://stackoverflow.com/questions/23624212/how-to-convert-a-float-into-hex
        if(f64 == 0): return "0000000000000000"                                 # workaround for zero. look into this later
        return hex(struct.unpack('<Q', struct.pack('<d', f64))[0])[2:]

    def hex_to_float64(self, h64):
        return struct.unpack("<d", struct.pack("Q",int("0x"+h64.hex(), 16)))[0]

    def hex_to_uint16(self,h16):
        return struct.unpack("<H",struct.pack("H",int("0x"+h16.hex(),16)))[0]

    def hex_to_uint32(self,h32):
        return struct.unpack("<I", struct.pack("I",int("0x"+h32.hex(), 16)))[0]

    def hex_to_float32(self,h32):
        # see https://forum.inductiveautomation.com/t/ieee-754-standard-converting-64-bit-hex-to-decimal/9324/3
        return struct.unpack("<f", struct.pack("I",int("0x"+h32.hex(), 16)))[0]        
 
    def float32_to_hex(self,f32):
        # see https://stackoverflow.com/questions/23624212/how-to-convert-a-float-into-hex
        if(f32 == 0): return "00000000"                                         # workaround for zero. look into this later
        return hex(struct.unpack('<I', struct.pack('<f', f32))[0])[2:]          #

    def string_to_hex(self,string):
        return string.encode('utf-8').hex()

    def make_header(self, command_name, body_size, resp=True):
        """
        Parameters
        command_name : name of the Nanonis function
        body_size    : size of the message body in bytes
        resp         : tell nanonis to send a response. response contains error
                       message so will nearly always want to receive it

        Returns
        hex_rep : hex representation of the header string
        """ 
        hex_rep = command_name.encode('utf-8').hex()                            # command name
        hex_rep += "{0:#0{1}}".format(0,(64 - len(hex_rep)))                    # command name (fixed 32)
        hex_rep += self.to_hex(body_size, 4)                                    # Body size (fixed 4)
        hex_rep += self.to_hex(resp, 2)                                         # Send response (fixed 2)
        hex_rep += "{0:#0{1}}".format(0, 4)                                     # not used (fixed 2)
        return hex_rep

class FolMe:
    def __init__(self, nanonisTCP):
        self.nanonisTCP = nanonisTCP

    def XYPosSet(self, X, Y, Wait_end_of_move=True):
        """
        This function moves the tip to the specified X and Y target coordinates
        (in meters). It moves at the speed specified by the "Speed" parameter
        in the Follow Me mode of the Scan Control module. This function will 
        return when the tip reaches its destination or if the movement stops.

        Parameters
        X : Set x position (m)
        Y : Set y position (m)
        Wait_end_of_move : False: Selects whether the function  immediately
                           True: Waits until tip stops moving
        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('FolMe.XYPosSet', body_size=20)
        
        ## arguments
        hex_rep += self.nanonisTCP.float64_to_hex(X)
        hex_rep += self.nanonisTCP.float64_to_hex(Y)
        hex_rep += self.nanonisTCP.to_hex(Wait_end_of_move,4)
        
        self.nanonisTCP.send_command(hex_rep)
        message =  self.nanonisTCP.receive_response(0)
        return message

    def XYPosGet(self, Wait_for_newest_data=False):
        """
        Returns the X,Y tip coordinates

        Parameters
        Wait_for_newest_data (uint32): 

        Returns
        xpos : x position of the tip (m)
        ypos : y position of the tip (m)

        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('FolMe.XYPosGet', body_size=4)        
        # arguments
        hex_rep += self.nanonisTCP.to_hex(Wait_for_newest_data,4)        
        self.nanonisTCP.send_command(hex_rep)        
        ## Receive Response
        response = self.nanonisTCP.receive_response(16)        
        xpos = self.nanonisTCP.hex_to_float64(response[0:8])
        ypos = self.nanonisTCP.hex_to_float64(response[8:16])
        return (xpos,ypos)
    
    def SpeedSet(self, speed, custom_speed):
        """
        Configures the tip speed when moving in Follow Me

        Parameters
        speed (float32)       : sets the surface speed in Follow Me
        custom_speed (uint32) : True:  speed setting is custom speed
                                False: speed setting is scan speed
                                
        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('FolMe.SpeedSet', body_size=8)        
        # arguments
        hex_rep += self.nanonisTCP.float32_to_hex(speed)
        hex_rep += self.nanonisTCP.to_hex(custom_speed,4)        
        self.nanonisTCP.send_command(hex_rep)        
        # Receive Response (check for errors)
        self.nanonisTCP.receive_response(0)
    
    def SpeedGet(self):
        """
        Returns the tip speed when moving in Follow Me mode

        Returns
        -------
        speed (float32)       : surface speed in Follow Me mode
        custom_speed (uint32) : True:  speed setting is custom speed
                                False: speed setting is scam speed
                                
        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('FolMe.SpeedGet', body_size=0)        
        self.nanonisTCP.send_command(hex_rep)        
        # Receive Response
        response = self.nanonisTCP.receive_response(8)        
        speed        = self.nanonisTCP.hex_to_float32(response[0:4])
        custom_speed = self.nanonisTCP.hex_to_uint32(response[4:8]) > 0
        return (speed,custom_speed)
        
    def Stop(self):
        """
        Stops the tip movement in follow me mode

        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('FolMe.Stop', body_size=0)        
        self.nanonisTCP.send_command(hex_rep)        
        # Receive Response (check errors)
        self.nanonisTCP.receive_response(0)
        
class Current:
    """
    Nanonis Current Module
    """
    def __init__(self,nanonisTCP):
        self.nanonisTCP = nanonisTCP
    
    def Get(self):
        """
        Returns the tunnelling current value

        Returns
        -------
        current : Current value (A)

        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('Current.Get', body_size=0)        
        self.nanonisTCP.send_command(hex_rep)        
        response = self.nanonisTCP.receive_response(4)        
        current = self.nanonisTCP.hex_to_float32(response[0:4])        
        return current
    
    def Get100(self):
        """
        Returns the current value of the "Current 100" module

        Returns
        -------
        current100 : Current 100 value (A)

        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('Current.100Get', body_size=0)        
        self.nanonisTCP.send_command(hex_rep)        
        response = self.nanonisTCP.receive_response(4)        
        current100 = self.nanonisTCP.hex_to_float32(response[0:4])        
        return current100
    
    def BEEMGet(self):
        """
        Returns the BEEM current value of the corresponding module in a BEEM
        system

        Returns
        -------
        currentBEEM : Current BEEM value (A)

        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('Current.BEEMGet', body_size=0)        
        self.nanonisTCP.send_command(hex_rep)        
        response = self.nanonisTCP.receive_response(4)        
        currentBEEM = self.nanonisTCP.hex_to_float32(response[0:4])        
        return currentBEEM
    
    def GainSet(self,gain_index):
        """
        Sets the gain of the current amplifier

        Parameters
        ----------
        gain_index : The index out of the list of gains which can be retrieved
                     by the function Current.GainsGet

        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('Current.GainSet', body_size=2)        
        ## Arguments
        hex_rep += self.nanonisTCP.to_hex(gain_index,2)        
        self.nanonisTCP.send_command(hex_rep)        
        self.nanonisTCP.receive_response(0)
    
    def GainsGet(self):
        """
        Returns the selectable gains of the current amplifier and the index of 
        the selected one

        Returns
        -------
        gains      : array of selectable gains
        gain_index : index of the selected gain in gains array

        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('Current.GainsGet', body_size=0)        
        self.nanonisTCP.send_command(hex_rep)        
        response = self.nanonisTCP.receive_response()        
        # gains_size      = self.nanonisTCP.hex_to_int32(response[0:4])         # Not needed since
        number_of_gains = self.nanonisTCP.hex_to_int32(response[4:8])           # We know the number of gains
        
        idx   = 8
        gains = []
        for g in range(number_of_gains):
            size = self.nanonisTCP.hex_to_int32(response[idx:idx+4])            # And the size of each next gain
            idx += 4
            gain = response[idx:idx+size].decode()
            idx += size
            gains.append(gain)
        
        gain_index = self.nanonisTCP.hex_to_uint16(response[idx:idx+2])        
        return [gains,gain_index]
    
    def CalibrSet(self,calibration,offset):
        """
        Sets the calibration and offset of the selected gain in the current 
        module

        Parameters
        ----------
        calibration : calibration factor (A/V)
        offset      : offset (A)

        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('Current.CalibrSet', body_size=16)        
        ## Arguments
        hex_rep += self.nanonisTCP.float64_to_hex(calibration)
        hex_rep += self.nanonisTCP.float64_to_hex(offset)        
        self.nanonisTCP.send_command(hex_rep)        
        self.nanonisTCP.receive_response(0)
    
    def CalibrGet(self):
        """
        Gets the calibration and offset of the selected gain in the current 
        module

        Returns
        -------
        callibtation : calibration (A/V)
        offset       : offset (A)

        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('Current.CalibrGet', body_size=0)        
        self.nanonisTCP.send_command(hex_rep)        
        response = self.nanonisTCP.receive_response(16)        
        calibration = self.nanonisTCP.hex_to_float64(response[0:8])
        offset      = self.nanonisTCP.hex_to_float64(response[8:16])
        
        return [calibration,offset]

class ZCtrl:
    def __init__(self, nanonisTCP):
        self.nanonisTCP = nanonisTCP

    def ZPosSet(self,zpos):
        """
        Sets the Z position of the tip

        Parameters
        ----------
        zpos : Z position (m)

        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('ZCtrl.ZPosSet', body_size=4)
        
        ## Arguments
        hex_rep += self.nanonisTCP.float32_to_hex(zpos)        
        self.nanonisTCP.send_command(hex_rep)        
        self.nanonisTCP.receive_response(0)

    def ZPosGet(self):
        """
        Returns the current Z position of the tip

        Returns
        -------
        zpos : the current z position of the tip

        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('ZCtrl.ZPosGet', body_size=0)
        self.nanonisTCP.send_command(hex_rep)
        response = self.nanonisTCP.receive_response(4)
        zpos = self.nanonisTCP.hex_to_float32(response[0:4])        
        return zpos  
    
    def SetpntSet(self,setpoint):
        """
        Sets the stpoint of the Z-Controller

        Parameters
        ----------
        setpoint : setpoint of the z-controller (A)

        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('ZCtrl.SetpntSet', body_size=4)
        ## Arguments
        hex_rep += self.nanonisTCP.float32_to_hex(setpoint)
        self.nanonisTCP.send_command(hex_rep)
        self.nanonisTCP.receive_response(0)
    
    def SetpntGet(self):
        """
        Returns the setpoint current of the Z-Controller

        Returns
        -------
        setpoint : setpoint current of the z-controller (A)

        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('ZCtrl.SetpntGet', body_size=0)
        self.nanonisTCP.send_command(hex_rep)
        response = self.nanonisTCP.receive_response(4)
        setpoint = self.nanonisTCP.hex_to_float32(response[0:4])
        return setpoint

class Bias:
    def __init__(self, nanonisTCP):
        self.nanonisTCP = nanonisTCP
    
    def Set(self, bias):
        """
        Set the tip voltage bias

        Parameters
        bias (float32): bias (V)
        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('Bias.Set', body_size=4)        
        ## Arguments
        hex_rep += self.nanonisTCP.float32_to_hex(bias)                         # bias (float 32)        
        self.nanonisTCP.send_command(hex_rep)        
        self.nanonisTCP.receive_response(0)
    
    def Get(self):
        """
        Returns the tip bias

        Returns
        bias (float32): bias (V)

        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('Bias.Get', body_size=0)        
        self.nanonisTCP.send_command(hex_rep)        
        response = self.nanonisTCP.receive_response(4)        
        bias = self.nanonisTCP.hex_to_float32(response[0:4])
        return bias

class BiasSpectr:
    """
    Nanonis Bias Spectroscopy Module
    """
    def __init__(self,nanonisTCP):
        self.nanonisTCP = nanonisTCP
        
    def Open(self):
        """
        Opens the Bias Spectroscopy Module

        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('BiasSpectr.Open', body_size=0)        
        self.nanonisTCP.send_command(hex_rep)        
        self.nanonisTCP.receive_response(0)
    
    def Start(self,get_data=True,save_base_name=""):
        """
        Starts a bias spectroscopy in the Bias Spectroscopy module.
        
        Before using this function, select the channels to record in the Bias
        Spectroscopy module.

        Parameters
        ----------
        get_data        : defines if the function returns the spectroscopy data
                          True: return data from this function
                          False: don't return data
        save_base_name  : Base name used by the saved files. Empty string 
                          keeps settings unchanged in nanonis
                    

        Returns
        -------
        if get_data  = False, this function returns None
        
        if get_data != False, this function returns:
            
        data_dict{
            '<channel_name>' : data for this channel
            }
        parameters  : List of fixed parameters and parameters (in that order).
                      To see the names of the returned parameters, use the 
                      BiasSpectr.PropsGet function

        """
        body_size = 4 + 4                                                       # 4 bytes for get_data (uint32) and 4 bytes for save_base_name_string_size (int)
        body_size += int(len(self.nanonisTCP.string_to_hex(save_base_name))/2)   # Variable size depending on the save_base_name string
        
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('BiasSpectr.Start', body_size=body_size)
        
        save_base_name_string_size = len(save_base_name)
        
        ## arguments
        hex_rep += self.nanonisTCP.to_hex(get_data,4)
        hex_rep += self.nanonisTCP.to_hex(save_base_name_string_size,4)
        if(save_base_name_string_size > 0):
            hex_rep += self.nanonisTCP.string_to_hex(save_base_name)
        
        self.nanonisTCP.send_command(hex_rep)
        
        if(not get_data == 1):
            response = self.nanonisTCP.receive_response(0)
            return
        
        # Receive Response
        response = self.nanonisTCP.receive_response()
        
        # channels_names_size = self.nanonisTCP.hex_to_int32(response[0:4])     # Useless
        number_of_channels  = self.nanonisTCP.hex_to_int32(response[4:8])
        
        idx = 8
        channel_names = []
        for i in range(number_of_channels):
            channel_name_size = self.nanonisTCP.hex_to_int32(response[idx:idx+4])
            idx += 4
            channel_names.append(response[idx:idx + channel_name_size].decode())
            idx += channel_name_size
        
        data_rows = self.nanonisTCP.hex_to_int32(response[idx:idx+4])
        idx += 4
        data_cols = self.nanonisTCP.hex_to_int32(response[idx:idx+4])
        
        data_dict = {}
        for i in range(data_rows):
            data = []
            for j in range(data_cols):
                idx += 4
                data.append(self.nanonisTCP.hex_to_float32(response[idx:idx+4]))
            data_dict[channel_names[i]] = np.array(data)
        
        idx += 4
        parameters = []
        number_of_parameters = self.nanonisTCP.hex_to_int32(response[idx:idx+4])
        for i in range(number_of_parameters):
            idx += 4
            parameter = self.nanonisTCP.hex_to_float32(response[idx:idx+4])
            parameters.append(parameter)

        return {"data_dict"  : data_dict,
                "parameters" : parameters}
    
    def Stop(self):
        """
        Stops the current Bias Spectroscopy measurement.

        """
        hex_rep = self.nanonisTCP.make_header('BiasSpectr.Stop', body_size=0)
        self.nanonisTCP.send_command(hex_rep)        
        self.nanonisTCP.receive_response(0)

class DigLines:
    """
    Nanonis Digital Lines Module
    """
    def __init__(self,nanonisTCP):
        self.nanonisTCP = nanonisTCP
    
    def _parse_port_line(self, port_str):
        """
        Parse port string like 'A1', 'B3', 'C8', 'D2'

        Returns
        -------
        port_index : int
            0=Port A, 1=Port B, 2=Port C, 3=Port D
        line : int
            Digital line number from 1 to 8
        """
        if not isinstance(port_str, str):
            raise TypeError("port must be a string like 'A1', 'B3', 'C8', or 'D2'")

        port_str = port_str.strip().upper()

        if len(port_str) < 2:
            raise ValueError("port must look like 'A1', 'B3', 'C8', or 'D2'")

        port_letter = port_str[0]
        line_str = port_str[1:]

        port_map = {
            'A': 0,
            'B': 1,
            'C': 2,
            'D': 3
        }

        if port_letter not in port_map:
            raise ValueError("port letter must be one of A, B, C, D")

        try:
            line = int(line_str)
        except ValueError:
            raise ValueError("digital line must be an integer from 1 to 8")

        if not (1 <= line <= 8):
            raise ValueError("digital line must be between 1 and 8")

        return port_map[port_letter], line

    def OutStatusSet(self, port="D8", status=1):
        """
        Set the status on the selected digital output.

        Parameters
        ----------
        port : str
            Port and line as string, e.g. 'A1', 'B3', 'C8', 'D2'
        status : int
            0=Inactive, 1=Active
        wait_until_finished : bool, optional
            If True, wait until all pulses are generated before returning

        Returns
        -------
        response : bytes
            Response body from Nanonis
        """
        port_index, line = self._parse_port_line(port)

        # Body layout:
        # Port                    : uint32  -> 4 bytes
        # Digital lines size      : uint32   -> 4 bytes
        # StatusGet               : uint32   -> 4 bytes
        body_size = 4 + 4 + 4

        ## Make Header
        hex_rep = self.nanonisTCP.make_header('DigLines.OutStatusSet', body_size=body_size)

        ## Arguments
        hex_rep += self.nanonisTCP.to_hex(port_index, 4)# uint32
        hex_rep += self.nanonisTCP.to_hex(line, 4)      # uint32
        hex_rep += self.nanonisTCP.to_hex(status, 4)    # uint32

        self.nanonisTCP.send_command(hex_rep)
        response = self.nanonisTCP.receive_response(0)
        return response

    def Pulse(self, port="D8", pulse_width=0.05, pulse_pause=0.4, number_of_pulses=1, wait_until_finished=True):
        """
        Configures and starts the pulse generator on the selected digital output.

        Parameters
        ----------
        port : str
            Port and line as string, e.g. 'A1', 'B3', 'C8', 'D2'
        pulse_width : float
            Pulse width in seconds
        pulse_pause : float
            Pulse pause in seconds
        number_of_pulses : int
            Number of pulses, valid range 1..32767
        wait_until_finished : bool, optional
            If True, wait until all pulses are generated before returning

        Returns
        -------
        response : bytes
            Response body from Nanonis
        """
        port_index, line = self._parse_port_line(port)

        if number_of_pulses < 1 or number_of_pulses > 32767:
            raise ValueError("number_of_pulses must be between 1 and 32767")

        digital_lines = [line]
        digital_lines_size = len(digital_lines)

        # Body layout:
        # Port                    : uint16  -> 2 bytes
        # Digital lines size      : int32   -> 4 bytes
        # Digital lines array     : uint8[] -> N bytes
        # Pulse width             : float32 -> 4 bytes
        # Pulse pause             : float32 -> 4 bytes
        # Number of pulses        : int32   -> 4 bytes
        # Wait until finished     : uint32  -> 4 bytes
        body_size = 2 + 4 + digital_lines_size + 4 + 4 + 4 + 4

        ## Make Header
        hex_rep = self.nanonisTCP.make_header('DigLines.Pulse', body_size=body_size)

        ## Arguments
        hex_rep += self.nanonisTCP.to_hex(port_index, 2)               # uint16
        hex_rep += self.nanonisTCP.to_hex(digital_lines_size, 4)       # int32

        for dl in digital_lines:
            hex_rep += self.nanonisTCP.to_hex(dl, 1)                   # uint8 array element

        hex_rep += self.nanonisTCP.float32_to_hex(pulse_width)                   # float32
        hex_rep += self.nanonisTCP.float32_to_hex(pulse_pause)                   # float32
        hex_rep += self.nanonisTCP.to_hex(number_of_pulses, 4)         # int32
        hex_rep += self.nanonisTCP.to_hex(int(wait_until_finished), 4) # uint32

        self.nanonisTCP.send_command(hex_rep)
        if wait_until_finished:
            response = self.nanonisTCP.receive_response(0)
            return response

class AtomTrack:
    """
    Nanonis Atom Tracking Module
    """
    def __init__(self,nanonisTCP):
        self.nanonisTCP = nanonisTCP

    def _validate_control(self, at_control):
        """
        Valid AT control values:
            0 = Modulation
            1 = Controller
            2 = Drift Measurement
        """
        if at_control not in (0, 1, 2):
            raise ValueError("at_control must be 0 (Modulation), 1 (Controller), or 2 (Drift Measurement)")

    def _validate_status(self, status):
        """
        Valid status values:
            0 = Off
            1 = On
        """
        if status not in (0, 1):
            raise ValueError("status must be 0 (Off) or 1 (On)")

    def CtrlSet(self, at_control, status):
        """
        Turns the selected Atom Tracking control On or Off.

        Parameters
        ----------
        at_control : int
            0 = Modulation
            1 = Controller
            2 = Drift Measurement
        status : int
            0 = Off
            1 = On

        Returns
        -------
        message : bytes
            Response body from Nanonis
        """
        self._validate_control(at_control)
        self._validate_status(status)
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('AtomTrack.CtrlSet', body_size=4)
        ## Arguments
        hex_rep += self.nanonisTCP.to_hex(at_control, 2)   # uint16
        hex_rep += self.nanonisTCP.to_hex(status, 2)       # uint16
        self.nanonisTCP.send_command(hex_rep)
        message = self.nanonisTCP.receive_response(0)
        return message

    def StatusGet(self, at_control):
        """
        Returns the status of the selected Atom Tracking control.

        Parameters
        ----------
        at_control : int
            0 = Modulation
            1 = Controller
            2 = Drift Measurement

        Returns
        -------
        status : int
            0 = Off
            1 = On
        """
        self._validate_control(at_control)

        ## Make Header
        hex_rep = self.nanonisTCP.make_header('AtomTrack.StatusGet', body_size=2)
        ## Arguments
        hex_rep += self.nanonisTCP.to_hex(at_control, 2)   # uint16
        self.nanonisTCP.send_command(hex_rep)
        response = self.nanonisTCP.receive_response(2)
        status = self.nanonisTCP.hex_to_uint16(response[0:2])
        return status

    def QuickCompStart(self, at_control):
        """
        Starts the Tilt or Drift compensation.

        Parameters
        ----------
        at_control : int
            0 = Tilt compensation
            1 = Drift compensation

        Returns
        -------
        message : bytes
            Response body from Nanonis
        """
        if at_control not in (0, 1):
            raise ValueError("at_control must be 0 (Tilt compensation) or 1 (Drift compensation)")
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('AtomTrack.QuickCompStart', body_size=2)
        ## Arguments
        hex_rep += self.nanonisTCP.to_hex(at_control, 2)   # uint16
        self.nanonisTCP.send_command(hex_rep)
        message = self.nanonisTCP.receive_response(0)
        return message

    def DriftComp(self):
        """
        Applies the Drift measurement to the Drift compensation and turns
        On the compensation.

        Returns
        -------
        message : bytes
            Response body from Nanonis
        """
        ## Make Header
        hex_rep = self.nanonisTCP.make_header('AtomTrack.DriftComp', body_size=0)
        self.nanonisTCP.send_command(hex_rep)
        message = self.nanonisTCP.receive_response(0)
        return message

class Pattern:
    """
    Nanonis Pattern Module
    """
    def __init__(self,nanonisTCP):
        self.nanonisTCP = nanonisTCP
        
    def GridGet(self):        
        """
        Returns the grid size parameters.

        Returns
        -------
        num_points_x   : number of points in x that defines the grid
        num_points_y   : number of points in y that defines the grid
        x              : x (m) coordinate of the centre of the grid 
        y              : y (m) coordinate of the centre of the grid 
        w              : the width (m) of the grid 
        h              : the height (m) of the grid 
        angle          : the rotation angle (deg) of the grid


        """
        hex_rep = self.nanonisTCP.make_header('Pattern.GridGet', body_size=0)        
        self.nanonisTCP.send_command(hex_rep)        
        response = self.nanonisTCP.receive_response()
        
        idx = 0
        args = []
        for i in range(2):                                              # Two int32 responses requiring 4 bytes each
            args.append(self.nanonisTCP.hex_to_int32(response[idx:idx+4]))
            idx += 4
        
        for i in range(5):                                              # Five float32 responses requiring 4 bytes each
            args.append(self.nanonisTCP.hex_to_float32(response[idx:idx+4]))
            idx += 4

        num_points_x, num_points_y, x, y, w, h, angle = args
        
        return [num_points_x, num_points_y, x, y, w, h, angle]

class Piezo:
    """
    Nanonis Piezo Module
    """
    def __init__(self,nanonisTCP):
        self.nanonisTCP = nanonisTCP

    def DriftCompGet(self):
        """
        Returns the drift compensation parameters

        Returns
        ----------
        on    : True: Turn compensation on
                False: Turn compensation off
        vx    : linear speed applied to the X piezo (m/s)
        vy    : linear speed applied to the Y piezo (m/s)
        vz    : linear speed applied to the Z piezo (m/s)
        xsat  : indicates if the X drift correction reached 10% of piezo range
        ysat  : indicates if the Y drift correction reached 10% of piezo range
        zsat  : indicates if the Z drift correction reached 10% of piezo range

        """
        hex_rep = self.nanonisTCP.make_header('Piezo.DriftCompGet', body_size=0)
        
        self.nanonisTCP.send_command(hex_rep)
        
        response = self.nanonisTCP.receive_response(32)
        
        status = self.nanonisTCP.hex_to_uint32(response[0:4])
        vx     = self.nanonisTCP.hex_to_float32(response[4:8])
        vy     = self.nanonisTCP.hex_to_float32(response[8:12])
        vz     = self.nanonisTCP.hex_to_float32(response[12:16])
        xsat   = self.nanonisTCP.hex_to_uint32(response[16:20])
        ysat   = self.nanonisTCP.hex_to_uint32(response[20:24])
        zsat   = self.nanonisTCP.hex_to_uint32(response[24:28])
        satlim = self.nanonisTCP.hex_to_float32(response[28:32])
        
        return [status,vx,vy,vz,xsat,ysat,zsat,satlim]

class Util:
    """
    Nanonis Ultilities Module
    """
    def __init__(self,nanonisTCP):
        self.nanonisTCP = nanonisTCP

    def SessionPathGet(self):
        """
        Returns the current Nanonis session path.

        Returns
        -------
        session_path : str
            Path to current session folder
        """

        # no arguments
        hex_rep = self.nanonisTCP.make_header('Util.SessionPathGet', body_size=0)

        self.nanonisTCP.send_command(hex_rep)

        # response: [size (int32), string, error]
        response = self.nanonisTCP.receive_response(-1)

        # first 4 bytes = string length
        size = self.nanonisTCP.hex_to_int32(response[0:4])

        # next 'size' bytes = string
        path_bytes = response[4:4 + size]

        session_path = path_bytes.decode('utf-8')

        return session_path