Scattering calculations for hard sphere coarse-grained mAb models

"""
Simulate a 7-bead mAb and compute scattering intensity.
"""

import argparse
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from pyfeasst import fstio
from pyfeasst import scattering

def parse():
    """ Parse arguments from command line or change their default values. """
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--feasst_install', type=str, default='../../../build/',
                        help='FEASST install directory (e.g., the path to build)')
    parser.add_argument('--fstprt', type=str, default='/feasst/plugin/chain/particle/cg7mab2.fstprt',
                        help='FEASST particle definition')
    parser.add_argument('--cubic_side_length', type=float, default=90,
                        help='cubic periodic boundary conditions')
    parser.add_argument('--tpc', type=int, default=int(1e5), help='trials per cycle')
    parser.add_argument('--equilibration_cycles', type=int, default=int(1e0),
                        help='number of cycles for equilibration')
    parser.add_argument('--production_cycles', type=int, default=int(1e2),
                        help='number of cycles for production')
    parser.add_argument('--hours_checkpoint', type=float, default=0.2, help='hours per checkpoint')
    parser.add_argument('--hours_terminate', type=float, default=1., help='hours until termination')
    parser.add_argument('--procs_per_node', type=int, default=2, help='number of processors')
    parser.add_argument('--run_type', '-r', type=int, default=0,
                        help='0: run, 1: submit to queue, 2: post-process')
    parser.add_argument('--seed', type=int, default=-1,
                        help='Random number generator seed. If -1, assign random seed to each sim.')
    parser.add_argument('--max_restarts', type=int, default=10, help='Number of restarts in queue')
    parser.add_argument('--num_nodes', type=int, default=1, help='Number of nodes in queue')
    parser.add_argument('--scratch', type=str, default=None,
                        help='Optionally write scheduled job to scratch/logname/jobid.')
    parser.add_argument('--queue_flags', type=str, default="", help='extra flags for queue (e.g., for slurm, "-p queue")')
    parser.add_argument('--node', type=int, default=0, help='node ID')
    parser.add_argument('--queue_id', type=int, default=-1, help='If != -1, read args from file')
    parser.add_argument('--queue_task', type=int, default=0, help='If > 0, restart from checkpoint')

    # Convert arguments into a parameter dictionary, and add argument-dependent parameters.
    args, unknown_args = parser.parse_known_args()
    assert len(unknown_args) == 0, 'An unknown argument was included: '+str(unknown_args)
    params = vars(args)
    params['script'] = __file__
    params['prefix'] = 'cg7'
    params['sim_id_file'] = params['prefix']+ '_sim_ids.txt'
    params['minutes'] = int(params['hours_terminate']*60) # minutes allocated on queue
    params['hours_terminate'] = 0.95*params['hours_terminate'] - 0.05 # terminate before queue
    params['procs_per_sim'] = 1
    params['num_sims'] = params['num_nodes']*params['procs_per_node']
    params['nums'] = [30, 3]
    return params, args

def sim_node_dependent_params(params):
    """ Set parameters that depent upon the sim or node here. """
    params['num_particles'] = params['nums'][params['sim']]

def write_feasst_script(params, script_file):
    """ Write fst script for a single simulation with keys of params {} enclosed. """
    with open(script_file, 'w', encoding='utf-8') as myfile:
        myfile.write("""
MonteCarlo
RandomMT19937 seed {seed}
Configuration cubic_side_length {cubic_side_length} particle_type0 {fstprt}
Potential Model HardSphere VisitModel VisitModelCell min_length 5.3283
ThermoParams beta 1 chemical_potential -1
Metropolis
TrialTranslate tunable_param 2 tunable_target_acceptance 0.2
TrialParticlePivot weight 0.5 tunable_param 0.2 tunable_target_acceptance 0.25 particle_type 0
Checkpoint checkpoint_file {prefix}{sim}_checkpoint.fst num_hours {hours_checkpoint} num_hours_terminate {hours_terminate}

# grand canonical ensemble initalization
TrialAdd particle_type 0
Run until_num_particles {num_particles}
Remove name TrialAdd

# canonical ensemble equilibration
Metropolis trials_per_cycle {tpc} cycles_to_complete {equilibration_cycles}
Tune
CheckEnergy trials_per_update {tpc} decimal_places 8
Log trials_per_write {tpc} output_file {prefix}{sim}_eq.txt
Run until complete
Remove name0 Tune name1 Log

# canonical ensemble production
Metropolis trials_per_cycle {tpc} cycles_to_complete {production_cycles}
Log trials_per_write {tpc} output_file {prefix}{sim}.txt
Movie trials_per_write {tpc} output_file {prefix}{sim}.xyz
PairDistribution trials_per_update 1000 trials_per_write {tpc} dr 0.025 output_file {prefix}{sim}_gr.csv print_intra true
Scattering trials_per_update 100 trials_per_write {tpc} num_frequency 4 output_file {prefix}{sim}_iq.csv
Energy trials_per_write {tpc} output_file {prefix}{sim}_en.txt
CPUTime trials_per_write {tpc} output_file {prefix}{sim}_cpu.txt
Run until complete
""".format(**params))

def post_process(params):
    iq3=pd.read_csv(params['prefix']+'1_iq.csv', comment="#")
    iq30=pd.read_csv(params['prefix']+'0_iq.csv', comment="#")
    grp3 = iq3.groupby('q', as_index=False)
    grp30 = iq30.groupby('q', as_index=False)
    plt.scatter(grp3.mean()['q'], grp30.mean()['i']/grp3.mean()['i'], label='direct I(10g/L)/I(1g/L)', marker='.')
    iq3rdfft = scattering.intensity(gr_file=params['prefix']+'1_gr.csv', iq_file=params['prefix']+'1_iq.csv', num_density=3/90**3, skip=10)
    iq30rdfft = scattering.intensity(gr_file=params['prefix']+'0_gr.csv', iq_file=params['prefix']+'0_iq.csv', num_density=30/90**3, skip=10)
    plt.scatter(iq3rdfft['q'], iq30rdfft['iq']/iq3rdfft['iq'], label='rdf ft I(10g/L)/I(1g/L)')
    plt.ylabel('S', fontsize=16)
    plt.xlabel('q(1/nm)', fontsize=16)
    plt.legend()
    #plt.savefig(params['prefix']+'.png', bbox_inches='tight', transparent='True')
    assert np.abs(grp30.mean()['i'][0]/grp3.mean()['i'][0] - 0.82777) < 0.05
    assert np.abs(iq30rdfft['iq'][0]/iq3rdfft['iq'][0] - 0.7784) < 0.05

if __name__ == '__main__':
    parameters, arguments = parse()
    fstio.run_simulations(params=parameters,
                          sim_node_dependent_params=sim_node_dependent_params,
                          write_feasst_script=write_feasst_script,
                          post_process=post_process,
                          queue_function=fstio.slurm_single_node,
                          args=arguments)