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
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)
init_mols = [g for g in init_mols if screen(g)]
logger.info(
f'{len(init_mols)} contain only elements from action space')
# Setup agent
agent = DQNFinalState(env,
gamma=args.gamma,
preprocessor=MorganFingerprints(
args.fingerprint_size),
batch_size=args.batch_size,
epsilon=args.epsilon,
memory_size=args.memory_size,
q_network_dense=args.hidden_layers,
epsilon_decay=args.epsilon_decay)
# Make a test directory
test_dir = os.path.join(
'rl_tests', f'{run_params["reward"]}_' +
datetime.now().isoformat().replace(":", "."))
if not os.path.isdir(test_dir):
os.makedirs(test_dir)
# Write the test parameters to the test directory
with open(os.path.join(test_dir, 'config.json'), 'w') as fp:
json.dump(run_params, fp)
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 generate_molecules(
agent: DQNFinalState,
episodes: int = 10,
n_steps: int = 32,
update_q_every: int = 10) -> Tuple[Set[str], DQNFinalState]:
"""Perform the RL experiment
Args:
agent (DQNFinalState): Molecular design agent
episodes (int): Number of episodes to run
n_steps (int): Maximum number of steps per episode
update_q_every (int): After how many updates to update the Q function
Returns:
([str]) List of molecules that were created
"""
# Prepare the output
output = set()
# Keep track of the smiles strings
for e in range(episodes):
current_state = agent.env.reset()
logger.info(f'Starting episode {e+1}/{episodes}')
for s in range(n_steps):
# Get action based on current state
action, _, _ = agent.action()
# Fix cluster action
new_state, reward, done, _ = agent.env.step(action)
# Check if it's the last step and flag as done
if s == n_steps:
logger.debug('Last step ... done')
done = True
# Add the state to the output
output.add(agent.env.state)
# Save outcome
agent.remember(current_state, action, reward, new_state,
agent.env.action_space.get_possible_actions(), done)
# Train model
agent.train()
# Update state
current_state = new_state
if done:
break
# Update the Q network after certain numbers of episodes and adjust epsilon
if e > 0 and e % update_q_every == 0:
agent.update_target_q_network()
agent.epsilon_adj()
# Clear out the memory: Too large to send back to client
agent.memory.clear()
# Convert the outputs back to SMILES strings
output = set(convert_nx_to_smiles(x) for x in output)
return output, agent
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()