# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# ## Recovery of species from decoded positions

# +
import random
import json

from ase import visualize
import numpy as np
import pandas

import milad
from milad import atomic
from milad import generate
from milad.play import asetools
from milad import invariants

import qm9_utils
# -

# Set the seeds for some kind of reproducibility
random.seed(1234)
np.random.seed(1234)

# +
from schnetpack import datasets

qm9data = datasets.QM9('data/qm9.db', download=True)
len(qm9data)
# -

# The settings for our reconstruction

# Number of attempts to reconstruct per structure
num_attempts = 3
# The multiple of the maximum radius to make the cutoff
cutoff = 5
# Optimiser to use
optimiser = milad.optimisers.StructureOptimiser() 
# Keep atoms from coming too close
optimiser.separation_force = atomic.SeparationForce(epsilon=1e-1, cutoff=0.55, power=2)
# Invariants to use
invs = invariants.read(invariants.COMPLEX_INVARIANTS) 

# Let's load the subset of indices from the QM9 database we want to use

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()}


def random_atom_collection(num_atoms: int, radius: float, numbers=None):
    atoms = milad.atomic.AtomsCollection(
        num_atoms,
        positions=generate.random_points_in_sphere(num_atoms, radius, centre=True))
    if numbers is not None:
        atoms.numbers = numbers
    
    return atoms


# +
columns=('QM9 ID', 'Size', 'Attempt #', 'RMSD', 'Result')

# Load results from reconstruction of positions only
no_species = pandas.read_pickle('structure_recovery_iterative_no_species.pickle')
# -

subframe = no_species.loc[(no_species['QM9 ID'] == 2)]
subframe.loc[subframe['RMSD'] == subframe['RMSD'].min()].iloc[0]['Result']


# + pycharm={"name": "#%%\n"}
def do_reconstruction(qm9data, test_set, result_set, num_attempts, cutoff: float):
    results = []
    
    # Let's create the descriptor we're going to use
    descriptor = qm9_utils.create_descriptor(invs, cutoff=cutoff, apply_cutoff=False)


    for size, indices in sorted(test_set.items()):
        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.
            assert max_radius < cutoff

            # Prepare the milad data type
            milad_molecule = asetools.ase2milad(system)
            num_atoms = milad_molecule.num_atoms
            
            fingerprint = descriptor(milad_molecule)
            
            # Get the result with the lowest RMSD
            subframe = result_set.loc[(result_set['QM9 ID'] == idx)]
#             result = subframe.loc[subframe['RMSD'] == subframe['RMSD'].min()].iloc[0]['Result']
#             reconstructed = result.value
            
            num_attempts = len(subframe)

            for attempt in range(num_attempts):
                print(f"...", end='')
                
                result = subframe.loc[subframe['Attempt #'] == attempt].iloc[0]['Result']
                reconstructed = result.value
                initial = reconstructed

                # Fix the positions, we only want to get the species
                mask = reconstructed.get_mask()
                mask.positions = reconstructed.positions

                result = optimiser.optimise(
                    descriptor=descriptor,
                    # Target the original fingerprint
                    target=fingerprint,
                    initial=initial,
                    mask=mask,
                )
                
                mask = result.value.get_mask()
                mask.numbers = result.value.numbers
                
                result = optimiser.optimise(
                    descriptor=descriptor,
                    # Target the original fingerprint
                    target=fingerprint,
                    initial=result.value,
                    mask=mask,
                )
                
                print(f'{result.rmsd:5.5}', end='')
                results.append((idx, size, attempt, result.rmsd, result))

            print()
        
    return results
# -

results = do_reconstruction(
    qm9data, 
    test_set,
    no_species,
    num_attempts, 
    cutoff=cutoff,
)
no_species_frame = pandas.DataFrame(results, columns=columns)

no_species_frame.to_pickle('data/species_recovery_from_decoded_positions.pickle')

StarcoderdataJupyterScriptsDedupFiltered