def test_get_input(self): cg = CrystalGraph(cutoff=4, bond_converter=GaussianDistance(np.linspace(0, 5, 100), 0.5)) inp = cg.get_input(self.structures[0]) self.assertEqual(len(inp), 7) shapes = [i.shape for i in inp] true_shapes = [(1, 28), (1, 704, 100), (1, 1, 2), (1, 704), (1, 704), (1, 28), (1, 704)] for i, j in zip(shapes, true_shapes): self.assertListEqual(list(i), list(j))
def test_check_dimension(self): gc = CrystalGraph(bond_converter=GaussianDistance(np.linspace(0, 5, 20), 0.5)) s = Structure(Lattice.cubic(3), ['Si'], [[0, 0, 0]]) graph = gc.convert(s) model = MEGNetModel(10, 2, nblocks=1, lr=1e-2, n1=4, n2=4, n3=4, npass=1, ntarget=1, graph_converter=CrystalGraph(bond_converter=gc), ) with self.assertRaises(Exception) as context: model.check_dimension(graph) self.assertTrue('The data dimension for bond' in str(context.exception))
def __init__(self, nfeat_edge=None, nfeat_global=None, nfeat_node=None, nblocks=3, lr=1e-3, n1=64, n2=32, n3=16, nvocal=95, embedding_dim=16, nbvocal=None, bond_embedding_dim=None, ngvocal=None, global_embedding_dim=None, npass=3, ntarget=1, act=softplus2, is_classification=False, loss="mse", l2_coef=None, dropout=None, graph_converter=None, optimizer_kwargs=None): # Build th MEG Model model = make_megnet_model(nfeat_edge=nfeat_edge, nfeat_global=nfeat_global, nfeat_node=nfeat_node, nblocks=nblocks, n1=n1, n2=n2, n3=n3, nvocal=nvocal, embedding_dim=embedding_dim, nbvocal=nbvocal, bond_embedding_dim=bond_embedding_dim, ngvocal=ngvocal, global_embedding_dim=global_embedding_dim, npass=npass, ntarget=ntarget, act=act, is_classification=is_classification, l2_coef=l2_coef, dropout=dropout) # Compile the model with the optimizer loss = 'binary_crossentropy' if is_classification else loss opt_params = {'lr': lr} if optimizer_kwargs is not None: opt_params.update(optimizer_kwargs) model.compile(Adam(**opt_params), loss) if graph_converter is None: graph_converter = CrystalGraph(cutoff=4, bond_converter=GaussianDistance( np.linspace(0, 5, 100), 0.5)) super().__init__(model=model, graph_converter=graph_converter)
def setUpClass(cls): cls.n_feature = 3 cls.n_bond_features = 10 cls.n_global_features = 2 class Generator(Sequence): def __init__(self, x, y): self.x = x self.y = y def __len__(self): return 10 def __getitem__(self, index): return self.x, self.y x_crystal = [np.array([1, 2, 3, 4]).reshape((1, -1)), np.random.normal(size=(1, 6, cls.n_bond_features)), np.random.normal(size=(1, 2, cls.n_global_features)), np.array([[0, 0, 1, 1, 2, 3]]), np.array([[1, 1, 0, 0, 3, 2]]), np.array([[0, 0, 1, 1]]), np.array([[0, 0, 0, 0, 1, 1]]), ] y = np.random.normal(size=(1, 2, 1)) cls.train_gen_crystal = Generator(x_crystal, y) x_mol = [np.random.normal(size=(1, 4, cls.n_feature)), np.random.normal(size=(1, 6, cls.n_bond_features)), np.random.normal(size=(1, 2, cls.n_global_features)), np.array([[0, 0, 1, 1, 2, 3]]), np.array([[1, 1, 0, 0, 3, 2]]), np.array([[0, 0, 1, 1]]), np.array([[0, 0, 0, 0, 1, 1]]), ] y = np.random.normal(size=(1, 2, 1)) cls.train_gen_mol = Generator(x_mol, y) cls.model = MEGNetModel(10, 2, nblocks=1, lr=1e-2, n1=4, n2=4, n3=4, npass=1, ntarget=1, graph_converter=CrystalGraph(bond_converter=GaussianDistance(np.linspace(0, 5, 10), 0.5)), ) cls.model2 = MEGNetModel(10, 2, nblocks=1, lr=1e-2, n1=4, n2=4, n3=4, npass=1, ntarget=2, graph_converter=CrystalGraph(bond_converter=GaussianDistance(np.linspace(0, 5, 10), 0.5)), )
def setUpClass(cls): cls.n_feature = 3 cls.n_bond_features = 10 cls.n_global_features = 2 def generator(x, y): while True: yield x, y x_crystal = [ np.array([1, 2, 3, 4]).reshape((1, -1)), np.random.normal(size=(1, 6, cls.n_bond_features)), np.random.normal(size=(1, 2, cls.n_global_features)), np.array([[0, 0, 1, 1, 2, 3]]), np.array([[1, 1, 0, 0, 3, 2]]), np.array([[0, 0, 1, 1]]), np.array([[0, 0, 0, 0, 1, 1]]), ] y = np.random.normal(size=(1, 2, 1)) cls.train_gen_crystal = generator(x_crystal, y) x_mol = [ np.random.normal(size=(1, 4, cls.n_feature)), np.random.normal(size=(1, 6, cls.n_bond_features)), np.random.normal(size=(1, 2, cls.n_global_features)), np.array([[0, 0, 1, 1, 2, 3]]), np.array([[1, 1, 0, 0, 3, 2]]), np.array([[0, 0, 1, 1]]), np.array([[0, 0, 0, 0, 1, 1]]), ] y = np.random.normal(size=(1, 2, 1)) cls.train_gen_mol = generator(x_mol, y) cls.model = MEGNetModel( 10, 2, nblocks=1, lr=1e-2, n1=4, n2=4, n3=4, npass=1, ntarget=1, graph_convertor=CrystalGraph( bond_convertor=GaussianDistance(np.linspace(0, 5, 10), 0.5)), )
def test_crystalgraph(self): cg = CrystalGraph(cutoff=4) graph = cg.convert(self.structures[0]) self.assertEqual(cg.cutoff, 4) keys = set(graph.keys()) self.assertSetEqual({"bond", "atom", "index1", "index2", "state"}, keys) cg2 = CrystalGraph(cutoff=6) self.assertEqual(cg2.cutoff, 6) graph2 = cg2.convert(self.structures[0]) self.assertListEqual(to_list(graph2["state"][0]), [0, 0]) graph3 = cg(self.structures[0]) np.testing.assert_almost_equal(graph["atom"], graph3["atom"])
def test_crystalgraph(self): cg = CrystalGraph(cutoff=4) graph = cg.convert(self.structures[0]) self.assertEqual(cg.cutoff, 4) keys = set(graph.keys()) self.assertSetEqual({"bond", "atom", "index1", "index2", "state"}, keys) cg2 = CrystalGraph(cutoff=6) self.assertEqual(cg2.cutoff, 6) graph2 = cg2.convert(self.structures[0]) self.assertListEqual(graph2['state'][0], [0, 0]) graph3 = cg(self.structures[0]) self.assertListEqual(graph['atom'], graph3['atom'])
def test_crystalgraph(self): cg = CrystalGraph() graph = cg.convert(self.structures[0]) self.assertEqual(cg.r, 4) keys = set(graph.keys()) self.assertSetEqual({"distance", "node", "index1", "index2", "state"}, keys) cg2 = CrystalGraph(r=6) self.assertEqual(cg2.r, 6) graph2 = cg2.convert(self.structures[0]) self.assertListEqual(graph2['state'][0], [0, 0]) graph3 = cg(self.structures[0]) self.assertListEqual(graph['node'], graph3['node'])
def prepare_model_megnet(individuals, epochs, outfile, excl=[]): # prepares model file # prepares Megnet model based on list of individuals # uses total energy per atom # excl - excluding particular stoichiometry - important for network learning structures = [] energies = [] adapt = AseAtomsAdaptor() empty = 0 if not excl: empty = 1 i = 0 for ind in individuals: struct_ase = ind.get_init_structure() chem_sym = struct_ase.get_chemical_symbols() e_tot = ind.e_tot struct_pymatgen = adapt.get_structure(struct_ase) flag = 1 if empty == 0 and chem_sym == excl: flag = 0 if flag == 1: structures.append(struct_pymatgen) energies.append(e_tot) i = i + 1 print("read data of " + str(i) + " structures total") # standard vales as taken from Megnet manual nfeat_bond = 100 nfeat_global = 2 r_cutoff = 5 gaussian_centers = np.linspace(0, r_cutoff + 1, nfeat_bond) gaussian_width = 0.5 distance_converter = GaussianDistance(gaussian_centers, gaussian_width) graph_converter = CrystalGraph(bond_converter=distance_converter, cutoff=r_cutoff) model = MEGNetModel(nfeat_bond, nfeat_global, graph_converter=graph_converter) # model training model.train(structures, energies, epochs=epochs) model.save_model(outfile)
def test_crystal_model_v2(self): cg = CrystalGraph() s = Structure(Lattice.cubic(3), ['Si'], [[0, 0, 0]]) with ScratchDir('.'): model = MEGNetModel(nfeat_edge=None, nfeat_global=2, nblocks=1, lr=1e-2, n1=4, n2=4, n3=4, npass=1, ntarget=1, graph_converter=cg, centers=np.linspace(0, 4, 10), width=0.5) model = model.train([s, s], [0.1, 0.1], epochs=2) t = model.predict_structure(s) self.assertTrue(t.shape == (1, ))
def test_train_pred(self): model = megnet_model(10, 2, n_blocks=1, lr=1e-2, n1=4, n2=4, n3=4, n_pass=1, n_target=1, graph_convertor=CrystalGraph(), distance_convertor=GaussianDistance(np.linspace(0, 5, 10), 0.5)) s = Structure.from_file(os.path.join(cwd, '../data/tests/cifs/BaTiO3_mp-2998_computed.cif')) structures = [s] * 4 targets = [0.1, 0.1, 0.1, 0.1] model.train(structures, targets, validation_structures=structures[:2], validation_targets=[0.1, 0.1], batch_size=2, epochs=1, verbose=2) preds = model.predict_structure(structures[0]) if os.path.isdir('callback'): shutil.rmtree('callback') self.assertTrue(np.size(preds) == 1)
def default_megnet_config( nfeat_bond: int = 100, r_cutoff: float = 5.0, gaussian_width: float = 0.5 ) -> dict: """Get sensible defaults for MEGNetModel configuration. These arguments are taken from the `MEGNet README file <https://github.com/materialsvirtuallab/megnet#training-a-new-megnetmodel-from-structures>`_. Examples: Create a MEGNetModel using these defaults: >>> model = MEGNetModel(**default_megnet_config()) """ gaussian_centres = np.linspace(0, r_cutoff + 1, nfeat_bond) graph_converter = CrystalGraph(cutoff=r_cutoff) return { "graph_converter": graph_converter, "centers": gaussian_centres, "width": gaussian_width, }
tf.compat.v1.disable_eager_execution() ## Import megnet related modules from megnet.callbacks import ManualStop, ReduceLRUponNan from megnet.data.crystal import CrystalGraph from megnet.data.graph import GaussianDistance, GraphBatchDistanceConvert from megnet.models import MEGNetModel ## Set GPU os.environ["CUDA_VISIBLE_DEVICES"] = "0" # 2. Model construction ## Graph converter crystal_graph = CrystalGraph(bond_converter=GaussianDistance( centers=np.linspace(0, 6, 100), width=0.5), cutoff=5.0) ## model setup model = MEGNetModel( nfeat_edge=100, nfeat_global=None, ngvocal=len(TRAIN_FIDELITIES), global_embedding_dim=16, nblocks=3, nvocal=95, npass=2, graph_converter=crystal_graph, lr=1e-3, ) # 3. Data loading and processing
def test_get_flat_data(self): cg = CrystalGraph(cutoff=4) graphs = [cg.convert(i) for i in self.structures] targets = [0.1, 0.2] inp = cg.get_flat_data(graphs, targets) self.assertListEqual([len(i) for i in inp], [2] * 6)
def test_convert(self): cg = CrystalGraph(cutoff=4) graph = cg.convert(self.structures[0]) self.assertListEqual(graph['atom'], [i.specie.Z for i in self.structures[0]])
def megnet_input(prop, ZeroVals, bond, nfeat_global, cutoff, width, *fraction): """ megnet_input(prop, ZeroVals, bond, nfeat_global, cutoff, width, *fraction) Extracts valid structures and targets and splits them into user specified datsets. Inputs: prop- Optical property of interest. ZeroVals- Exclude/Include zero optical property values. bond- MEGNet feature bond. nfeat_global- MEGNet feature global. cutoff- MEGNet MEGNet radial cutoff. width- MEGNet gaussian width. *fraction- Fraction of data to split into training and validation sets. Passing an extra argument to split data based on quantity is permissible. Outputs: 1- Featurised structures for training with MEGNet. 2- Valid structures and targets. 3- Inputs for extraction of activations. 4- Pool, test, training and validation sets. """ logging.info("Get graph inputs to MEGNet ...") print("Bond features = ", bond) print("Global features = ", nfeat_global) print("Radial cutoff = ", cutoff) print("Gaussian width = ", width) gaussian_centers = np.linspace(0, cutoff, bond) distance_converter = GaussianDistance(gaussian_centers, width) graph_converter = CrystalGraph(bond_converter=distance_converter) model = MEGNetModel(bond, nfeat_global, graph_converter=graph_converter) datafile = "%s_data.pkl" % prop inputs = pd.read_pickle(datafile) print("\nNumber of input entries found for %s data = %s" % (prop, len(inputs))) if ZeroVals == False: logging.info( "Excluding zero optical property values from the dataset ...") mask = np.array( [i for i, val in enumerate(inputs[prop]) if abs(val) == 0.]) structures = np.delete(inputs["structure"].to_numpy(), mask) targets = np.delete(inputs[prop].to_numpy(), mask) print("Remaining number of entries = %s" % len(targets)) else: logging.info("Zero optical property values will be included ...") structures = inputs["structure"].to_numpy() targets = inputs[prop].to_numpy() # Get the valid structures and targets i.e exclude isolated atoms logging.info("Obtaining valid structures and targets ...") valid_structures = [] valid_targets = [] activations_input_full = [] for s, t in zip(structures, targets): try: activations_input_full.append( StructureGraph.get_input(graph_converter, s)) except: print("Skipping structure with isolated atom ...") continue valid_structures.append(s) valid_targets.append(t) print("Number of invalid structures = %s" % (len(targets) - len(valid_targets))) print("\nTotal number of entries available for analysis = %s" % len(valid_targets)) pool_frac = fraction[0][0] if len(fraction) == 1: if (fraction[0][0] + fraction[0][1]) == 1.: # For train-test split and k-fold cross-validation test_frac = fraction[0][1] logging.info("The pool is the same as the training set ...") print("Requested pool: %s%%" % (pool_frac * 100)) print("Requested test set: %s%%" % (test_frac * 100)) # Data split is based on percentages pool_boundary = int(len(valid_targets) * pool_frac) Xpool = np.array(valid_structures[0:pool_boundary]) ypool = np.array(valid_targets[0:pool_boundary]) Xtest = np.array(valid_structures[pool_boundary:]) ytest = np.array(valid_targets[pool_boundary:]) logging.info("The pool is the same as the training set ...") print("Pool:", ypool.shape) print("Test set:", ytest.shape) return (model, activations_input_full, valid_structures, valid_targets, Xpool, ypool, Xtest, ytest) elif (fraction[0][0] + fraction[0][1]) < 1.: # For repeat active learning val_frac = fraction[0][1] test_frac = np.round(1 - pool_frac, decimals=2) pool_boundary = int(len(valid_targets) * pool_frac) Xpool = np.array(valid_structures[0:pool_boundary]) ypool = np.array(valid_targets[0:pool_boundary]) Xtest = np.array(valid_structures[pool_boundary:]) ytest = np.array(valid_targets[pool_boundary:]) val_boundary = int(pool_boundary * val_frac) Xtrain = Xpool[:-val_boundary] ytrain = ypool[:-val_boundary] Xval = Xpool[-val_boundary:] yval = ypool[-val_boundary:] print("Requested validation set: %s%% of pool" % (val_frac * 100)) print("Training set:", ytrain.shape) print("Validation set:", yval.shape) print("Test set:", ytest.shape) return (model, activations_input_full, valid_structures, valid_targets, Xpool, ypool, Xtest, ytest, Xtrain, ytrain, Xval, yval) else: return (model, activations_input_full, np.array(valid_structures), np.array(valid_targets))
def __init__(self, nfeat_edge: int = None, nfeat_global: int = None, nfeat_node: int = None, nblocks: int = 3, lr: float = 1e-3, n1: int = 64, n2: int = 32, n3: int = 16, nvocal: int = 95, embedding_dim: int = 16, nbvocal: int = None, bond_embedding_dim: int = None, ngvocal: int = None, global_embedding_dim: int = None, npass: int = 3, ntarget: int = 1, act: Callable = softplus2, is_classification: bool = False, loss: str = "mse", metrics: List[str] = None, l2_coef: float = None, dropout: float = None, graph_converter: StructureGraph = None, target_scaler: Scaler = DummyScaler(), optimizer_kwargs: Dict = None, dropout_on_predict: bool = False): """ Args: nfeat_edge: (int) number of bond features nfeat_global: (int) number of state features nfeat_node: (int) number of atom features nblocks: (int) number of MEGNetLayer blocks lr: (float) learning rate n1: (int) number of hidden units in layer 1 in MEGNetLayer n2: (int) number of hidden units in layer 2 in MEGNetLayer n3: (int) number of hidden units in layer 3 in MEGNetLayer nvocal: (int) number of total element embedding_dim: (int) number of embedding dimension nbvocal: (int) number of bond types if bond attributes are types bond_embedding_dim: (int) number of bond embedding dimension ngvocal: (int) number of global types if global attributes are types global_embedding_dim: (int) number of global embedding dimension npass: (int) number of recurrent steps in Set2Set layer ntarget: (int) number of output targets act: (object) activation function l2_coef: (float or None) l2 regularization parameter is_classification: (bool) whether it is a classification task loss: (object or str) loss function metrics: (list or dict) List or dictionary of Keras metrics to be evaluated by the model during training and testing dropout: (float) dropout rate graph_converter: (object) object that exposes a "convert" method for structure to graph conversion target_scaler: (object) object that exposes a "transform" and "inverse_transform" methods for transforming the target values optimizer_kwargs (dict): extra keywords for optimizer, for example clipnorm and clipvalue """ # Build the MEG Model model = make_megnet_model(nfeat_edge=nfeat_edge, nfeat_global=nfeat_global, nfeat_node=nfeat_node, nblocks=nblocks, n1=n1, n2=n2, n3=n3, nvocal=nvocal, embedding_dim=embedding_dim, nbvocal=nbvocal, bond_embedding_dim=bond_embedding_dim, ngvocal=ngvocal, global_embedding_dim=global_embedding_dim, npass=npass, ntarget=ntarget, act=act, is_classification=is_classification, l2_coef=l2_coef, dropout=dropout, dropout_on_predict=dropout_on_predict) # Compile the model with the optimizer loss = 'binary_crossentropy' if is_classification else loss opt_params = {'lr': lr} if optimizer_kwargs is not None: opt_params.update(optimizer_kwargs) model.compile(Adam(**opt_params), loss, metrics=metrics) if graph_converter is None: graph_converter = CrystalGraph(cutoff=4, bond_converter=GaussianDistance( np.linspace(0, 5, 100), 0.5)) super().__init__(model=model, target_scaler=target_scaler, graph_converter=graph_converter)
def __init__(self, nfeat_edge, nfeat_global, nfeat_node=None, nblocks=3, lr=1e-3, n1=64, n2=32, n3=16, nvocal=95, embedding_dim=16, npass=3, ntarget=1, act=softplus2, is_classification=False, loss="mse", l2_coef=None, dropout=None, graph_convertor=None, optimizer_kwargs=None ): int32 = 'int32' if nfeat_node is None: x1 = Input(shape=(None,), dtype=int32) # only z as feature x1_ = Embedding(nvocal, embedding_dim)(x1) else: x1 = Input(shape=(None, nfeat_node)) x1_ = x1 x2 = Input(shape=(None, nfeat_edge)) x3 = Input(shape=(None, nfeat_global)) x4 = Input(shape=(None,), dtype=int32) x5 = Input(shape=(None,), dtype=int32) x6 = Input(shape=(None,), dtype=int32) x7 = Input(shape=(None,), dtype=int32) if l2_coef is not None: reg = l2(l2_coef) else: reg = None # two feedforward layers def ff(x, n_hiddens=[n1, n2]): out = x for i in n_hiddens: out = Dense(i, activation=act, kernel_regularizer=reg)(out) return out # a block corresponds to two feedforward layers + one MEGNetLayer layer # Note the first block does not contain the feedforward layer since # it will be explicitly added before the block def one_block(a, b, c, has_ff=True): if has_ff: x1_ = ff(a) x2_ = ff(b) x3_ = ff(c) else: x1_ = a x2_ = b x3_ = c out = MEGNetLayer( [n1, n1, n2], [n1, n1, n2], [n1, n1, n2], pool_method='mean', activation=act, kernel_regularizer=reg)( [x1_, x2_, x3_, x4, x5, x6, x7]) x1_temp = out[0] x2_temp = out[1] x3_temp = out[2] if dropout: x1_temp = Dropout(dropout)(x1_temp) x2_temp = Dropout(dropout)(x2_temp) x3_temp = Dropout(dropout)(x3_temp) return x1_temp, x2_temp, x3_temp x1_ = ff(x1_) x2_ = ff(x2) x3_ = ff(x3) for i in range(nblocks): if i == 0: has_ff = False else: has_ff = True x1_1 = x1_ x2_1 = x2_ x3_1 = x3_ x1_1, x2_1, x3_1 = one_block(x1_1, x2_1, x3_1, has_ff) # skip connection x1_ = Add()([x1_, x1_1]) x2_ = Add()([x2_, x2_1]) x3_ = Add()([x3_, x3_1]) # set2set for both the atom and bond node_vec = Set2Set(T=npass, n_hidden=n3, kernel_regularizer=reg)([x1_, x6]) edge_vec = Set2Set(T=npass, n_hidden=n3, kernel_regularizer=reg)([x2_, x7]) # concatenate atom, bond, and global final_vec = Concatenate(axis=-1)([node_vec, edge_vec, x3_]) if dropout: final_vec = Dropout(dropout)(final_vec) # final dense layers final_vec = Dense(n2, activation=act, kernel_regularizer=reg)(final_vec) final_vec = Dense(n3, activation=act, kernel_regularizer=reg)(final_vec) if is_classification: final_act = 'sigmoid' loss = 'binary_crossentropy' else: final_act = None loss = loss out = Dense(ntarget, activation=final_act)(final_vec) model = Model(inputs=[x1, x2, x3, x4, x5, x6, x7], outputs=out) opt_params = {'lr': lr} if optimizer_kwargs is not None: opt_params.update(optimizer_kwargs) model.compile(Adam(**opt_params), loss) if graph_convertor is None: graph_convertor = CrystalGraph(cutoff=4, bond_convertor=GaussianDistance(np.linspace(0, 5, 100), 0.5)) super().__init__(model=model, graph_convertor=graph_convertor)
def train(): # Parse args args = parse_args() radius = args.radius n_works = args.n_works warm_start = args.warm_start output_path = args.output_path graph_file = args.graph_file prop_col = args.property learning_rate = args.learning_rate embedding_file = args.embedding_file k_folds = list(map(int, args.k_folds.split(","))) print("args is : {}".format(args)) print("Local devices are : {}, \n\n Available gpus are : {}".format( device_lib.list_local_devices(), K.tensorflow_backend._get_available_gpus())) # prepare output path if not os.path.exists(output_path): os.makedirs(output_path, exist_ok=True) # Get a crystal graph with cutoff radius A cg = CrystalGraph( bond_convertor=GaussianDistance(np.linspace(0, radius + 1, 100), 0.5), cutoff=radius, ) if graph_file is not None: # load graph data with gzip.open(graph_file, "rb") as f: valid_graph_dict = pickle.load(f) idx_list = list(range(len(valid_graph_dict))) valid_idx_list = [ idx for idx, graph in valid_graph_dict.items() if graph is not None ] else: # load structure data with gzip.open(args.input_file, "rb") as f: df = pd.DataFrame(pickle.load(f))[["structure", prop_col]] idx_list = list(range(len(df))) # load embedding data for transfer learning if embedding_file is not None: with open(embedding_file) as json_file: embedding_data = json.load(json_file) # Calculate and save valid graphs valid_idx_list = list() valid_graph_dict = dict() for idx in idx_list: try: graph = cg.convert(df["structure"].iloc[idx]) if embedding_file is not None: graph["atom"] = [embedding_data[i] for i in graph["atom"]] valid_graph_dict[idx] = { "graph": graph, "target": df[prop_col].iloc[idx], } valid_idx_list.append(idx) except RuntimeError: valid_graph_dict[idx] = None # Save graphs with gzip.open(os.path.join(output_path, "graphs.pkl.gzip"), "wb") as f: pickle.dump(valid_graph_dict, f) # Split data kf = KFold(n_splits=args.cv, random_state=18012019, shuffle=True) for fold, (train_val_idx, test_idx) in enumerate(kf.split(idx_list)): print(fold) if fold not in k_folds: continue fold_output_path = os.path.join(output_path, "kfold_{}".format(fold)) fold_model_path = os.path.join(fold_output_path, "model") if not os.path.exists(fold_model_path): os.makedirs(fold_model_path, exist_ok=True) train_idx, val_idx = train_test_split(train_val_idx, test_size=0.25, random_state=18012019, shuffle=True) # Calculate valid train validation test ids and save it valid_train_idx = sorted(list(set(train_idx) & (set(valid_idx_list)))) valid_val_idx = sorted(list(set(val_idx) & (set(valid_idx_list)))) valid_test_idx = sorted(list(set(test_idx) & (set(valid_idx_list)))) np.save(os.path.join(fold_output_path, "train_idx.npy"), valid_train_idx) np.save(os.path.join(fold_output_path, "val_idx.npy"), valid_val_idx) np.save(os.path.join(fold_output_path, "test_idx.npy"), valid_test_idx) # Prepare training graphs train_graphs = [valid_graph_dict[i]["graph"] for i in valid_train_idx] train_targets = [ valid_graph_dict[i]["target"] for i in valid_train_idx ] # Prepare validation graphs val_graphs = [valid_graph_dict[i]["graph"] for i in valid_val_idx] val_targets = [valid_graph_dict[i]["target"] for i in valid_val_idx] # Normalize targets or not if args.normalize: y_scaler = StandardScaler() train_targets = y_scaler.fit_transform( np.array(train_targets).reshape(-1, 1)).ravel() val_targets = y_scaler.transform( np.array(val_targets).reshape((-1, 1))).ravel() else: y_scaler = None # Initialize model if warm_start is None: # Set up model if learning_rate is None: learning_rate = 1e-3 model = MEGNetModel( 100, 2, nblocks=args.n_blocks, nvocal=95, npass=args.n_pass, lr=learning_rate, loss=args.loss, graph_convertor=cg, is_classification=True if args.type == "classification" else False, nfeat_node=None if embedding_file is None else 16, ) initial_epoch = 0 else: # Model file model_list = [ m_file for m_file in os.listdir( os.path.join(warm_start, "kfold_{}".format(fold), "model")) if m_file.endswith(".hdf5") ] if args.type == "classification": model_list.sort( key=lambda m_file: float( m_file.split("_")[3].replace(".hdf5", "")), reverse=False, ) else: model_list.sort( key=lambda m_file: float( m_file.split("_")[3].replace(".hdf5", "")), reverse=True, ) model_file = os.path.join(warm_start, "kfold_{}".format(fold), "model", model_list[-1]) # Load model from file if learning_rate is None: full_model = load_model( model_file, custom_objects={ "softplus2": softplus2, "Set2Set": Set2Set, "mean_squared_error_with_scale": mean_squared_error_with_scale, "MEGNetLayer": MEGNetLayer, }, ) learning_rate = K.get_value(full_model.optimizer.lr) # Set up model model = MEGNetModel( 100, 2, nblocks=args.n_blocks, nvocal=95, npass=args.n_pass, lr=learning_rate, loss=args.loss, graph_convertor=cg, is_classification=True if args.type == "classification" else False, nfeat_node=None if embedding_file is None else 16, ) model.load_weights(model_file) initial_epoch = int(model_list[-1].split("_")[2]) print("warm start from : {}, \nlearning_rate is {}.".format( model_file, learning_rate)) # Train model.train_from_graphs( train_graphs, train_targets, val_graphs, val_targets, batch_size=args.batch_size, epochs=args.max_epochs, verbose=2, initial_epoch=initial_epoch, use_multiprocessing=False if n_works <= 1 else True, workers=n_works, dirname=fold_model_path, y_scaler=y_scaler, save_best_only=args.save_best_only, )
db = connect(db_name) ###### megnet example hyper-parameters from megnet.models import MEGNetModel from megnet.data.graph import GaussianDistance from megnet.data.crystal import CrystalGraph import numpy as np nfeat_bond = 100 nfeat_global = 2 r_cutoff = 5 gaussian_centers = np.linspace(0, r_cutoff + 1, nfeat_bond) gaussian_width = 0.5 distance_converter = GaussianDistance(gaussian_centers, gaussian_width) graph_converter = CrystalGraph(bond_converter=distance_converter, cutoff=r_cutoff) model = MEGNetModel(nfeat_bond, nfeat_global, graph_converter=graph_converter) ######################################### def cvt_fmt_graph(rows): structures = [] props = [] for row in rows: structures.append( pymatgen_io_ase.AseAtomsAdaptor.get_structure(row.toatoms())) props.append(row.data[predict_item] / 100) # props.append(abs(row.data[predict_item]/10)) graphs_valid = [] targets_valid = []
result = keras.losses.mean_squared_error(y_true, y_pred) # result = K.print_tensor(result, message='losses') return result # === megnet start === # from megnet.models import MEGNetModel from megnet.data.graph import GaussianDistance from megnet.data.crystal import CrystalGraph from megnet.utils.preprocessing import StandardScaler from megnet.callbacks import ReduceLRUponNan, ManualStop, XiaotongCB import numpy as np gc = CrystalGraph(bond_converter=GaussianDistance( np.linspace(0, 5, 100), 0.5), cutoff=4) model = MEGNetModel(100, 2, graph_converter=gc, lr=1e-4, loss=examine_loss) # , metrics=[examine_loss]) INTENSIVE = False # U0 is an extensive quantity scaler = StandardScaler.from_training_data(structures, targets, is_intensive=INTENSIVE) model.target_scaler = scaler # callbacks = [ReduceLRUponNan(patience=500), ManualStop(), XiaotongCB()] # change structures to megnet predictable structures mp_strs = [] train_graphs, train_targets = model.get_all_graphs_targets(structures, targets) train_nb_atoms = [len(i['atom']) for i in train_graphs] train_targets = [model.target_scaler.transform(i, j) for i, j in zip(train_targets, train_nb_atoms)]
ME += e error_lst.append(e) if abs(e) > 0.5: targets[i] = model.predict_structure(structures[i]).ravel() # targets[i] = (model.predict_structure(structures[i]).ravel() + targets[i])/2 ME /= sz f = open(str(sz) + 'txt', 'wb') pickle.dump(error_lst, f) f.close() # for i in range(idx, idx + sz): # targets[i] += ME idx += sz model = MEGNetModel(10, 2, nblocks=1, lr=1e-4, n1=4, n2=4, n3=4, npass=1, ntarget=1, graph_converter=CrystalGraph(bond_converter=GaussianDistance(np.linspace(0, 5, 10), 0.5))) ep = 5000 callback = tf.keras.callbacks.EarlyStopping(monitor="val_loss", patience=50, restore_best_weights=True) for s in test_structures: test_input.append(model.graph_converter.graph_to_input(model.graph_converter.convert(s))) if training_mode == 0: # PBE -> HSE ... -> part EXP, one by one idx = 0 for i in range(len(data_size)): model.train(structures[idx:idx+data_size[i]], targets[idx:idx+data_size[i]], epochs=ep) idx += data_size[i] prediction(model) elif training_mode == 1: # all training set together
epochs = 5 batch_size = 56 Xtrain = inputs.iloc[0:boundary]['structure'] ytrain = inputs.iloc[0:boundary]['band_gap'] Xtest = inputs.iloc[boundary:]['structure'] ytest = inputs.iloc[boundary:]['band_gap'] nfeat_bond = 10 nfeat_global = 2 r_cutoff = 5 gaussian_centers = np.linspace(0, 5, 10) gaussian_width = 0.5 distance_convertor = GaussianDistance(gaussian_centers, gaussian_width) bond_convertor = CrystalGraph(bond_convertor=distance_convertor, cutoff=r_cutoff) graph_convertor = CrystalGraph( bond_convertor=GaussianDistance(np.linspace(0, 5, 10), 0.5)) model = MEGNetModel(nfeat_bond, nfeat_global, graph_convertor=graph_convertor) model.from_file('fitted_gap_model.hdf5') model.train(Xtrain, ytrain, epochs=epochs, batch_size=batch_size, validation_structures=Xtest, validation_targets=ytest, scrub_failed_structures=True) model.save_model('fitted_gap_model.hdf5')
for i in range(queries): # Train MegNet print('Query ', i) print('Sample from y test for consistency', len(ytest), ytest[0]) Xactive = np.concatenate((Xpool, Xunlab)) yactive = np.concatenate((ypool, yunlab)) training.active(i, prop, model, 'entropy', batch, epochs, Xpool, ypool, Xtest, ytest) # Get the activations for the active set activations = [] gaussian_centers = np.linspace(0, cutoff, bond) distance_converter = GaussianDistance(gaussian_centers, width) graph_converter = CrystalGraph(bond_converter=distance_converter) for s in Xactive: activations.append(StructureGraph.get_input(graph_converter, s)) # Obtain latent points tsne_active = latent.active( i, prop, perp, layer, 'entropy', activations, Xactive, Xpool, ypool, Xtest, val_frac, ndims, niters) # Split the data tsne_pool = tsne_active[:len(ypool)] tsne_unlab = tsne_active[len(ypool):] cut = int(len(tsne_pool)*split_pool) tsne_train = tsne_pool[:cut] ytrain = ypool[:cut] tsne_val = tsne_pool[cut:] yval = ypool[cut:]
if abs(e) > cut_value: targets[it][i] = prdc # targets[i] = (model.predict_structure(structures[i]).ravel() + targets[i])/2 logging.info('Data count: {dc}, std orig dft value: {std_orig}, std of model output: {std_model}'.format( dc=len(targets_lst), std_orig=np.std(targets_lst), std_model=np.std(prediction_lst))) logging.info('Data count: {dc}, Mean orig: {mean_orig}, Mean_model: {mean_model}'.format( dc=len(targets_lst), mean_orig=np.mean(targets_lst), mean_model=np.mean(prediction_lst))) f = open(dump_model_name + '_'+ it + '.txt', 'wb') # to store and analyze the error pickle.dump(error_lst, f) f.close() # model = MEGNetModel(10, 2, nblocks=3, lr=1e-3, # n1=4, n2=4, n3=4, npass=1, ntarget=1, # graph_converter=CrystalGraph(bond_converter=GaussianDistance(np.linspace(0, 5, 10), 0.5))) model = MEGNetModel(nfeat_edge=10, nfeat_global=2, graph_converter=CrystalGraph(bond_converter=GaussianDistance(np.linspace(0, 5, 10), 0.5))) model.save_model(dump_model_name+'_1by1_init_randomly' + '.hdf5') init_model_tag = 'EGPHS' ep = 5000 callback = tf.keras.callbacks.EarlyStopping(monitor="val_loss", patience=10, restore_best_weights=True) for s in test_structures: test_input.append(model.graph_converter.graph_to_input(model.graph_converter.convert(s))) db_short_full_dict = {'G': 'gllb-sc', 'H': 'hse', 'S': 'scan', 'P': 'pbe', 'E': 'E1'} def construct_dataset_from_str(db_short_str): s = [] t = [] for i in range(len(db_short_str)):