helixmc-run

#!/usr/bin/env python
# -*- coding: UTF-8 -*-
#This file is part of HelixMC.
#    Copyright (C) 2013  Fang-Chieh Chou <fcchou@stanford.edu>
#
#    This program is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program.  If not, see <http://www.gnu.org/licenses/>.

from os.path import isfile
import numpy as np
import helixmc
from helixmc import util
from helixmc.pose import HelixPose
from helixmc.random_step import RandomStepSimple, RandomStepAgg, symmetrize_WC
from helixmc.score import ScoreTweezers
import time
import argparse


def check_file_exist(filename):
    if not isfile(filename):
        raise IOError("File '%s' does not exist!" % filename)

start_time = time.time()
#############################################
#Parsing input arguments
parser = argparse.ArgumentParser(
    description='HelixMC: Base-pair level Monte-Carlo simulation package for'
    'Nucleic Acids.')
parser.add_argument(
    '-n_step', type=int, required=True, help='Number of MC steps',
    metavar='int')
group = parser.add_mutually_exclusive_group(required=True)
group.add_argument(
    '-params', type=str,
    help='File of database bp-step parameters', metavar='str')
group.add_argument(
    '-gaussian_params', type=str,
    help='File for mutlivariate gaussian parameters', metavar='str')
parser.add_argument(
    '-no_gaussian_sampling', action='store_true',
    help='No sampling assuming mutlivariate gaussian')
parser.add_argument(
    '-symmetrize', action='store_true',
    help='Symmetrize the input parameter set.')
parser.add_argument(
    '-seq', type=str,
    help='Full sequence or repeating unit. Input the sequence or a fasta file.'
    ' Assumes random sequence if not given.', metavar='str')
parser.add_argument(
    '-n_bp', type=int, help='Number of base-pairs', metavar='int')
parser.add_argument(
    '-force', type=float, default=0,
    help='External force along z-direction (pN)', metavar='float')
parser.add_argument(
    '-out', type=str, default='MC_data.npz',
    help='Name of output file (in .npz format)', metavar='str')
parser.add_argument(
    '-relax_step', type=int, default=120,
    help='Number of relax MC steps before saving data', metavar='int')
parser.add_argument(
    '-link_relax_step', type=int, default=50,
    help='Number of MC steps to relax after link constraint', metavar='int')
parser.add_argument(
    '-compute_exact_link', action='store_true',
    help='Compute exact linking number (slow)')
parser.add_argument(
    '-compute_fuller_link', action='store_true',
    help="Compute Fuller's linking number")
parser.add_argument(
    '-check_fuller', type=str,
    help='Print out the fuller vs exact writhe to specified file for checking.'
    ' Overrides the standard behavior.', metavar='str')
parser.add_argument(
    '-target_link', type=float,
    help='Constrain the linking number to stay close'
    ' to input value in radians', metavar='float')
parser.add_argument(
    '-torsional_stiffness', type=float, default=0,
    help='Stiffness of the torsional trap (for link constraint, pN.Å)',
    metavar='float')
parser.add_argument(
    '-xy_stiffness', type=float, default=0,
    help='Stiffness of the xy trap (pN/Å)', metavar='float')
parser.add_argument(
    '-save_all_frames', action='store_true',
    help='Save all frames sampled in the run.')
parser.add_argument(
    '-out_frame', type=str,
    help='File for final frame saving (.npz)', metavar='str')
parser.add_argument(
    '-in_frame', type=str,
    help='File for imported saved frame (.npz)', metavar='str')
parser.add_argument(
    '-frame0', type=str,
    help='File for frame of first base-pair (3x3 matrix in .npy)',
    metavar='str')
parser.add_argument(
    '-const_seed', action='store_true',
    help='Use constant seed in random generator')

args = parser.parse_args()
if args.params is None and args.no_gaussian_sampling:
    raise argparse.ArgumentError(
        '-params must be specified for non-gaussian sampling!!!')
if args.seq is None and args.n_bp is None:
    raise argparse.ArgumentError(
        'At least one of -n_bp and -seq must be specified!!!')

is_link_constrained = (
    args.target_link is not None and args.torsional_stiffness != 0)
fuller_check_out = args.check_fuller
is_check_fuller = (fuller_check_out is not None)
if args.const_seed:
    helixmc.constant_seed()

#############################################
#Initial setup
if args.params is not None:
    if args.params[-3:] == 'npz':
        random_step = RandomStepAgg(
            args.params, gaussian_sampling=not args.no_gaussian_sampling)
        if args.symmetrize:
            symmetrize_WC(random_step)
    else:
        if args.symmetrize:
            raise argparse.ArgumentError(
                '-symmetrize is True but user did not input '
                'a seq-dependent database file!')
        if args.seq is not None:
            raise argparse.ArgumentError(
                '-seq is specified but user did not input '
                'a seq-dependent database file!')
        random_step = RandomStepSimple.load_params(
            args.params, gaussian_sampling=not args.no_gaussian_sampling)
else:
    if args.seq is not None:
        raise argparse.ArgumentError(
            '-seq is specified but user did not input'
            ' a seq-dependent database file!')
    random_step = RandomStepSimple.load_gaussian_params(args.gaussian_params)
    if args.no_gaussian_sampling:
        raise argparse.ArgumentError(
            '-no_gaussian_sampling when -gaussian_params is being input!')

is_seq_dep = False
n_bp = args.n_bp
if args.seq is not None:
    is_seq_dep = True
    seq_unit = ''
    try:
        seq_unit = util.read_seq_from_fasta(args.seq)
    except ValueError:
        seq_unit = args.seq
    if n_bp is None:
        seq = seq_unit
        n_bp = len(seq)
    else:
        seq = ''
        while len(seq) < n_bp:
            seq += seq_unit
        seq = seq[:n_bp]
    seq_id = np.empty(n_bp - 1, dtype=int)
    for i in xrange(n_bp-1):
        name = seq[i:i+2]
        seq_id[i] = random_step.name2rand_idx(name)

if args.in_frame is not None:
    check_file_exist(args.in_frame)
    pose = HelixPose.from_file(args.in_frame)
else:
    frame0 = None
    if args.frame0 is not None:
        check_file_exist(args.frame0)
        frame0 = np.load(args.frame0)
    # The params should be in (n_bp-1, 6)
    params = np.tile(random_step.params_avg, (n_bp-1, 1))
    pose = HelixPose(params, frame0)

scorefxn = ScoreTweezers(
    args.force, args.torsional_stiffness, args.target_link, args.xy_stiffness)
scorefxn_prerun = ScoreTweezers(args.force, xy_stiffness=args.xy_stiffness)


#############################################
#MC Workhorse function
def MC_move(pose, moving_step, scorefxn):
    if is_seq_dep:
        next_step_data = random_step(seq_id[moving_step])
    else:
        next_step_data = random_step()
    if scorefxn.is_empty:
        pose.update(moving_step, *next_step_data)
        return True
    else:
        score_old = scorefxn(pose)
        pose.update_trial(moving_step, *next_step_data)
        score_new = scorefxn(pose)
        if util.MC_acpt_rej(score_old, score_new):
            pose.accept_update()
            return True
        else:
            pose.reject_update()
            return False
#############################################
#initial relax
for j in xrange(args.relax_step):
    for i in xrange(n_bp-1):
        MC_move(pose, i, scorefxn_prerun)
#############################################
#Trap ramping
if is_link_constrained:
    pose.compute_tw_wr = True
    curr_target = pose.link_fuller
    scorefxn_prerun = ScoreTweezers(
        args.force, args.torsional_stiffness, curr_target, args.xy_stiffness)

    incr = np.radians(20) * np.sign(args.target_link - curr_target)
    cutoff = abs(incr * 0.5)
    while abs(curr_target - args.target_link) > cutoff:
        for i in xrange(n_bp-1):
            MC_move(pose, i, scorefxn_prerun)
            if abs(curr_target - pose.link_fuller) <= cutoff:
                curr_target += incr
                scorefxn_prerun = ScoreTweezers(
                    args.force, args.torsional_stiffness,
                    curr_target, args.xy_stiffness)
            if abs(curr_target - args.target_link) <= cutoff:
                break
    for j in xrange(args.link_relax_step):
        for i in xrange(n_bp-1):
            MC_move(pose, i, scorefxn)
#############################################
#Actual sampling
accept = 0
all_coord = np.empty((args.n_step, 3))
all_frame = np.empty((args.n_step, 3, 3))
all_writhe = np.empty(args.n_step)
all_twist = np.empty(args.n_step)

if is_check_fuller:
    out = open(fuller_check_out, 'w')
    for j in xrange(args.n_step):
        for i in xrange(n_bp-1):
            is_accept = MC_move(pose, i, scorefxn)
            accept += is_accept
        out.write('%f %f\n' % (pose.writhe_fuller, pose.writhe_exact))
    out.close()
else:
    for j in xrange(args.n_step):
        for i in xrange(n_bp-1):
            is_accept = MC_move(pose, i, scorefxn)
            accept += is_accept
        all_coord[j] = pose.coord_terminal
        all_frame[j] = pose.frame_terminal
        if args.compute_exact_link:
            all_twist[j] = pose.twist
            all_writhe[j] = pose.writhe_exact
        elif args.compute_fuller_link:
            all_twist[j] = pose.twist
            all_writhe[j] = pose.writhe_fuller
        if args.save_all_frames:
            pose.write2disk('frame%d' % j)
    if args.compute_exact_link or args.compute_fuller_link:
        np.savez(
            args.out, coord_terminal=all_coord, frame_terminal=all_frame,
            writhe=all_writhe, twist=all_twist)
    else:
        np.savez(args.out, coord_terminal=all_coord, frame_terminal=all_frame)
#############################################
if args.out_frame is not None:
    pose.write2disk(args.out_frame)

print 'Total time = %.3f s' % (time.time() - start_time)
if args.n_step > 0:
    print 'Accept rate = %.2f %%' % (
        100 * float(accept) / (args.n_step * (n_bp - 1)))