import numpy as np import pyrosetta as pyr from pyrosetta import teaching as pyrt from pyrosetta.toolbox import get_hbonds def get_scorefxn(this_type): if this_type == "hbonds": weights = [1.0, 0.55, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0] order = [pyrt.fa_atr, pyrt.fa_rep, pyrt.fa_sol, pyrt.lk_ball_wtd, pyrt.fa_elec, pyrt.hbond_lr_bb, pyrt.hbond_sr_bb, pyrt.hbond_bb_sc, pyrt.hbond_sc] elif this_type == "basic": weights = [1.0, 0.55, 1.0, 1.0, 1.0, 1.25] order = [pyrt.fa_atr, pyrt.fa_rep, pyrt.fa_sol, pyrt.lk_ball_wtd, pyrt.fa_elec, pyrt.pro_close] weights = weights[:-1] order = order[:-1] elif this_type == "all": stuff = "(fa_atr 1) (fa_rep 0.55) (fa_sol 1) (fa_intra_rep 0.005) (fa_intra_sol_xover4 1) (lk_ball_wtd 1) (fa_elec 1) (pro_close 1.25) (hbond_sr_bb 1) (hbond_lr_bb 1) (hbond_bb_sc 1) (hbond_sc 1) (dslf_fa13 1.25) (omega 0.4) (fa_dun 0.7) (p_aa_pp 0.6) (yhh_planarity 0.625) (ref 1) (rama_prepro 0.45)" things = stuff.strip().split(")") weights = [] order = [] for i_thing_count,thing in enumerate(things): if len(thing) > 0: if i_thing_count == 0: useful_stuff = thing[1:].split() else: useful_stuff = thing[2:].split() order.append(eval("pyrt." + useful_stuff[0])) weights.append(float(useful_stuff[1])) assert len(weights) == len(order) scorefxn_custom = pyrt.ScoreFunction() for thing, wt in zip(order,weights): scorefxn_custom.set_weight(thing, wt) print(scorefxn_custom) return scorefxn_custom, order, weights if __name__ == "__main__": pyr.init() native_file = "fip35_best_native.clean.pdb" this_type = "all" print("Computing Local Frustration") scorefxn_custom, order, weights = get_scorefxn(this_type) native_pose = pyr.pose_from_pdb(native_file) total_E = scorefxn_custom(native_pose) hbond_set = get_hbonds(native_pose) print(hbond_set.show(native_pose)) print("### Using eval_ci_2b ###") emap = pyrt.EMapVector() scorefxn_custom.eval_ci_2b(native_pose.residue(9), native_pose.residue(21), native_pose, emap) for thing,wt in zip(order,weights): print("%s: %f" % (thing, emap[thing])) print print("### using eval_cd_2b ###") emap = pyrt.EMapVector() scorefxn_custom.eval_cd_2b(native_pose.residue(9), native_pose.residue(21), native_pose, emap) for thing,wt in zip(order,weights): print("%s: %f" % (thing, emap[thing]))