From 953e47b01a130224cac56911ac4405637fc815ad Mon Sep 17 00:00:00 2001 From: Evan Komp Date: Sun, 25 Aug 2024 10:43:38 -0700 Subject: [PATCH] sampling to test script --- scripts/fireprot_zero.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/scripts/fireprot_zero.py b/scripts/fireprot_zero.py index b5d7a1d..c42fc15 100644 --- a/scripts/fireprot_zero.py +++ b/scripts/fireprot_zero.py @@ -106,13 +106,16 @@ def apply_mutation(sequence: str, mutation: str) -> str: def evaluate_model(model: NOMELTModel, wt_df: pd.DataFrame, variant_df: pd.DataFrame) -> dict: results = [] + counter = 0 + for _, wt_row in tqdm(wt_df.iterrows(), total=len(wt_df), desc="Evaluating proteins"): wt_sequence = wt_row['sequence'] variants = variant_df[variant_df['uniprot_id'] == wt_row['uniprot_id']] - + if counter > 3: + break if len(variants) == 0: continue - + counter +=1 mutated_sequences = [apply_mutation(wt_sequence, mut) for mut in variants['mutation']] wt_score, variant_scores = model.score_variants(wt_sequence, mutated_sequences, batch_size=5, indels=False)