from modeller import *
from modeller.optimizers import actions
from modeller.scripts import complete_pdb

# Load our custom steepest descent optimizer
from steepest_descent import steepest_descent

env = environ()
env.io.atom_files_directory = '../atom_files'
env.libs.topology.read(file='$(LIB)/top_heav.lib')
env.libs.parameters.read(file='$(LIB)/par.lib')

# Read in the initial structure:
code = '1fdx'
mdl = complete_pdb(env, code)
atmsel = selection(mdl)

# Generate the restraints:
mdl.restraints.make(atmsel, restraint_type='stereo', spline_on_site=False)

# Optimize with our custom optimizer:
opt = steepest_descent(max_iterations=80)
opt.optimize(atmsel, actions=actions.trace(5))