text: import time as t
  
  import numpy as np
  
  import matplotlib.pyplot as plt
  from matplotlib.widgets import Slider
  
  from functools import partial
  import copy
  
  from easy_gui import EasyPlotGUI
  import interface_ui
  
  from numba import jit
  
  
  class FdBpm(EasyPlotGUI):
      def __init__(self):
          super().__init__(interface_ui)
          self.slider_realtime = True
          self.window_title = "FD-BPM"
          self.icon_path = r"fdbpm-logo.png"
  
      def create_space(self):
          # window properties
          if not hasattr(self, 'NUM_SAMPLES') or not hasattr(self, 'LENGTH'):
              self.NUM_SAMPLES = int(100)
              self.LENGTH = int(1E3)   # set to 1E3 for better precision
          self.cmap = "jet"
          # Physical Properties
          if not hasattr(self, 'L'):
              self.L = 240E-6
          if not hasattr(self, 'dy'):
              self.dy = 1e-6  # set to 1E6
          if not hasattr(self, 'l_ambda'):
              self.l_ambda = 0.8e-6
  
          self.dx = self.L/self.NUM_SAMPLES
          self.x = np.linspace(-self.L/2, +self.L/2, self.NUM_SAMPLES)
  
          self.light = np.zeros(self.x.shape)
  
      def create_source(self, waist=8e-6, offset=0, plotOn=True):
  
          self.light = self.gauss(self.x-offset, waist)
  
          self.light_offset = offset
  
          if plotOn:
              plt.plot(self.x, self.light)
              plt.show()
  
      def gauss_light(self, fwhm=20E-6, offset=0):
          """
          Create a gaussian beam in amplitude.
  
          :math:`E = e^{-((x-x_0)/w)^{2P}}`
  
          The waist is defined as fwhm/sqrt(2*log(2)) and correspond to the 1/e
          field value and 1/:math:`e^2` intensity value.
  
          Parameters
          ----------
          fwhm : float
              Full width at half maximum (for intensity not amplitude) (µm).
          offset : float, optional
              Light offset from center in µm. 0 by default.
  
          Returns
          -------
          field : array
              Amplitude values over x in µm.
  
          Notes
          -----
          This methods uses the x and dist_x variables defined in :class:`Bpm`.
          """
          if not hasattr(self, 'fwhm'):
              self.fwhm = fwhm
          # such as I=1/e^2 in intensity
          spot_size = self.fwhm / np.sqrt(2 * np.log(2))
          if spot_size != 0:
              field = np.exp(-(self.x / spot_size)**2)
              field = np.roll(field, int(round(offset / self.dx)))
          else:
              field = 0 * self.x  # Avoid division by zero error
          self.light_offset = offset
          self.light = field
          return field
  
      def gauss(self, x, w0):
          return np.array(np.exp(-(x/w0)**2))
  
      def create_guides(self, width=25E-6, offset=0, plotOn=False):
          if not hasattr(self, 'n_env'):
              self.n_env = 1.004
          if not hasattr(self, 'guides'):
              self.guides = np.ones((self.NUM_SAMPLES,))*self.n_env
          if not hasattr(self, 'dn'):
              self.dn = 0.058
  
          mask = np.logical_and(self.x > -width/2+offset,
                                self.x < width/2+offset)
          self.guides[mask] += self.dn  # square guide
          # avg_guide = (np.max(guides)+np.min(guides))/2
          # avg_guide = (np.min(guides))
          self.avg_guide = (np.mean(self.guides))
  
          if plotOn:
              plt.plot(self.x, self.guides)
              plt.show()
  
          return (self.guides, self.avg_guide)
  
      def make_tri_matrix(self):
  
          if not hasattr(self, 'guides'):
              self.n_env = 1.004
              self.guides = np.ones((self.NUM_SAMPLES,))*self.n_env
              self.avg_guide = self.guides
  
          k0 = 2*np.pi/self.l_ambda
          # k = k0*self.guides
          k = k0/self.guides
          k_bar = k*self.avg_guide
          # k_bar = k/self.avg_guide
  
          self.h = self.dy
          self.ro = self.dy/(self.dx**2)
          self.A = 1j/(2*k_bar)
          self.B = 1j*(k**2-k_bar**2)/(2*k_bar)
          a = -self.ro*self.A
          b = 2*(1+self.ro*self.A)-self.h*self.B
  
          tridiag_matrix = np.zeros(
              (self.NUM_SAMPLES, self.NUM_SAMPLES), dtype='complex_')
  
          index = np.arange(self.NUM_SAMPLES)
  
          index0 = index[:-1]
          tridiag_matrix[index0, index0+1] = a[index0]
  
          index1 = index[1:]
          tridiag_matrix[index1, index1-1] = a[index1]
  
          tridiag_matrix[index, index] = b[index]
  
          # TODO: scipy class sparse matrices
  
          self.tridiag_matrix = tridiag_matrix
  
      @staticmethod
      @jit(forceobj=True)
      def core_propag_jit(A, B, ro, light, tridiag_matrix, h, LENGTH, NUM_SAMPLES):
  
          propag = np.zeros((int(LENGTH), int(NUM_SAMPLES)))
          d = np.zeros((NUM_SAMPLES), dtype='complex_')
  
          index = np.arange(1, NUM_SAMPLES-1)
  
          # cannot be transferred to numpy cause of np.linalg.solve in between lines
          for n in range(int(LENGTH)):
              d[index] = (2*(1-ro*A[index])+h*B[index])*light[index] + \
                  ro*A[index]*(light[index-1]+light[index+1])
              d[0] = (2*(1-ro*A[0])+h*B[0])*light[0]+ro*A[0]*(light[1])
              d[-1] = (2*(1-ro*A[-1])+h*B[-1])*light[-1]+ro*A[-1]*(light[-2])
  
              # light = np.linalg.lstsq(tridiag_matrix,d)[0]
              light = np.linalg.solve(tridiag_matrix, d)  # fastest option
              # propag[n,:] = np.abs(light)
              # propag[n, :] = (light*light.conjugate()).real
              propag[n, :] = light.real**2 + light.imag**2
              # propag[n,:] = np.real(light) # Oscillations
              # propag[n,:] = np.angle(light) # interesting
          return propag
  
      def calculate_propagation(self, plotOn=True, timing=True):
  
          if (not hasattr(self, 'tridiag_matrix')):
              self.make_tri_matrix()
  
          if(timing == True):
              start = t.time()
  
          propag = self.core_propag_jit(
              self.A, self.B, self.ro, self.light, self.tridiag_matrix, self.h, self.LENGTH, self.NUM_SAMPLES)
  
          if(timing == True):
              end = t.time()
              print("Time elapsed:%0.4f seconds" % (end-start))
  
          self.propag = propag
  
          if(plotOn == True):
              propag_img = propag[::int(
                  np.floor(self.LENGTH/self.NUM_SAMPLES)), :]
  
              fig, ax = plt.subplots()
              ax.imshow(propag_img, cmap=self.cmap, interpolation='bilinear', extent=[
                        -self.L/2*1E6, +self.L/2*1E6, self.LENGTH*self.dy*1E6, 0], aspect='auto')
              ax.set_xlabel(r"x ($\mu$m)")
              ax.set_ylabel(r"Length ($\mu$m)")
  
              plt.show()
  
          return self.propag
  
      def plot_moving_source(self):
          plt.close('all')
  
          fig, ax = plt.subplots()
          plt.subplots_adjust(left=0.25, bottom=0.25)
          if int(np.floor(self.LENGTH/self.NUM_SAMPLES)) > 0:
              propag_img = self.calculate_propagation(
                  plotOn=False)[::int(np.floor(self.LENGTH/self.NUM_SAMPLES)), :]
          else:
              propag_img = self.calculate_propagation(plotOn=False)
          img = ax.imshow(propag_img, cmap=self.cmap, interpolation='bilinear', extent=[
                          -self.L/2*1E6, +self.L/2*1E6, self.LENGTH*self.dy*1E3, 0], aspect='auto')
          ax.set_xlabel(r"x ($\mu$m)")
          ax.set_ylabel(r"Length (mm)")
          ax.margins(x=0)
  
          axcolor = 'lightgoldenrodyellow'
          ax_position = plt.axes([0.25, 0.1, 0.65, 0.03], facecolor=axcolor)
          plt.plot(self.x, self.guides, 'orange')
  
          slider_position = Slider(ax_position, 'Light Offset', -self.L/2, +self.L/2,
                                   valinit=self.light_offset, valstep=self.dx*1E-2, valfmt="%2.0e")
  
          def update(val):
              pos = np.float(slider_position.val)
              # self.create_source(offset=pos,plotOn=False)
              self.gauss_light(offset=pos)
              img.set_data(self.calculate_propagation(plotOn=False))
              fig.canvas.draw_idle()
  
          slider_position.on_changed(update)
  
          plt.show()
  
      def reset_variables(self):
          # pass
          delattr(self, 'NUM_SAMPLES')
          delattr(self, 'LENGTH')
          delattr(self, 'dy')
          delattr(self, 'dx')
          delattr(self, 'l_ambda')
          delattr(self, 'L')
          delattr(self, 'fwhm')
          delattr(self, 'dn')
          delattr(self, 'light')
          delattr(self, 'light_offset')
          delattr(self, 'n_env')
          delattr(self, 'x')
          delattr(self, 'guides')
          delattr(self, 'avg_guide')
          delattr(self, 'h')
          delattr(self, 'ro')
          delattr(self, 'A')
          delattr(self, 'B')
          delattr(self, 'tridiag_matrix')
  
      def update_interactivity(self):
          self.ui.doubleSpinBox_lambda.setValue(self.l_ambda*1E6)
          self.ui.spinBox_propLength.setValue(self.LENGTH*self.dy*1E6)
          self.ui.spinBox_windowSize.setValue(self.L*1E6)
          self.ui.spinBox_numSamples.setValue(self.NUM_SAMPLES)
  
          self.ui.SliderPos.setMinimum(-self.L/2*1E6)
          self.ui.SliderPos.setMaximum(+self.L/2*1E6)
          self.ui.SliderPos.setValue(self.light_offset*1E6)
          self.ui.label_slider.setText(
              str(np.format_float_positional(self.light_offset*1E6, 3))+" um")
  
          self.ui.actionFree_Space_Propagation.triggered.connect(
              partial(self.examples, 'free_space', chart_update=True))
          self.ui.actionOne_Waveguide.triggered.connect(
              partial(self.examples, 'one_guide', chart_update=True))
          self.ui.actionThree_Waveguides.triggered.connect(
              partial(self.examples, '3_guides', chart_update=True))
  
          # self.ui.SliderPos.setTickInterval(self.dx)
  
          def update_label():
              pos = self.ui.SliderPos.value()
              self.ui.label_slider.setText(str(pos)+" um")
  
          def update_input_values():
              dy = copy.deepcopy(self.dy)
              fwhm = copy.deepcopy(self.fwhm)
              self.reset_variables()
              self.dy = dy
              self.fwhm = fwhm
              self.l_ambda = 1E-6*np.float(self.ui.doubleSpinBox_lambda.value())
              self.LENGTH = 1/self.dy * \
                  np.float(self.ui.spinBox_propLength.value()*1E-6)
              self.L = 1E-6*np.float(self.ui.spinBox_windowSize.value())
              self.NUM_SAMPLES = int(self.ui.spinBox_numSamples.value())
              self.create_space()
              self.gauss_light()
              self.create_guides()
              self.ui.SliderPos.setValue(self.light_offset*1E6)
              self.ui.SliderPos.setMinimum(-self.L/2*1E6)
              self.ui.SliderPos.setMaximum(+self.L/2*1E6)
              self.update_graph()
              # FIXME: update doesn't refresh guides
  
          def changed_value():
              pos = self.ui.SliderPos.value()
              self.gauss_light(offset=pos*1E-6)
              self.update_graph(slider=True)
              self.ui.label_slider.setText(str(pos)+" um")
  
          self.ui.pushButton.clicked.connect(update_input_values)
  
          if self.slider_realtime == True:
              self.ui.SliderPos.valueChanged.connect(changed_value)
  
          else:
              self.ui.SliderPos.sliderReleased.connect(changed_value)
              self.ui.SliderPos.valueChanged.connect(update_label)
  
      def examples(self, name, chart_update=False):
          if chart_update:
              self.reset_variables()
          if name == 'free_space':
  
              self.NUM_SAMPLES = 101
              self.LENGTH = 1E2
              self.dy = 1E-4
              self.l_ambda = 1.5E-6
              self.L = 1000E-6
  
              self.create_space()
              self.gauss_light(fwhm=20E-6)
              self.create_guides(width=0)
  
          elif name == 'one_guide':
              self.NUM_SAMPLES = 101
              self.LENGTH = 1E2
              self.dy = 1E-4
              self.l_ambda = 1.5E-6
              self.L = 500E-6
  
              self.create_space()
              self.gauss_light()
              self.dn = 0.058
              self.n_env = 1.004
              self.create_guides()
          elif name == '3_guides':
              self.NUM_SAMPLES = 101
              self.LENGTH = 500
              self.dy = 1E-5
              self.l_ambda = 1.55E-6
              self.L = 60E-6
              self.create_space()
              plotOn = True
              offset = 14E-6
              self.gauss_light(fwhm=12E-6, offset=-offset+1E-7)
              self.create_guides(width=10E-6, offset=-offset)
              self.create_guides(width=10E-6)
              self.create_guides(width=10E-6, offset=offset, plotOn=False)
  
          if chart_update:
              self.ui.doubleSpinBox_lambda.setValue(self.l_ambda*1E6)
              self.ui.spinBox_propLength.setValue(self.LENGTH*self.dy*1E6)
              self.ui.spinBox_windowSize.setValue(self.L*1E6)
              self.ui.spinBox_numSamples.setValue(self.NUM_SAMPLES)
              self.ui.SliderPos.setMinimum(-self.L/2*1E6)
              self.ui.SliderPos.setMaximum(+self.L/2*1E6)
              self.ui.SliderPos.setValue(self.light_offset*1E6)
              self.update_graph(slider=False)
  
      def update_graph(self, slider=False):
  
          
          ax = self.ax  # no canvas in the object
  
          if int(np.floor(self.LENGTH/self.NUM_SAMPLES)) > 0:
              propag_img = self.calculate_propagation(
                  plotOn=False)[::int(np.floor(self.LENGTH/self.NUM_SAMPLES)), :]
          else:
              propag_img = self.calculate_propagation(plotOn=False)
  
          if slider == False:
              ax.clear()
              self.img = ax.imshow(propag_img, cmap=self.cmap, interpolation='bilinear', extent=[
                                   -self.L/2*1E6, +self.L/2*1E6, self.LENGTH*self.dy*1E3, 0], aspect='auto')
              ax.set_xlabel(r"x ($\mu$m)")
              ax.set_ylabel(r"Length (mm)")
          else:
              # ax.clear()
              self.img.set_data(propag_img)
          plt.tight_layout(pad=0.)
          self.draw()
  
  
  if __name__ == "__main__":
      # %matplotlib qt
      fd = FdBpm()
  
      fd.examples('one_guide')
      fd.show_gui()
  