def test_parallel(use_delta):
benzene_traj = ase.io.read(os.path.join(test_dir, 'benzene_test', 'benzene.xyz'), '0')
# Break symmetry
positions = benzene_traj.get_positions()
positions[0, 1] += 0.05
positions[3, 1] += 0.05
benzene_traj.set_positions(positions)
density_getter = xc.utils.SiestaDensityGetter(binary=True)
rho, unitcell, grid = density_getter.get_density(os.path.join(test_dir, 'benzene_test', 'benzene.RHOXC'))
positions = benzene_traj.get_positions() / Bohr
species = benzene_traj.get_chemical_symbols()
if use_delta:
drho, _, _ = density_getter.get_density(os.path.join(test_dir, 'benzene_test', 'benzene.DRHO'))
rho_const = (rho - drho)
benzene_nxc = xc.NeuralXC(os.path.join(test_dir, 'benzene_test', 'dbenzene'))
benzene_nxc.initialize(unitcell, grid, positions, species)
benzene_nxc.projector = xc.projector.DeltaProjector(benzene_nxc.projector)
benzene_nxc.projector.set_constant_density(rho_const, positions, species)
else:
benzene_nxc = xc.NeuralXC(os.path.join(test_dir, 'benzene_test', 'benzene'))
benzene_nxc.initialize(unitcell, grid, positions, species)
V_serial = benzene_nxc.get_V(rho, calc_forces=False)[1]
benzene_nxc.max_workers = 4
V_parallel = benzene_nxc.get_V(rho, calc_forces=False)[1]
assert np.allclose(V_serial, V_parallel, atol=1e-6, rtol=1e-5)
def merge_driver(chained, merged):
""" Converts the tensorflow estimator inside a NXCPipeline to a simple
numpy base estimator"""
nxc_tf = xc.NeuralXC(chained)
pipeline = nxc_tf._pipeline
label, estimator = pipeline.steps[-1]
_, npestimator = pipeline.steps[-2]
if not isinstance(npestimator, NumpyNetworkEstimator):
raise Exception('Something went wrong. Second to last pipeline element'\
+' must be NumpyNetworkEstimator')
if not isinstance(estimator, NumpyNetworkEstimator):
if not isinstance(estimator, NetworkWrapper):
raise Exception('Something went wrong. Last pipeline element'\
+' must be an estimator')
else:
convert_tf(tf_path=chained, np_path=merged)
chained = merged
nxc_tf = xc.NeuralXC(chained)
pipeline = nxc_tf._pipeline
label, estimator = pipeline.steps[-1]
_, npestimator = pipeline.steps[-2]
if not npestimator.trunc:
npestimator = npestimator.trunc_after(-1)
pipeline.steps[-2] = (label, ChainedEstimator([npestimator,
estimator]).merge())
pipeline.steps = pipeline.steps[:-1]
nxc_tf._pipeline.save(merged, True, True)
def test_mybox():
benzene_nxc = xc.NeuralXC(
os.path.join(test_dir, 'benzene_test', 'benzene.jit'))
benzene_traj = ase.io.read(
os.path.join(test_dir, 'benzene_test', 'benzene.xyz'), '0')
density_getter = xc.utils.SiestaDensityGetter(binary=True)
rho, unitcell, grid = density_getter.get_density(
os.path.join(test_dir, 'benzene_test', 'benzene.RHOXC'))
grid_np = np.array(grid)
positions = benzene_traj.get_positions() / Bohr
# Break symmetries
positions[:, 0] += 0.02
positions[:, 1] += 0.01
positions[:, 2] += 0.12
species = benzene_traj.get_chemical_symbols()
a = np.linalg.norm(unitcell, axis=1) / grid[:3]
benzene_nxc.initialize(unitcell=unitcell,
grid=grid,
positions=positions,
species=species)
basis_models = benzene_nxc.basis_models
projector_models = benzene_nxc.projector_models
unitcell = torch.from_numpy(unitcell).double()
grid = torch.from_numpy(grid).double()
positions = torch.from_numpy(positions).double()
a = torch.from_numpy(a).double()
for pos, spec in zip(positions, species):
c_jit = 0
box_lim = [0, int(grid_np[0] / 2) + 5, grid_np[0]]
for ibox in range(2):
for jbox in range(2):
for kbox in range(2):
my_box = np.zeros([3, 2])
my_box[:, 1] = grid
my_box[0, 0] = box_lim[ibox]
my_box[0, 1] = box_lim[ibox + 1]
my_box[1, 0] = box_lim[jbox]
my_box[1, 1] = box_lim[jbox + 1]
my_box[2, 0] = box_lim[kbox]
my_box[2, 1] = box_lim[kbox + 1]
my_box = my_box.astype(int)
rho_jit = torch.from_numpy(
rho[my_box[0, 0]:my_box[0, 1], my_box[1, 0]:my_box[1,
1],
my_box[2, 0]:my_box[2, 1]]).double()
my_box = torch.from_numpy(my_box).double()
rad, ang, box = basis_models[spec](pos, unitcell, grid,
my_box)
rsize = rad.size()
if not rsize[-1]: continue
c_jit += projector_models[spec](rho_jit, pos, unitcell,
grid, rad, ang,
box).detach().numpy()
assert np.allclose(c_jit, np.load(os.path.join(test_dir,
'my_box_ref.npy')))
def test_neuralxc_benzene():
benzene_nxc = xc.NeuralXC(os.path.join(test_dir, 'benzene_test', 'benzene'))
benzene_traj = ase.io.read(os.path.join(test_dir, 'benzene_test', 'benzene.xyz'), '0')
density_getter = xc.utils.SiestaDensityGetter(binary=True)
rho, unitcell, grid = density_getter.get_density(os.path.join(test_dir, 'benzene_test', 'benzene.RHOXC'))
positions = benzene_traj.get_positions() / Bohr
species = benzene_traj.get_chemical_symbols()
benzene_nxc.initialize(unitcell, grid, positions, species)
V, forces = benzene_nxc.get_V(rho, calc_forces=True)[1]
V = V / Hartree
forces = forces / Hartree * Bohr
def convert_tf(tf_path, np_path):
""" Converts the tensorflow estimator inside a NXCPipeline to a simple
numpy base estimator"""
nxc_tf = xc.NeuralXC(tf_path)
pipeline = nxc_tf._pipeline
C = {}
basis = pipeline.get_basis_instructions()
for sym in basis:
if len(sym) > 2: continue
C[sym] = np.zeros([1, 1, basis[sym]['n'] * basis[sym]['l']**2])
D = nxc_tf.symmetrizer.get_symmetrized(C)
nxc_tf._pipeline.predict(D)
nxc_tf._pipeline.save(np_path, True, True)
def test_radial_model():
from pyscf import dft, gto
mol = gto.M(atom='O 0 0 0; H 0 1 0 ; H 0 0 1', basis='6-31g*')
mf = dft.RKS(mol)
mf.xc = 'PBE'
mf.grids.level = 5
mf.kernel()
model = xc.NeuralXC(test_dir[:-len('neuralxc/tests/')] +
'/examples/models/NXC-W01/nxc_w01_radial.jit')
rho = pyscf.dft.numint.get_rho(mf._numint, mol, mf.make_rdm1(), mf.grids)
model.initialize(grid_coords=mf.grids.coords,
grid_weights=mf.grids.weights,
positions=np.array([[0, 0, 0], [0, 1, 0], [0, 0, 1]]) /
Bohr,
species=['O', 'H', 'H'])
res = model.get_V(rho)[0]
assert np.allclose(res, np.load(test_dir + '/rad_energy.npy'))
def eval_driver(hdf5,
model='',
plot=False,
savefig='',
cutoff=0.0,
predict=False,
dest='prediction'):
""" Evaluate fitted NXCPipeline on dataset and report statistics
"""
hdf5 = hdf5
if predict:
hdf5.append(hdf5[1])
cutoff = 0
else:
cutoff = cutoff
datafile = h5py.File(hdf5[0], 'r')
if not model == '':
model = xc.NeuralXC(model)._pipeline
basis = model.get_basis_instructions()
basis_key = basis_to_hash(basis)
else:
basis_key = ''
data = load_sets(datafile, hdf5[1], hdf5[2], basis_key, cutoff)
results = {}
if not model == '':
symmetrizer_instructions = model.get_symmetrize_instructions()
symmetrizer_instructions.update({'basis': basis})
species = [''.join(find_attr_in_tree(datafile, hdf5[1], 'species'))]
spec_group = SpeciesGrouper(basis, species)
symmetrizer = symmetrizer_factory(symmetrizer_instructions)
print('Symmetrizer instructions', symmetrizer_instructions)
pipeline = NXCPipeline(
[('spec_group', spec_group),
('symmetrizer', symmetrizer)] + model.steps,
basis_instructions=basis,
symmetrize_instructions=symmetrizer_instructions)
targets = data[:, -1].real
predictions = pipeline.predict(data)[0]
if predict:
np.save(dest, predictions)
return 0
dev = (predictions.flatten() - targets.flatten())
else:
if predict:
raise Exception('Must provide a model to make predictions')
dev = data[:, -1].real
# predictions = load_sets(datafile, hdf5[1], hdf5[1], basis_key, cutoff)[:,-1].flatten()
# targets = load_sets(datafile, hdf5[2], hdf5[2], basis_key, cutoff)[:,-1].flatten()
predictions = datafile[hdf5[1] + '/energy'][:]
targets = datafile[hdf5[2] + '/energy'][:]
try:
force_base = datafile[hdf5[1] + '/forces'][:]
force_ref = datafile[hdf5[2] + '/forces'][:]
force_results = {
'force_mae': np.mean(np.abs(force_ref - force_base)),
'force_std': np.std(force_ref - force_base),
'force_max': np.max(force_ref - force_base)
}
results.update(force_results)
except Exception:
pass
dev0 = np.abs(dev - np.mean(dev))
results.update({
'mean deviation': np.mean(dev).round(4),
'rmse': np.std(dev).round(4),
'mae': np.mean(dev0).round(4),
'max': np.max(dev0).round(4)
})
pprint(results)
if plot:
if model == '':
plt.figure(figsize=(10, 8))
plt.subplot(2, 1, 1)
plt.hist(dev.flatten())
plt.xlabel('Target energies [eV]')
plt.subplot(2, 1, 2)
targets -= np.mean(targets)
predictions -= np.mean(predictions)
maxlim = np.max([np.max(targets), np.max(predictions)])
minlim = np.min([np.max(targets), np.min(predictions)])
plt.plot(targets.flatten(), predictions.flatten(), ls='', marker='.')
plt.plot([minlim, maxlim], [minlim, maxlim],
ls='-',
marker='',
color='grey')
plt.xlabel('$E_{ref}[eV]$')
plt.ylabel('$E_{pred}[eV]$')
plt.show()
return results
def test_force_correction(use_delta):
benzene_traj = ase.io.read(os.path.join(test_dir, 'benzene_test', 'benzene.xyz'), '0')
density_getter = xc.utils.SiestaDensityGetter(binary=True)
rho, unitcell, grid = density_getter.get_density(os.path.join(test_dir, 'benzene_test', 'benzene.DRHO'))
positions = benzene_traj.get_positions() / Bohr
species = benzene_traj.get_chemical_symbols()
if use_delta:
drho, _, _ = density_getter.get_density(os.path.join(test_dir, 'benzene_test', 'benzene.DRHO'))
benzene_nxc = xc.NeuralXC(os.path.join(test_dir, 'benzene_test', 'dbenzene'))
benzene_nxc.initialize(unitcell, grid, positions, species)
benzene_nxc.projector = xc.projector.DeltaProjector(benzene_nxc.projector)
benzene_nxc.projector.set_constant_density(drho, positions, species)
else:
benzene_nxc = xc.NeuralXC(os.path.join(test_dir, 'benzene_test', 'benzene'))
benzene_nxc.initialize(unitcell, grid, positions, species)
def get_V_shifted(self, rho, unitcell, grid, positions, positions_shifted, species, calc_forces=False):
""" only defined to calculate basis set contribution to forces with numerical derivatives.
For finite difference, the descriptors should be calculated using the original positions,
whereas V will then be built with displaced atoms
"""
projector = xc.projector.DensityProjector(unitcell, grid, self._pipeline.get_basis_instructions())
if use_delta:
projector = xc.projector.DeltaProjector(projector)
projector.set_constant_density(drho, positions, species)
symmetrize_dict = {'basis': self._pipeline.get_basis_instructions()}
symmetrize_dict.update(self._pipeline.get_symmetrize_instructions())
symmetrizer = xc.symmetrizer.symmetrizer_factory(symmetrize_dict)
C = projector.get_basis_rep(rho, positions, species)
D = symmetrizer.get_symmetrized(C)
dEdC = symmetrizer.get_gradient(self._pipeline.get_gradient(D))
E = self._pipeline.predict(D)[0]
return E, projector.get_V(dEdC, positions_shifted, species, calc_forces, rho)
V, forces = benzene_nxc.get_V(rho, calc_forces=True)[1]
forces = forces[:-3] # no stress
assert np.allclose(np.sum(forces, axis=0), np.zeros(3), atol=1e-6)
for incr_atom in [0, 1, 2, 3]:
for incr_dx in range(3):
incr = 0.00001
incr_idx = 1
pp = np.array(positions)
pm = np.array(pp)
pp[incr_atom, incr_idx] += incr
pm[incr_atom, incr_idx] -= incr
Vp = get_V_shifted(benzene_nxc, rho, unitcell, grid, positions, pp, species)[1]
Vm = get_V_shifted(benzene_nxc, rho, unitcell, grid, positions, pm, species)[1]
dv = (unitcell[0, 0] / grid[0])**3
fp = dv * np.sum(Vp * rho)
fm = dv * np.sum(Vm * rho)
forces_fd = (fp - fm) / (2 * incr)
assert np.allclose(-forces_fd, forces[incr_atom, incr_idx], atol=incr)