# -*-Python-*-
# Created by bgrierson at 21 Dec 2016 07:29
# This is a demo to show how to set entries in a TRANSP namelist
# After we create the namelist, open in text editor to view how
# the arrays are indexed
from pylab import random
# Set up the OMFIT tree to store the output from this script
root['OUTPUTS'].setdefault('NAMELISTS', OMFITtree())
# Create a comment
txt = '''
! TRANSP namelist comment
'''
# Now create a TRANSP namelist
nml = OMFITtranspNamelist('12345TR.DAT', fromString=txt)
# Set a float
nml['TINIT'] = 1.0
# Set a boolean
nml['NLMDIF'] = True
# Set a float array
nml['TQMODA'] = np.array([1.1])
# Set an int array
nml['NQMODA'] = np.array([4, 1])
# Set a multi-dimensional array which will create
# TEST(1,1) = ?
# TEST(1,2) = ?
# TEST(1,3) = ?
# TEST(2,1) = ?
# TEST(2,2) = ?
# TEST(2,3) = ?
nml['RAND'] = random((2, 3))
# Set a multi-dimensional array where you only set some of the values
var = np.zeros((2, 3))
var[0, 1] = 2.0
nml['SPARSE'] = var
# Create a sub-namelist with entries defined in the namelist in the myriad ways
# that a TRANSP namelist can be formed
nubeam_txt = '''
! NUBEAM
NPTCLS=100
EINJA = 81e3, 81e3
PINJA(1) = 1, 2, 3
PINJA(4) = 4, 5, 6
'''
nubeam = OMFITtranspNamelist('nubeam.nml', fromString=nubeam_txt)
nml['NUBEAM'] = NamelistName()
nml['NUBEAM'].update(nubeam)
# Place the namelist in the tree
root['OUTPUTS']['NAMELISTS']['TRANSP'] = nml
# Create a sub-namelist with entries defined sparsely
sparray_txt = '''
! NUBEAM
NPTCLS=11
EINJA = 81e3, 81e3
PINJA(2) = 2.0
PINJA(4) = 4.0
'''
spa = OMFITtranspNamelist('sparray.nml', fromString=sparray_txt)
nml['SPARRAY'] = NamelistName()
nml['SPARRAY'].update(spa)
# Create a sparray directly
a = sparray(shape=(10,), default=np.nan)
a[(2,)] = 1.0
nml['SPARRAY']['SPARSE'] = a
# Can still add a variable after placing in tree
root['OUTPUTS']['NAMELISTS']['TRANSP']['FTIME'] = 2.0
# Add an update_time block
update_time_txt = '''
! Update in time
~update_time = {}
lpredict_te = True
'''.format(
1.8
)
nml = OMFITtranspNamelist('update_time.nml', fromString=update_time_txt)
root['OUTPUTS']['NAMELISTS']['TRANSP'].update(nml)