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)