#*  **************************************************************************
#*
#*  CDTK, Chemical Dynamics Toolkit
#*  A modular system for chemical dynamics applications and more
#*
#*  Copyright (C) 2011, 2012, 2013, 2014, 2015, 2016
#*  Oriol Vendrell, DESY, <oriol.vendrell@desy.de>
#*
#*  Copyright (C) 2017, 2018, 2019
#*  Ralph Welsch, DESY, <ralph.welsch@desy.de>
#*
#*  Copyright (C) 2020, 2021, 2022, 2023
#*  Ludger Inhester, DESY, ludger.inhester@cfel.de
#*
#*  This file is part of CDTK.
#*
#*  CDTK is free software: you can redistribute it and/or modify
#*  it under the terms of the GNU General Public License as published by
#*  the Free Software Foundation, either version 3 of the License, or
#*  (at your option) any later version.
#*
#*  This program is distributed in the hope that it will be useful,
#*  but WITHOUT ANY WARRANTY; without even the implied warranty of
#*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#*  GNU General Public License for more details.
#*
#*  You should have received a copy of the GNU General Public License
#*  along with this program.  If not, see <http://www.gnu.org/licenses/>.
#*
#*  **************************************************************************
#! /usr/bin/env python
import sys
import os
import string
import re
import numpy as np
import CDTK.Tools.Conversion as conv
symbols = []
basis_set      = ''
input_energy   = ''
input_gradient = ''
input_cplmat   = ''
input_init_cplmat = ''
state = 0
hf_charge = 0
input_mo       = ''
MOLCAS = 'molcas'
# Energy calculation
MOLCAS_ENERGY_INP = 'molcas_energy.inp'
MOLCAS_ENERGY_OUT = 'molcas_energy.out'
MOLCAS_ENERGY_XYZ = 'molcas_energy.xyz'
# Gradient calculation
MOLCAS_GRADIENT_INP = 'molcas_gradient.inp'
MOLCAS_GRADIENT_OUT = 'molcas_gradient.out'
MOLCAS_GRADIENT_XYZ = 'molcas_gradient.xyz'
# Non-adiabatic coupling matrix calculation
MOLCAS_CPLMAT_INP = 'molcas_cplmat.inp'
MOLCAS_CPLMAT_OUT = 'molcas_cplmat.out'
MOLCAS_T_XYZ    = 'molcas_cplmat_T.xyz'    # wavefunction at T
MOLCAS_TpDT_XYZ = 'molcas_cplmat_TpDT.xyz' # wavefunction at T+DT
# Molecular orbital calculation 
MOLCAS_MO_INP = 'molcas_mo.inp'
MOLCAS_MO_OUT = 'molcas_mo.out'
MOLCAS_MO_ORBITAL = 'molcas_mo.ScfOrb'
# Molecular spectra calculation TAS
MOLCAS_TAS_INP = 'molcas_tas.inp'
MOLCAS_TAS_OUT = 'molcas_tas.out'
MOLCAS_TAS_XYZ = 'molcas_tas.xyz'
ZERO = 0.0
# Hartree-Fock-Koopmans
ab_initio_setup = {}
[docs]
def runMolcasHFKoopmansEnergy(X_au,**opts):
    """
    Return the Koopmans energy of the given molecule
    Loading an excess electron to the neutral molecule
    Run Molcas for given molecular coordinates.
    Input:
        - X_au: 3N array, molecular coordinates in bohr
    Output: if state == 0, EHF (Hartree Fock)
            if state == 1, ECAS
            ...
    optional arguments:
    - unit: Unit of coordinates
            default: Bohr
            option : Angstrom
    """
    retAll = opts.get('returnAll',False)
    retstates = opts.get('returnStates',[])
    retEHF = opts.get('returnEHF',False)
    unit = opts.get('unit','Bohr')
    X = X_au.copy()
    
    try:
        Fileref = open(MOLCAS_ENERGY_INP,'w')
    except:
        msg = 'File '+MOLCAS_ENERGY_INP+' could not be accessed'
        raise ValueError(msg)
    try:
        Filexyz = open(MOLCAS_ENERGY_XYZ,'w')
    except:
        msg = 'File '+MOLCAS_ENERGY_XYZ+' could not be accessed'
        raise ValueError(msg)
    
    natoms = X.size/3
    X.shape = (natoms,3)
    Filexyz.write('%i\n' % (natoms) )
    if unit is 'Bohr':
        Filexyz.write("Bohr\n")
    elif unit is 'Angstrom':
        Filexyz.write("Angstrom\n")    
    for i,atom in enumerate(X):
        Filexyz.write('%s  %12.7f  %12.7f  %12.7f\n' % (symbols[i],atom[0],atom[1],atom[2]))
    Filexyz.close()
    Fileref.write("&SEWARD\n")
    Fileref.write("  Title=Koopmans molecular energy calculation\n")
    Fileref.write("  coord=%s\n" % (MOLCAS_ENERGY_XYZ))
    Fileref.write("  group=nosym\n")
    Fileref.write("  basis=%s\n" % (ab_initio_setup['basis_set']))
    # SCF
    Fileref.write("&SCF\n")
    Fileref.write("  charge=%i\n" % (ab_initio_setup['scf_charge']))
    Fileref.write("  PRORbitals=2 100.0 3\n") # Orbital with orbitals energies smaller than 100 [au]. in format [3]
    Fileref.close()
    MEXEC = MOLCAS + ' ' + MOLCAS_ENERGY_INP  + ' > ' + MOLCAS_ENERGY_OUT
    status = os.system(MEXEC)
    
    EHF, orbE, NOCCN = getHFKoopmans_Energy(MOLCAS_ENERGY_OUT)
    if EHF is None:
        return None
    if retstates:
        Epes = []
        for s in retstates:
            pes = float(EHF) + float(orbE[NOCCN+s])
            Epes.append(pes)
        return np.array(Epes)   
 
    elif retAll:
        Epes = []
        idx = np.arange(len(orbE),dtype=int)
        for i in idx:
            pes = EHF + float(orbE[NOCCN+i])
            Epes.append(pes)
        return np.array(Epes)
    elif retEHF:
        return EHF
    if state == 0:
        return EHF
    if state > 0:
        return EHF + float(orbE[NOCCN+state-1]) 
[docs]
def getHFKoopmans_Energy(QOUT):
    """
    Get HF and orbital energies from a Molcas output file
    """
    SCF_Energy = ''
    NOCC   = ''
    NOCCN  = ''
    OrbEi  = ''
    
    try:
        StateF=open(QOUT,'r')
    except:
        msg = 'File '+QOUT+' could not be accessed for reading'
        raise ValueError(msg)
    Line=StateF.readline()
    while Line:
        if Line.find('Total SCF energy')>=0:
            StateF.readline()
            SCF_Energy=SCF_Energy+Line.split()[4]
            E0=float(SCF_Energy)
        Line=StateF.readline()
   
    StateF=open(QOUT,'r')
    Line=StateF.readline()
    while Line:
        if Line.find('Convergence problem!') >= 0:
            return None,None
        Line=StateF.readline()
    StateF=open(QOUT,'r')
    Line=StateF.readline()
    while Line:
        if Line.find('Occupied orbitals')>=0:
            StateF.readline()
            NOCC=NOCC+string.split(Line)[2]
            NOCCN=int(NOCC)
        elif Line.find('nOcc=')>=0:
            StateF.readline()
            NOCC=NOCC+string.split(Line)[1]
            NOCCN=int(NOCC)
        # Nosymm is assumed in the calculation
        if Line.find('Molecular orbitals for symmetry species 1: a') >= 0: # initial line
            OrbEi=[]
            StateF.readline()
            while Line.find('--') < 0: # final line
                Line=StateF.readline()
                if Line.find('Energy')>=0:
                    # Two types of orbital energy output.
                    # Type 0 - for few  orbitals. horizontal format
                    # Type 1 - for many orbitals. vertical   format
                    if Line.split().index('Energy') == 0:
                        TypeOrbEOutput = 0
                    elif Line.split().index('Energy') == 1:
                        TypeOrbEOutput = 1
                    else:
                        raise ValueError('Orbital energy output type not defined')
                    if TypeOrbEOutput == 0:
                        OrbEi=OrbEi + Line.split()[1:]
                        Line=StateF.readline()
                    elif TypeOrbEOutput == 1:
                        break
            if TypeOrbEOutput == 1:
                while Line.find('--') < 0: # final line
                    Line=StateF.readline()
                    if len(Line.split()) >= 3:
                        posOccNum = Line.split()[2] # position for occ. num.
                        if posOccNum == '2.0000' or posOccNum == '1.0000':
                            # an effective line for orbital energies
                            OrbEi.append(Line.split()[1])
            return (E0, OrbEi, NOCCN)   
        Line=StateF.readline() 
[docs]
def runMolcasHFKoopmansGradient(X_au,**opts):
    """
    Return the Hartree-Fock-Koopmans energy gradient of given molecule
    Run Molcas for given molecular coordinates.
    
    For the case of simultaneous computation of [gradient]s for [mult]i-[elec]tronic [s]tates.
    Input:
        - X_au: 3N array, molecular coordinates in bohr
    Output:
        - GCAS: The analytical gradient
 
    optional arguments:
    - unit: Unit of coordinates
            default: Bohr
            option : Angstrom
    - nstates_g: number of electronic states requiring gradient
            default = 1
    - state_g: the certain electronic state(s) requiring gradient
            default = None
    """
    retAll = opts.get('returnAll',False)
    retstates = opts.get('returnStates',[])
    nstates_g = opts.get('nstates_g',1)
    state_g  = opts.get('state_g',None)
    unit = opts.get('unit','Bohr')
    if state_g is None:
        raise ValueError('The certain electronic state requiring analytical gradient not defined')
    if len(state_g) != nstates_g:
        raise ValueError('Inconsistency in state_g and nstates_g')
    
    X = X_au.copy()
 
    ndim = X.size
    natoms = ndim/3 
    X.shape = (natoms,3)
    
    # gradient for one electronic state
    GCAS_s = []
    # gradient for required electronic state(s)
    GCAS = np.zeros((nstates_g,ndim),float)
    try:
        Fileref = open(MOLCAS_GRADIENT_INP,'w')
    except:
        msg = 'File '+MOLCAS_GRADIENT_INP+' could not be accessed'
        raise ValueError(msg)
    try:
        Filexyz = open(MOLCAS_GRADIENT_XYZ,'w')
    except:
        msg = 'File '+MOLCAS_GRADIENT_XYZ+' could not be accessed'
        raise ValueError(msg)
    Filexyz.write('%i\n' % (natoms) )
    Filexyz.write("Bohr\n")
    for i,atom in enumerate(X):
        Filexyz.write('%s  %12.7f  %12.7f  %12.7f\n' % (symbols[i],atom[0],atom[1],atom[2]) )
    # [SEWARD]. once
    Fileref.write("&SEWARD\n")
    Fileref.write("  Title=Hartree-Fock-Koopmans gradient\n")
    Fileref.write("  coord=%s\n" % (MOLCAS_GRADIENT_XYZ))
    Fileref.write("  group=nosym\n")
    Fileref.write("  basis=%s\n" % (ab_initio_setup['basis_set']))
    # SCF
    Fileref.write("&SCF\n")
    Fileref.write("charge=%i\n" % (ab_initio_setup['scf_charge']))
    # [RASSCF] [MCLR] [ALASKA]. x nstates_g
    for state_g_s in state_g:
        # RASSCF
        # INPORB
        # Molcas 7.6
        #Fileref.write("! ln -fs $Project.ScfOrb INPORB\n")
        # Molcas 7.8
        Fileref.write(">>LINK FORCE $Project.ScfOrb INPORB\n")
        Fileref.write("&RASSCF\n")
        Fileref.write("  Title=CASSCF Procedure\n")
        Fileref.write("  charge=%s\n" % (ab_initio_setup['casscf_charge']))
        Fileref.write("  spin=%s\n" % (ab_initio_setup['casscf_spin']))
        Fileref.write("  NACTEL=%3i 0 0; Inactive=%3i; Ras2=%3i\n" % (ab_initio_setup['nactel'], ab_initio_setup['inactive'], ab_initio_setup['ras2']) )
        Fileref.write("  CIRoot=%3i %3i 1\n" % (ab_initio_setup['ciroot'], ab_initio_setup['ciroot']))
        Fileref.write("  CIONLY\n")
        Fileref.write("  CIMX=0\n")
        state_bottom  = 0
        if state_g_s != state_bottom: 
            mo_roof  = ab_initio_setup['nocc'] + 1
            mo_alter = mo_roof + state_g_s
            Fileref.write("  ALTER=1; 1 %3i %3i\n" % (mo_alter, mo_roof))
        Fileref.write("  LUMORB\n")
        # ALASKA
        Fileref.write("&ALASKA\n")
        Fileref.write("  PNEW\n")
    Filexyz.close()
    Fileref.close()
 
    MEXEC = MOLCAS + ' ' + MOLCAS_GRADIENT_INP  + ' > ' + MOLCAS_GRADIENT_OUT 
    status = os.system(MEXEC)
    GCAS0 = getHFKoopmans_Gradient(MOLCAS_GRADIENT_OUT, nstates_g)
    GCAS = np.reshape(np.array(GCAS0),(nstates_g,ndim))
    return GCAS 
   
[docs]
def getHFKoopmans_Gradient(QOUT, nstates_g):
    """
    Get analytical gradients for a given ROOT from Molcas output file
    For the case of simultaneous computation of [gradient]s for [mult]i-[elec]tronic [s]tates.
    Input:
        nstates_g -  number of electronic states requiring gradient
            default = 1
    """
    StateF=open(QOUT,'r')
    Line=StateF.readline()
    gradT      = []
    elec_state = 0
    while Line and (elec_state <= nstates_g):
        # compute the gradient for each electronic state
        if Line.find('Molecular gradients') >= 0:
            elec_state += 1
            grad_analy = []
            StateF.readline()
            StateF.readline()
            StateF.readline()
            StateF.readline()
            StateF.readline()
            StateF.readline()
            StateF.readline()
            Line=StateF.readline()
            while Line.find('--------') < 0:
                grad0 = map(float,Line.split()[1:])
                grad_analy.append(grad0)
                Line=StateF.readline()
            gradT.append(grad_analy) # gradient for each electronic state
        Line=StateF.readline()
    return np.array(gradT) 
[docs]
def runMolcasHFKoopmansCplmat(X_au_T,X_au_TpDT,DT,phase_track,state_block,**opts):
    """
    Return the non-adiabatic coupling matrix of the given molecule
    
    Run Molcas for given molecular coordinates.
    
    Input:
        - X_au_T:    3N array, molecular coordinates in bohr at T
        - X_au_TpDT: 3N array, molecular coordinates in bohr at T+DT
        - DT: Classical time-step
        - phase_track: track of electronic wavefunction phase
        - state_block: block of states to be considered for overlap
        
    Output: VD - The non-adiabatic coupling matrix 
    optional arguments:
    - unit: Unit of coordinates
            default: Bohr
            option : Angstrom
    - nstates: number of electronic states
            default = 1
    - is_init: whether to be the initializing classical step 
      for wavefunction overlap calculation of non-adiabatic coupling
            default: False
    - is_reinit: whether to be the re-initializing classical step
      for wavefunction overlap calculation of non-adiabatic coupling
      applied for Hartree-Fock-Koopmans method
            default: False
    """
    nstates = opts.get('nstates',1)
    unit = opts.get('unit','Bohr')
    is_init = opts.get('is_init',False)
    is_reinit = opts.get('is_reinit',False)
    NBlock = len(state_block) # Number of orbitals in the block
    
    X_T    = X_au_T.copy()
    X_TpDT = X_au_TpDT.copy() 
 
    natoms = X_T.size/3
    X_T.shape    = (natoms,3)
    X_TpDT.shape = (natoms,3)
    # Molecular geometry at T
    try:
        Fileref = open(MOLCAS_T_XYZ,'w')
    except:
        msg = 'File '+MOLCAS_T_XYZ+' could not be accessed'
        raise ValueError(msg)
 
    Fileref.write('%i\n' % (natoms) )
    if unit is 'Bohr':
        Fileref.write("Bohr\n")
    elif unit is 'Angstrom':
        Fileref.write("Angstrom\n")
    for i,atom in enumerate(X_T):
        Fileref.write('%s  %12.7f  %12.7f  %12.7f\n' % (symbols[i],atom[0],atom[1],atom[2]) )
    Fileref.close()
    # Molecular geometry at T+DT
    try:
        Fileref = open(MOLCAS_TpDT_XYZ,'w')
    except:
        msg = 'File '+MOLCAS_TpDT_XYZ+' could not be accessed'
        raise ValueError(msg)
    
    Fileref.write('%i\n' % (natoms) )
    if unit is 'Bohr':
        Fileref.write("Bohr\n")
    elif unit is 'Angstrom':
        Fileref.write("Angstrom\n")
    for i,atom in enumerate(X_TpDT):
        Fileref.write('%s  %12.7f  %12.7f  %12.7f\n' % (symbols[i],atom[0],atom[1],atom[2]) )
    Fileref.close()
    
    # Input for non-adiabatic coupling
    # IDX for JOBIPHS
    idx_s = []
    niphs = 2 # for two geometries
    for idx in range(niphs):
        if (idx+1) < 10:
            idx_s.append('00' + str(idx+1))
        elif 10 <= (idx+1) < 100:
            idx_s.append('0'  + str(idx+1))
        elif 100 <= (idx+1) < 1000:
            idx_s.append(str(idx+1))
        else:
            raise ValueError('JobIph not defined')
    #
    # RASSCF wave function for each electronic state
    # For molecular geometry ONE at TIME 0. TWO at TIME DT
    try:
        Fileref = open(MOLCAS_CPLMAT_INP,'w')
    except:
        msg = 'File '+MOLCAS_CPLMAT_INP+' could not be accessed'
        raise ValueError(msg)
    # MOLCAS_HOME
    Fileref.write(">> EXPORT MOLCAS_HOME=%s\n" % (ab_initio_setup['molcas_home']))
    if is_init:
        for geom in range(2):
            # Molcas input for non-adiabatic coupling at TIME 0+DT/2
            # SEWARD
            if is_reinit:
                Fileref.write("! rm *Orb* *molden* *clu* *File *grid\n")
            else:
                if geom == 1:
                    # should not search for molecular orbital files from previous RASSCF procedure
                    Fileref.write("! rm *Orb* *molden* *clu* *File *grid\n")
            Fileref.write("&SEWARD\n")
            Fileref.write("  Title=Hartree-Fock-Koopmans wavefunction\n")
            Fileref.write("  coord=")
            if geom == 0:
                Fileref.write("%s\n" % (MOLCAS_T_XYZ))
            elif geom == 1:
                Fileref.write("%s\n" % (MOLCAS_TpDT_XYZ))
            Fileref.write("  group=nosym\n")
            Fileref.write("  basis=%s\n" % (ab_initio_setup['basis_set']))
            # SCF
            Fileref.write("&SCF\n")
            Fileref.write("  charge=%s\n" % (ab_initio_setup['scf_charge']))
            Fileref.write("> COPY molcas_cplmat.ScfOrb $MOLCAS_HOME/\n")
            # RASSCF
            # define molecular orbitals used for wavefunction overlap calculation
            # [typeindex]
            state_start = state_block[0]
            state_end   = state_block[-1]
            # [mo_floor ... mo_roof]
            mo_roof  = ab_initio_setup['nocc'] + state_end 
            mo_floor = ab_initio_setup['nocc'] + state_start
            Fileref.write("! python $MOLCAS_HOME/typidxalter.py %6i %6i $MOLCAS_HOME/molcas_cplmat.ScfOrb \n" % (mo_floor, mo_roof))
            # 
            Fileref.write("&RASSCF\n")
            Fileref.write("  Title=CASSCF Procedure\n")
            Fileref.write("  charge=%s\n" % (ab_initio_setup['casscf_charge']))
            Fileref.write("  spin=%s\n" % (ab_initio_setup['casscf_spin']))
            if ab_initio_setup['has_occ_sing']:
                NACTel = 2 * NBlock - 1
            else:
                NACTel = 2 * NBlock
            NInact = ab_initio_setup['nocc'] - NBlock + 1
            Fileref.write("  NACTEL=%3i 0 0; Inactive=%3i; Ras2=%3i\n" % (NACTel, NInact, NBlock) )
            Fileref.write("  CIRoot=%3i %3i 1\n" % (NBlock, NBlock))
            Fileref.write("  CIONLY\n")
            Fileref.write("  CIMX=0\n")
            Fileref.write("  LUMORB\n")
            # JobIph
            iph_s = idx_s[geom]
            Fileref.write("> COPY $Project.JobIph JOB%s\n" % (iph_s)) 
            Fileref.write("\n")
    else:    
        # Molcas input for non-adiabatic coupling
        # Molcas input for non-adiabatic coupling at TIME T+DT/2
        # SEWARD
        Fileref.write("! rm *Orb* *molden* *clu* *File *grid\n")
        # [geom 0] is stored at previous classical step
        geom = 1
        Fileref.write("&SEWARD\n")
        Fileref.write("  Title=Hartree-Fock-Koopmans wavefunction\n")
        Fileref.write("  coord=")
        Fileref.write("%s\n" % (MOLCAS_TpDT_XYZ))
        Fileref.write("  group=nosym\n")
        Fileref.write("  basis=%s\n" % (ab_initio_setup['basis_set']))
        # SCF
        Fileref.write("&SCF\n")
        Fileref.write("  charge=%s\n" % (ab_initio_setup['scf_charge']))
        Fileref.write("> COPY molcas_cplmat.ScfOrb $MOLCAS_HOME/\n")
        # RASSCF
        # define molecular orbitals used for wavefunction overlap calculation
        # [typeindex]
        state_start = state_block[0]
        state_end   = state_block[-1]
        # [mo_floor ... mo_roof]
        mo_roof  = ab_initio_setup['nocc'] + state_end
        mo_floor = ab_initio_setup['nocc'] + state_start  
        Fileref.write("! python $MOLCAS_HOME/typidxalter.py %6i %6i $MOLCAS_HOME/molcas_cplmat.ScfOrb \n" % (mo_floor, mo_roof))
        # 
        Fileref.write("&RASSCF\n")
        Fileref.write("  Title=CASSCF Procedure\n")
        Fileref.write("  charge=%s\n" % (ab_initio_setup['casscf_charge']))
        Fileref.write("  spin=%s\n" % (ab_initio_setup['casscf_spin']))
        if ab_initio_setup['has_occ_sing']:
            NACTel = 2 * NBlock - 1
        else:
            NACTel = 2 * NBlock
        NInact = ab_initio_setup['nocc'] - NBlock + 1
        Fileref.write("  NACTEL=%3i 0 0; Inactive=%3i; Ras2=%3i\n" % (NACTel, NInact, NBlock) )
        Fileref.write("  CIRoot=%3i %3i 1\n" % (NBlock, NBlock))
        Fileref.write("  CIONLY\n")
        Fileref.write("  CIMX=0\n")
        Fileref.write("  LUMORB\n")
        # JobIph
        iph_s = idx_s[geom]
        Fileref.write("> COPY $MOLCAS_HOME/JOB%s JOB%s\n" % (idx_s[0],idx_s[0]))
        Fileref.write("> COPY $Project.JobIph JOB%s\n" % (iph_s)) 
        Fileref.write("\n")
    # Overall RASSI for collected RASSCF wave functions
    # Overlap of wavefunctions
    # RASSI at TIME 0+DT/2 and T+DT/2
    Fileref.write("&RASSI\n")
    Fileref.write("  NR OF JOBIPHS=2 %3i %3i\n" % (NBlock, NBlock))
    for i in range(NBlock):
        Fileref.write(" %2i " % (i+1))
    Fileref.write(";")
    for i in range(NBlock):
        Fileref.write(" %2i " % (i+1))
    Fileref.write("\n")
    Fileref.write("  overlaps\n")
    Fileref.write("  ONEL\n")
    Fileref.write("\n")
    
    # Overwrite wavefunction at TIME with wavefunction at TIMEpDT
    iph_T    = idx_s[0]
    iph_TpDT = idx_s[1]
    Fileref.write(">> RM JOB%s\n" % (iph_T))
    Fileref.write("\n")
    Fileref.write(">> COPY JOB%s $MOLCAS_HOME/JOB%s\n" % (iph_TpDT,iph_T))
    Fileref.close()
    
    MEXEC = MOLCAS + ' ' + MOLCAS_CPLMAT_INP  + ' > ' + MOLCAS_CPLMAT_OUT
    status = os.system(MEXEC)
  
    # The non-adiabatic coupling matrix V.D at T+DT/2
    # The blockweis non-adiabatic coupling matrix
    # [NBlock] x [NBlock]   sub-block
    idx_block_init = state_block[0]
    Block_VD = getHFKoopmans_Cplmat(MOLCAS_CPLMAT_OUT,DT,phase_track,idx_block_init,nstates_block=NBlock,nstates=nstates)
    # Construct the full skew-Hermitian non-adiabatic coupling matrix from the block
    # [nstates] x [nstates] full-matrix
    VD = np.zeros((nstates,nstates),float) 
    for i in range(NBlock):
        for j in range(NBlock):
            if i != j:
                # compute corresponding position in the full matrix
                i_fm = idx_block_init + i 
                j_fm = idx_block_init + j
                VD[i_fm,j_fm] = Block_VD[i,j] 
    if VD is None:
        return None
    else:
        return VD 
[docs]
def getHFKoopmans_Cplmat(QOUT,DT,phase_track,idx_block_init,**opts):
    """
    Return the non-adiabatic coupling matrix at T+DT/2
    calculated from overlap of wavefunctions at T and T+DT
    Input:
    - QOUT: Molcas output of wavefunction overlap calculation
    - DT: Classical time-step
    - phase_track: track of electronic wavefunction phase
    - idx_block_init: the initial state of the block
    optional arguments:
    - nstates_block: number of electronic states in the block
            default = 2 
    - nstates: number of electronic states
            default = 1
    """
    nstates_block = opts.get('nstates_block',2) 
    nstates       = opts.get('nstates',1)
    try:
        StateF=open(QOUT,'r')
    except:
        msg = 'File '+QOUT+' could not be accessed for reading'
        raise ValueError(msg)
    Line=StateF.readline()
    Line=StateF.readline()
    DIJ0 = []
    DDIJ = []
    VD = []
    while Line:
        if Line.find('OVERLAP MATRIX FOR THE ORIGINAL STATES') >= 0:
            AIJ = []
            StateF.readline()
            Line = StateF.readline()
            while Line.find('--- Stop Module') < 0:
                AIJ  = AIJ + Line.split()[0:]
                Line = StateF.readline()
        Line=StateF.readline()
    DIJ = []
    idx_start_overlap = nstates_block * (nstates_block + 1) / 2 # exclude self-overlap
    idx_r = idx_start_overlap
    for i in range(nstates_block):
        for j in range(nstates_block):
            DIJ0 = float(AIJ[idx_r+j])
            DIJ.append(DIJ0)
        idx_r += nstates_block + i + 1
    
    DDIJ = np.reshape(np.array(DIJ),(nstates_block,nstates_block))
    # adjust the phase of electronic wavefunction
    # blockwise
    phase_current = np.ones(nstates, dtype=float)
    for i in range(nstates_block):
        # compute corresponding position in the full matrix
        i_fm = idx_block_init + i
        if DDIJ[i,i] < 0.0:
            phase_current[i_fm] = -1.0 * phase_track[i_fm]
        else:
            phase_current[i_fm] =  1.0 * phase_track[i_fm]
    for i in range(nstates_block):
        for j in range(nstates_block):
            i_fm = idx_block_init + i
            j_fm = idx_block_init + j
            DDIJ[i,j] = DDIJ[i,j] * phase_track[i_fm] * phase_current[j_fm]
    phase_track = phase_current.copy()        
    # calculate the non-adiabtic coupling matrix
    VD = np.zeros((nstates_block,nstates_block),float)
    for i in range(nstates_block):
        for j in range(nstates_block):
             VD[i,j] = DDIJ[i,j] - DDIJ[j,i]
    VD = VD / 2.0 / DT
 
    return VD 
[docs]
def runMolcasHFKoopmansAOCoeff(X_au,state_current,**opts):
    """
    Return the coefficients of atomic orbitals 
    for the molecular orbital of current electronic state
    Input:
        - X_au: 3N array, molecular coordinates in bohr
        - state_current: current electronic state
    
    """
    unit = opts.get('unit','Bohr')
    X = X_au.copy()
    
    try:
        Fileref = open(MOLCAS_MO_INP,'w')
    except:
        msg = 'File '+MOLCAS_MO_INP+' could not be accessed'
        raise ValueError(msg)
    
    natoms = X.size/3
    X.shape = (natoms,3)
    # Molcas input for molecular orbital
    Fileref.write("&SEWARD\n")
    Fileref.write("  Title=Molecular orbital calculation\n")
    Fileref.write("  coord\n")
    Fileref.write('%i\n' % (natoms) )
    if unit is 'Bohr':
        Fileref.write("Bohr\n")
    elif unit is 'Angstrom':
        Fileref.write("Angstrom\n")    
    for i,atom in enumerate(X):
        Fileref.write('%s  %12.7f  %12.7f  %12.7f\n' % (symbols[i],atom[0],atom[1],atom[2]) )
    #    
    Fileref.write("  group=nosym\n")
    Fileref.write("  basis=%s\n" % (ab_initio_setup['basis_set']))
    # SCF
    Fileref.write("&SCF\n")
    Fileref.write("  charge=%s\n" % (ab_initio_setup['scf_charge']))
    Fileref.close()
    
    MEXEC = MOLCAS + ' ' + MOLCAS_MO_INP  + ' > ' + MOLCAS_MO_OUT
    status = os.system(MEXEC)
    NOCC     = getAO_NOCC (MOLCAS_MO_OUT) # the number of occupied orbitals
    AO_coeff = getAO_Coeff(MOLCAS_MO_ORBITAL, NOCC, state_current) # the coefficients of atomic orbitals
    return np.array(AO_coeff) 
[docs]
def getAO_NOCC(QOUT):
    """
    Return the number of occupied orbitals
    """
    NOCC   = ''
    NOCCN  = ''
    
    try:
        StateF=open(QOUT,'r')
    except:
        msg = 'File '+QOUT+' could not be accessed for reading'
        raise ValueError(msg)
    Line=StateF.readline()
    while Line:
        if Line.find('Convergence problem!') >= 0:
            raise ValueError('Convergence problem')
        Line=StateF.readline()
    StateF=open(QOUT,'r')
    Line=StateF.readline()
    while Line:
        if Line.find('Occupied orbitals')>=0:
            StateF.readline()
            NOCC=NOCC+string.split(Line)[2]
            NOCCN=int(NOCC)
            return NOCCN
        elif Line.find('nOcc=')>=0:
            StateF.readline()
            NOCC=NOCC+string.split(Line)[1]
            NOCCN=int(NOCC)
            return NOCCN 
        Line=StateF.readline() 
[docs]
def getAO_Coeff(QOUT, NOCC, state_current):
    """
    Return coefficients of atomic orbitals 
    for the molecular orbital of current electronic state [NOCC - state_current]. Koopmans theorem.
    """
    try:
        StateF=open(QOUT,'r')
    except:
        msg = 'File '+QOUT+' could not be accessed for reading'
        raise ValueError(msg)
    # the Koopmans molecular orbital of the current electronic state
    orbital_current = NOCC + state_current + 1
    
    # the atomic orbital coefficients of the current molecular orbital
    AO_coeff = []
   
    Line=StateF.readline()
    while Line:
        if Line.find('* ORBITAL') >= 0:
            if int(Line.split()[3]) == orbital_current:
                # the atomic orbital coefficients for
                # the current molecular orbital
                Line=StateF.readline()
                while Line.find('* ORBITAL') < 0 and Line.find('#OCC') < 0: # next molecular orbital
                    # Molcas output.
                    if len(Line[0 :18]) > 1 : AO_coeff.append(float(Line[0 :18]))
                    if len(Line[18:36]) > 1 : AO_coeff.append(float(Line[18:36]))
                    if len(Line[36:54]) > 1 : AO_coeff.append(float(Line[36:54]))
                    if len(Line[54:72]) > 1 : AO_coeff.append(float(Line[54:72]))
                    Line=StateF.readline()
                return AO_coeff    
        Line=StateF.readline()