import numpy as np import time from multivariate_gpr_cv import MultivariateGaussianProcessCV from sklearn.metrics import mean_squared_error, mean_absolute_error from sklearn.metrics.pairwise import euclidean_distances import os.path from sacred import Experiment import dft_utils as du from sacred.commandline_options import CommandLineOption from sacred.observers import MongoObserver import random class MongoAdvancedOption(CommandLineOption): """Enables a Mongo database to be specified with added authentication options.""" short_flag = 'M' arg = 'DATABASE' arg_description = 'A string containing the database host+authentication, database name and collection name separated by ";"' @classmethod def apply(cls, args, run): # run.config contains the configuration. You can read from there. split_args = args.split(';') print(split_args) mongo = MongoObserver.create(url=split_args[0], db_name=split_args[1], collection=None if split_args[2] == '' else split_args[2]) run.observers.append(mongo) ex = Experiment('Machine learning DFT') @ex.config def config(): """ Config of default experiment parameters""" results_dir = '' train_dir = 'train' test_dir = 'test' train_inds = range(0,5000) test_inds = list(range(0, 240)) train_inds_file = None test_inds_file = None state_id = 0 output_file = None n_jobs = None energy_type = 'pbe0' descriptor_type = 'pot' gaussian_width = 0.36 grid_spacing = 0.20 grid_file = None density_kernel = 'rbf' energy_kernel = 'rbf' density_kernel_params = {} energy_kernel_params = {} # rbf kernel # we've searched a much larger range of alpha and gamma parameters, here for convenience and to save computational cost, we just provide a combination that does well for this model. # one can easily search alpha and gamma with extending these list density_alpha_params = [2.21221629e-10] density_gamma_params = [901.0484834870775, 146.31015683675886, 79.82089280705694] energy_alpha_params = [2.21221629e-10] energy_gamma_params = [49157.554040919036] use_true_densities = False verbose = 2 plot_cv_errors = False if verbose > 0: print('Parameters:') print('results_dir:', results_dir) print('train_dir:', train_dir) print('test_dir:', test_dir) print('state_id:', state_id) print('output_file', output_file) print('n_jobs', n_jobs) print('energy_type', energy_type) print('descriptor_type', descriptor_type) print('grid_spacing:', grid_spacing) print('gaussian_width:', gaussian_width) print('grid_file:', grid_file) print('density_kernel:', density_kernel) print('energy_kernel:', energy_kernel) print('train_inds_file:', train_inds_file) print('test_inds_file:', test_inds_file) print('density_kernel_params:', density_kernel_params) print('energy_kernel_params:', energy_kernel_params) print('density_alpha_params:', density_alpha_params) print('density_gamma_params:', density_gamma_params) print('energy_alpha_params:', energy_alpha_params) print('energy_gamma_params:', energy_gamma_params) print('plot_cv_errors:', plot_cv_errors) print('use_true_densities:', use_true_densities) print('verbose:', verbose) if verbose > 1: print('train_inds:', train_inds) print('test_inds:', test_inds) @ex.command def transform_positions(results_dir, train_dir): """Performs a transformation operation on the positions of the atoms in order to align them. Args: config: Dictionary containing the necessary configuration parameters. """ positions = np.load(os.path.join(results_dir, train_dir, 'pos.npy')) charges = np.loadtxt(os.path.join(results_dir, train_dir, 'a_num.txt')) base = np.load(os.path.join(results_dir, train_dir, 'base_pos.npy')) heavy = [] if len(charges) <= 3: heavy = np.ones(len(charges), dtype=np.bool_) else: heavy = np.where(charges > 1)[0] heavy = np.array(heavy, dtype=np.bool_) for i in range(positions.shape[0]): positions[i] = du.transform_molecule(positions[i], base, heavy) pos_new = positions np.save(os.path.join(results_dir, train_dir, 'trans_pos.npy'), pos_new) return pos_new @ex.command def calculate_descriptors(descriptor_type): if descriptor_type == 'pot': calculate_potentials() else: raise Exception('Descriptor unsupported!') @ex.command def calculate_potentials(results_dir, train_dir, test_dir, state_id, grid_spacing, grid_file, gaussian_width, verbose): """Calculates the potential as a sum of artificial Gaussians based on the atom positions as saves them on disk. Args: config: Dictionary containing the necessary configuration parameters. """ if verbose > 0: print('Calculating potentials') if train_dir == test_dir: work_dirs = [train_dir] else: work_dirs = [train_dir, test_dir] all_positions = [] for work_dir in work_dirs: positions = np.load(os.path.join(results_dir, work_dir, 'pos_' + work_dir + '.npy')) all_positions.append(np.copy(positions)) cat_positions = np.concatenate(all_positions, axis=0) if verbose > 0: print('Positions shape:', cat_positions.shape) max_pos = np.max(cat_positions, axis=0) max_pos = np.max(max_pos, axis=0) if verbose > 1: print(max_pos) min_pos = np.min(cat_positions, axis=0) min_pos = np.min(min_pos, axis=0) if verbose > 1: print(min_pos) max_pos = np.ceil(max_pos) + 1 min_pos = np.floor(min_pos) - 1 max_pos = np.array([16.,15.,13.]) min_pos = np.array([4.,5.,7.]) if verbose > 0: print('Min, max coordinates:', min_pos, max_pos) max_range = max_pos - min_pos steps = np.round(max_range / grid_spacing).astype(np.int) if (grid_file is None): grid_range = [np.linspace(ma, mm, s) for ma, mm, s in zip(max_pos, min_pos, steps)] elif(grid_file == ''): if verbose > 0: print('Saving potentials grid') grid_range = [np.linspace(ma, mm, s) for ma, mm, s in zip(max_pos, min_pos, steps)] np.save(os.path.join(results_dir, train_dir, 'grid_range_' + str(state_id) + '.npy'), grid_range) elif os.path.exists(os.path.join(results_dir, train_dir, grid_file + '_' + state_id + '.npy')): grid_range = np.load(os.path.join(results_dir, train_dir, grid_file + '_' + state_id + '.npy')) else: raise RuntimeError('Invalid grid file: \'' + grid_file + '\'') Y, X, Z = np.meshgrid(grid_range[0], grid_range[1], grid_range[2]) if verbose > 1: print(X.shape) X = X.flatten()[:, np.newaxis] Y = Y.flatten()[:, np.newaxis] Z = Z.flatten()[:, np.newaxis] grid = np.concatenate((X, Y, Z), axis=1) if verbose > 0: print('Potential positions shape', grid.shape) for i, work_dir in enumerate(work_dirs): charges = np.loadtxt(os.path.join(results_dir, work_dir, 'a_num.txt')) positions = all_positions[i] potentials = du.calculate_potential(positions, charges, gaussian_width, grid, verbose=verbose) if verbose > 1: print(potentials.shape) np.save(os.path.join(results_dir, work_dir, 'pot_' + '0' + '.npy'), potentials) @ex.command def descriptors_to_density(results_dir, train_dir, train_inds, train_inds_file, state_id, n_jobs, descriptor_type, density_alpha_params, density_gamma_params, density_kernel, density_kernel_params, plot_cv_errors, verbose): """Trains independent KRR models to predict the basis coefficients of the density from the descriptors Args: config: Dictionary containing the necessary configuration parameters. """ if verbose > 0: print('Training descriptors to density model') if train_inds_file is not None: train_inds = list(np.load(train_inds_file)) descriptors = np.load(os.path.join(results_dir, train_dir, descriptor_type + '_' + '0' + '.npy')) coefs = np.load(os.path.join(results_dir, train_dir, 'densities_s'+ str(state_id) + '.npy')) descriptors = descriptors[train_inds, :] coefs = coefs[train_inds, :] if verbose > 1: print(coefs.shape) dist = euclidean_distances(descriptors , None, squared=False) for gamma_i in range(1): density_kr = MultivariateGaussianProcessCV(cv_nfolds=5, krr_param_grid={"alpha": density_alpha_params, "gamma": density_gamma_params, }, id=state_id, verbose=verbose, n_jobs=n_jobs, cluster_params=['-l h_vmem=100G'], kernel=density_kernel, kernel_params=density_kernel_params) start = time.time() density_kr.fit(descriptors, coefs, dist=dist) end = time.time() test_dir = 'test' coefs_pred = density_kr.predict(descriptors) rmse = np.mean(np.linalg.norm(coefs - coefs_pred, axis=1)) for work_dir in [test_dir]: descriptors_t = np.load(os.path.join(results_dir, work_dir, descriptor_type + '_' + '0' + '.npy')) coefs_t = np.load(os.path.join(results_dir, work_dir, 'densities_s' + str(state_id) + 'test.npy')) coefs_pred_t = density_kr.predict(descriptors_t) print('RMSE:', np.mean(np.linalg.norm(coefs_t - coefs_pred_t, axis=1))) print('Coefs test norm:', np.mean(np.linalg.norm(coefs_t, axis=1))) print('Coef test predictions norm:', np.mean(np.linalg.norm(coefs_pred_t, axis=1))) print('Coefs norm:', np.mean(np.linalg.norm(coefs, axis=1))) print('Coef predictions norm:', np.mean(np.linalg.norm(coefs_pred, axis=1))) rmse_t=np.mean(np.linalg.norm(coefs_t - coefs_pred_t, axis=1)) density_kr.save(os.path.join(results_dir, train_dir, 'density_kr_' + str(state_id))) ex.info['density_kr_alpha'] = density_kr.alphas_[0] ex.info['density_kr_gamma'] = density_kr.gammas_[0] print('Alphas = ' + str(density_kr.alphas_)) print('Gammas = ' + str(density_kr.gammas_)) if verbose > 1: print('Elapsed train', end - start) if plot_cv_errors: density_kr.plot_cv_error() return np.min(density_kr.errors) @ex.command def predict_density(results_dir, train_dir, test_dir, state_id, n_jobs, descriptor_type, density_alpha_params, density_gamma_params, verbose): """Predicts the density coefficients using the learned KRR models Args: config: Dictionary containing the necessary configuration parameters. """ if verbose > 0: print('Predicting density coefficients') for work_dir in [train_dir, test_dir]: descriptors = np.load(os.path.join(results_dir, work_dir, descriptor_type + '_' + '0' + '.npy')) if work_dir=='train': coefs = np.load(os.path.join(results_dir, work_dir, 'densities_s' + str(state_id) + '.npy')) else: coefs = np.load(os.path.join(results_dir, work_dir, 'densities_s' + str(state_id) + work_dir + '.npy')) if verbose > 1: print(coefs.shape) density_kr = MultivariateGaussianProcessCV(cv_nfolds=5, krr_param_grid={"alpha": density_alpha_params, "gamma": density_gamma_params}, id=state_id, verbose=verbose, n_jobs=n_jobs, cluster_params=['-l h_vmem=100G']) density_kr.load(os.path.join(results_dir, train_dir, 'density_kr_' + str(state_id))) start = time.time() coefs_pred = density_kr.predict(descriptors) end = time.time() if verbose > 1: print('Elapsed predict', end - start) start = time.time() np.save(os.path.join(results_dir, work_dir, 'coefs_pred_' + str(state_id) + '.npy'), coefs_pred) end = time.time() if verbose > 1: print('Elapsed save', end - start) if verbose > 0: print('RMSE:', np.mean(np.linalg.norm(coefs - coefs_pred, axis=1))) print('Coefs norm:', np.mean(np.linalg.norm(coefs, axis=1))) print('Coef predictions norm:', np.mean(np.linalg.norm(coefs_pred, axis=1))) return np.mean(np.linalg.norm(coefs - coefs_pred, axis=1)) @ex.command def density_to_energy(results_dir, train_dir, energy_type, train_inds, train_inds_file, state_id, n_jobs, energy_alpha_params, energy_gamma_params, energy_kernel, energy_kernel_params, plot_cv_errors, use_true_densities, verbose): """Trains a KRR model that predicts the energy from the density coefficients Args: config: Dictionary containing the necessary configuration parameters. """ if verbose > 0: print('Training density to energy model') use_true_densities=True if use_true_densities: coefs_pred = np.append(np.load(os.path.join(results_dir, train_dir, 'densities_s0.npy')),np.load(os.path.join(results_dir, 'train', 'densities_s1.npy')),axis=0) coefs_pred = np.append(coefs_pred,np.load(os.path.join(results_dir, 'train', 'densities_s2.npy')),axis=0) else: coefs_pred = np.load(os.path.join(results_dir, train_dir, 'coefs_pred_' + str(state_id) + '.npy')) coefs_predicted = np.append(np.load(os.path.join(results_dir, train_dir, 'coefs_pred_' + str(state_id) + '.npy')),np.load(os.path.join(results_dir, 'train', 'coefs_pred_' + str(state_id) + '.npy')),axis=0) coefs_predicted = np.append(coefs_predicted ,np.load(os.path.join(results_dir, 'train', 'coefs_pred_' + str(state_id) + '.npy')),axis=0) if train_inds_file is not None: train_inds = list(np.load(train_inds_file)) energies = np.append(np.load(os.path.join(results_dir, train_dir, energy_type + '_s0.npy')),np.load(os.path.join(results_dir, train_dir, energy_type + '_s1.npy'))) energies = np.append(energies,np.load(os.path.join(results_dir, train_dir, energy_type + '_s2.npy'))) energies = np.reshape(energies, (-1, 1)) if verbose > 1: print(energies.shape) gamma_list=[] for gamma_i in [0]: energy_pot_kr = MultivariateGaussianProcessCV(cv_nfolds=5, krr_param_grid={"alpha": energy_alpha_params, "gamma": energy_gamma_params, }, id=state_id + .5, verbose=verbose, n_jobs=n_jobs, cluster_params=['-l h_vmem=100G'], kernel=energy_kernel, kernel_params=energy_kernel_params, delta_learning=(energy_type == 'diff')) energy_pot_kr.fit(coefs_pred, energies) energy_pred = energy_pot_kr.predict(coefs_pred) print('Energy train error:', mean_absolute_error(energies, energy_pred)) if mean_absolute_error(energies,energy_pred) < 10.0 : if mean_absolute_error(energies,energy_pred) > 0. : energy_pot_kr.save(os.path.join(results_dir, train_dir, str(energy_type) + str('_kr_') + '0')) if verbose > 1: print('Alphas = ' + str(energy_pot_kr.alphas_)) print('Gammas = ' + str(energy_pot_kr.gammas_)) energy_pred = energy_pot_kr.predict(coefs_predicted) return(mean_absolute_error(energies, energy_pred)) @ex.command def descriptors_to_energy(results_dir, train_dir, energy_type, train_inds, train_inds_file, state_id, n_jobs, energy_alpha_params, energy_gamma_params, energy_kernel, descriptor_type, energy_kernel_params, plot_cv_errors, verbose): """Trains a KRR model that predicts the energy from the density coefficients Args: config: Dictionary containing the necessary configuration parameters. """ if verbose > 0: print('Training descriptors to energy model') if train_inds_file is not None: train_inds = list(np.load(train_inds_file)) descriptors = np.load(os.path.join(results_dir, train_dir, descriptor_type + '_' + '0' + '.npy')) energies = np.load(os.path.join(results_dir, train_dir, energy_type + '.npy')) descriptors = descriptors[train_inds, :] energies = energies[train_inds] energies = np.reshape(energies, (-1, 1)) if verbose > 1: print(energies.shape) if verbose > 1: print(descriptors.shape) energy_pot_kr = MultivariateGaussianProcessCV(cv_nfolds=5, krr_param_grid={"alpha": energy_alpha_params, "gamma": energy_gamma_params, }, id=state_id + .5, verbose=verbose, n_jobs=n_jobs, cluster_params=['-l h_vmem=100G'], kernel=energy_kernel, kernel_params=energy_kernel_params, delta_learning=(energy_type == 'diff')) energy_pot_kr.fit(descriptors, energies) if plot_cv_errors: energy_pot_kr.plot_cv_error() start = time.time() energy_pot_kr.save(os.path.join(results_dir, train_dir, energy_type + '_pot_kr_' + str(state_id))) end = time.time() if verbose > 1: print('Elapsed save', end - start) energy_pred = energy_pot_kr.predict(descriptors) if verbose > 0: print('Energy train error:', mean_absolute_error(energies, energy_pred)) if verbose > 1: print('Alphas = ' + str(energy_pot_kr.alphas_)) print('Gammas = ' + str(energy_pot_kr.gammas_)) return mean_absolute_error(energies, energy_pred) @ex.command def test(results_dir, train_dir, test_dir, energy_type, test_inds, test_inds_file, state_id, output_file, energy_alpha_params, energy_gamma_params, use_true_densities, _config, verbose): """Evaluates the model using the data in test for s2 config: Dictionary containing the necessary configuration parameters. """ if verbose > 0: print('Testing the model') if test_inds_file is not None: test_inds = list(np.load(test_inds_file)) energies = np.load(os.path.join(results_dir, 'test', energy_type + '_s' + str(state_id) + '_test.npy')) if use_true_densities: coefs_pred = np.load(os.path.join(results_dir, test_dir, 'densities_s' + str(state_id) + 'test.npy')) else: coefs_pred = np.load(os.path.join(results_dir, 'test', 'coefs_pred_' + str(state_id) + '.npy')) energies = np.reshape(energies, (-1, 1)) for gamma_i in [0]: energy_kr = MultivariateGaussianProcessCV(cv_nfolds=5, krr_param_grid={"alpha": energy_alpha_params, "gamma": energy_gamma_params, }, id=state_id + .5, verbose=verbose, cluster_params=['-l h_vmem=100G']) energy_kr.load(os.path.join(results_dir, train_dir, energy_type + '_kr_' + '0')) energies_pred = energy_kr.predict(coefs_pred) energies_pred = np.reshape(energies_pred, (-1, 1)) errors_pred = np.abs(energies - energies_pred) np.savetxt(os.path.join(results_dir, train_dir, energy_type + '_kr_' + str(state_id)+'_pred.txt'),energies_pred) ex.info['corr'] = np.corrcoef(energies.T, energies_pred.T)[0][1] ex.info['errors'] = errors_pred ex.info['preds'] = energies_pred if output_file is None: mae=mean_absolute_error(energies, energies_pred) print('Energies') print('Correlation: ', np.corrcoef(energies.T, energies_pred.T)[0][1]) print('RMSE: ', np.sqrt(mean_squared_error(energies, energies_pred))) print('MAE: ', mean_absolute_error(energies, energies_pred)) print('Max: ', np.max(np.abs(energies - energies_pred))) else: if verbose > 0: print('Writing to file:', os.path.join(results_dir, test_dir, output_file + '_' + energy_type + '_' + str(state_id) + '.npy')) f = open(os.path.join(results_dir, test_dir, output_file + '_' + energy_type + '_' + str(state_id) + '.npy'), 'w') f.write('Config:\n') f.write(str(_config)) f.write('\n') f.write('Training sample inds:' + str(test_inds) + '\n') f.write('Energies\n') corr = np.corrcoef(energies.T, energies_pred.T)[0][1] if np.isnan(corr): corr = 0 if verbose > 0: print('Correlation is ', corr) f.write('Correlation: ' + str(corr) + '\n') f.write('RMSE: ' + str(np.sqrt(mean_squared_error(energies, energies_pred))) + '\n') f.write('MAE: ' + str(mean_absolute_error(energies, energies_pred)) + '\n') f.write('Max: ' + str(np.max(np.abs(energies - energies_pred))) + '\n') np.save(os.path.join(results_dir, test_dir, 'errors_pred_' + energy_type + '_' + str(state_id) + '.npy'), errors_pred) return mean_absolute_error(energies, energies_pred) @ex.command def train_full(): """Performs the entire training procedure to predict the energies, including generating the desriptors config: Dictionary containing the necessary configuration parameters. """ print('Starting full training process') calculate_potentials(state_id=0) dens_cv_err_s0 = descriptors_to_density(state_id=0) dens_err_s0 = predict_density(state_id=0) calculate_potentials(state_id=1) dens_cv_err_s1 = descriptors_to_density(state_id=1) dens_err_s1 = predict_density(state_id=1) calculate_potentials(state_id=2) dens_cv_err_s2 = descriptors_to_density(state_id=2) dens_err_s2 = predict_density(state_id=2) en_err = density_to_energy() ex.info['energy_train_err'] = en_err @ex.command def run_density(): """Performs the descriptor to density prediction procedure config: Dictionary containing the necessary configuration parameters. """ calculate_potentials(state_id=0) dens_cv_err_s0 = descriptors_to_density(state_id=0) dens_err_s0 = predict_density(state_id=0) calculate_potentials(state_id=1) dens_cv_err_s1 = descriptors_to_density(state_id=1) dens_err_s1 = predict_density(state_id=1) calculate_potentials(state_id=2) dens_cv_err_s2 = descriptors_to_density(state_id=2) dens_err_s2 = predict_density(state_id=2) ex.info['test_err_s0'] = dens_err_s0 ex.info['test_err_s1'] = dens_err_s1 ex.info['test_err_s2'] = dens_err_s2 return dens_err_s2 @ex.command def run_energy(): """Performs the density to energy prediction procedure config: Dictionary containing the necessary configuration parameters. """ density_to_energy() mae = test(state_id=2) return mae @ex.automain def run_train_test(): """Performs trains and tests the ML_MSHK model, assuming the descriptors have already been created config: Dictionary containing the necessary configuration parameters. """ train_full() mae = test(state_id=2) return mae