Python写的调用Gromacs分子动力学MD的主函数

代码:

def main(args):
    parser = argparse.ArgumentParser(
        description='Physical validation suite for GROMACS.',
        prog='gmx_physicalvalidation.py',
        formatter_class=argparse.RawTextHelpFormatter,
        epilog='Use --tests for details about the available tests and their arguments.'
    )
    parser.add_argument('json', type=open,
                        metavar='systems.json',
                        help='Json file containing systems and tests to be ran.')
    parser.add_argument('--tests', default=False, action='store_true',
                        help='Print details about the available tests and their arguments and exit.')
    parser.add_argument('-v', '--verbosity', type=int,
                        metavar='v', default=0,
                        help='Verbosity level. Default: 0 (quiet).')
    parser.add_argument('--gmx', type=str, metavar='exe', default=None,
                        help=('GROMACS executable. Default: Trying to use \'gmx\'.\n' +
                              'Note: If -p is used, the GROMACS executable is not needed.'))
    parser.add_argument('--bindir', type=str, metavar='dir', default=None,
                        help=('GROMACS binary directory.\n' +
                              'If set, trying to use \'bindir/gmx\' instead of plain \'gmx\'\n' +
                              'Note: If --gmx is set, --bindir and --suffix are ignored.'))
    parser.add_argument('--suffix', type=str, metavar='_s', default=None,
                        help=('Suffix of the GROMACS executable.\n' +
                              'If set, trying to use \'gmx_s\' instead of plain \'gmx\'\n' +
                              'Note: If --gmx is set, --bindir and --suffix are ignored.'))
    group = parser.add_mutually_exclusive_group()
    group.add_argument('-p', '--prepare', action='store_true',
                       default=False,
                       help=('Only prepare simulations and output a \'run.sh\' file\n' +
                             'containing the necessary commands to run the systems.\n' +
                             'This allows to separate running simulations from analyzing them,\n' +
                             'useful e.g. to analyze the results on a different machine.\n' +
                             'Default: If none of \'-p\', \'-r\' or \'-a\' is given,\n' +
                             '         the systems are prepared, ran and analyzed in one call.'))
    group.add_argument('-r', '--run', action='store_true',
                       default=False,
                       help=('Only prepare and run simulations.\n' +
                             'Default: If none of \'-p\', \'-r\' or \'-a\' is given,\n' +
                             '         the systems are prepared, ran and analyzed in one call.'))
    group.add_argument('-a', '--analyze', action='store_true',
                       default=False,
                       help=('Only analyze previously ran simulations.\n' +
                             'This requires that the systems have been prepared using this program.\n' +
                             'Default: If none of \'-p\', \'-r\' or \'-a\' is given,\n' +
                             '         the systems are prepared, ran and analyzed in one call.'))
    parser.add_argument('-s', '--system', action='append', dest='systems',
                        metavar='system',
                        help=('Specify which system to run.\n' +
                              '\'system\' needs to match a system defined in the json file.\n' +
                              'Several systems can be specified by chaining several \'-s\' arguments.\n' +
                              'Defaults: If no \'-s\' argument is given, all systems and tests\n' +
                              '          defined in the json file are ran.\n' +
                              'Note: \'system\' can be a regular expression matching more than one system.'))
    parser.add_argument('--mpicmd', type=str, metavar='cmd', default=None,
                        help='MPI command used to invoke run command')
    parser.add_argument('--wd', '--working_dir', type=str,
                        metavar='dir', default=None,
                        help='Working directory (default: current directory)')
    parser.add_argument('--nobackup', default=False, action='store_true',
                        help='Do not create backups of files or folders.')
    
    if '--tests' in args:
        message = ('Physical validation suite for GROMACS\n'
                   'Available tests and their arguments\n'
                   '=====================================\n\n'
                   'The following tests can be specified in the json file:\n\n')
        for test in all_tests:
            message += '  * ' + test + '\n'

        message += '\nAll tests accept additional arguments:\n'
    
        for test, test_cls in all_tests.items():
            message += '\n'
            test_help = test_cls.parser().format_help()
            test_help = test_help.replace(test_cls.__name__, test).replace('usage: ', '')
            test_help = test_help.replace('  -h, --help            show this help message and exit\n', '')
            test_help = test_help.replace(' [-h]', '')
            test_help = test_help.replace(' ' * (len(test_cls.__name__) + 7), ' ' * len(test))
    
            first_line = test_help.split('\n')[0]
            separator = '-' * 70
            message += test_help.replace(first_line, separator + '\n' + first_line)

        sys.stderr.write(message)
        return

    args = parser.parse_args(args)

    # the input files are expected to be located where this script is
    source_path = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'systems')
    # target directory can be current or user chosen
    if args.wd is None:
        target_path = os.getcwd()
    else:
        if not os.path.exists(args.wd):
            os.makedirs(args.wd)
        target_path = args.wd

    # parse simulation stage to perform
    do_all = not (args.prepare or args.run or args.analyze)
    do_prepare = do_all or args.prepare or args.run
    write_script = args.prepare
    do_run = do_all or args.run
    do_analysis = do_all or args.analyze

    # get ordered dict of systems from combination of json file and user choices
    systems = parse_systems(args.json, args.systems, source_path, args.analyze)

    # prepare GROMACS interface
    if args.gmx:
        gmx = args.gmx
    else:
        gmx = 'gmx'
        if args.suffix:
            gmx += args.suffix
        if args.bindir:
            gmx = os.path.join(args.bindir, gmx)
    gmx_interface = None
    gmx_parser = None
    if do_run or do_analysis:
        gmx_interface = GromacsInterface(exe=gmx)
        gmx_parser = GromacsParser(exe=gmx)

    if do_prepare:
        nsystems = len(systems)
        n = 0
        runs = []  # this will contain all information needed to run the system
        for system_name, system in systems.items():
            n += 1
            print('\rPreparing run files for systems... [{:d}/{:d}] '.format(n, nsystems), end='')
            sys.stdout.flush()  # py2 compatibility
            system_dir = system['dir']
            system_dirs = []  # list of directories with subsystems
            # prepare the base system
            input_dir = os.path.join(source_path, system_dir, 'input')
            target_dir = os.path.join(target_path, system_dir)
            mkdir_bk(target_dir, nobackup=args.nobackup)
            basedir = os.path.join(target_dir, 'base')
            mkdir_bk(basedir, nobackup=args.nobackup)
            for suffix in ['mdp', 'gro', 'top']:
                shutil.copy2(os.path.join(input_dir, 'system.' + suffix),
                             basedir)
            system_dirs.append(basedir)
            # call prepare method of chosen tests
            for test_name, test in system['tests'].items():
                for test_args in test['args']:
                    system_dirs.extend(
                        all_tests[test_name].prepare_parser(input_dir, target_dir, system_name,
                                                            args.nobackup, test_args)
                    )

            # save run information
            for d in system_dirs:
                runs.append({
                    'dir': d,
                    'grompp_args': system['grompp_args'],
                    'mdrun_args': system['mdrun_args']
                })
        # end of loop over systems
        print('-- done.')

        if write_script:
            print('Writing run script... ', end='')
            sys.stdout.flush()  # py2 compatibility
            script_file = os.path.join(target_path, 'run_simulations.sh')
            if not args.nobackup:
                file_bk(script_file)
            with open(script_file, 'w') as f:
                f.write('# This file was created by the physical validation suite for GROMACS.\n')
                f.write('\n# Define run variables\n')
                f.write('WORKDIR=' + os.path.abspath(target_path) + '\n')
                f.write('GROMPPCMD="' + os.path.abspath(gmx) + ' grompp"\n')
                f.write('MDRUNCMD="' + os.path.abspath(gmx) + ' mdrun"\n')
                f.write('\n# Run systems\n')
                f.write('startpath=$PWD\n')
                f.write('cd $WORKDIR\n')
                for run in runs:
                    for cmd in basic_run_cmds(directory=os.path.relpath(os.path.abspath(run['dir']),
                                                                        os.path.abspath(target_path)),
                                              grompp_args=run['grompp_args'],
                                              mdrun_args=run['mdrun_args']):
                        f.write(cmd + '\n')
                    f.write('\n')
                f.write('cd $startpath\n')
            print('-- done.')
            print('Run script written to ' + script_file)
            print('Adapt script as necessary and run simulations. Make sure to preserve the folder structure!')
            print('Once all simulations have ran, analyze the results using `make check-phys-analyze` or '
                  'using the `-a` flag of `gmx_physicalvalidation.py`.')
        # end if write_script

        if do_run:
            nruns = len(runs)
            # send messages from GROMACS to log
            gmx_log = open(os.path.join(target_path, 'physicalvalidation_gmx.log'), 'w')
            for n, run in enumerate(runs):
                print('\rRunning (sub)systems... [{:d}/{:d}] '.format(n+1, nruns), end='')
                sys.stdout.flush()  # py2 compatibility
                gmx_interface.grompp(mdp='system.mdp',
                                     top='system.top',
                                     gro='system.gro',
                                     tpr='system.tpr',
                                     cwd=run['dir'],
                                     args=run['grompp_args'],
                                     stdout=gmx_log,
                                     stderr=gmx_log)
                gmx_interface.mdrun(tpr='system.tpr',
                                    deffnm='system',
                                    cwd=run['dir'],
                                    args=run['mdrun_args'],
                                    stdout=gmx_log,
                                    stderr=gmx_log,
                                    mpicmd=args.mpicmd)
            gmx_log.close()
            print('-- done.')
        # end if do_run
    # end if do_prepare

    if do_analysis:
        title = 'GROMACS PHYSICAL VALIDATION RESULTS'
        width = 70
        indent = int((width - len(title))/2)
        print()
        print(' ' * indent + '=' * len(title))
        print(' ' * indent + title)
        print(' ' * indent + '=' * len(title))
        print()
        passed = True
        for system_name, system in systems.items():
            system_dir = system['dir']
            # save system data if re-used for different test
            # massively reduces run time of multiple tests
            system_data = {
                'reduced': None,
                'full': None
            }
            # system directory
            target_dir = os.path.join(target_path, system_dir)

            print('Analyzing system ' + system_name)

            # call analyze method of chosen tests
            for test_name, test in system['tests'].items():
                for test_args in test['args']:
                    try:
                        result = all_tests[test_name].analyze_parser(gmx_parser, target_dir,
                                                                     system_name, system_data,
                                                                     args.verbosity, test_args)
                    except Exception as err:
                        print('    ' + all_tests[test_name].__name__ + ' FAILED (Exception in evaluation)')
                        print('    '*2 + type(err).__name__ + ': ' + str(err))
                        passed = False
                    else:
                        for line in result['message'].split('\n'):
                            print('    ' + line)

                        passed = passed and result['test']
            # end loop over tests
            print()
        # end loop over systems
        return int(not passed)
    # end if do_analysis

    # assuming everything is ok if we ended up here
    return 0


if __name__ == "__main__":
    return_value = main(sys.argv[1:])
    sys.exit(return_value)

  • 1
    点赞
  • 0
    收藏
    觉得还不错? 一键收藏
  • 打赏
    打赏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

温柔的行子

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值