"""

Generates next generation of population by probabilistically choosing mating pool based on fitness, then

probabilistically reproducing molecules in the mating pool and randomly mutating the children

:param population: list of RDKit molecules

:param fitness: probability distribution of same length as population, returned by pop_fitness

:param mutation_rate: hyperparameter determining the likelihood of mutations occuring

:return: new_population

"""

mating_pool = []

for i in range(len(population)):

mating_pool.append(np.random.choice(population, p=fitness))

new_population = []

for n in range(len(population)):

parent_A = random.choice(mating_pool)

parent_B = random.choice(mating_pool)

# print Chem.MolToSmiles(parent_A),Chem.MolToSmiles(parent_B)

new_child = co.crossover(parent_A, parent_B)

# print new_child

if new_child is not None:

new_child = mu.mutate(new_child, mutation_rate)

# print("after mutation",new_child)

if new_child is not None:

new_population.append(new_child)

"""

Calculates the fitness (ie SOAP similarity score) of the population by generating conformers for each of the

population molecules, then evaluating their SOAP descriptors and calculating its similarity score with the SOAP

descriptor of the binding ligand 'field'

Conformer generation and similarity calculation are the computational bottlenecks - might be worth splitting the

task up with MPI. see return_borders.py in helper.py if you want to do that - make sure you only run the

reproduction on the master node (since there is randomness), then broadcast to the other nodes

:param population: list of RDKit molecule objects

:param tgt_atoms: list of ASE atom objects of the target ligand field - from read_xyz, second is optional if separate sites

:param tgt_species: list of the atomic species present in the target ligand field - from read_xyz

:param rcut, sigma: SOAP parameters

:param max_score: Maximum SOAP similarity found so far

:return: fitness, max_score, fit_mean, fit_std

"""

t0 = time.time()

# loop over RDKit mols and turn them into lists of ASE atom objects for dscribe SOAP atomic feature generation

population_ase = []

num_list = []

species = ['C']

bad_mols = []

for m in population:

m = Chem.AddHs(m)

conf_result = AllChem.EmbedMolecule(m, maxAttempts=1000)

if conf_result != 0:

bad_mols.append(m)

continue

m = Chem.RemoveHs(m)

num_list.append(len(m.GetAtoms()))

for i, atom in enumerate(m.GetAtoms()):

symbol = atom.GetSymbol()

conf = m.GetConformer()

population_ase.append(Atoms(symbol, [conf.GetPositions()[i]]))

if symbol not in species: # find unique atomic species for SOAP generation

species.append(symbol)

if bad_mols != []:

for bm in bad_mols:

bm = Chem.RemoveHs(bm)

try:

population.remove(bm) # filter out molecules which have no conformers

except:

continue

# Check that we also include the atom types present in the ligand targets

for atom in tgt_species:

if atom not in species:

species.append(atom)

t1 = time.time()

print('Time taken to generate conformers: {}'.format(t1-t0))

# Generate SOAP descriptors using dscribe

soap_generator = SOAP(species=species, periodic=False, rcut=rcut, nmax=8, lmax=6, sigma=sigma, sparse=True)

soap = soap_generator.create(population_ase)

tgt_soap = soap_generator.create(tgt_atoms)

if tgt_atoms2 is not None:

tgt_soap2 = [normalize(soap_generator.create(tgt_atoms2), copy=False)]

# normalize SOAP atom descriptors and group by molecule

soap = normalize(soap, copy=False)

tgt_soap = [normalize(tgt_soap, copy=False)]

soap = split_by_lengths(soap, num_list)

t2 = time.time()

print('Time taken to generate SOAP descriptors: {}'.format(t2-t1))

# TODO make REMatch kernel args as input args

if kernel == 'rematch':

soap_similarity = REMatchKernel(metric="polynomial", degree=3, gamma=1, coef0=0, alpha=0.1, threshold=1e-3, normalize_kernel=True)

elif kernel == 'average':

soap_similarity = AverageKernel(metric="polynomial", degree=3, gamma=1, coef0=0, normalize_kernel=True)

if tgt_atoms2 is not None:

fitness1 = soap_similarity.create(soap, tgt_soap)

fitness1.flatten()

fitness2 = soap_similarity.create(soap, tgt_soap2)

fitness2.flatten()

# calculate fitness score as product of the two fitnesses

fitness = np.multiply(fitness1, fitness2)

fitness = np.array([f[0] for f in fitness])

else:

fitness = soap_similarity.create(soap, tgt_soap)

fitness = fitness.flatten()

t3 = time.time()

print('Time taken to calculate fitness: {}'.format(t3-t2))

# update max_score, include new champion

if np.amax(fitness) > max_score[0]:

max_score = [np.amax(fitness), Chem.MolToSmiles(population[np.argmax(fitness)])]

#Print the top 5 scores and corresponding molecules for a particular generation

top_scores = np.flip(fitness[np.argsort(fitness)[-5:]])

# print(top_scores)

for i in range(5):

print("Mol {}: {} (fitness = {:.3f})".format(i, Chem.MolToSmiles(population[np.argsort(fitness)[-i-1]]), top_scores[i]))

fitness = fitness / np.sum(fitness)

"""

Reads in a .csv file and generates a population of RDKit molecules, as well as reading in target ligand coordinates

:param args: system arguments parsed into main - should contain args.csv, args.tgt (and args.tgt2)

:return: population, tgt_atoms, tg_species

or population, tgt_atoms, tgt_species, tgt_atoms2, tgt_species2

"""

population = []

csv = pd.read_csv(args.csv, header=0)

for i, row in csv.iterrows():

population.append(Chem.MolFromSmiles(row['SMILES']))

tgt_atoms, _, _, tgt_species = read_xyz(args.tgt)

if args.tgt2 is not None:

tgt_atoms2, _, _, tgt_species2 = read_xyz(args.tgt2)

tgt_species = list(set().union(tgt_species, tgt_species2)) # creates a single tgt_species list

return population, tgt_atoms, tgt_species, tgt_atoms2

else:

"""

Runs genetic algorithm - initialise_system generates the starting population. Program then loops over generations,

calculates the fitness of the population, prints+saves the highest fitness individual ('champion'),

and then updates the population based on the fitness.

:param args: system arguments parsed into program

:return:

"""

t0 = time.time()

co.average_size = args.tgt_size # read what this does

co.size_stdev = args.size_stdev

if args. tgt2 is not None:

population, tgt_atoms, tgt_species, tgt_atoms2 = initialise_system(args)

else:

population, tgt_atoms, tgt_species = initialise_system(args)

print('\nInitial Population Size: {}'.format(len(population)))

print('No. of generations: {}'.format(args.n_gens))

print('Mutation rate: {}'.format(args.mut_rate))

print('')

max_score = [-999, '']

f = open('champions.dat', 'w')

for generation in range(args.n_gens):

print('\nGeneration #{}, population size: {}'.format(generation, len(population)))

print('Calculating fitness...')

if args.tgt2 is not None:

fitness, max_score = pop_fitness(population, args.rcut, args.sigma, args.kernel, tgt_atoms, tgt_species, tgt_atoms2, max_score)

else:

fitness, max_score = pop_fitness(population, args.rcut, args.sigma, args.kernel, tgt_atoms, tgt_species, None, max_score)

print('Producing next generation...')

population = reproduce(population, fitness, args.mut_rate)

# Think you might want to print out the best-k individuals from each generation - wil leave that to you

print('Champion fitness = {}, smiles = {}'.format(max_score[0], max_score[1]))

f.write(max_score[1] + '\t' + str(max_score[0]) + '\n')

f.flush()

t1 = time.time()