# Homology modeling by the automodel class
#
# Demonstrates how to refine only a part of the model.
#
# You may want to use the more exhaustive "loop" modeling routines instead.
#
from modeller import *
from modeller.automodel import *    # Load the automodel class
from modeller.parallel import task

# Override the 'select_atoms' routine in the 'automodel' class:
# (To build an all-hydrogen model, derive from allhmodel rather than automodel
# here.)
class MyModel(automodel):
    def select_atoms(self):
        return selection(self.residue_range('560:B', '581:B'))

env = environ()
# directories for input atom files
env.io.atom_files_directory = ['.', '../atom_files']
# selected atoms do not feel the neighborhood
env.edat.nonbonded_sel_atoms = 2
# Read in HETATM records from template PDBs
env.io.hydrogen = env.io.hetatm = True

class MyModel(automodel):
    def special_restraints(self, aln):
        rsr = self.restraints
        at = self.atoms
#       Add some restraints from a file:
#       rsr.append(file='constraints.rsr')

#       Residues 298 through 306 should be an alpha helix:
        rsr.add(secondary_structure.alpha(self.residue_range('564:B', '574:B')))

# Be sure to use 'MyModel' rather than 'automodel' here!
a = MyModel(env,
            alnfile  = 'mult.ali', # alignment filename
            knowns   = ('BETA_CORR'),   # codes of the templates
            sequence = 'SEQ', 
            assess_methods = (assess.DOPE, assess.GA341))              # code of the target

a.starting_model= 1                # index of the first model
a.ending_model  = 30                # index of the last model
                                   # (determines how many models to calculate)

# Very thorough VTFM optimization:
a.library_schedule = autosched.slow
a.max_var_iterations = 1000

# Thorough MD optimization:
a.md_level = refine.slow

# Repeat the whole cycle 2 times and do not stop unless obj.func. > 1E6
a.repeat_optimization = 2
a.max_molpdf = 1e6

a.make()                           # do homology modeling