# -*-Python-*-
# Created by bgrierson at 09 Mar 2018 12:11
"""
This script is an example of handling the situation that a script is given many inputs
where some are lists and some can be single values, but all are required for a code
with the same number of elements.
The trival example is a general need for the following cases
>>> run_myplot(shot=[123,456], time=1.0)
>>> run_myplot(shot=123, time=[1.0,2.0])
A non-trivial example is
>>> run_myplot(shot=[123,456], time=[1.0, 2.0])
where we need to know if we are making 2 plots or 4 plots (all permutations)
A Note on Numpy
---------------
This example provides explicit and verbose code for accomplishing this
organizational task. However, the majority of cases in which one simply wants
to run similar code across multiple inputs can be accomplished in 2 lines
using numpy.
The trivial examples above are accomplished with
>>> vec_plot = np.vectorize(run_myplot, otypes='O')
>>> vec_plot([123,456], 1.0)
>>> vec_plot([123,456], [1.0, 2.0])
Note, I assume your plotting function returns nothing (None) or a pyplot
object (like and Axes, for Figure instance). If you do not specify otypes,
the first combination of outputs is first used to determine the output type
and the whole thing is run n+1 times (fine for calculations, not ideal for plots).
Permutation is simply,
>>> vec_plot = np.vectorize(run_myplot, otypes='O')
>>> vec_plotnp.tile([123,456]), np.repeat([1.0,2.0]))
If you have a lot of arguments, and you want to permute them, you probably
have some specific for loops in mind and should write it all out - maybe
using something like the logic below. Still, if you are the type that likes
one liners you can do it like so,
>>> vec_plot(*np.meshgrid([123, 456], [1.0, 2.0, 3.0], [1, 2]))
and if you have non-numeric arguments, you can combine numpy and itertools,
which is probably getting a little too fancy for the average tastes,
>>> vec_plot(*np.transpose(list(itertools.product([123, 456], [1.0, 2.0, 3.0], ['A01', 'A02']))))
These examples are demonstrated at the end of the tutorial.
defaultVars parameters
----------------------
:param case: which case to run
;param perm: Permute or not
"""
defaultVars(case=1, perm=False)
if case == 1:
shots = 123
times = 1.0
runs = 'A01'
if case == 2:
shots = 123
times = [1.0, 2.0]
runs = 'A01'
if case == 3:
shots = [123, 456]
times = [1.0, 2.0]
runs = ['A01']
if case == 4:
shots = [123, 456]
times = [1.0, 2.0]
runs = ['A01', 'A02']
if case == 5:
# this is nonsense if we're not permuting, don't know how to plot, should fail
perm = True
shots = [123, 456]
times = [1.0, 2.0]
runs = ['A01', 'A02', 'A03']
if case == 6:
# this demonstrates that we can use different number of times per shot, etc.
# permutation does not make sense in this case
perm = False
shots = [123, 123, 456]
times = [1.0, 2.0, 3.0]
runs = 'A01'
def my_func(shot=None, time=None, run=None):
"""
This is a dummy function that would most likely be
path_to_script_in_tree.run in a real OMFIT application.
:return: None
"""
print(' my_func(shot={:}, time={:}, run={:})'.format(shot, time, run))
printi('\n\n\nRunning input list tutorial case {:}\n'.format(case) + '=' * 40)
# The list of inputs (each a list of its own)
shots = np.atleast_1d(shots)
times = np.atleast_1d(times)
runs = np.atleast_1d(runs)
inList = [shots, times, runs]
# Check the lengths
n = []
for el in inList:
n.append(len(el))
printi('Length of each input')
print(' {:}'.format(n))
# Check if length of all elements are equal
lens = len(set(iter(n)))
len_equal = lens <= 1
printi('Length Check:')
print(' {}'.format(len_equal))
myList = []
if len_equal and not perm:
printi(' Single values or all flat lists to zip')
myList = inList
elif len_equal and perm:
printi(' All flat lists and permuting before zip')
myTuples = list(itertools.product(*inList))
for i in range(len(inList)):
myList.append([item[i] for item in myTuples])
elif not len_equal and not perm:
# Only two allowable options: list has max(n) elements or list is 0 or 1D
for i, el in enumerate(inList):
if len(el) == np.max(n):
printi(' List element has appropriate number of entries')
elif len(el) == 1:
printi(' List element is single; duplicating')
print(' {:}'.format(el))
inList[i] = el.repeat(np.max(n))
else:
printe('Error, list element is neither the same as others, nor 1D')
raise OMFITexception('Inputs error')
myList = inList
elif not len_equal and perm:
printi(' Permuting before zip')
myTuples = list(itertools.product(*inList))
for i in range(len(inList)):
myList.append([item[i] for item in myTuples])
printi(' Final result')
print(' {:}'.format(myList))
printi("\nExplicit loop would call...")
for shot, time, run in zip(myList[0], myList[1], myList[2]):
my_func(shot, time, run)
# re-do entire tutorial using numpy vectorize (note this function returns None)
vec_func = np.vectorize(my_func, otypes='O')
printi("\nVectorized function would use...")
# it is probably cleanest and most intuitive to simply demand this be true
if not perm:
vec_func(shots, times, runs)
# but if you really want to catch all possible combos without
# writing them out explicitly in the argument lists, you caaaaan...
else:
vec_func(*np.transpose(list(itertools.product(shots, times, runs))))
printi('Input List Tutorial Completed!')
