text: # -*- encoding: utf-8 -*-
  '''
  @Description: Calculate the weighting potential and electric field wiht fenics      
  @Date       : 2021/08/31 15:04:25
  @Author     : tanyuhang
  @version    : 1.0
  '''
  
  import fenics
  import mshr
  import sys
  import numpy as np
  import ROOT
  from array import array
  
  
  #Calculate the weighting potential and electric field
  class FenicsCal:
  
      def __init__(self,my_d,fen_dic):
          self.p_electric = []
          self.w_p_electric = []
          self.det_model = fen_dic['name']
          self.fl_x=my_d.l_x/fen_dic['xyscale']  
          self.fl_y=my_d.l_y/fen_dic['xyscale']  
          self.fl_z=my_d.l_z
          self.tol = 1e-14
          m_sensor_box=self.fenics_space(my_d)  
          self.mesh3D = mshr.generate_mesh(m_sensor_box,fen_dic['mesh'])  
          self.V = fenics.FunctionSpace(self.mesh3D, 'P', 1)
          self.fenics_p_electric(my_d)
          self.fenics_p_w_electric(my_d)
  
      def fenics_space(self,my_d):
          """
          @description: 
              Define the fenics solver space 
          @param:
              None
          @Returns:
              Fenics Box structure
          @Modify:
              2021/08/31
          """
          if "plugin3D" in self.det_model:
              self.sensor_range_confirm(my_d)
              m_sensor =  mshr.Box(fenics.Point(self.sx_l,self.sy_l, 0), 
                                   fenics.Point(self.sx_r,self.sy_r, self.fl_z))
              for i in range(len(my_d.e_tr)):
                  e_t_i = my_d.e_tr[i]
                  elec_n=mshr.Cylinder(fenics.Point(e_t_i[0],e_t_i[1],e_t_i[3]), 
                                       fenics.Point(e_t_i[0],e_t_i[1],e_t_i[4]),
                                       e_t_i[2],e_t_i[2])
                  m_sensor =m_sensor - elec_n 
          elif "planar3D" in self.det_model:
              m_sensor =  mshr.Box(fenics.Point(0, 0, 0), 
                                   fenics.Point(self.fl_x, self.fl_y, self.fl_z))
          else:
              print("sensor model is wrrong")
              sys.exit()
          return m_sensor            
          
      def sensor_range_confirm(self,my_d):
          """
          @description:
              confirm the sensor range 
              at x,y axias to avoide the the big sensor size
              which will lead to the mesh complicate         
          @param:
              xv_min - fenics sensor x left value
              xv_max - fenics sensor x right value
              yv_min - fenics sensor y left value
              yv_max - fenics sensor y right value          
          @Modify:
              2021/08/31
          """    
          xv_list=[]
          yv_list=[]
          rest_length=50 #um
          length=0
          for i in range(len(my_d.e_tr)):
              e_t_i = my_d.e_tr[i]
              xv_list.append(e_t_i[0])
              yv_list.append(e_t_i[1])
              ele_radius= e_t_i[2]
          while length == 0:
              xv_max = max(xv_list)+ele_radius+rest_length 
              xv_min = min(xv_list)-ele_radius-rest_length 
              yv_max = max(yv_list)+ele_radius+rest_length
              yv_min = min(yv_list)-ele_radius-rest_length 
              if xv_max >= yv_max:
                  yv_max = xv_max
              else:
                  xv_max = yv_max
              if xv_min <= yv_min:
                  yv_min = xv_min
              else:
                  xv_min = yv_min
              if (xv_max > my_d.l_x or xv_min <0 
                 or yv_max > my_d.l_y or yv_min < 0):
                  rest_length -= 1
              else:
                  length=1
          self.sx_l=xv_min 
          self.sx_r=xv_max 
          self.sy_l=yv_min 
          self.sy_r=yv_max 
  
      def fenics_p_electric(self,my_d):    
          """
          @description:
              Solve poisson equation to get potential and electric field
          @Modify:
              2021/08/31
          """
          if  "plugin3D" in self.det_model:
              bc_l=[]
              bc_l = self.boundary_definition_3D(my_d,"Possion")          
          elif "planar3D" in self.det_model:
              bc_l = self.boundary_definition_2D(my_d,"Possion")
  
          u = fenics.TrialFunction(self.V)
          v = fenics.TestFunction(self.V)
          f = fenics.Constant(self.f_value(my_d))
          a = fenics.dot(fenics.grad(u), fenics.grad(v))*fenics.dx
          L = f*v*fenics.dx
          # Compute solution
          self.u = fenics.Function(self.V)
          fenics.solve(a == L, self.u, bc_l,
                       solver_parameters=dict(linear_solver='gmres',
                                              preconditioner='ilu'))
          #calculate electric field
          W = fenics.VectorFunctionSpace(self.mesh3D, 'P', 1)
          self.E_field = fenics.project(fenics.as_vector((self.u.dx(0),
                                                          self.u.dx(1),
                                                          self.u.dx(2))),W)
  
      def fenics_p_w_electric(self,my_d):  
          """
          @description:
              Solve Laplace equation to 
              get weighting potential and weighting electric field
          @Modify:
              2021/08/31
          """
          if  "plugin3D" in self.det_model:
              bc_l = []
              bc_l = self.boundary_definition_3D(my_d,"Laplace")
          elif "planar3D" in self.det_model:
              bc_l = self.boundary_definition_2D(my_d,"Laplace")
          # Define variational problem
          u_w = fenics.TrialFunction(self.V)
          v_w = fenics.TestFunction(self.V)
          f_w = fenics.Constant(0)
          a_w = fenics.dot(fenics.grad(u_w), fenics.grad(v_w))*fenics.dx
          L_w = f_w*v_w*fenics.dx
          # Compute solution
          self.u_w = fenics.Function(self.V)
          fenics.solve(a_w == L_w, self.u_w, bc_l)
  
      def boundary_definition_3D(self,my_d,model):
          """
          @description:
              Get boundary definition of 3D detector with Possion and Laplace
          @Modify:
              2021/08/31
          """
          bc_l = []
          p_ele,n_ele=self.model_para(my_d,model)
          for i in range (len(my_d.e_tr)):
              e_i = my_d.e_tr[i]
              str_e = "x[0]>={e_0}-{e_2} && x[0]<={e_0}+"\
                      +"{e_2} && x[1]>={e_1}-{e_2} && "\
                      +"x[1]<={e_1}+{e_2} && x[2]>={e_3} \
                      && x[2]<={e_4} && on_boundary"
              elec_p = str_e.format(e_0=e_i[0], e_1=e_i[1],
                                    e_2=e_i[2], e_3=e_i[3],
                                    e_4=e_i[4])
              if e_i[5] == "p":
                  bc = fenics.DirichletBC(self.V, p_ele, elec_p)
              else:
                  bc = fenics.DirichletBC(self.V, n_ele, elec_p)
              bc_l.append(bc)
          return bc_l
  
      def boundary_definition_2D(self,my_d,model):
          """
          @description:
              Get boundary definition of 2D detector with Possion and Laplace
          @Modify:
              2021/08/31
          """
          p_ele,n_ele=self.model_para(my_d,model)
          u_D = fenics.Expression('x[2]<tol ? p_1:p_2',
                                  degree=2, tol=1E-14,
                                  p_1=p_ele, p_2=n_ele)
          def boundary(x, on_boundary):
              return abs(x[2])<self.tol or abs(x[2]-self.fl_z)<self.tol
          bc_l = fenics.DirichletBC(self.V, u_D, boundary)
          return bc_l
      
      def model_para(self,my_d,model):
          """
          @description:
              Choose Possion and Laplace parameters
          @Modify:
              2021/08/31
          """       
          bc_l=[]
          if model == "Possion":
              p_ele = my_d.v_voltage
              n_ele = 0.0
          elif model == "Laplace":
              p_ele = 0.0
              n_ele = 1.0
          else:
              print("The input fenics solver model is wrong")
          return p_ele,n_ele
  
      def f_value(self,my_d):
          """
          @description: 
              Cal f_value of Poisson equation
          @param:
              perm_mat -- Dielectric constant of using material
                       -- 11.7 Silicon
                       -- 9.76 Silicon carbide
          @Modify:
              2021/08/31
          """
          if my_d.mater == 0:
              perm_mat = 11.7  
          elif my_d.mater == 1:
              perm_mat = 9.76  
          else:
              print("material is wrong")            
          e0 = 1.60217733e-19
          perm0 = 8.854187817e-12   #F/m
          f_value = e0*my_d.d_neff*1e6/perm0/perm_mat
          return f_value
          
      def get_e_field(self,px,py,pz):
          """
          @description: 
              Get eletric field at the px,py,pz position
          @param:
              out_range -- out_range = False
                        -- Position (x,y,z) don't exit in sensor fenics range
          @reture:
              Eletric field along x,y,z direction
          @Modify:
              2021/08/31
          """
          out_range=self.judge_fenics_range(px,py,pz)
          if out_range:   
              x_value,y_value,z_value = 0,0,0
          else:
              scale_px=px%self.fl_x
              scale_py=py%self.fl_y          
              try:
                  x_value,y_value,z_value = self.E_field(scale_px,scale_py,pz)
              except RuntimeError:
                  x_value,y_value,z_value = 0,0,0
          return x_value,y_value,z_value
  
      def get_w_p(self,px,py,pz):
          """
          @description: 
              Get weighting potential at the (x,y,z) position
          @param:
              out_range -- out_range = False
                        -- Position (x,y,z) don't exit in sensor fenics range
          @reture:
              Get weighting potential at (x,y,z) position
          @Modify:
              2021/08/31
          """
          out_range=self.judge_fenics_range(px,py,pz)
          if out_range:   
              f_w_p = 1.0
          else:
              scale_px=px%self.fl_x
              scale_py=py%self.fl_y   
              try:
                  f_w_p = self.u_w(scale_px,scale_py,pz)
              except RuntimeError:
                  f_w_p = 0.0
          return f_w_p
  
      def get_potential(self,px,py,pz):
          """
          @description: 
              Get potential at the (x,y,z) position
          @param:
              out_range -- out_range = False
                        -- Position (x,y,z) don't exit in sensor fenics range
          @reture:
              Get potential at (x,y,z) position
          @Modify:
              2021/08/31
          """
          out_range=self.judge_fenics_range(px,py,pz)
          if out_range:   #
              f_w_p = 0
          else:
              scale_px=px%self.fl_x
              scale_py=py%self.fl_y   
              try:
                  f_w_p = self.u(scale_px,scale_py,pz)
              except RuntimeError:
                  f_w_p = 0.0
          return f_w_p
  
      def judge_fenics_range(self,px,py,pz):
          """
          @description: 
             Judge whether (x,y,z) position is in sensor fenics range
          @reture:
              False: not
              Ture:  in
          @Modify:
              2021/08/31
          """
          if "plugin3D" in self.det_model:
              if (px < self.sx_l or px > self.sx_r 
                     or py < self.sy_l or py > self.sy_r):
                  out_range=True
              else:
                  out_range=False
          elif "planar3D" in self.det_model:
              out_range=False
          return out_range
      
      def __del__(self):
          pass
  
  
  class FenicsCal2D:
  
      def __init__(self,det):
          
          self.det = det
  
          # poential & field
          self.potential_value_2d = []
      
          self.electric_field_x_value = [ [] for n in range(self.det.ny+1) ]
          self.electric_field_y_value = [ [] for n in range(self.det.ny) ]
  
          self.electric_field_x_position = [ [] for n in range(self.det.ny+1) ]
          self.electric_field_y_position = [ [] for n in range(self.det.ny) ]
  
          # weighting poential & weighting field
          self.weighting_potential_value_2d = []
      
          self.weighting_electric_field_x_value = [ [] for n in range(self.det.ny+1) ]
          self.weighting_electric_field_y_value = [ [] for n in range(self.det.ny) ]
  
          self.weighting_electric_field_x_position = [ [] for n in range(self.det.ny+1) ]
          self.weighting_electric_field_y_position = [ [] for n in range(self.det.ny) ]
  
          self.solve()
          #self.draw()
  
      def cal_possion(self):
          
          width = self.det.det_width
          thin = self.det.det_thin
          
          nx = self.det.nx
          ny = self.det.ny
  
          # Create mesh and function space
          mesh = fenics.RectangleMesh(fenics.Point(0, 0), fenics.Point(width, thin), nx, ny)
          V = fenics.FunctionSpace(mesh, "P", 1)
  
          # Define boundary condition
          u_D = fenics.Expression('x[1] < tol? det_voltage : 0', degree = 2,tol = 1E-14,det_voltage = self.det.bias_voltage)
  
          def boundary(x, on_boundary):
              return abs(x[1])<1E-14 or abs(x[1]-thin)<1E-14
  
          bc = fenics.DirichletBC(V, u_D, boundary)
  
          # Define variational problem
          u = fenics.TrialFunction(V)
          v = fenics.TestFunction(V)
  
          f = fenics.Expression(self.det.doping_epr,degree=2)
          a = fenics.dot(fenics.grad(u), fenics.grad(v))*fenics.dx
          L = f*v*fenics.dx #+ g*v*ds
  
          # Compute solution
          u = fenics.Function(V)
          fenics.solve(a == L, u, bc)
  
          potential_value_1d = u.compute_vertex_values()
          potential_value_2d = np.array(potential_value_1d).reshape(ny+1,nx+1)
  
          self.potential_value_2d = potential_value_2d
  
          # print(self.potential_value_2d)
  
  
      def cal_weighting_possion(self):
  
          width = self.det.det_width
          thin = self.det.det_thin
          
          nx = self.det.nx
          ny = self.det.ny
  
          # Create mesh and function space
          mesh = fenics.RectangleMesh(fenics.Point(0, 0), fenics.Point(width, thin), nx, ny)
          V = fenics.FunctionSpace(mesh, "P", 1)
  
          # Define boundary condition
          u_D = fenics.Expression('x[1] < tol? det_voltage : 0', degree = 2,tol = 1E-14,det_voltage = self.det.bias_voltage/abs(self.det.bias_voltage))
  
          def boundary(x, on_boundary):
              return abs(x[1])<1E-14 or abs(x[1]-thin)<1E-14
  
          bc = fenics.DirichletBC(V, u_D, boundary)
  
          # Define variational problem
          u = fenics.TrialFunction(V)
          v = fenics.TestFunction(V)
  
          f = fenics.Constant(0)
          a = fenics.dot(fenics.grad(u), fenics.grad(v))*fenics.dx
          L = f*v*fenics.dx #+ g*v*ds
  
          # Compute solution
          u = fenics.Function(V)
          fenics.solve(a == L, u, bc)
  
          weighting_potential_value_1d = u.compute_vertex_values()
          weighting_potential_value_2d = np.array(weighting_potential_value_1d).reshape(ny+1,nx+1)
  
          self.weighting_potential_value_2d = weighting_potential_value_2d
  
          return weighting_potential_value_2d
  
      def cal_electric_field(self):
  
          width = self.det.det_width
          thin = self.det.det_thin
          
          nx = self.det.nx
          ny = self.det.ny
  
          x_step = width/nx
          y_step = thin/ny
  
          # x direction
          for j in range(ny+1):
              for i in range(nx):
  
                  # electric field
                  tmp_xpos = 0.5*x_step*(2*i+1)
                  tmp_xef = (self.potential_value_2d[j][i] - self.potential_value_2d[j][i+1])/x_step
                  self.electric_field_x_position[j].append(tmp_xpos)
                  self.electric_field_x_value[j].append(tmp_xef)
  
                  # weighting field
                  tmp_wxpos = 0.5*x_step*(2*i+1)
                  tmp_xwef = (self.weighting_potential_value_2d[j][i] - self.weighting_potential_value_2d[j][i+1])/x_step
                  self.weighting_electric_field_x_position[j].append(tmp_wxpos)
                  self.weighting_electric_field_x_value[j].append(tmp_xwef)
  
  
          # y direction 
          for j in range(ny):
              for i in range(nx+1):
  
                  # electric field
                  tmp_ypos = 0.5*y_step*(2*j+1)
                  tmp_yef = (self.potential_value_2d[j][i]- self.potential_value_2d[j+1][i])/y_step
                  self.electric_field_y_position[j].append(tmp_ypos)
                  self.electric_field_y_value[j].append(tmp_yef)
  
                  # weighting field
                  tmp_wypos = 0.5*y_step*(2*j+1)
                  tmp_ywef = (self.weighting_potential_value_2d[j][i] - self.weighting_potential_value_2d[j+1][i])/y_step
                  self.weighting_electric_field_y_position[j].append(tmp_wypos)
                  self.weighting_electric_field_y_value[j].append(tmp_ywef)     
  
  
  
      def cal_point_field(self,px_point,py_point,input_value):
  
          width = self.det.det_width
          thin = self.det.det_thin
          
          nx = self.det.nx
          ny = self.det.ny
  
          x_step = width/nx
          y_step = thin/ny
  
          #Interpolation method 
          rex_value=px_point%x_step
          nx_value=int(px_point/x_step)
          rey_value=py_point%y_step
          ny_value=int(py_point/y_step)
  
          if(rex_value>x_step/2):
              e_v_x1=rex_value-x_step/2
              nx1_v=nx_value
              nx2_v=nx_value+1
          else:
              e_v_x1=rex_value+x_step/2
              e_v_x2=x_step-e_v_x1
              nx1_v=nx_value-1
              nx2_v=nx_value
  
          if(rey_value>y_step/2):
              e_v_y1=rey_value-y_step/2
              ny1_v=ny_value
              ny2_v=ny_value+1
          else:
              e_v_y1=rey_value+y_step/2
              e_v_y2=y_step-e_v_y1
              ny1_v=ny_value-1
              ny2_v=ny_value
  
          if (nx_value<=0):
              r_u=0
              nx1_v=nx2_v
          elif (nx_value>=nx-1):
              r_u=0
              nx2_v=nx1_v
          else:
              r_u=e_v_x1/x_step
  
          if (ny_value<=0):
              r_t=0
              ny1_v=ny2_v
          elif (ny_value>=ny-1):
              r_t=0
              ny2_v=ny1_v
          else:
              r_t=e_v_y1/y_step
  
          value_11=input_value[ny1_v][nx1_v]
          value_21=input_value[ny1_v][nx2_v]
          value_12=input_value[ny1_v][nx1_v]
          value_22=input_value[ny1_v][nx2_v]
          out_field=0.0
          out_field=(1-r_u)*(1-r_t)*value_11
          out_field+=r_u*(1-r_t)*value_21
          out_field+=r_u*r_t*value_22
          out_field+=(1-r_u)*r_t*value_12
  
          return out_field        
  
  
      def solve(self):
  
          self.cal_possion()
          self.cal_weighting_possion()
          self.cal_electric_field()
  
      def draw(self):
  
          cutline = int(self.det.nx/2.0)
  
          # plot electric field at x = middle
          ep = array( 'd' )
          ev = array( 'd' )
  
          for i in range(self.det.ny):
  
              ep.append(self.electric_field_y_position[i][cutline])
  
              efx = self.cal_point_field(self.det.det_width/2.0, self.electric_field_y_position[i][cutline], self.electric_field_x_value)
              efy = self.cal_point_field(self.det.det_width/2.0, self.electric_field_y_position[i][cutline], self.electric_field_y_value)
              ef = np.array([efx,efy])
  
              ef_value = np.linalg.norm(ef)#*1e4 #V/cm
  
              ev.append(ef_value)
  
          # print(ep)
          # print(ev)
  
          g_e = ROOT.TGraph(self.det.ny,ep,ev)
  
          g_e.SetLineColor(600)
          g_e.SetLineWidth(4)
          g_e.SetTitle( 'Electric Field at Cut Line' )
          g_e.GetXaxis().SetTitle( 'Depth [um]' )
          g_e.GetXaxis().SetRangeUser(0,self.det.det_thin)
          g_e.GetYaxis().SetTitle( 'E [V/um]' )
  
          g_e.GetXaxis().CenterTitle()
          g_e.GetXaxis().SetTitleOffset(1.8)
          g_e.GetXaxis().SetTitleSize(0.05)
          g_e.GetXaxis().SetLabelSize(0.05)
          #g_e.GetXaxis().SetNdivisions(505)
  
          g_e.GetYaxis().CenterTitle()
          g_e.GetYaxis().SetTitleOffset(1.8)
          g_e.GetYaxis().SetTitleSize(0.05)
          g_e.GetYaxis().SetLabelSize(0.05)
          #g_e.GetYaxis().SetNdivisions(505)
         
          c = ROOT.TCanvas( 'c', 'c',500, 500 )
          c.SetGrid()
          c.SetLeftMargin(0.18)
          c.SetBottomMargin(0.18)
          # c.Divide(1,2)
  
          c.cd()
          g_e.Draw()
          c.Modified()
          c.Update()
          c.SaveAs("./fig/%s_%s_electricfield_%dV.pdf"%(self.det.mat_name,self.det.det_model,self.det.bias_voltage))
  