def test_reward():
env = Molecule()
assert env._state is None
assert env.reward() == 0
env = Molecule(init_mol='C')
env.reward()
def test_serialization():
moldqn = DQNFinalState(Molecule(), MorganFingerprints(), batch_size=8)
res = pkl.dumps(moldqn)
moldqn_2 = pkl.loads(res)
assert np.isclose(moldqn.action_network.get_weights()[0],
moldqn_2.action_network.get_weights()[0]).all()
# Test after training
moldqn.env.reset()
for i in range(10):
action, _, _ = moldqn.action()
new_state, reward, done, _ = moldqn.env.step(action)
# Save outcome
moldqn.remember(moldqn.env.state, action, reward,
new_state, moldqn.env.action_space.get_possible_actions(), done)
# Make a training step
moldqn.train()
# Test out the serialization
res = pkl.dumps(moldqn)
moldqn_2 = pkl.loads(res)
assert len(moldqn_2.optimizer.get_weights()) == 0
} for r in ['ic50', 'QED', 'SA', 'cycles']])
elif args.reward == "oneshot":
reward = rewards['oneshot']
elif args.reward == "tuned":
reward = LogisticCombination(rewards['ic50'], rewards['oneshot'])
else:
raise ValueError(f'Reward function not defined: {args.reward}')
run_params['maximize'] = reward.maximize
# Set up environment
action_space = MoleculeActions(elements,
allow_removal=not args.no_backtrack)
init_mol = args.initial_molecule
if init_mol is not None:
init_mol = convert_smiles_to_nx(init_mol)
env = Molecule(action_space, reward=reward, init_mol=init_mol)
logger.debug('using environment: %s' % env)
# Load in initial molecules, if defined
init_mols = None
if args.initial_molecule_file is not None:
with open(args.initial_molecule_file) as fp:
init_mols = [convert_smiles_to_nx(x) for x in json.load(fp)]
logger.info(f'Read in {len(init_mols)} seed molecules')
# Make sure the molecules only contain elements from the list
def screen(g):
atomic_nums = nx.get_node_attributes(g, 'atomic_num').values()
return all(
pt.GetElementSymbol(z) in elements for z in atomic_nums
if z > 1)
def run(self):
# Output files
output_files = {
'simulation':
open(os.path.join(self.output_dir, 'simulation_records.jsonld'),
'w')
}
# Get the reference energies
# TODO (wardlt): I write many of the "send out and wait" patterns, should I build a utility
# or is this a fundamental issue?
self.logger.info(f'Starting Thinker process')
elems = ['H', 'C', 'N', 'O', 'F']
for elem in elems:
self.queues.send_inputs(elem,
spec,
multiplicity[elem],
method='compute_reference_energy')
ref_energies = {}
for _ in elems:
result = self.queues.get_result()
ref_energies[result.args[0]] = result.value
self.logger.info(
f'Computed reference energies for: {", ".join(ref_energies.keys())}'
)
# Run the initial molecules
for mol in self.initial_molecules:
self.queues.send_inputs(mol,
spec,
ref_energies,
compute_config,
method='compute_atomization_energy')
for _ in self.initial_molecules:
result = self.queues.get_result()
if result.success:
self.database[result.args[0]] = result.value
else:
logging.warning(
'Calculation failed! See simulation outputs and Parsl log file'
)
print(result.json(), file=output_files['simulation'])
output_files['simulation'].flush()
self.logger.info(
f'Computed initial population of {len(self.database)} molecules')
# Make the initial RL agent
agent = DQNFinalState(
Molecule(reward=SklearnReward(None, maximize=False)),
MorganFingerprints(),
q_network_dense=(1024, 512, 256),
epsilon_decay=0.995)
for i in range(self.n_evals // self.n_parallel):
self.logger.info(f'Starting design loop step {i}')
# Train the machine learning model
gf = GroupFeaturizer()
model = Pipeline([('group', gf),
('lasso', BayesianRidge(normalize=True))])
mols, atoms = zip(*self.database.items())
model.fit(
mols,
atoms) # negative so that the RL optimizes a positive value
self.logger.info(
f'Fit a model with {len(mols)} training points and {len(gf.known_groups_)} groups'
)
# Use RL to generate new molecules
agent.env.reward_fn.model = model
self.queues.send_inputs(agent, method='generate_molecules')
result = self.queues.get_result()
new_molecules, agent = result.value
self.logger.info(
f'Generated {len(new_molecules)} candidate molecules')
# Assign them scores
self.queues.send_inputs(model,
new_molecules,
method='compute_score')
result = self.queues.get_result()
scores = result.value
self.logger.info(f'Assigned scores to all molecules')
# Pick a set of calculations to run
# Greedy selection for now
task_options = [{
'smiles': s,
'pred_atom': e
} for s, e in zip(new_molecules, scores)]
selections = greedy_selection(task_options, self.n_parallel,
lambda x: -x['pred_atom'])
self.logger.info(f'Selected {len(selections)} new molecules')
# Run the selected simulations
for task in selections:
self.queues.send_inputs(task['smiles'],
spec,
ref_energies,
method='compute_atomization_energy')
# Wait for them to return
for _ in selections:
result = self.queues.get_result()
if result.success:
self.database[result.args[0]] = result.value
else:
logging.warning(
'Calculation failed! See simulation outputs and Parsl log file'
)
print(result.json(), file=output_files['simulation'])
output_files['simulation'].flush()
def test_moldqn():
agent = DQNFinalState(
Molecule(reward=SklearnReward(FakeEstimator().predict)),
MorganFingerprints())
output = generate_molecules(agent)
assert 'C' in output # There is a (0.25)^10 chance that this will fail
def run(self):
# Get the reference energies
# TODO (wardlt): I write many of the "send out and wait" patterns, should I build a utility
# or is this a fundamental issue?
self.logger.info(f'Starting Thinker process')
elems = ['H', 'C', 'N', 'O', 'F']
for elem in elems:
self.queues.send_inputs(elem,
spec,
multiplicity[elem],
method='compute_reference_energy')
ref_energies = {}
for _ in elems:
result = self.queues.get_result()
ref_energies[result.args[0]] = result.value
self._ref_energies = ref_energies
self.logger.info(
f'Computed reference energies for: {", ".join(ref_energies.keys())}'
)
# Run the initial molecules
for mol in self.initial_molecules:
self.queues.send_inputs(mol,
spec,
ref_energies,
compute_config,
method='compute_atomization_energy')
for _ in self.initial_molecules:
result = self.queues.get_result()
if result.success:
self.database[result.args[0]] = result.value
else:
logging.warning(
'Calculation failed! See simulation outputs and Parsl log file'
)
print(result.json(), file=self.output_file)
self.output_file.flush()
self.logger.info(
f'Computed initial population of {len(self.database)} molecules')
# Make the initial RL agent
agent = DQNFinalState(
Molecule(reward=SklearnReward(None, maximize=False)),
MorganFingerprints(),
q_network_dense=(1024, 512, 256),
epsilon_decay=0.995)
# Launch the "simulator" thread
design_thread = Thread(target=self.simulation_dispatcher)
design_thread.start()
# Perform the design loop iteratively
step_number = 0
while len(self.database) < self.n_evals:
self.logger.info(f'Generating new molecules')
# Train the machine learning model
gf = GroupFeaturizer()
model = Pipeline([('group', gf),
('lasso', BayesianRidge(normalize=True))])
database_copy = self.database.copy()
mols, atoms = zip(*database_copy.items())
model.fit(mols, atoms)
self.logger.info(
f'Fit a model with {len(mols)} training points and {len(gf.known_groups_)} groups'
)
# Use RL to generate new molecules
agent.env.reward_fn.model = model
self.queues.send_inputs(agent,
method='generate_molecules',
topic='ML')
result = self.queues.get_result(topic='ML')
new_molecules, agent = result.value
self.logger.info(
f'Generated {len(new_molecules)} candidate molecules')
print(result.json(exclude={'inputs', 'value'}),
file=self.rl_records)
self.rl_records.flush()
# Assign them scores
self.queues.send_inputs(model,
new_molecules,
method='compute_score',
topic='ML')
result = self.queues.get_result(topic='ML')
scores = result.value
print(result.json(exclude={'inputs'}), file=self.screen_records)
self.screen_records.flush()
self.logger.info(f'Assigned scores to all molecules')
# Pick a set of calculations to run
# Greedy selection for now
task_options = [{
'smiles': s,
'pred_atom': e
} for s, e in zip(new_molecules, scores)]
selections = greedy_selection(task_options, self.n_parallel,
lambda x: -x['pred_atom'])
self.logger.info(f'Selected {len(selections)} new molecules')
# Add requested simulations to the queue
for rank, task in enumerate(selections):
self._task_queue.put(
((-step_number, rank),
task['smiles'])) # Sort by recency and then by best
step_number += 1 # Increment the loop
self.logger.info('No longer generating new candidates')
self._gen_done.set()