#!/usr/bin/python
import sys
import numpy as np
from . import common as PPU
from . import core, cpp_utils
from . import fieldFFT as fFFT
from . import io
from .defaults import d3
verbose = 1
# ===== constants
Fmax_DEFAULT = 10.0
Vmax_DEFAULT = 10.0
# ==== PP Relaxation
[docs]
def Gauss(Evib, E0, w):
return np.exp(-0.5 * ((Evib - E0) / w) ** 2)
[docs]
def symGauss(Evib, E0, w):
return Gauss(Evib, E0, w) - Gauss(Evib, -E0, w)
[docs]
def meshgrid3d(xs, ys, zs):
Xs, Ys, Zs = np.zeros()
Xs, Ys = np.meshgrid(xs, ys)
[docs]
def trjByDir(n, d, p0):
trj = np.zeros((n, 3))
trj[:, 0] = p0[0] + (np.arange(n)[::-1]) * d[0]
trj[:, 1] = p0[1] + (np.arange(n)[::-1]) * d[1]
trj[:, 2] = p0[2] + (np.arange(n)[::-1]) * d[2]
return trj
[docs]
def shift_positions(R, s):
"""
Shifts positions in R by s; returns the result. Needed especially to shift atoms according to the grid origin.
Arguments:
R: list of positional vectors or a 2D array. The first index donoting the item (typically an atom)
the second index determines the coordinate
s: vector that determines the required shift
Returns:
Rs: Rs[ia,i] = R[ia,i] - s[i]
"""
Rs = np.array(R).copy()
Rs[:, 0] += s[0]
Rs[:, 1] += s[1]
Rs[:, 2] += s[2]
return Rs
[docs]
def relaxedScan3D(xTips, yTips, zTips, trj=None, bF3d=False):
if verbose > 0:
print(">>BEGIN: relaxedScan3D()")
if verbose > 0:
print(" zTips : ", zTips)
ntips = len(zTips)
rTips = np.zeros((ntips, 3))
rs = np.zeros((ntips, 3))
fs = np.zeros((ntips, 3))
nx = len(zTips)
ny = len(yTips)
nz = len(xTips)
if bF3d:
fzs = np.zeros((nx, ny, nz, 3))
else:
fzs = np.zeros((nx, ny, nz))
PPpos = np.zeros((nx, ny, nz, 3))
if trj is None:
trj = np.zeros((ntips, 3))
trj[:, 2] = zTips[::-1]
for ix, x in enumerate(xTips):
sys.stdout.write("\033[K")
sys.stdout.flush()
sys.stdout.write(f"\rrelax ix: {ix}")
sys.stdout.flush()
for iy, y in enumerate(yTips):
rTips[:, 0] = trj[:, 0] + x
rTips[:, 1] = trj[:, 1] + y
rTips[:, 2] = trj[:, 2]
core.relaxTipStroke(rTips, rs, fs)
if bF3d:
fzs[:, iy, ix, 0] = (fs[:, 0].copy())[::-1]
fzs[:, iy, ix, 1] = (fs[:, 1].copy())[::-1]
fzs[:, iy, ix, 2] = (fs[:, 2].copy())[::-1]
else:
fzs[:, iy, ix] = (fs[:, 2].copy())[::-1]
PPpos[:, iy, ix, 0] = rs[::-1, 0]
PPpos[:, iy, ix, 1] = rs[::-1, 1]
PPpos[:, iy, ix, 2] = rs[::-1, 2]
if verbose > 0:
print("<<<END: relaxedScan3D()")
return fzs, PPpos
[docs]
def relaxedScan3D_omp(xTips, yTips, zTips, trj=None, bF3d=False, tip_spline=None):
if verbose > 0:
print(">>BEGIN: relaxedScan3D_omp()")
if verbose > 0:
print(" zTips : ", zTips)
nz = len(zTips)
ny = len(yTips)
nx = len(xTips)
rTips = np.zeros((nx, ny, nz, 3))
rs = np.zeros((nx, ny, nz, 3))
fs = np.zeros((nx, ny, nz, 3))
rTips[:, :, :, 0] = xTips[:, None, None]
rTips[:, :, :, 1] = yTips[None, :, None]
rTips[:, :, :, 2] = zTips[::-1][None, None, :]
core.relaxTipStrokes_omp(rTips, rs, fs, tip_spline=tip_spline)
# rs[:,:,:,:] = rs[:,:,::-1,:].transpose(2,1,0,3).copy()
# fs[:,:,:,:] = fs[:,:,::-1,:]
# fzs = fs[:,:,::-1,2].transpose(2,1,0).copy()
rs = rs[:, :, ::-1, :].transpose(2, 1, 0, 3).copy()
if bF3d:
fzs = fs[:, :, ::-1, :].transpose(2, 1, 0, 3).copy()
else:
fzs = fs[:, :, ::-1, 2].transpose(2, 1, 0).copy()
if verbose > 0:
print("<<<END: relaxedScan3D_omp()")
return fzs, rs
# ==== Forcefield grid generation
[docs]
def prepareArrays(FF, Vpot, parameters):
if parameters.gridN[0] <= 0:
PPU.autoGridN()
if FF is None:
gridN = parameters.gridN
FF = np.zeros((gridN[2], gridN[1], gridN[0], 3))
else:
gridN = np.shape(FF)
parameters.gridN = gridN
core.setFF_Fpointer(FF)
if Vpot:
V = np.zeros((gridN[2], gridN[1], gridN[0]))
core.setFF_Epointer(V)
else:
V = None
return FF, V
[docs]
def computeLJ(geomFile, speciesFile, geometry_format=None, save_format=None, computeVpot=False, Fmax=Fmax_DEFAULT, Vmax=Vmax_DEFAULT, ffModel="LJ", parameters=None):
if verbose > 0:
print(">>>BEGIN: computeLJ()")
# --- load species (LJ potential)
FFparams = PPU.loadSpecies(speciesFile)
elem_dict = PPU.getFFdict(FFparams)
# print elem_dict
# --- load atomic geometry
atoms, nDim, lvec = io.loadGeometry(geomFile, format=geometry_format, parameters=parameters)
atomstring = io.primcoords2Xsf(PPU.atoms2iZs(atoms[0], elem_dict), [atoms[1], atoms[2], atoms[3]], lvec)
if verbose > 0:
print(parameters.gridN, parameters.gridO, parameters.gridA, parameters.gridB, parameters.gridC)
iZs, Rs, Qs = PPU.parseAtoms(atoms, elem_dict, autogeom=False, PBC=parameters.PBC, lvec=lvec, parameters=parameters)
# --- prepare LJ parameters
iPP = PPU.atom2iZ(parameters.probeType, elem_dict)
# --- prepare arrays and compute
FF, V = prepareArrays(None, computeVpot, parameters=parameters)
if verbose > 0:
print("FFLJ.shape", FF.shape)
core.setFF_shape(np.shape(FF), lvec, parameters=parameters)
# shift atoms to the coordinate system in which the grid origin is zero
Rs0 = shift_positions(Rs, -lvec[0])
if ffModel == "Morse":
REs = PPU.getAtomsRE(iPP, iZs, FFparams)
core.getMorseFF(Rs0, REs) # THE MAIN STUFF HERE
elif ffModel == "vdW":
vdWDampKind = parameters.vdWDampKind
if vdWDampKind == 0:
cLJs = PPU.getAtomsLJ(iPP, iZs, FFparams)
core.getVdWFF(Rs0, cLJs) # THE MAIN STUFF HERE
else:
REs = PPU.getAtomsRE(iPP, iZs, FFparams)
core.getVdWFF_RE(Rs0, REs, kind=vdWDampKind) # THE MAIN STUFF HERE
else:
cLJs = PPU.getAtomsLJ(iPP, iZs, FFparams)
core.getLennardJonesFF(Rs0, cLJs) # THE MAIN STUFF HERE
# --- post porces FFs
if Fmax is not None:
if verbose > 0:
print("Clamp force >", Fmax)
io.limit_vec_field(FF, Fmax=Fmax)
if (Vmax is not None) and computeVpot:
if verbose > 0:
print("Clamp potential >", Vmax)
V[V > Vmax] = Vmax # remove too large values
# --- save to files ?
if save_format is not None:
if verbose > 0:
print("computeLJ Save ", save_format)
io.save_vec_field("FF" + ffModel, FF, lvec, data_format=save_format, head=atomstring, atomic_info=(atoms[:4], lvec))
if computeVpot:
io.save_scal_field("E" + ffModel, V, lvec, data_format=save_format, head=atomstring, atomic_info=(atoms[:4], lvec))
if verbose > 0:
print("<<<END: computeLJ()")
return FF, V, nDim, lvec
[docs]
def computeDFTD3(input_file, df_params="PBE", geometry_format=None, save_format=None, compute_energy=False, parameters=None):
"""
Compute the Grimme DFT-D3 force field and optionally save to a file. See also :meth:`.add_dftd3`.
Arguments:
input_file: str. Path to input file. Supported formats are .xyz, .xsf, and .cube.
save_format: str or None. If not None, then the generated force field is saved to files FFvdW_{x,y,z} in format
that can be either 'xsf' or 'npy'.
compute_energy: bool. In addition to force, also compute the energy. The energy is saved to file Evdw if save format
is not None.
df_params: str or dict. Functional-specific scaling parameters. Can be a str with the
functional name or a dict with manually specified parameters.
Returns:
FF: np.ndarray of shape (nx, ny, nz, 3). Force field.
V: np.ndarray of shape (nx, ny, nz) or None. Energy, if compute_energy == True.
lvec: np.ndarray of shape (4, 3). Origin and lattice vectors of the force field.
"""
# Load atomic geometry
atoms, nDim, lvec = io.loadGeometry(input_file, format=geometry_format, parameters=parameters)
elem_dict = PPU.getFFdict(PPU.loadSpecies())
iZs, Rs, _ = PPU.parseAtoms(atoms, elem_dict, autogeom=False, PBC=parameters.PBC, lvec=lvec, parameters=parameters)
iPP = PPU.atom2iZ(parameters.probeType, elem_dict)
# Compute coefficients for each atom
df_params = d3.get_df_params(df_params)
coeffs = core.computeD3Coeffs(Rs, iZs, iPP, df_params)
# Compute the force field
FF, V = prepareArrays(None, compute_energy, parameters=parameters)
core.setFF_shape(np.shape(FF), lvec, parameters=parameters)
core.getDFTD3FF(shift_positions(Rs, -lvec[0]), coeffs)
# Save to file
if save_format is not None:
atom_string = io.primcoords2Xsf(PPU.atoms2iZs(atoms[0], elem_dict), atoms[1:4], lvec)
io.save_vec_field("FFvdW", FF, lvec, data_format=save_format, head=atom_string, atomic_info=(atoms[:4], lvec))
if compute_energy:
io.save_scal_field("EvdW", V, lvec, data_format=save_format, head=atom_string, atomic_info=(atoms[:4], lvec))
return FF, V, lvec
[docs]
def computeELFF_pointCharge(geomFile, geometry_format=None, tip="s", save_format=None, computeVpot=False, Fmax=Fmax_DEFAULT, Vmax=Vmax_DEFAULT, parameters=None):
if verbose > 0:
print(">>>BEGIN: computeELFF_pointCharge()")
tipKinds = {"s": 0, "pz": 1, "dz2": 2}
tipKind = tipKinds[tip]
if verbose > 0:
print(" ========= get electrostatic forcefiled from the point charges tip=%s %i " % (tip, tipKind))
# --- load atomic geometry
FFparams = PPU.loadSpecies()
elem_dict = PPU.getFFdict(FFparams)
# print elem_dict
atoms, nDim, lvec = io.loadGeometry(geomFile, format=geometry_format, parameters=parameters)
atomstring = io.primcoords2Xsf(PPU.atoms2iZs(atoms[0], elem_dict), [atoms[1], atoms[2], atoms[3]], lvec)
# --- prepare arrays and compute
if verbose > 0:
print(parameters.gridN, parameters.gridA, parameters.gridB, parameters.gridC)
_, Rs, Qs = PPU.parseAtoms(atoms, elem_dict=elem_dict, autogeom=False, PBC=parameters.PBC, lvec=lvec, parameters=parameters)
FF, V = prepareArrays(None, computeVpot, parameters=parameters)
core.setFF_shape(np.shape(FF), lvec, parameters=parameters)
# shift atoms to the coordinate system in which the grid origin is zero
Rs0 = shift_positions(Rs, -lvec[0])
core.getCoulombFF(Rs0, Qs * PPU.CoulombConst, kind=tipKind) # THE MAIN STUFF HERE
# --- post porces FFs
if Fmax is not None:
if verbose > 0:
print("Clamp force >", Fmax)
io.limit_vec_field(FF, Fmax=Fmax)
if (Vmax is not None) and computeVpot:
if verbose > 0:
print("Clamp potential >", Vmax)
V[V > Vmax] = Vmax # remove too large values
# --- save to files ?
if save_format is not None:
if verbose > 0:
print("computeLJ Save ", save_format)
io.save_vec_field("FFel", FF, lvec, data_format=save_format, head=atomstring, atomic_info=(atoms[:4], lvec))
if computeVpot:
io.save_scal_field("Vel", V, lvec, data_format=save_format, head=atomstring, atomic_info=(atoms[:4], lvec))
if verbose > 0:
print("<<<END: computeELFF_pointCharge()")
return FF, V, nDim, lvec
[docs]
def computeElFF(V, lvec, nDim, tip, computeVpot=False, tilt=0.0, sigma=None, deleteV=True, parameters=None):
rho = None
multipole = None
if sigma is None:
sigma = parameters.sigma
if isinstance(tip, (list, np.ndarray)):
rho = tip
elif isinstance(tip, dict):
multipole = tip
else:
if tip in {"s", "px", "py", "pz", "dx2", "dy2", "dz2", "dxy", "dxz", "dyz"}:
rho = None
multipole = {tip: 1.0}
elif tip.endswith(".xsf"):
rho, lvec_tip, nDim_tip, tiphead = io.loadXSF(tip)
if any(nDim_tip != nDim):
sys.exit("Error: Input file for tip charge density has been specified, but the dimensions are incompatible with the Hartree potential file!")
rho *= -1 # Negative charge density from positive electron density
Fel_x, Fel_y, Fel_z, Vout = fFFT.potential2forces_mem(V, lvec, nDim, rho=rho, sigma=sigma, multipole=multipole, doPot=computeVpot, tilt=tilt, deleteV=deleteV)
FFel = io.packVecGrid(Fel_x, Fel_y, Fel_z)
del Fel_x, Fel_y, Fel_z
return FFel, Vout
[docs]
def loadValenceElectronDict():
valElDict_ = None
namespace = {}
try:
fname_valelec_dict = "valelec_dict.py"
namespace = {}
exec(open(fname_valelec_dict).read(), namespace)
print(" : ", namespace["valElDict"])
valElDict_ = namespace["valElDict"]
print("Valence electrons loaded from local file : ", fname_valelec_dict)
except:
pass
if valElDict_ is None:
namespace = {}
fname_valelec_dict = cpp_utils.PACKAGE_PATH / "defaults" / "valelec_dict.py"
exec(open(fname_valelec_dict).read(), namespace)
valElDict_ = namespace["valElDict"]
print("Valence electrons loaded from default location : ", fname_valelec_dict)
if verbose > 0:
print(" Valence Electron Dict : \n", valElDict_)
return valElDict_
def _getAtomsWhichTouchPBCcell(Rs, elems, nDim, lvec, Rcut, bSaveDebug, fname=None):
inds, Rs_ = PPU.findPBCAtoms3D_cutoff(Rs, np.array(lvec[1:]), Rcut=Rcut) # find periodic images of PBC images of atom of radius Rcut which touch our cell
elems = [elems[i] for i in inds] # atomic number of all relevant peridic images of atoms
if bSaveDebug:
io.saveGeomXSF(fname + "_TouchCell_debug.xsf", elems, Rs_, lvec[1:], convvec=lvec[1:], bTransposed=True) # for debugging - mapping PBC images of atoms to the cell
Rs_ = Rs_.transpose().copy()
return Rs_, elems
[docs]
def getAtomsWhichTouchPBCcell(fname, Rcut=1.0, bSaveDebug=True, geometry_format=None, parameters=None):
atoms, nDim, lvec = io.loadGeometry(fname, format=geometry_format, parameters=parameters)
Rs = np.array(atoms[1:4]) # get just positions x,y,z
elems = np.array(atoms[0])
Rs, elems = _getAtomsWhichTouchPBCcell(Rs, elems, nDim, lvec, Rcut, bSaveDebug, fname)
return Rs, elems
[docs]
def subtractCoreDensities(
rho, lvec, elems=None, Rs=None, fname=None, valElDict=None, Rcore=0.7, bSaveDebugDens=False, bSaveDebugGeom=True, head=io.XSF_HEAD_DEFAULT, parameters=None
):
nDim = rho.shape
if fname is not None:
elems, Rs = getAtomsWhichTouchPBCcell(fname, Rcut=Rcore, bSaveDebug=bSaveDebugDens)
if valElDict is None:
valElDict = loadValenceElectronDict()
print("subtractCoreDensities valElDict ", valElDict)
print("subtractCoreDensities elems ", elems)
cRAs = np.array([(-valElDict[elem], Rcore) for elem in elems])
V = np.linalg.det(lvec[1:]) # volume of triclinic cell
N = nDim[0] * nDim[1] * nDim[2]
dV = V / N # volume of one voxel
if verbose > 0:
print("V : ", V, " N: ", N, " dV: ", dV)
if verbose > 0:
print("sum(RHO): ", rho.sum(), " Nelec: ", rho.sum() * dV, " voxel volume: ", dV) # check sum
core.setFF_shape(rho.shape, lvec, parameters=parameters) # set grid sampling dimension and shape
core.setFF_Epointer(rho) # set pointer to array with density data (to write into)
if verbose > 0:
print(">>> Projecting Core Densities ... ")
core.getDensityR4spline(shift_positions(Rs, -lvec[0]), cRAs.copy()) # Do the job ( the Projection of atoms onto grid )
if verbose > 0:
print("sum(RHO), Nelec: ", rho.sum(), rho.sum() * dV) # check sum
if bSaveDebugDens:
io.saveXSF("rho_subCoreChg.xsf", rho, lvec, head=head)