<jupyter_start><jupyter_text>IntroductionThis notebook demonstrates the creation of an artificial neural network that identifies the crystal structure of individual particles. It uses `pythia` to generate the descriptors and [keras](http://keras.io) to construct the ANNs. Several types of descriptors are used to compare their ability to distinguish various structures. Example dataWe keep some lightly-thermalized example data of a few structures in the pythia project. [libgetar](https://github.com/glotzerlab/libgetar) is required to use these examples (`pip install gtar`), but it should be straightforward to adapt this notebook to work with your own data.<jupyter_code>%config HistoryManager.enabled=False
import os
os.environ['OMP_NUM_THREADS'] = os.environ['GOTO_NUM_THREADS'] = '1'
import freud
freud.parallel.set_num_threads(1)
import multiprocessing
Nproc = min(10, multiprocessing.cpu_count())
import numpy as np

import matplotlib, matplotlib.pyplot as pp
%matplotlib inline

import gtar
import pythia

import functools
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.models import Sequential, model_from_json
from tensorflow.keras.layers import Dense, Dropout, Activation, BatchNormalization

gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    tf.config.experimental.set_memory_growth(gpu, True)

try:
    from tqdm.notebook import tqdm as track_progress
except ImportError:
    def track_progress(target):
        for v in target:
            yield v<jupyter_output><empty_output><jupyter_text>In this repository, structure filenames are generated by the interactions that generated them and then the structure name.<jupyter_code>!wget -c https://bitbucket.org/glotzer/pythia/downloads/example_structures.zip -O example_structures.zip
!unzip -u example_structures.zip
from collections import defaultdict
import os
import re

filename_pattern = re.compile(r'^(?P<data_set>[a-zA-Z]+)_(?P<structure>[a-zA-Z0-9\-]+)\.sqlite$')

remap_names = {'cP54-K4Si23': 'ClaI'}

structure_fnames = defaultdict(list)
fname_structures = {}
for fname in os.listdir('.'):
    match = filename_pattern.match(fname)

    if match:
        structure = match.group('structure')
        structure = remap_names.get(structure, structure)

        structure_fnames[structure].append(fname)
        fname_structures[fname] = structure

sorted_structures = list(sorted(structure_fnames))
structure_indices = {struc: i for (i, struc) in enumerate(sorted_structures)}
all_fnames = sum((structure_fnames[struc] for struc in sorted_structures), [])

print('Found structures:')
structure_fnames<jupyter_output>Found structures:<jupyter_text>Descriptor parametersModels will be created for each descriptor mode specified in `all_modes`. Local environment spherical harmonic descriptors ('amean')- neigh_max = maximum number of neighbors to look at for bonds- lmax = maximum order of spherical harmonics to look at Voronoi cell angle histogram descriptors ('voronoi_angle_histogram')- bins = Number of bins to histogram over Normalized radial distance ('normalized_radial_distance')- neighbors = number of neighbors to look at Pairwise bond angle and distance matrix singular values ({'neighborhood_range_angle_singvals',          'neighborhood_range_distance_singvals'})- neigh_min = minimum number of neighbors to look at- neigh_max = maximum number of neighbors to look at Vector of Steinhardt order parameters ('steinhardt_q')- neighbors = number of neighbors to look at- lmax = maximum order of spherical harmonics to look at Spherical harmonic bispectrum ('bispectrum_sphs')- neighbors = number of neighbors to look at- lmax = maximum order of spherical harmonics to look at<jupyter_code>params = dict(
    amean=dict(neigh_max=12, lmax=12),
    voronoi_angle_histogram=dict(bins=32),
    normalized_radial_distance=dict(neighbors=16),
    neighborhood_range_angle_singvals=dict(neigh_min=4, neigh_max=12),
    neighborhood_angle_sorted=dict(neighbors=16),
    neighborhood_range_distance_singvals=dict(neigh_min=4, neigh_max=12),
    neighborhood_distance_sorted=dict(neighbors=16),
    steinhardt_q=dict(neighbors=12, lmax=20),
    bispectrum_sphs=dict(neighbors=4, lmax=6)
    )

all_modes = ['normalized_radial_distance',
             'neighborhood_angle_sorted',
             'neighborhood_distance_sorted',
             'voronoi_angle_histogram',
             'amean',
             'steinhardt_q',
             'bispectrum_sphs',
            ]<jupyter_output><empty_output><jupyter_text>Supervised LearningHere we create an artificial neural network and use it to identify the structure of individual particles.We start by generating the per-particle descriptors that will be fed into the network.<jupyter_code># %%time

from collections import namedtuple
import json
import re
import sys
import time
import functools
import multiprocessing

FakeBox = namedtuple('FakeBox', ['Lx', 'Ly', 'Lz', 'xy', 'xz', 'yz'])

def computeParticleDescriptors(box, positions, mode, seed=13, limit=None,
                               average_over_neighbors=None, neighbor_average_rounds=None):
    if '|' in mode:
        single_mode_descriptors = []
        for minor_mode in mode.split('|'):
            single_mode_descriptors.append(
                computeParticleDescriptors(box, positions, minor_mode, seed, limit,
                                          average_over_neighbors, neighbor_average_rounds, nsplits))
        return np.hstack(single_mode_descriptors)

    kwargs = dict(params[mode])
    if mode == 'amean':
        descriptors = pythia.spherical_harmonics.abs_neighbor_average(
            box, positions, **kwargs)
    elif mode == 'voronoi_angle_histogram':
        descriptors = pythia.voronoi.angle_histogram(box, positions, **kwargs)
    elif mode == 'normalized_radial_distance':
        descriptors = pythia.bonds.normalized_radial_distance(box, positions, **kwargs)
    elif mode == 'neighborhood_distance_singvals':
        descriptors = pythia.bonds.neighborhood_distance_singvals(box, positions, **kwargs)
    elif mode == 'neighborhood_distance_sorted':
        descriptors = pythia.bonds.neighborhood_distance_sorted(box, positions, **kwargs)
    elif mode == 'neighborhood_range_distance_singvals':
        descriptors = pythia.bonds.neighborhood_range_distance_singvals(box, positions, **kwargs)
    elif mode == 'neighborhood_angle_singvals':
        descriptors = pythia.bonds.neighborhood_angle_singvals(box, positions, **kwargs)
    elif mode == 'neighborhood_range_angle_singvals':
        descriptors = pythia.bonds.neighborhood_range_angle_singvals(box, positions, **kwargs)
    elif mode == 'neighborhood_angle_sorted':
        descriptors = pythia.bonds.neighborhood_angle_sorted(box, positions, **kwargs)
    elif mode == 'steinhardt_q':
        descriptors = pythia.spherical_harmonics.steinhardt_q(box, positions, **kwargs)
    elif mode == 'bispectrum_sphs':
        descriptors = pythia.spherical_harmonics.bispectrum(box, positions, **kwargs)
    else:
        raise NotImplementedError('Unknown descriptor mode {}'.format(mode))

    if average_over_neighbors:
        neighbor_average_rounds = neighbor_average_rounds or 1

        fbox = freud.box.Box.from_box(box)
        nn = freud.locality.NearestNeighbors(2, average_over_neighbors)
        nn.compute(fbox, positions, positions)
        nlist = nn.nlist
        for _ in range(neighbor_average_rounds):
            bond_descriptors = descriptors[nlist.index_i]
            bond_descriptors += descriptors[nlist.index_j]
            descriptors = np.add.reduceat(bond_descriptors, nlist.segments)
            descriptors /= average_over_neighbors + 1

    if limit is not None:
        np.random.seed(seed)
        randidx = np.arange(len(descriptors))
        np.random.shuffle(randidx)
        descriptors = descriptors[randidx[:limit]]

    return descriptors

def globalDescriptors(fname, mode, limit=None):
    with gtar.GTAR(fname, 'r') as traj:
        ((boxRec, posRec), frames) = traj.framesWithRecordsNamed(['box', 'position'])
        positions = traj.getRecord(posRec, frames[-1])
        box = traj.getRecord(boxRec, frames[-1])

    box = FakeBox(*box)

    return computeParticleDescriptors(box, positions, mode, limit=limit)

def makeDataset(mode):
    train_inputs = []
    train_outputs = []

    for structure in track_progress(sorted_structures):
        for fname in structure_fnames[structure]:
            descriptors = globalDescriptors(fname, mode)
            train_inputs.append(descriptors)
            train_outputs.append(np.ones((len(descriptors),), dtype=np.uint32)*structure_indices[structure])

    train_inputs = np.concatenate(train_inputs, axis=0)
    train_outputs = np.concatenate(train_outputs)

    print(train_inputs.shape, train_outputs.shape)
    return (train_inputs, train_outputs)

datasets = {}
for mode in all_modes:
    datasets[mode] = makeDataset(mode)<jupyter_output><empty_output><jupyter_text>Many ML methods don't naively work very well with imbalanced datasets. This cell resamples the data to get roughly-equal numbers of each structure as training data.<jupyter_code>def balanceOutputs(train_inputs, train_outputs, plot=False):
    if plot:
        pp.hist(train_outputs, bins=len(sorted_structures), range=(-.5, len(sorted_structures) - .5))
        pp.title('Before balancing')

    hist = np.bincount(train_outputs, minlength=len(sorted_structures))
    min_count = np.min(hist)
    proba = (hist/np.sum(hist))[train_outputs]
    sortidx = np.argsort(np.random.uniform(0, 1, proba.shape)*proba)[:len(sorted_structures)*min_count]

    train_inputs = train_inputs[sortidx]
    train_outputs = train_outputs[sortidx]

    print('Data shape: {}'.format(train_inputs.shape))

    if plot:
        pp.figure()
        pp.hist(train_outputs, bins=len(sorted_structures), range=(-.5, len(sorted_structures) - .5))
        pp.title('After balancing')
        pp.show()

    return (train_inputs, train_outputs)

balanced_data = {}
for mode in all_modes:
    balanced_data[mode] = balanceOutputs(*datasets[mode])<jupyter_output>Data shape: (14000, 15)
Data shape: (14000, 256)
Data shape: (14000, 256)
Data shape: (14000, 32)
Data shape: (14000, 1521)
Data shape: (14000, 10)
Data shape: (14000, 350)<jupyter_text>We use a standard cross-validation scheme of withholding some fraction of the training data to ensure we haven't overfit.<jupyter_code>import numpy.random

def splitTrainingValidation(train_inputs, train_outputs, fraction=.75):
    idx = np.arange(train_inputs.shape[0])
    np.random.shuffle(idx)
    trainStop = int(fraction*train_inputs.shape[0])
    trainIdx = idx[:trainStop]
    valIdx = idx[trainStop:]

    trainIns = train_inputs[trainIdx]
    trainOuts = train_outputs[trainIdx]
    valIns = train_inputs[valIdx]
    valOuts = train_outputs[valIdx]
    print(trainOuts.shape)

    return (trainIns, trainOuts, valIns, valOuts)

splitData = {}
for mode in all_modes:
    splitData[mode] = splitTrainingValidation(*balanced_data[mode])<jupyter_output>(10500,)
(10500,)
(10500,)
(10500,)
(10500,)
(10500,)
(10500,)<jupyter_text>Here we generate a simple feedforward neural network for each set of descriptors we test. The networks have a single hidden layer.We import keras inside the model functions to make it easier to use the code with python's `multiprocessing` module, for example, where you'll need to avoid importing keras before forking.<jupyter_code>GLOBAL_EXTRA_PARAMS = dict(
    structure_names=sorted_structures, widths=[32],
    compile_kwargs=dict(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
)

def save_model(model, modelName, mode):
    with gtar.GTAR(modelName, 'w') as traj:
        for minor_mode in mode.split('|'):
            path = 'modes/{}/params.json'.format(minor_mode)
            traj.writeStr(path, json.dumps(params[minor_mode]))
        traj.writeStr('label_names.json', json.dumps(sorted_structures))
        traj.writeStr('extra_params.json', json.dumps(GLOBAL_EXTRA_PARAMS))
        traj.writeStr('mode.txt', mode)
        traj.writeStr('model.json', json.dumps(model.to_json()))

        for (i, weights) in enumerate(model.get_weights()):
            traj.writePath('frames/{}/weights.f32.uni'.format(i), weights)
            traj.writePath('frames/{}/weight_shape.u32.uni'.format(i), weights.shape)

def load_model(modelName):
    from tensorflow import keras
    from tensorflow.keras.models import model_from_json

    with gtar.GTAR(modelName, 'r') as traj:
        model_text = traj.readStr('model.json')

        weights = []
        for (layer_index, (weight_array, shape)) in traj.recordsNamed(['weights', 'weight_shape']):
            weights.append(weight_array.reshape(shape))

    model = model_from_json(json.loads(model_text))
    model.compile(**GLOBAL_EXTRA_PARAMS['compile_kwargs'])
    model.set_weights(weights)
    return model

def make_model(input_dim, num_classes, widths):
    from tensorflow import keras
    from tensorflow.keras.models import Sequential
    from tensorflow.keras.layers import Dense, Dropout, Activation, BatchNormalization

    model = Sequential()
    model.add(BatchNormalization(input_shape=(input_dim,)))
    for width in widths:
        model.add(Dense(width, activation='relu'))
        model.add(Dropout(.25))
    model.add(Dense(num_classes, activation='softmax'))

    model.compile(**GLOBAL_EXTRA_PARAMS['compile_kwargs'])

    return model

def getNamedModel(input_dim, num_classes, widths, mode, force_new=True):
    modelName = 'pythia_ann_{}.zip'.format(mode)

    if not os.path.exists(modelName) or force_new:
        model = make_model(input_dim, num_classes, widths)
        newModel = True
    else:
        model = load_model(modelName)
        newModel = False

    return model, modelName

modelData = {}
for mode in all_modes:
    (trainIns, trainOuts, valIns, valOuts) = splitData[mode]
    modelData[mode] = getNamedModel(
        trainIns.shape[-1], len(structure_fnames), GLOBAL_EXTRA_PARAMS['widths'], mode)
%%time

for mode in all_modes:
    (trainIns, trainOuts, valIns, valOuts) = splitData[mode]
    net_trainOuts = keras.utils.to_categorical(trainOuts, len(sorted_structures))
    (model, modelName) = modelData[mode]
    if True:
        train_hist = model.fit(trainIns, net_trainOuts, epochs=512, batch_size=4096, verbose=0,
                               validation_split=0.3,
            callbacks=[keras.callbacks.EarlyStopping('val_loss', patience=128, restore_best_weights=True)])
        save_model(model, modelName, mode)
        print('Finished training after {} epochs with accuracy {}'.format(
                train_hist.epoch[-1], train_hist.history['val_accuracy'][-1]))
import itertools

for mode in all_modes:
    (trainIns, trainOuts, valIns, valOuts) = splitData[mode]
    (model, modelName) = modelData[mode]

    predictions = model.predict_classes(valIns, verbose=False)

    confusion = {}
    for (true_class, pred_class) in itertools.product(range(len(sorted_structures)), range(len(sorted_structures))):
        match = np.logical_and(predictions == pred_class, valOuts == true_class)
        frac = float(np.sum(match))/np.sum(valOuts == true_class)
        confusion[(true_class, pred_class)] = frac

    matrix = []
    for row in range(len(sorted_structures)):
        rowvals = []
        for col in range(len(sorted_structures)):
            rowvals.append(confusion[(row, col)])
        matrix.append(rowvals)
    matrix = np.array(matrix, dtype=np.float32)

    # print(np.array_str(matrix, precision=3))

    def confusion_matrix(indices):
        pp.figure(figsize=(8, 8))

        compressed_indices = {idx: new for (new, idx) in enumerate(indices)}

        for (row, col) in itertools.product(indices, indices):
            pct = np.round(matrix[row, col]*1e2).astype(int)
            if pct > 100/len(sorted_structures):
                (row, col) = (compressed_indices[row], compressed_indices[col])
                pp.text(col, row, '{}%'.format(pct), ha='center', va='center', size=6)

        m = matrix[indices]
        m = m[:, indices]

        pp.imshow(m, extent=(-.5, len(indices) - .5, len(indices) - .5, -.5), origin='upper')
        pp.xticks(np.arange(len(indices)), [sorted_structures[idx] for idx in indices], rotation=40, ha='left')
        pp.yticks(np.arange(len(indices)), [sorted_structures[idx] for idx in indices])
        pp.gca().xaxis.tick_top(); pp.gca().xaxis.set_label_position('top')
        pp.xlabel('Predicted type')
        pp.ylabel('True type')
        pp.colorbar()
        pp.text(0.5, -.03, mode, transform=pp.gca().transAxes, size=18,
                horizontalalignment='center', verticalalignment='top')

    confusion_matrix(list(range(len(sorted_structures))))<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.10804229033024472, 'fraction_numerical': 0.011344737467331908, 'mean_word_length': 4.164723926380368, 'pattern_counts': {'":': 0, '<': 28, '<?xml version=': 0, '>': 29, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '122681', 'license': 'BSD-3-Clause', 'n_tokens_mistral': 5464, 'n_tokens_neox': 5172, 'n_words': 1191, 'path': 'examples/Supervised Learning.ipynb', 'repo_name': 'glotzerlab/pythia', 'size': '16836'}