text_plain_output_1.png
import ase as ase

import dscribe as ds

import pandas as pd

import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

structure = pd.read_csv('../input/structures.csv')

rcut = 10.0
g2_params = [[1, 2], [0.1, 2], [0.01, 2], [1, 6], [0.1, 6], [0.01, 6]]
g4_params = [[1, 4, 1], [0.1, 4, 1], [0.01, 4, 1], [1, 4, -1], [0.1, 4, -1], [0.01, 4, -1]]
g3_params = None
g5_params = None

tmp_structure = structure.loc[structure.molecule_name == 'dsgdb9nsd_000001', :].copy()
species = tmp_structure.atom.unique()
acsf = ds.descriptors.ACSF(species=species, rcut=rcut, g2_params=g2_params, g3_params=g3_params, g4_params=g4_params, g5_params=g5_params)

molecule_atoms = tmp_structure.loc[:, 'atom']
molecule_positions = tmp_structure.loc[:, ['x', 'y', 'z']]
molecule_system = ase.atoms.Atoms(symbols=molecule_atoms, positions=molecule_positions)

acsf_features = acsf.create(molecule_system, n_jobs=1)
acsf_features[0]

def create_feature_labels(species, rcut, g2_params=None, g3_params=None, g4_params=None, g5_params=None, transform_to_symbols=True):
    def get_atom_id(atom_nr, tranform_to_symbols):
        if transform_to_symbols == True:
            atom_id = nr_to_symbol[atom_nr]
        else:
            atom_id = atom_nr
        return atom_id
    feature_label = []
    g_params = {'g1': [rcut], 'g2': g2_params, 'g3': g3_params, 'g4': g4_params, 'g5': g5_params}
    tmp_system = ase.Atoms(species, [[0, 0, 0]] * len(species))
    nr_to_symbol = {number: symbol for symbol, number in zip(tmp_system.get_chemical_symbols(), tmp_system.get_atomic_numbers())}
    atomic_numbers = sorted(tmp_system.get_atomic_numbers())
    for atom_nr in atomic_numbers:
        atom_id = get_atom_id(atom_nr, transform_to_symbols)
        for g in ['g1', 'g2', 'g3']:
            params = g_params[g]
            if params is not None:
                for para in params:
                    feature_label.append(f'feat_acsf_{g}_{atom_id}_{para}')
    for atom_nr in atomic_numbers:
        atom_id = get_atom_id(atom_nr, transform_to_symbols)
        for i in range(0, atom_nr + 1):
            if i in atomic_numbers:
                atom_id_2 = get_atom_id(i, transform_to_symbols)
                for g in ['g4', 'g5']:
                    params = g_params[g]
                    if params is not None:
                        for para in params:
                            feature_label.append(f'feat_acsf_{g}_{atom_id}_{atom_id_2}_{para}')
    return feature_label
def calculate_symmetric_functions(df_structure, rcut, g2_params=None, g3_params=None, g4_params=None, g5_params=None):
    species = df_structure.atom.unique()
    acsf = ds.descriptors.ACSF(species=species, rcut=rcut, g2_params=g2_params, g3_params=g3_params, g4_params=g4_params, g5_params=g5_params)
    structure_molecules = df_structure.molecule_name.unique()
    acsf_feature_labels = create_feature_labels(species=species, rcut=rcut, g2_params=g2_params, g3_params=g3_params, g4_params=g4_params, g5_params=g5_params)
    df_structure = df_structure.reindex(columns=df_structure.columns.tolist() + acsf_feature_labels)
    df_structure = df_structure.sort_values(['molecule_name', 'atom_index'])
    acsf_structure_chunks = calculate_acsf_in_chunks(structure_molecules, df_structure, acsf, acsf_feature_labels)
    acsf_structure = pd.DataFrame().append(acsf_structure_chunks)
    return acsf_structure
def calculate_acsf_in_chunks(structure_molecules, df_structure, acsf, acsf_feature_labels, step_size=2000):
    mol_counter = 0
    max_counter = len(structure_molecules)
    all_chunks = []
    tic = time.time()
    while mol_counter * step_size < max_counter:
        tmp_molecules = structure_molecules[mol_counter * step_size:(mol_counter + 1) * step_size]
        tmp_structure = df_structure.loc[df_structure.molecule_name.isin(tmp_molecules), :].copy()
        tmp_results = calculate_acsf_multiple_molecules(tmp_molecules, tmp_structure, acsf, acsf_feature_labels)
        all_chunks.append(tmp_results.copy())
        mol_counter += 1
    return all_chunks
def calculate_acsf_multiple_molecules(molecule_names, df_structure, acsf, acsf_feature_labels):
    counter = 0
    tic = time.time()
    for molecule_name in molecule_names:
        df_molecule = df_structure.loc[df_structure.molecule_name == molecule_name, :]
        acsf_values = calculate_acsf_single_molecule(df_molecule, acsf)
        df_structure.loc[df_structure.molecule_name == molecule_name, acsf_feature_labels] = copy.copy(acsf_values)
        counter += 1
    return df_structure
def calculate_acsf_single_molecule(df_molecule, acsf):
    molecule_atoms = df_molecule.loc[:, 'atom']
    molecule_positions = df_molecule.loc[:, ['x', 'y', 'z']]
    molecule_system = ase.atoms.Atoms(symbols=molecule_atoms, positions=molecule_positions)
    return acsf.create(molecule_system, n_jobs=1)

acsf_structure = calculate_symmetric_functions(structure.loc[structure.molecule_name == 'dsgdb9nsd_000001', :].copy(), rcut, g2_params=g2_params, g4_params=g4_params)
acsf_structure.head()