#read saved configuration
configuration0 = nlread('piezoelectric_tensor_ci.hdf5', object_id='gID000')[0]

# Get the fractional coordinates of read configuration
R0 = configuration0.fractionalCoordinates()

# Get the elements
elements = configuration0.elements()

# Get the lattice
lattice = configuration0.bravaisLattice()

# From the lattice extract unit cell volume and length of lattice vector in z-direction
volume = lattice.unitCellVolume()

vectors = lattice.primitiveVectors()
c = vectors[2][2]

# Get the calculator
calculator = configuration0.calculator()

# Number of atoms
numberOfAtoms = len(R0)

# Fractional displacement in the +/- z direction
delta_z = 0.01

# Array for storing the calculated Born Charges
bornCharges = numpy.zeros(numberOfAtoms)


# Loop over atoms in unit cell
for nAtom in range(numberOfAtoms):
    # Loop over displacement in positive/negative z-direction
    
    # List with polarization values
    Pz = []
    for s in [1,-1]:
        # Make a copy of the initial coordinates
        R = R0.copy()
        
        # Displace z-coordinate if atom 'nAtom'
        R[nAtom,2] += s*delta_z
        
        # Make a new configuration with the displaced atom
        configuration = BulkConfiguration(bravais_lattice=lattice,
                                                elements=elements,
                                                fractional_coordinates=R)
        
        # Set the calculator using the saved configuration as initial state
        configuration.setCalculator(calculator,initial_state=configuration0)
        
        # Update the configuration (DFT calculation)
        configuration.update()
        
        # Calculate polarization. Only use fine k-sampling in the z-direction
        polarization = Polarization(configuration=configuration,
                                    kpoints_a=MonkhorstPackGrid(2,2,2),
                                    kpoints_b=MonkhorstPackGrid(2,2,2),
                                    kpoints_c=MonkhorstPackGrid(5,5,20),
                                    )
        
        # Print polarization
        nlprint(polarization)
        
        # Get the total cartesian polarization
        Pt = polarization.totalCartesianPolarization()
        
        # Append the z-component to the Pz list
        Pz.append(Pt[2])
    
    # Make a finite difference approximation for the derivative
    dP = (Pz[0]-Pz[1])/(2*delta_z*c)
    
    # Calculate Born charge
    born_charge = volume*dP
    
    # Add the value (in units of electron charge) to the list 'bornCharges'
    bornCharges[nAtom] = born_charge.inUnitsOf(elementary_charge) 



# Finally print out the results
print('')
print('+------------------------------+')
print('| Born effective charges (e)   |')
print('+------------------------------+')

for nAtom in range(numberOfAtoms):
    print('  %2s' %elements[nAtom].symbol() + '  :      %4.4f       ' %bornCharges[nAtom]    )

print('+------------------------------+')
print('  Sum :      %4.4f       ' %numpy.sum(bornCharges))
print('+------------------------------+')