-
Notifications
You must be signed in to change notification settings - Fork 0
/
2dsolve.py
executable file
·82 lines (72 loc) · 2.82 KB
/
2dsolve.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Written by Rui Xu, Jan. 2018
# Email: [email protected]
import sys
import time
import disp_tmp as disp
import utils
import logging
import argparse
import numpy as np
from scipy import interp
from scipy.optimize import root
__author__ = 'ruix'
def main(args):
""" parse command line arguments"""
tstart = time.clock()
if not args.log :
logging.basicConfig(level=args.loglevel or logging.INFO)
else:
logging.basicConfig(filename='log', filemode='w', level=logging.DEBUG)
logger = logging.getLogger(__name__)
""" read plasma parameters """
param = utils.read_param(args)
""" iterate through wavenumber """
dk = (param['kend'][0]-param['kstart'][0])/(param['ksteps'][0]-1)
#fzeta = np.empty(int(param['ksteps'][0]),dtype=complex)
fzeta = np.empty((int(param['ntheta'][0]), int(param['ksteps'][0])),dtype=complex)
wave_k = np.empty(int(param['ksteps'][0]))
zeta_guess = complex(param['omega_r'][0],param['omega_i'][0])
theta_arr = np.linspace(param['thetai'][0], param['thetae'][0],
int(param['ntheta'][0]))
for i in range(len(theta_arr)):
param['theta'][0] = theta_arr[i]
for n in range(int(param['ksteps'][0])):
if i>0 and n==0:
#zeta_guess = complex(param['omega_r'][0],param['omega_i'][0])
zeta_guess = fzeta[i-1][0]
logger.info('%d th iteration in %d ksteps \n' ,n,param['ksteps'][0])
wave_k[n] = param['kstart'][0]+n*dk
""" find dispersion relation root """
data = (args, param, wave_k[n])
try:
sol = root(disp.det,(zeta_guess.real,zeta_guess.imag), \
args=data,method='lm',tol=param['sol_err'][0])
fzeta[i][n] = complex(sol.x[0],sol.x[1])
logger.info("solution: zeta= %1.1e +%1.1e i",fzeta[i][n].real,fzeta[i][n].imag)
except ValueError:
logger.info('ERROR in root finding: wave_k =%f',wave_k[n])
sys.exit()
zeta_guess = fzeta[i][n]
""" save results to output file """
param['wave_k'] = wave_k
param['fzeta'] = fzeta
param['theta'] = theta_arr
if args.output :
np.save(args.output+'.npy',param)
else:
np.save('output.npy', param)
tend = time.clock()
logger.info('\n Total time lapse: %1.2f\n',tend-tstart)
return 0
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Specify command line arguments')
parser.add_argument('-i','--input', help='Input file name',required=False)
parser.add_argument('-o','--output',help='Output file name')
parser.add_argument('-l','--log', action='store',help='log file')
parser.add_argument('-v', '--verbose', help='Verbose (debug) logging',
action='store_const', const=logging.DEBUG,dest='loglevel')
args = parser.parse_args()
#print (parser.parse_args())
sys.exit(main(args))