text: import visa
  import os
  import numpy as np
  import time
  import matplotlib.pyplot as plt
  from scipy import stats
  
  from fft import FFT, Window
  from koheron import connect
  
  # Both the Koheron Alpha and the Keysight 33600A must be
  # disciplined on the Tektronix MCA3027 OCXO 10 MHz output.
  
  class Keysight33600A:
      def __init__(self):
          import visa
          rm = visa.ResourceManager('@py')
          self.inst = rm.open_resource('TCPIP::A-33600-00000.local::inst0::INSTR')
          self.inst.write('OUTP1:LOAD INF') # High impedance output
          self.inst.write('OUTP2:LOAD INF') # High impedance output
          self.channel = 1
  
      def set_channel(self, channel):
          self.channel = channel
  
      def set_sinus(self, freq, v_pp, offset):
          self.inst.write('SOUR' + str(int(self.channel)) + ':APPLy:SIN ' + str(freq / 1E3) + ' KHZ, '
          	            + str(v_pp) + ' VPP, '
          	            + str(offset) + ' V\n')
  
      def set_dc_offset(self, offset):
          self.inst.write('SOUR' + str(int(self.channel)) + ':APPLy:DC DEF, DEF, ' + str(offset) + ' V\n')
  
  def calibrate_transfer_function(gene, driver):
      n_pts = driver.n_pts
      fs = driver.get_fs()
      print(n_pts, fs)
  
      freqs = np.linspace(1e6, 62e6, num=200)
      freqs = np.round(freqs / fs * n_pts) * fs / n_pts
  
      peak_power = 0.0 * freqs
  
      for i, freq in enumerate(freqs):
          n = np.uint32(freq / fs * n_pts)
          gene.set_sinus(freq, 0.5, 0)
  
          time.sleep(0.5)
          psd = driver.read_psd_raw()
          peak_power[i] = psd[n-1]
          #peak_power[i] = np.max(psd)
          print(i, freq, peak_power[i])
  
      H = peak_power / peak_power[0] # Power transfer function
  
      # We fit 1 / H(f) with a degree 5 polynomial
      p = np.polyfit(freqs, 1 / H, 5)
  
      H_db = 10 * np.log10(H)
      H_fit_db = 10 * np.log10(1 / np.polyval(p, freqs))
      residuals_db = H_db - H_fit_db
  
      f, axarr = plt.subplots(2, sharex=True)
      axarr[1].set_xlabel('Frequency (MHz)')
      axarr[0].set_ylabel('Transfer function (dB)')
      axarr[1].set_ylabel('Residuals (dB)')
      axarr[0].semilogx(freqs * 1e-6, H_db)
      axarr[0].semilogx(freqs * 1e-6, H_fit_db)
      axarr[1].semilogx(freqs * 1e-6, residuals_db)
      plt.show()
  
      return p
  
  def calibrate_offset_gain(gene, driver, channel):
      n_cal_pts = 100
      voltages = np.linspace(-1, 1, n_cal_pts)
      off_raw = 0 * voltages
  
      for i, Voff in enumerate(voltages):
          gene.set_dc_offset(Voff)
          time.sleep(0.5)
          off_raw[i] = driver.get_adc_raw_data(100)[channel]
          print(i, off_raw[i])
  
      gain_lsb, offset_lsb, r_value, p_value, std_err = stats.linregress(voltages, off_raw)
      print(gain_lsb, offset_lsb)
  
      gain_volts, offset_volts, r_value, p_value, std_err = stats.linregress(off_raw, voltages)
      print(gain_volts, offset_volts)
  
      residuals = off_raw - (gain_lsb * voltages + offset_lsb)
  
      f, axarr = plt.subplots(2, sharex=True)
      axarr[0].plot(voltages, off_raw)
      axarr[0].plot(voltages, gain_lsb * voltages + offset_lsb)
      axarr[1].plot(voltages, residuals)
      axarr[0].set_ylabel("ADC raw value")
      axarr[1].set_ylabel("ADC INL")
      axarr[1].set_xlabel("Input voltage (V)")
      plt.show()
  
      return gain_lsb, offset_lsb
  
  if __name__=="__main__":
      gene = Keysight33600A()
      host = os.getenv('HOST', '192.168.1.18')
      client = connect(host, 'fft', restart=False)
      driver = FFT(client)
  
      driver.set_fft_window(Window.FLAT_TOP)
  
      for channel_under_test in range(2):
          gene.set_channel(channel_under_test + 1)
          driver.set_in_channel(channel_under_test)
          time.sleep(1)
  
          gain, offset = calibrate_offset_gain(gene, driver, channel_under_test)
          p = calibrate_transfer_function(gene, driver)
  
          cal_coeffs = np.float32(np.concatenate((np.array([gain, offset]), p)))
          print(cal_coeffs)