<jupyter_start><jupyter_text>Load starting pool of RFP variantsTest a single mutation<jupyter_code>rfp_known_structures = pd.read_csv('../lambo/assets/fpbase/rfp_known_structures.csv')
print('valid protein names')
print(rfp_known_structures.Name)
protein_name = 'DsRed.M1'
pdb_id = rfp_known_structures[rfp_known_structures.Name == protein_name].pdb_id.item().lower()
chain_id = rfp_known_structures[rfp_known_structures.Name == protein_name].longest_chain.item()

work_dir = './tmp'
parent_pdb_path = f'../lambo/assets/foldx/{pdb_id}_{chain_id}/wt_input_Repair.pdb'
tokenizer = ResidueTokenizer()
base_candidate = FoldedCandidate(work_dir, parent_pdb_path, [], 
                                   tokenizer=tokenizer, skip_minimization=True,
                                   chain=chain_id, wild_name=protein_name)

# substitute G at the first position
mutation_list = [base_candidate.new_mutation(seq_idx=0, mutant_residue='G', mutation_type='sub')]
# run FoldX to get new structure
new_candidate = base_candidate.new_candidate(mutation_list)

print(f'---- wild type {protein_name}-{chain_id} ----\n')
print(f'residues: {new_candidate.wild_residue_seq}\n')
print(f'stability: {-new_candidate.wild_total_energy}\n')
print(f'surface area: {new_candidate.wild_surface_area}\n')

print(f'---- mutant {new_candidate.uuid} ----\n')
print(f'residues: {new_candidate.mutant_residue_seq}\n')
print(f'stability: {-new_candidate.mutant_total_energy}\n')
print(f'surface area: {new_candidate.mutant_surface_area}\n')<jupyter_output>---- wild type DsRed.M1-A ----

residues: TEDVIKEFMQFKVRMEGSVNGHYFEIEGEGEGKPYEGTQTAKLQVTKGGPLPFAWDILSPQFSKAYVKHPADIPDYMKLSFPEGFTWERSMNFEDGGVVEVQQDSSLQDGTFIYKVKFKGVNFPADGPVMQKKTAGWEPSTEKLYPQDGVLKGEISHALKLKDGGHYTCDFKTVYKAKKPVQLPGNHYVDSKLDITNHNEDYTVVEQYEHAEARHSG

stability: 71.485

surface area: 10385.128377983554

---- mutant 7d013b57203f414fa78798afdf927a15 ----

residues: GEDVIKEFMQFKVRMEGSVNGHYFEIEGEGEGKPYEGTQTAKLQVTKGGPLPFAWDILSPQFSKAYVKHPADIPDYMKLSFPEGFTWERSMNFEDGGVVEVQQDSSLQDGTFIYKVKFKGVNFPADGPVMQKKTAGWEPSTEKLYPQDGVLKGEISHALKLKDGGHYTCDFKTVYKAKKPVQLPGNHYVDSKLDITNHNEDYTVVEQYEHAEARHSG

stability: 70.5536

surface area: 10439.58313654544<jupyter_text>Exhaustively evaluate all 1-edit mutantsResidue substitutions only<jupyter_code># the full mutation sweep is pretty slow (~ 1 day), saving results is recommended
asset_path = Path(f'./assets/{protein_name}-{chain_id}_1-edit_obj-vals.csv')
if os.path.exists(asset_path):
    df = pd.read_csv(asset_path)

else:
    wild_seq, _ = pdb_to_residues(parent_pdb_path, chain_id)
    records = []
    pos_range = range(len(wild_seq))
    for res, pos in itertools.product(AMINO_ACIDS, pos_range):
        mutation_list = [base_candidate.new_mutation(pos, res)]
        print(mutation_list[0])
        candidate = base_candidate.new_candidate(mutation_list)
        records.append({
            'seq_idx': pos,
            'wt_residue': wild_seq[pos],
            'mutant_residue': res,
            'mutant_uuid': candidate.uuid,
            'mutant_surface_area': candidate.mutant_surface_area,
            'mutant_total_energy': candidate.mutant_total_energy
        })

        # comment out to run full loop
        if pos == 10:
            break
            
    df = pd.DataFrame(records)
    os.makedirs(asset_path.parent, exist_ok=True)
    df.to_csv(asset_path, index=None)<jupyter_output><empty_output><jupyter_text>Evaluate fitness landscape<jupyter_code># post-processing
mutant_surface_area = torch.tensor(df.mutant_surface_area)
sa_min = mutant_surface_area.min()
sa_range = mutant_surface_area.max() - sa_min
x_data = 2 * (mutant_surface_area - sa_min) / sa_range - 1

mutant_stability = -torch.tensor(df.mutant_total_energy)
stab_min = mutant_stability.min()
stab_range = mutant_stability.max() - stab_min
y_data = 2 * (mutant_stability - stab_min) / stab_range - 1

obj = torch.stack([x_data, y_data], dim=-1)
pareto_mask = pareto.is_non_dominated(obj)

ref_x = 2 * (base_candidate.wild_surface_area - sa_min) / sa_range - 1
ref_y = 2 * (-base_candidate.wild_total_energy - stab_min) / stab_range - 1

# fit OLS model
lin_reg = LinearRegression().fit(x_data.view(-1, 1), y_data)
r_squared = lin_reg.score(x_data.view(-1, 1), y_data)
reg_x = np.linspace(-1, 1, 100)
reg_y = lin_reg.coef_ * reg_x + lin_reg.intercept_

# detailed OLS results
X2 = sm.add_constant(x_data.view(-1, 1).numpy())
est = sm.OLS(y_data.numpy(), X2)
est2 = est.fit()
print(est2.summary())
fig = plt.figure(figsize=(13, 8))
ax = fig.add_subplot(1, 1, 1)

ax.plot(reg_x, reg_y, color='black', linestyle='--', linewidth=2, zorder=1, label='OLS fit')
ax.scatter(x_data[~pareto_mask], y_data[~pareto_mask], linewidth=2,
           s=64, facecolors='none', edgecolors='cornflowerblue', zorder=2, label='dominated')
ax.scatter(x_data[pareto_mask], y_data[pareto_mask], 
           s=64, label='non-dominated', color='darkorange', zorder=3)
ax.scatter(ref_x, ref_y, color='red', s=64, marker='x', label=protein_name,
           linewidth=4, zorder=2)

ax.set_xlabel('SASA')
ax.set_ylabel('Stability')
ax.set_title(f'{protein_name} 1-edit Mutants')
ax.legend(ncol=4)
plt.tight_layout()

# plt.savefig(f'./figures/{protein_name}-{chain_id}_1-edit_obj-vals.pdf')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12471395881006865, 'fraction_numerical': 0.023073989321128908, 'mean_word_length': 5.127336448598131, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 12, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '3', 'dir': 'jupyter-structured-clean-dedup', 'id': '518720', 'license': 'Apache-2.0', 'n_tokens_mistral': 2210, 'n_tokens_neox': 2039, 'n_words': 291, 'path': 'notebooks/foldx_demo.ipynb', 'repo_name': 'samuelstanton/lambo', 'size': '5244'}