def test_to_list(self): x = 1 y = [1] z = tuple([1, 2, 3]) v = np.array([1, 2, 3]) for k in [x, y, z, v]: self.assertTrue(type(to_list(k)), list)
def get_flat_data(self, graphs, targets): """ Expand the graph dictionary to form a list of features and targets tensors This is useful when the model is trained on assembled graphs on the fly Args: graphs: (list of dictionary) list of graph dictionary for each structure targets: (list of float or list) correpsonding target values for each structure Returns: tuple(node_features, edges_features, global_values, index1, index2, targets) """ atoms = [] bonds = [] states = [] index1 = [] index2 = [] final_targets = [] for g, t in zip(graphs, targets): if isinstance(g, dict): atoms.append(np.array(g['atom'])) bonds.append(np.array(g['bond'])) states.append(g['state']) index1.append(g['index1']) index2.append(g['index2']) final_targets.append(to_list(t)) return atoms, bonds, states, index1, index2, final_targets
def index_rep_from_structure(structure, r=4): """ Take a pymatgen structure and convert it to a index-type graph representation The graph will have node, distance, index1, index2, where node is a vector of Z number of atoms in the structure, index1 and index2 mark the atom indices forming the bond and separated by distance Args: structure: (pymatgen structure) r: (float) distance cutoff Returns: (dictionary) """ atom_i_segment_id = [ ] # index list for the center atom i for all bonds (row index) atom_i_j_id = [] # index list for atom j atom_number = [] all_neighbors = structure.get_all_neighbors(r, include_index=True) distances = [] for k, n in enumerate(all_neighbors): atom_number.append(structure.sites[k].specie.Z) if len(n) < 1: index = None else: _, distance, index = list(zip(*n)) index = np.array(index) distance = np.array(distance) if index is not None: ind = np.argsort(index) it = itemgetter(*ind) index = it(index) index = to_list(index) index = [int(i) for i in index] distance = distance[ind] distances.append(distance) atom_i_segment_id.extend([k] * len(index)) atom_i_j_id.extend(index) else: pass if len(distances) < 1: return None else: return { 'distance': np.concatenate(distances), 'index1': atom_i_segment_id, 'index2': atom_i_j_id, 'node': atom_number }
def test_to_list(self): x = 1 y = [1] self.assertListEqual(to_list(x), [1]) self.assertListEqual(to_list(y), y)