From 7d311766d4b528c58ce398be703d55e9cb21d579 Mon Sep 17 00:00:00 2001 From: "oz.mendelsohn" Date: Tue, 9 Feb 2021 17:27:06 +0200 Subject: [PATCH 01/14] Added support to windows --- sgdml/__init__.py | 2 +- sgdml/predict.py | 5 ++++- sgdml/torchtools.py | 4 ++-- sgdml/train.py | 5 ++++- sgdml/utils/desc.py | 6 ++++-- sgdml/utils/perm.py | 6 ++++-- 6 files changed, 19 insertions(+), 9 deletions(-) diff --git a/sgdml/__init__.py b/sgdml/__init__.py index fcda3e5..7c2075c 100755 --- a/sgdml/__init__.py +++ b/sgdml/__init__.py @@ -22,7 +22,7 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. -__version__ = '0.4.10' +__version__ = '0.4.10-oz' MAX_PRINT_WIDTH = 100 LOG_LEVELNAME_WIDTH = 7 # do not modify diff --git a/sgdml/predict.py b/sgdml/predict.py index 40f39b7..066c685 100644 --- a/sgdml/predict.py +++ b/sgdml/predict.py @@ -32,7 +32,10 @@ import multiprocessing as mp -Pool = mp.get_context('fork').Pool +if sys.platform == 'win32': + from multiprocessing.pool import ThreadPool as Pool +else: + Pool = mp.get_context('fork').Pool import timeit from functools import partial diff --git a/sgdml/torchtools.py b/sgdml/torchtools.py index 2bca694..a566f64 100644 --- a/sgdml/torchtools.py +++ b/sgdml/torchtools.py @@ -87,7 +87,7 @@ def __init__(self, model, lat_and_inv=None, batch_size=None, max_memory=None): self._xs_train, self._Jx_alphas = ( nn.Parameter( - xs.repeat(1, n_perms)[:, perm_idxs].reshape(-1, desc_siz), + xs.repeat(1, n_perms)[:, perm_idxs.type(torch.LongTensor)].reshape(-1, desc_siz), requires_grad=False, ) for xs in ( @@ -146,7 +146,7 @@ def set_alphas(self, R_d_desc, alphas): xs = torch.from_numpy(R_d_desc_alpha).to(self._dev) self._Jx_alphas = nn.Parameter( - xs.repeat(1, self.n_perms)[:, self.perm_idxs].reshape(-1, dim_d), + xs.repeat(1, self.n_perms)[:, self.perm_idxs.type(torch.LongTensor)].reshape(-1, dim_d), requires_grad=False, ) diff --git a/sgdml/train.py b/sgdml/train.py index 97608ef..aedc6ad 100755 --- a/sgdml/train.py +++ b/sgdml/train.py @@ -31,7 +31,10 @@ import multiprocessing as mp -Pool = mp.get_context('fork').Pool +if sys.platform == 'win32': + from multiprocessing.pool import ThreadPool as Pool +else: + Pool = mp.get_context('fork').Pool import timeit from functools import partial diff --git a/sgdml/utils/desc.py b/sgdml/utils/desc.py index 5f0e716..6bd0d4e 100755 --- a/sgdml/utils/desc.py +++ b/sgdml/utils/desc.py @@ -26,8 +26,10 @@ import scipy as sp import multiprocessing as mp - -Pool = mp.get_context('fork').Pool +if sys.platform == 'win32': + from multiprocessing.pool import ThreadPool as Pool +else: + Pool = mp.get_context('fork').Pool from functools import partial from scipy import spatial diff --git a/sgdml/utils/perm.py b/sgdml/utils/perm.py index 5f173a5..12e6d9c 100755 --- a/sgdml/utils/perm.py +++ b/sgdml/utils/perm.py @@ -25,8 +25,10 @@ from __future__ import print_function import multiprocessing as mp - -Pool = mp.get_context('fork').Pool +if sys.platform == 'win32': + from multiprocessing.pool import ThreadPool as Pool +else: + Pool = mp.get_context('fork').Pool import sys import timeit From d9bd6b46d694802205a827dd6cc49164a5b7c360 Mon Sep 17 00:00:00 2001 From: "oz.mendelsohn" Date: Tue, 9 Feb 2021 17:28:01 +0200 Subject: [PATCH 02/14] Added python api for model. model.train(), model.test() --- sgdml/core.py | 1161 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 1161 insertions(+) create mode 100644 sgdml/core.py diff --git a/sgdml/core.py b/sgdml/core.py new file mode 100644 index 0000000..abfd6eb --- /dev/null +++ b/sgdml/core.py @@ -0,0 +1,1161 @@ +try: + from ase.io import read +except ImportError: + raise ImportError('Optional ASE dependency not found! Please run \'pip install sgdml[ase]\' to install it.') +import numpy as np +import os +import sgdml +from functools import partial +from . import MAX_PRINT_WIDTH, DONE, NOT_DONE +from .utils import ui, io, perm, desc +from .cli import _batch, _online_err +from .train import GDMLTrain +from .predict import GDMLPredict +import time + + +def dummy_callback(done=None, **kwargs): + pass + + +class AssistantError(Exception): + pass + + +class Dataset: + def __init__(self, dataset_path: str, to_file=False, + name=None, + theory='unknown', + overwrite=False, + r_unit='', e_unit='', + verbose=True): + self.type = 'd' + self.code_version = sgdml.__version__ + self.name = alias_str(name if not name is None else os.path.splitext(os.path.basename(dataset_path))[0]) + self.theory = alias_str(theory) + + mols = read(dataset_path, index=':') + dataset_file_name = self.name + '.npz' + dataset_exists = os.path.isfile(dataset_file_name) + if dataset_exists and overwrite and verbose: + print(ui.color_str('[INFO]', bold=True) + ' Overwriting existing dataset file.') + if not dataset_exists or overwrite and verbose: + print('Writing dataset to \'{}\'...'.format(dataset_file_name)) + else: + if verbose: + print(ui.color_str('[FAIL]', fore_color=ui.RED, bold=True) + + ' Dataset \'{}\' already exists.'.format(dataset_file_name)) + return + + # filter incomplete outputs from trajectory + mols = [mol for mol in mols if mol.get_calculator() is not None] + + lattice, R, z, E, F = None, None, None, None, None + + calc = mols[0].get_calculator() + if verbose: + print("\rNumber geometries: {:,}".format(len(mols))) + print("\rAvailable properties: " + ', '.join(calc.results)) + print() + + if 'forces' not in calc.results: + if verbose: + print(ui.color_str('[FAIL]', fore_color=ui.RED, bold=True) + ' Forces are missing in the input file!') + return + + lattice = np.array(mols[0].get_cell()) + if not np.any(lattice): + if verbose: + print(ui.color_str('[INFO]', bold=True) + ' No lattice vectors specified.') + + Z = np.array([mol.get_atomic_numbers() for mol in mols]) + all_z_the_same = (Z == Z[0]).all() + if not all_z_the_same: + if verbose: + print(ui.color_str('[FAIL]', fore_color=ui.RED, bold=True) + ' Order of atoms changes accross dataset.') + return + + lattice = np.array(mols[0].get_cell()) + if not np.any(lattice): # all zeros + lattice = None + + R = np.array([mol.get_positions() for mol in mols]) + z = Z[0] + + E = np.array([mol.get_potential_energy() for mol in mols]) + F = np.array([mol.get_forces() for mol in mols]) + + self.F_min, self.F_max = np.min(F.ravel()), np.max(F.ravel()) + self.F_mean, self.F_var = np.mean(F.ravel()), np.var(F.ravel()) + if r_unit != '': + self.r_unit = r_unit + if e_unit != '': + self.e_unit = e_unit + + if E is not None: + self.E = E + self.E_min, self.E_max = np.min(E), np.max(E) + self.E_mean, self.E_var = np.mean(E), np.var(E) + else: + if verbose: + print(ui.color_str('[INFO]', bold=True) + ' No energy labels found in dataset.') + self.z = z + self.R = R + self.F = F + if lattice is not None: + self.lattice = lattice + + self.md5 = io.dataset_md5(self) + if to_file: + self.dataset_file_name = dataset_file_name + np.save(dataset_file_name, mols) + + def __setitem__(self, key, item): + self.__dict__[key] = item + + def __getitem__(self, key): + return self.__dict__[key] + + def __repr__(self): + return repr(self.__dict__) + + def __len__(self): + return len(self.__dict__) + + def __delitem__(self, key): + del self.__dict__[key] + + def clear(self): + return self.__dict__.clear() + + def copy(self): + return self.__dict__.copy() + + def has_key(self, k): + return k in self.__dict__ + + def update(self, *args, **kwargs): + return self.__dict__.update(*args, **kwargs) + + def keys(self): + return self.__dict__.keys() + + def values(self): + return self.__dict__.values() + + def items(self): + return self.__dict__.items() + + def pop(self, *args): + return self.__dict__.pop(*args) + + def __cmp__(self, dict_): + return self.__cmp__(self.__dict__, dict_) + + def __contains__(self, item): + return item in self.__dict__ + + def __iter__(self): + return iter(self.__dict__) + + def __str__(self, title_str='Dataset properties'): + text = '' + text += ui.white_bold_str(title_str) + '\n' + + n_mols, n_atoms, _ = self['R'].shape + text += ' {:<18} {} ({:i', f_pred_norm, f_norm)) / np.pi + cos_mae, cos_mae_sum, cos_rmse, cos_rmse_sum = _online_err(cos_err, n_atoms, n_done, cos_mae_sum, + cos_rmse_sum) + + sps = n_done / (time.time() - t) # examples per second + disp_str = 'energy %.3f/%.3f, ' % (e_mae, e_rmse) if self.use_E else '' + disp_str += 'forces %.3f/%.3f' % (f_mae, f_rmse) + disp_str = ('{} errors (MAE/RMSE): '.format('Test' if is_test else 'Validation') + + disp_str) + sec_disp_str = '@ %.1f geo/s' % sps if b_range is not None else '' + + self.callback(n_done, + len(test_idxs), + disp_str=disp_str, + sec_disp_str=sec_disp_str, + newline_when_done=False, + ) + if not self.use_E: + e_mae, e_rmse = 0, 0 + try: + self.callback(e_mae=e_mae, e_rmse=e_rmse, + f_mae=f_mae, f_rmse=f_rmse, + mag_mae=mag_mae, mag_rmse=mag_rmse, + cos_mae=cos_mae, cos_rmse=cos_rmse, iter=i) + except: + pass + + if is_test: + self.callback(DONE, disp_str='Testing on {:,} points'.format(n_test), + sec_disp_str=sec_disp_str, + ) + else: + self.callback(DONE, disp_str=disp_str, sec_disp_str=sec_disp_str) + + if self.use_E: + e_rmse_pct = (e_rmse / e_err['rmse'] - 1.0) * 100 + f_rmse_pct = (f_rmse / f_err['rmse'] - 1.0) * 100 + + # if func_called_directly and n_models == 1: + if is_test: + print(ui.white_bold_str('\nTest errors (MAE/RMSE)')) + + r_unit = 'unknown unit' + e_unit = 'unknown unit' + f_unit = 'unknown unit' + if 'r_unit' in dataset and 'e_unit' in dataset: + r_unit = dataset['r_unit'] + e_unit = dataset['e_unit'] + f_unit = str(dataset['e_unit']) + '/' + str(dataset['r_unit']) + + format_str = ' {:<18} {:>.4f}/{:>.4f} [{}]' + if self.use_E: + ui.print_two_column_str(format_str.format('Energy:', e_mae, e_rmse, e_unit), + 'relative to expected: {:+.1f}%'.format(e_rmse_pct), ) + + ui.print_two_column_str(format_str.format('Forces:', f_mae, f_rmse, f_unit), + 'relative to expected: {:+.1f}%'.format(f_rmse_pct), ) + + print(format_str.format(' Magnitude:', mag_mae, mag_rmse, r_unit)) + print(format_str.format(' Angle:', cos_mae, cos_rmse, '0-1, lower is better')) + + model_needs_update = (overwrite + or (is_test and self.n_test < len(test_idxs)) + or (is_validation and not is_model_validated) + ) + if model_needs_update: + if is_validation and overwrite: + self.n_test = 0 # flag the model as not tested + + if is_test: + self.n_test = len(test_idxs) + self.md5_test = dataset.md5 + + if self.use_E: + self.e_err = {'mae': e_mae, 'rmse': e_rmse, } + + self.f_err = {'mae': f_mae, 'rmse': f_rmse} + + if is_test and self.n_test > 0: + print('Expected errors were updated in model file.') + + else: + add_info_str = ( + 'the same number of' if self.n_test == len(test_idxs) else 'only {:,}'.format(len(test_idxs))) + print('This model has previously been tested on {:,} points, ' + 'which is why the errors for the current test run with {} points have ' + 'NOT been used to update the model file.\n'.format(self.n_test, add_info_str) + ) + F_rmse.append(f_rmse) + + def __str__(self): + text = ui.white_bold_str('Model properties') + '\n' + + text += ' {:<18} {} \n'.format('Dataset:', self.dataset_name) + + n_atoms = len(self.z) + text += ' {:<18} {: 1000: + R_train_sync_mat = self.R_train[np.random.choice(n_train, 1000, replace=False), :, :] + print('Symmetry search has been restricted to a random subset of 1000/{:d} ' + 'training points for faster convergence.'.format(n_train)) + + # TOOD: PBCs disabled when matching (for now). + # task['perms'] = perm.find_perms( + # R_train_sync_mat, train_dataset['z'], lat_and_inv=lat_and_inv, max_processes=self._max_processes, + # ) + self.perms = perm.find_perms(R_train_sync_mat, + train_dataset['z'], + lat_and_inv=None, + callback=self.callback, + max_processes=self.max_processes) + + # NEW + + USE_FRAG_PERMS = False + + if USE_FRAG_PERMS: + frag_perms = perm.find_frag_perms(R_train_sync_mat, + train_dataset['z'], + lat_and_inv=None, + max_processes=self.max_processes) + self.perms = np.vstack((self.perms, frag_perms)) + self.perms = np.unique(self.perms, axis=0) + + print('| Keeping ' + str(self.perms.shape[0]) + ' unique permutations.') + + # NEW + + else: + self.perms = np.arange(train_dataset.R.shape[1])[None, :] # no symmetries + else: + self.perms = self.model0['perms'] # TODO: change to obj syntax + print('Reusing permutations from initial model.') + + if self.model0 is not None: + + n_train, n_atoms = self.R_train.shape[:2] + + if 'alphas_F' in self.model0: + print('Reusing alphas from initial model.') + + # Pad existing alphas, if this training dataset is larger than the one in self.model0 + alphas0_F_padding = np.ones(((n_train - m0_n_train) * n_atoms * 3,)) * np.mean( + self.model0['alphas_F']) # TODO: update to obj syntax + self.alphas0_F = np.append(self.model0['alphas_F'], alphas0_F_padding) # TODO: update to obj syntax + + if 'alphas_E' in self.model0: + # Pad existing alphas, if this training dataset is larger than the one in model0 + alphas0_E_padding = np.ones(((n_train - m0_n_train) * n_atoms,)) * np.mean(self.model0['alphas_E']) + self.alphas0_E = np.append(self.model0['alphas_E'], alphas0_E_padding) # TODO: update to obj syntax + # Which atoms can we keep, if we exclude all symmetric ones? + n_perms = self.perms.shape[0] + if self.use_cprsn and n_perms > 1: + _, cprsn_keep_idxs = np.unique(np.sort(self.perms, axis=0), axis=1, return_index=True) + self.cprsn_keep_atoms_idxs = cprsn_keep_idxs + + def update_model_valid_index(self, n_valid, valid_dataset): + md5_valid = io.dataset_md5(valid_dataset) + if self.model0 is not None: + m0_idxs_valid = self.model0['idxs_valid'] # TODO: Change the dict syntax + m0_n_valid = m0_idxs_valid.shape[0] + else: + m0_n_valid = 0 + + excl_idxs = (self.idxs_train if self.md5_train == self.md5_valid + else np.array([], dtype=np.uint)) # TODO: TEST CASE: differnt test and val sets and m0 + excl_idxs = np.concatenate((self.m0_excl_idxs, excl_idxs)).astype(np.uint) + if 'E' in valid_dataset: + idxs_valid = draw_strat_sample(valid_dataset['E'], n_valid - m0_n_valid, excl_idxs) + else: + idxs_valid_all = np.setdiff1d(np.arange(valid_dataset['F'].shape[0]), excl_idxs, assume_unique=True) + idxs_valid = np.random.choice(idxs_valid_all, n_valid - m0_n_valid, replace=False) + # TODO: m0 handling, zero handling + self.idxs_valid = idxs_valid + self.md5_valid = md5_valid + + def update_hyperparameters(self, train_dataset=None, valid_dataset=None, **kwargs): + for k, v in kwargs.items(): + self[k] = v + if train_dataset is not None: + self.update_model_index(self.n_train, train_dataset) + if self.use_E: + self.E_train = train_dataset.E[self.idxs_train] + + if valid_dataset is not None: + self.update_model_valid_index(self.n_valid, valid_dataset) + + def __setitem__(self, key, item): + self.__dict__[key] = item + + def __getitem__(self, key): + return self.__dict__[key] + + def __repr__(self): + return repr(self.__dict__) + + def __len__(self): + return len(self.__dict__) + + def __delitem__(self, key): + del self.__dict__[key] + + def clear(self): + return self.__dict__.clear() + + def copy(self): + return self.__dict__.copy() + + def has_key(self, k): + return k in self.__dict__ + + def update(self, *args, **kwargs): + return self.__dict__.update(*args, **kwargs) + + def keys(self): + return self.__dict__.keys() + + def values(self): + return self.__dict__.values() + + def items(self): + return self.__dict__.items() + + def pop(self, *args): + return self.__dict__.pop(*args) + + def __cmp__(self, dict_): + return self.__cmp__(self.__dict__, dict_) + + def __contains__(self, item): + return item in self.__dict__ + + def __iter__(self): + return iter(self.__dict__) + + +class alias_str(str): + def __init__(self, strg): + super().__init__() + self.str = strg + + def astype(self, astype): + return astype(self.str) + + def __str__(self): + return self.str + + +def draw_strat_sample(T, n, excl_idxs=None): + """ + Draw sample from dataset that preserves its original distribution. + + The distribution is estimated from a histogram were the bin size is + determined using the Freedman-Diaconis rule. This rule is designed to + minimize the difference between the area under the empirical + probability distribution and the area under the theoretical + probability distribution. A reduced histogram is then constructed by + sampling uniformly in each bin. It is intended to populate all bins + with at least one sample in the reduced histogram, even for small + training sizes. + + Parameters + ---------- + T : :obj:`numpy.ndarray` + Dataset to sample from. + n : int + Number of examples. + excl_idxs : :obj:`numpy.ndarray`, optional + Array of indices to exclude from sample. + + Returns + ------- + :obj:`numpy.ndarray` + Array of indices that form the sample. + """ + + if len(excl_idxs) == 0: + excl_idxs = None + + if n == 0: + return np.array([], dtype=np.uint) + + if T.size == n: # TODO: this only works if excl_idxs=None + assert excl_idxs is None + return np.arange(n) + + if n == 1: + idxs_all_non_excl = np.setdiff1d( + np.arange(T.size), excl_idxs, assume_unique=True + ) + return np.array([np.random.choice(idxs_all_non_excl)]) + + # Freedman-Diaconis rule + h = 2 * np.subtract(*np.percentile(T, [75, 25])) / np.cbrt(n) + n_bins = int(np.ceil((np.max(T) - np.min(T)) / h)) if h > 0 else 1 + n_bins = min( + n_bins, int(n / 2) + ) # Limit number of bins to half of requested subset size. + + bins = np.linspace(np.min(T), np.max(T), n_bins, endpoint=False) + idxs = np.digitize(T, bins) + + # Exclude restricted indices. + if excl_idxs is not None and excl_idxs.size > 0: + idxs[excl_idxs] = n_bins + 1 # Impossible bin. + + uniq_all, cnts_all = np.unique(idxs, return_counts=True) + + # Remove restricted bin. + if excl_idxs is not None and excl_idxs.size > 0: + excl_bin_idx = np.where(uniq_all == n_bins + 1) + cnts_all = np.delete(cnts_all, excl_bin_idx) + uniq_all = np.delete(uniq_all, excl_bin_idx) + + # Compute reduced bin counts. + reduced_cnts = np.ceil(cnts_all / np.sum(cnts_all, dtype=float) * n).astype(int) + reduced_cnts = np.minimum( + reduced_cnts, cnts_all + ) # limit reduced_cnts to what is available in cnts_all + + # Reduce/increase bin counts to desired total number of points. + reduced_cnts_delta = n - np.sum(reduced_cnts) + + while np.abs(reduced_cnts_delta) > 0: + # How many members can we remove from an arbitrary bucket, without any bucket with more than one member going to zero? + max_bin_reduction = np.min(reduced_cnts[np.where(reduced_cnts > 1)]) - 1 + + # Generate additional bin members to fill up/drain bucket counts of subset. This array contains (repeated) bucket IDs. + outstanding = np.random.choice( + uniq_all, + min(max_bin_reduction, np.abs(reduced_cnts_delta)), + p=(reduced_cnts - 1) / np.sum(reduced_cnts - 1, dtype=float), + replace=True, + ) + uniq_outstanding, cnts_outstanding = np.unique( + outstanding, return_counts=True + ) # Aggregate bucket IDs. + + outstanding_bucket_idx = np.where( + np.in1d(uniq_all, uniq_outstanding, assume_unique=True) + )[ + 0 + ] # Bucket IDs to Idxs. + reduced_cnts[outstanding_bucket_idx] += ( + np.sign(reduced_cnts_delta) * cnts_outstanding + ) + reduced_cnts_delta = n - np.sum(reduced_cnts) + + # Draw examples for each bin. + idxs_train = np.empty((0,), dtype=int) + for uniq_idx, bin_cnt in zip(uniq_all, reduced_cnts): + idx_in_bin_all = np.where(idxs.ravel() == uniq_idx)[0] + idxs_train = np.append( + idxs_train, np.random.choice(idx_in_bin_all, bin_cnt, replace=False) + ) + + return idxs_train + + +def select(validated_models, overwrite=False, max_processes=None, model_file=None, command=None, + **kwargs): # noqa: C901 + any_model_not_validated = False + any_model_is_tested = False + + if len(validated_models) > 1: + use_E = True + + rows = [] + data_names = ['sig', 'MAE', 'RMSE', 'MAE', 'RMSE'] + for i, model in enumerate(validated_models): + use_E = model['use_E'] + + if i == 0: + idxs_train = set(model['idxs_train']) + md5_train = model['md5_train'] + idxs_valid = set(model['idxs_valid']) + md5_valid = model['md5_valid'] + else: + + if (md5_train != model['md5_train'] or md5_valid != model['md5_valid'] + or idxs_train != set(model['idxs_train']) or idxs_valid != set(model['idxs_valid'])): + raise AssistantError( + '{} contains models trained or validated on different datasets.'.format(1)) + + e_err = {'mae': 0.0, 'rmse': 0.0} + try: + if model['use_E']: + e_err = model['e_err'].item() + f_err = model['f_err'].item() + except: + if model['use_E']: + e_err = model['e_err'] + f_err = model['f_err'] + + is_model_validated = not (np.isnan(f_err['mae']) or np.isnan(f_err['rmse'])) + if not is_model_validated: + any_model_not_validated = True + + is_model_tested = model['n_test'] > 0 + if is_model_tested: + any_model_is_tested = True + + rows.append([model['sig'], e_err['mae'], e_err['rmse'], f_err['mae'], f_err['rmse']]) + + if any_model_not_validated: + print('One or more models in the given directory have not been validated yet.\n' + + 'This is required before selecting the best performer.') + print() + return + + if any_model_is_tested: + print( + 'One or more models in the given directory have already been tested. This means that their recorded expected errors are test errors, not validation errors. However, one should never perform model selection based on the test error!\n' + + 'Please run the validation command (again) with the overwrite option \'-o\', then this selection command.') + return + + f_rmse_col = [row[4] for row in rows] + best_idx = f_rmse_col.index(min(f_rmse_col)) # idx of row with lowest f_rmse + best_sig = rows[best_idx][0] + + rows = sorted(rows, key=lambda col: col[0]) # sort according to sigma + print(ui.white_bold_str('Cross-validation errors')) + print(' ' * 7 + 'Energy' + ' ' * 6 + 'Forces') + print((' {:>3} ' + '{:>5} ' * 4).format(*data_names)) + print(' ' + '-' * 27) + format_str = ' {:>3} ' + '{:5.2f} ' * 4 + format_str_no_E = ' {:>3} - - ' + '{:5.2f} ' * 2 + for row in rows: + if use_E: + row_str = format_str.format(*row) + else: + row_str = format_str_no_E.format(*[row[0], row[3], row[4]]) + + if row[0] != best_sig: + row_str = ui.gray_str(row_str) + print(row_str) + print() + + sig_col = [row[0] for row in rows] + if best_sig == min(sig_col) or best_sig == max(sig_col): + print('The optimal sigma lies on the boundary of the search grid.\n' + + 'Model performance might improve if the search grid is extended in direction sigma {} {:d}.'.format( + '<' if best_idx == 0 else '>', best_sig)) + + else: # only one model available + print('Skipping model selection step as there is only one model to select.') + + best_idx = 0 + best_model = validated_models[best_idx] + + return best_model + + +if __name__ == '__main__': + dataset_path = 'cspbbr3-300K-train-4676all.xyz' + dataset = Dataset(dataset_path, name='test_dataset', overwrite=True) + print(dataset['name']) + # for k in dataset.keys(): + # print(k) From 427ac2c599c0e554ce7e56d515b702dc46570b60 Mon Sep 17 00:00:00 2001 From: "oz.mendelsohn" Date: Thu, 11 Feb 2021 10:32:22 +0200 Subject: [PATCH 03/14] Added "import sys" for Windows support --- sgdml/utils/desc.py | 1 + sgdml/utils/perm.py | 1 + 2 files changed, 2 insertions(+) diff --git a/sgdml/utils/desc.py b/sgdml/utils/desc.py index 6bd0d4e..6614257 100755 --- a/sgdml/utils/desc.py +++ b/sgdml/utils/desc.py @@ -24,6 +24,7 @@ import numpy as np import scipy as sp +import sys import multiprocessing as mp if sys.platform == 'win32': diff --git a/sgdml/utils/perm.py b/sgdml/utils/perm.py index 12e6d9c..41efd5d 100755 --- a/sgdml/utils/perm.py +++ b/sgdml/utils/perm.py @@ -23,6 +23,7 @@ # SOFTWARE. from __future__ import print_function +import sys import multiprocessing as mp if sys.platform == 'win32': From c1a26f0003b30dc19359561d03575d394861c50a Mon Sep 17 00:00:00 2001 From: "oz.mendelsohn" Date: Thu, 11 Feb 2021 10:33:22 +0200 Subject: [PATCH 04/14] Update documentation and model.__str__ bug fix --- sgdml/core.py | 449 +++++++++++++++++++++++++++++++++++++++++++------- 1 file changed, 390 insertions(+), 59 deletions(-) diff --git a/sgdml/core.py b/sgdml/core.py index abfd6eb..1fea8ca 100644 --- a/sgdml/core.py +++ b/sgdml/core.py @@ -23,12 +23,52 @@ class AssistantError(Exception): class Dataset: + """ + Dataset object + + This object is a custom dataset holder for the sGDML. both object API and + directory API for compatibility with the CLI + """ def __init__(self, dataset_path: str, to_file=False, name=None, theory='unknown', overwrite=False, r_unit='', e_unit='', verbose=True): + + """ + Constructor for the Data object. + using ase.io.read as a general file format loader. + + Parameters + ---------- + dataset_path : str + Path to dataset file. + to_file : bool, optional + True: save Dataset object to file + False: only return the object + name : str, optional + A custom name to the dataset + theory : str, optional + The source method for the dataset + overwrite : bool, optional + True: rewrite existing dataset file if present + False: will abort if existing dataset file is present + r_unit : str, optional + Distance units of the Dataset + e_unit : str, optional + Energy units of the Dataset + verbose : bool, optional + True: will display the process information + False: will not display any informtion + + + Returns + ------- + sgdml.core.Dataset + type :obj:`sgdml.core.Task`. + """ + self.type = 'd' self.code_version = sgdml.__version__ self.name = alias_str(name if not name is None else os.path.splitext(os.path.basename(dataset_path))[0]) @@ -109,6 +149,8 @@ def __init__(self, dataset_path: str, to_file=False, if to_file: self.dataset_file_name = dataset_file_name np.save(dataset_file_name, mols) + if verbose: + print(ui.color_str('[INFO]', bold=True) + f' save dataset to {dataset_file_name}') def __setitem__(self, key, item): self.__dict__[key] = item @@ -246,17 +288,73 @@ def __str__(self, title_str='Dataset properties'): class Model: + """ + Model object + + This object serve as recipes for model creation, + summarizing the configuration of one particular training run. + Training and test points are sampled from the provided dataset, + without replacement. If the same dataset if given for training + and testing, the subsets are drawn without overlap. + + Each model also contains a choice for the hyper-parameters of the + training process and the MD5 fingerprints of the used datasets. + """ def __init__(self, trainer=None, max_processes=None, use_torch=False, dataset=None, n_train=0, save_dataset=False, valid_dataset=None, n_valid=0, save_valid_dataset=False, - callback=dummy_callback, model=None, + model=None, callback=dummy_callback, **kwargs): + """ + Constructor for the Model object. + The resulting Task object is pre-trained and pre-tested. + + Parameters + ---------- + trainer : :obj:`sgdml.train.GDMLTrain` + predefined GDMLTrain object. overwrite max_processes and use_torch Parameters + dataset : :obj:'sgdml.core.Dataset' + Data structure of custom type :obj:`dataset` containing + training dataset. + n_train : int + number of trainig point to sample from the training dataset. + save_dataset: bool, optional + True: save the dataset object into the model for quick + change of the n_train via the update_hyperparameters method. + False: not saving the training dataset into the Model object. + valid_dataset : :obj:'sgdml.core.Dataset', optional + Data structure of custom type :obj:`dataset` containing + training dataset. + n_valid : int, optional + number of validation point to sample from the validation dataset. + save_valid_dataset: bool, optional + True: save the dataset object into the model for quick + change of the n_valid via the update_hyperparameters() method. + False: not saving the validation dataset into the Model object. + model : :obj:'sgdml.core.Model', optional + A precreate Model object to use as template for a new Model object. + overwrite every other option. + callback : :obj:'sgdml.utils.ui.callback', optional + A callback object for monitoring the training and testing processes. + kwargs : + Parameters for the creation of :obj: 'sgdml.core.Task' + object as a holder for the model hyperparameters + Returns + ------- + sgdml.core.Model + pre-trained and pre-tested type :obj:`sgdml.core.Model`. + + Raises + ------ + ValueError + If a reconstruction of the potential energy surface is requested, + but the energy labels are missing in the dataset. + """ self.max_processes = max_processes self.use_torch = use_torch self.trainer = trainer - if model is not None: for k, v in model.items(): self[k] = v @@ -282,15 +380,39 @@ def __init__(self, trainer=None, max_processes=None, use_torch=False, def update_hyperparameters(self, train_dataset=None, n_train=None, valid_dataset=None, n_valid=None, **kwargs): + """ + Update the Task hyperparameters. + + Parameters + ---------- + train_dataset :obj:`sgdml.core.Dataset` + Data structure of custom type :obj:`dataset` containing a new + train dataset + n_train: int + update value for number of training point for the model + valid_dataset :obj:`sgdml.core.Dataset` + Data structure of custom type :obj:`dataset` containing a new + validation dataset + n_valid: int + update value for number of validation point for the model + + """ + if n_train is not None: - train_dataset = self.dataset if (train_dataset is None and self.dataset is not None) else train_dataset + if train_dataset is None and self.dataset is not None: + train_dataset = self.dataset + kwargs['n_train'] = n_train if n_valid is not None: - valid_dataset = self.valid_dataset if (valid_dataset is None and self.valid_dataset is not None) \ - else valid_dataset + if valid_dataset is None and self.valid_dataset is not None: + valid_dataset = self.valid_dataset + kwargs['n_valid'] = n_valid self.task.update_hyperparameters(train_dataset=train_dataset, valid_dataset=valid_dataset, **kwargs) + for k, v in self.task.items(): + self[k] = v def train(self, callback=None): + del_trainer = False callback = callback if callback is not None else self.callback if self.trainer is None: self.trainer = GDMLTrain(self.max_processes, self.use_torch) @@ -301,51 +423,142 @@ def train(self, callback=None): if del_trainer: del self.trainer self.trainer = None - self.predictor = GDMLPredict(self, - max_processes=self.max_processes, use_torch=self.use_torch) - - def prepare_parallel(self, b_size): - num_workers, batch_size = 0, 0 + self.predictor = None + + + def prepare_parallel(self, n_bulk=1, n_reps=1): + + """ + run the prepare_parallel predictor :obj: 'sgdml.predict.GDMLpredict' method + to find and set the optimal parallelization parameters for the + currently loaded model, running on a particular system. The result + also depends on the number of geometries `n_bulk` that will be + passed at once when calling the `predict` function. + + This function runs a benchmark in which the prediction routine is + repeatedly called `n_reps`-times (default: 1) with varying parameter + configurations, while the runtime is measured for each one. The + optimal parameters are then cached for fast retrival in future + calls of this function. + + Parameters + ---------- + n_bulk : int, optional + Number of geometries that will be passed to the + `predict` function in each call (performance + will be optimized for that exact use case). + n_reps : int, optional + Number of repetitions (bigger value: more + accurate, but also slower). + """ + if self.predictor is None: + self.predictor = GDMLPredict(self, + max_processes=self.max_processes, use_torch=self.use_torch) if not self.use_torch: - if num_workers == 0 or batch_size == 0: - self.callback(NOT_DONE, disp_str='Optimizing parallelism') - - gps, is_from_cache = self.predictor.prepare_parallel(n_bulk=b_size, return_is_from_cache=True) - num_workers, batch_size, bulk_mp = (self.predictor.num_workers, - self.predictor.chunk_size, - self.predictor.bulk_mp, - ) - - self.callback(DONE, - disp_str='Optimizing parallelism' - + (' (from cache)' if is_from_cache else ''), - sec_disp_str='%d workers %s/ chunks of %d' - % (num_workers, '[MP] ' if bulk_mp else '', batch_size), - ) - else: - self.predictor._set_num_workers(num_workers) - self.predictor._set_batch_size(batch_size) - self.predictor._set_bulk_mp(bulk_mp) - - def predict(self, r): - return self.predictor.predict(r) - - def test(self, dataset=None, n_test=None, use_torch=None, overwrite=False): - dataset = self.valid_dataset if (dataset is None and self.valid_dataset is not None) else dataset - # NOTE: this function runs a validation if n_test < 0 and test with all points if n_test == 0 + self.callback(NOT_DONE, disp_str='Optimizing parallelism') + + gps, is_from_cache = self.predictor.prepare_parallel(n_bulk=n_bulk,n_reps=n_reps, return_is_from_cache=True) + num_workers, batch_size, bulk_mp = (self.predictor.num_workers, + self.predictor.chunk_size, + self.predictor.bulk_mp, + ) + + self.callback(DONE, + disp_str='Optimizing parallelism' + + (' (from cache)' if is_from_cache else ''), + sec_disp_str='%d workers %s/ chunks of %d' + % (num_workers, '[MP] ' if bulk_mp else '', batch_size), + ) + else: + print('Skipping multi-CPU benchmark, since torch is enabled.')# TODO: clarity! + + def predict(self, R, R_desc=None, R_d_desc=None): + """ + Predict energy and forces for multiple geometries. This function + can run on the GPU, if the optional PyTorch dependency is + installed and `use_torch=True` was speciefied during + initialization of this class. + + Optionally, the descriptors and descriptor Jacobians for the + same geometries can be provided, if already available from some + previous calculations. + + Note + ---- + The order of the atoms in `R` is not arbitrary and must + be the same as used for training the model. + + Parameters + ---------- + R : :obj:`numpy.ndarray` + An 2D array of size M x 3N containing the + Cartesian coordinates of each atom of M + molecules. + R_desc : :obj:`numpy.ndarray`, optional + An 2D array of size M x D containing the + descriptors of dimension D for M + molecules. + R_d_desc : :obj:`numpy.ndarray`, optional + A 2D array of size M x D x 3N containing of the + descriptor Jacobians for M molecules. The descriptor + has dimension D with 3N partial derivatives with + respect to the 3N Cartesian coordinates of each atom. + + Returns + ------- + :obj:`numpy.ndarray` + Energies stored in an 1D array of size M. + :obj:`numpy.ndarray` + Forces stored in an 2D arry of size M x 3N. + """ + if self.predictor is None: + self.predictor = GDMLPredict(self, + max_processes=self.max_processes, use_torch=self.use_torch) + return self.predictor.predict(R, R_desc=None, R_d_desc=None) + + def test(self, dataset=None, n_test=None, is_test=False, overwrite=False): + + """ + test or validated the model on a dataset + + Note + ---- + this function runs a validation if: is_test is False + + Parameters + ---------- + dataset : :obj:`sgdml.core.Dataset`, optional + dataset to test the model on. if non is given the model will test on the saved validation dataset if present + n_test : int, optional + number of datapoint to test the model on. + if None is given: validation: for validation n_valid will be used if present + else will use the full validation dataset will be used + test: for test the full validation\test dataset will be used + is_test : bool, optional + True: preform a test of the model. + False: preform a validation of the model + overwrite : bool, optional + if to overwrite previous test/validation + """ + + if dataset is None and self.valid_dataset is not None: + dataset = self.valid_dataset + if is_test: + print('Running test on validation dataset') + # NOTE: this function runs a validation if: + # is_test is False + is_validation = not is_test if n_test is None: - try: - n_test = 0 if self.n_valid is None else self.n_valid - except AttributeError: + if is_validation: + try: + n_test = 0 if self.n_valid is None else self.n_valid + except AttributeError: + n_test = 0 + if is_test: n_test = 0 - is_validation = True - is_test = False - else: - is_test = True - is_validation = False + F_rmse = [] - use_torch = self.use_torch if use_torch is None else use_torch if not np.array_equal(self.z, dataset.z): raise AssistantError('Atom composition or order in dataset does not match the one in model.') @@ -377,8 +590,7 @@ def test(self, dataset=None, n_test=None, use_torch=None, overwrite=False): n_data = dataset.F.shape[0] n_data_eff = n_data - len(excl_idxs) - if ( - n_test == 0 and n_data_eff != 0): # test on all data points that have not been used for training or testing + if (n_test == 0 and n_data_eff != 0): # test on all data points that have not been used for training or testing n_test = n_data_eff print('Test set size was automatically set to {:,} points.'.format(n_test)) @@ -409,7 +621,7 @@ def test(self, dataset=None, n_test=None, use_torch=None, overwrite=False): b_size = min(1000, len(test_idxs)) - if not use_torch: + if self.use_torch: self.prepare_parallel(b_size) n_atoms = z.shape[0] @@ -588,8 +800,11 @@ def __str__(self): text += ' {:<18} {}{:,} points'.format('Validated on:', '' if is_valid else '[pending] ', n_valid) + \ ' from \'' + ui.unicode_str(self.md5_valid) + '\'' + '\n' if 'n_test' in self.keys(): - n_test = int(self.n_test) - is_test = True + if self.n_test != 0: + n_test = int(self.n_test) + is_test = True + else: + is_test = False else: is_test = False if is_test: @@ -670,6 +885,19 @@ def __iter__(self): class Task: + + """ + Task object + + This object serve as recipes for model creation, + summarizing the configuration of one particular training run. + Training and test points are sampled from the provided dataset, + without replacement. If the same dataset if given for training + and testing, the subsets are drawn without overlap. + + Each task also contains a choice for the hyper-parameters of the + training process and the MD5 fingerprints of the used datasets. + """ def __init__(self, train_dataset, n_train, valid_dataset=None, @@ -688,6 +916,62 @@ def __init__(self, train_dataset, interact_cut_off=None, # TODO: document me callback=None, ): # TODO: document me + """ + Constructor for the Task object. + The resulting Task object is pre-trained and pre-tested. + + Parameters + ---------- + train_dataset : :obj:`sgdml.core.Dataset` + Data structure of custom type :obj:`dataset` containing + train dataset. + n_train : int + Number of training points to sample. + valid_dataset : :obj:`sgdml.core.Dataset`, optional + Data structure of custom type :obj:`dataset` containing + validation dataset. + n_valid : int, optional + Number of validation points to sample. + sig : int + Hyper-parameter (kernel length scale). + lam : float, optional + Hyper-parameter lambda (regularization strength). + use_sym : bool, optional + True: include symmetries (sGDML), False: GDML. + use_E : bool, optional + True: reconstruct force field with corresponding potential energy surface, + False: ignore energy during training, even if energy labels are available + in the dataset. The trained model will still be able to predict + energies up to an unknown integration constant. Note, that the + energy predictions accuracy will be untested. + use_E_cstr : bool, optional + True: include energy constraints in the kernel, + False: default (s)GDML. + use_cprsn : bool, optional + True: compress kernel matrix along symmetric degrees of + freedom, + False: train using full kernel matrix + + Returns + ------- + sgdml.core.Task + pre-trained and pre-tested type :obj:`sgdml.core.Task`. + + Raises + ------ + ValueError + If a reconstruction of the potential energy surface is requested, + but the energy labels are missing in the dataset. + """ + + if use_E and 'E' not in train_dataset: + raise ValueError( + 'No energy labels found in dataset!\n' + + 'By default, force fields are always reconstructed including the\n' + + 'corresponding potential energy surface (this can be turned off).\n' + + 'However, the energy labels are missing in the provided dataset.\n' + ) + self.type = 't' self.code_version = sgdml.__version__ self.dataset_name = train_dataset.name @@ -701,7 +985,7 @@ def __init__(self, train_dataset, self.max_processes = max_processes self.solver_name = solver self.solver_tol = solver_tol - self.n_inducing_pts_init = n_inducing_pts_init + self.n_inducing_pts_init = alias_int(n_inducing_pts_init) self.interact_cut_off = interact_cut_off self.model0 = model0 self.callback = callback @@ -760,6 +1044,21 @@ def __init__(self, train_dataset, self.e_unit = train_dataset.e_unit def update_model_index(self, n_train, train_dataset): + + """ + Update the Task train indices while taking into account indices of the base model0 + and permutations in the training dataset + + Parameters + ---------- + n_train : int + Number of training points to sample. + train_dataset : :obj:`sgdml.core.Dataset` + Data structure of custom type :obj:`dataset` containing + train dataset. + + """ + md5_train = io.dataset_md5(train_dataset) m0_excl_idxs = np.array([], dtype=np.uint) m0_n_train, m0_n_valid = 0, 0 @@ -857,6 +1156,21 @@ def update_model_index(self, n_train, train_dataset): self.cprsn_keep_atoms_idxs = cprsn_keep_idxs def update_model_valid_index(self, n_valid, valid_dataset): + + """ + Update the Task validation indices while taking into account indices of the base model0. + If using the same dataset for the training and validation will omit training indices. + + Parameters + ---------- + n_valid : int + Number of training points to sample. + valid_dataset : :obj:`sgdml.core.Dataset` + Data structure of custom type :obj:`dataset` containing + train dataset. + + """ + md5_valid = io.dataset_md5(valid_dataset) if self.model0 is not None: m0_idxs_valid = self.model0['idxs_valid'] # TODO: Change the dict syntax @@ -877,6 +1191,19 @@ def update_model_valid_index(self, n_valid, valid_dataset): self.md5_valid = md5_valid def update_hyperparameters(self, train_dataset=None, valid_dataset=None, **kwargs): + """ + Update the Task hyperparameters. + + Parameters + ---------- + train_dataset :obj:`sgdml.core.Dataset` + Data structure of custom type :obj:`dataset` containing a new + train dataset + valid_dataset :obj:`sgdml.core.Dataset` + Data structure of custom type :obj:`dataset` containing a new + validation dataset + + """ for k, v in kwargs.items(): self[k] = v if train_dataset is not None: @@ -948,6 +1275,18 @@ def __str__(self): return self.str +class alias_int(int): + def __init__(self, num): + super().__init__() + self.num = num + + def astype(self, astype): + return astype(self.num) + + def copy(self): + return self.num + + def draw_strat_sample(T, n, excl_idxs=None): """ Draw sample from dataset that preserves its original distribution. @@ -1151,11 +1490,3 @@ def select(validated_models, overwrite=False, max_processes=None, model_file=Non best_model = validated_models[best_idx] return best_model - - -if __name__ == '__main__': - dataset_path = 'cspbbr3-300K-train-4676all.xyz' - dataset = Dataset(dataset_path, name='test_dataset', overwrite=True) - print(dataset['name']) - # for k in dataset.keys(): - # print(k) From 1005cf6e65e8d757cd6cd3ca553d3638d01ff65c Mon Sep 17 00:00:00 2001 From: "oz.mendelsohn" Date: Tue, 16 Feb 2021 14:30:04 +0200 Subject: [PATCH 05/14] update gitignore --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index 6a6df3b..9f1710e 100644 --- a/.gitignore +++ b/.gitignore @@ -11,3 +11,4 @@ /*.egg-info /*.egg sgdml/_bmark_cache.npz +/sgdml/solvers/iterative.py From fbb304614412057dcf33bcf208a8116f147a0904 Mon Sep 17 00:00:00 2001 From: "oz.mendelsohn" Date: Mon, 22 Feb 2021 12:43:53 +0200 Subject: [PATCH 06/14] Update pull request --- .gitignore | 2 +- sgdml/__init__.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.gitignore b/.gitignore index 9f1710e..1fc1fcd 100644 --- a/.gitignore +++ b/.gitignore @@ -11,4 +11,4 @@ /*.egg-info /*.egg sgdml/_bmark_cache.npz -/sgdml/solvers/iterative.py + diff --git a/sgdml/__init__.py b/sgdml/__init__.py index 7c2075c..fcda3e5 100755 --- a/sgdml/__init__.py +++ b/sgdml/__init__.py @@ -22,7 +22,7 @@ # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. -__version__ = '0.4.10-oz' +__version__ = '0.4.10' MAX_PRINT_WIDTH = 100 LOG_LEVELNAME_WIDTH = 7 # do not modify From 3ea0b5d1c60b720a153fdb409f6e4ab932683610 Mon Sep 17 00:00:00 2001 From: "oz.mendelsohn" Date: Tue, 5 Oct 2021 11:25:41 +0300 Subject: [PATCH 07/14] add GPU memory limiter to improve performance --- sgdml/torchtools.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/sgdml/torchtools.py b/sgdml/torchtools.py index a566f64..9027c23 100644 --- a/sgdml/torchtools.py +++ b/sgdml/torchtools.py @@ -100,13 +100,12 @@ def __init__(self, model, lat_and_inv=None, batch_size=None, max_memory=None): const_memory = 2 * self._xs_train.nelement() * self._xs_train.element_size() if max_memory is None: + memory_reduce_factor = 0.9 if torch.cuda.is_available(): - max_memory = min( - [ - torch.cuda.get_device_properties(i).total_memory - for i in range(torch.cuda.device_count()) - ] - ) + max_memory = int(min( + [torch.cuda.get_device_properties(i).total_memory*memory_reduce_factor + for i in range(torch.cuda.device_count()) ] + )) else: max_memory = int( 2 ** 30 * 32 From 7acd7951adcc03ffc5d11b1c9c4161876866c330 Mon Sep 17 00:00:00 2001 From: "oz.mendelsohn" Date: Tue, 5 Oct 2021 11:27:16 +0300 Subject: [PATCH 08/14] Add print statement for number of pools --- sgdml/train.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/sgdml/train.py b/sgdml/train.py index aedc6ad..5c5e819 100755 --- a/sgdml/train.py +++ b/sgdml/train.py @@ -36,6 +36,8 @@ else: Pool = mp.get_context('fork').Pool +from sgdml.dummy_pool import Pool as dPool + import timeit from functools import partial @@ -1242,7 +1244,11 @@ def _assemble_kernel_mat( glob['desc_func'] = desc start = timeit.default_timer() - pool = Pool(self._max_processes) + print(f'starting a Pool of processes: {self._max_processes}') + if self._max_processes == 1: + pool = dPool(self._max_processes) + else: + pool = Pool(self._max_processes) todo, done = K_n_cols, 0 for done_wkr in pool.imap_unordered( From 96d648058bb35d5018e39b7c416edc6b9737811d Mon Sep 17 00:00:00 2001 From: "oz.mendelsohn" Date: Tue, 5 Oct 2021 11:28:19 +0300 Subject: [PATCH 09/14] Add print statement for number of pools --- sgdml/predict.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/sgdml/predict.py b/sgdml/predict.py index 066c685..ae577bd 100644 --- a/sgdml/predict.py +++ b/sgdml/predict.py @@ -37,6 +37,8 @@ else: Pool = mp.get_context('fork').Pool +from sgdml.dummy_pool import Pool as dPool + import timeit from functools import partial @@ -504,8 +506,12 @@ def _reset_mp(self): self.pool.close() self.pool.join() self.pool = None - - self.pool = Pool(processes=self.num_workers) + print(f'starting a Pool of processes: {self._max_processes}') + if self._max_processes == 1: + pool = dPool(processes=self.num_workers) + else: + pool = Pool(processes=self.num_workers) + self.pool = pool self.num_workers = self.pool._processes def _set_chunk_size(self, chunk_size=None): From fd2f250063518d9b7ae8ece0297f8a9c353a84e9 Mon Sep 17 00:00:00 2001 From: "oz.mendelsohn" Date: Tue, 5 Oct 2021 11:30:23 +0300 Subject: [PATCH 10/14] Bug and edge cases fixes --- sgdml/core.py | 440 ++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 410 insertions(+), 30 deletions(-) diff --git a/sgdml/core.py b/sgdml/core.py index 1fea8ca..452cebc 100644 --- a/sgdml/core.py +++ b/sgdml/core.py @@ -6,15 +6,16 @@ import os import sgdml from functools import partial -from . import MAX_PRINT_WIDTH, DONE, NOT_DONE +from . import __version__, MAX_PRINT_WIDTH, DONE, NOT_DONE from .utils import ui, io, perm, desc from .cli import _batch, _online_err from .train import GDMLTrain from .predict import GDMLPredict +import timeit import time -def dummy_callback(done=None, **kwargs): +def dummy_callback(done=None, dummy=None, **kwargs): pass @@ -66,21 +67,23 @@ def __init__(self, dataset_path: str, to_file=False, Returns ------- sgdml.core.Dataset - type :obj:`sgdml.core.Task`. + type :obj:`sgdml.core.Dataset`. """ self.type = 'd' self.code_version = sgdml.__version__ self.name = alias_str(name if not name is None else os.path.splitext(os.path.basename(dataset_path))[0]) self.theory = alias_str(theory) - mols = read(dataset_path, index=':') + dataset_file_name = self.name + '.npz' dataset_exists = os.path.isfile(dataset_file_name) - if dataset_exists and overwrite and verbose: + if (dataset_exists and overwrite) and verbose: print(ui.color_str('[INFO]', bold=True) + ' Overwriting existing dataset file.') - if not dataset_exists or overwrite and verbose: - print('Writing dataset to \'{}\'...'.format(dataset_file_name)) + if (not dataset_exists or overwrite): + if verbose: + # print('Writing dataset to \'{}\'...'.format(dataset_file_name)) + pass else: if verbose: print(ui.color_str('[FAIL]', fore_color=ui.RED, bold=True) @@ -304,6 +307,7 @@ def __init__(self, trainer=None, max_processes=None, use_torch=False, dataset=None, n_train=0, save_dataset=False, valid_dataset=None, n_valid=0, save_valid_dataset=False, model=None, callback=dummy_callback, + solver=None, **kwargs): """ Constructor for the Model object. @@ -359,11 +363,16 @@ def __init__(self, trainer=None, max_processes=None, use_torch=False, for k, v in model.items(): self[k] = v return + if isinstance(trainer, Model): + for k, v in trainer.items(): + self[k] = v + return if 'task' not in kwargs.keys(): self.task = Task(dataset, n_train, valid_dataset=valid_dataset, n_valid=n_valid, - callback=callback, + callback=callback, max_processes = self.max_processes, + solver=solver, **kwargs) else: self.task = kwargs['task'] @@ -376,6 +385,7 @@ def __init__(self, trainer=None, max_processes=None, use_torch=False, self.dataset = dataset if save_valid_dataset: self.valid_dataset = valid_dataset + self.solver = solver def update_hyperparameters(self, train_dataset=None, n_train=None, valid_dataset=None, n_valid=None, @@ -402,6 +412,9 @@ def update_hyperparameters(self, train_dataset=None, n_train=None, if train_dataset is None and self.dataset is not None: train_dataset = self.dataset kwargs['n_train'] = n_train + if 'idxs_train' in kwargs.keys(): + if train_dataset is None and self.dataset is not None: + train_dataset = self.dataset if n_valid is not None: if valid_dataset is None and self.valid_dataset is not None: valid_dataset = self.valid_dataset @@ -411,18 +424,27 @@ def update_hyperparameters(self, train_dataset=None, n_train=None, for k, v in self.task.items(): self[k] = v + self['type'] = 'm' + def train(self, callback=None): del_trainer = False callback = callback if callback is not None else self.callback if self.trainer is None: self.trainer = GDMLTrain(self.max_processes, self.use_torch) del_trainer = True + self.task.solver_name = self.solver if self.solver is not None else 'analytic' + for k, v in self.task.items(): + self.task[k] = self[k] self.model_dict = self.trainer.train(self.task, callback=callback) for k, v in self.model_dict.items(): self[k] = v if del_trainer: del self.trainer self.trainer = None + global glob + + if 'glob' in globals(): + del glob self.predictor = None @@ -572,8 +594,7 @@ def test(self, dataset=None, n_test=None, is_test=False, overwrite=False): if self.use_E: e_err = self.e_err f_err = self.f_err - - is_model_validated = not (np.isnan(f_err['mae']) or np.isnan(f_err['rmse'])) + is_model_validated = not (np.isnan(np.array(f_err['mae'])) or np.isnan(np.array(f_err['rmse']))) if dataset.md5 != self.md5_valid: raise AssistantError('Fingerprint of provided validation dataset does not match the one in model file.') @@ -751,6 +772,361 @@ def test(self, dataset=None, n_test=None, is_test=False, overwrite=False): ) F_rmse.append(f_rmse) + def solve(self, **kwargs): + + del_trainer = False + if self.trainer is None: + self.trainer = GDMLTrain(self.max_processes, self.use_torch) + del_trainer = True + + # TODO: change to if self.solver == None + assert self.solver == 'analytic' or self.solver == 'cg' or not self.solver == None # or solver == 'fk' + + n_train, n_atoms = self.R_train.shape[:2] + + self.desc = desc.Desc( + n_atoms, + interact_cut_off=self.interact_cut_off, + max_processes=self.max_processes, + ) + + n_perms = self.perms.shape[0] + tril_perms = np.array([self.desc.perm(p) for p in self.perms]) + + dim_i = 3 * n_atoms + dim_d = self.desc.dim + + perm_offsets = np.arange(n_perms)[:, None] * dim_d + self.tril_perms_lin = (tril_perms + perm_offsets).flatten('F') + + # TODO: check if all atoms are in span of lattice vectors, otherwise suggest that + # rows and columns might have been switched. + lat_and_inv = None + if 'lattice' in self: + try: + lat_and_inv = (self.lattice, np.linalg.inv(self.lattice)) + except np.linalg.LinAlgError: + raise ValueError( # TODO: Document me + 'Provided dataset contains invalid lattice vectors (not invertible). Note: Only rank 3 lattice vector matrices are supported.' + ) + + # # TODO: check if all atoms are within unit cell + # for r in task['R_train']: + # r_lat = lat_and_inv[1].dot(r.T) + # if not (r_lat >= 0).all(): + # # raise ValueError( # TODO: Document me + # # 'Some atoms appear outside of the unit cell! Please check lattice vectors in dataset file.' + # # ) + # pass + + R = self.R_train.reshape(n_train, -1) + self.R_desc, self.R_d_desc = self.desc.from_R( + R, + lat_and_inv=lat_and_inv, + callback=partial(self.callback, disp_str='Generating descriptors and their Jacobians'), + ) + + # Generate label vector. + E_train_mean = None + y = self.F_train.ravel().copy() + if self.use_E and self.use_E_cstr: + E_train = self.E_train.ravel().copy() + E_train_mean = np.mean(E_train) + + y = np.hstack((y, -E_train + E_train_mean)) + # y = np.hstack((n*Ft, (1-n)*Et)) + y_std = np.std(y) + y /= y_std + + n_train, dim_d = self.R_d_desc.shape[:2] + n_atoms = int((1 + np.sqrt(8 * dim_d + 1)) / 2) + dim_i = 3 * n_atoms + + # Compress kernel based on symmetries + col_idxs = np.s_[:] + if 'cprsn_keep_atoms_idxs' in self: + cprsn_keep_idxs_lin = (np.arange(dim_i).reshape(n_atoms, -1)[self.cprsn_keep_atoms_idxs, :].ravel()) + + col_idxs = (cprsn_keep_idxs_lin[:, None] + np.arange(n_train) * dim_i).T.ravel() + + if self.callback is not None: + self.callback = partial( + self.callback, + disp_str='Assembling kernel matrix', + ) + self.K = self.trainer._assemble_kernel_mat( + self.R_desc, + self.R_d_desc, + self.tril_perms_lin, + self.sig, + self.desc, + use_E_cstr=self.use_E_cstr, + col_idxs=col_idxs, + callback=self.callback, + ) + start = timeit.default_timer() + + alphas = self.solver(np.copy(self.K), np.copy(y.copy()), callback=self.callback, **kwargs).copy() + self.y = y + self.std = y_std + stop = timeit.default_timer() + if self.callback is not None: + dur_s = (stop - start) / 2 + sec_disp_str = 'took {:.1f} s'.format(dur_s) if dur_s >= 0.1 else '' + self.callback(DONE, + disp_str='Training on {:,} points'.format(self.n_train), + sec_disp_str=sec_disp_str, + ) + self.alphas_F = alphas + if self.use_E_cstr: + self.alphas_E = alphas[-n_train:] + self.alphas_F = alphas[:-n_train] + + # Recover integration consta nt. + # Note: if energy constraints are included in the kernel (via 'use_E_cstr'), do not + # compute the integration constant, but simply set it to the mean of the training energies + # (which was subtracted from the labels before training). + # compatibility with original code: + + + if 'cprsn_keep_atoms_idxs' in self: + cprsn_keep_idxs = self.cprsn_keep_atoms_idxs + + R_d_desc_full = self.desc.d_desc_from_comp(self.R_d_desc).reshape( + n_train, dim_d, n_atoms, 3 + ) + R_d_desc_full = R_d_desc_full[:, :, cprsn_keep_idxs, :].reshape( + n_train, dim_d, -1 + ) + + r_d_desc_alpha = np.einsum( + 'kji,ki->kj', R_d_desc_full, self.alphas_F.reshape(n_train, -1) + ) + + else: + + # TOOD: why not use 'd_desc_dot_vec'? + + i, j = np.tril_indices(n_atoms, k=-1) + alphas_F_exp = self.alphas_F.reshape(-1, n_atoms, 3) + + r_d_desc_alpha = np.einsum( + 'kji,kji->kj', self.R_d_desc, alphas_F_exp[:, j, :] - alphas_F_exp[:, i, :] + ) + + self.R_d_desc_alpha = r_d_desc_alpha + self.type = 'm' + self.c = 0.0 + self.R_desc = self.R_desc.T + if self.use_E: + c = (self._recov_int_const() + if E_train_mean is None + else E_train_mean) + if c is None: + # Something does not seem right. Turn off energy predictions for this model, only output force predictions. + self.use_E = False + else: + self.c = c + + if del_trainer: + del self.trainer + self.trainer = None + + + def _recov_int_const(self + ): # TODO: document e_err_inconsist return + """ + Estimate the integration constant for a force field model. + + The offset between the energies predicted for the original training + data and the true energy labels is computed in the least square sense. + Furthermore, common issues with the user-provided datasets are self + diagnosed here. + + Parameters + ---------- + R_desc : :obj:`numpy.ndarray`, optional + An 2D array of size M x D containing the + descriptors of dimension D for M + molecules. + R_d_desc : :obj:`numpy.ndarray`, optional + A 2D array of size M x D x 3N containing of the + descriptor Jacobians for M molecules. The descriptor + has dimension D with 3N partial derivatives with + respect to the 3N Cartesian coordinates of each atom. + + Returns + ------- + float + Estimate for the integration constant. + + Raises + ------ + ValueError + If the sign of the force labels in the dataset from + which the model emerged is switched (e.g. gradients + instead of forces). + ValueError + If inconsistent/corrupted energy labels are detected + in the provided dataset. + ValueError + If different scales in energy vs. force labels are + detected in the provided dataset. + """ + gdml = GDMLPredict(self, max_processes=self.max_processes) # , use_torch=self._use_torch + + n_train = self.E_train.shape[0] + R = self.R_train.reshape(n_train, -1) + E_pred, _ = gdml.predict(R, R_desc=self.R_desc.T, R_d_desc=self.R_d_desc) + E_ref = np.squeeze(self.E_train) + + e_fact = np.linalg.lstsq( + np.column_stack((E_pred, np.ones(E_ref.shape))), E_ref, rcond=-1 + )[0][0] + corrcoef = np.corrcoef(E_ref, E_pred)[0, 1] + + # import matplotlib.pyplot as plt + # sidx = np.argsort(E_ref) + # plt.plot(E_ref[sidx]) + # c = np.sum(E_ref - E_pred) / E_ref.shape[0] + # plt.plot(E_pred[sidx]+c) + # plt.show() + # sys.exit() + + # import matplotlib.pyplot as plt + # sidx = np.argsort(F_ref) + # plt.plot(F_ref[sidx]) + # c = np.sum(F_ref - F_pred) / F_ref.shape[0] + # plt.plot(F_pred[sidx],'--') + # plt.show() + # sys.exit() + + if np.sign(e_fact) == -1: + print('The provided dataset contains gradients instead of force labels (flipped sign). Please correct!\n' + + ui.color_str('Note:', bold=True) + + 'Note: The energy prediction accuracy of the model will thus neither be validated nor tested in the following steps!' + ) + return None + + if corrcoef < 0.95: + print('Inconsistent energy labels detected!\n' + + 'The predicted energies for the training data are only weakly correlated with the reference labels (correlation coefficient {:.2f}) which indicates that the issue is most likely NOT just a unit conversion error.\n\n'.format( + corrcoef + ) + + ui.color_str('Troubleshooting tips:\n', bold=True) + + ui.wrap_indent_str( + '(1) ', + 'Verify the correct correspondence between geometries and labels in the provided dataset.', + ) + + '\n' + + ui.wrap_indent_str( + '(2) ', 'Verify the consistency between energy and force labels.' + ) + + '\n' + + ui.wrap_indent_str(' - ', 'Correspondence correct?') + + '\n' + + ui.wrap_indent_str(' - ', 'Same level of theory?') + + '\n' + + ui.wrap_indent_str(' - ', 'Accuracy of forces?') + + '\n' + + ui.wrap_indent_str( + '(3) ', + 'Is the training data spread too broadly (i.e. weakly sampled transitions between example clusters)?', + ) + + '\n' + + ui.wrap_indent_str( + '(4) ', 'Are there duplicate geometries in the training data?' + ) + + '\n' + + ui.wrap_indent_str( + '(5) ', 'Are there any corrupted data points (e.g. parsing errors)?' + ) + + '\n\n' + + ui.color_str('Note:', bold=True) + + ' The energy prediction accuracy of the model will thus neither be validated nor tested in the following steps!' + ) + return None + + if np.abs(e_fact - 1) > 1e-1: + print('Different scales in energy vs. force labels detected!\n' + + 'The integrated forces differ from the energy labels by factor ~{:.2f}, meaning that the trained model will likely fail to predict energies accurately.\n\n'.format( + e_fact + ) + + ui.color_str('Troubleshooting tips:\n', bold=True) + + ui.wrap_indent_str( + '(1) ', 'Verify consistency of units in energy and force labels.' + ) + + '\n' + + ui.wrap_indent_str( + '(2) ', + 'Is the training data spread too broadly (i.e. weakly sampled transitions between example clusters)?', + ) + + '\n\n' + + ui.color_str('Note:', bold=True) + + ' The energy prediction accuracy of the model will thus neither be validated nor tested in the following steps!' + ) + return None + + # Least squares estimate for integration constant. + return np.sum(E_ref - E_pred) / E_ref.shape[0] + + def to_dict(self): + model = { + 'type': 'm', + 'code_version': __version__, + 'dataset_name': self['dataset_name'], + 'dataset_theory': self['dataset_theory'], + 'solver_name': self.solver, + 'solver_tol': self['solver_tol'], + 'norm_y_train': self.norm_y_train, + 'n_inducing_pts_init': self.n_inducing_pts_init if self.solver == 'cg' else None, + 'z': self['z'], + 'idxs_train': self['idxs_train'], + 'md5_train': self['md5_train'], + 'idxs_valid': self['idxs_valid'], + 'md5_valid': self['md5_valid'], + 'n_test': 0, + 'md5_test': None, + 'f_err': self.f_err, + + 'R_desc': self.R_desc, + 'R_d_desc_alpha': self.R_d_desc_alpha, + 'interact_cut_off': self['interact_cut_off'], + 'c': self.c, + 'std': self.std, + 'sig': self['sig'], + 'lam': self['lam'], + 'alphas_F': self.alphas_F, + 'perms': self['perms'], + 'tril_perms_lin': self.tril_perms_lin, + 'use_E': self['use_E'], + 'use_cprsn': self['use_cprsn'], + } + + if 'solver_resid' in self.keys(): + model['solver_resid'] = self.solver_resid # residual of solution (cg solver) + + if 'solver_iters' in self.keys(): + model[ + 'solver_iters' + ] = self.solver_iters # number of iterations performed to obtain solution (cg solver) + + if 'inducing_pts_idxs' in self.keys(): + model['inducing_pts_idxs'] = self.inducing_pts_idxs + + if self['use_E']: + model['e_err'] = self.e_err, + if self['use_E_cstr']: + model['alphas_E'] = self.alphas_E + + if 'lattice' in self.keys(): + model['lattice'] = self['lattice'] + + if 'r_unit' in self.keys() and 'e_unit' in self.keys(): + model['r_unit'] = self['r_unit'] + model['e_unit'] = self['e_unit'] + return model + def __str__(self): text = ui.white_bold_str('Model properties') + '\n' @@ -983,9 +1359,9 @@ def __init__(self, train_dataset, self.use_sym = use_sym self.use_cprsn = use_cprsn self.max_processes = max_processes - self.solver_name = solver + self.solver_name = solver if solver is not None else 'analytic' self.solver_tol = solver_tol - self.n_inducing_pts_init = alias_int(n_inducing_pts_init) + self.n_inducing_pts_init = alias_int(n_inducing_pts_init) if n_inducing_pts_init is not None else None self.interact_cut_off = interact_cut_off self.model0 = model0 self.callback = callback @@ -998,8 +1374,8 @@ def __init__(self, train_dataset, self.idxs_valid = None self.md5_valid = None if self.use_E: - self.e_err = {'mae': None, 'rmse': None, } - self.f_err = {'mae': None, 'rmse': None} + self.e_err = {'mae': np.NaN, 'rmse': np.NaN, } + self.f_err = {'mae': np.NaN, 'rmse': np.NaN} self.R_train = None self.F_train = None @@ -1022,7 +1398,7 @@ def __init__(self, train_dataset, raise ValueError('Provided training and/or validation dataset(s) do(es) not match the ones in the initial ' 'model.') - self.update_model_index(n_train, train_dataset) + self.update_model_index(train_dataset, n_train) if valid_dataset is not None: self.update_model_valid_index(n_valid, valid_dataset) if use_E: @@ -1043,7 +1419,7 @@ def __init__(self, train_dataset, self.r_unit = train_dataset.r_unit self.e_unit = train_dataset.e_unit - def update_model_index(self, n_train, train_dataset): + def update_model_index(self, train_dataset,n_train=None, idxs_train=None): """ Update the Task train indices while taking into account indices of the base model0 @@ -1051,11 +1427,11 @@ def update_model_index(self, n_train, train_dataset): Parameters ---------- - n_train : int - Number of training points to sample. train_dataset : :obj:`sgdml.core.Dataset` Data structure of custom type :obj:`dataset` containing train dataset. + n_train : int + Number of training points to sample. """ @@ -1075,18 +1451,19 @@ def update_model_index(self, n_train, train_dataset): callback = partial(self.callback, disp_str='Sampling training and validation subsets') callback(NOT_DONE) - if 'E' in train_dataset: - idxs_train = draw_strat_sample(train_dataset.E, n_train - m0_n_train, m0_excl_idxs) - else: - idxs = np.arange(train_dataset.F.shape[0]) - not_use_idx = np.setdiff1d(idxs, self.m0_excl_idxs, assume_unique=True) # DONE: m0 handling - idxs_train = np.random.choice(not_use_idx, n_train - m0_n_train, replace=False) + if idxs_train is None: + if 'E' in train_dataset: + idxs_train = draw_strat_sample(train_dataset.E, n_train - m0_n_train, m0_excl_idxs) + else: + idxs = np.arange(train_dataset.F.shape[0]) + not_use_idx = np.setdiff1d(idxs, self.m0_excl_idxs, assume_unique=True) # DONE: m0 handling + idxs_train = np.random.choice(not_use_idx, n_train - m0_n_train, replace=False) + if self.model0 is not None: + idxs_train = np.concatenate((m0_idxs_train, idxs_train)).astype(np.uint) if self.callback is not None: self.callback(DONE) - if self.model0 is not None: - idxs_train = np.concatenate((m0_idxs_train, idxs_train)).astype(np.uint) self.idxs_train = idxs_train self.md5_train = md5_train self.R_train = train_dataset.R[self.idxs_train, :, :] @@ -1206,8 +1583,11 @@ def update_hyperparameters(self, train_dataset=None, valid_dataset=None, **kwarg """ for k, v in kwargs.items(): self[k] = v + if 'n_inducing_pts_init' in kwargs.keys(): + self['n_inducing_pts_init'] = alias_int(kwargs['n_inducing_pts_init']) if train_dataset is not None: - self.update_model_index(self.n_train, train_dataset) + idxs_train = kwargs['idxs_train'] if 'idxs_train' in kwargs.keys() else None + self.update_model_index(train_dataset, self.n_train, idxs_train) if self.use_E: self.E_train = train_dataset.E[self.idxs_train] @@ -1314,9 +1694,9 @@ def draw_strat_sample(T, n, excl_idxs=None): :obj:`numpy.ndarray` Array of indices that form the sample. """ - - if len(excl_idxs) == 0: - excl_idxs = None + if excl_idxs is not None: + if len(excl_idxs) == 0: + excl_idxs = None if n == 0: return np.array([], dtype=np.uint) From cbbadac5f6b0653089ef97d29652fdbe1566c978 Mon Sep 17 00:00:00 2001 From: "oz.mendelsohn" Date: Tue, 5 Oct 2021 11:33:32 +0300 Subject: [PATCH 11/14] Dummy loop for easy integration one only one thread is needed --- sgdml/dummy_pool.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 sgdml/dummy_pool.py diff --git a/sgdml/dummy_pool.py b/sgdml/dummy_pool.py new file mode 100644 index 0000000..cbe4d69 --- /dev/null +++ b/sgdml/dummy_pool.py @@ -0,0 +1,18 @@ +class Pool(): + def __init__(self, i): + print(f'Dummpy pool of {i} threads') + + def imap_unordered(self, func, iterable, processes=None): + self._processes = 1 + for iter in iterable: + yield func(iter) + + def imap(self, func, iterable): + for iter in iterable: + yield func(iter) + + def close(self): + pass + + def join(self): + pass \ No newline at end of file From 5cdecc7dde2d3d1b29248552e9f211d2e6bbacd5 Mon Sep 17 00:00:00 2001 From: "oz.mendelsohn" Date: Tue, 5 Oct 2021 11:35:05 +0300 Subject: [PATCH 12/14] Add print statement for number of pools --- sgdml/utils/desc.py | 11 ++++++++++- sgdml/utils/perm.py | 11 ++++++++--- 2 files changed, 18 insertions(+), 4 deletions(-) diff --git a/sgdml/utils/desc.py b/sgdml/utils/desc.py index 6614257..d6815de 100755 --- a/sgdml/utils/desc.py +++ b/sgdml/utils/desc.py @@ -32,6 +32,8 @@ else: Pool = mp.get_context('fork').Pool +from sgdml.dummy_pool import Pool as dPool + from functools import partial from scipy import spatial import timeit @@ -337,7 +339,14 @@ def from_R(self, R, lat_and_inv=None, callback=None): # Generate descriptor and their Jacobians start = timeit.default_timer() - pool = Pool(self.max_processes) + + print(f'starting a Pool of processes: {self.max_processes}') + if self.max_processes == 1: + pool = dPool(self.max_processes) + else: + pool = Pool(self.max_processes) + # pool = Pool(self.max_processes) + coff = None if self.coff_dist is None else (self.coff_dist, self.coff_slope) diff --git a/sgdml/utils/perm.py b/sgdml/utils/perm.py index 41efd5d..f947cb7 100755 --- a/sgdml/utils/perm.py +++ b/sgdml/utils/perm.py @@ -31,7 +31,8 @@ else: Pool = mp.get_context('fork').Pool -import sys +from sgdml.dummy_pool import Pool as dPool + import timeit from functools import partial @@ -43,7 +44,6 @@ from .. import DONE, NOT_DONE from .desc import Desc -from . import ui glob = {} @@ -205,7 +205,12 @@ def bipartite_match(R, z, lat_and_inv=None, max_processes=None, callback=None): callback = partial(callback, disp_str='Bi-partite matching') start = timeit.default_timer() - pool = Pool(max_processes) + print(f'starting a Pool of processes: {max_processes}') + if max_processes == 1: + pool = dPool(max_processes) + else: + pool = Pool(max_processes) + # pool = Pool(max_processes) match_perms_all = {} for i, match_perms in enumerate( From 8526f314f366e64ee5cf0755585e1a9bb65b046c Mon Sep 17 00:00:00 2001 From: "oz.mendelsohn" Date: Tue, 5 Oct 2021 11:35:18 +0300 Subject: [PATCH 13/14] bug fix --- sgdml/utils/ui.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/sgdml/utils/ui.py b/sgdml/utils/ui.py index 446f89c..542bb0d 100755 --- a/sgdml/utils/ui.py +++ b/sgdml/utils/ui.py @@ -211,10 +211,12 @@ def warn_str(str): def unicode_str(s): - - if sys.version[0] == '3': - return str(s, 'utf-8', 'ignore') - else: + try: + if sys.version[0] == '3': + return str(s, 'utf-8', 'ignore') + else: + return str(s) + except TypeError: return str(s) From 1e63a3fe689ee2df51d7d3872a12105db068e223 Mon Sep 17 00:00:00 2001 From: Claude Date: Wed, 4 Feb 2026 16:46:43 +0000 Subject: [PATCH 14/14] Add comprehensive CLAUDE.md documentation for AI assistants This file documents the sGDML codebase structure, development workflows, key conventions, and implementation details to help AI assistants understand and work with the codebase effectively. https://claude.ai/code/session_016R5rsk8kNKeyRVWn8ZFig4 --- CLAUDE.md | 294 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 294 insertions(+) create mode 100644 CLAUDE.md diff --git a/CLAUDE.md b/CLAUDE.md new file mode 100644 index 0000000..2400856 --- /dev/null +++ b/CLAUDE.md @@ -0,0 +1,294 @@ +# CLAUDE.md - sGDML Codebase Guide + +## Project Overview + +**sGDML** (Symmetric Gradient Domain Machine Learning) is a Python library for reconstructing accurate molecular force fields from ab initio molecular dynamics (AIMD) data. It uses kernel-based machine learning to predict energies and forces for molecular systems while incorporating molecular symmetries for improved accuracy and data efficiency. + +**Official website:** http://sgdml.org/ +**Documentation:** http://sgdml.org/doc/ + +### Key Features +- Machine learning of molecular force fields from quantum chemistry data +- Symmetry-aware predictions (sGDML) or symmetry-agnostic (GDML) +- CPU multiprocessing and optional GPU acceleration via PyTorch +- ASE (Atomic Simulation Environment) integration for molecular dynamics +- Command-line interface and Python API + +## Repository Structure + +``` +sGDML/ +├── sgdml/ # Main Python package +│ ├── __init__.py # Package init, version, logging config +│ ├── cli.py # Command-line interface implementation +│ ├── core.py # Core classes: Dataset, Model, Task (object-oriented API) +│ ├── train.py # GDMLTrain class for training models +│ ├── predict.py # GDMLPredict class for inference +│ ├── get.py # Dataset/model download utility +│ ├── torchtools.py # PyTorch GPU implementation +│ ├── dummy_pool.py # Single-threaded pool for single-process execution +│ ├── intf/ # Interfaces to external tools +│ │ └── ase_calc.py # ASE calculator integration +│ ├── solvers/ # Linear system solvers +│ │ └── analytic.py # Analytic (direct) solver +│ └── utils/ # Utility modules +│ ├── io.py # File I/O, validation, XYZ format handling +│ ├── ui.py # Terminal UI, progress bars, formatting +│ ├── desc.py # Descriptor computation +│ └── perm.py # Permutation/symmetry detection +├── scripts/ # Dataset conversion scripts +│ ├── sgdml_dataset_from_extxyz.py +│ ├── sgdml_dataset_from_aims.py +│ ├── sgdml_dataset_from_ipi.py +│ ├── sgdml_dataset_via_ase.py +│ ├── sgdml_dataset_to_extxyz.py +│ └── sgdml_datasets_from_model.py +├── setup.py # Package installation config +├── setup.cfg # Flake8, isort configuration +├── pyproject.toml # Black formatter configuration +└── README.md # Project documentation +``` + +## Key Concepts and Data Structures + +### File Types (stored as `.npz` NumPy archives) +- **Dataset (`type='d'`):** Contains molecular geometries, atomic numbers, energies, and forces +- **Task (`type='t'`):** Training configuration with hyperparameters and sampled indices +- **Model (`type='m'`):** Trained model with kernel coefficients and metadata + +### Core Classes + +#### `GDMLTrain` (train.py) +- Creates training tasks from datasets +- Assembles kernel matrices +- Trains models using analytic or iterative solvers +- Handles symmetry detection and compression + +#### `GDMLPredict` (predict.py) +- Loads trained models +- Predicts energies and forces +- Supports CPU multiprocessing and GPU (PyTorch) +- Auto-tunes parallel parameters with `prepare_parallel()` + +#### Object-Oriented API (core.py) +- `Dataset`: Dataset wrapper with ASE file format support +- `Task`: Training task configuration +- `Model`: Trainable model with predict/test methods + +## Development Setup + +### Requirements +- Python 3.7+ +- NumPy >= 1.19 +- SciPy >= 1.1 + +### Optional Dependencies +- **PyTorch:** GPU acceleration (`pip install sgdml[torch]`) +- **ASE >= 3.16.2:** Molecular dynamics integration (`pip install sgdml[ase]`) + +### Installation (Development) +```bash +git clone https://github.com/stefanch/sGDML.git +cd sGDML +pip install -e . +``` + +## Code Style and Conventions + +### Formatting +- **Black:** String normalization and numeric underscore normalization disabled +- **Flake8:** Max complexity 12, ignores E501 (line length), W503, E741 +- **isort:** Multi-line output style 3, trailing comma enabled + +### Code Patterns +- Use `np.load(path, allow_pickle=True)` when loading `.npz` files +- Models/tasks/datasets are dict-like objects with string type identifiers +- Shared memory arrays via `multiprocessing.RawArray` for parallel workers +- Callback functions for progress reporting (see `utils/ui.py`) + +### Multiprocessing +- Uses `fork` context on Unix, `ThreadPool` on Windows +- Global `glob` dict for sharing data with worker processes +- Single-process fallback via `dummy_pool.Pool` when `max_processes=1` + +## CLI Commands + +The package provides two CLI entry points: + +### `sgdml` - Main CLI +```bash +sgdml all [n_test] # Full training pipeline +sgdml create # Create training tasks +sgdml train # Train models +sgdml validate # Validate models +sgdml select # Select best model +sgdml test [n_test] # Test model +sgdml show # Display file info +sgdml reset # Clear caches +``` + +### `sgdml-get` - Download Utility +```bash +sgdml-get dataset [name] # Download benchmark dataset +sgdml-get model [name] # Download pre-trained model +``` + +### Key CLI Options +- `-o, --overwrite`: Overwrite existing files +- `-p, --max_processes`: Limit parallel processes +- `--torch`: Enable GPU acceleration +- `--gdml`: Disable symmetries (use GDML instead of sGDML) +- `--no_E`: Train forces only without energies +- `-s, --sig`: Kernel length scale hyperparameter(s) + +## Python API Usage + +### Basic Prediction +```python +import numpy as np +from sgdml.predict import GDMLPredict +from sgdml.utils import io + +# Load model and geometry +model = np.load('model.npz') +r, z = io.read_xyz('geometry.xyz') + +# Create predictor and predict +gdml = GDMLPredict(model) +gdml.prepare_parallel() # Optimize for performance +e, f = gdml.predict(r) +``` + +### Training Pipeline +```python +from sgdml.train import GDMLTrain +import numpy as np + +# Load dataset +dataset = np.load('dataset.npz', allow_pickle=True) + +# Create trainer and task +gdml_train = GDMLTrain() +task = gdml_train.create_task( + train_dataset=dataset, + n_train=200, + valid_dataset=dataset, + n_valid=1000, + sig=50, + use_sym=True, +) + +# Train model +model = gdml_train.train(task) +np.savez_compressed('model.npz', **model) +``` + +### ASE Integration +```python +from sgdml.intf.ase_calc import SGDMLCalculator +from ase.md import VelocityVerlet +from ase import Atoms + +calc = SGDMLCalculator('model.npz') +atoms = Atoms(...) +atoms.calc = calc + +# Run MD simulation +dyn = VelocityVerlet(atoms, timestep=0.5) +dyn.run(1000) +``` + +## Important Implementation Details + +### Kernel Matrix Assembly +- Implemented in `train.py:_assemble_kernel_mat_wkr()` +- Uses Matern 5/2 kernel with gradient information +- Exploits symmetry for matrix compression +- Parallel assembly via multiprocessing + +### Descriptor Computation +- Inverse pairwise distances as descriptors +- Jacobians computed for force predictions +- Efficient representation using lower triangular indices + +### Symmetry Detection +- Automatic permutation detection in `utils/perm.py` +- Uses graph matching algorithms +- Symmetries reduce training data requirements + +### Energy Integration +- Integration constant recovered via least squares +- Consistency checks for energy/force label alignment +- Optional energy constraints in kernel (`use_E_cstr`) + +## Testing and Validation + +### Model Validation Workflow +1. Train multiple models with different sigma values +2. Validate each on held-out validation set +3. Select model with lowest force RMSE +4. Final test on separate test set + +### Error Metrics +- MAE (Mean Absolute Error) +- RMSE (Root Mean Square Error) +- Force magnitude and angle errors + +## File Format Details + +### Dataset Structure +```python +{ + 'type': 'd', + 'name': 'molecule_name', + 'theory': 'DFT/B3LYP', + 'z': np.array([...]), # Atomic numbers (n_atoms,) + 'R': np.array([...]), # Positions (n_samples, n_atoms, 3) + 'E': np.array([...]), # Energies (n_samples,) + 'F': np.array([...]), # Forces (n_samples, n_atoms, 3) + 'md5': 'fingerprint', + # Optional: 'lattice', 'r_unit', 'e_unit' +} +``` + +### Model Structure +```python +{ + 'type': 'm', + 'z': np.array([...]), # Atomic numbers + 'R_desc': np.array([...]), # Training descriptors + 'R_d_desc_alpha': np.array([...]), # Coefficients + 'alphas_F': np.array([...]), # Force alphas + 'sig': 50, # Length scale + 'perms': np.array([...]), # Permutations + 'c': 0.0, # Integration constant + 'std': 1.0, # Label standard deviation + # ... metadata fields +} +``` + +## Common Tasks for AI Assistants + +### Adding New Functionality +1. Check if feature fits CLI (`cli.py`) or API (`core.py`, `train.py`, `predict.py`) +2. Follow existing callback patterns for progress reporting +3. Update docstrings following NumPy style +4. Handle both dict-based and object-based data structures + +### Debugging Training Issues +1. Check dataset integrity with `sgdml show ` +2. Verify energy/force consistency (watch for sign issues) +3. Review symmetry detection output +4. Check kernel assembly for numerical issues + +### Performance Optimization +1. Use `prepare_parallel()` before predictions +2. Consider GPU acceleration for large systems +3. Monitor memory usage during kernel assembly +4. Adjust `max_processes` for available CPU cores + +## References + +1. Chmiela et al., "Machine Learning of Accurate Energy-conserving Molecular Force Fields", Science Advances (2017) +2. Chmiela et al., "Towards Exact Molecular Dynamics Simulations with Machine-Learned Force Fields", Nature Communications (2018) +3. Chmiela et al., "sGDML: Constructing Accurate and Data Efficient Molecular Force Fields Using Machine Learning", Computer Physics Communications (2019)