<jupyter_start><jupyter_text>Reconstruction of structures from reconstructed momentsThis notebook takes a bunch of structure from the QM9 dataset, creates the corresponding invariants and tries to reconstruct the corresponding molecule by first finding peaks in reconstruction from moments<jupyter_code>import random
import json

import numpy as np
import pandas

import milad
from milad.play import asetools
from milad import invariants
from milad import reconstruct
from milad import zernike

import qm9_utils
random.seed(1234)
np.random.seed(1234)
from schnetpack import datasets

qm9data = datasets.QM9('data/qm9.db', download=True)
len(qm9data)<jupyter_output><empty_output><jupyter_text>The settings for our reconstruction<jupyter_code># Number of attempts to reconstruct per structure
num_attempts = 3

cutoff = 5

# Invariants to use
invs = invariants.read(invariants.COMPLEX_INVARIANTS)<jupyter_output><empty_output><jupyter_text>Let's load the subset of indices from the QM9 database we want to use<jupyter_code>with open('data/qm9_subset.json', 'r') as subset:
    test_set = json.load(subset)
test_set = {int(key): value for key, value in test_set.items()}
columns=('QM9 ID', 'Size', 'Attempt #', 'RMSD', 'Result')
def do_reconstruction(
    qm9data, 
    test_set,
    num_attempts: int,
    with_species: bool,
    species=None
):
    results = []
    
    # Let's create the descriptor we're going to use
    descriptor = qm9_utils.create_descriptor(invs, cutoff, apply_cutoff=False)
    # Create and re-use the grid query used when finding peaks in the reconstructed moments
    grid_query = zernike.ZernikeMoments.create_reconstruction_query(
        zernike.ZernikeMoments.get_grid(31), invs.max_order)
    
    for size, indices in sorted(test_set.items(), reverse=True):
        size = int(size)
        print(f"Size {size}:")
        for idx in indices:
            print(f"Idx={idx}:", end='')

            # Get the system
            system = qm9data.get_atoms(idx=idx)
            max_radius = asetools.prepare_molecule(system) # Centre the molecule and find the max radius.

            
            # Prepare the milad data type
            milad_molecule = asetools.ase2milad(system)
            num_atoms = milad_molecule.num_atoms
            
            if not with_species:
                # Species all fixed
                milad_molecule.numbers = 1.
            
            # Calculate the ground truth
            fingerprint = descriptor(milad_molecule)

            for attempt in range(num_attempts):
                print(f"\t", end='')
                
                initial = qm9_utils.create_initial_atoms(num_atoms, max_radius, include_species=with_species)
                

                # Finally, optimise the structure wrt to the fingerprint
                result = reconstruct.find_iteratively(
                    descriptor,
                    fingerprint, # Target the fingerprint
                    initial,
                    find_species=with_species,
                    grid_query=grid_query,
                    minsep=0.55,
                    verbose=False,
                )
                
                print(f'{result.rmsd:5.5}', end='')
                
                results.append((idx, size, attempt, result.rmsd, result))

            print()
        
    return results
results = do_reconstruction(
    qm9data, 
    test_set,
    num_attempts, 
    with_species=True,
    species=qm9_utils.species,
)
with_species_frame = pandas.DataFrame(results, columns=columns)
with_species_frame.to_pickle('structure_recovery_iterative_with_species.pickle')
results = do_reconstruction(
    qm9data, 
    test_set,
    num_attempts,
    with_species=False,
)
no_species_frame = pandas.DataFrame(results, columns=columns)
no_species_frame.to_pickle('structure_recovery_iterative_no_species.pickle')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.07712016397642839, 'fraction_numerical': 0.008967460927491674, 'mean_word_length': 2.711026615969582, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '308963', 'license': 'MIT', 'n_tokens_mistral': 1226, 'n_tokens_neox': 1132, 'n_words': 282, 'path': 'notebooks/Reconstruct structures iteratively.ipynb', 'repo_name': 'muhrin/milad-paper-2021', 'size': '3903'}