Add pseudo to qp_convert

scripts/ezfio_interface/qp_convert_output_to_ezfio.py

@@ -273,8 +273,126 @@ def write_ezfio(res, filename):
     # \_|  |___/\___|\__,_|\__,_|\___/
     #
 
-    ezfio.set_pseudo_do_pseudo(False)
 
+    # INPUT
+    # {% for lanel,zcore, l_block in l_atom  $}
+    #       #local l_block l=0}
+    #       {label} GEN {zcore} {len(l_block)-1 #lmax_block}
+    #       {% for l_param in l_block%}
+    #                {len(l_param) # list of parameter aka n_max_bock_max(n)}
+    #                {% for coef,n,zeta for l_param}
+    #                    {coef,n, zeta}
+
+    # OUTPUT
+
+    # Local are 1 array padded by max(n_max_block) when l == 0  (output:k_loc_max)
+    # v_k[n-2][atom] = value
+
+    #No Local are 2 array padded with max of lmax_block when l!=0 (output:lmax+1)  and max(n_max_block)whem l !=0 (kmax)
+    # v_kl[l][n-2][atom] = value
+
+    def pad(array, size, value=0):
+      new_array= array
+      for add in xrange(len(array), size):
+         new_array.append(value)
+
+      return new_array
+
+    def parse_str(pseudo_str):
+       '''Return 4d array  atom,l,n, attribute (attribute is coef, n, zeta)'''
+       matrix = []
+       array_l_max_block = []
+       array_z_remove = [] 
+
+       for block in [b for b in pseudo_str.split('\n\n') if b]:
+         #First element is header, the rest are l_param 
+         array_party = [i for i in re.split(r"\n\d+\n",block) if i]
+
+         z_remove, l_max_block = map(int,array_party[0].split()[-2:])
+         array_l_max_block.append(l_max_block)
+         array_z_remove.append(z_remove)
+
+         matrix_3d.append([ [coef_n_zeta.split()[1:] for coef_n_zeta in l.split('\n')] for l in array_party[1:] ])
+
+       return (matrix, array_l_max_block, array_z_remove)
+
+    def get_local_stuff(matrix):
+       
+         matrix_local_unpad = [atom[0] for atom in matrix]
+         k_loc_max = max(len(i) for i in matrix_local_unpad)
+
+         matrix_local = [pad(ll,k_loc_max,[0.,0,0.]) for ll in matrix_local_unpad]
+
+         m_coef =[ [float(i[0]) for i in atom] for atom in matrix_local]
+         m_n =   [ [int(i[1])-2 for i in atom] for atom in matrix_local]
+         m_zeta = [ [float(i[2]) for i in atom] for atom in matrix_local]
+
+         return (k_loc_max, m_coef,m_n,m_zeta)
+
+    def get_non_local_stuff(matrix):
+
+        matrix_unlocal_unpad = [atom[1:] for atom in matrix]
+        l_max_block = max(len(i) for i in matrix_unlocal_unpad)
+        k_max = max( [len(item) for row  in matrix_unlocal_unpad  for item in row ] )
+
+        matrix_unlocal_semipaded =[ [pad(item,k_max,[0.,0,0.]) for item in row] for row in matrix_unlocal_unpad]
+
+        empty_row = [ [0.,0,0.] for k in range(l_max_block)]
+        matrix_unlocal = [pad(ll,l_max_block,empty_row) for ll in matrix_unlocal_semipaded]
+
+        m_coef_noloc = [ [  [ float(k[0])  for k in j ] for j in i   ]  for i in matrix_local ]
+        m_n_noloc =    [ [  [ int(k[1])-2  for k in j ] for j in i   ]  for i in matrix_local ]
+        m_zeta_noloc = [ [  [ float(k[2])  for k in j ] for j in i   ]  for i in matrix_local ]
+        
+        return (l_max_block,k_max, m_coef_noloc,m_n_noloc,m_zeta_noloc)
+
+    try:
+        pseudo_str = res_file.get_pseudo()    
+    except:
+	ezfio.set_pseudo_do_pseudo(False)
+    else:
+	ezfio.set_pseudo_do_pseudo(True)
+        matrix, array_l_max_block, array_z_remove = parse(pseudo_str)
+
+        # ~#~#~#~#~#~#~#~#~#~#~#~#~#~#~#~#~#~#~#~#~ #
+        # Z _ e f f , a l p h a / b e t a _ e l e c #
+        # ~#~#~#~#~#~#~#~#~#~#~#~#~#~#~#~#~#~#~#~#~ #
+
+        ezfio.pseudo_charge_remove = array_z_remove
+        ezfio.nuclei_nucl_charge = ezfio.nuclei_nucl_charge - array_z_remove
+
+        import math
+        num_elec = sum(ezfio.nuclei_nucl_charge)
+
+        ezfio.electrons_elec_alpha_num = int(math.ceil(num_elec / 2.))
+        ezfio.electrons_elec_beta_num = int(math.floor(num_beta / 2.))
+
+        # Change all the array 'cause EZFIO
+        #   v_kl (v, l) => v_kl(l,v)
+        #    v_kl => zip(*_v_kl)
+        # [[7.0, 79.74474797, -49.45159098], [1.0, 5.41040609, -4.60151975]]
+        # [(7.0, 1.0), (79.74474797, 5.41040609), (-49.45159098, -4.60151975)]
+
+        # ~#~#~#~#~ #
+        # L o c a l #
+        # ~#~#~#~#~ #
+
+        klocmax, m_coef,m_n,m_zeta = get_local_stuff(matrix)
+        ezfio.pseudo_pseudo_klocmax = klocmax
+
+        ezfio.pseudo_pseudo_v_k = zip(*m_coef)
+        ezfio.pseudo_pseudo_n_k = zip(*m_n)
+        ezfio.pseudo_pseudo_dz_k = zip(*m_zeta)
+ 
+        # ~#~#~#~#~#~#~#~#~ #
+        # N o n _ L o c a l #
+        # ~#~#~#~#~#~#~#~#~ #
+
+        l_max_block,k_max, m_coef_noloc,m_n_noloc,m_zeta_noloc = get_non_local_stuff(matrix)
+
+        ezfio.pseudo_pseudo_v_kl = zip(*m_coef_noloc)
+        ezfio.pseudo_pseudo_n_kl = zip(*m_n_noloc)
+        ezfio.pseudo_pseudo_dz_kl = zip(*m_zeta_noloc)
 
 def get_full_path(file_path):
     file_path = os.path.expanduser(file_path)