#!/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
from .logging_utils import get_logger
logger = get_logger("HighLevel")
# ===== 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 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):
logger.info(">>BEGIN: relaxedScan3D()")
logger.debug(f" zTips : {zTips}")
ntips = len(zTips)
rTips = np.zeros((ntips, 3))
rs = np.zeros((ntips, 3))
fs = np.zeros((ntips, 3))
nx = len(xTips)
ny = len(yTips)
nz = len(zTips)
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[ix, iy, :, 0] = fs[::-1, 0]
fzs[ix, iy, :, 1] = fs[::-1, 1]
fzs[ix, iy, :, 2] = fs[::-1, 2]
else:
fzs[:, iy, ix] = fs[::-1, 2]
PPpos[ix, iy, :, 0] = rs[::-1, 0]
PPpos[ix, iy, :, 1] = rs[::-1, 1]
PPpos[ix, iy, :, 2] = rs[::-1, 2]
logger.info("<<END: relaxedScan3D()")
return fzs, PPpos
[docs]
def relaxedScan3D_omp(xTips, yTips, zTips, trj=None, bF3d=False, tip_spline=None):
logger.info(">>BEGIN: relaxedScan3D_omp()")
logger.debug(f" zTips : {zTips}")
nx = len(xTips)
ny = len(yTips)
nz = len(zTips)
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, :].copy()
if bF3d:
fzs = fs[:, :, ::-1, :].copy()
else:
fzs = fs[:, :, ::-1, 2].copy()
logger.info("<<END: relaxedScan3D_omp()")
return fzs, rs
# ==== Forcefield grid generation
[docs]
def setFF(FF=None, computeVpot=False, n=None, lvec=None, parameters=None):
# Find gridN
if FF is not None:
gridN = np.array(np.shape(FF)[0:3], dtype=np.int32)
if parameters is not None:
parameters.gridN = gridN
if n is not None:
n[0:3] = gridN
else:
if n is not None:
gridN = np.array(n, dtype=np.int32)
elif parameters is not None:
if parameters.gridN[0] <= 0:
gridN = PPU.autoGridN(parameters)
else:
gridN = parameters.gridN
else:
raise ValueError("FF dimensions not set !!")
# Create a new array for FF if needed
FF = np.zeros((gridN[0], gridN[1], gridN[2], 3))
logger.debug(f"setFF() gridN: {gridN}")
core.setGridN(gridN)
# Set pointer to FF
if len(FF.shape) == 4 and FF.shape[-1] == 3:
logger.debug("setFF: Creating a pointer to a vector field")
core.setFF_Fpointer(FF)
elif len(FF.shape) == 3 or FF.shape[-1] == 1:
logger.debug("setFF: Creating a pointer to a scalar field")
core.setFF_Epointer(FF)
if computeVpot:
logger.warning("setFF: computeVpot required but ignored because FF itself is scalar!")
computeVpot = False
else:
raise ValueError("setFF: Array dimensions wrong for both vector and array field !!")
# Create a scalar (potential) field if required
if computeVpot:
logger.debug("setFF: Creating a pointer to a scalar field")
V = np.zeros((gridN[0], gridN[1], gridN[2]))
core.setFF_Epointer(V)
else:
V = None
# Set lattice vectors or grid geometry
if (lvec is None) and (parameters is not None):
lvec = np.array(
[
parameters.gridA,
parameters.gridB,
parameters.gridC,
],
dtype=np.float64,
).copy()
lvec = np.array(lvec, dtype=np.float64)
if lvec.shape == (3, 3):
lvec = lvec.copy()
elif lvec.shape == (4, 3):
lvec = lvec[1:, :].copy()
else:
raise ValueError("lvec matrix has a wrong format !!")
logger.debug(f"setFF: lvec = {lvec}")
core.setGridCell(lvec)
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):
logger.info(">>>BEGIN: computeLJ()")
# --- load species (LJ potential)
FFparams = PPU.loadSpecies(speciesFile)
elem_dict = PPU.getFFdict(FFparams)
# --- 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)
logger.debug(f"{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 = setFF(None, computeVpot, lvec=lvec, parameters=parameters)
logger.debug(f"FFLJ.shape {FF.shape}")
# 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:
logger.debug(f"Clamp force > {Fmax}")
io.limit_vec_field(FF, Fmax=Fmax)
if (Vmax is not None) and computeVpot:
logger.debug(f"Clamp potential > {Vmax}")
V[V > Vmax] = Vmax # remove too large values
# --- save to files ?
if save_format is not None:
logger.info(f"Saving Lennard-Jones force field to {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))
logger.info("<<<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 = setFF(None, compute_energy, lvec=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):
logger.info(">>>BEGIN: computeELFF_pointCharge()")
tipKinds = {"s": 0, "pz": 1, "dz2": 2}
tipKind = tipKinds[tip]
logger.debug(f" ========= get electrostatic forcefield from the point charges tip={tip} {tipKind} ")
# --- load atomic geometry
FFparams = PPU.loadSpecies()
elem_dict = PPU.getFFdict(FFparams)
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
logger.debug(f"{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 = setFF(None, computeVpot, lvec=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:
logger.debug(f"Clamp force > {Fmax}")
io.limit_vec_field(FF, Fmax=Fmax)
if (Vmax is not None) and computeVpot:
logger.debug(f"Clamp potential > {Vmax}")
V[V > Vmax] = Vmax # remove too large values
# --- save to files ?
if save_format is not None:
logger.debug(f"computeELFF 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))
logger.info("<<<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.loadXSFData(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)
logger.debug(f" : {namespace['valElDict']}")
valElDict_ = namespace["valElDict"]
logger.debug(f"Valence electrons loaded from local file : {fname_valelec_dict}")
except:
pass
if valElDict_ is None:
from .defaults import valelec_dict
valElDict_ = valelec_dict.valElDict
logger.debug(f"Valence electrons loaded from defaults")
logger.debug(f" 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()
logger.debug(f"subtractCoreDensities valElDict {valElDict}")
logger.debug(f"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
logger.debug(f"V : {V}, N: {N}, dV: {dV}")
logger.debug(f"sum(RHO): {rho.sum()}, Nelec: {rho.sum() * dV}, voxel volume: {dV}") # check sum
# set sampling grid (dimension, shape, and pointer)
setFF(rho, lvec=lvec, parameters=parameters)
logger.debug(">>> Projecting Core Densities ... ")
core.getDensityR4spline(shift_positions(Rs, -lvec[0]), cRAs.copy()) # Do the job ( the Projection of atoms onto grid )
logger.debug(f"sum(RHO), Nelec: {rho.sum()}, {rho.sum() * dV}") # check sum
if bSaveDebugDens:
io.saveXSFData("rho_subCoreChg.xsf", rho, lvec, head=head)