mrdna

#! /usr/bin/env python
# -*- coding: utf-8 -*-
import sys
if sys.version_info[0] != 3:
    raise Exception("This script must be run with python3")
    
import argparse
import re
import pathlib
from mrdna import __version__ as __version__
from mrdna.simulate import multiresolution_simulation as simulate

"""Easy multiresolution simulations of DNA nanotechnology objects using ARBD"""

parser = argparse.ArgumentParser(prog="mrdna",
				 description=__doc__,
                                 formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)

parser.add_argument('-o','--output-prefix', type=str, default=None,
                    help="Name for your job's output")
parser.add_argument('-d','--directory', type=str,  default=None,
                    help='Directory for simulation; does not need to exist yet')
parser.add_argument('-g','--gpu', type=int, default=0,
                    help='GPU for simulation; check nvidia-smi for availability')
parser.add_argument('--output-period', type=float, default=1e4,
                    help='Simulation steps between DCD frames')
# parser.add_argument('--minimization-steps', type=float, default=0,
#                     help='Simulation steps between DCD frames')
parser.add_argument('--coarse-local-twist', action='store_true',
                    help='Use local twist for coarsest representation?')
parser.add_argument('--coarse-steps', type=float, default=1e7,
                    help='Simulation steps for coarse model (200 fs/step)')
parser.add_argument('--fine-steps', type=float, default=1e7,
                    help='Simulation steps for fine model (50 fs/step)')
parser.add_argument('--backbone-scale', type=float, default=1.0,
                    help='Factor to scale DNA backbone in atomic model; try 0.25 to avoid clashes for atomistic simulations')
parser.add_argument('--debug', action='store_true',
                    help='Run through the python debugger?')

parser.add_argument('--draw-cylinders', action='store_true',
                    help='Whether or not to draw the cylinders')
parser.add_argument('--draw-tubes', action='store_true',
                    help='Whether or not to draw the tubes')

parser.add_argument('--version', action='version',
                    version='%(prog)s {}'.format(__version__),
                    help='Print the version of mrdna')

parser.add_argument('input_file', type=str,
                    help="""Any of the following:
                    (1) a cadnano JSON file;
                    (2) a vHelix Maya (.ma) file;
                    (3) an atomic PDB file""")

args = parser.parse_args()

def main():
    infile = pathlib.Path(args.input_file)
    try:
        prefix = infile.stem
        extension = infile.suffix
    except:
        raise Exception("Unrecognized input file '{}'".format(infile))
        
    if extension == '.json':
        from mrdna.readers import read_cadnano as read_model
    elif extension == '.ma':
        from mrdna.readers import read_vhelix as read_model
    elif extension == '.pdb':
        from mrdna.readers import read_atomic_pdb as read_model
    else:
        raise Exception("Unrecognized input file '{}'".format(infile))
                    
    model = read_model( str(infile) )

    if args.output_prefix is not None:
        prefix = args.output_prefix
    
    if args.draw_cylinders is True:
        output_name = prefix + ".cylinders.tcl"
        model.vmd_cylinder_tcl( output_name )
        if prefix == infile.stem and extension == ".pdb":
            print("Are you sure you want to overwrite '{}'? (you may provide '-o' option) [y/N]".format(infile))
            if input() not in ("y","Y","yes","Y","Yes","YES"):
                import sys
                sys.exit(1)
        model._clear_beads()
        model._generate_atomic_model(scale=args.backbone_scale)            
        model.write_atomic_ENM( prefix )
        model.atomic_simulate( output_name = prefix )
    elif args.draw_tubes is True:
        output_name = prefix + ".tubes.tcl"
        model.vmd_tube_tcl( output_name )
        if prefix == infile.stem and extension == ".pdb":
            print("Are you sure you want to overwrite '{}'? (you may provide '-o' option) [y/N]".format(infile))
            if input() not in ("y","Y","yes","Y","Yes","YES"):
                import sys
                sys.exit(1)
        model._clear_beads()
        model._generate_atomic_model(scale=args.backbone_scale)            
        model.write_atomic_ENM( prefix )
        model.atomic_simulate( output_name = prefix )
    else:
        run_args = dict(
            model = model,
            output_name = prefix,
            job_id = "job-" + prefix,
            directory = args.directory,
            gpu = args.gpu,
            minimization_output_period = int(args.output_period),
            coarse_local_twist = args.coarse_local_twist,
            coarse_output_period = int(args.output_period),
            fine_output_period = int(args.output_period),
            minimization_steps = 0, # int(args.minimization_steps),
            coarse_steps = int(args.coarse_steps),
            fine_steps = int(args.fine_steps),
            backbone_scale = args.backbone_scale
        )

        simulate( **run_args )

if __name__ == '__main__':

    if args.debug:

        class Restart(Exception):
            """Causes a debugger to be restarted for the debugged python program."""
            pass

        import traceback
        import sys
        from pdb import Pdb
        pdb = Pdb()

        while True:
            try:
                # pdb._runscript(mainpyfile)
                self._user_requested_quit = False
                pdb.runcall(main)
                if pdb._user_requested_quit:
                    break
                print("The program finished and will be restarted")
            except Restart:
                print("Restarting...")
                print("\t" + " ".join(args))
            except SystemExit:
                # In most cases SystemExit does not warrant a post-mortem session.
                # print("The program exited via sys.exit(). Exit status:", end=' ')
                print("The program exited via sys.exit(). Exit status:")
                print(sys.exc_info()[1])
            except SyntaxError:
                traceback.print_exc()
                sys.exit(1)
            except:
                traceback.print_exc()
                print("Uncaught exception. Entering post mortem debugging")
                print("Running 'cont' or 'step' will restart the program")
                t = sys.exc_info()[2]
                pdb.interaction(None, t)
                print("Post mortem debugger finished. mrdna will be restarted")

    else:
        main()