def propose_sequences(
self, measured_sequences: pd.DataFrame
) -> Tuple[np.ndarray, np.ndarray]:
"""Propose top `sequences_batch_size` sequences for evaluation."""
if self.num_actions == 0:
# indicates model was reset
self.initialize_data_structures()
else:
# set state to best measured sequence from prior batch
last_round_num = measured_sequences["round"].max()
last_batch = measured_sequences[
measured_sequences["round"] == last_round_num
]
_last_batch_seqs = last_batch["sequence"].tolist()
_last_batch_true_scores = last_batch["true_score"].tolist()
last_batch_seqs = _last_batch_seqs
if self.recomb_rate > 0 and len(last_batch) > 1:
last_batch_seqs = self._recombine_population(last_batch_seqs)
measured_batch = []
for seq in last_batch_seqs:
if seq in _last_batch_seqs:
measured_batch.append(
(_last_batch_true_scores[_last_batch_seqs.index(seq)], seq)
)
else:
measured_batch.append((np.mean(self.model.get_fitness([seq])), seq))
measured_batch = sorted(measured_batch)
sampled_seq = self.Thompson_sample(measured_batch)
self.state = string_to_one_hot(sampled_seq, self.alphabet)
# generate next batch by picking actions
self.initial_uncertainty = None
samples = set()
prev_cost = self.model.cost
all_measured_seqs = set(measured_sequences["sequence"].tolist())
while self.model.cost - prev_cost < self.model_queries_per_batch:
uncertainty, new_state_string, _ = self.pick_action(all_measured_seqs)
all_measured_seqs.add(new_state_string)
samples.add(new_state_string)
if self.initial_uncertainty is None:
self.initial_uncertainty = uncertainty
if uncertainty > 2 * self.initial_uncertainty:
# reset sequence to starting sequence if we're in territory that's too
# uncharted
sampled_seq = self.Thompson_sample(measured_batch)
self.state = string_to_one_hot(sampled_seq, self.alphabet)
self.initial_uncertainty = None
if len(samples) < self.sequences_batch_size:
random_sequences = generate_random_sequences(
self.seq_len, self.sequences_batch_size - len(samples), self.alphabet
)
samples.update(random_sequences)
# get predicted fitnesses of samples
samples = list(samples)
preds = np.mean(self.model.get_fitness(samples), axis=1)
# train ensemble model before returning samples
self.train_models()
return samples, preds
def test_rna():
# Since ViennaRNA is an optional dependency, only test if installed
try:
problem = flexs.landscapes.rna.registry()["C20_L100_RNA1+2"]
landscape = flexs.landscapes.RNABinding(**problem["params"])
test_seqs = s_utils.generate_random_sequences(100, 100, s_utils.RNAA)
landscape.get_fitness(test_seqs)
except ImportError:
warnings.warn(
"Skipping RNABinding landscape test since" "ViennaRNA not installed."
)
def test_rosetta():
# Since PyRosetta is an optional dependency, only test if installed
try:
problem = flexs.landscapes.rosetta.registry()["3msi"]
landscape = flexs.landscapes.RosettaFolding(**problem["params"])
seq_length = len(landscape.wt_pose.sequence())
test_seqs = s_utils.generate_random_sequences(seq_length, 100, s_utils.AAS)
landscape.get_fitness(test_seqs)
except ImportError:
warnings.warn(
"Skipping RosettaFolding landscape test since PyRosetta not installed."
)
def test_additive_aav_packaging():
problem = flexs.landscapes.additive_aav_packaging.registry()["heart"]
landscape = flexs.landscapes.AdditiveAAVPackaging(**problem["params"])
test_seqs = s_utils.generate_random_sequences(90, 100, s_utils.AAS)
landscape.get_fitness(test_seqs)
def test_tf_binding():
problem = flexs.landscapes.tf_binding.registry()["SIX6_REF_R1"]
landscape = flexs.landscapes.TFBinding(**problem["params"])
test_seqs = s_utils.generate_random_sequences(8, 100, s_utils.DNAA)
landscape.get_fitness(test_seqs)
def __init__(
self,
landscape: flexs.Landscape,
rounds: int,
sequences_batch_size: int,
model_queries_per_batch: int,
starting_sequence: str,
alphabet: str,
log_file: Optional[str] = None,
model: Optional[flexs.Model] = None,
num_experiment_rounds: int = 10,
num_model_rounds: int = 1,
):
"""
Args:
num_experiment_rounds: Number of experiment-based rounds to run. This is by
default set to 10, the same number of sequence proposal of rounds run.
num_model_rounds: Number of model-based rounds to run.
"""
tf.config.run_functions_eagerly(False)
name = f"DynaPPO_Agent_{num_experiment_rounds}_{num_model_rounds}"
if model is None:
model = DynaPPOEnsemble(
len(starting_sequence),
alphabet,
)
model.train(
s_utils.generate_random_sequences(len(starting_sequence), 10,
alphabet),
[0] * 10,
)
super().__init__(
model,
name,
rounds,
sequences_batch_size,
model_queries_per_batch,
starting_sequence,
log_file,
)
self.alphabet = alphabet
self.num_experiment_rounds = num_experiment_rounds
self.num_model_rounds = num_model_rounds
env = DynaPPOEnvMut(
alphabet=self.alphabet,
starting_seq=starting_sequence,
model=model,
landscape=landscape,
max_num_steps=model_queries_per_batch,
)
validate_py_environment(env, episodes=1)
self.tf_env = tf_py_environment.TFPyEnvironment(env)
encoder_layer = tf.keras.layers.Lambda(lambda obs: obs["sequence"])
actor_net = actor_distribution_network.ActorDistributionNetwork(
self.tf_env.observation_spec(),
self.tf_env.action_spec(),
preprocessing_combiner=encoder_layer,
fc_layer_params=[128],
)
value_net = value_network.ValueNetwork(
self.tf_env.observation_spec(),
preprocessing_combiner=encoder_layer,
fc_layer_params=[128],
)
self.agent = ppo_agent.PPOAgent(
self.tf_env.time_step_spec(),
self.tf_env.action_spec(),
optimizer=tf.keras.optimizers.Adam(learning_rate=1e-5),
actor_net=actor_net,
value_net=value_net,
num_epochs=10,
summarize_grads_and_vars=False,
)
self.agent.initialize()
def __init__(
self,
landscape: flexs.Landscape,
rounds: int,
sequences_batch_size: int,
model_queries_per_batch: int,
starting_sequence: str,
alphabet: str,
log_file: Optional[str] = None,
model: Optional[flexs.Model] = None,
num_experiment_rounds: int = 10,
num_model_rounds: int = 1,
env_batch_size: int = 4,
):
"""
Args:
num_experiment_rounds: Number of experiment-based rounds to run. This is by
default set to 10, the same number of sequence proposal of rounds run.
num_model_rounds: Number of model-based rounds to run.
env_batch_size: Number of epsisodes to batch together and run in parallel.
"""
tf.config.run_functions_eagerly(False)
name = f"DynaPPO_Agent_{num_experiment_rounds}_{num_model_rounds}"
if model is None:
model = DynaPPOEnsemble(
len(starting_sequence),
alphabet,
)
# Some models in the ensemble need to be trained on dummy dataset before
# they can predict
model.train(
s_utils.generate_random_sequences(len(starting_sequence), 10,
alphabet),
[0] * 10,
)
super().__init__(
model,
name,
rounds,
sequences_batch_size,
model_queries_per_batch,
starting_sequence,
log_file,
)
self.alphabet = alphabet
self.num_experiment_rounds = num_experiment_rounds
self.num_model_rounds = num_model_rounds
self.env_batch_size = env_batch_size
env = DynaPPOEnv(self.alphabet, len(starting_sequence), model,
landscape, env_batch_size)
self.tf_env = tf_py_environment.TFPyEnvironment(env)
actor_net = actor_distribution_network.ActorDistributionNetwork(
self.tf_env.observation_spec(),
self.tf_env.action_spec(),
fc_layer_params=[128],
)
value_net = value_network.ValueNetwork(self.tf_env.observation_spec(),
fc_layer_params=[128])
print(self.tf_env.action_spec())
self.agent = ppo_agent.PPOAgent(
time_step_spec=self.tf_env.time_step_spec(),
action_spec=self.tf_env.action_spec(),
optimizer=tf.keras.optimizers.Adam(learning_rate=1e-5),
actor_net=actor_net,
value_net=value_net,
num_epochs=10,
summarize_grads_and_vars=False,
)
self.agent.initialize()