#!/usr/bin/env python
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import numpy
import h5py
import pyscf.gto

def load_chkfile_key(chkfile, key):
    return load(chkfile, key)
def load(chkfile, key):
    fh5 = h5py.File(chkfile, 'r')
    val = fh5[key].value
    fh5.close()
    return val

def dump_chkfile_key(chkfile, key, value):
    dump(chkfile, key, value)
def dump(chkfile, key, value):
    if h5py.is_hdf5(chkfile):
        fh5 = h5py.File(chkfile)
        if key in fh5:
            del(fh5[key])
    else:
        fh5 = h5py.File(chkfile, 'w')
    fh5[key] = value
    fh5.close()


###########################################
def load_scf(chkfile):
    fh5 = h5py.File(chkfile, 'r')
    mol = pyscf.gto.Mole()
    mol.verbose = 0
    mol.output = '/dev/null'
    moldic = eval(fh5['mol'].value)
    mol.build(False, False, **moldic)
    scf_rec = {
        'hf_energy': fh5['scf/hf_energy'].value,
        'mo_energy': fh5['scf/mo_energy'].value,
        'mo_occ'   : fh5['scf/mo_occ'   ].value,
        'mo_coeff' : fh5['scf/mo_coeff' ].value,
    }
    fh5.close()
    return mol, scf_rec

def dump_scf(mol, chkfile, hf_energy, mo_energy, mo_coeff, mo_occ):
    '''save temporary results'''
    if h5py.is_hdf5(chkfile):
        fh5 = h5py.File(chkfile)
        if 'scf' in fh5:
            del(fh5['scf'])
    else:
        fh5 = h5py.File(chkfile, 'w')
    if 'mol' not in fh5:
        fh5['mol'] = format(mol.pack())
    fh5['scf/hf_energy'] = hf_energy
    fh5['scf/mo_energy'] = mo_energy
    fh5['scf/mo_occ'   ] = mo_occ
    fh5['scf/mo_coeff' ] = mo_coeff
    fh5.close()


@@ -422,37 +422,36 @@
 kfile):%0A        
-fh5 =
+with
  h5py.File(chkfi
@@ -444,34 +444,52 @@
 5py.File(chkfile
-)%0A
+, 'r+') as fh5:%0A    
         if key i
@@ -499,32 +499,36 @@
 h5:%0A            
+    
 del(fh5%5Bkey%5D)%0A  
@@ -525,16 +525,45 @@
 5%5Bkey%5D)%0A
+            fh5%5Bkey%5D = value%0A
     else
@@ -564,37 +564,36 @@
   else:%0A        
-fh5 =
+with
  h5py.File(chkfi
@@ -596,25 +596,41 @@
 hkfile, 'w')
-%0A
+ as fh5:%0A        
     fh5%5Bkey%5D
@@ -638,32 +638,16 @@
 = value%0A
-    fh5.close()%0A
 %0A%0A######
@@ -707,29 +707,28 @@
 kfile):%0A    
-fh5 =
+with
  h5py.File(c
@@ -733,27 +733,34 @@
 (chkfile
-, 'r')%0A
+) as fh5:%0A  
     
+  
 mol = py
@@ -770,24 +770,28 @@
 .gto.Mole()%0A
+    
     mol.verb
@@ -798,24 +798,28 @@
 ose = 0%0A    
+  
+  
 mol.output =
@@ -831,16 +831,20 @@
 v/null'%0A
+    
     mold
@@ -871,24 +871,28 @@
 .value)%0A    
+  
+  
 mol.build(Fa
@@ -913,16 +913,20 @@
 moldic)%0A
+    
     scf_
@@ -941,16 +941,20 @@
         
+   
+ 
 'hf_ener
@@ -978,32 +978,36 @@
 energy'%5D.value,%0A
+    
         'mo_ener
@@ -1043,24 +1043,28 @@
 ue,%0A        
+   
+ 
 'mo_occ'   :
@@ -1096,24 +1096,28 @@
 ue,%0A        
+    
 'mo_coeff' :
@@ -1148,31 +1148,11 @@
 lue,
-%0A   
  %7D%0A
-    fh5.close()%0A
     
@@ -1303,37 +1303,36 @@
 kfile):%0A        
-fh5 =
+with
  h5py.File(chkfi
@@ -1334,17 +1334,110 @@
 chkfile)
-%0A
+ as fh5:%0A            if 'mol' not in fh5:%0A                fh5%5B'mol'%5D = format(mol.pack())%0A    
         
@@ -1445,32 +1445,36 @@
 f 'scf' in fh5:%0A
+    
             del(
@@ -1485,30 +1485,24 @@
 'scf'%5D)%0A    
-else:%0A
         fh5 
@@ -1504,100 +1504,232 @@
  fh5
- = h5py.File(chkfile, 'w')%0A    if 'mol' not in fh5:%0A        fh5%5B'mol'%5D = format(mol.pack())%0A
+%5B'scf/hf_energy'%5D = hf_energy%0A            fh5%5B'scf/mo_energy'%5D = mo_energy%0A            fh5%5B'scf/mo_occ'   %5D = mo_occ%0A            fh5%5B'scf/mo_coeff' %5D = mo_coeff%0A    else:%0A        with h5py.File(chkfile, 'w') as fh5:%0A        
     
@@ -1753,32 +1753,40 @@
 y'%5D = hf_energy%0A
+        
     fh5%5B'scf/mo_
@@ -1802,24 +1802,32 @@
 = mo_energy%0A
+        
     fh5%5B'scf
@@ -1844,24 +1844,32 @@
  %5D = mo_occ%0A
+        
     fh5%5B'scf
@@ -1891,26 +1891,10 @@
 mo_coeff
-%0A    fh5.close()
 %0A%0A

fix problems of reading incomplete chkfile