Source code documentation

Framework

omfit_tree.reload_python(moduleName, quiet=False)[source]

This function extends the Python builtin reload function to easily reload OMFIT classes

>>> reload_python(omfit_classes.omfit_onetwo)
Parameters
  • moduleName – module or module name

  • quiet – bool Suppress print statements listing reloaded objects

class omfit_tree.OMFITmaintree(*args, **kwargs)[source]

Bases: omfit_classes.omfit_base.OMFITproject

Parameters
  • filename – ‘directory/bla/OMFITsave.txt’ or ‘directory/bla.zip’ where the OMFITtree will be saved (if ‘’ it will be saved in the same folder of the parent OMFITtree)

  • only – list of strings used to load only some of the branches from the tree (eg. [“[‘MainSettings’]”,”[‘myModule’][‘SCRIPTS’]”]

  • modifyOriginal – by default OMFIT will save a copy and then overwrite previous save only if successful. If modifyOriginal=True and filename is not .zip, will write data directly at destination, which will be faster but comes with the risk of deleting a good save if the new save fails for some reason

  • readOnly – will place entry in OMFITsave.txt of the parent so that this OMFITtree can be loaded, but will not save the actual content of this subtree. readOnly=True is meant to be used only after this subtree is deployed where its fileneme says it will be. Using this feature could result in much faster projects save if the content of this tree is large.

  • quiet – Verbosity level

  • developerMode – load OMFITpython objects within the tree as modifyOriginal

  • serverPicker – take server/tunnel info from MainSettings[‘SERVER’]

  • remote – access the filename in the remote directory

  • server – if specified the file will be downsync from the server

  • tunnel – access the filename via the tunnel

  • **kw – Extra keywords are passed to the SortedDict class

start(filename='')[source]
projectName()[source]
onlyRunningCopy(deletePID=False)[source]
Parameters

delete – whether to remove PID from list of running OMFIT processes (this should be done only at exit)

Returns

return True/False wether this is the only running copy of OMFIT on this computer

reset()[source]
newProvenanceID()[source]
newProjectID()[source]
addMainSettings(updateUserSettings=False, restore='')[source]
add_bindings_to_main_settings()[source]

Take the descriptions and events from global_events_bindings and insert them in self[‘MainSettings’][‘SETUP’][‘KeyBindings’][desc] = <event>

apply_bindings()[source]

Take the descriptions and events from self[‘MainSettings’][‘SETUP’][‘KeyBindings’] and use them to update the global_event_bindings

save(quiet=None, skip_save_errors=False)[source]

Writes the content of the OMFIT tree to the filesystem using the same filename and zip options of the last saveas

Parameters
  • quiet – whether to print save progress to screen

  • skip_save_errors – skip errors when saving objects

saveas(filename, zip=None, quiet=None, skip_save_errors=False)[source]

Writes the content of the OMFIT tree to the filesystem and permanently changes the name of the project

Parameters
  • filename – project filename to save to

  • zip – whether the save should occur as a zip file

  • quiet – whether to print save progress to screen

  • skip_save_errors – skip errors when saving objects

loadModule(filename, location=None, withSubmodules=True, availableModulesList=None, checkLicense=True, developerMode=None, depth=0, quiet=False, startup_lib=None, **kw)[source]

Load a module in OMFIT

Parameters
  • filename

    • full path to the module to be loaded

    • if just the module name is provided, this will be loaded from the public modules

    • remote/branch:module format will load a module from a specific git remote and branch

    • module:remote/branch format will load a module from a specific git remote and branch

  • location – string with the location where to place the module in the OMFIT tree

  • withSubmodules – load submodules or not

  • availableModulesList – list of available modules generated by OMFIT.availableModules() If this list is not passed, then the availableModulesList is generated internally

  • checkLicense – Check license files at load

  • developerMode – Load module with developer mode option (ie. scripts loaded as modifyOriginal) if None then default behavior is set by OMFIT['MainSettings']['SETUP']['developer_mode_by_default'] Note: public OMFIT installation cannot be loaded in developer mode

  • depth – parameter used internally by for keeping track of the recursion depth

  • quiet – load modules silently or not

  • startup_lib – Used internally for executing OMFITlib_startup scripts

  • **kw – additional keywords passed to OMFITmodule() class

Returns

instance of the loaded module

load(filename, persistent_projectID=False)[source]

loads an OMFITproject in OMFIT

Parameters
  • filename – filename of the project to load (if -1 then the most recent project is loaded)

  • persistent_projectID – whether the projectID should remain the same

updateCWD()[source]
updateMainSettings()[source]
saveMainSettingsSkeleton()[source]

This utility function updates the MainSettingsSkeleton for the current OMFIT installation

availableModules(quiet=None, directories=None, same_path_as=None, force=False)[source]

Index available OMFIT modules in directories

Parameters
  • quiet – verbosity

  • directories – list of directories to index. If None this is taken from OMFIT[‘MainSettings’][‘SETUP’][‘modulesDir’]

  • same_path_as – sample OMFITsave.txt path to set directory to index

  • force – force rebuild of .modulesInfoCache file

Returns

This method returns a dictionary with the available OMFIT modules. Each element in the dictionary is a dictionary itself with the details of the available modules.

quit(deepClean=False)[source]

Cleanup current OMFIT session directory OMFITsessionDir and PID from OMFITsessionsDir Also close all SSH related tunnels and connections

Parameters

deepClean – if deepClean is True, then the OMFITtmpDir and OMFITsessionsDir get deleted

reset_connections(server=None, mds_cache=True)[source]

Reset connections, stored passwords, and MDS+ cache

Parameters

server – only reset SSH connections to this server

recentProjects(only_read=False)[source]

This routine keeps track of the recent projects and is also in charge of deleting auto-saves if they do not appear in the project list.

Parameters

read_only – only read projects file (don’t do any maintenance on it)

showExecDiag()[source]

display execution diagram

class omfit_tree.OMFITmainscratch(*args, **kwargs)[source]

Bases: omfit_classes.omfit_base.OMFITtmp

initialize()[source]
load_framework_gui(item)[source]
omfit_tree.locationFromRoot(location)[source]
Parameters

location – tree location

Returns

tree location from the closest root

omfit_tree.rootLocation(location, *args, **kw)[source]
Parameters

location – tree location

Returns

tree location of the root

omfit_tree.OMFITobject_fromS3(filename, s3bucket)[source]

Recovers OMFITobject from S3 and reloads it with the right class and original keyword parameters

Returns

object loaded from S3

class omfit_tree.OMFITfileASCII(filename, **kw)[source]

Bases: omfit_classes.omfit_ascii.OMFITascii

Use of this class is deprecated: use OMFITascii instead

class omfit_tree.OMFIT_Ufile(*args, **kwargs)[source]

Bases: omfit_classes.omfit_ufile.OMFITuFile

Use of this class is deprecated: use OMFITuFile instead

class omfit_tree.OMFITdict(*args, **kwargs)[source]

Bases: omfit_classes.sortedDict.SortedDict

Use of this class is deprecated: use SortedDict instead

class omfit_tree.chebyfit(x, y, yerr, m=18)[source]

Bases: omfit_classes.utils_fit.fitCH

Use of this class is deprecated: use fitCH instead

exception omfit_tree.OMFITdependenceException[source]

Bases: RuntimeError

Use of this class is deprecated: use RuntimeError instead

omfit_classes.sortedDict.sortHuman(inStr)[source]

Sort the given list the way that humans expect

omfit_classes.sortedDict.get_bases(clss, tp=[])[source]

Returns a list of strings describing the dependencies of a class

omfit_classes.sortedDict.parseBuildLocation(inv)[source]

DEPRECATED: use parseLocation and buildLocation functions instead

Function to handle locations in the OMFIT tree (i.e. python dictionaries)

Parameters

inv – input location

Returns

  • if inv is a string, then the dictionary path is split and a list is returned (Note that this function strips the root name)

  • if it’s a list, then the dictionary path is assembled and a string is returned (Note that this function assumes that the root name is missing)

omfit_classes.sortedDict.parseLocation(inv)[source]

Parse string representation of the dictionary path and return list including root name This function can parse things like: OMFIT[‘asd’].attributes[u’aiy’ ][“[ ‘bla’][‘asa’]”][3][1:5]

Parameters

inv – string representation of the dictionary path

Returns

list of dictionary keys including rootname

omfit_classes.sortedDict.traverseLocation(inv)[source]

returns list of locations to reach input location

Parameters

inv – string representation of the dictionary path

Returns

list of locations including rootname to reach input location

omfit_classes.sortedDict.buildLocation(inv)[source]

Assemble list of keys into dictionary path string

Parameters

inv – list of dictionary keys including rootname

Returns

string representation of the dictionary path

omfit_classes.sortedDict.setLocation(location, value, globals=None, locals=None)[source]

Takes a string or a list of stirings output by parseLocation() and set the leaf to the value provided

Parameters
  • location – string or a list of stirings output by parseLocation()

  • value – value to set the leaf

  • globals – global namespace for the evaluation of the location

  • locals – local namespace for the evaluation of the location

Returns

value

omfit_classes.sortedDict.dirbaseLocation(location)[source]

Takes a string or a list of stirings output by parseLocation() and returns two strings for convenient setting of dictionary locations

>> d, b=dirbaseLocation(“OMFIT[‘dir’][‘base’]”) >> eval(d)[b] d = OMFIT[‘dir’] b = ‘base’

Parameters

location – string or a list of stirings output by parseLocation()

Returns

two string, the first one with the path leading to the leaf, the second with the name of the leaf

omfit_classes.sortedDict.traverse(self, string='', level=100, split=True, onlyDict=False, onlyLeaf=False, skipDynaLoad=False, noSubmodules=False, traverse_classes=(<class 'collections.abc.MutableMapping'>, ))[source]

Returns a string or list of strings describing the path of every entry/subentries in the dictionary

Parameters
  • string – string to be appended in front of all entries

  • level – maximum depth

  • split – split the output string into a list of strings

  • onlyDict – return only dictionary entries (can be a tuple of classes)

  • onlyLeaf – return only non-dictionary entries (can be a tuple of classes)

  • skipDynaLoad – skip entries that have .dynaLoad==True

  • noSubmodules – controls whether to traverse submodules or not

  • traverse_classes – tuple of classes to traverse

Returns

string or list of string

omfit_classes.sortedDict.treeLocation(obj, memo=None)[source]

Identifies location in the OMFIT tree of an OMFIT object

NOTE: Typical users should not need to use this function as part of their modules. If you find yourself using this function in your modules, it is likely that OMFIT already provides the functionality that you are looking for in some other way. We recommend reaching out the OMFIT developers team to see if there is an easy way to get what you want.

Parameters
  • obj – object in the OMFIT tree

  • memo – used internally to avoid infinite recursions

Returns

string with tree location

omfit_classes.sortedDict.recursiveUpdate(A, B, overwrite=True, **kw)[source]

Recursive update of dictionary A based on data from dictionary B

Parameters
  • A – dictionary A

  • B – dictionary B

  • overwrite – force overwrite of duplicates

Returns

updated dictionary

omfit_classes.sortedDict.pretty_diff(d, ptrn={})[source]

generate “human readable” dictionary output from SortedDict.diff()

omfit_classes.sortedDict.prune_mask(what, ptrn)[source]

prune dictionary structure based on mask The mask can be in the form of of a pretty_diff dictionary or a list of traverse strings

omfit_classes.sortedDict.dynaLoad(f)[source]

Decorator which calls dynaLoader method

Parameters

f – function to decorate

Returns

decorated function

omfit_classes.sortedDict.dynaLoadKey(f)[source]

Decorator which calls dynaLoad method only if key is not found

Parameters

f – function to decorate

Returns

decorated function

omfit_classes.sortedDict.dynaSave(f)[source]

Decorator which calls dynaSaver method

Parameters

f – function to decorate

Returns

decorated function

omfit_classes.sortedDict.dynaLoader(self, f=None)[source]

Call load function if object has dynaLoad attribute set to True After calling load function the dynaLoad attribute is set to False

omfit_classes.sortedDict.dynaSaver(self)[source]

This function is meant to be called in the .save() function of objects of the class OMFITobject that support dynamic loading. The idea is that if an object has not been loaded, then its file representation has not changed and the original file can be resued. This function returns True/False to say if it was successful at saving. If True, then the original .save() function can return, otherwise it should go through saving the data from memory to file.

class omfit_classes.sortedDict.SortedDict(*args, **kwargs)[source]

Bases: dict

A dictionary that keeps its keys in the order in which they’re inserted

Parameters
  • data – A dict object or list of (key,value) tuples from which to initialize the new SortedDict object

  • **kw – Optional keyword arguments given below

kw:
param caseInsensitive

(bool) If True, allows self[‘UPPER’] to yield self[‘upper’].

param sorted

(bool) If True, keep keys sorted alphabetically, instead of by insertion order.

param limit

(int) keep only the latest limit number of entries (useful for data cashes)

param dynaLoad

(bool) Not sure what this does

index(item)[source]

returns the index of the item

pop(k[, d])v, remove specified key and return the corresponding value.[source]

If key is not found, d is returned if given, otherwise KeyError is raised

popitem()(k, v), remove and return some (key, value) pair as a[source]

2-tuple; but raise KeyError if D is empty.

items()a set-like object providing a view on D’s items[source]
iteritems()[source]
keys(filter=None, matching=False)[source]

returns the sorted list of keys in the dictionary

Parameters
  • filter – regular expression for filtering keys

  • matching – boolean to indicate whether to return the keys that match (or not)

Returns

list of keys

iterkeys()[source]
values()an object providing a view on D’s values[source]
itervalues()[source]
update([E, ]**F)None.  Update D from dict/iterable E and F.[source]

If E is present and has a .keys() method, then does: for k in E: D[k] = E[k] If E is present and lacks a .keys() method, then does: for k, v in E: D[k] = v In either case, this is followed by: for k in F: D[k] = F[k]

recursiveUpdate(other, overwrite=False)[source]
setdefault(key, default)[source]

The method setdefault() is similar to get(), but will set dict[key]=default if key is not already in dict

Parameters
  • key – key to be accessed

  • default – default value if key does not exist

Returns

value

get(key, default)[source]

Return the value for key if key is in the dictionary, else default.

value_for_index(index)[source]

Returns the value of the item at the given zero-based index

insert(index, key, value)[source]

Inserts the key, value pair before the item with the given index

copy()[source]

Returns a copy of this object

clear()None.  Remove all items from D.[source]
moveUp(index)[source]

Shift up in key list the item at a given index

Parameters

index – index to be shifted

Returns

None

moveDown(index)[source]

Shift down in key list the item at a given index

Parameters

index – index to be shifted

Returns

None

across(what='', sort=False, returnKeys=False)[source]

Aggregate objects across the tree

Parameters
  • what – string with the regular expression to be cut across

  • sort – sorting of results alphabetically

  • returnKeys – return keys of elements in addition to objects

Returns

list of objects or tuple with with objects and keys

>> OMFIT[‘test’]=OMFITtree() >> for k in range(5): >> OMFIT[‘test’][‘aaa’+str(k)]=OMFITtree() >> OMFIT[‘test’][‘aaa’+str(k)][‘var’]=k >> OMFIT[‘test’][‘bbb’+str(k)]=-1 >> print(OMFIT[‘test’].across(“[‘aaa*’][‘var’]”))

sort(key=None, **kw)[source]
Parameters

key – function that returns a string that is used for sorting or dictionary key whose content is used for sorting

>> tmp=SortedDict() >> for k in range(5): >> tmp[k]={} >> tmp[k][‘a’]=4-k >> # by dictionary key >> tmp.sort(key=’a’) >> # or equivalently >> tmp.sort(key=lambda x:tmp[x][‘a’])

Parameters

**kw – additional keywords passed to the underlying list sort command

Returns

sorted keys

sort_class(class_order=[<class 'dict'>])[source]

sort items based on their class

Parameters

class_order – list containing order of class

Returns

sorted keys

diff(other, ignoreComments=False, ignoreContent=False, skipClasses=(), noloadClasses=(), precision=0.0, order=True, favor_my_order=False, modify_order=False, quiet=True)[source]

Comparison of a SortedDict

Parameters
  • other – other dictionary to compare to

  • ignoreComments – ignore keys that start and end with “__” (e.g. “__comment__”)

  • ignoreContent – ignore content of the objects

  • skipClasses – list of class of objects to ignore

  • noloadClasses – list of class of objects to not load

  • precision – relative precision to which the comparison is carried out

  • order – does the order of the keys matter

  • favor_my_order – favor my order of keys

  • modify_order – update order of input dictionaries based on diff

  • quiet – verbosity of the comparison

Returns

comparison dictionary

pretty_diff(other, ignoreComments=False, ignoreContent=False, skipClasses=(), noloadClasses=(), precision=0.0, order=True, favor_my_order=False, quiet=True)[source]

Comparison of a SortedDict in human readable format

Parameters
  • other – other dictionary to compare to

  • ignoreComments – ignore keys that start and end with “__” (e.g. “__comment__”)

  • ignoreContent – ignore content of the objects

  • skipClasses – list of class of objects to ignore

  • noloadClasses – list of class of objects to not load

  • precision – relative precision to which the comparison is carried out

  • order – does the order of the keys matter

  • favor_my_order – favor my order of keys

  • quiet – verbosity of the comparison

Returns

comparison dictionary in pretty to print format

diffKeys(other)[source]
Parameters

other – other dictionary to compare to

Returns

floating point to indicate the ratio of keys that are similar

changeKeysCase(case=None, recursive=False)[source]

Change all the keys in the dictionary to be upper/lower case

Parameters
  • case – ‘upper’ or ‘lower’

  • recursive – apply this recursively

Returns

None

traverse(string='', level=100, onlyDict=False, onlyLeaf=False, skipDynaLoad=False)[source]

Equivalent to the tree command in UNIX

Parameters
  • string – prepend this string

  • level – depth

  • onlyDict – only subtrees and not the leafs

Returns

list of strings containing the dictionary path to each object

walk(function, **kw)[source]

Walk every member and call a function on the keyword and value

Parameters
  • functionfunction(self,kid,**kw)

  • **kw – kw passed to the function

Returns

self

safe_del(key)[source]

Delete key entry only if it is present

Parameters

key – key to be deleted

flatten()[source]

The hierarchical structure of the dictionaries is flattened

Returns

SortedDict populated with the flattened content of the dictionary

setFlat(key, value)[source]

recursively searches dictionary for key in order to set a value raises KeyError if key could not be found, so this method cannot be used to set new entries in the dictionary

Parameters
  • key – key to be set

  • value – value to set

check_location(location, value=[])[source]

check if location exist and equals value (if value is specified)

Parameters
  • location – location as string

  • value – value for which to return equal

Returns

True/False

>> root[‘SETTINGS’].check_location(“[‘EXPERIMENT’][‘shot’]”, 133221)

todict()[source]

Return a standard Python dictionary representation of the SortedDict

Returns

dictionary

class omfit_classes.sortedDict.OMFITdataset(data_vars=None, coords=None, attrs=None)[source]

Bases: object

Subclassing from this class is like subclassing from the xarray.Dataset class but without having to deal with the hassle of inheriting from xarrays (internally this class uses class composition rather than subclassing).

Also this class makes it possible to use the OMFIT dynamic loading capabilities. All classes that subclass OMFITdataset must define the .dynaLoad attribute.

NOTE: Classes that subclass from OMFITdataset will be identified as an xarray.Dataset when using isinstance(…, xarray.Dataset) within OMFIT

Parameters
  • data_vars – see xarray.Dataset

  • coords – see xarray.Dataset

  • attrs – see xarray.Dataset

to_dataset()[source]

Return an xarray.Dataset representation of the data

Returns

xarray.Dataset

from_dataset(dataset)[source]

Create from xarray.Dataset representation of the data

omfit_classes.sortedDict.size_tree_objects(location)[source]

Returns file sizes of objects in the dictionary based on the size of their filename attribute

Parameters

location – string of the tree location to be analyzed

Returns

dictionary with locations sorted by size

omfit_classes.sortedDict.sorted_join_lists(a, b, favor_order_of_a=False, case_insensitive=False)[source]

Join two lists in a way that minimizes the distance between them and the merged list

Parameters
  • a – first list

  • b – second list

  • favor_order_of_a – favor order of list a over order of list b

  • case_insensitive – merge list in a case-insensitive way

Returns

merged list

omfit_classes.exceptions_omfit.print_last_exception(file=<_io.TextIOWrapper name='<stderr>' mode='w' encoding='UTF-8'>)[source]

This function prints the last exception that has occurred

omfit_classes.exceptions_omfit.print_stack()[source]
class omfit_classes.exceptions_omfit.doNotReportException[source]

Bases: object

Exceptions that inherit from this class will not trigger an email sent be set to the OMFIT developer team

exception omfit_classes.exceptions_omfit.EndOMFITpython(message='', *args, **kw)[source]

Bases: KeyboardInterrupt, omfit_classes.exceptions_omfit.doNotReportException

Class used to stop the running python script without reporting it to the OMFIT developer team

exception omfit_classes.exceptions_omfit.EndAllOMFITpython(message='', *args, **kw)[source]

Bases: KeyboardInterrupt, omfit_classes.exceptions_omfit.doNotReportException

Class used to stop the entire python workflow without reporting it to the OMFIT developer team

exception omfit_classes.exceptions_omfit.OMFITexception(message='', *args, **kw)[source]

Bases: Exception, omfit_classes.exceptions_omfit.doNotReportException

Class used to raise an exception in a user’s script without reporting it to the OMFIT developer team

exception omfit_classes.exceptions_omfit.ReturnCodeException[source]

Bases: RuntimeError

Class used to raise an exception when a code return code is !=0

omfit_classes.exceptions_omfit.signalHandler(signal=None, frame=None)[source]
class omfit_plot.Figure(*args, **kw)[source]

Bases: matplotlib.figure.Figure

figsize2-tuple of floats, default: :rc:`figure.figsize`

Figure dimension (width, height) in inches.

dpifloat, default: :rc:`figure.dpi`

Dots per inch.

facecolordefault: :rc:`figure.facecolor`

The figure patch facecolor.

edgecolordefault: :rc:`figure.edgecolor`

The figure patch edge color.

linewidthfloat

The linewidth of the frame (i.e. the edge linewidth of the figure patch).

frameonbool, default: :rc:`figure.frameon`

If False, suppress drawing the figure background patch.

subplotparsSubplotParams

Subplot parameters. If not given, the default subplot parameters :rc:`figure.subplot.*` are used.

tight_layoutbool or dict, default: :rc:`figure.autolayout`

If False use subplotpars. If True adjust subplot parameters using .tight_layout with default padding. When providing a dict containing the keys pad, w_pad, h_pad, and rect, the default .tight_layout paddings will be overridden.

constrained_layoutbool, default: :rc:`figure.constrained_layout.use`

If True use constrained layout to adjust positioning of plot elements. Like tight_layout, but designed to be more flexible. See /tutorials/intermediate/constrainedlayout_guide for examples. (Note: does not work with add_subplot or ~.pyplot.subplot2grid.)

colorbar(mappable, cax=None, ax=None, use_gridspec=True, **kw)[source]

Customized default grid_spec=True for consistency with tight_layout.

printlines(filename, squeeze=False)[source]

Print all data in line pyplot.plot(s) to text file.The x values will be taken from the line with the greatest number of points in the (first) axis, and other lines are interpolated if their x values do not match. Column labels are the line labels and xlabel.

Parameters

filename (str) – Path to print data to

Return type

bool

dmp(filename=None)[source]
Parameters

filename – file where to dump the h5 file

Returns

OMFITdmp object of the current figure

data_managment_plan_file(filename=None)[source]

Output the contents of the figure (self) to a hdf5 file given by filename

Parameters

filename – The path and basename of the file to save to (the extension is stripped, and ‘.h5’ is added) For the purpose of the GA data managment plan, these files can be uploaded directly to https://diii-d.gat.com/dmp

Returns

OMFITdmp object of the current figure

savefig(filename, saveDMP=True, PDFembedDMP=True, *args, **kw)[source]

Revised method to save the figure and the data to netcdf at the same time

Parameters
  • filename – filename to save to

  • saveDMP – whether to save also the DMP file. Failing to save as DMP does not raise an exception.

  • PDFembedDMP – whether to embed DMP file in PDF file (only applies if file is saved as PDF)

  • *args – arguments passed to the original matplotlib.figure.Figure.savefig function

  • **kw – kw arguments passed to the original matplotlib.figure.Figure.savefig function

Returns

Returns dmp object if save of DMP succeded. Returns None if user decided not to save DMP.

script()[source]
Returns

string with Python script to reproduce the figure (with DATA!)

OMFITpythonPlot(filename=None)[source]

generate OMFITpythonPlot script from figure (with DATA!)

Parameters

filename – filename for OMFITpythonPlot script

Returns

OMFITpythonPlot object

omfit_plot.close(which_fig)[source]

Wrapper for pyplot.close that closes FigureNotebooks when closing ‘all’

Close a figure window.

figNone or int or str or .Figure

The figure to close. There are a number of ways to specify this:

  • None: the current figure

  • .Figure: the given .Figure instance

  • int: a figure number

  • str: a figure name

  • ‘all’: all figures

class omfit_plot.savedFigure(fig)[source]

Bases: object

class omfit_plot.DraggableColorbar(cbar, mappable)[source]

Bases: object

connect()[source]

connect to all the events we need

on_press(event)[source]

on button press we will see if the mouse is over us and store some data

key_press(event)[source]
on_motion(event)[source]

on motion we will move the rect if the mouse is over us

on_release(event)[source]

on release we reset the press data

disconnect()[source]

disconnect all the stored connection ids

class omfit_plot.OMFITfigure(obj, figureButtons=True)[source]

Bases: object

addOMFITfigureToolbar(figureButtons=True)[source]
pin(event=None, fig=None, savefig=True, PDFembedDMP=True)[source]
Parameters

savefig – if False, save figure as object, if True, save figure as image.

email(event=None, fig=None, ext='PDF', saveDMP=False, PDFembedDMP=False)[source]
Parameters
  • ext – default ‘PDF’. figure format, e.g. PDF, PNG, JPG, etc.

  • saveDMP – default False, save HDF5 binary file [might have more data than shown in figure; but can be used for DIII-D Data Management Plan (DMP)]

  • PDFembedDMP – default False, embed DMP file in PDF

openPDF(event=None, fig=None, PDFembedDMP=False)[source]
help()[source]
crosshair(force=None)[source]
get(event=None)[source]
getObj(obj)[source]
selectAxes()[source]
selectPick(k)[source]
closePopup(event=None)[source]

this function closes the popup

poPopup(event)[source]

this function creates the popup which can then be populated

mvPopup(event)[source]
button_press_callback(event)[source]

when a mouse button is pressed

button_release_callback(event)[source]

when a mouse button is released

pick(event)[source]

this fucntion takes care of detecting which object was selected

closePropwindow(event=None)[source]

close the properties editing window

selectProperty(property)[source]

open the properties editing window

getProperty(property)[source]

retrieve the value of the property as seen by the user

setProperty(property, text)[source]
button_manager(event)[source]
objDelete(event=None)[source]
objCopy(event=None)[source]
objLegend(event=None)[source]
objText(event=None)[source]
objPaste(event=None)[source]
objAutoZoom(event=None, ax='')[source]
objSelect(event=None, forceDisable=False)[source]
key_press_callback(event)[source]
key_release_callback(event)[source]
static save_figure(self, saveDMP=True, PDFembedDMP=True, *args)[source]
pan(*args)[source]
zoom(*args)[source]
superzoomed = False
superzoom(event)[source]

Enlarge or restore the selected axis.

property active
class omfit_plot.FigureNotebook(nfig=0, name='', labels=[], geometry='710x710', figsize=(1, 1))[source]

Bases: object

Notebook of simple tabs containing figures.

Parameters
  • nfig – Number of initial tabs

  • name – String to display as the window title

  • labels – Sequence of labels to be used as the tab labels

  • geometry – size of the notebook

  • figsize – tuple, minimum maintained figuresize when resizing tabs

on_tab_change(event=None)[source]
email(event=None)[source]
openPDF(event=None)[source]
close()[source]

Close the FigureNotebook master window or a tab

add_figure(label='', num=None, fig=None, **fig_kwargs)[source]

Return the figure canvas for the tab with the given label, creating a new tab if that label does not yet exist. If fig is passed, then that fig is inserted into the figure.

static save_figure(self, _self, *args)[source]
subplots(nrows=1, ncols=1, label='', **subplots_kwargs)[source]

Adds a figure and axes using pyplot.subplots. Refer to pyplot.subplots for documentation

draw(ntab=None)[source]

Draw the canvas in the specified tab. None draws all.

savefig(filename='', **kw)[source]

Call savefig on each figure, with its label appended to filename

Parameters
  • filename – The fullpath+base of the filename to save

  • *kw – Passed to Figure.savefig

omfit_plot.figure(num=None, figsize=None, dpi=None, facecolor=None, edgecolor=None, frameon=True, FigureClass=<class 'omfit_plot.Figure'>, **kw)[source]
omfit_plot.colorbar(mappable=None, cax=None, ax=None, use_gridspec=True, **kw)[source]

Modified pyplot colorbar for default use_gridspec=True.

ORIGINAL DOCUMENTATION

Add a colorbar to a plot.

Function signatures for the pyplot interface; all but the first are also method signatures for the ~.Figure.colorbar method:

colorbar(**kwargs)
colorbar(mappable, **kwargs)
colorbar(mappable, cax=cax, **kwargs)
colorbar(mappable, ax=ax, **kwargs)
mappable

The matplotlib.cm.ScalarMappable (i.e., ~matplotlib.image.AxesImage, ~matplotlib.contour.ContourSet, etc.) described by this colorbar. This argument is mandatory for the .Figure.colorbar method but optional for the .pyplot.colorbar function, which sets the default to the current image.

Note that one can create a .ScalarMappable “on-the-fly” to generate colorbars not attached to a previously drawn artist, e.g.

fig.colorbar(cm.ScalarMappable(norm=norm, cmap=cmap), ax=ax)
cax~matplotlib.axes.Axes, optional

Axes into which the colorbar will be drawn.

ax~matplotlib.axes.Axes, list of Axes, optional

Parent axes from which space for a new colorbar axes will be stolen. If a list of axes is given they will all be resized to make room for the colorbar axes.

use_gridspecbool, optional

If cax is None, a new cax is created as an instance of Axes. If ax is an instance of Subplot and use_gridspec is True, cax is created as an instance of Subplot using the gridspec module.

colorbar~matplotlib.colorbar.Colorbar

See also its base class, ~matplotlib.colorbar.ColorbarBase.

Additional keyword arguments are of two kinds:

axes properties:

fractionfloat, default: 0.15

Fraction of original axes to use for colorbar.

shrinkfloat, default: 1.0

Fraction by which to multiply the size of the colorbar.

aspectfloat, default: 20

Ratio of long to short dimensions.

padfloat, default: 0.05 if vertical, 0.15 if horizontal

Fraction of original axes between colorbar and new image axes.

anchor(float, float), optional

The anchor point of the colorbar axes. Defaults to (0.0, 0.5) if vertical; (0.5, 1.0) if horizontal.

panchor(float, float), or False, optional

The anchor point of the colorbar parent axes. If False, the parent axes’ anchor will be unchanged. Defaults to (1.0, 0.5) if vertical; (0.5, 0.0) if horizontal.

colorbar properties:

Property

Description

extend

{‘neither’, ‘both’, ‘min’, ‘max’} If not ‘neither’, make pointed end(s) for out-of- range values. These are set for a given colormap using the colormap set_under and set_over methods.

extendfrac

{None, ‘auto’, length, lengths} If set to None, both the minimum and maximum triangular colorbar extensions with have a length of 5% of the interior colorbar length (this is the default setting). If set to ‘auto’, makes the triangular colorbar extensions the same lengths as the interior boxes (when spacing is set to ‘uniform’) or the same lengths as the respective adjacent interior boxes (when spacing is set to ‘proportional’). If a scalar, indicates the length of both the minimum and maximum triangular colorbar extensions as a fraction of the interior colorbar length. A two-element sequence of fractions may also be given, indicating the lengths of the minimum and maximum colorbar extensions respectively as a fraction of the interior colorbar length.

extendrect

bool If False the minimum and maximum colorbar extensions will be triangular (the default). If True the extensions will be rectangular.

spacing

{‘uniform’, ‘proportional’} Uniform spacing gives each discrete color the same space; proportional makes the space proportional to the data interval.

ticks

None or list of ticks or Locator If None, ticks are determined automatically from the input.

format

None or str or Formatter If None, ~.ticker.ScalarFormatter is used. If a format string is given, e.g., ‘%.3f’, that is used. An alternative ~.ticker.Formatter may be given instead.

drawedges

bool Whether to draw lines at color boundaries.

label

str The label on the colorbar’s long axis.

The following will probably be useful only in the context of indexed colors (that is, when the mappable has norm=NoNorm()), or other unusual circumstances.

Property

Description

boundaries

None or a sequence

values

None or a sequence which must be of length 1 less than the sequence of boundaries. For each region delimited by adjacent entries in boundaries, the color mapped to the corresponding value in values will be used.

If mappable is a ~.contour.ContourSet, its extend kwarg is included automatically.

The shrink kwarg provides a simple way to scale the colorbar with respect to the axes. Note that if cax is specified, it determines the size of the colorbar and shrink and aspect kwargs are ignored.

For more precise control, you can manually specify the positions of the axes objects in which the mappable and the colorbar are drawn. In this case, do not use any of the axes properties kwargs.

It is known that some vector graphics viewers (svg and pdf) renders white gaps between segments of the colorbar. This is due to bugs in the viewers, not Matplotlib. As a workaround, the colorbar can be rendered with overlapping segments:

cbar = colorbar()
cbar.solids.set_edgecolor("face")
draw()

However this has negative consequences in other circumstances, e.g. with semi-transparent images (alpha < 1) and colorbar extensions; therefore, this workaround is not used by default (see issue #1188).

omfit_plot.subplot(*args, **kw)[source]

Add a subplot to the current figure.

Wrapper of .Figure.add_subplot with a difference in behavior explained in the notes section.

Call signatures:

subplot(nrows, ncols, index, **kwargs)
subplot(pos, **kwargs)
subplot(**kwargs)
subplot(ax)
args : int, (int, int, *index), or .SubplotSpec, default: (1, 1, 1)

The position of the subplot described by one of

  • Three integers (nrows, ncols, index). The subplot will take the index position on a grid with nrows rows and ncols columns. index starts at 1 in the upper left corner and increases to the right. index can also be a two-tuple specifying the (first, last) indices (1-based, and including last) of the subplot, e.g., fig.add_subplot(3, 1, (1, 2)) makes a subplot that spans the upper 2/3 of the figure.

  • A 3-digit integer. The digits are interpreted as if given separately as three single-digit integers, i.e. fig.add_subplot(235) is the same as fig.add_subplot(2, 3, 5). Note that this can only be used if there are no more than 9 subplots.

  • A .SubplotSpec.

projection{None, ‘aitoff’, ‘hammer’, ‘lambert’, ‘mollweide’, ‘polar’, ‘rectilinear’, str}, optional

The projection type of the subplot (~.axes.Axes). str is the name of a custom projection, see ~matplotlib.projections. The default None results in a ‘rectilinear’ projection.

polarbool, default: False

If True, equivalent to projection=’polar’.

sharex, sharey~.axes.Axes, optional

Share the x or y ~matplotlib.axis with sharex and/or sharey. The axis will have the same limits, ticks, and scale as the axis of the shared axes.

labelstr

A label for the returned axes.

.axes.SubplotBase, or another subclass of ~.axes.Axes

The axes of the subplot. The returned axes base class depends on the projection used. It is ~.axes.Axes if rectilinear projection is used and .projections.polar.PolarAxes if polar projection is used. The returned axes is then a subplot subclass of the base class.

**kwargs

This method also takes the keyword arguments for the returned axes base class; except for the figure argument. The keyword arguments for the rectilinear base class ~.axes.Axes can be found in the following table but there might also be other keyword arguments if another projection is used.

Properties: adjustable: {‘box’, ‘datalim’} agg_filter: a filter function, which takes a (m, n, 3) float array and a dpi value, and returns a (m, n, 3) array alpha: float or None anchor: 2-tuple of floats or {‘C’, ‘SW’, ‘S’, ‘SE’, …} animated: bool aspect: {‘auto’} or num autoscale_on: bool autoscalex_on: bool autoscaley_on: bool axes_locator: Callable[[Axes, Renderer], Bbox] axisbelow: bool or ‘line’ box_aspect: None, or a number clip_box: .Bbox clip_on: bool clip_path: Patch or (Path, Transform) or None contains: unknown facecolor or fc: color figure: .Figure frame_on: bool gid: str in_layout: bool label: object navigate: bool navigate_mode: unknown path_effects: .AbstractPathEffect picker: None or bool or callable position: [left, bottom, width, height] or ~matplotlib.transforms.Bbox prop_cycle: unknown rasterization_zorder: float or None rasterized: bool or None sketch_params: (scale: float, length: float, randomness: float) snap: bool or None title: str transform: .Transform url: str visible: bool xbound: unknown xlabel: str xlim: (bottom: float, top: float) xmargin: float greater than -0.5 xscale: {“linear”, “log”, “symlog”, “logit”, …} xticklabels: unknown xticks: unknown ybound: unknown ylabel: str ylim: (bottom: float, top: float) ymargin: float greater than -0.5 yscale: {“linear”, “log”, “symlog”, “logit”, …} yticklabels: unknown yticks: unknown zorder: float

Creating a subplot will delete any pre-existing subplot that overlaps with it beyond sharing a boundary:

import matplotlib.pyplot as plt
# plot a line, implicitly creating a subplot(111)
plt.plot([1, 2, 3])
# now create a subplot which represents the top plot of a grid
# with 2 rows and 1 column. Since this subplot will overlap the
# first, the plot (and its axes) previously created, will be removed
plt.subplot(211)

If you do not want this behavior, use the .Figure.add_subplot method or the .pyplot.axes function instead.

If the figure already has a subplot with key (args, kwargs) then it will simply make that subplot current and return it. This behavior is deprecated. Meanwhile, if you do not want this behavior (i.e., you want to force the creation of a new subplot), you must use a unique set of args and kwargs. The axes label attribute has been exposed for this purpose: if you want two subplots that are otherwise identical to be added to the figure, make sure you give them unique labels.

In rare circumstances, .add_subplot may be called with a single argument, a subplot axes instance already created in the present figure but not in the figure’s list of axes.

.Figure.add_subplot .pyplot.subplots .pyplot.axes .Figure.subplots

plt.subplot(221)

# equivalent but more general
ax1=plt.subplot(2, 2, 1)

# add a subplot with no frame
ax2=plt.subplot(222, frameon=False)

# add a polar subplot
plt.subplot(223, projection='polar')

# add a red subplot that shares the x-axis with ax1
plt.subplot(224, sharex=ax1, facecolor='red')

# delete ax2 from the figure
plt.delaxes(ax2)

# add ax2 to the figure again
plt.subplot(ax2)
class omfit_plot.quickplot(x, y, ax=None)[source]

Bases: object

quickplot plots lots of data quickly

Parameters
  • x – x values

  • y – y values

  • ax – ax to plot on

It assumes the data points are dense in the x dimension compared to the screen resolution at all points in the plot. It will resize when the axes are clicked on.

get_ax_size()[source]
plot()[source]

API

omfit_classes.OMFITx.repr(value)[source]

repr modified to work for GUI functions

omfit_classes.OMFITx.repr_eval(location, preentry=None, collect=False)[source]

evaluate location and provide representation

class omfit_classes.OMFITx.GUI(pythonFile, relativeLocations, **kw)[source]

Bases: object

This class creates a new GUI. It is used internally by OMFIT when a OMFITpythonGUI object is executed

Parameters

pythonFile – OMFITpythonGUI object to be executed

Returns

None

update()[source]
omfit_classes.OMFITx.inputbox(prompt='Input box')[source]

Open a Dialog box to prompt for user input. Note: this is a blocking call.

Return input string if user chooses submit and None otherwise

omfit_classes.OMFITx.UpdateGUI(top=None)[source]

Function used to update users GUIs

Parameters

top – TopLevel tk GUI to be updated

omfit_classes.OMFITx.Refresh()[source]

Force a refresh of the OMFIT GUI by issuing a TkInter .update()

omfit_classes.OMFITx.CloseAllGUIs()[source]

Function for closing all users GUIs

omfit_classes.OMFITx.CompoundGUI(pythonFile, title=None, **kw)[source]

This method allows the creation of nested GUI.

Parameters
  • pythonFile – is meant to be an OMFITpythonGUI object in the OMFIT tree

  • title – title to appear in the compound GUI frame. If None, the location of the pythonFile object in the OMFIT tree will be shown. If an empty string, the compound GUI title is suppressed.

Returns

None

omfit_classes.OMFITx.Tab(name='')[source]

This method creates a Tab under which the successive GUI elements will be placed

Parameters

name – Name to assign to the TAB

Returns

None

omfit_classes.OMFITx.Entry(location, lbl=None, comment=None, updateGUI=False, help='', preentry=None, postcommand=None, check=[], multiline=False, norm=None, default=[], delete_if_default=False, url='', kwlabel={}, **kw)[source]

This method creates a GUI element of the entry type The background of the GUI gets colored green/red depending on whether the input by the user is a valid Python entry

Parameters
  • location – location in the OMFIT tree (notice that this is a string)

  • lbl – Label which is put on the left of the entry

  • comment – A comment which appears on top of the entry

  • updateGUI – Force a re-evaluation of the GUI script when this parameter is changed

  • help – help provided when user right-clicks on GUI element (adds GUI button)

  • preentry – function to pre-process the data at the OMFIT location to be displayed in the entry GUI element

  • postcommand – command to be executed after the value in the tree is updated. This command will receive the OMFIT location string as an input

  • check – function that returns whether what the user has entered in the entry GUI element is a valid entry. This will make the background colored yellow, and users will not be able to set the value.

  • default – Set the default value if the tree location does not exist (adds GUI button)

  • delete_if_default – Delete tree entry if the value is the default value

  • multiline – Force display of button for multiple-line text entry

  • norm – normalizes numeric variables (overrides preentry or postcommand)

  • url – open url in web-browser (adds GUI button)

  • kwlabel – keywords passed to ttk.Label

Returns

associated ttk.Entry object

omfit_classes.OMFITx.ComboBox(location, options, lbl=None, comment=None, updateGUI=False, state='readonly', help='', postcommand=None, check=[], default=[], url='', kwlabel={}, **kw)[source]

This method creates a GUI element of the combobox type. The background of the GUI gets colored green/red depending on whether the input by the user is a valid Python entry

Notice that this method can be used to set multiple entries at once: ComboBox([“root[‘asd’]”,”root[‘dsa’]”,”root[‘aaa’]”,],{‘’:[0,0,0],’a’:[1,1,0],’b’:[1,0,’***’]},’Test multi’,default=[0,0,0]) which comes very handy when complex/exclusive switch combinations need to be set in a namelist file, for example. Use the string *** to leave parameters unchanged.

Parameters
  • location – location in the OMFIT tree (notice that this is either a string or a list of strings)

  • options – possible options the user can choose from. This can be a list or a dictionary.

  • lbl – Label which is put on the left of the entry

  • comment – A comment which appears on top of the entry

  • updateGUI – Force a re-evaluation of the GUI script when this parameter is changed

  • state

    • ‘readonly’ (default) the user can not type in whatever he wants

    • ’normal’ allow user to type in

    • ’search’ allow searching for entries

  • help – help provided when user right-clicks on GUI element (adds GUI button)

  • postcommand – command to be executed after the value in the tree is updated. This command will receive the OMFIT location string as an input

  • check – function that returns whether what the user has entered in the entry GUI element is a valid entry. This will make the background colored yellow, and users will not be able to set the value.

  • default – Set the default value if the tree location does not exist (adds GUI button)

  • url – open url in web-browser (adds GUI button)

  • kwlabel – keywords passed to ttk.Label

Returns

associated TkInter combobox object

class omfit_classes.OMFITx.same_row[source]

Bases: object

Environment to place GUI elements on the same row

For example to place two buttons side by side:

>> with OMFITx.same_row(): >> OMFITx.Button(‘Run’, lambda: None) >> OMFITx.Button(‘Plot’, lambda: None)

class omfit_classes.OMFITx.same_tab(tab_name)[source]

Bases: object

Environment to place GUI elements within the same tab

For example to place two buttons in the same tab named ‘test’

>> with OMFITx.same_tab(‘test’): >> OMFITx.Button(‘Run’, lambda: None) >> OMFITx.Button(‘Plot’, lambda: None)

omfit_classes.OMFITx.Button(lbl, command, help='', url='', closeGUI=False, updateGUI=False, **kw)[source]

This method creates a GUI element of the button type

Parameters
  • lbl – the text to be written on the button

  • command – the command to be executed

  • help – help provided when user right-clicks on GUI element (adds GUI button)

  • url – open url in web-browser (adds GUI button)

  • updateGUI – Force a re-evaluation of the GUI script when this parameter is changed

  • closeGUI – Close current GUI after executing the command

  • **kw – extra keywords are passed to the ttk.Button

Returns

associated ttk.Button object

omfit_classes.OMFITx.Label(lbl, align='center', **kw)[source]

This method creates a GUI element of the label type

Parameters
  • lbl – the text to be written in the label

  • align – alignment of the text in the frame

Returns

associated ttk.Label object

omfit_classes.OMFITx.Separator(lbl=None, kwlabel={}, **kw)[source]

This method creates a TkInter separator object

Parameters
  • lbl – text to be written between separator lines

  • kwlabel – keywords passed to ttk.Label

  • **kw – keywords passed to ttk.Label

Returns

associated ttk.Label object

omfit_classes.OMFITx.FilePicker(location, lbl=None, comment=None, updateGUI=False, help='', postcommand=None, localRemote=True, transferRemoteFile=True, directory=False, action='open', tree=True, default=[], url='', kwlabel={}, init_directory_location=None, init_pattern_location=None, favorite_list_location=None, pattern_list_location=None, reveal_location=None, **kw)[source]

This method creates a GUI element of the filePicker type, which allows to pick a file/directory

Parameters
  • location – location in the OMFIT tree (notice that this is a string)

  • lbl – label to be shown near the button

  • help – help provided when user right-clicks on GUI element (adds GUI button)

  • postcommand – command to be executed after the value in the tree is updated. This command will receive the OMFIT location string as an input

  • updateGUI – Force a re-evaluation of the GUI script when this parameter is changed

  • localRemote – True: both, ‘local’: only local, ‘remote’: only remote

  • transferRemoteFile

    controls what goes into location

    • string with local filename (if transferRemoteFile==True)

    • string with the filename (if transferRemoteFile==False)

    • tuple with the filename,server,tunnel (if transferRemoteFile==None)

    if transferRemoteFile=True, then the file is transferred to a temporary folder

    if transferRemoteFile is a string, then it will be interpreted as the directory where to move the file

  • directory – whether it’s a directory or a file

  • action – ‘open’ or ‘save’

  • tree – load from OMFIT tree location

  • url – open url in web-browser (adds GUI button)

  • kwlabel – keywords passed to ttk.Label

  • init_directory_location – The contents of this location are used to set the initial directory for file searches. If a file name is specified the directory will be determined from the file name and this input ignored. Otherwise, if set this will be used to set the initial directory.

  • init_pattern_location – The default pattern is ‘*’. If this is specified then the contents of the tree location will replace the default intial pattern.

  • favorite_list_location – OMFIT tree location which contains a possibly empty list of favorite file directories. To keep with the general omfit approach this should be a string.

  • pattern_list_location – OMFIT tree location which contains a possibly empty list of favorite search patterns. To keep with the general omfit approach this should be a string.

  • reveal_location – location used for creation of the help (this is used internally by OMFIT, should not be used by users)

  • **kw – keywords passed to Entry object

Returns

associated ttk.Entry object

omfit_classes.OMFITx.ObjectPicker(location, lbl=None, objectType=None, objectKW={}, postcommand=None, unset_postcommand=None, kwlabel={}, init_directory_location=None, init_pattern_location=None, favorite_list_location=None, pattern_list_location=None, **kw)[source]

This helper method creates a GUI element of the objectPicker type, which allows to load objects in the tree.

If an object is already present at the location, then a button allows picking of a different object.

Notice that this GUI element will always call an updateGUI

Parameters
  • location – location in the OMFIT tree (notice that this is a string)

  • lbl – label to be shown near the button/object picker

  • objectType – class of the object that one wants to load (e.g. OMFITnamelist, OMFITgeqdsk, …) if objectType is None then the object selected with Tree is deepcopied

  • objectKW – keywords passed to the object

  • postcommand – command to be executed after the value in the tree is updated. This command will receive the OMFIT location string as an input.

  • unset_postcommand – command to be executed after the value in the tree is deleted. This command will receive the OMFIT location string as an input.

  • kwlabel – keywords passed to ttk.Label

  • init_directory_location – The contents of this location are used to set the initial directory for file searches. If a file name is specified the directory will be determined from the file name and this input ignored. Otherwise, if set this will be used to set the initial directory.

  • init_pattern_location – The default pattern is ‘*’. If this is specified then the contents of the tree location will replace the default intial pattern.

  • favorite_list_location – OMFIT tree location which contains a possibly empty list of favorite file directories. To keep with the general omfit approach this should be a string.

  • pattern_list_location – OMFIT tree location which contains a possibly empty list of favorite search patterns. To keep with the general omfit approach this should be a string.

  • **kw – extra keywords are pased to the FilePicker object

Returns

associated ttk.Entry object

omfit_classes.OMFITx.ModulePicker(location, modules=None, lbl=None, *args, **kw)[source]

This method creates a GUI element of the combobox type for the selection of modules within the OMFIT project.

Parameters
  • location – location in the OMFIT tree (notice that this is either a string or a list of strings)

  • modules – string or list of strings with IDs of the allowed modules. If modules is None all modules in OMFIT are listed

  • lbl – label to be shown near the combobox

  • load – list of two elements lists with module name and location where modules can be loaded eg. [[‘OMFITprofiles’,”root[‘OMFITprofiles’]”],[‘EFIT’,”OMFITmodules[-2][‘EFIT’]”],] Setting load=True will set loading of the modules as submodules

  • *args – arguments passed to OMFITx.ComboBox

  • **kw – keywords passed to OMFITx.ComboBox

Returns

returns from OMFITx.ComboBox

omfit_classes.OMFITx.TreeLocationPicker(location, lbl=None, comment=None, kwlabel={}, default=[], help='', url='', updateGUI=False, postcommand=None, check=None, base=None, **kw)[source]

This method creates a GUI element used to select a tree location The label of the GUI turns green/red if the input by the user is a valid OMFIT tree entry (non existing tree entries are allowed) The label of the GUI turns green/red if the input by the user does or doesn’t satisfy the check (non valid tree entries are NOT allowed)

Parameters
  • location – location in the OMFIT tree (notice that this is a string)

  • lbl – Label which is put on the left of the entry

  • comment – A comment which appears on top of the entry

  • kwlabel – keywords passed to ttk.Label

  • default – Set the default value if the tree location does not exist (adds GUI button)

  • help – help provided when user right-clicks on GUI element (adds GUI button)

  • url – open url in web-browser (adds GUI button)

  • updateGUI – Force a re-evaluation of the GUI script when this parameter is changed

  • postcommand – command to be executed after the value in the tree is updated. This command will receive the OMFIT location string as an input

  • check – function that returns whether what the user has entered in the entry GUI element is a valid entry This will make the label colored yellow, and users will not be able to set the value.

  • base – object in location with respect to which relative locations are evaluated

  • **kw – keywords passed to OneLineText object

Returns

associated ttk.Entry object

omfit_classes.OMFITx.CheckBox(location, lbl=None, comment=None, useInt=False, mapFalseTrue=[], updateGUI=False, help='', postcommand=None, default=[], url='', kwlabel={}, **kw)[source]

This method creates a GUI element of the checkbutton type

This method accepts a list of locations, labels and defaults

Parameters
  • location – location in the OMFIT tree (notice that this is a string)

  • lbl – Label which is put on the left of the entry

  • comment – A comment which appears on top of the entry

  • useInt – Use integers (1 or 0) instead of boolean (True or False)

  • mapFalseTrue – a 2 elements list, the first one for the unchecked, the second one for the checked button

  • updateGUI – Force a re-evaluation of the GUI script when this parameter is changed

  • help – help provided when user right-clicks on GUI element (adds GUI button)

  • postcommand – command to be executed after the value in the tree is updated. This command will receive the OMFIT location string as an input

  • default – Set the default value if the tree location does not exist (adds GUI button)

  • url – open url in web-browser (adds GUI button)

  • kwlabel – keywords passed to ttk.Label

  • **kw – extra keywords are pased to the Checkbutton object

Returns

associated TkInter checkbutton object

>>> OMFITx.CheckBox(["OMFIT['ck']","OMFIT['ck1']"],['hello','asd'],default=[False,True])
>>> OMFITx.CheckBox("OMFIT['ck']",'hello',default=False)
omfit_classes.OMFITx.ListEditor(location, options, lbl=None, default=None, unique=True, ordered=True, updateGUI=False, postcommand=None, only_valid_options=False, help='', url='', show_delete_button=False, max=None)[source]

GUI element to add or remove objects to a list Note: multiple items selection possible with the Shift and Ctrl keys

Parameters
  • location – location in the OMFIT tree (notice that this is a string).

  • options – possible options the user can choose from. This can be a tree location, a list, or a dictionary. If a dictinoary, then keys are shown in the GUI and values are set in the list. In order to use “show_delete_button”, this must be a string giving the location of a list in the tree.

  • lbl – Label which is put on the left of the entry

  • default – Set the default value if the tree location does not exist

  • unique – Do not allow repetitions in the list

  • ordered – Keep the same order as in the list of options If false, then buttons to move elements up/down are shown

  • updateGUI – Force a re-evaluation of the GUI script when this parameter is changed

  • postcommand – function to be called after a button is pushed. It is called as postcommand(location=location,button=button) where button is in [‘add’,’add_all’,’remove’,’remove_all’]

  • only_valid_options – list can only contain valid options

  • help – help provided when user right-clicks on GUI element (adds GUI button)

  • url – open url in web-browser (adds GUI button)

  • show_delete_button – bool: Show an additional button for deleting items from the left hand list

  • max – allow at most MAX choices

omfit_classes.OMFITx.Slider(location, start_stop_step, lbl=None, comment=None, digits=None, updateGUI=False, help='', preentry=None, postcommand=None, norm=None, default=[], url='', kwlabel={}, refresh_every=100, **kw)[source]

This method creates a GUI element of the slider type

Parameters
  • location – location in the OMFIT tree (notice that this is a string)

  • start_stop_step – list of tree elements with start/stop/step of the slider

  • lbl – Label which is put on the left of the entry

  • comment – A comment which appears on top of the entry

  • digits – How many digits to use (if None uses 3 digits if start_stop_step has floats or else 0 digits if these are integers)

  • updateGUI – Force a re-evaluation of the GUI script when this parameter is changed

  • help – help provided when user right-clicks on GUI element (adds GUI button)

  • preentry – function to pre-process the data at the OMFIT location to be displayed in the entry GUI element

  • postcommand – command to be executed after the value in the tree is updated. This command will receive the OMFIT location string as an input

  • default – Set the default value if the tree location does not exist (adds GUI button)

  • norm – normalizes numeric variables (overrides preentry or postcommand)

  • url – open url in web-browser (adds GUI button)

  • refresh_every – how often to call postcommand function (in ms)

  • kwlabel – keywords passed to ttk.Label

Returns

associated TtkScale object

omfit_classes.OMFITx.Lock(location, value=[], checkLock=False, clear=False)[source]

The lock method prevents users from using a GUI element which would affect a specific location in the OMFIT tree

Parameters
  • location – location in the OMFIT tree (notice that this is a string or a list of strings) If location is None, then all locks are cleared.

  • checkLock – False=set the lock | True=return the lock value

  • value – lock location at this value

  • clear – clear or set the lock

Returns

None if checkLock=False, otherwise True/False depending on value of the lock

omfit_classes.OMFITx.TitleGUI(title=None)[source]

Sets the title to the user gui window (if it’s not a compound GUI)

Parameters

title – string containing the title

Returns

None

omfit_classes.OMFITx.ShotTimeDevice(postcommand=None, showDevice=True, showShot=True, showTime=True, showRunID=False, multiShots=False, multiTimes=False, showSingleTime=False, checkDevice=None, checkShot=None, checkTime=None, checkRunID=None, subMillisecondTime=False, stopIfNotSet=True, updateGUI=True)[source]

This high level GUI allows setting of DEVICE/SHOT/TIME of each module (sets up OMFIT MainSettings if root[‘SETTINGS’][‘EXPERIMENT’][‘XXX’] is an expression)

Parameters
  • postcommand – command to be executed every time device,shot,time are changed (location is passed to postcommand)

  • showDevice – True/False show device section or list of suggested devices

  • showShot – True/False show shot section or list with list of suggested shots

  • showTime – True/False show time section or list with list of suggested times

  • showRunID – True/False show runID Entry

  • multiShots – True/False show single/multi shots

  • multiTimes – True/False show single/multi times

  • showSingleTime – True/False if multiTimes, still show single time

  • checkDevice – check if device user input satisfies condition

  • checkShot – check if shot user input satisfies condition

  • checkTime – check if time user input satisfies condition

  • checkRunID – check if runID user input satisfies condition

  • subMillisecondTime – Allow floats as times

  • stopIfNotSet – Stop GUI visualization if shot/time/device are not set

Returns

None

omfit_classes.OMFITx.CloseGUI()[source]

Function for closing the active user GUI

omfit_classes.OMFITx.End(what='single')[source]

End execution of OMFITpython script

Parameters

what

  • ‘single’ terminates the running script

  • ’all’ terminates the whole workflow

omfit_classes.OMFITx.Open(object)[source]

Open OMFITascii object in editor or OMFITweb in browser File extension behaviour can be specified in OMFIT[‘MainSettings’][‘SETUP’][‘EXTENSIONS’]

Parameters

object – OMFIT object or filename to be opened in external editor

omfit_classes.OMFITx.Figure(toolbar=True, returnFigure=False, fillX=False, **kw)[source]

Embed a matplotlib figure in an OMFIT GUI

Parameters
  • toolbar – [True] show/hide the figure toolbar

  • returnFigure – [False] function returns figure f or axis f.use_subplot(111)

  • fillX – [False] fill X dimension of screen

  • figsize – (5*2./3., 4*2./3.) figure size

  • **kw – keyword arguments passed to pyplot.Figure

omfit_classes.OMFITx.Dialog(*args, **kw)[source]

Display a dialog box and wait for user input

Parameters
  • message – the text to be written in the label

  • answers – list of possible answers

  • icon – “question”, “info”, “warning”, “error”

  • title – title of the frame

  • options – dictionary of True/False options that are displayed as checkbuttons in the dialog

  • entries – dictionary of string options that are displayed as entries in the dialog

Returns

return the answer chosen by the user (a dictionary if options keyword was passed)

omfit_classes.OMFITx.clc(tag=None)[source]
clear console (possible tags are)

INFO : forest green HIST : dark slate gray WARNING : DarkOrange2 HELP : PaleGreen4 STDERR : red3 STDOUT : black DEBUG : gold4 PROGRAM_OUT : blue PROGRAM_ERR : purple

Parameters

tag – specific tag to clear

omfit_classes.OMFITx.EditASCIIobject(location, lbl=None, comment=None, updateGUI=False, help='', postcommand=None, url='', **kw)[source]

This method creates a GUI element that edits ASCII files in the OMFIT tree Sample usage:

OMFITx.EditASCIIobject("root['INPUTS']['TRANSP']", 'edit namelist', postcommand=lambda location:eval(location).load())
Parameters
  • location – location of the ASCII OMFITobject in the OMFIT tree (notice that this is a string)

  • lbl – Label which is put on the left of the entry

  • comment – A comment which appears on top of the entry

  • updateGUI – Force a re-evaluation of the GUI script when this parameter is changed

  • help – help provided when user right-clicks on GUI element (adds GUI button)

  • postcommand – command to be executed after the value in the tree is updated. This command will receive the OMFIT location string as an input

  • url – open url in web-browser (adds GUI button)

Returns

associated ttk.Button object

class omfit_classes.OMFITx.FileDialog(directory=None, serverPicker='', server='localhost', tunnel='', pattern='*', default='', master=None, lockServer=False, focus='filterDirs', favorite_list_location=None, pattern_list_location=None, is_dir=False, title='File Browser')[source]

Bases: object

Standard remote file selection dialog – no checks on selected file.

Parameters
  • directory – directory where to start browsing

  • serverPicker – serverPicker wins over server/tunnel settings serverpicker=None will reuse latest server/tunnel that the user browsed to

  • server – server

  • tunnel – tunnel

  • pattern – glob regular expression for files selection

  • default – default filename selection

  • master – Tkinter master GUI

  • lockServer – allow users to change server settings

  • focus – what to focus in the GUI (‘filterDirs’,’filterFiles’)

  • favorite_list_location – OMFIT tree location which contains a possibly empty list of favorite file directories. To keep with the general omfit approach this should be a string.

  • pattern_list_location – OMFIT tree location which contains a possibly empty list of favorite search patterns. To keep with the general omfit approach this should be a string.

  • is_dir – (bool) Whether the requested file is a directory

go(directory=None)[source]
quit(how=None)[source]
dirs_double_event()[source]
dirs_back_event()[source]
files_select_event()[source]
ok_command()[source]
remote_command(command)[source]
filter_command(dir=None)[source]
get_filter()[source]
get_selection()[source]
set_filter(dir, pat)[source]
set_selection(file)[source]
manage_list(fav_list, obj, op)[source]
class omfit_classes.OMFITx.LoadFileDialog(*args, **kw)[source]

Bases: omfit_classes.OMFITx.FileDialog

File selection dialog which checks that the file exists.

ok_command()[source]
class omfit_classes.OMFITx.SaveFileDialog(directory=None, serverPicker='', server='localhost', tunnel='', pattern='*', default='', master=None, lockServer=False, focus='filterDirs', favorite_list_location=None, pattern_list_location=None, is_dir=False, title='File Browser')[source]

Bases: omfit_classes.OMFITx.FileDialog

File selection dialog which checks that the file exists before saving.

ok_command()[source]
omfit_classes.OMFITx.remoteFile(parent=None, transferRemoteFile=True, remoteFilename=None, server=None, tunnel=None, init_directory_location=None, init_pattern_location=None, favorite_list_location=None, pattern_list_location=None, is_dir=False)[source]

Opens up a dialogue asking filename, server/tunnel for remote file transfer This function is mostly used within the framework; for use in OMFIT GUI scripts please consider using the OMFITx.FilePicker and OMFITx.ObjectPicker functions instead.

Parameters
  • parent – Tkinter parent GUI

  • transferRemoteFile – [True,False,None] if True the remote file is transferred to the OMFITcwd directory

  • remoteFilename – initial string for remote filename

  • server – initial string for server

  • tunnel – initial string for tunnel

  • init_directory_location – The contents of this location are used to set the initial directory for file searches. If a file name is specified the directory will be determined from the file name and this input ignored. Otherwise, if set this will be used to set the initial directory.

  • init_pattern_location – The default pattern is ‘*’. If this is specified then the contents of the tree location will replace the default intial pattern.

  • favorite_list_location – OMFIT tree location which contains a possibly empty list of favorite file directories. To keep with the general omfit approach this should be a string.

  • pattern_list_location – OMFIT tree location which contains a possibly empty list of favorite search patterns. To keep with the general omfit approach this should be a string.

Returns

is controlled with transferRemoteFile parameter

  • string with local filename (if transferRemoteFile==True)

  • string with the filename (if transferRemoteFile==False)

  • tuple with the filename,server,tunnel (if transferRemoteFile==None)

omfit_classes.OMFITx.remote_sysinfo(server, tunnel='', quiet=False)[source]

This function retrieves information from a remote server (like the shell which is running there):

{'ARG_MAX': 4611686018427387903,
 'QSTAT': '',
 'SQUEUE': '/opt/slurm/default/bin/squeue',
 'environment': OMFITenv([]),
 'id': 6216321643098941518,
 'login': ['.cshrc', '.login'],
 'logout': ['.logout'],
 'shell': 'csh',
 'shell_path': '/bin/csh',
 'sysinfo': 'csh\nARG_MAX=4611686018427387903\nQSTAT=\nSQUEUE=/opt/slurm/default/bin/squeue\necho: No match.'
}

Information from the remote server is stored in a dictionary

Parameters
  • server – remote server

  • tunnel – via tunnel

  • quiet – suppress output or not

Returns

dictionary with info from the server

omfit_classes.OMFITx.manage_user_errors(command, reportUsersErrorByEmail=False, **kw)[source]

This method wraps around users calls of _OMFITpython scripts and manages printout of errors

Parameters

command – command to be executed

Returns

whatever the command returned

omfit_classes.OMFITx.execute(command_line, interactive_input=None, ignoreReturnCode=False, std_out=None, std_err=None, quiet=False, arguments='', script=None, use_bang_command='OMFIT_run_command.sh', progressFunction=None, extraButtons=None)[source]

This function allows execution of commands on the local workstation.

Parameters
  • command_line – string to be executed locally

  • interactive_input – interactive input to be passed to the command

  • ignoreReturnCode – ignore return code of the command

  • std_out – if a list is passed (e.g. []), the stdout of the program will be put there line by line

  • std_err – if a list is passed (e.g. []), the stderr of the program will be put there line by line

  • quiet – print command to screen or not

  • arguments – arguments that are passed to the command_line

  • script – string with script to be executed. script option substitutes %s with the automatically generated name of the script if script is a list or a tuple, then the first item should be the script itself and the second should be the script name

  • use_bang_command – Execute commands via OMFIT_run_command.sh script (useful to execute scripts within a given shell: #!/bin/…) If use_bang_command is a string, then the run script will take that filename. Notice that setting use_bang_command=True is not safe for multiple processes running in the same directory.

  • progressFunction – user function to which the std-out of the process is passed and returns values from 0 to 100 to indicate progress towards completion

  • extraButtons – dictionary with key/function that is used to add extra buttons to the GUI. The function receives a dictionary with the process std_out and pid

Returns

return code of the command

omfit_classes.OMFITx.remote_execute(server, command_line, remotedir, tunnel=None, interactive_input=None, ignoreReturnCode=False, std_out=None, std_err=None, quiet=False, arguments='', script=None, forceRemote=False, use_bang_command='OMFIT_run_command.sh', progressFunction=None, extraButtons=None, xterm=False)[source]

This function allows execution of commands on remote workstations. It has the logic to check if the remote workstation is the local workstation and in that case executes locally.

Parameters
  • server – server to connect and execute the command

  • command_line – string to be executed remotely (NOTE that if server=’’, the command is executed locally in the local directory)

  • remotedir – remote working directory, if remote directory does not exist it will be created

  • tunnel – tunnel to go through to connect to the server

  • interactive_input – interactive input to be passed to the command

  • ignoreReturnCode – ignore return code of the command

  • std_out – if a list is passed (e.g. []), the stdout of the program will be put there line by line

  • std_err – if a list is passed (e.g. []), the stderr of the program will be put there line by line

  • quiet – print command to screen or not

  • arguments – arguments that are passed to the command_line

  • script – string with script to be executed. script option substitutes %s with the automatically generated name of the script if script is a list or a tuple, then the first item should be the script itself and the second should be the script name

  • forceRemote – force remote connection even if server is localhost

  • use_bang_command – execute commands via OMFIT_run_command.sh script (useful to execute scripts within a given shell: #!/bin/…) If use_bang_command is a string, then the run script will take that filename. Notice that setting use_bang_command=True is not safe for multiple processes running in the same directory.

  • progressFunction – user function to which the std-out of the process is passed and returns values from 0 to 100 to indicate progress towards completion

  • extraButtons – dictionary with key/function that is used to add extra buttons to the GUI. The function receives a dictionary with the process std_out and pid

  • xterm – if True, launch command in its own xterm

Returns

return code of the command

omfit_classes.OMFITx.remote_upsync(server, local, remote, tunnel=None, ignoreReturnCode=False, keepRelativeDirectoryStructure='', quiet=False)[source]

Function to upload files/directories to remote server (possibly via tunnel connection)

NOTE: this function relies on rsync. There is no way to arbitrarily rename files with rsync. All rsync can do is move files to a different directory.

Parameters
  • server – server to connect and execute the command

  • local – local file(s) (string or list strings) to upsync

  • remote – remote directory or file to save files to

  • tunnel – tunnel to go through to connect to the server

  • ignoreReturnCode – whether to ignore return code of the rsync command

  • keepRelativeDirectoryStructure – string with common based directory of the locals files to be removed (usually equals local_dir)

  • quiet – print command to screen or not

Returns

return code of the rsync command (or True if keepRelativeDirectoryStructure and ignoreReturnCode and some rsync fail)

omfit_classes.OMFITx.remote_downsync(server, remote, local, tunnel=None, ignoreReturnCode=False, keepRelativeDirectoryStructure='', quiet=False)[source]

Function to download files/directories from remote server (possibly via tunnel connection)

NOTE: this function relies on rsync. There is no way to arbitrarily rename files with rsync. All rsync can do is move files to a different directory.

Parameters
  • server – server to connect and execute the command

  • remote – remote file(s) (string or list strings) to downsync

  • local – local directory or file to save files to

  • tunnel – tunnel to go through to connect to the server

  • ignoreReturnCode – whether to ignore return code of the rsync command

  • keepRelativeDirectoryStructure – string with common based directory of the remote files to be removed (usually equals remote_dir)

  • quiet – print command to screen or not

  • use_scp – (bool) If this flag is True remote_downsync will be executed with “scp” instead of “rsync”. Use for increased download speed. (default: False)

Returns

return code of the rsync command (or True if keepRelativeDirectoryStructure and ignoreReturnCode and some rsync fail)

omfit_classes.OMFITx.mdsvalue(server, treename=None, shot=None, TDI=None)[source]
Parameters
  • server – MDS+ server to connect to

  • treename – name of the MDS+ tree

  • shot – shot number

  • TDI – TDI string to be executed

Returns

result of TDI command

class omfit_classes.OMFITx.IDL(module_root, server=None, tunnel=None, executable=None, workdir=None, remotedir=None, clean=False)[source]

Bases: object

This class provides a live IDL session via the pidly module: https://pypi.python.org/pypi/pyIDL/ In practice this class wraps the pidly.IDL session so that it can handle SERVERS remote connections (including tunneling) and directory management the OMFIT way. The IDL executable is taken from the idl entry of this server under OMFIT[‘MainSettings’][‘SERVER’].

Local and remote working directories are specified in root[‘SETTINGS’][‘SETUP’][‘workDir’] and root[‘SETTINGS’][‘REMOTE_SETUP’][‘workDir’].

Server and tunnel are specified in root[‘SETTINGS’][‘REMOTE_SETUP’][‘server’] and root[‘SETTINGS’][‘REMOTE_SETUP’][‘tunnel’].

If the tunnel is an empty string, the connection to the remote server is direct. If server is an empty string, everything will occur locally and the remote working directory will be ignored.

Parameters
  • module_root – root of the module (e.g. root)

  • server – override module server

  • tunnel – override module tunnel

  • executable – override the executable is taken from the idl entry of this server under OMFIT[‘MainSettings’][‘SERVER’].

  • workdir – override module local working directory

  • remotedir – override module remote working directory

  • clean

    clear local/remote working directories

    • ”local”: clean local working directory only

    • ”local_force”: force clean local working directory only

    • ”remote”: clean remote working directory only

    • ”remote_force”: force clean remote working directory only

    • True: clean both

    • ”force”: force clean both

    • False: clean neither [DEFAULT]

>>> idl=OMFITx.IDL(OMFIT['EFIT'])
>>> idl('$pwd')
>>> idl('x = total([1, 1], /int)')
>>> print(idl.x)
>>> tmp=OMFITgeqdsk(OMFITsrc+'/../samples/g133221.01000')
>>> idl.upsync([tmp])
>>> idl.downsync(['g133221.01000'])
upsync(inputs=[], keepRelativeDirectoryStructure=False, ignoreReturnCode=False, quiet=False)[source]

Function used to upload files from the local working directory to remote IDL directory

Parameters
  • inputs – list of input objects or path to files, which will be deployed in the local or remote working directory. To deploy objects with a different name one can specify tuples (inputObject,’deployName’)

  • ignoreReturnCode – whether to ignore return code of the rsync command

  • quiet – print command to screen or not

downsync(outputs=[], ignoreReturnCode=False, keepRelativeDirectoryStructure=False, quiet=False)[source]

Function used to download files from the remote IDL directory to the local working directory

Parameters
  • outputs – list of output files which will be fetched from the remote directory

  • keepRelativeDirectoryStructure – [True/False] keep relative directory structure of the remote files

  • ignoreReturnCode – whether to ignore return code of the rsync command

  • quiet – print command to screen or not

omfit_classes.OMFITx.initWorkdir(module_root, server=None, tunnel=None, workdir=None, remotedir=None, clean=True, quiet=False)[source]

High level function to simplify initialization of directories within a module. This function will: 1) Create and clear local and remote working directories 2) Change directory to local working directory

Server and tunnel are specified in root[‘SETTINGS’][‘REMOTE_SETUP’][‘server’] and root[‘SETTINGS’][‘REMOTE_SETUP’][‘tunnel’]

Local and remote working directories are specified in root[‘SETTINGS’][‘SETUP’][‘workDir’] and root[‘SETTINGS’][‘REMOTE_SETUP’][‘workDir’]

Parameters
  • module_root – root of the module

  • server – string that overrides module server

  • tunnel – string that overrides module tunnel

  • workdir – string that overrides module local working directory

  • remotedir – string that overrides module remote working directory

  • clean

    clear local/remote working directories

    • ”local”: clean local working directory only

    • ”local_force”: force clean local working directory only

    • ”remote”: clean remote working directory only

    • ”remote_force”: force clean remote working directory only

    • True: clean both

    • ”force”: force clean both

    • False: clean neither

  • quiet – print command to screen or not

Returns

strings for local and remote directories (None if there was a problem in either one)

omfit_classes.OMFITx.executable(module_root=None, inputs=[], outputs=[], clean=True, interactive_input=None, server=None, tunnel=None, executable=None, workdir=None, remotedir=None, ignoreReturnCode=True, std_out=None, std_err=None, quiet=False, queued=False, keepRelativeDirectoryStructure=True, arguments='', script=None, forceRemote=False, progressFunction=None, use_bang_command='OMFIT_run_command.sh', extraButtons=None, xterm=False, clean_after=False)[source]

High level function that simplifies local/remote execution of software within a module.

This function will: 1. cd to the local working directory 2. Clear local/remote working directories [True] by default 3. Deploy the the “input” objects to local working directory 4. Upload files them remotely 5. Executes the software 6. Download “output” files to local working directory

Executable command is specified in root[‘SETTINGS’][‘SETUP’][‘executable’]

Local and remote working directories are specified in root[‘SETTINGS’][‘SETUP’][‘workDir’] and root[‘SETTINGS’][‘REMOTE_SETUP’][‘workDir’].

Server and tunnel are specified in root[‘SETTINGS’][‘REMOTE_SETUP’][‘server’] and root[‘SETTINGS’][‘REMOTE_SETUP’][‘tunnel’].

If the tunnel is an empty string, the connection to the remote server is direct. If server is an empty string, everything will occur locally and the remote working directory will be ignored.

Parameters
  • module_root – root of the module (e.g. root) used to set default values for ‘executable’,’server’,’tunnel’,’workdir’, ‘remotedir’ if module_root is None or module_root is OMFIT then ‘executable’,’server’,’tunnel’,’workdir’, ‘remotedir’ must be specified

  • inputs – list of input objects or path to files, which will be deployed in the local or remote working directory. To deploy objects with a different name one can specify tuples (inputObject,’deployName’)

  • outputs – list of output files which will be fetched from the remote directory

  • clean

    clear local/remote working directories

    • ”local”: clean local working directory only

    • ”local_force”: force clean local working directory only

    • ”remote”: clean remote working directory only

    • ”remote_force”: force clean remote working directory only

    • True: clean both [DEFAULT]

    • ”force”: force clean both

    • False: clean neither

  • arguments – arguments which will be passed to the executable

  • interactive_input – interactive input to be passed to the executable

  • server – override module server

  • tunnel – override module tunnel

  • executable – override module executable

  • workdir – override module local working directory

  • remotedir – override module remote working directory

  • ignoreReturnCode – ignore return code of executable

  • std_out – if a list is passed (e.g. []), the stdout of the program will be put there line by line; if a string is passed and bool(queued), this should indicate the path of the file that gives the stdout of the queued job

  • std_err – if a list is passed (e.g. []), the stderr of the program will be put there line by line; if a string is passed and bool(queued), this should indicate the path of the file that gives the stdout of the queued job

  • quiet – if True, suppress output to the command box

  • keepRelativeDirectoryStructure – [True/False] keep relative directory structure of the remote files

  • script – string with script to be executed. script option requires %s in the command line location where you want the script filename to appear if script is a list or a tuple, then the first item should be the script itself and the second should be the script name

  • forceRemote – force remote connection even if server is localhost

  • progressFunction – user function to which the std-out of the process is passed and returns values from 0 to 100 to indicate progress towards completion

  • queued – If cast as bool is True, invokes manage_job, using queued as qsub_findID keyword of manage_job, and also takes over std_out and std_err

  • use_bang_command – Execute commands via OMFIT_run_command.sh script (useful to execute scripts within a given shell: #!/bin/…) If use_bang_command is a string, then the run script will take that filename. Notice that setting use_bang_command=True is not safe for multiple processes running in the same directory.

  • extraButtons – dictionary with key/function that is used to add extra buttons to the GUI. The function receives a dictionary with the process std_out and pid

  • xterm – if True, launch the command in its own xterm

  • clean_after – (bool) If this flag is True, the remote directory will be removed once the outputs have been transferred to the local working directory. The remote directory have OMFIT in it’s name. (default: False)

  • use_scp – (bool) If this flag is True, the remote downsync of data will use the “scp” command instead of “rsync”. This should be used for increased download speed. (default: False)

Returns

return code of the command

omfit_classes.OMFITx.remote_python(module_root=None, python_script=None, target_function=None, namespace={}, executable=None, forceRemote=False, pickle_protocol=2, clean_local=False, **kw)[source]

Execute a Python target_function that is self-contained in a Python python_script, Useful to execute a Python module as a separate process on a local (or remote) workstation. This fuction relies on the OMFITx.executable function and additional keyword arguments are passed to it.

Parameters
  • module_root – root of the module (e.g. root) used to set default values for ‘executable’,’server’,’tunnel’,’workdir’, ‘remotedir’ if module_root is None or module_root is OMFIT then ‘executable’,’server’,’tunnel’,’workdir’, ‘remotedir’ must be specified

  • python_script – OMFITpythonTask (or string) to execute

  • target_function – function in the python_script that will be called

  • namespace – dictionary with variables passed to the target_function

  • executable – python executable (if None then is set based on SERVER)

  • forceRemote – force remote connection even if server is localhost

  • pickle_protocol – pickle protocol version (use 2 for Python2/3 compatibility)

  • clean_local – (bool) If Flag is True, this cleans and deletes the local working directory after result to be returned has been loaded into memory. The directory must have OMFIT somewhere in the name as a safety measure. (default: False).

  • **kw – additional arguments are passed to the underlying OMFITx.executable function

Returns

returns the output of the target_function

class omfit_classes.OMFITx.manage_job(module_root, qsub_job, server=None, tunnel=None, remotedir=None, qsub_findID='(?i)\\\'[^\\\']+\\\'|"[^"]+"|([0-9]{3,}[\\[\\]\\.\\w-]*)\\s*', qstat_command=None, qstat_findStatus='(?i)\\s+[rqwecpd]{1,2}\\s+', qdel_command=None)[source]

Bases: object

identify_queuing_system()[source]

This function identifies qstat and qdel commands to be used, if these were not specified

Returns

None

qstat(quiet=True, sleep=0, std_out=None, std_err=None)[source]

qstat command (or equivalent)

Parameters
  • quiet

    controls the output that is displayed to screen

    • False: prints the full output of the command

    • select: prints only the line involving the right jobID

    • True: nothing gets printed

  • sleep – grace time in seconds before checking

Returns

the status of the job, if the jobID is found in the output of qstat. Otherwise None.

qdel(std_out=None, std_err=None)[source]

qdel command (or equivalent)

wait(timeout=None, sleep=10)[source]

Wait for a job to finish

Parameters
  • timeout – Timeout in seconds after which the wait function will return

  • sleep – Interval in seconds between checks

Returns

string with the last seen job status. If the job was not there, then a ‘?’ is returned

wait_print(output_file='', error_file='', progressFunction=None, extraButtons=None)[source]

Wait for a job to finish and follow the output file and print the error file that the job generates

Parameters
  • output_file – output file that will be followed until the job ends (on the std_out)

  • error_file – error file that will be printed when the job ends (on the std_err)

  • progressFunction – user function to which the std-out of the process is passed and returns values from 0 to 100 to indicate progress towards completion

  • extraButtons – dictionary with key/function that is used to add extra buttons to the GUI. The function receives a dictionary with the process std_out and pid

omfit_classes.OMFITx.submit_job(module_root, batch_command, environment='', partition='', partition_flag=None, ntasks=1, nproc_per_task=1, job_time='1', memory='2GB', batch_type='SLURM', batch_option='', out_name='Job', **kw)[source]

Launch a (SLURM,PBS) job

Parameters
  • module_root – The module instance from which servers, etc. are culled

  • batch_command – A multi-line string or list of strings that should be executed

  • environment – A string to be executed to set up the environment before launching the batch job

  • partition – A string to be inserted into the batch script that indicates which partition(s) (comma separated) to run on; if None, execute batch_command serially

  • partition_flag – A string to be inserted before the partion names which matches the system configuration (e.g. -p, –qos)

  • nproc_per_task – Number of processors to be used by each line of batch_command

  • job_time – Max wall time of each line of batch_command - see sbatch –time option (default 1 minute)

  • memory – Max memory usage of each cpu utilized by batch_command - see sbatch –mem-per-cpu option (default 2GB)

  • batch_type – Type of batch system (SLURM, PBS)

  • batch_option

    A string specifying any additional batch options in the file header; It is inserted in raw form after the other batch options, so should include #SBATCH or #PBS if it is a batch type option, and it could be a multiline string of options

    (expected to contain the relevant #{SBATCH,PBS})

  • out_name – Name used for the output and error files

  • **kw – All other keywords are passed to OMFITx.executable

Returns

None

omfit_classes.OMFITx.job_array(module_root, batch_lines, environment='', partition='', ntasks=1, nproc_per_task=1, job_time='1', memory='2GB', batch_type='SLURM', partition_flag=None, batch_option='', **kw)[source]

Launch a (SLURM,PBS) job array

Parameters
  • module_root – The module instance from which servers, etc. are culled

  • batch_lines – A multi-line string or list of strings that should be executed in parallel

  • environment – A string to be executed to set up the environment before launching the batch job

  • partition – A string to be inserted into the batch script that indicates which partition(s) (comma separated) to run on; if None, execute batch_lines serially

  • partition_flag – A string to be inserted before the partion names which matches the system configuration (e.g. -p, –qos)

  • nproc_per_task – Number of processors to be used by each line of batch_lines

  • job_time – Max wall time of each line of batch_lines - see sbatch –time option (default 1 minute)

  • memory – Max memory usage of each cpu utilized by batch_lines - see sbatch –mem-per-cpu option (default 2GB)

  • batch_type – Type of batch system (SLURM, PBS)

  • batch_option

    A string specifying any additional batch options in the file header; It is inserted in raw form after the other batch options, so should include #SBATCH or #PBS if it is a batch type option, and it could be a multiline string of options

    (expected to contain the relevant #{SBATCH,PBS})

  • **kw – All other keywords are passed to OMFITx.executable

Returns

None

omfit_classes.OMFITx.ide(module_root, ide, script=None, inputs=[], outputs=[], clean=True, interactive_input=None, server=None, tunnel=None, executable=None, workdir=None, remotedir=None, ignoreReturnCode=True, std_out=None, std_err=None, keepRelativeDirectoryStructure=True, arguments='')[source]

High level function that simplifies local/remote execution of Integrated Development Environments (IDEs)

This function will: 1. cd to the local working directory 2. Clear local/remote working directories 3. Deploy the the “input” objects to local working directory 4. Upload files them remotely 5. Executes the IDE 6. Download “output” files to local working directory

The IDE executable, server and tunnel, and local and remote working directory depend on the MainSettings[‘SERVER’][ide]

If the tunnel is an empty string, the connection to the remote server is direct. If server is an empty string, everything will occur locally and the remote working directory will be ignored.

Parameters
  • module_root – root of the module (e.g. root) or OMFIT itself

  • ide – what IDE to execute (e.g. idl or matlab)

  • inputs – list of input objects or path to files, which will be deployed in the local or remote working directory. To deploy objects with a different name one can specify tuples (inputObject,’deployName’)

  • outputs – list o output files which will be fetched from the remote directory

  • clean

    clear local/remote working directories

    • ”local”: clean local working directory only

    • ”local_force”: force clean local working directory only

    • ”remote”: clean remote working directory only

    • ”remote_force”: force clean remote working directory only

    • True: clean both

    • ”force”: force clean both

    • False: clean neither

  • arguments – arguments which will be passed to the executable

  • interactive_input – interactive input to be passed to the executable

  • server – override module server

  • tunnel – override module tunnel

  • executable – override module executable

  • workdir – override module local working directory

  • remotedir – override module remote working directory

  • ignoreReturnCode – ignore return code of executable

  • std_out – if a list is passed (e.g. []), the stdout of the program will be put there line by line

  • std_err – if a list is passed (e.g. []), the stderr of the program will be put there line by line

  • script – string with script to be executed. script option requires that %s in the command line location where you want the script filename to appear

Returns

return code of the command

class omfit_classes.OMFITx.archive(module_root, server=None, tunnel=None, remotedir=None, storedir=None, store_command=None, restore_command=None)[source]

Bases: object

High level function that simplifies archival of simulations files

Parameters
  • module_root – root of the module (e.g. root)

  • server – override module server

  • tunnel – override module tunnel

  • storedir – directory where ZIP files are stored if storedir is None, then this will be sought under module_root[‘SETTINGS’][‘REMOTE_SETUP’][‘storedir’] and finally under SERVER[server][‘storedir’]

  • store_command – (optional) user-defined store command issued to store data (eg. for HPSS usage). Strings {remotedir} {storedir} and {filename} are substituted with actual remotedir, storedir and filename if store_command is None, then this will be sought under module_root[‘SETTINGS’][‘REMOTE_SETUP’][‘store_command’] and finally under SERVER[server][‘store_command’]

  • restore_command – (optional) user-defined restore command issued to restore data (eg. for HPSS usage). Strings {remotedir} {storedir} and {filename} are substituted with actual remotedir, storedir and filename if restore_command is None, then this will be sought under module_root[‘SETTINGS’][‘REMOTE_SETUP’][‘restore_command’] and finally under SERVER[server][‘restore_command’]

store(remotedir, filename, quiet=False, background=False, force=False, store_command=None)[source]

Store remotedir to filename ZIP file

Parameters
  • remotedir – remote directory to archive (usually: root[‘SETTINGS’][‘REMOTE_SETUP’][‘workdir’]) This parameter needs to be specified because the working directory can change.

  • filename – filename to be used for archival

  • quiet – print store process to screen

  • background – put creation of ZIP archive in background (ignored if store_commnad is used)

  • force – force store even if remotedir does not have OMFIT substring in it

  • store_command – store command to be used

Returns

object instance

restore(remotedir, quiet=False, background=False, force=False, restore_command=None)[source]

Restore filename ZIP file to remotedir

Parameters
  • remotedir – remote directory to deflate to (usually: root[‘SETTINGS’][‘REMOTE_SETUP’][‘workdir’]) This parameter needs to be specified because the working directory can change.

  • quiet – print restore process to screen

  • background – put restore of ZIP archive in background (ignored if restore_commnad is used)

  • force – force restore even if remotedir does not have OMFIT substring in it

  • restore_command – restore command to be used

Returns

object instance

omfit_classes.OMFITx.glob(server='', file_filter='*', workDir='./', tunnel='', search_upstream=False)[source]

Returns list of files in remote directory

Params server

remote server

Params file_filter

regular expression to filter files on (* by default)

Params workdir

remote working directory

Params tunnel

remote tunnel to use

Params search_upstream

T/F: Search for the file in parent directories until it is found or / is reached.

omfit_classes.OMFITx.email(to, subject, message, attachments=None, sysinfo=True, **kw)[source]

Send an email

Parameters
  • to – must be one of 1) a single address as a string 2) a string of comma separated multiple address 3) a list of string addresses

  • subject – String

  • message – String

  • attachments – List of OMFITobjects or path to files

  • sysinfo – Include system info at the bottom of the message

  • **kw – Extra arguments passed to the send_email utility function

Returns

string that user can decide to print to screen

omfit_classes.OMFITx.f2py(filename, modname)[source]

Run f2py on filename to get modname.so and return modname

omfit_classes.OMFITx.titleGUI(*args, **kw)[source]
omfit_classes.OMFITx.lock(*args, **kw)[source]
omfit_classes.OMFITx.Checkbutton(*args, **kw)[source]
omfit_classes.OMFITx.CheckButton(*args, **kw)[source]
omfit_classes.OMFITx.compoundGUI(*args, **kw)[source]
omfit_classes.OMFITx.closeGUI(*args, **kw)[source]
omfit_classes.OMFITx.closeAllGUIs(*args, **kw)[source]
omfit_classes.OMFITx.updateGUI(*args, **kw)[source]
omfit_classes.OMFITx.end(*args, **kw)[source]
omfit_classes.OMFITx.mainSettings_ShotTimeDevice_GUI(*args, **kw)[source]
omfit_classes.OMFITx.AskUser(*args, **kw)[source]
omfit_classes.OMFITx.custom_ttk_style(type, **kw)[source]

Generates ttk styles dynamiecally and buffers them

Parameters
  • type – one of the ttk_styles

  • kw – ttk style configuration attributes

Returns

dynamically generated ttk style name

omfit_classes.OMFITx.update_gui_theme(location=None)[source]

One place to set the ttk theme, reset custom styles, and update main gui elements.

Returns

Utilities

utils_math

omfit_classes.utils_math.np_errors(**kw)[source]
omfit_classes.utils_math.np_ignored(f)
omfit_classes.utils_math.np_raised(f)
omfit_classes.utils_math.np_printed(f)
omfit_classes.utils_math.np_warned(f)
omfit_classes.utils_math.get_array_hash(y)[source]

Get the hash for an array. A hash is an fixed sized integer that identifies a particular value, comparing this integer is faster than comparing the whole array (if you already have it stored).

Example:

y1 = np.arange(3)
y2 = np.arange(3) + 1
assert(get_array_hash(y1) == get_array_hash(y2)) # will raise an error
Parameters

y – np.ndarray.

Returns

hash.

omfit_classes.utils_math.ismember(A, B)[source]

Mimics the Matlab ismember() function to look for occurrences of A into B

Parameters
  • A – number or list/array

  • B – number or list/array

Returns

returns lia, locb lists. lia: returns ‘True’ where the data in A is found in B and ‘False’ elsewhere. locb: contains the lowest index in B for each value in A that is a member of B. while it contains ‘None’ elsewhere (where A is not a member of B)

omfit_classes.utils_math.map1d(vecfun, ndarr, axis=- 1, args=[], **kwargs)[source]

Map a vector operation to a single axis of any multi-dimensional array.

input:
param vecfun

A function that takes a 1D vector input.

param ndarr

A np.ndarray.

param axis

Axis of the array on which vecfun will operate.

param *args

Additional arguments passed to vecfun.

param **kwargs

Additional key word arguments are also passed to vecfun.

output:

A new array of same dimensionality as ndarr.

omfit_classes.utils_math.uniform_resample_2D_line(X0, Y0, n)[source]

Resampling of 2D line with uniform distribution of points along the line

Parameters
  • X0 – input x coordinates of 2D path

  • Y0 – input y coordinates of 2D path

  • n – number of points for output 2d path

Returns

tuple of x,y uniformly resampled path

omfit_classes.utils_math.stepwise_data_index(y, broaden=0)[source]

This function returns the indices that one must use to reproduce the step-wise data with the minimum number of points. In the ascii-art below, it returns the indices of the crosses. The original data can then be reproduced by nearest-neighbor interpolation:

Y ^
  |              x.....x
  |        x....x
  |  x....x             x......x
  |
  0-----------------------------> X

Hint: can also be used to compress linearly varying data:

i = stepwise_data_index(np.gradient(y))

The original data can then be reproduced by linear interpolation.

Parameters
  • y – input array

  • broaden – return broaden points around the x’s

Returns

indices for compression

omfit_classes.utils_math.pack_points(n, x0, p)[source]

Packed points distribution between -1 and 1

Parameters
  • n – number of points

  • x0 – pack points around x0, a float between -1 and 1

  • p – packing proportional to p factor >0

Returns

packed points distribution between -1 and 1

omfit_classes.utils_math.simplify_polygon(x, y, tolerance=None, preserve_topology=True)[source]

Returns a simplified representation of a polygon

Parameters
  • x – array of x coordinates

  • y – array of y coordinates

  • tolerance – all points in the simplified object will be within the tolerance distance of the original geometry if tolerance is None, then a tolerance guess is returned

  • preserve_topology – by default a slower algorithm is used that preserves topology

Returns

x and y coordinates of simplified polygon geometry if tolerance is None, then a tolerance guess is returned

omfit_classes.utils_math.torecarray(data2D, names=None)[source]

Converts a 2D list of lists, or a dictionary of dictionaries with 1d elements to a recarray

Parameters
  • data2D – 2D input data

  • names – recarray columns names

Returns

recarray

omfit_classes.utils_math.dictdict_2_dictlist(data)[source]

dictionary of dictionaries with 1d elements, to dictionary of lists

Parameters

data – dictionary of dictionaries

Returns

dictionary of dictionaries

omfit_classes.utils_math.flatten_iterable(l)[source]

flattens iterable of iterables structures

Parameters

l – input iterable structure

Returns

generator with items in the structure

omfit_classes.utils_math.lists_loop_generator(*argv, **kwargs)[source]

This function generates a list of items combining the entries from multiple lists

Parameters
  • *args – multiple 1D lists (of equal length, or of 0 or 1 element length)

  • permute – output all element permutations

Returns

list of items combining the entries from multiple lists

Examples:

shots = 123
times = [1,2]
runs = 'a01'
shots, times, runs = lists_loop_generator(shots,times,runs)
for shot, time, run in zip(shots,times,runs):
    myCode(shot,time,run)

shots = [123,446]
times = [1,2]
runs = 'a01'
shots, times, runs = lists_loop_generator(shots,times,runs)

now permute will return lists of length 2 or length 4

omfit_classes.utils_math.is_uncertain(var)[source]
Parameters

var – Variable or array to test

Returns

True if input variable or array is uncertain

omfit_classes.utils_math.wrap_array_func(func)[source]
omfit_classes.utils_math.chunks(l, n)[source]

Yield successive n-sized chunks from l

omfit_classes.utils_math.unsorted_unique(x)[source]

Find the unique elements of an array The order of the elements is according to their first appearance

Parameters

x – input array or list

Returns

unique array or list (depending on whether the input was an array or a list)

omfit_classes.utils_math.closestIndex(my_list, my_number=0)[source]

Given a SORTED iterable (a numeric array or list of numbers) and a numeric scalar my_number, find the index of the number in the list that is closest to my_number

Parameters
  • my_list – Sorted iterable (list or array) to search for number closest to my_number

  • my_number – Number to get close to in my_list

Returns

Index of my_list element closest to my_number

Note

If two numbers are equally close, returns the index of the smallest number.

omfit_classes.utils_math.order_axes(M, order)[source]

arbitrary sort axes of a multidimensional array

Parameters
  • M – input multidimensional array

  • order – integers with list of axes to be ordered

Returns

M with reordered axes

omfit_classes.utils_math.qgriddata(px, py, data, X, Y, func=<function sum>, fill=nan)[source]
Parameters
  • px

  • py

  • data

  • X

  • Y

  • func

Returns

class omfit_classes.utils_math.RectBivariateSplineNaN(Z, R, Q, *args, **kw)[source]

Bases: object

ev(*args)[source]
class omfit_classes.utils_math.interp1e(x, y, *args, **kw)[source]

Bases: externalImports.interp1d

Shortcut for scipy.interpolate.interp1d with fill_value=’extrapolate’ and bounds_error=False as defaults

Interpolate a 1-D function.

x and y are arrays of values used to approximate some function f: y = f(x). This class returns a function whose call method uses interpolation to find the value of new points.

Parameters: x : (N,) array_like

A 1-D array of real values.

y(…,N,…) array_like

A N-D array of real values. The length of y along the interpolation axis must be equal to the length of x.

kindstr or int, optional

Specifies the kind of interpolation as a string or as an integer specifying the order of the spline interpolator to use. The string has to be one of ‘linear’, ‘nearest’, ‘nearest-up’, ‘zero’, ‘slinear’, ‘quadratic’, ‘cubic’, ‘previous’, or ‘next’. ‘zero’, ‘slinear’, ‘quadratic’ and ‘cubic’ refer to a spline interpolation of zeroth, first, second or third order; ‘previous’ and ‘next’ simply return the previous or next value of the point; ‘nearest-up’ and ‘nearest’ differ when interpolating half-integers (e.g. 0.5, 1.5) in that ‘nearest-up’ rounds up and ‘nearest’ rounds down. Default is ‘linear’.

axisint, optional

Specifies the axis of y along which to interpolate. Interpolation defaults to the last axis of y.

copybool, optional

If True, the class makes internal copies of x and y. If False, references to x and y are used. The default is to copy.

bounds_errorbool, optional

If True, a ValueError is raised any time interpolation is attempted on a value outside of the range of x (where extrapolation is necessary). If False, out of bounds values are assigned fill_value. By default, an error is raised unless fill_value="extrapolate".

fill_valuearray-like or (array-like, array_like) or “extrapolate”, optional
  • if a ndarray (or float), this value will be used to fill in for requested points outside of the data range. If not provided, then the default is NaN. The array-like must broadcast properly to the dimensions of the non-interpolation axes.

  • If a two-element tuple, then the first element is used as a fill value for x_new < x[0] and the second element is used for x_new > x[-1]. Anything that is not a 2-element tuple (e.g., list or ndarray, regardless of shape) is taken to be a single array-like argument meant to be used for both bounds as below, above = fill_value, fill_value.

    New in version 0.17.0.

  • If “extrapolate”, then points outside the data range will be extrapolated.

    New in version 0.17.0.

assume_sortedbool, optional

If False, values of x can be in any order and they are sorted first. If True, x has to be an array of monotonically increasing values.

Attributes: fill_value

Methods: __call__

See Also: splrep, splev

Spline interpolation/smoothing based on FITPACK.

UnivariateSpline : An object-oriented wrapper of the FITPACK routines. interp2d : 2-D interpolation

Calling interp1d with NaNs present in input values results in undefined behaviour.

Input values x and y must be convertible to float values like int or float.

Examples: >>> import matplotlib.pyplot as plt >>> from scipy import interpolate >>> x = np.arange(0, 10) >>> y = np.exp(-x/3.0) >>> f = interpolate.interp1d(x, y)

>>> xnew = np.arange(0, 9, 0.1)
>>> ynew = f(xnew)   # use interpolation function returned by `interp1d`
>>> plt.plot(x, y, 'o', xnew, ynew, '-')
>>> plt.show()

Initialize a 1-D linear interpolation class.

dtype
class omfit_classes.utils_math.interp1dPeriodic(t, y, *args, **kw)[source]

Bases: object

1D linear interpolation for periodic functions

Parameters
  • t – array Independent variable: Angle (rad). Doesn’t have to be on any particular interval as it will be unwrapped.

  • y – array matching length of t Dependent variable: values as a function of t

class omfit_classes.utils_math.uinterp1d(x, y, kind='linear', axis=- 1, copy=True, bounds_error=True, fill_value=nan, assume_sorted=False, std_pow=2, **kw)[source]

Bases: omfit_classes.utils_math.interp1e

Adjusted scipy.interpolate.interp1d (documented below) to interpolate the nominal_values and std_devs of an uncertainty array.

NOTE: uncertainty in the x data is neglected, only uncertainties in the y data are propagated.

Parameters

std_pow – float. Uncertainty is raised to this power, interpolated, then lowered. (Note std_pow=2 interpolates the variance, which is often used in fitting routines).

Additional arguments and key word arguments are as in interp1e (documented below).

Examples

>> x = np.linspace(0,2*np.pi,30) >> u = unumpy.uarray(np.cos(x),np.random.rand(len(x))) >> >> fi = utils.uinterp1d(x,u,std_pow=2) >> xnew = np.linspace(x[0],x[-1],1e3) >> unew = fi(xnew) >> >> f = figure(num=’uniterp example’) >> f.clf() >> ax = f.use_subplot(111) >> uerrorbar(x,u) >> uband(xnew,unew)

interp1e Documentation:

Shortcut for scipy.interpolate.interp1d with fill_value=’extrapolate’ and bounds_error=False as defaults

Interpolate a 1-D function.

x and y are arrays of values used to approximate some function f: y = f(x). This class returns a function whose call method uses interpolation to find the value of new points.

Parameters: x : (N,) array_like

A 1-D array of real values.

y(…,N,…) array_like

A N-D array of real values. The length of y along the interpolation axis must be equal to the length of x.

kindstr or int, optional

Specifies the kind of interpolation as a string or as an integer specifying the order of the spline interpolator to use. The string has to be one of ‘linear’, ‘nearest’, ‘nearest-up’, ‘zero’, ‘slinear’, ‘quadratic’, ‘cubic’, ‘previous’, or ‘next’. ‘zero’, ‘slinear’, ‘quadratic’ and ‘cubic’ refer to a spline interpolation of zeroth, first, second or third order; ‘previous’ and ‘next’ simply return the previous or next value of the point; ‘nearest-up’ and ‘nearest’ differ when interpolating half-integers (e.g. 0.5, 1.5) in that ‘nearest-up’ rounds up and ‘nearest’ rounds down. Default is ‘linear’.

axisint, optional

Specifies the axis of y along which to interpolate. Interpolation defaults to the last axis of y.

copybool, optional

If True, the class makes internal copies of x and y. If False, references to x and y are used. The default is to copy.

bounds_errorbool, optional

If True, a ValueError is raised any time interpolation is attempted on a value outside of the range of x (where extrapolation is necessary). If False, out of bounds values are assigned fill_value. By default, an error is raised unless fill_value="extrapolate".

fill_valuearray-like or (array-like, array_like) or “extrapolate”, optional
  • if a ndarray (or float), this value will be used to fill in for requested points outside of the data range. If not provided, then the default is NaN. The array-like must broadcast properly to the dimensions of the non-interpolation axes.

  • If a two-element tuple, then the first element is used as a fill value for x_new < x[0] and the second element is used for x_new > x[-1]. Anything that is not a 2-element tuple (e.g., list or ndarray, regardless of shape) is taken to be a single array-like argument meant to be used for both bounds as below, above = fill_value, fill_value.

    New in version 0.17.0.

  • If “extrapolate”, then points outside the data range will be extrapolated.

    New in version 0.17.0.

assume_sortedbool, optional

If False, values of x can be in any order and they are sorted first. If True, x has to be an array of monotonically increasing values.

Attributes: fill_value

Methods: __call__

See Also: splrep, splev

Spline interpolation/smoothing based on FITPACK.

UnivariateSpline : An object-oriented wrapper of the FITPACK routines. interp2d : 2-D interpolation

Calling interp1d with NaNs present in input values results in undefined behaviour.

Input values x and y must be convertible to float values like int or float.

Examples: >>> import matplotlib.pyplot as plt >>> from scipy import interpolate >>> x = np.arange(0, 10) >>> y = np.exp(-x/3.0) >>> f = interpolate.interp1d(x, y)

>>> xnew = np.arange(0, 9, 0.1)
>>> ynew = f(xnew)   # use interpolation function returned by `interp1d`
>>> plt.plot(x, y, 'o', xnew, ynew, '-')
>>> plt.show()

Initialize a 1-D linear interpolation class.

dtype
class omfit_classes.utils_math.uinterp1e(x, y, kind='linear', axis=- 1, copy=True, assume_sorted=False, std_pow=2, **kw)[source]

Bases: omfit_classes.utils_math.uinterp1d

Adjusted unterp1d to extrapolate

Arguments and key word arguments are as in uinterp1d (documented below).

Uinterp1d Documentation

Adjusted scipy.interpolate.interp1d (documented below) to interpolate the nominal_values and std_devs of an uncertainty array.

NOTE: uncertainty in the x data is neglected, only uncertainties in the y data are propagated.

Parameters

std_pow – float. Uncertainty is raised to this power, interpolated, then lowered. (Note std_pow=2 interpolates the variance, which is often used in fitting routines).

Additional arguments and key word arguments are as in interp1e (documented below).

Examples

>> x = np.linspace(0,2*np.pi,30) >> u = unumpy.uarray(np.cos(x),np.random.rand(len(x))) >> >> fi = utils.uinterp1d(x,u,std_pow=2) >> xnew = np.linspace(x[0],x[-1],1e3) >> unew = fi(xnew) >> >> f = figure(num=’uniterp example’) >> f.clf() >> ax = f.use_subplot(111) >> uerrorbar(x,u) >> uband(xnew,unew)

interp1e Documentation:

Shortcut for scipy.interpolate.interp1d with fill_value=’extrapolate’ and bounds_error=False as defaults

Interpolate a 1-D function.

x and y are arrays of values used to approximate some function f: y = f(x). This class returns a function whose call method uses interpolation to find the value of new points.

Parameters: x : (N,) array_like

A 1-D array of real values.

y(…,N,…) array_like

A N-D array of real values. The length of y along the interpolation axis must be equal to the length of x.

kindstr or int, optional

Specifies the kind of interpolation as a string or as an integer specifying the order of the spline interpolator to use. The string has to be one of ‘linear’, ‘nearest’, ‘nearest-up’, ‘zero’, ‘slinear’, ‘quadratic’, ‘cubic’, ‘previous’, or ‘next’. ‘zero’, ‘slinear’, ‘quadratic’ and ‘cubic’ refer to a spline interpolation of zeroth, first, second or third order; ‘previous’ and ‘next’ simply return the previous or next value of the point; ‘nearest-up’ and ‘nearest’ differ when interpolating half-integers (e.g. 0.5, 1.5) in that ‘nearest-up’ rounds up and ‘nearest’ rounds down. Default is ‘linear’.

axisint, optional

Specifies the axis of y along which to interpolate. Interpolation defaults to the last axis of y.

copybool, optional

If True, the class makes internal copies of x and y. If False, references to x and y are used. The default is to copy.

bounds_errorbool, optional

If True, a ValueError is raised any time interpolation is attempted on a value outside of the range of x (where extrapolation is necessary). If False, out of bounds values are assigned fill_value. By default, an error is raised unless fill_value="extrapolate".

fill_valuearray-like or (array-like, array_like) or “extrapolate”, optional
  • if a ndarray (or float), this value will be used to fill in for requested points outside of the data range. If not provided, then the default is NaN. The array-like must broadcast properly to the dimensions of the non-interpolation axes.

  • If a two-element tuple, then the first element is used as a fill value for x_new < x[0] and the second element is used for x_new > x[-1]. Anything that is not a 2-element tuple (e.g., list or ndarray, regardless of shape) is taken to be a single array-like argument meant to be used for both bounds as below, above = fill_value, fill_value.

    New in version 0.17.0.

  • If “extrapolate”, then points outside the data range will be extrapolated.

    New in version 0.17.0.

assume_sortedbool, optional

If False, values of x can be in any order and they are sorted first. If True, x has to be an array of monotonically increasing values.

Attributes: fill_value

Methods: __call__

See Also: splrep, splev

Spline interpolation/smoothing based on FITPACK.

UnivariateSpline : An object-oriented wrapper of the FITPACK routines. interp2d : 2-D interpolation

Calling interp1d with NaNs present in input values results in undefined behaviour.

Input values x and y must be convertible to float values like int or float.

Examples: >>> import matplotlib.pyplot as plt >>> from scipy import interpolate >>> x = np.arange(0, 10) >>> y = np.exp(-x/3.0) >>> f = interpolate.interp1d(x, y)

>>> xnew = np.arange(0, 9, 0.1)
>>> ynew = f(xnew)   # use interpolation function returned by `interp1d`
>>> plt.plot(x, y, 'o', xnew, ynew, '-')
>>> plt.show()

Initialize a 1-D linear interpolation class.

dtype
class omfit_classes.utils_math.URegularGridInterpolator(points, values, method='linear', bounds_error=True, fill_value=nan, std_pow=2, **kw)[source]

Bases: scipy.interpolate.interpolate.RegularGridInterpolator

Adjusted scipy.interpolate.RegularGridInterpolator (documented below) to interpolate the nominal_values and std_devs of an uncertainty array.

Parameters

std_pow – float. Uncertainty is raised to this power, interpolated, then lowered. (Note std_pow=2 interpolates the variance, which is often used in fitting routines).

Additional arguments and key word arguments are as in RegularGridInterpolator.

Examples

Make some sample 2D data

>> x = np.linspace(0,2*np.pi,30) >> y = np.linspace(0,2*np.pi,30) >> z = np.cos(x[:,np.newaxis]+y[np.newaxis,:]) >> u = unumpy.uarray(np.cos(x[:,np.newaxis]+y[np.newaxis,:]),np.random.rand(*z.shape))

Form interpolator

>> fi = URegularGridInterpolator((x,y),u,std_pow=2)

interpolate along the diagonal of (x,y)

>> xnew = np.linspace(x[0],x[-1],1e3) >> unew = fi(zip(xnew,xnew))

Compare original and interpolated values.

>> f = figure(num=’URegularGridInterpolator example’) >> f.clf() >> ax = f.use_subplot(111) >> uerrorbar(x,u.diagonal()) >> uband(xnew,unew)

Note the interpolated uncertainty between points is curved by std_pow=2. The curve is affected by the uncertainty of nearby off diagonal points (not shown).

RegularGridInterpolator Documentation

Interpolation on a regular grid in arbitrary dimensions

The data must be defined on a regular grid; the grid spacing however may be uneven. Linear and nearest-neighbor interpolation are supported. After setting up the interpolator object, the interpolation method (linear or nearest) may be chosen at each evaluation.

Parameters: points : tuple of ndarray of float, with shapes (m1, ), …, (mn, )

The points defining the regular grid in n dimensions.

valuesarray_like, shape (m1, …, mn, …)

The data on the regular grid in n dimensions.

methodstr, optional

The method of interpolation to perform. Supported are “linear” and “nearest”. This parameter will become the default for the object’s __call__ method. Default is “linear”.

bounds_errorbool, optional

If True, when interpolated values are requested outside of the domain of the input data, a ValueError is raised. If False, then fill_value is used.

fill_valuenumber, optional

If provided, the value to use for points outside of the interpolation domain. If None, values outside the domain are extrapolated.

Methods: __call__

Notes: Contrary to LinearNDInterpolator and NearestNDInterpolator, this class avoids expensive triangulation of the input data by taking advantage of the regular grid structure.

If any of points have a dimension of size 1, linear interpolation will return an array of nan values. Nearest-neighbor interpolation will work as usual in this case.

New in version 0.14.

Examples: Evaluate a simple example function on the points of a 3-D grid:

>>> from scipy.interpolate import RegularGridInterpolator
>>> def f(x, y, z):
...     return 2 * x**3 + 3 * y**2 - z
>>> x = np.linspace(1, 4, 11)
>>> y = np.linspace(4, 7, 22)
>>> z = np.linspace(7, 9, 33)
>>> data = f(*np.meshgrid(x, y, z, indexing='ij', sparse=True))

data is now a 3-D array with data[i,j,k] = f(x[i], y[j], z[k]). Next, define an interpolating function from this data:

>>> my_interpolating_function = RegularGridInterpolator((x, y, z), data)

Evaluate the interpolating function at the two points (x,y,z) = (2.1, 6.2, 8.3) and (3.3, 5.2, 7.1):

>>> pts = np.array([[2.1, 6.2, 8.3], [3.3, 5.2, 7.1]])
>>> my_interpolating_function(pts)
array([ 125.80469388,  146.30069388])

which is indeed a close approximation to [f(2.1, 6.2, 8.3), f(3.3, 5.2, 7.1)].

See also: NearestNDInterpolator : Nearest neighbor interpolation on unstructured

data in N dimensions

LinearNDInterpolatorPiecewise linear interpolant on unstructured data

in N dimensions

References: .. [1] Python package regulargrid by Johannes Buchner, see

2

Wikipedia, “Trilinear interpolation”, https://en.wikipedia.org/wiki/Trilinear_interpolation

3

Weiser, Alan, and Sergio E. Zarantonello. “A note on piecewise linear and multilinear table interpolation in many dimensions.” MATH. COMPUT. 50.181 (1988): 189-196. https://www.ams.org/journals/mcom/1988-50-181/S0025-5718-1988-0917826-0/S0025-5718-1988-0917826-0.pdf

omfit_classes.utils_math.cumtrapz(y, x=None, dx=1.0, axis=- 1, initial=0)[source]

This is a convenience function for scipy.integrate.cumtrapz. Notice that here initial=0 which is what one most often wants, rather than the initial=None, which is the default for the scipy function.

Cumulatively integrate y(x) using the composite trapezoidal rule. This is the right way to integrated derivatives quantities which were calculated with gradient. If a derivative was obtained with the diff command, then the cumsum command should be used for its integration.

Parameters
  • y – Values to integrate.

  • x – The coordinate to integrate along. If None (default), use spacing dx between consecutive elements in y.

  • dx – Spacing between elements of y. Only used if x is None

:param axis : Specifies the axis to cumulate. Default is -1 (last axis).

:param initialIf given, uses this value as the first value in the returned result.

Typically this value should be 0. If None, then no value at x[0] is returned and the returned array has one element less than y along the axis of integration.

Returns

The result of cumulative integration of y along axis. If initial is None, the shape is such that the axis of integration has one less value than y. If initial is given, the shape is equal to that of y.

omfit_classes.utils_math.deriv(x, y)[source]

This function returns the derivative of the 2nd order lagrange interpolating polynomial of y(x) When re-integrating, to recover the original values y use cumtrapz(dydx,x)

Parameters
  • x – x axis array

  • y – y axis array

Returns

dy/dx

omfit_classes.utils_math.reverse_enumerate(l)[source]
omfit_classes.utils_math.factors(n)[source]

get all the factors of a number

> print(list(factors(100)))

omfit_classes.utils_math.greatest_common_delta(input_int_array)[source]

Given an array of integers it returns the greatest uniform delta step

Parameters

input_int_array – array of integers

Returns

greatest uniform delta step

omfit_classes.utils_math.mad(x, axis=- 1, std_norm=True)[source]

Median absolute deviation, defined as 1.4826 * np.median(np.abs(np.median(x)-x))

Parameters
Returns

Median absolute deviation of the input data

omfit_classes.utils_math.mad_outliers(data, m=3.0, outliers_or_valid='valid')[source]

Function to identify outliers data based on median absolute deviation (mad) distance. Note: used median absolute deviation defined as 1.4826 * np.median(np.abs(np.median(x)-x))

Parameters
  • data – input data array (if a dictionary of arrays, the mad_outliers function is applied to each of the values in the dictionary)

  • m – mad distance multiplier from the median after which a point is considered an outlier

  • outliers_or_valid – return valid/outlier points (valid is default)

Returns

boolean array indicating which data points are a within m mad from the median value (i.e. the valid points)

omfit_classes.utils_math.bin_outliers(data, mincount, nbins, outliers_or_valid='valid')[source]

Function to identify outliers data based on binning of data. The algorythm bins the data in nbins and then considers valid data only the data that falls within the bins that have at least mincount counts.

Parameters
  • data – input data array (if a dictionary of arrays, the bin_outliers function is applied to each of the values in the dictionary)

  • mincount – minimum number of counts within a bin for data to be considered as valid

  • nbins – number of bins for binning of data

  • outliers_or_valid – return valid/outlier points (valid is default)

Returns

boolean array indicating which data points are a valid or not

omfit_classes.utils_math.powerlaw_fit(x, y)[source]

evaluates multidimensional power law for y based on inputs x

Parameters
  • x – 2D array of inputs [N,M]

  • y – 1D array of output [M]

Returns

power law coefficient p, fitting function f(p,x)

omfit_classes.utils_math.exp_no_overflow(arg, factor=1.0, minpad=1000.0, extra_pad=1000.0, return_big_small=False)[source]

Performs exp(arg) but first limits the value of arg to prevent floating math errors. Checks sys.float_info to so it can avoid floating overflow or underflow. Can be informed of factors you plan on multiplying with the exponential result later in order to make the limits more restrictive (the limits are “padded”) and avoid over/underflow later on in your code.

Parameters
  • arg – Argument of exponential function

  • factor – Factor that might be multiplied in to exponential function later. Adjusts limits to make them more restrictive and prevent overflow later on. The adjustment to limits is referred to as padding.

  • minpad – Force the padding to be at least a certain size.

  • extra_pad – Extra padding beyond what’s determined by minpad and factor

  • return_big_small – T/F: flag to just return big and small numbers. You may be able to speed up execution in repeated calls by getting appropriate limits, doing your own cropping, and looping exp() instead of looping exp_no_overflow(). Even in this case, exp_no_overflow() can help you pick good limits. Or if you don’t have time for any of that, you can probably use -70 and 70 as the limits on arg, which will get you to order 1e30.

Returns

exp(arg) with no math errors, or (big, small) if return_big_small is set

omfit_classes.utils_math.dimens(x)[source]

Returns the number of dimensions in a mixed list/array object. Handles mutli-level nested iterables.

From Craig Burgler on stack exchange https://stackoverflow.com/a/39255678/6605826

This would be useful, for example on a case like: x = [np.array([1, 2, 3]), np.array([2, 5, 10, 20])]

Parameters

x – iterable (maybe nested iterable)

Returns

int Generalized dimension count, considering that some elements in a list might be iterable and others not.

omfit_classes.utils_math.safe_divide(numerator, denominator, fill_value=0)[source]

Division function to safely compute the ratio of two lists/arrays. The fill_value input parameter specifies what value should be filled in for the result whenever the denominator is 0.

Parameters
  • numerator – numerator of the division

  • denominator – denominator of the division

  • fill_value – fill value when denominator is 0

Returns

division with fill_value where nan or inf would have been instead

omfit_classes.utils_math.nannanargmin(x, axis=1)[source]

Performs nanargmin along an axis on a 2D array while dodging errors from empty rows or columns

argmin finds the index where the minimum value occurs. nanargmin ignores NaNs while doing this.

However, if nanargmin is operating along just one axis and it encounters a row that’s all NaN, it raises a ValueError, because there’s no valid index for that row. It can’t insert NaN into the result, either, because the result should be an integer array (or else it couldn’t be used for indexing), and NaN is a float. This function is for cases where we would like nanargmin to give valid results where possible and clearly invalid indices for rows that are all NaN. That is, it returns -N, where N is the row length, if the row is all NaN.

Parameters
  • x – 2D float array Input data to process

  • axis – int 0 or 1

Returns

1D int array indices of the minimum value of each row or column. Rows/columns which are all NaN will be have -N, where N is the size of the relevant dimension (so -N is invalid).

omfit_classes.utils_math.calcz(x, y, consistent_reconstruction=True)[source]

Calculate Z: the inverse normalized scale-length The function is coded in such a way to avoid NaN and Inf where y==0

z = -dy/dx/y

Parameters
  • x – x axis array

  • y – y axis array

  • consistent_reconstruction – calculate z so that integration of z with integz exactly generates original profile

>> z_cF =calcz(rho,ne,consistent_reconstruction=False) >> z_cT=calcz(rho,ne,consistent_reconstruction=True) >> pyplot.plot(rho,ne) >> pyplot.plot(rho,integz(rho,z_cT,rho[0],ne[0],rho),’.’) >> pyplot.plot(rho,integz(rho,z_cF,rho[0],ne[0],rho),’x’)

Returns

z calculated at x

omfit_classes.utils_math.mergez(x0, z0, x1, z1, core, nml, edge, x)[source]

Perform merge of two Z profiles

  • rho < core: Z goes linearly to 0 on axis

  • core < rho < nml: profile 0

  • nml < rho < edge: linearly connect Z at nml and edge

  • rho > edge: profile 1

Parameters
  • x0 – abscissa of first profile

  • z0 – values of first Z profile

  • x1 – abscissa of second profile

  • z1 – values of second Z profile

  • core – abscissa of core boundary

  • nml – abscissa of nml boundary

  • edge – abscissa of edge boundary

Param

returns tuple of merged x and z points

omfit_classes.utils_math.integz(x0, z0, xbc, ybc, x, clipz=True)[source]

Integrate inverse scale-length Z to get profile

Parameters
  • x0 – abscissa of the Z profile

  • z0 – Z profile

  • xbc – abscissa of the boundary condition

  • ybc – value of the boundary condition

  • x – integrate the Z profile at these points

  • clipz – use clipped z for extrapolation (alternatively linear)

Returns

integrated Z profile at points defined by x

omfit_classes.utils_math.detect_peaks(x, mph=None, mpd=1, threshold=0, edge='rising', kpsh=False, valley=False, show=False, ax=None)[source]

Detect peaks in data based on their amplitude and other features.

x1D array_like

data.

mph{None, number}, optional (default = None)

detect peaks that are greater than minimum peak height.

mpdpositive integer, optional (default = 1)

detect peaks that are at least separated by minimum peak distance (in number of data).

thresholdpositive number, optional (default = 0)

detect peaks (valleys) that are greater (smaller) than threshold in relation to their immediate neighbors.

edge{None, ‘rising’, ‘falling’, ‘both’}, optional (default = ‘rising’)

for a flat peak, keep only the rising edge (‘rising’), only the falling edge (‘falling’), both edges (‘both’), or don’t detect a flat peak (None).

kpshbool, optional (default = False)

keep peaks with same height even if they are closer than mpd.

valleybool, optional (default = False)

if True (1), detect valleys (local minima) instead of peaks.

showbool, optional (default = False)

if True (1), plot data in matplotlib figure.

ax : a matplotlib.axes.Axes instance, optional (default = None).

ind1D array_like

indeces of the peaks in x.

The detection of valleys instead of peaks is performed internally by simply negating the data: ind_valleys = detect_peaks(-x)

The function can handle NaN’s

See this IPython Notebook 1.

1

http://nbviewer.ipython.org/github/demotu/BMC/blob/master/notebooks/DetectPeaks.ipynb

>> from detect_peaks import detect_peaks >> x = np.random.randn(100) >> x[60:81] = np.nan >> # detect all peaks and plot data >> ind = detect_peaks(x, show=True) >> print(ind)

>> x = np.sin(2*np.pi*5*np.linspace(0, 1, 200)) + np.random.randn(200)/5. >> # set minimum peak height = 0 and minimum peak distance = 20 >> detect_peaks(x, mph=0, mpd=20, show=True)

>> x = [0, 1, 0, 2, 0, 3, 0, 2, 0, 1, 0] >> # set minimum peak distance = 2 >> detect_peaks(x, mpd=2, show=True)

>> x = np.sin(2*np.pi*5*np.linspace(0, 1, 200)) + np.random.randn(200)/5. >> # detection of valleys instead of peaks >> detect_peaks(x, mph=0, mpd=20, valley=True, show=True)

>> x = [0, 1, 1, 0, 1, 1, 0] >> # detect both edges >> detect_peaks(x, edge=’both’, show=True)

>> x = [-2, 1, -2, 2, 1, 1, 3, 0] >> # set threshold = 2 >> detect_peaks(x, threshold = 2, show=True)

omfit_classes.utils_math.find_feature(y, x=None, M=0.01, k=5.0, xmin=None, xmax=None, retall=False)[source]

Identify edges and center of a sharp feature (pedestal, peak, or well) in a otherwise smooth profile.

Parameters
  • y – np.ndarray. Parameter of interest, such as T_e

  • x – np.ndarray. Position basis, such as psi_N (Default is np.arange(len(y))

  • M – float. Gaussian smoothing sigma in units of position basis

  • k – float. Difference of gaussians factor (second smoothing sigma is M*k)

  • xmin – float. Lower limit on range of X values to consider in search

  • xmax – float. Upper limit on range of X values to consider in search

  • retall – bool. Return gaussian smoothed functions

:returns : ( Inner edge, outer edge, center, value at inner edge, value at outer edge )

Examples

Here, is a simple 1D example using find_feature to identify the steep gradient region of a tanh.

>> x = np.linspace(0,1,1000) >> center = 0.8 >> halfwidth = 0.05 >> y = np.tanh((center-x)/halfwidth)+1 >> dydx = deriv(x,y) >> indx = find_feature(dydx,x=x) >> xpt = x[indx] >> ypt = dydx[indx]

>> foundwidth = x[indx[2]]-x[indx[0]] >> print(“Width within {:.2f}%”.format((1-2*halfwidth/foundwidth)*100))

Note we chose M to set the appropriate scale of interest for the above problem. We can see the sensitivity to this by scanning the tanh width in a 2D example.

>> xs = x.reshape(1,-1) >> halfwidths = np.linspace(0.01,0.1,50).reshape(-1,1) >> ys = np.tanh((center-xs)/halfwidths)+1 >> dydxs = np.gradient(ys)[1]/np.gradient(x)

M sets the scale for the steepness of interest using one M for a range of bump sizes isolates only the scale of interest >> indxs = np.apply_along_axis(find_feature,1,dydxs,x=x,M=0.01,k=5) Tracking the scale of the bump with M approximates tanh widths >> #indxs = [find_feature(dy,x=x,M=2*hw/10.,k=5) for dy,hw in zip(dydxs,halfwidths)]

>> # found peak and edge points of steep gradient region >> foundwidths = map(lambda indx: x[indx[2]]-x[indx[0]], indxs) >> xpts = [x[i] for i in indxs] >> ypts = [yy[i] for yy,i in zip(ys,indxs)] >> dypts= [yy[i] for yy,i in zip(dydxs,indxs)]

>> # tanh half width points >> # Note np.tanh(1) = 0.76, and d/dxtanh = sech^2 = 0.42 >> ihws = np.array([(np.abs(x-center-hw).argmin(), >> np.abs(x-center+hw).argmin()) for hw in halfwidths[:,0]]) >> xhws = [x[i] for i in ihws] >> yhws = [yy[i] for yy,i in zip(ys,ihws)] >> dyhws= [yy[i] for yy,i in zip(dydxs,ihws)]

Visualize the comparison between tanh widths and the identified region of steep gradient >> close(‘all’) >> f,ax = pyplot.subplots(3,1) >> ax[0].set_ylabel(‘y = np.tanh((c-x)2/w)+1’) >> ax[0].set_xlabel(‘x’) >> ax[1].set_ylabel(‘dy/dx’) >> ax[1].set_xlabel(‘x’) >> for i in [0,24,49]: … l, = ax[0].plot(x,ys[i]) … w1, = ax[0].plot(xhws[i],yhws[i],marker=’o’,ls=’’,fillstyle=’none’,color=l.get_color()) … w2, = ax[0].plot(xpts[i],ypts[i],marker=’x’,ls=’’,fillstyle=’none’,color=l.get_color()) … l, = ax[1].plot(x,dydxs[i]) … w1, = ax[1].plot(xhws[i],dyhws[i],marker=’o’,ls=’’,fillstyle=’none’,color=l.get_color()) … w2, = ax[1].plot(xpts[i],dypts[i],marker=’x’,ls=’’,fillstyle=’none’,color=l.get_color()) >> ax[-1].set_ylabel(‘Edge Width’) >> ax[-1].set_xlabel(‘Tanh Width’) >> ax[-1].plot([0,2*halfwidths[-1,0]],[0,2*halfwidths[-1,0]]) >> ax[-1].plot(2*halfwidths[:,0],foundwidths,marker=’o’,lw=0)

omfit_classes.utils_math.parabola(x, y)[source]

y = a*x^2 + b*x + c

Returns

a,b,c

omfit_classes.utils_math.parabolaMax(x, y, bounded=False)[source]

Calculate a parabola through x,y, then return the extremum point of the parabola

Parameters
  • x – At least three abcissae points

  • y – The corresponding ordinate points

  • bounded – False, ‘max’, or ‘min’ - False: The extremum is returned regardless of location relative to x (default) - ‘max’ (‘min’): The extremum location must be within the bounds of x, and if not return the location and value of max(y) (min(y))

Returns

x_max,y_max - The location and value of the extremum

omfit_classes.utils_math.parabolaCycle(X, Y, ix)[source]
omfit_classes.utils_math.parabolaMaxCycle(X, Y, ix, bounded=False)[source]

Calculate a parabola through X[ix-1:ix+2],Y[ix-1:ix+2], with proper wrapping of indices, then return the extremum point of the parabola

Parameters
  • X – The abcissae points: an iterable to be treated as periodic

  • y – The corresponding ordinate points

  • ix – The index of X about which to find the extremum

  • bounded – False, ‘max’, or ‘min’ - False: The extremum is returned regardless of location relative to x (default) - ‘max’ (‘min’): The extremum location must be within the bounds of x, and if not return the location and value of max(y) (min(y))

Returns

x_max,y_max - The location and value of the extremum

omfit_classes.utils_math.paraboloid(x, y, z)[source]

z = ax*x^2 + bx*x + ay*y^2 + by*y + c

NOTE: This function uses only the first 5 points of the x, y, z arrays to evaluate the paraboloid coefficients

Returns

ax,bx,ay,by,c

omfit_classes.utils_math.smooth(x, window_len=11, window='hann', axis=0)[source]

smooth the data using a window with requested size.

This method is based on the convolution of a scaled window with the signal. The signal is prepared by introducing reflected copies of the signal (with the window size) in both ends so that transient parts are minimized in the beginning and end part of the output signal.

input:
param x

the input signal

param window_len

the dimension of the smoothing window; should be an odd integer; is ignored if window is an array

param window

the window function to use, see scipy.signal.get_window documentation for list of available windows ‘flat’ or ‘boxcar’ will produce a moving average smoothing Can also be an array, in which case it is used as the window function and window_len is ignored

param axis

The axis that is smoothed

output:

the smoothed signal

Examples

>> t=np.linspace(-2,2,100) >> x=np.sin(t)+randn(len(t))*0.1 >> y=smooth(x,11)

see also:

scipy.signal.get_window, np.convolve, scipy.signal.lfilter

omfit_classes.utils_math.smoothIDL(data, n_smooth=3, timebase=None, causal=False, indices=False, keep=False)[source]

An efficient top hat smoother based on the smooth IDL routine. The use of cumsum-shift(cumsum) means that execution time is 2xN flops compared to 2 x n_smooth x N for a convolution. If supplied with a timebase, the shortened timebase is returned as the first of a tuple.

Parameters
  • data – (timebase, data) is a shorthand way to pass timebase

  • n_smooth – smooth bin size

  • timebase – if passed, a tuple with (timebase,smoothed_data) gets returned

  • causal – If True, the smoothed signal never preceded the input, otherwise, the smoothed signal is “centred” on the input (for n_smooth odd) and close (1//2 timestep off) for even

  • indices – if True, return the timebase indices instead of the times

  • keep – Better to throw the partially cooked ends away, but if you want to keep them use keep=True. This is useful for quick filtering applications so that original and filtered signals are easily compared without worrying about timebase

>> smoothIDL([1,2,3,4],3) np.array([ 2., 3.]) >> smoothIDL([1.,2.,3.,4.,5.],3) np.array([ 2., 3., 4.]) >> smoothIDL([1,2,3,4,5],timebase=np.array([1,2,3,4,5]),n_smooth=3, causal=False) (np.array([2, 3, 4]), np.array([ 2., 3., 4.])) >> smoothIDL([0,0,0,3,0,0,0],timebase=[1,2,3,4,5,6,7],n_smooth=3, causal=True) ([3, 4, 5, 6, 7], np.array([ 0., 1., 1., 1., 0.])) >> smoothIDL([0,0,0,3,0,0,0],timebase=[1,2,3,4,5,6,7],n_smooth=3, causal=True, indices=True) ([2, 3, 4, 5, 6], np.array([ 0., 1., 1., 1., 0.])) >> smoothIDL([0, 0, 0, 0, 5, 0, 0, 0, 0, 0, 0], 5, keep=1) np.array([ 0., 0., 1., 1., 1., 1., 1., 0., 0., -1., -1.])

omfit_classes.utils_math.pcs_rc_smooth(xx, yy, s)[source]

RC lowpass smoothing filter using same formula as in the PCS. Can be used to reproduce PCS lowpass filtering.

Parameters
  • xx – array Independent variable (like time). Does not have to be evenly spaced.

  • yy – array matching xx Depending variable

  • s – float Smoothing time constant in units matching xx

Returns

array Smoothed/filtered version of yy

omfit_classes.utils_math.get_time_windows_for_causal_smooth(t_base, t_target, t_window)[source]
omfit_classes.utils_math.causal_smooth_equidistant(t, y, t_window, t_mask=None, t_target=None, return_full_filtered=False)[source]
omfit_classes.utils_math.smooth_by_convolution(yi, xi=None, xo=None, window_size=None, window_function='gaussian', axis=0, causal=False, interpolate=False, std_dev=2)[source]

Convolution of a non-uniformly discretized array with window function.

The output values are np.nan where no points are found in finite windows (weight is zero). The gaussian window is infinite in extent, and thus returns values for all xo.

Supports uncertainties arrays. If the input –does not– have associated uncertainties, then the output will –not– have associated uncertainties.

Parameters
  • yi – array_like (…,N,…). Values of input array

  • xi – array_like (N,). Original grid points of input array (default y indicies)

  • xo – array_like (M,). Output grid points of convolution array (default xi)

  • window_size – float. Width of passed to window function (default maximum xi step). For the Gaussian, sigma=window_size/4. and the convolution is integrated across +/-4.*sigma.

  • window_function – str/function. Accepted strings are ‘hanning’,’bartlett’,’blackman’,’gaussian’, or ‘boxcar’. Function should accept x and window_size as arguments and return a corresponding weight.

  • axis – int. Axis of y along which convolution is performed

  • causal – int. Forces f(x>0) = 0.

  • interpolate – False or integer number > 0 Paramter indicating to interpolate data so that there are`interpolate` number of data points within a time window. This is useful in presence of sparse data, which would result in stair-case output if not interpolated. The integer value sets the # of points per window size.

  • std_dev – str/int Accepted strings are ‘none’, ‘propagate’, ‘population’, ‘expand’, ‘deviation’, ‘variance’. Only ‘population’ and ‘none’ are valid if yi is not an uncertainties array (i.e. std_devs(yi) is all zeros). Setting to an integer will convolve the error uncertainties to the std_dev power before taking the std_dev root. std_dev = ‘propagate’ is true propagation of errors (slow if not interpolating) std_dev = ‘population’ is the weighted “standard deviation” of the points themselves (strictly correct for the boxcar window) std_dev = ‘expand’ is propagation of errors weighted by w~1/window_function std_dev = ‘deviation’ is equivalent to std_dev=1 std_dev = ‘variance’ is equivalent to std_dev=2

Returns

convolved array on xo

>> M=300 >> ninterp = 10 >> window=[‘hanning’,’gaussian’,’boxcar’][1] >> width=0.05 >> f = figure(num=’nu_conv example’) >> f.clf() >> ax = f.use_subplot(111) >> >> xo=np.linspace(0,1,1000) >> >> x0=xo >> y0=(x0>0.25)&(x0<0.75) >> pyplot.plot(x0,y0,’b-‘,label=’function’) >> >> x1=np.cumsum(rand(M)) >> x1=(x1-min(x1))/(max(x1)-min(x1)) >> y1=(x1>0.25)&(x1<0.75) >> pyplot.plot(x1,y1,’b.’,ms=10,label=’subsampled function’) >> if window==’hanning’: >> ys=smooth(interp1e(x0,y0)(xo),int(len(xo)*width)) >> pyplot.plot(xo,ys,’g-‘,label=’smoothed function’) >> yc=smooth(interp1e(x1,y1)(xo),int(len(xo)*width)) >> pyplot.plot(xo,yc,’m–’,lw=2,label=’interp subsampled then convolve’) >> >> y1=unumpy.uarray(y1,y1*0.1) >> a=time.time() >> yo=nu_conv(y1,xi=x1,xo=xo,window_size=width,window_function=window,std_dev=’propagate’,interpolate=ninterp) >> print(‘nu_conv time: {:}’.format(time.time()-a)) >> ye=nu_conv(y1,xi=x1,xo=xo,window_size=width,window_function=window,std_dev=’expand’,interpolate=ninterp) >> >> uband(x1,y1,color=’b’,marker=’.’,ms=10) >> uband(xo,yo,color=’r’,lw=3,label=’convolve subsampled’) >> uband(xo,ye,color=’c’,lw=3,label=’convolve with expanded-error propagation’) >> >> legend(loc=0) >> pyplot.ylim([-0.1,1.1]) >> pyplot.title(‘%d points , %s window, %3.3f width, interpolate %s’%(M,window,width, ninterp))

omfit_classes.utils_math.trismooth(*args, **kw)[source]

Smooth with triangle kernel, designed to mimic smooth in reviewplus :param args:

y: float array, y, window_size: float array, float y, x, window_size: float array, float array, float OR y, window_size[optional] as OMFITmdsValue, float If x or window_size are not defined, smooth_by_convolution will pick values automatically. window_size is the half width fo the kernel.

The default window_size set by smooth_by_convolution will probably be too small, unless you are providing a much higher resolution output grid with xo. You should probably set window_size by providing at least two arguments.

Parameters

kw – Keywords passed to smooth_by_convolution.

Returns

float array or uarray or tuple of arrays ysmo = result from smooth_by_convolution; array (default) or uarray (std_dev is modified) if OMFITmdsValue:

(x, ysmo)

else:

ysmo

omfit_classes.utils_math.cicle_fourier_smooth(R, Z, keep_M_harmonics, equalAngle=False, doPlot=False)[source]

smooth closed periodic boundary by zeroing out high harmonics components

Parameters
  • R – x coordinates of the boundary

  • Z – y coordinates of the boundary

  • keep_M_harmonics – how many harmonics to keep

  • equalAngle – use equal angle interpolation, and if so, how many points to use

  • doPlot – plot plasma boundary before and after

Returns

smoothed R and Z coordinates

omfit_classes.utils_math.nu_conv(yi, xi=None, xo=None, window_size=None, window_function='gaussian', axis=0, causal=False, interpolate=False, std_dev=2)

Convolution of a non-uniformly discretized array with window function.

The output values are np.nan where no points are found in finite windows (weight is zero). The gaussian window is infinite in extent, and thus returns values for all xo.

Supports uncertainties arrays. If the input –does not– have associated uncertainties, then the output will –not– have associated uncertainties.

Parameters
  • yi – array_like (…,N,…). Values of input array

  • xi – array_like (N,). Original grid points of input array (default y indicies)

  • xo – array_like (M,). Output grid points of convolution array (default xi)

  • window_size – float. Width of passed to window function (default maximum xi step). For the Gaussian, sigma=window_size/4. and the convolution is integrated across +/-4.*sigma.

  • window_function – str/function. Accepted strings are ‘hanning’,’bartlett’,’blackman’,’gaussian’, or ‘boxcar’. Function should accept x and window_size as arguments and return a corresponding weight.

  • axis – int. Axis of y along which convolution is performed

  • causal – int. Forces f(x>0) = 0.

  • interpolate – False or integer number > 0 Paramter indicating to interpolate data so that there are`interpolate` number of data points within a time window. This is useful in presence of sparse data, which would result in stair-case output if not interpolated. The integer value sets the # of points per window size.

  • std_dev – str/int Accepted strings are ‘none’, ‘propagate’, ‘population’, ‘expand’, ‘deviation’, ‘variance’. Only ‘population’ and ‘none’ are valid if yi is not an uncertainties array (i.e. std_devs(yi) is all zeros). Setting to an integer will convolve the error uncertainties to the std_dev power before taking the std_dev root. std_dev = ‘propagate’ is true propagation of errors (slow if not interpolating) std_dev = ‘population’ is the weighted “standard deviation” of the points themselves (strictly correct for the boxcar window) std_dev = ‘expand’ is propagation of errors weighted by w~1/window_function std_dev = ‘deviation’ is equivalent to std_dev=1 std_dev = ‘variance’ is equivalent to std_dev=2

Returns

convolved array on xo

>> M=300 >> ninterp = 10 >> window=[‘hanning’,’gaussian’,’boxcar’][1] >> width=0.05 >> f = figure(num=’nu_conv example’) >> f.clf() >> ax = f.use_subplot(111) >> >> xo=np.linspace(0,1,1000) >> >> x0=xo >> y0=(x0>0.25)&(x0<0.75) >> pyplot.plot(x0,y0,’b-‘,label=’function’) >> >> x1=np.cumsum(rand(M)) >> x1=(x1-min(x1))/(max(x1)-min(x1)) >> y1=(x1>0.25)&(x1<0.75) >> pyplot.plot(x1,y1,’b.’,ms=10,label=’subsampled function’) >> if window==’hanning’: >> ys=smooth(interp1e(x0,y0)(xo),int(len(xo)*width)) >> pyplot.plot(xo,ys,’g-‘,label=’smoothed function’) >> yc=smooth(interp1e(x1,y1)(xo),int(len(xo)*width)) >> pyplot.plot(xo,yc,’m–’,lw=2,label=’interp subsampled then convolve’) >> >> y1=unumpy.uarray(y1,y1*0.1) >> a=time.time() >> yo=nu_conv(y1,xi=x1,xo=xo,window_size=width,window_function=window,std_dev=’propagate’,interpolate=ninterp) >> print(‘nu_conv time: {:}’.format(time.time()-a)) >> ye=nu_conv(y1,xi=x1,xo=xo,window_size=width,window_function=window,std_dev=’expand’,interpolate=ninterp) >> >> uband(x1,y1,color=’b’,marker=’.’,ms=10) >> uband(xo,yo,color=’r’,lw=3,label=’convolve subsampled’) >> uband(xo,ye,color=’c’,lw=3,label=’convolve with expanded-error propagation’) >> >> legend(loc=0) >> pyplot.ylim([-0.1,1.1]) >> pyplot.title(‘%d points , %s window, %3.3f width, interpolate %s’%(M,window,width, ninterp))

class omfit_classes.utils_math.FilteredNDInterpolator(points, values, size_scales=None, check_finite=False, std_dev=True, filter_type='median')[source]

Bases: object

A class for smoothing ND data. Useful for down-sampling and gridding.

Mean or median filter interpolate in N dimensions. Calculates the mean or median of all values within a size_scale “sphere” of the interpolation point. Unlike linear interpolation, you can incorporate information from all your data when down-sampling to a regular grid.

Of course, it would be better to do a weight function (gaussian, hanning, etc) convolution, but the “volume” elements required for integration get very computationally expensive in multiple dimensions. In that case, try linear interpolation followed by regular-grid convolution (ndimage processing).

Parameters
  • points – np.ndarray of floats shape (M, D), where D is the number of dimensions

  • values – np.ndarray of floats shape (M,)

  • size_scales – tuple (D,) scales of interest in each dimension

  • check_finite – bool. Cleans non-finite points and/or values prior to interpolation

  • std_dev – bool. Estimate uncertainty of mean or median interpolation

  • filter_type – str. Accepts ‘mean’ or ‘median’

Examples

This example shows the interpolator is reasonably fast, but note it is nowhere near as scalable to 100k+ points as linear interpolation + ndimage processing.

>> ngrid = 100 >> width = (0.05, 0.05) >> noise_factor = 0.5

>> def func(x, y): >> return x * (1 - x) * np.cos(4 * np.pi * x) * np.sin(4 * np.pi * y ** 2) ** 2

>> x = np.linspace(0, 1, ngrid) >> y = np.linspace(0, 1, ngrid) >> grid_x, grid_y = np.meshgrid(x, y) >> xy = np.array(np.meshgrid(x, y)).T.reshape(-1, 2)

>> fig, axs = pyplot.subplots(3, 2, sharex=True, sharey=True, figsize=(8, 9)) >> axs[0, 0].set_title(‘Original’) >> data = func(grid_x, grid_y) >> im = axs[0, 0].imshow(data, extent=(0, 1, 0, 1), origin=’lower’, aspect=’auto’) >> kw = dict(extent=(0, 1, 0, 1), origin=’lower’, aspect=’auto’, clim=im.get_clim())

>> axs[0, 1].set_title(‘Noisy’) >> data = data + (np.random.random(data.shape) - 0.5) * np.ptp(data) * noise_factor >> im = axs[0, 1].imshow(data, **kw )

>> ns = [400, 1600, 6400] >> times, nums = [], [] >> for npts, axi in zip(ns, axs[1:]): >> print(‘npts = {:}’.format(npts)) >> points = np.random.rand(npts, 2) >> values = func(points[:, 0], points[:, 1]) >> values = values + (np.random.random((npts,)) - 0.5) * np.ptp(values) * noise_factor

>> strt = datetime.datetime.now() >> # ci = ConvolutionNDInterpolator(points, values, xi=xy, window_size=width, window_function=’boxcar’) >> ci = FilteredNDInterpolator(points, values, size_scales=width, filter_type=’mean’)(grid_x, grid_y) >> ctime = datetime.datetime.now() - strt >> print(‘ > Convolution took {:}’.format(ctime)) >> times.append(ctime.total_seconds()) >> nums.append(npts)

>> ax = axi[1] >> ax.set_title(‘Interpolated {:} pts’.format(npts)) >> yi = LinearNDInterpolator(points, values)(grid_x, grid_y) >> im = ax.imshow(nominal_values(yi), **kw) >> l, = ax.plot(points[:, 0], points[:, 1], ‘k.’, ms=1)

>> ax = axi[0] >> ax.set_title(‘Convolved {:} pts’.format(npts)) >> im = ax.imshow(nominal_values(ci), **kw) >> l, = ax.plot(points[:, 0], points[:, 1], ‘k.’, ms=1)

>>fig.tight_layout() >>axs[0, 0].set_xlim(0, 1) >>axs[0, 0].set_ylim(0, 1)

>>fig, axs = pyplot.subplots() >>axs.plot([0] + nums, [0] + times, marker=’o’) >>axs.set_xlabel(‘# of points’) >>axs.set_ylabel(‘time (s)’) >>print(‘sec/k ~ {:.3e}’.format(1e3 * np.mean(np.array(times) / ns)))

This example shows that the median filter is “edge preserving” and contrasts both filters with a more sophisticated convolution. Note the boxcar convolution is not identical to the mean filter for low sampling because the integral weights isolated points more than very closely spaced points.

>> npts = 50 >> window = [‘hanning’, ‘gaussian’, ‘boxcar’][1] >> width = 0.05 >> ped = 0.1 >> pedwidth = 0.03 >> fun = lambda x: (0.9-x)*2*(x < 0.9) + 1.0 - np.tanh((x-(1.0-ped))/(0.25*pedwidth))

>> fig, axs = pyplot.subplots(3, sharex=True, sharey=True) >> axs[0].set_title(‘Pedestal width {:}, Filter Length Scale {:}’.format(pedwidth, width)) >> for npts, ax in zip([20, 50, 200], axs): >> x0 = np.linspace(0, 1, 1000)

>> x1 = np.cumsum(rand(npts)) >> x1 = (x1 - min(x1)) / (max(x1) - min(x1)) >> y1 = fun(x1) + np.random.random(x1.shape) * 0.4 * fun(x1) >> y1 = uarray(y1, y1 * 0.1)

>> t0 = datetime.datetime.now() >> yc = nu_conv(y1, xi=x1, xo=x0, window_size=width, window_function=window, std_dev=’propagate’) >> tc = datetime.datetime.now() - t0 >> print(‘nu_conv time: {:}’.format(tc)) >> t0 = datetime.datetime.now() >> yd = FilteredNDInterpolator(x1.reshape(-1, 1), y1, size_scales=width / 2.)(x0) >> td = datetime.datetime.now() - t0 >> print(‘median filter time: {:}’.format(td)) >> t0 = datetime.datetime.now() >> ym = FilteredNDInterpolator(x1.reshape(-1, 1), y1, size_scales=width / 2., filter_type=’mean’)(x0) >> tm = datetime.datetime.now() - t0 >> print(‘median filter time: {:}’.format(td))

>> #ax.plot(x0, fun(x0), label=’function’) >> uband(x0, yc, ax=ax, lw=3, label=’{:} convolution, {:.2f}s’.format(window.title(), tc.total_seconds())) >> uband(x0, ym, ax=ax, lw=3, label=’Mean Filter, {:.2f}s’.format(tm.total_seconds())) >> uband(x0, yd, ax=ax, lw=3, label=’Median Filter, {:.2f}s’.format(td.total_seconds())) >> uerrorbar(x1, y1, ax=ax, markersize=4, ls=’-‘, lw=0.5, label=’{:} data points’.format(npts), color=’k’) >> ax.legend(loc=0, frameon=False) >> ax.set_ylim([-0.1, 5]) >> fig.tight_layout()

exception omfit_classes.utils_math.WeightedAvgBadInput[source]

Bases: ValueError, omfit_classes.exceptions_omfit.doNotReportException

weighted_avg got some invalid values and has halted.

omfit_classes.utils_math.weighted_avg(x, err=None, minerr=None, return_stddev_mean=False, return_stddev_pop=False, nan_policy='ignore')[source]

Weighted average using uncertainty

Propagates errors and calculates weighted standard deviation. While nu_conv does these for a sliding window vs. time, this function is simpler and does calculations for a single mean of an array.

Parameters
  • x – 1D float array The input data to be averaged

  • err – 1D float array Uncertainty in x. Should have the same units as x. Should have the same length as x. Special case: a single value of err will be used to propagate errors for the standard deviation of the mean, but will produce uniform (boring) weights. If no uncertainty is provided (err==None), then uniform weighting is used.

  • minerr – float Put a floor on the uncertainties before calculating weight. This prevents a few points with unreasonbly low error from stealing the calculation. Should be a scalar with same units as x.

  • return_stddev_mean – bool Flag for whether the standard deviation of the mean (propagated uncertainty) should be returned with the mean.

  • return_stddev_pop – bool Flag for whether the standard deviation of the population (weighted standard deviation) should be returned with the mean.

  • nan_policy – str ‘nan’: return NaN if there are NaNs in x or err ‘error’: raise an exception ‘ignore’: perform the calculation on the non-nan elements only (default)

Returns

float or tuple mean if return_stddev_mean = False and return_stddev_pop = False (mean, xpop, xstdm) if return_stddev_mean = True and return_stddev_pop = True (mean, xpop) if return_stddev_mean = False and return_stddev_pop = True (mean, xstdm) if return_stddev_mean = True and return_stddev_pop = False where xstdm and xpop are the propagated error and the weighted standard deviation.

omfit_classes.utils_math.firFilter(time_s, data, cutoff_hz, sharpness=0.5, full_output=False)[source]

Filter signal data with FIR filter (lowpass, highpass)

Parameters
  • time_s – time basis of the data signal in seconds

  • data – data signal

  • cutoff_hz – a list of two elements with low and high cutoff frequencies [lowFreq,highFreq]

  • sharpness – sharpness of the filter (taps of the FIR filter = sample_rate/2. * sharpness)

  • full_output – return filter window and frequency in addition to filtered data

Returns

filtered data or tuple with filtered data and filter window and frequency as a tuple

omfit_classes.utils_math.butter_smooth(xx, yy, timescale=None, cutoff=None, laggy=False, order=1, nan_screen=True, btype='low')[source]

Butterworth smoothing lowpass filter.

Similar to firFilter, but with a notable difference in edge effects (butter_smooth may behave better at the edges of an array in some cases).

Parameters
  • xx – 1D array Independent variable. Should be evenly spaced for best results.

  • yy – array matching dimension of xx

  • timescale – float [specifiy either timescale or cutoff] Smoothing timescale. Units should match xx.

  • cutoff – float [specify either timescale or cutoff] Cutoff frequency. Units should be inverse of xx. (xx in seconds, cutoff in Hz; xx in ms, cutoff in kHz, etc.)

  • laggy – bool True: causal filter: smoothed output lags behind input False: acausal filter: uses information from the future so that smoothed output doesn’t lag input

  • order – int Order of butterworth filter. Lower order filters seem to have a longer tail after the ELM which is helpful for detecting the tail of the ELM.

  • nan_screen – bool Perform smoothing on only the non-NaN part of the array, then pad the result out with NaNs to maintain length

  • btype – string low or high. For smoothing, always choose low. You can do a highpass filter instead of lowpass by choosing high, though.

Returns

array matching dimension of xx Smoothed version of yy

class omfit_classes.utils_math.fourier_boundary(nf, rb, zb, symmetric=False)[source]

Bases: object

The routines to make and plot Fourier representation of a boundary from a list of points

make_fourier(nf, rb, zb)[source]
fourier_coef()[source]
reconstruct_boundary()[source]

Uses the fourier representation and amin, r0 and z0 parameters to reconstruct the boundary

omfit_classes.utils_math.fourier(y, t=None)[source]

Calculate fourier transform

Parameters
  • y – signal

  • t – time basis of the signal (uniformly spaced)

Returns

tuple with Fourier transform of the signal and frequency basis

>> t=r_[0:1:0.1] >> y=np.cos(2*np.pi*t) >> >> Y,W=fourier(y,t) >> >> ax=pyplot.subplot(2,1,1) >> ax.plot(t,y) >> ax.set_xlabel(‘t [s]’) >> ax.set_ylabel(‘$y$’) >> >> ax=pyplot.subplot(2,1,2) >> ax.plot(W/2./np.pi,Y) >> ax.set_xlabel(‘f [Hz]’) >> ax.set_ylabel(‘$Y$’)

omfit_classes.utils_math.windowed_FFT(t, y, nfft='auto', overlap=0.95, hanning_window=True, subtract_mean=True, real_input=None)[source]

Bin data into windows and compute FFT for each. Gives amplitude vs. time and frequency. Useful for making spectrograms.

input:
param t

1D time vector in ms

param y

1D parameter as a function of time

param nfft

Number of points in each FFT bin. More points = higher frequency resolution but lower time resolution.

param overlap

Fraction of window that overlaps previous/next window.

param hanning_window

Apply a Hanning window to y(t) before each FFT.

param subtract_mean

Subtract the average y value (of the entire data set, not individually per window) before calculating FFT

param real_input

T/F: Use rfft instead of fft because all inputs are real numbers. Set to None to decide automatically.

output:
return

spectral density (time,frequency)

return

array of times at the center of each FFT window

return

frequency

return

amplitude(time, positive frequency)

return

power(time, positive frequency)

return

positive-only frequency array

return

nfft (helpful if you let it be set automatically)

more on overlap for windows 0, 1, 2:

overlap = 0.0 : no overlap, window 0 ends where window 1 begins overlap = 0.5 : half overlap, window 0 ends where window 2 begins overlap = 0.99: this is probably as high as you should go. It will look nice and smooth, but will take longer. overlap = 1.0 : FAIL, all the windows would stack on top of each other and infinite windows would be required.

more on nfft:

Set nfft=’auto’ and the function will pick a power of two that should give a reasonable view of the dataset. It won’t choose nfft < 16.

omfit_classes.utils_math.noise_estimator(t, y, cutoff_timescale=None, cutoff_omega=None, window_dt=0, dt_var_thresh=1e-09, avoid_jumps=False, debug=False, debug_plot=None, restrict_cutoff=False)[source]

Estimates uncertainty in a signal by assuming that high frequency (above cutoff frequency) variation is random noise

Parameters
  • t – 1D float array with regular spacing (ms) Time base for the input signal

  • y – 1D float array with length matching t (arbitrary) Input signal value as a function of t

  • cutoff_timescale – float (ms) For a basic RC filter, this would be tau = R*C. Define either this or cutoff_omega.

  • cutoff_omega – float (krad/s) The cutoff angular frequency, above which variation is assumed to be noise. Overrides cutoff_timescale if specified. cutoff_timescale = 1.0/cutoff_omega

  • window_dt – float or None (ms) Window half width for windowed_fft, used in some strategies for time varying noise. If <= 0, one scalar noise estimate is made for the entire range of t, using FFT methods. Set to None to choose window size automatically based on cutoff_timescale. This option does not seem to be as good as the standard method.

  • debug – bool Flag for activating debugging tests, plots (unless debug_plot is explicitly disabled), and reports. Also returns all four estimates instead of just the best one.

  • debug_plot – bool [optional] By default, debug_plot copies the value of debug, but you can set it to False to disable plots and still get other debugging features. Setting debug_plot without debug is not supported and will be ignored.

  • dt_var_thresh – float (fraction) Threshold for variability in dt. t must be evenly spaced, so nominally dt will be a constant. Because of various real world issues, there could be some slight variation. As long as this is small, everything should be fine. You can adjust the threshold manually, but be careful: large variability in the spacing of the time base breaks assumptions in the signal processing techniques used by this method. If std(dt)/mean(dt) exceeds this threshold, an exception will be raised.

  • avoid_jumps – bool Large jumps in signal will have high frequency components which shouldn’t be counted with high frequency noise. You SHOULD pick a time range to avoid these jumps while estimating noise, but if you don’t want to, you can try using this option instead. If this flag is true, an attempt will be made to identify time periods when jumps are happening. The time derivative, dy/dt, is evaluated at times where there are no detected jumps and interpolated back onto the full time base before being integrated to give a new signal. The new signal should have a continuous derivative with no spikes, such that its high frequency component should now be just noise. This will prevent the high frequency components of a jump in y from bleeding into the noise estimate near the jump. The noise estimate during the jump may not be accurate, but you were never going to get a good fix on that, anyway. The big problem that is solved by this option is that the jump causes spurious spectral noise which extends well before and after the jump itself. Cutting the jump out somehow confines the problem to the relatively narrow time range when the jump itself is happening.

  • restrict_cutoff – bool Some versions of scipy throw an error if cutoff_frequency > nyquist frequency, and others do not. If your version hates high frequency cutoffs, set this to True and cutoff will be reduced to nyqusit - df/2.0, where df is the frequency increment of the FFT, if cutoff >= nyquist.

Returns

1D uncertain float array with length matching t, or set of four such arrays with different estimates. Lowpass smoothed y with uncertainty (dimensions and units match input y) The standard estimate is a hilbert envelope of the high frequency part of the signal, times a constant for correct normalization:

ylf = butterworth_lowpass(y, cutoff_frequency)
yhf = butterworth_highpass(y, cutoff_frequency)
uncertainty = smooth(hilbert_env(yhf)) * norm_factor
return = unumpy.uarray(ylf, uncertainty)

where smoothing of the envelope is accomplished by a butterworth lowpass filter using cutoff_frequency. norm_factor = np.sqrt(0.5) = std_dev(sin(x)) There are other estimates (accessible by setting the debug flag) based on the fluctuation amplitude in the windowed FFT above the cutoff frequency.

omfit_classes.utils_math.ihs_opti_transform(x, shift=False, lmbda=None, off=None)[source]

Calculate the optimal coefficient for the inverse hyperbolic sine transformation of the data

Parameters
  • x – input data

  • shift – calculate and return offset coefficients

  • lmbda – initial guess for lmbda parameter

  • off – initial guess for offset parameter

Returns

lmbda and offset coefficients (offset=0 if shift==False)

omfit_classes.utils_math.ihs(x, lmbda=1.0, off=0.0)[source]

Inverse hyperbolic sine transform: y=np.arcsinh(lmbda*(x-off))/lmbda

Parameters
  • x – input data

  • lmbda – transformation coefficient

  • off – offset coefficient

Returns

transformed data

omfit_classes.utils_math.iihs(y, lmbda=1.0, off=0.0)[source]

Inverse of Inverse hyperbolic sine transform: x=np.sinh(y*lmbda)/lmbda+off

Parameters
  • y – transformed data

  • lmbda – transformation coefficient

  • off – offset coefficient

Returns

original input data

omfit_classes.utils_math.array_snippet(a)[source]

Gives a string summarizing array or list contents: either elements [0, 1, 2, …, -3, -2, -1] or all elements :param a: Array, list, or other compatible iterable to summarize :return: String with summary of array, or just str(a) for short arrays

omfit_classes.utils_math.ordinal(n, fmt='%d%s')[source]

Return ordinal of an integer ‘1st’, ‘2nd’, ‘3rd’, ‘4th’, ‘5th’, … Based on: https://stackoverflow.com/a/20007730/6605826

Parameters
  • n – integer

  • fmt – string format to use (“%d%s” is the default)

Returns

string with ordinal number

omfit_classes.utils_math.array_info(inv)[source]
Parameters

inv – input variable (recarray, np.ndarray)

Returns

string summarizing arrays statistics

exception omfit_classes.utils_math.RomanError[source]

Bases: Exception

exception omfit_classes.utils_math.RomanOutOfRangeError[source]

Bases: omfit_classes.utils_math.RomanError

exception omfit_classes.utils_math.RomanNotIntegerError[source]

Bases: omfit_classes.utils_math.RomanError

exception omfit_classes.utils_math.RomanInvalidNumeralError[source]

Bases: omfit_classes.utils_math.RomanError

omfit_classes.utils_math.toRoman(n)[source]

convert integer to Roman numeral

omfit_classes.utils_math.fromRoman(s)[source]

convert Roman numeral to integer

omfit_classes.utils_math.atomic_element(symbol_A=None, symbol=None, name=None, Z=None, Z_ion=None, mass=None, A=None, abundance=None, use_D_T=True, return_most_abundant=True)[source]

Returns dictionary with name, symbol, symbol_A, Z, mass, A, abundance information of all the elements that match a a given query. Most of the information was gathered from: http://www.sisweb.com/referenc/source/exactmas.htm returns dictionary with name, symbol, symbol_A, Z, mass, A, abundance information of all the elements that match a a given query

Parameters
  • symbol_A – string Atomic symbol followed by the mass number in parenthesis eg. H(2) for Deuterium

  • symbol

    string Atomic symbol

    • can be followed by the mass number eg. H2 for Deuterium

    • can be prepreceded the mass number and followed by the ion charge number eg. 2H1 for Deuterium

  • name – string Long name of the atomic element

  • Z – int Atomic number (proton count in nucleus)

  • Z_ion – int Charge number of the ion (if not specified, it is assumed Z_ion = Z)

  • mass – float Mass of the atomic element in AMU For matching, it will be easier to use A

  • A – int Mass of the atomic element rounded to the closest integer

  • abundance – float Abundance of the atomic element as a fraction 0 < abundance <= 1

  • use_D_T – bool Whether to use deuterium and tritium for isotopes of hydrogen

  • return_most_abundant – bool Whether only the most abundant element should be returned for a query that matches multiple isotopes

Returns

dictionary with all the elements that match a query

omfit_classes.utils_math.element_symbol(z, n=None)[source]

Return the element symbol given the charge (and possibly the number of neutrons)

Parameters
  • z – Ion’s charge (number of protons)

  • n – Ion’s number of nucleons

Returns

Element symbol (1 or 2 character string)

Note

If z==-1, then ‘elec’ is returned

omfit_classes.utils_math.element_label(**kw)[source]

Return a LaTeX formatted string with the element’s symbol, charge state as superscript, and optionally mass number as subscript. :param z_n or zn: int

Nuclear charge / atomic number

Parameters
  • or za (z_a) – int [optional] Ionic charge, including any electrons which may be bound. Defaults to displaying fully-stripped ion if not specified (z_a = z_n).

  • zamin – int Minimum for a range of Zs in a bundled charge state (like in SOLPS-ITER)

  • zamax – int Minimum for a range of Zs in a bundled charge state (like in SOLPS-ITER)

  • or am (a) – int [optional] Mass number. Provides additional filter on results.

Returns

string Prettily formatted symbol

omfit_classes.utils_math.z_from_symbol(symbol)[source]

Given an atomic element symbol it returns it charge Z

Parameters

symbol – element symbol

Returns

element charge Z

omfit_classes.utils_math.splinet(t, y, x, tau)[source]

Tension spline evaluator By VICTOR AGUIAR NUMERICAL ANALYSIS OF KINCAID

Parameters
  • t – nodes location

  • y – value at the nodes

  • x – return values at

  • tau – tension

class omfit_classes.utils_math.CLSQTensionSpline(x, y, t, tau=1, w=None, xbounds=(None, None), min_separation=0, xy_constraints=(), xyprime_constraints=(), optimize_knot_type='xy')[source]

Bases: object

Constrained least square tension spline.

Parameters
  • x – np.ndarray. Data grid.

  • y – np.ndarray. Data values.

  • t – int or np.ndarray. Knot number or locations.

  • tau – float. Tension (higher is smoother).

  • w – np.ndarray. Data weights (usually ~1/std_devs)

  • xbounds – tuple. Minimum and maximum x of knot locations.

  • min_separation – float. Minumum separation between knot locations.

  • xy_constraints – list of tuples. Spline is constrained to have these values at these locations.

  • xyprime_constraints – Spline is constrained to have these derivatives as these locations.

  • optimize_knot_type – str. choose ‘xy’ to simultaneously optimize knot (x,y) values, ‘x’ to optimize x and y separately, and ‘y’ to simply use the prescribed knot locations.

Examples

Using the same data from the LSQUnivariateSpline examples,

>> x = np.linspace(-3, 3, 50) >> y = np.exp(-x**2) + 0.1 * np.random.randn(50)

We can fit a tension spline. We can even set some boundary constraints.

>> t = [-1, 0, 1] >> spl = CLSQTensionSpline(x, y, t, tau=0.1, xy_constraints=[(-3,0)]) >> xs = np.linspace(-3, 3, 1000) >> pyplot.subplots() >> pyplot.plot(x, y, marker=’o’, lw=0) >> pyplot.plot(xs, spl(xs))

Note the xknots are optimized by default. We can compare to the un-optimized knot locations, but (for historical reasons) we wont be able to set constraints.

>> sply = CLSQTensionSpline(x, y, t, tau=0.1, xy_constraints=[(-3,0)], optimize_knot_type=’y’) >> pyplot.plot(xs, sply(xs))

Parameters
  • x – np.ndarray. Data grid.

  • y – np.ndarray. Data values.

  • t – int or np.ndarray. Knot number or locations.

  • tau – float. Tension (higher is smoother).

  • w – np.ndarray. Data weights (usually ~1/std_devs)

  • xbounds – tuple. Minimum and maximum x of knot locations.

  • min_separation – float. Minumum separation between knot locations.

  • xy_constraints – list of tuples. Spline is constrained to have these values at these locations.

  • xyprime_constraints – Spline is constrained to have these derivatives as these locations.

  • optimize_knot_type – str. choose ‘xy’ to simultaneously optimize knot (x,y) values, ‘x’ to optimize x and y separately, and ‘y’ to simply use the prescribed knot locations.

get_knots()[source]

Get the x locations of knots

Returns

tuple. list of (x, y) tuples of knot location and value.

get_knot_locations()[source]

Get the x locations of knots

Returns

tuple. (x, y) knot location and value.

get_knot_values()[source]

Get the x locations of knots

Returns

tuple. (x, y) knot location and value.

class omfit_classes.utils_math.AutoUnivariateSpline(x, y, w=None, bbox=(None, None), k=3, ext=0, check_finite=False, max_interior_knots=None, verbose=False)[source]

Bases: scipy.interpolate.fitpack2.UnivariateSpline

Class of univariate spline with smoothing optimized for reduced chi squared assuming w=1/std_devs.

If no weight is supplied, it is taken to be 1/std_devs(y) and these should all be finite for this to make sense.

Examples

If we have some curve of data with meaningful deviations beyond the errorbars,

>> nx = 50 >> x = np.linspace(-3, 3, nx) >> y = np.exp(-x**2) + 0.3 * (np.random(nx) - 0.5) >> u = unumpy.uarray(y, 0.1 * np.ones(nx)) >> fig, ax = pyplot.subplots() >> uerrorbar(x, u, marker=’o’, label=’data’, ax=ax)

The scipy default spline uses smoothing s = len(x),

>> xs = np.linspace(min(x), max(x), 1000) >> spl = UnivariateSpline(x, y, w = 1/std_devs(u)) >> ax.plot(xs, spl(xs), label=’default spline’)

but the unceratinties spline optimizes the smoothing for reduced chi squared ~1,

>> uspl = AutoUnivariateSpline(x, u) >> ax.plot(xs, uspl(xs), label=’auto spline’) >> ax.legend()

The figure shows the uncertainties spline more accurately captures the meaningful deviations beyond the errorbars. In numbers,

>> default_smooth = spl._data[6] >> default_rchisq = spl.get_residual() / (nx - (len(spl.get_coeffs()) + len(spl.get_knots()) - 2)) >> print(‘Default smoothing is {:.1f} results in reduced chi squared {:.1f}’.format(default_smooth, default_rchisq)) >> print(‘Optimal smoothing of {:.1f} results in reduced chi squared {:.1f}’.format(uspl.get_smoothing_factor(), … uspl.get_reduced_chisqr()))

If the difference is not large, try running again (the deviations are random!).

To see how the optimizer arrived at the result, you can get the full evolution. Remember, it is targeting a reduced chi squared of unity.

>> s, f = uspl.get_evolution(norm=False) >> fig, ax = pyplot.subplots() >> ax.plot(s, f, marker=’o’, ls=’’) # all points tested >> ax.plot([uspl.get_smoothing_factor()], [uspl.get_reduced_chisqr() - 1], marker=’s’) # final value >> ax.set_xscale(‘log’) >> ax.set_xlabel(‘s’) >> ax.set_ylabel(‘Reduced chi squared - 1’)

Parameters
  • x – np.ndarray. Must be increasing

  • y – unumpy.uarray. Uncertainties array from uncertainties.unumpy.

  • w – np.ndarray. Optionally overrides uncertainties from y. Assumed to be 1/std_devs of gaussian errors.

  • bbox – (2,) array_like. 2-sequence specifying the boundary of the approximation interval.

  • k – int. Degree of the smoothing spline. Must be 1 <= k <= 5. Default is k=3, a cubic spline.

  • ext – int or str. Controls the extrapolation mode for elements not in the knot interval. Default 0. if ext=0 or ‘extrapolate’, return the extrapolated value. if ext=1 or ‘zeros’, return 0 if ext=2 or ‘raise’, raise a ValueError if ext=3 of ‘const’, return the boundary value.

  • check_finite – bool. Whether to check that the input arrays contain only finite numbers.

  • max_interior_knots – int. Maximum number of interior knots in a successful optimization. Use this to enforce over smoothing.

  • verbose – bool. Print warnings

get_reduced_chisqr()[source]
Returns

float. Reduced chi squared of the spline.

get_smoothing_factor()[source]
Returns

float. Smoothing factor using in UnivariateSpline.

get_evolution(norm=True)[source]
Parameters

norm – bool. Normalize s to the UnivariateSpline default npts.

Returns

tuple. sev, fev where sev is a record of all smoothing values tried in the optimization and fev is the corresponding reduced chi squared costs.

class omfit_classes.utils_math.CLSQUnivariateSpline(x, y, w=None, bbox=(None, None), k=3, ext=0, check_finite=False, t=None, optimize_knots=True, min_separation=0, xy_constraints=(), xyprime_constraints=(), maxiter=100, verbose=False)[source]

Bases: scipy.interpolate.fitpack2.LSQUnivariateSpline

Constrained least square univariate spline. This class sacrifices the generality of UnivariateSpline’s smoothing but enables the ability to constrain values and/or derivatives of the spline.

The constraints are used in an optimization of the knot locations, not fundamentally enforced in the underlying equations. Thus, constraints far from the natural spline will cause errors.

Examples

Using the same data from the LSQUnivariateSpline examples,

>> x = np.linspace(-2, 2, 50) >> nominal_values = np.exp(-x**2) >> std_devs = 0.1 * np.ones_like(nominal_values) >> y = nominal_values + np.random.normal(loc=0, scale=std_devs)

If we simply use this to fit a spline it will use AutoUnivariateSpline to optimize the knot locations.

>> spl = CLSQUnivariateSpline(x, y, w=1./std_devs) >> xs = np.linspace(-2, 2, 1000) >> pyplot.subplots() >> pyplot.errorbar(x, y, std_devs, marker=’o’, ls=’’) >> pyplot.plot(xs, spl(xs), label=’unconstrained’)

But the new part of this class is that is enables additional constraints on the spline. For example, we can request the spline have zero derivative at the left boundary.

>> splc = CLSQUnivariateSpline(x, y, w=1/std_devs, min_separation=0.1, xyprime_constraints=[(-2,0)]) >> pyplot.plot(xs, splc(xs), label=’constrained’) >> pyplot.legend()

Initialize a instance of a constrained least square univariate spline.

Parameters
  • x – (N,) array_like. Input dimension of data points. Must be increasing

  • y – (N,) array_like. Input dimension of data points

  • w – (N,) array_like. Weights for spline fitting. Must be positive. Default is equal weighting.

  • bbox – (2,) array_like. 2-sequence specifying the boundary of the approximation interval.

  • k – int. Degree of the smoothing spline. Must be 1 <= k <= 5. Default is k=3, a cubic spline.

  • ext – int or str. Controls the extrapolation mode for elements not in the knot interval. Default 0. if ext=0 or ‘extrapolate’, return the extrapolated value. if ext=1 or ‘zeros’, return 0 if ext=2 or ‘raise’, raise a ValueError if ext=3 of ‘const’, return the boundary value.

  • check_finite – bool. Whether to check that the input arrays contain only finite numbers.

  • t – (M,) array_like or integer. Interior knots of the spline in ascending order (t in LSQUnivariateSplien) or maximum number of interior knots (max_interior_knots in AutoUnivariateSpline).

  • optimize_knots – bool. Allow optimizer to change knot locations after initial guess from t or AutoUnivariateSpline.

  • min_separation – float. Minimum separation between knot locations if not explicitely specified by t.

  • xy_constraints – list of tuples. Spline is constrained to have these values at these locations.

  • xyprime_constraints – Spline is constrained to have these derivatives as these locations. x and y separately, and ‘y’ to simply use the prescribed knot locations.

  • maxiter – int. Maximum number of iterations for spline coeff optimization under constraints.

  • verbose – bool. Print warnings

get_reduced_chisqr()[source]
Returns

float. Reduced chi squared of the spline.

class omfit_classes.utils_math.MonteCarloSpline(x, y, w=None, bbox=(None, None), k=3, ext=0, check_finite=False, t=None, optimize_knots=True, min_separation=0, xy_constraints=(), xyprime_constraints=(), maxiter=200, n_trials=100)[source]

Bases: omfit_classes.utils_math.CLSQUnivariateSpline

Monte Carlo Uncertainty propagation through python spline fits.

The concept follows https://gist.github.com/thriveth/4680e3d3cd2cfe561a57 by Th/oger Rivera-Thorsen (thriveth), and essentially forms n_trials unique spline instances with randomly perturbed data assuming w=1/std_devs of gaussian noise.

Note, calling instances of this class returns unumpy.uarrays of Variable objects using the uncertainties package.

Examples

Using the same data from the LSQUnivariateSpline examples,

>> x = np.linspace(-2, 2, 50) >> nominal_values = np.exp(-x**2) >> std_devs = 0.1 * np.ones_like(y) >> y = nominal_values + np.random.normal(loc=0, scale=std_devs)

We can fit a monte carlo spline to get an errorbar on the interpolation.

>> spl = MonteCarloSpline(x, y, w=1/std_devs) >> xs = np.linspace(-2, 2, 1000) >> fig, ax = pyplot.subplots() >> eb = ax.errorbar(x, y, std_devs, marker=’o’, ls=’’) >> ub = uband(xs, spl(xs), label=’unconstrained’)

The individual Monte Carlo splines are CLSQUnivariateSpline instances, so we can set hard constraints as well.

>> splc = MonteCarloSpline(x, y, w=1/std_devs, min_separation=0.1, xyprime_constraints=[(-2,0)]) >> ub = uband(xs, splc(xs), label=’constrained’) >> ax.legend()

Note, this class is a child of the scipy.interpolate.LSQUnivariateSpline class, and has all of the standard spline class methods. Where appropriate, these methods dig into the montecarlo trials to return uncertainties. For example,

>> print(‘knots are fixed at {}’.format(splc.get_knots())) >> print(‘coeffs vary around {}’.format(splc.get_coeffs()))

Initialize a instance of a MonteCarlo constrained least square univariate spline.

Parameters
  • x – (N,) array_like. Input dimension of data points. Must be increasing

  • y – (N,) array_like. Input dimension of data points

  • w – (N,) array_like. Weights for spline fitting. Must be positive. Default is equal weighting.

  • bbox – (2,) array_like. 2-sequence specifying the boundary of the approximation interval.

  • k – int. Degree of the smoothing spline. Must be 1 <= k <= 5. Default is k=3, a cubic spline.

  • ext – int or str. Controls the extrapolation mode for elements not in the knot interval. Default 0. if ext=0 or ‘extrapolate’, return the extrapolated value. if ext=1 or ‘zeros’, return 0 if ext=2 or ‘raise’, raise a ValueError if ext=3 of ‘const’, return the boundary value.

  • check_finite – bool. Whether to check that the input arrays contain only finite numbers.

  • t – (M,) array_like or integer. Interior knots of the spline in ascending order or maximum number of interior knots.

  • optimize_knots – bool. Allow optimizer to change knot locations after initial guess from t or AutoUnivariateSpline.

  • min_separation – float. Minimum separation between knot locations if not explicitely specified by t.

  • xy_constraints – list of tuples. Spline is constrained to have these values at these locations.

  • xyprime_constraints – Spline is constrained to have these derivatives as these locations. x and y separately, and ‘y’ to simply use the prescribed knot locations.

  • maxiter – int. Maximum number of iterations for spline coeff optimization under constraints.

  • n_trials – int. Number of Monte Carlo spline iterations used to form errorbars.

antiderivative(n=1)[source]

Construct a new spline representing the antiderivative of this spline.

Parameters

n – int. Order of antiderivative to evaluate.

derivative(n=1)[source]

Construct a new spline representing the derivative of this spline.

Parameters

n – int. Order of antiderivative to evaluate.

derivatives(x)[source]

Return all derivatives of the spline at the point x.

get_coeffs()[source]

Return spline coefficients.

integral(a, b)[source]

Return definite integral of the spline between two given points.

roots()[source]

Return the zeros of the spline.

set_smoothing_factor(s)[source]

Continue spline computation with the given smoothing factor s and with the knots found at the last call.

omfit_classes.utils_math.contourPaths(x, y, Z, levels, remove_boundary_points=False, smooth_factor=1)[source]
Parameters
  • x – 1D x coordinate

  • y – 1D y coordinate

  • Z – 2D data

  • levels – levels to trace

  • remove_boundary_points – remove traces at the boundary

  • smooth_factor – smooth contours by cranking up grid resolution

Returns

list of segments

omfit_classes.utils_math.get_contour_generator(X, Y, Z, mask, corner_mask, nchunk)[source]
omfit_classes.utils_math.remove_adjacent_duplicates(x)[source]

Remove adjacent duplicate rows in a 2D array :param x: original array :return: array with adjacent duplicates removed

omfit_classes.utils_math.streamPaths(xm, ym, u, v, start_points, minlength=0, maxlength=10000000000.0, bounds_error=True)[source]

Integrate vector field and returns stream line

Params xm

uniformly spaced x grid

Params xm

uniformly spaced y grid

Params u

vector field in the x direction

Params v

vector field in the y direction

Params start_points

2D array of seed x,y coordinate points used to start integration

Params minlength

reject trajectories shorter than this length

Params maxlength

stop tracing trajectory when this length is reached

Parameters

bounds_error – raise error if trajectory starts outside of bounds

omfit_classes.utils_math.line_intersect(path1, path2, return_indices=False)[source]

intersection of two 2D paths

Parameters
  • path1 – array of (x,y) coordinates defining the first path

  • path2 – array of (x,y) coordinates defining the second path

  • return_indices – return indices of segments where intersection occurred

Returns

array of intersection points (x,y)

omfit_classes.utils_math.intersect3d_path_surface(path, surface)[source]

3D intersections of a path and a surface (list of triangles)

Parameters
  • path – list of points in 3D [[X0,Y0,Z0],[X1,Y1,X1],…,[Xn,Yn,Xn]]

  • surface – list of three points in 3D [[[X00,Y0,Z0],[X01,Y01,X01],[X02,Y02,X02]],…,[[Xn0,Yn0,Zn0],[Xn1,Yn1,Xn1],[Xn2,Yn2,Xn2]]]

Returns

list of intersections

omfit_classes.utils_math.intersect3d_path_triangle(path, triangle)[source]

3D intersections of a path and a triangle

Parameters
  • path – list of points in 3D [[X0,Y0,Z0],[X1,Y1,X1],…,[Xn,Yn,Xn]]

  • triangle – three points in 3D [[X0,Y0,Z0],[X1,Y1,X1],[X2,Y2,X2]]

Returns

list of intersections

omfit_classes.utils_math.intersect3d_line_triangle(line, triangle)[source]

3D intersection of a line and a triangle https://stackoverflow.com/questions/42740765/intersection-between-line-and-triangle-in-3d

Parameters
  • line – two points in 3D [[X0,Y0,Z0],[X1,Y1,X1]]

  • triangle – three points in 3D [[X0,Y0,Z0],[X1,Y1,X1],[X2,Y2,X2]]

Returns

None if no intersection is found, or 3D point of intersection

omfit_classes.utils_math.get_s_on_wall(rp, zp, rlim, zlim, slim)[source]

Given R,Z of a point p and curve lim and parameter slim counting distance along the curve, find s at point p.

Primarily intended for mapping back from R,Z to s. Simple interpolation doesn’t work because the rlim, zlim arrays do not monotonically increase.

If the point is not on the curve, return s coordinate at point closest to (rp, zp).

Units are quoted in meters, but it should work if all the units are changed consistently.

Parameters
  • rp – float or 1D float array R coordinate(s) of point(s) of interest in meters

  • zp – float or 1D float array

  • rlim – 1D float array R values of points along the wall in meters

  • zlim – 1D float array Z values corresponding to rlim

  • slim – 1D float array s values corresponding to rlim

Returns

float or 1D float array s at the point(s) on the wall closest to (rp, zp)

omfit_classes.utils_math.point_to_line(px, py, x1, y1, x2, y2, return_closest_point=False)[source]

Calculate minimum distance from a set of points to a set of line segments.

The segments might be defined by consecutive vertices in a polyline. The closest point is closest to the SEGMENT, not the line extended to infinity.

The inputs can be arrays or scalars (that get forced into 1 element arrays). All arrays longer than 1 must have matching length. If (px, py) and (x1, x2, y1, y2) have the same length, the comparison is done for (px[0], py[0]) vs (x1[0], y1[0], x2[0], y2[0]). That is, line 0 (x1[0], …) is only compared to point 0.

All inputs should have matching units.

Parameters
  • px – 1D float array-like X coordinates of test points

  • py – 1D float array-like Y coordinates of test points

  • x1 – 1D float array-like X-coordinates of the first endpoint of each line segment.

  • x2 – 1D float array-like X-coordinates of the second endpoint of each line segment.

  • y1 – 1D float array-like Y-coordinates of the first endpoint of each line segment.

  • y2 – 1D float array-like Y-coordinates of the second endpoint of each line segment.

  • return_closest_point – bool Return the coordinates of the closest points instead of the distances.

Returns

array or tuple of arrays if return_closest_point = True:

tuple of two 1D float arrays with the X and Y coordinates of the closest point on each line segment to each point.

if return_closest_point = False:

1D float array giving the shortest distances between the points and the line segments.

omfit_classes.utils_math.point_in_poly(x, y, poly)[source]

Determine if a point is inside a given polygon or not. Polygon is a list of (x,y) pairs. This function returns True or False. The algorithm is called the “Ray Casting Method”. Source: http://geospatialpython.com/2011/01/point-in-polygon.html , retrieved 20160105 18:39 :param x, y: floats

Coordinates of the point to test

Parameters

poly – List of (x,y) pairs defining a polygon.

Returns

bool Flag indicating whether or not the point is within the polygon.

To test:

polygon = [(0,10),(10,10),(10,0),(0,0)] point_x = 5 point_y = 5 inside = point_in_poly(point_x, point_y, polygon) print(inside)

omfit_classes.utils_math.centroid(x, y)[source]

Calculate centroid of polygon

Parameters
  • x – x coordinates of the polygon

  • y – y coordinates of the polygon

Returns

tuple with x and y coordinates of centroid

omfit_classes.utils_math.usqrt(y)[source]

Handle uncertainties if needed

omfit_classes.utils_math.ulog(y)[source]

Handle uncertainties if needed

omfit_classes.utils_math.gaussian_filter1d(input, sigma, axis=- 1, order=0, output=None, mode='reflect', cval=0.0, truncate=4.0, causal=False)[source]

One-dimensional Gaussian filter. Parameters ———- %(input)s sigma : scalar

standard deviation for Gaussian kernel

%(axis)s order : int, optional

An order of 0 corresponds to convolution with a Gaussian kernel. A positive order corresponds to convolution with that derivative of a Gaussian.

%(output)s %(mode)s %(cval)s truncate : float, optional

Truncate the filter at this many standard deviations. Default is 4.0.

causalbool or sequence of bools

Remove all forward weightings.

gaussian_filter1d : np.ndarray Examples ——– >> from scipy.ndimage import gaussian_filter1d >> gaussian_filter1d([1.0, 2.0, 3.0, 4.0, 5.0], 1) np.array([ 1.42704095, 2.06782203, 3. , 3.93217797, 4.57295905]) >> gaussian_filter1d([1.0, 2.0, 3.0, 4.0, 5.0], 4) np.array([ 2.91948343, 2.95023502, 3. , 3.04976498, 3.08051657]) >> from matplotlib import pyplot >> np.random.seed(280490) >> x = np.random.randn(101).cumsum() >> y3 = gaussian_filter1d(x, 3) >> y6 = gaussian_filter1d(x, 6) >> pyplot.plot(x, ‘k’, label=’original data’) >> pyplot.plot(y3, ‘–’, label=’filtered, sigma=3’) >> pyplot.plot(y6, ‘:’, label=’filtered, sigma=6’) >> pyplot.legend() >> pyplot.grid() >> pyplot.show()

Causality example, >> x = np.arange(0,100) >> y = 1.* (x > 40) * (x < 60) >> fig, ax = pyplot.subplots() >> ax.plot(x, y, ‘x-‘) >> ax.plot(x, gaussian_filter1d(y,3.)) >> ax.plot(x, gaussian_filter1d(y,3., causal=True)) >> ax.set_ylim(0,1.5)

omfit_classes.utils_math.gaussian_filter(input, sigma, order=0, output=None, mode='reflect', cval=0.0, truncate=4.0, causal=False)[source]

Multidimensional Gaussian filter. Parameters ———- %(input)s sigma : scalar or sequence of scalars

Standard deviation for Gaussian kernel. The standard deviations of the Gaussian filter are given for each axis as a sequence, or as a single number, in which case it is equal for all axes.

orderint or sequence of ints, optional

The order of the filter along each axis is given as a sequence of integers, or as a single number. An order of 0 corresponds to convolution with a Gaussian kernel. A positive order corresponds to convolution with that derivative of a Gaussian.

%(output)s %(mode_multiple)s %(cval)s truncate : float

Truncate the filter at this many standard deviations. Default is 4.0.

causalbool or sequence of bools

Remove all forward weightings.

gaussian_filternp.ndarray

Returned array of same shape as input.

The multidimensional filter is implemented as a sequence of one-dimensional convolution filters. The intermediate arrays are stored in the same data type as the output. Therefore, for output types with a limited precision, the results may be imprecise because intermediate results may be stored with insufficient precision. Examples ——– >> from scipy.ndimage import gaussian_filter >> a = np.arange(50, step=2).reshape((5,5)) >> a np.array([[ 0, 2, 4, 6, 8],

[10, 12, 14, 16, 18], [20, 22, 24, 26, 28], [30, 32, 34, 36, 38], [40, 42, 44, 46, 48]])

>> gaussian_filter(a, sigma=1) np.array([[ 4, 6, 8, 9, 11],

[10, 12, 14, 15, 17], [20, 22, 24, 25, 27], [29, 31, 33, 34, 36], [35, 37, 39, 40, 42]])

>> from scipy import misc >> from matplotlib import pyplot >> fig = pyplot.figure() >> pyplot.gray() # show the filtered result in grayscale >> ax1 = fig.add_subplot(121) # left side >> ax2 = fig.add_subplot(122) # right side >> ascent = misc.ascent() >> result = gaussian_filter(ascent, sigma=5) >> ax1.imshow(ascent) >> ax2.imshow(result) >> pyplot.show()

Here is a nice little demo of the added OMFIT causal feature,

>> nn = 24 >> a = np.zeros((nn, nn)) >> a[int(nn//2),int(nn//2)] = 1 >> fig, axs = pyplot.subplots(2, 2) >> ax = axs[0, 0] >> ax.imshow(scipy.ndimage.gaussian_filter(a, 3, mode=’nearest’), origin=’lower’) >> ax.set_title(‘scipy version’) >> ax = axs[0, 1] >> ax.imshow(gaussian_filter(a, 3, mode=’nearest’), origin=’lower’) >> ax.set_title(‘OMFIT version’) >> ax = axs[1, 0] >> ax.imshow(gaussian_filter(a, 3, causal=True, mode=’nearest’), origin=’lower’) >> ax.set_title(‘causal True’) >> ax = axs[1, 1] >> ax.imshow(gaussian_filter(a,3, causal=(True, False), mode=’nearest’), origin=’lower’) >> ax.set_title(‘causal True, False’)

omfit_classes.utils_math.bicoherence(s1, s2=None, fs=1.0, nperseg=None, noverlap=None, **kwargs)[source]

Compute the bicoherence between two signals of the same lengths s1 and s2 using the function scipy.signal.spectrogram.

Sourced from https://stackoverflow.com/questions/4194554/function-for-computing-bicoherence

Parameters
  • s1 – ndarray. Time series of measurement values

  • s2 – ndarray. Time series of measurement values (default of None uses s1)

  • fs – Sampling frequency of the x time series. Defaults to 1.0.

  • nperseg – int. Length of each segment. Defaults to None, but if window is str or tuple, is set to 256, and if window is array_like, is set to the length of the window.

  • noverlap – int. Number of points to overlap between segments. If None, noverlap = nperseg // 8. Defaults to None.

  • kwargs – All additional key word arguments are passed to signal.spectrogram

Return f, bicoherence

array of frequencies and matrix of bicoherence at those frequencies

utils_fusion

omfit_classes.utils_fusion.gyroradius(T, Bt, Z, m)[source]

This function calculates plasma gyroradius

Parameters
  • T – Ion temperature [eV]

  • Bt – Magnetic field [T]

  • Z – charge

  • m – mass [AMU]

Returns

gyroradius [m]

omfit_classes.utils_fusion.banana_width(T, Bt, Z, m, epsilon, q)[source]

This function estimates the banana orbit width

Parameters
  • T – Temperature [eV]

  • Bt – Magnetic field [T]

  • Z – Charge

  • m – Mass [AMU]

  • epsilon – Inverse aspect ratio

  • q – Safety factor

Returns

Estimate of banana orbit width [m]

omfit_classes.utils_fusion.lambda_debye(**keyw)[source]

Debye length [m]

Needs ne [m^-3], te [eV]

Formula: \(\sqrt{\frac{\varepsilon_0 T_e}{q_e n_e}}\)

omfit_classes.utils_fusion.bmin_classical(**keyw)[source]

Classical distance of minimum approach [m]

Needs zeff [-], te [eV]

omfit_classes.utils_fusion.bmin_qm(**keyw)[source]

Quantum mechanical distance of minimum approach [m]

Needs te [eV]

omfit_classes.utils_fusion.ln_Lambda(**keyw)[source]

Coulomb logarithm [-]

Needs te [eV], zeff [-], ne [m^-3]

omfit_classes.utils_fusion.nu_e(**keyw)[source]

Electron collision frequency [1/s] Eq. (A9) in UW CPTC 09-6R

Needs te [eV], zeff [-], ne [m^-3]

omfit_classes.utils_fusion.vTe(**keyw)[source]

Electron thermal velocity [m/s]

Needs te [eV]

omfit_classes.utils_fusion.lambda_e(**keyw)[source]

Electron mean free path [m]

Needs te [eV], zeff [-], ne [m^-3]

omfit_classes.utils_fusion.omega_transit_e(**keyw)[source]

Electron transit frequency [1/s]

Needs q [-], R0 [m], te [eV]

omfit_classes.utils_fusion.epsilon(**keyw)[source]

Inverse aspect ratio [-]

Needs (rho [m], R0 [m]) or (r_minor, R_major)

omfit_classes.utils_fusion.f_c(**keyw)[source]

Flow-weighted fraction of circulating particles

Needs epsilon inputs

omfit_classes.utils_fusion.f_t(**keyw)[source]

Flow-weighted fraction of trapped particles

Needs epsilon inputs

omfit_classes.utils_fusion.fsa_B2_over_fsa_b_dot_B(**keyw)[source]

Approximation of geometric factor <B_0^2>/<(b.grad(B_0))^2>. [m^-2]

Needs R0 [m], q [-], epsilon inputs

omfit_classes.utils_fusion.nu_star_e(**keyw)[source]

Electron collisionality parameter [-]

Needs R0 [m], q [-], te [eV], zeff [-], ne [m^-3], epsilon inputs

omfit_classes.utils_fusion.omega_transit_e_tau(**keyw)[source]

Dimensionless omega_transit [-]

Needs te [eV], zeff [-], ne [m^-3], q [-], R0 [m]

omfit_classes.utils_fusion.M_visc_e(**keyw)[source]

Dimensionless electron viscosity matrix M_e [-] 2 x 2 x len(ne)

Needs te [eV], zeff [-], ne [m^-3], q [-], R0 [m]

omfit_classes.utils_fusion.N_fric_e(**keyw)[source]

Dimensionless electron friction matrix N_e [-] 2 x 2 x len(zeff)

Needs zeff [-]

omfit_classes.utils_fusion.inverse_N(**keyw)[source]

Inverse of the electron friction matrix [-]

Needs zeff

omfit_classes.utils_fusion.Spitzer_factor(**keyw)[source]

Spitzer resistivity factor [-]

Needs zeff [-]

omfit_classes.utils_fusion.inverse_NM(**keyw)[source]

Inverse of the sum N_fric_e + M_fric_e [-]

Needs te [eV], zeff [-], ne [m^-3], q [-], R0 [m]

omfit_classes.utils_fusion.resistive_factor(**keyw)[source]

Resistive factor [-]

Needs te [eV], zeff [-], ne [m^-3], q [-], R0 [m]

omfit_classes.utils_fusion.eta_0(**keyw)[source]

Reference resistivity [ohm-m]

Needs ne [m^-3], te [eV], zeff [-]

omfit_classes.utils_fusion.eta_par_neo(**keyw)[source]

Parallel neoclassical resistivity [ohm-m]

Needs te [eV], zeff [-], ne [m^-3], q [-], R0 [m]

omfit_classes.utils_fusion.q_rationals(x, q, nlist, mlist=None, doPlot=False)[source]

Evaluate rational flux surfaces (m/n) give a q profile

Parameters
  • x – x axis over which q profile is defined

  • q – q profile

  • nlist – list of n mode number

  • mlist – list of m mode number (if None returns all possible (m/n) rationals)

  • doPlot – Plot (either True or matplotlib axes)

Returns

dictionary with (m,n) keys and x intersections as values

omfit_classes.utils_fusion.tokamak(tokamak, output_style='OMFIT', allow_not_recognized=True, translation_dict=None)[source]

function that sanitizes user input tokamak in a format that is recognized by other codes

Parameters
  • tokamak – user string of the tokamak

  • output_style – format of the tokamak used for the output one of [‘OMFIT’,’TRANSP’,’GPEC’]

  • allow_not_recognized – allow a user to enter a tokamak which is not recognized

  • translation_dict – dictionary used for further translation. This is handy for example in situations where we want to get the same string back independently of whether it is a older tokamak or its upgraded version. For example tokamak(‘NSTX-U’, translation_dict={‘NSTXU’:’NSTX’})

Returns

sanitized string

omfit_classes.utils_fusion.is_device(devA, devB)[source]

Compare strings or lists of strings for equivalence in tokamak name

Parameters
  • devA – A string or list of strings

  • devB – A string or list of strings

Returns

True or False

Example: is_device(‘D3D’,[‘DIII-D’,’MAST’])

omfit_classes.utils_fusion.device_specs(device='DIII-D')[source]

Returns a dictionary of information that is specific to a particular tokamak

Parameters

device – The name of a tokamak. It will be evaluated with the tokamak() function so variation in spelling and capitalization will be tolerated. This function has some additional translation capabilities for associating MDSplus servers with tokamaks; for example, “EAST_US”, which is the entry used to access the eastdata.gat.com MDSplus mirror, will be translated to EAST.

Returns

A dictionary with as many static device measurements as are known

omfit_classes.utils_fusion.nclass_conductivity(psi_N=None, Te=None, ne=None, Ti=None, q=None, eps=None, R=None, fT=None, volume=None, Zeff=None, nis=None, Zis=None, Zdom=None, version='osborne', return_info_pack=False, plot_slice=None, sigma_compare=None, sigma_compare_label='Input for comparison', spitzer_compare=None, spitzer_compare_label='Spitzer input for comparison', charge_number_to_use_in_ion_collisionality='Koh', charge_number_to_use_in_ion_lnLambda='Zavg')[source]

Calculation of neoclassical conductivity

See: O. Sauter, et al., Phys. Plasmas 6, 2834 (1999); doi:10.1063/1.873240 Neoclassical conductivity appears in equations: 5, 7, 13a, and unnumbered equations in the conclusion

Other references:

S Koh, et al., Phys. Plasmas 19, 072505 (2012); doi: 10.1063/1.4736953

for dealing with ion charge number when there are multiple species

T Osborne, “efit.py Kinetic EFIT Method”, private communication (2013);

this is a word file with a description of equations used to form the current profile constraint

O Sauter, et al., Phys. Plasmas 9, 5140 (2002); doi:10.1063/1.1517052

this has corrections for Sauter 1999 but it also has a note on what Z to use in which equations; it argues that ion equations should use the charge number of the main ions for Z instead of the ion effective charge number from Koh 2012

Sauter website

Accurate neoclassical resistivity, bootstrap current and other transport coefficients (Fortran 90 subroutines and matlab functions): has some code that was used to check the calculations in this script (BScoeff.m, nustar.m, sigmaneo.m, jdotB_BS.m)

GACODE NEO source

Calculations from NEO (E. A. Belli)

Update August 2021add new set of analytical formulae for the computation of the neoclassical condactivity from

A.Redl, et al., Phys. Plasma 28, 022502 (2021) https://doi.org/10.1063/5.0012664 and all relevant variables are mentioned as neo_2021

This function was initially written as part of the Kolemen Group Automatic Kinetic EFIT Project (auto_kEFIT).

Parameters
  • psi_N – position basis for all profiles, required only for plotting (normalized poloidal magnetic flux)

  • Te – electron temperature in eV as a function of time and position (time should be first axis, then position)

  • ne – electron density in m^-3 (vs. time and psi)

  • Ti – ion temperature in eV

  • Zeff – [optional if nis and Zis are provided] effective charge state of ions = sum_j(n_j (Z_j)^2)/sum_j(n_j Z_j) where j is ion species (this is probably a sum over deuterium and carbon)

  • nis – [optional if Zeff is provided] list of ion densities in m^-3

  • Zis – [optional if Zeff is provided] ion charge states (list of scalars)

  • Zdom – [might be optional] specify the charge number of the dominant ion species. Defaults to the one with the highest total number of particles (volume integral of ni). If using the estimation method where only Zeff is provided, then Zdom is assumed to be 1 if not provided.

  • q – safety factor

  • eps – inverse aspect ratio

  • R – major radius of the geometric center of each flux surface

  • fT – trapped particles fraction

  • volume – [not needed if Zdom is provided, unlikely to cause trouble if not provided even when “needed”] volume enclosed by each flux surface, used to identify dominant ion species if dominant ion species is not defined explicitly by doing a volume integral (so we need this so we can get dV/dpsiN). If volume is needed but not provided, it will be crudely estimated. Errors in the volume integral are very unlikely to affect the selection of the dominant ion species (it would have to be a close call to begin with), so it is not critical that volume be provided with high quality, if at all.

  • return_info_pack – Boolean: If true, returns a dictionary full of many intermediate variables from this calculation instead of just conductivity

  • plot_slice – Set to the index of the timeslice to plot in order to plot one timeslice of the calculation, including input profiles and intermediate quantities. Set to None for no plot (default)

  • sigma_compare – provide a conductivity profile for comparison in Ohm^-1 m^-1

  • sigma_compare_label – plot label to use with sigma_compare

  • spitzer_compare – provide another conductivity profile for comparison (so you can compare neoclassical and spitzer) (Ohm^1 m^1)

  • spitzer_compare_label – plot label to use with spitzer_compare

  • charge_number_to_use_in_ion_collisionality

    instruction for replacing single ion species charge number Z in nuistar equation when going to multi-ion species plasma. Options are: [‘Koh’, ‘Dominant’, ‘Zeff’, ‘Zavg’, ‘Koh_avg’]

    Dominant uses charge number of ion species which contributed the most electrons (recommended by Sauter 2002) Koh uses expression from Koh 2012 page 072505-11 evaluated for dominant ion species (recommended by Koh 2012) Koh_avg evaluates Koh for all ion species and then averages over species Zeff uses Z_eff (No paper recommends using this but it appears to be used by ONETWO) Zavg uses ne/sum(ni) (Koh 2012 recommends using this except for collision frequency)

    Use Koh for best agreement with TRANSP

  • charge_number_to_use_in_ion_lnLambda

    instruction for replacing single ion species charge number Z in lnLambda equation when going to multi-ion species plasma. Options are: [‘Koh’, ‘Dominant’, ‘Zeff’, ‘Zavg’, ‘Koh_avg’]

    Use Koh for best agreement with TRANSP

Returns

neoclassical conductivity in (Ohm^-1 m^-1) as a function of time and input psi_N (after interpolation/extrapolation). If output with “return_info_pack”, the return is a dictionary containing several intermediate variables which are used in the calculation (collisionality, lnLambda, etc.)

omfit_classes.utils_fusion.nclass_conductivity_from_gfile(psi_N=None, Te=None, ne=None, Ti=None, gEQDSK=None, Zeff=None, nis=None, Zis=None, Zdom=None, return_info_pack=False, plot_slice=None, charge_number_to_use_in_ion_collisionality='Koh', charge_number_to_use_in_ion_lnLambda='Zavg')[source]

WRAPPER FOR nclass_conductivity THAT EXTRACTS GFILE STUFF AND INTERPOLATES FOR YOU Calculation of neoclassical conductivity See: O. Sauter, et al., Phys. Plasmas 6, 2834 (1999); doi:10.1063/1.873240 Neoclassical conductivity appears in equations: 5, 7, 13a, and unnumbered equations in the conclusion

This function was initially written as part of the Kolemen Group Automatic Kinetic EFIT Project (auto_kEFIT).

Parameters
  • psi_N – position basis for all non-gfile profiles

  • Te – electron temperature in eV as a function of time and position (time should be first axis, then position)

  • ne – electron density in m^-3 (vs. time and psi)

  • Ti – ion temperature in eV

  • Zeff – [optional if nis and Zis are provided] effective charge state of ions = sum_j(n_j (Z_j)^2)/sum_j(n_j Z_j) where j is ion species (this is probably a sum over deuterium and carbon)

  • nis – [optional if Zeff is provided] list of ion densities in m^-3

  • Zis – [optional if Zeff is provided] ion charge states (list of scalars)

  • Zdom – [might be optional] specify the charge number of the dominant ion species. Defaults to the one with the highest total number of particles (volume integral of ni). If using the estimation method where only Zeff is provided, then Zdom is assumed to be 1 if not provided.

  • gEQDSK – an OMFITcollection of g-files or a single g-file as an instance of OMFITgeqdsk

  • return_info_pack – Boolean: If true, returns a dictionary full of many intermediate variables from this calculation instead of just conductivity

  • plot_slice – Set to the index of the timeslice to plot in order to plot one timeslice of the calculation, including input profiles and intermediate quantities. Set to None for no plot (default)

  • charge_number_to_use_in_ion_collisionality

    instruction for replacing single ion species charge number Z in nuistar equation when going to multi-ion species plasma. Options are: [‘Koh’, ‘Dominant’, ‘Zeff’, ‘Zavg’, ‘Koh_avg’]

    Dominant uses charge number of ion species which contributed the most electrons (recommended by Sauter 2002) Koh uses expression from Koh 2012 page 072505-11 evaluated for dominant ion species (recommended by Koh 2012) Koh_avg evaluates Koh for all ion species and then averages over species Zeff uses Z_eff (No paper recommends using this but it appears to be used by ONETWO) Zavg uses ne/sum(ni) (Koh 2012 recommends using this except for collision frequency)

    Use Koh for best agreement with TRANSP

  • charge_number_to_use_in_ion_lnLambda

    instruction for replacing single ion species charge number Z in lnLambda equation when going to multi-ion species plasma. Options are: [‘Koh’, ‘Dominant’, ‘Zeff’, ‘Zavg’, ‘Koh_avg’]

    Use Koh for best agreement with TRANSP

Returns

neoclassical conductivity in (Ohm^-1 m^-1) as a function of time and input psi_N (after interpolation/extrapolation). If output with “return_info_pack”, the return is a dictionary containing several intermediate variables which are used in the calculation (collisionality, lnLambda, etc.)

omfit_classes.utils_fusion.sauter_bootstrap(psi_N=None, Te=None, Ti=None, ne=None, p=None, nis=None, Zis=None, Zeff=None, gEQDSKs=None, R0=None, device='DIII-D', psi_N_efit=None, psiraw=None, R=None, eps=None, q=None, fT=None, I_psi=None, nt=None, version='osborne', debug_plots=False, return_units=False, return_package=False, charge_number_to_use_in_ion_collisionality='Koh', charge_number_to_use_in_ion_lnLambda='Zavg', dT_e_dpsi=None, dT_i_dpsi=None, dn_e_dpsi=None, dnis_dpsi=None)[source]

Sauter’s formula for bootstrap current

See: O. Sauter, et al., Phys. Plasmas 6, 2834 (1999); doi:10.1063/1.873240

Other references:

S Koh, et al., Phys. Plasmas 19, 072505 (2012); doi: 10.1063/1.4736953

for dealing with ion charge number when there are multiple species

T Osborne, “efit.py Kinetic EFIT Method”, private communication (2013);

this is a word file with a description of equations used to form the current profile constraint

O Sauter, et al., Phys. Plasmas 9, 5140 (2002); doi:10.1063/1.1517052

this has corrections for Sauter 1999 but it also has a note on what Z to use in which equations; it argues that ion equations should use the charge number of the main ions for Z instead of the ion effective charge number from Koh 2012

Sauter website

Accurate neoclassical resistivity, bootstrap current and other transport coefficients (Fortran 90 subroutines and matlab functions): has some code that was used to check the calculations in this script (BScoeff.m, nustar.m, sigmaneo.m, jdotB_BS.m)

GACODE NEO source

Calculations from NEO (E. A. Belli)

Y R Lin-Liu, et al., “Zoo of j’s”, DIII-D physics memo (1996);

got hardcopy from Sterling Smith & photocopied

Update August 2021add new set of analytical formulae for the computation of the neoclassical condactivity from

A.Redl, et al., Phys. Plasma 28, 022502 (2021) https://doi.org/10.1063/5.0012664 and all relevant variables are mentioned as neo_2021

This function was initially written as part of the Kolemen Group Automatic Kinetic EFIT Project (auto_kEFIT).

Parameters
  • psi_N – normalized poloidal magnetic flux as a position coordinate for input profiles Te, Ti, ne, etc.

  • Te – electron temperature in eV, first dimension: time, second dimension: psi_N

  • Ti – ion temperature in eV, 2D with dimensions matching Te (time first)

  • ne – electron density in m^-3, dimensions matching Te

  • p – total pressure in Pa, dimensions matching Te

  • Zeff – [optional if nis and Zis are provided] effective charge state of ions = sum_j(n_j (Z_j)^2)/sum_j(n_j Z_j) where j is ion species (this is probably a sum over deuterium and carbon)

  • nis – [optional if Zeff is provided] list of ion densities in m^-3

  • Zis – [optional if Zeff is provided] ion charge states (list of scalars)

  • R0 – [optional if device is provided and recognized] The geometric center of the tokamak’s vacuum vessel in m. (For DIII-D, this is 1.6955 m (Osborne, Lin-Liu))

  • device – [used only if R0 is not provided] The name of a tokamak for the purpose of looking up R0

  • gEQDSKs

    a collection of g-files from which many parameters will be derived. The following quantities are taken from g-files if ANY of the required ones are missing:

    param psi_N_efit

    [optional] psi_N for the EFIT quantities if different from psi_N for kinetic profiles

    param nt

    [optional] number of time slices in equilibrium data (if you don’t want to tell us, we will measure the shape of the array)

    param psiraw

    poloidal flux before normalization (psi_N is derived from this).

    param R

    major radius coordinate R of each flux surface’s geometric center in m

    param q

    safety factor (inverse rotation transform)

    param eps

    inverse aspect ratio of each flux surface: a/R

    param fT

    trapped particle fraction on each flux surface

    param I_psi

    also known as F = R*Bt, averaged over each flux surface

  • version

    which quantity to return: ‘jB_fsa’ is the object directly from Sauter’s paper: 2nd term on RHS of last equation in conclusion. ‘osborne’ is jB_fsa w/ |I_psi| replaced by R0. Motivated by memo from T. Osborne about kinetic EFITs ‘jboot1’ is 2nd in 1st equation of conclusion of Sauter 1999 w/ correction from Sauter 2002 erratum. ‘jboot1BROKEN’ is jboot1 without correction from Sauter 2002 (THIS IS FOR TESTING/COMPARISON ONLY) ‘neo_2021’ a new set of analytical coefficients from A.Redl, et al. (the new set of analytical formulae consists of the same analytical structure as the ‘jboot1’ and ‘jboot1BROKEN’ )

    You should use jboot1 if you want <J.B> You should use osborne if you want J *** Put this into current_to_efit_form() to make an EFIT You should use jboot1 or jB_fsa to compare to Sauter’s paper, equations 1 and 2 of the conclusion You should use jboot1BROKEN to compare to Sauter 1999 without the 2002 correction

  • debug_plots – plot internal quantities for debugging

  • return_units – If False: returns just the current profiles in one 2D array. If True: returns a 3 element tuple containing the current profiles, a plain string containing the units, and a formatted string containing the units

  • return_package – instead of just a current profile, return a dictionary containing the current profile as well as other information

  • charge_number_to_use_in_ion_collisionality

    instruction for replacing single ion species charge number Z in nuistar equation when going to multi-ion species plasma. Options are: [‘Koh’, ‘Dominant’, ‘Zeff’, ‘Zavg’, ‘Koh_avg’]

    Dominant uses charge number of ion species which contributed the most electrons (recommended by Sauter 2002) Koh uses expression from Koh 2012 page 072505-11 evaluated for dominant ion species (recommended by Koh 2012) Koh_avg evaluates Koh for all ion species and then averages over species Zeff uses Z_eff (No paper recommends using this but it appears to be used by ONETWO) Zavg uses ne/sum(ni) (Koh 2012 recommends using this except for collision frequency)

    Use Koh for best agreement with TRANSP Use Zavg for best agreement with recommendations by Koh 2012

  • charge_number_to_use_in_ion_lnLambda

    instruction for replacing single ion species charge number Z in lnLambda equation when going to multi-ion species plasma. Options are: [‘Koh’, ‘Dominant’, ‘Zeff’, ‘Zavg’, ‘Koh_avg’]

    Use Koh for best agreement with TRANSP Use Koh for best agreement with recommendations by Koh 2012

Return jB

flux surface averaged j_bootstrap * B with some modifications according to which version you select

Return units

[only if return_units==True] a string with units like “A/m^2”

Return units_format

[only if return_units==True] a TeX formatted string with units like “$A/m^2$” (can be included in plot labels directly)

This is first equation in the conclusion of Sauter 1999 (which is equation 5 with stuff plugged in) (with correction from the erratum (Sauter 2002):

<j_par * B> = sigma_neo * <E_par * B> - I(psi) p_e *
                  [L_31 * p/p_e * d_psi(ln(p)) + L_32 * d_psi(ln(T_e))
                  + L_34 * alpha * (1-R_pe)/R_pe * d_psi(ln(T_i))]

The second equation in the conclusion is nicer looking:

<j_par * B> = sigma_new * <E_par * B> - I(psi) p *
               [L_31 d_psi(ln(n_e)) + R_pe * (L_31 + L_32) * d_psi(ln(T_e))
               + (1-R_pe)*(1+alpha*L_34/L_31)*L_31 * d_psi(ln(T_i))]

In both equations, the first term is ohmic current and the second term is bootstrap current. The second equation uses some approximations which make the result much smoother. The first equation had an error in the original Sauter 1999 paper that was corrected in the 2002 erratum.

  • < > denotes a flux surface average (mentioned on page 2835)

  • j_par is the parallel current (parallel to the magnetic field B) (this is what we’re trying to find)

  • B is the total magnetic field

  • sigma_neo is the neoclassical conductivity given by equation 7 on page 2835 or equation 13 on page 2837 (this is mentioned as neoclassical resistivity on page 2836, but the form of the equation clearly shows that it is conductivity, the reciprocal of resistivity. Also the caption of figure 2 confirms that conductivity is what is meant.)

  • E_par is the parallel electric field

  • I(psi) = R * B_phi (page 2835)

  • p_e is the electron pressure

  • L_31, L_32, and L_34 are given by equations 8, 9, and 10 respectively (eqns on page 2835). Also they are given again by eqns 14-16 on pages 2837-2838

  • p is the total pressure

  • d_psi() is the derivative with respect to psi (not psi_N)

  • T_e is the electron temperature

  • alpha is given by equation 11 on page 2835 or by eqn 17a on page 2838

  • T_i is the ion temperature

  • R_pe = p_e/p

  • f_T the trapped particle fraction appears in many equations and is given by equation 12 on page 2835 but also in equation 18b with nu_i* in equation 18c

useful quantities are found in equation 18

omfit_classes.utils_fusion.current_to_efit_form(r0, inv_r, cross_sec, total_current, x)[source]

Conversion of current density to EFIT constraint format. Adapted from currentConstraint.py by O. Meneghini

Parameters
  • r0 – major radius of the geometric center of the vacuum vessel (1.6955 m for DIII-D) (scalar)

  • inv_r – flux surface average (1/R); units should be reciprocal of r0 (function of position or function of position and time)

  • cross_sec – cross sectional area of the plasma in m^2 (scalar or function of time

  • total_current – total plasma current in A (scalar or function of time)

  • x – input current density to be converted in A/m^2 (function of position or function of position and time)

Returns

x normalized to EFIT format (function of position or function of position and time)

omfit_classes.utils_fusion.estimate_ohmic_current_profile(cx_area, sigma, itot, jbs=None, ibs=None, jdriven=None, idriven=None)[source]

Estimate the profile of ohmic current using total current, the profile of bootstrap and driven current, and neoclassical conductivity. The total Ohmic current profile is calculated by integrating bootstrap and driven current and subtracting this from the total current. The Ohmic current profile is assigned assuming flat loop voltage and the total is scaled to match the estimated total Ohmic current.

All inputs should be on the same coordinate system with the same dimensions, except itot, ibs, and idriven should lack the position axis. If inputs have more than one dimension, position should be along the axis with index = 1 (the second dimension).

This function was initially written as part of the Kolemen Group Automatic Kinetic EFIT Project (auto_kEFIT).

Parameters
  • cx_area – Cross sectional area enclosed by each flux surface as a function of psin in m^2

  • sigma – Neoclassical conductivity in Ohm^-1 m^-1

  • itot – Total plasma current in A

  • jbs – [optional if ibs is provided] Bootstrap current density profile in A/m^2. If this comes from sauter_bootstrap(), the recommended version is ‘osborne’

  • ibs – [optional if jbs is provided] Total bootstrap current in A

  • jdriven – [optional if idriven is provided] Driven current density profile in A/m^2

  • idriven – [optional if jdriven is provided] Total driven current in A

Returns

Ohmic current profile as a function of psin in A/m^2

omfit_classes.utils_fusion.intrinsic_rotation(geo_a, geo_R, geo_Rx, R_mp, rho, I_p_sgn, Te, Ti, q, fc, B0, rho_ped, rho_sep=1.0, C_phi=1.0, d_c=1.0)[source]

Tim Stoltzfus-Dueck & Arash Ashourvan intrinsic rotation model

Parameters
  • geo_a – [m] plasma minor radius evaluated at the midplane

  • geo_R – [m] plasa major radius evaluated at the midplane

  • geo_Rx – [m] radial position of the X point

  • R_mp – [m] midplane radial coordinate from on-axis to the separatrix (LFS)

  • rho – normalised sqrt(toroidal flux)

  • I_p_sgn – sign of I_p to get the correct rotation direction, positive rotation is alway co-current

  • Te – [eV] electron temperature profile

  • Ti – [eV] ion temperature profile

  • q – safety factor/q profile

  • fc – Flux surface averaged passing particles fraction profile

  • B0 – [T] Magnetic field on axis

  • rho_ped – rho value at pedestal top

  • rho_sep – rho value at separatrix (/pedestal foot)

  • C_phi – constant that translates Te scale length to turbulence scale length. default value = 1.75, range: [1.0,2.0]/[0.5,4]

  • d_c – ballooning parameter for turbulence, where 0.0 is symmetric in ballooning angle, 2.0 is all at LFS. default value = 1.0, range: [0.0,2.0]

Return omega_int

[rad/s] intrinsic plasma rotation at pedestal top

omfit_classes.utils_fusion.standardize_gas_species_name(name, output_form='simple', on_fail='raise')[source]

Standardizes gas species names that could come from different/unknown sources.

These include common impurity gas molecules, so nitrogen and deuterium translate to N2 and D2, since they form diatomic molecules.

For example, N2, N$_2$, N_2, and nitrogen all mean the same thing. This function should accept any one of those and turn it into the format you request.

Intended to handle exotic capitaliation; e.g. corrects NE or ne into Ne.

Parameters
  • name – str The name of the species

  • output_form

    str simple: the shortest, simplest, unambiguous and correct form. E.g.: N2 latex: latex formatting for subscripts. E.g.: N$_2$ name: the name of the species. E.g.: nitrogen markup: symbol with punctuation like _ to indicate underscores, but no $. E.g.: N_2 atom: just the symbol for the atom without molecular information. E.g.: N

    This isn’t recommended as an output format; it’s here to simplify lookup in case someone indicates nitrogen gas by N instead of N2. Doesn’t work for mixed-element molecules.

  • on_fail

    str Behavior on lookup failure.

    raise: raise OMFITexception print: print an error message and return None quiet: quietly return None

Returns

str

omfit_classes.utils_fusion.lookup_d3d_gas_species(shot, valve=None, output_form=None)[source]

Retrieves information on the gas species loaded into DIII-D gas injector(s)

Parameters
  • shot – int Shot to look up. The species loaded into each valve can change from shot to shot.

  • valve – str [optional] Name of the gas valve or injector, like ‘A’ or ‘PFX1’. ‘GASA’ and ‘A’ are interchangeable. The valve name is not case sensitive. If provided, returns a string with the gas species for this valve, or raises OMFITexception on failure. If not provided, returns a dictionary of all the gas valves and species in the database.

  • name_format

    str [optional] Reformat the name of the gas species to conform to a preferred standard. Options (with examples) are:

    simple: N2 latex: N$_2$ name: nitrogen markup: N_2

    Set to None to get the string as written in the database.

Returns

str or dict Either the species of a specific valve, or a dictionary of all recorded valves and species. Either the return value itself or the dictionary values will be strings that are padded with spaces. If a valve is not in use, its gas species will be recorded as ‘None ‘.

omfit_classes.utils_fusion.east_gas_injection(shot=None, valve=None, species=None, main_ion_species='D', server='EAST', tank_temperature=293.15, verbose=True, plot_total=False, plot_flow=False, tsmo=0.5, axs=None)[source]

Calculates EAST gas injection amounts based on tank pressure

Whereas DIII-D gas commands in V are often close to proportional to gas flows (because they’re really inputs to an external flow controller built into the valve assembly), EAST gas commands are much more raw. An EAST command of 5 V is NOT roughly double a command of 2.5 V, for example. 2.5 V might not even open the valve, and there’d be no way to be sure from just the command. There is no flow sensor inside the EAST injectors like there is at DIII-D (the source of the gasb_cal measurement instead of the gascgasb command). Lastly, the EAST reservoirs behind the injectors are small and so the flow rate vs. voltage is probably not constant. The “tank” here isn’t really a huge tank of gas, it’s a length of pipe between the big tank of gas and the injector itself. So, letting out enough gas to control a shot really can affect it significantly.

To get an actual measurement, we can turn to the pressure in the gas tank that feeds the injector and watch its decrease over time to get a flow rate. This script doesn’t calculate a flow rate, it provides the integral of the flow rate. Since the data are noisy, some smoothing is recommended before differentiation to find the flow rate; we leave the choice of smoothing strength to the end user.

Limitations:

1. COOLING: We assumed constant temperature so far, but the gas tank clearly is cooled by letting gas out because the pressure very slowly rebounds after the shot, presumably from the tank warming up again. Remember the “tank” is actually just a little length of pipe behind the injector. To see how much error this causes, just look at how much tank pressure rebounds after seeding stops. The time history usually extends long after the shot. It seems like a small enough error that it hasn’t been corrected yet.

2. NEEDS POST-PROCESSING: most users are probably interested in flow rates and will have to take derivatives of the outputs of this function to get them, including smoothing to defeat noise in the signals.

3. INCOMPLETE INFO: the electrons_per_molecule_puffed dictionary only lists a few species so far.

Parameters
  • shot – int EAST shot number, like 85293

  • valve – str Like OU1. Also accepts names like OUPEV1 and VOUPEV1; PEV and leading V will be removed. There are different naming conventions in different contexts and this function tries to parse them all.

  • species

    str Species in the tank, like Ne. Diluted species are accepted; for example, “50% Ne” will be split at % and give 0.5 of the molecules as Ne

    and 0.5 of the molecules as the main ion species (probably D2).

  • main_ion_species – str Species of main ions of the plasma, like D

  • server – str MDSplus server to use as the source of the data. Intended to allow choice between ‘EAST’ and ‘EAST_US’.

  • tank_temperature – float Temperature of the gas in the tank in K

  • verbose – bool Print information like gas totals

  • plot_total – bool Plot total seeding amount vs time

  • plot_flow – bool Plot flow rate (derivative of total seeding) vs time

  • tsmo – float Smoothing timescale in seconds, used in plots only. Very important for viewing flow rate.

  • axs – Axes instance or 1D array of Axes instances Axes used with plot_total or plot_flow. If none are provided, a new figure will be created. Number of Axes must be >= plot_flow+plot_total.

Returns

tuple of 1D arrays and a str 1D array: time in s 1D array: impurity electrons added 1D array: fuel electrons added 1D array: total molecules added str: primary species added

omfit_classes.utils_fusion.calc_h98_d3d(shot, calc_tau=False, calc_ptot=False, pinj_only=False, estimated_signals=None, data=None)[source]

H98 confinement quality calcuation valid for DIII-D

There are some other H98 calculations for DIII-D:

h_thh98y2 is calcualted by the Transport code and is often regarded as pretty accurate, but it will turn itself off at the slightest hint of impurity seeding. If there’s noise in one of the valves (GASB had an issue in 2021), this will be interpreted as seeding and the calculation will disable itself. So, it was pretty much always disabled because of this in years with noisy gas signals. The database that drives the 98y2 scaling goes up to Zeff=3, so a tiny bit of impurity seeding is tolerable. Even if Zeff is too high, it might still be interesting to see what H98 would be if it could be trustworthy.

h98y2_aot is calculated from Automatic OneTwo runs and can overestimate H98. The AOT calculation is usually availble and doesn’t fail because of impurity seeding, but it’s probably less accurate than this calculation.

Parameters
  • shot – int

  • calc_tau – bool Calculate tauth instead of gathering tauth. Might be useful if the confinement code stops early and doesn’t write tauth.

  • calc_ptot – bool Calculate total power instead of gathering it. Useful if ptot isn’t written due to an error in another code

  • pinj_only – bool Approximate ptot by pinj, ignoring ohmic and ECH power. This might be okay.

  • estimated_signals – dict Data for estimating missing signals. Should contain a ‘time’ key whose value is a 1d float array. Should contain one additional key for each pointname that is being estimated, whose value is a 1d float array with length matching that of time.

  • data – dict [optional] Pass in an empty dict and data used in the calculation will be written to it. Intermediate quantities may be captured for inspection in this way.

Returns

tuple containing two 1D float arrays Time in ms H98 as a function of time (unitless)

omfit_classes.utils_fusion.available_efits_from_guess(scratch_area, device, shot, default_snap_list=None, format='{tree}', **kw)[source]

Attempts to form a reasonable list of likely EFITs based on guesses

Parameters
  • scratch_area – dict Scratch area for storing results to reduce repeat calls. Mainly included to match call sigure of available_efits_from_rdb(), since OMFITmdsValue already has caching.

  • device – str Device name

  • shot – int Shot number

  • default_snap_list – dict [optional] Default set of EFIT treenames. Newly discovered ones will be added to the list.

  • **kw

    quietly accepts and ignores other keywords for compatibility with other similar functions

Returns

(dict, str) Dictionary keys will be descriptions of the EFITs

Dictionary values will be the formatted identifiers. For now, the only supported format is just the treename. If lookup fails, the dictionary will be {‘’: ‘’} or will only contain default results, if any.

String will contain information about the discovered EFITs

omfit_classes.utils_fusion.available_EFITs(scratch_area, device, shot, allow_rdb=True, allow_mds=True, allow_guess=True, **kw)[source]

Attempts to look up a list of available EFITs using various sources

Parameters
  • scratch_area – dict Scratch area for storing results to reduce repeat calls.

  • device – str Device name

  • shot – int Shot number

  • allow_rdb – bool Allow connection to DIII-D RDB to gather EFIT information (only applicable for select devices) (First choice for supported devices)

  • allow_mds – bool Allow connection to MDSplus to gather EFIT information (only applicable to select devices) (First choice for non-RDB devices, second choice for devices that normally support RDB)

  • allow_guess – bool Allow guesses based on common patterns of EFIT availability on specific devices (Last resort, only if other options fail)

  • **kw

    Keywords passed to specific functions. Can include:

    param default_snap_list

    dict [optional] Default set of EFIT treenames. Newly discovered ones will be added to the list.

    param format

    str Instructions for formatting data to make the EFIT tag name. Provided for compatibility with available_efits_from_rdb() because the only option is ‘{tree}’.

Returns

(dict, str) Dictionary keys will be descriptions of the EFITs

Dictionary values will be the formatted identifiers. If lookup fails, the dictionary will be {‘’: ‘’} or will only contain default results, if any.

String will contain information about the discovered EFITs

utils_plot

omfit_classes.utils_plot.autofmt_sharexy(trim_xlabel=True, trim_ylabel=True, fig=None)[source]
omfit_classes.utils_plot.autofmt_sharey(trim_xlabel=True, fig=None, wspace=0)[source]

Prunes y-tick labels and y-axis labels from all but the first cols axes and moves cols (optionally) closer together.

Parameters
  • trim_xlabel – bool. prune right ytick label to prevent overlap.

  • fig – Figure. Defaults to current figure.

  • wspace – Horizontal spacing between axes.

omfit_classes.utils_plot.autofmt_sharex(trim_ylabel=True, fig=None, hspace=0)[source]

Prunes x-tick labels and x-axis labels from all but the last row axes and moves rows (optionally) closer together.

Parameters
  • trim_ylabel – bool. prune top ytick label to prevent overlap.

  • fig – Figure. Defaults to current figure.

  • hspace – Vertical spacing between axes.

omfit_classes.utils_plot.uerrorbar(x, y, ax=None, **kwargs)[source]

Given arguments y or x,y where x and/or y have uncertainties, feed the appropriate terms to matplotlib’s errorbar function.

If y or x is more than 1D, it is flattened along every dimension but the last.

Parameters
  • x – array of independent axis values

  • y – array of values with uncertainties, for which shaded error band is plotted

  • ax – The axes instance into which to plot (default: pyplot.gca())

  • **kwargs – Passed to ax.errorbar

Returns

list. A list of ErrorbarContainer objects containing the line, bars, and caps of each (x,y) along the last dimension.

class omfit_classes.utils_plot.Uband(line, bands)[source]

Bases: object

This class wraps the line and PollyCollection(s) associated with a banded errorbar plot for use in the uband function.

Parameters
  • line – Line2D A line of the x,y nominal values

  • bands – list of PolyCollections The fill_between and/or fill_betweenx PollyCollections spanning the std_devs of the x,y data

omfit_classes.utils_plot.uband(x, y, ax=None, fill_kwargs=None, **kwargs)[source]

Given arguments x,y where either or both have uncertainties, plot x,y using pyplot.plot of the nominal values and surround it with with a shaded error band using matplotlib’s fill_between and/or fill_betweenx.

If y or x is more than 1D, it is flattened along every dimension but the last.

Parameters
  • x – array of independent axis values

  • y – array of values with uncertainties, for which shaded error band is plotted

  • ax – The axes instance into which to plot (default: pyplot.gca())

  • fill_kwargs – dict. Passed to pyplot.fill_between

  • **kwargs – Passed to pyplot.plot

Returns

list. A list of Uband objects containing the line and bands of each (x,y) along the last dimension.

omfit_classes.utils_plot.hardcopy(fn, bbox_inches='tight', fig=None, **keyw)[source]
omfit_classes.utils_plot.set_fontsize(fig=None, fontsize='+0')[source]

For each text object of a figure fig, set the font size to fontsize

Parameters
  • fig – matplotlib.figure object

  • fontsize – can be an absolute number (e.g 10) or a relative number (-2 or +2)

Returns

None

omfit_classes.utils_plot.user_lines_cmap_cycle()[source]

return colormap chosen by the user for representation of lines

omfit_classes.utils_plot.user_image_cmap_cycle()[source]

return colormap chosen by the user for representation of images

omfit_classes.utils_plot.color_cycle(n=10, k=None, cmap_name=None)[source]

Utility function to conveniently return the color of an index in a colormap cycle

Parameters
  • n – number of uniformly spaced colors, or array defining the colors’ spacings

  • k – index of the color (if None an array of colors of length n will be returned)

  • cmap_name – name of the colormap

Returns

color of index k from colormap cmap_name made of n colors, or array of colors of length n if k is None Note: if n is an array, then the associated ScalarMappable object is also returned (e.g. for use in a colorbar)

omfit_classes.utils_plot.cycle_cmap(length=50, cmap=None, start=None, stop=None, ax=None)[source]

Set default color cycle of matplotlib based on colormap

Note that the default color cycle is not changed if ax parameter is set; only the axes’s color cycle will be changed

Parameters
  • length – The number of colors in the cycle

  • cmap – Name of a matplotlib colormap

  • start – Limit colormap to this range (0 < start < stop 1)

  • stop – Limit colormap to this range (0 < start < stop 1)

  • ax – If ax is not None, then change the axes’s color cycle instead of the default color cycle

Returns

color_cycle

omfit_classes.utils_plot.contrasting_color(line_or_color)[source]

Given a matplotlib color specification or a line2D instance or a list with a line2D instance as the first element, pick and return a color that will contrast well. More complicated than just inversion as inverting blue gives yellow, which doesn’t display well on a white background.

Parameters

line_or_color – matplotlib color spec, line2D instance, or list w/ line2D instance as the first element

Returns

4 element array RGBA color specification for a contrasting color

omfit_classes.utils_plot.associated_color(line_or_color)[source]

Given a matplotlib color specification or a line2D instance or a list with a line2D instance as the first element, pick and return a color that will look thematically linked to the first color, but still distinguishable.

Parameters

line_or_color – matplotlib color spec, line2D instance, or list w/ line2D instance as the first element

Returns

4 element array RGBA color specification for a related, similar (but distinguishable) color

omfit_classes.utils_plot.blur_image(im, n, ny=None)[source]

blurs the image by convolving with a gaussian kernel of typical size n. The optional keyword argument ny allows for a different size in the y direction.

omfit_classes.utils_plot.pcolor2(*args, fast=False, **kwargs)[source]

Plots 2D data as a patch collection. Differently from matplotlib.pyplot.pcolor the mesh is extended by one element so that the number of tiles equals the number of data points in the Z matrix. The X,Y grid does not have to be rectangular.

Parameters
  • *args – Z or X,Y,Z data to be plotted

  • fast – bool Use pcolorfast instead of pcolor. Speed improvements may be dramatic. However, pcolorfast is marked as experimental and may produce unexpected behavior.

  • **kwargs – these arguments are passed to matplotlib.pyplot.pclor

Returns

None

omfit_classes.utils_plot.image(*args, **kwargs)[source]

Plots 2D data as an image.

Much faster than pcolor/pcolor2(fast=False), but the data have to be on a rectangular X,Y grid

Parameters
  • *args – Z or X,Y,Z data to be plotted

  • **kwargs – these arguments are passed to pcolorfast

omfit_classes.utils_plot.meshgrid_expand(xdata, ydata)[source]

returns the veritices of the mesh, if the xdata and ydata were the centers of the mesh xdata and ydata are 2D matrices, which could for example be generated by np.meshgrid

Parameters
  • xdata – center of the mesh

  • ydata – center of the mesh

Returns

omfit_classes.utils_plot.map_HBS_to_RGB(H, B, S=1.0, cmap=None)[source]

map to a RGB colormap separate HUE, BRIGHTNESS and SATURATIONS arrays

Parameters
  • H – HUE data (any shape array)

  • B – BRIGHTNESS data (any shape array)

  • S – SATURATION data (any shape array)

  • cmap – matplotlib.colormap to be used

Returns

RGB array (shape of input array with one more dimension of size 3 (RGB) )

omfit_classes.utils_plot.RGB_to_HEX(R, G, B)[source]

Convert color from numerical RGB to hexadecimal representation

Parameters
  • R – integer 0<x<255 or float 0.0<x<1.0

  • G – integer 0<x<255 or float 0.0<x<1.0

  • B – integer 0<x<255 or float 0.0<x<1.0

Returns

hexadecimal representation of the color

omfit_classes.utils_plot.plotc(*args, **kwargs)[source]

Plot the various curves defined by the arguments [X],Y,[Z] where X is the x value, Y is the y value, and Z is the color. If one argument is given it is interpreted as Y; if two, then X, Y; if three then X, Y, Z. If all three are given, then it is passed to plotc and the labels are discarded. If Z is omitted, a rainbow of colors is used, with blue for the first curve and red for the last curve. A different color map can be given with the cmap keyword (see http://wiki.scipy.org/Cookbook/Matplotlib/Show_colormaps for other options). If X is omitted, then the (ith) index of Y is used as the x value.

Parameters
  • *args

  • **kwargs

Returns

omfit_classes.utils_plot.title_inside(string, x=0.5, y=0.9, ax=None, **kwargs)[source]

Write the title of a figure inside the figure axis rather than outside

Parameters
  • string – title string

  • x – x location of the title string (default 0.5, that is centered)

  • y – y location of the title string (default 0.875)

  • ax – axes to operate on

  • **kwargs – additional keywords passed to pyplot.title

Returns

pyplot.title text object

omfit_classes.utils_plot.increase_resolution(*args, **kwargs)[source]

This function takes 1 (Z) or 3 (X,Y,Z) 2D tables and interpolates them to higher resolution by bivariate spline interpolation. If 1 (3) table(s) is(are) provided, then the second(fourth) argument is the resolution increase, which can be a positive or negative integer integer: res=res0*2^n or a float which sets the grid size in the units provided by the X and Y tables

class omfit_classes.utils_plot.infoScatter(x, y, annotes, axis=None, tol=5, func=None, all_on=False, suppress_canvas_draw=False, **kw)[source]

Bases: object

improved version of: http://wiki.scipy.org/Cookbook/Matplotlib/Interactive_Plotting

Callback for matplotlib to display an annotation when points are clicked on

Parameters
  • x – x of the annotations

  • y – y of the annotations

  • annotes – list of string annotations

  • axis – axis on which to operate on (default to current axis)

  • tol – vicinity in pixels where to look for annotations

  • func – function to call with signature: func(x,y,annote,visible,axis)

  • all_on – Make all of the text visible to begin with

  • suppress_canvas_draw – Do not actively draw the canvas if all_on is True, makes plotting faster there are many subplots

  • **kw – extra keywords passed to matplotlib text class

drawAnnote(axis, x, y, annote, redraw_canvas=True)[source]

Draw the annotation on the plot

drawSpecificAnnote(annote)[source]
omfit_classes.utils_plot.infoPoint(fig=None)[source]

print x,y coordinates where the user clicks

Parameters

fig – matplotlib figure

omfit_classes.utils_plot.XKCDify(ax, mag=1.0, f1=50, f2=0.01, f3=15, bgcolor='w', xaxis_loc=None, yaxis_loc=None, xaxis_arrow='+', yaxis_arrow='+', ax_extend=0.1, expand_axes=False, ylabel_rot=78)[source]

XKCD plot generator by, Jake Vanderplas; Modified by Sterling Smith

This is a script that will take any matplotlib line diagram, and convert it to an XKCD-style plot. It will work for plots with line & text elements, including axes labels and titles (but not axes tick labels).

The idea for this comes from work by Damon McDougall

http://www.mail-archive.com/matplotlib-users@lists.sourceforge.net/msg25499.html

This adjusts all lines, text, legends, and axes in the figure to look like xkcd plots. Other plot elements are not modified.

Parameters
  • ax – Axes instance the axes to be modified.

  • mag – float the magnitude of the distortion

  • f2, f3 (f1,) – int, float, int filtering parameters. f1 gives the size of the window, f2 gives the high-frequency cutoff, f3 gives the size of the filter

  • yaxis_log (xaxis_loc,) – float The locations to draw the x and y axes. If not specified, they will be drawn from the bottom left of the plot

  • xaxis_arrow – str where to draw arrows on the x axes. Options are ‘+’, ‘-‘, ‘+-‘, or ‘’

  • yaxis_arrow – str where to draw arrows on the y axes. Options are ‘+’, ‘-‘, ‘+-‘, or ‘’

  • ax_extend – float How far (fractionally) to extend the drawn axes beyond the original axes limits

  • expand_axes – bool if True, then expand axes to fill the figure (useful if there is only a single axes in the figure)

  • ylabel_rot – float number of degrees to rotate the y axis label

omfit_classes.utils_plot.autoscale_y(ax, margin=0.1)[source]

Rescales the y-axis based on the data that is visible given the current xlim of the axis. Created by eldond at 2017 Mar 23 20:26

This function was taken from an answer by DanKickstein on stackoverflow.com http://stackoverflow.com/questions/29461608/matplotlib-fixing-x-axis-scale-and-autoscale-y-axis http://stackoverflow.com/users/1078391/danhickstein

I don’t think this function considers shaded bands such as would be used to display error bars. Increasing the margin may be a good idea when dealing with such plots.

Parameters
  • ax – a matplotlib axes object

  • margin – The fraction of the total height of the y-data to pad the upper and lower ylims

omfit_classes.utils_plot.set_linearray(lines, values=None, cmap='viridis', vmin=None, vmax=None)[source]

Set colors of lines to colormapping of values.

Other good sequential colormaps are YlOrBr and autumn. A good diverging colormap is bwr.

Parameters
  • lines (list) – Lines to set colors.

  • values (array like) – Values corresponding to each line. Default is indexing.

  • cmap (str) – Valid matplotlib colormap name.

  • vmax (float) – Upper bound of colormapping.

  • vmin (float) – Lower bound of colormapping.

Returns

ScalarMappable. A mapping object used for colorbars.

omfit_classes.utils_plot.pi_multiple(x, pos=None)[source]

Provides a string representation of x that is a multiple of the fraction pi/’denominator’.

See multiple_formatter documentation for more info.

omfit_classes.utils_plot.multiple_formatter(denominator=24, number=3.141592653589793, latex='\\pi')[source]

Returns a tick formatting function that creates tick labels in multiples of ‘number’/’denominator’.

Code from https://stackoverflow.com/a/53586826/6605826

Parameters
  • denominator – The denominator of the fraction that tick labels are created in multiples of.

  • number – The numerator of the fraction that tick labels are created in multiples of

  • latex – The latex string used to represent ‘number’

omfit_classes.utils_plot.convert_ticks_to_pi_multiple(axis=None, major=2, minor=4)[source]

Given an axis object, force its ticks to be at multiples of pi, with the labels formatted nicely […,-2pi,-pi,0,pi,2pi,…]

Parameters
  • axis – An axis object, such as pyplot.gca().xaxis

  • major – int Denominator of pi for major tick marks. 2: major ticks at 0, pi/2., pi, … Can’t be greater than 24.

  • minor – int Denominator of pi for minor tick marks. 4: minor ticks at 0, pi/4., pi/2., …

Returns

None

omfit_classes.utils_plot.is_colorbar(ax)[source]

Guesses whether a set of Axes is home to a colorbar

https://stackoverflow.com/a/53568035/6605826

Parameters

ax – Axes instance

Returns

bool True if the x xor y axis satisfies all of the following and thus looks like it’s probably a colorbar: No ticks, no tick labels, no axis label, and range is (0, 1)

omfit_classes.utils_plot.tag_plots_abc(fig=None, axes=None, corner=[1, 1], font_size='medium', skip_suspected_colorbars=True, start_at=0, **annotate_kw)[source]

Tag plots with (a), (b), (c), …

Parameters
  • fig – Specify a figure instance instead of letting the function pick the most recent one

  • axes – Specify a plot axes instance or list/array of plot axes instances instead of letting the function use fig.get_axes()

  • corner – Which corner does the tag go in? [0, 0] for bottom left, [1, 0] for bottom right, etc.

  • font_size – Font size of the annotation.

  • skip_suspected_colorbars – bool Try to detect axes which are home to colorbars and skip tagging them. An Axes instance is suspected of having a colorbar if either the xaxis or yaxis satisfies all of these conditions: - Length of tick list is 0 - Length of tick label list is 0 - Length of axis label is 0 - Axis range is (0,1)

  • start_at – int Offset value for skipping some numbers. Useful if you aren’t doing real subfigs, but two separate plots and placing them next to each other in a publication. Set to 1 to start at (b) instead of (a), for example.

  • annotate_kw – dict Additional keywords passed to annotate(). Keywords used by settings such as corner, etc. will be overriden.

omfit_classes.utils_plot.mark_as_interactive(ax, interactive=True)[source]

Mark an axis as interactive or not

Parameters
  • ax – axis

  • interactive – boolean

Returns

axis

class omfit_classes.utils_plot.View1d(data, coords=None, dims=None, name=None, dim=None, axes=None, dynamic_ylim=False, use_uband=False, cornernote_options=None, plot_options=None, **indexers)[source]

Bases: object

Plot 2D or 3D data as line-plots with interactive navigation through the alternate dimensions. Navigation uses the 4 arrow keys to traverse up to 2 alternate dimensions.

The data must be on a regular grid, and is formed into a xarray DataArray if not already.

Uses matplotlib line plot for float/int data, OMFIT uerrrorbar for uncertainty variables.

Examples:

The view1d can be used to interactively explore data. For usual arrays it draws line slices.

>> t = np.arange(20) >> s = np.linspace(0,2*np.pi,60) >> y = np.sin(np.atleast_2d(s).T+np.atleast_2d(t)) >> da = xarray.DataArray(y,coords=SortedDict([(‘space’,s),(‘time’,t)]),name=’sine’) >> v = View1d(da.transpose(‘time’,’space’),dim=’space’,time=10)

For uncertainties arrays, it draws errorbars using the uerrorbar function. Multiple views with the same dimensions can be linked for increased speed (eliminate redundant calls to redraw).

>> y_u = unumpy.uarray(y+(random(y.shape)-0.5),random(y.shape)) >> da_u = xarray.DataArray(y_u,coords=SortedDict([(‘space’,s),(‘time’,t)]),name=’measured’) >> v_u = View1d(da_u,dim=’space’,time=10,axes=pyplot.gca()) >> v.link(v_u) # v will remain connected to keypress events and drive vu

Variable dependent axis data can be viewed if x and y share a regular grid in some coordinates,

>> x = np.array([s+(random(s.shape)-0.5)*0.2 for i in t]).T >> da_x = xarray.DataArray(x,coords=SortedDict([(‘space’,s),(‘time’,t)]),name=’varspace’) >> ds = da_u.to_dataset().merge(da_x.to_dataset()) >> v_x = View1d(ds,name=’measured’,dim=’varspace’,time=10,axes=pyplot.gca()) >> v.link(v_x)

Parameters
  • data – DataArray or array-like 2D or 3D data values to be viewed.

  • coords – dict-like Dictionary of Coordinate objects that label values along each dimension.

  • dims – tuple Dimension names associated with this array.

  • name – string Label used in legend. Empty or beginning with ‘_’ produces no legend label. If the data is a DataArray it will be renamed before plotting. If the data is a Dataset, the name specifies which of its existing data_vars to plot.

  • dim – string, DataArray Dimension plotted on x-axis. If DataArray, must have same dims as data.

  • axes – Axes instance The axes plotting is done in.

  • dynamic_ylim – bool Re-scale y limits of axes when new slices are plotted.

  • use_uband – bool Use uband instead of uerrorbar to plot uncertainties variables.

  • cornernote_options – dict Key word arguments passed to cornernote (such as root, shot, device). If this is present, then cornernote will be updated with the new time if there is only one time to show, or the time will be erased from the cornernote if more than one time is shown by this View1d instance (such as by freezing one slice).

  • plot_options – dict Key word arguments passed to plot/uerrorbar/uband.

  • **indexers – dict Dictionary with keys given by dimension names and values given by arrays of coordinate index values. Must include all dimensions other, than the fundamental.

use_uband(use=True)[source]

Toggle use of uband instead of uerrorbar for plotting function

set_label_indexers_visible(visible=True)[source]

Include the current indexers in line labels.

key_command(event, draw=True, **plot_options)[source]

Use arrows to navigate up to 2 extra dimensions by incrementing the slice indexes.

Use w/e to write/erase slices to persist outside on navigation.

isel(draw=True, **indexers)[source]

Re-slice the data along its extra dimensions using indexes.

Link all actions in this view to the given view.

Parameters
  • view – View1d. A view that will be driven by this View’s key press responses.

  • disconnect – bool. Disconnect key press events from driven view.

Unlink actions in this view from controlling the given view.

class omfit_classes.utils_plot.View2d(data, coords=None, dims=None, name=None, axes=None, quiet=False, use_uband=False, contour_levels=0, imag_options={}, plot_options={'ls': '-', 'marker': ''}, **indexers)[source]

Bases: object

Plot 2D data with interactive slice viewers attached to the 2D Axes. Left clicking on the 2D plot refreshes the line plot slices, right clicking overplots new slices.

The original design of this viewer was for data on a rectangular grid, for which x and y are 1D arrays defining the axes but may be irregularly spaced. In this case, the line plot points correspond to the given data. If x or y is a 2D array, the data is assumed irregular and interpolated to a regular grid using scipy.interpolate.griddata.

Example:

Explore a basic 2D np array without labels,

>> x = np.linspace(-1, 1, 200) >> y = np.linspace(-2, 2, 200) >> xx, yy = meshgrid(x, y) >> z = np.exp(-xx**2 - yy**2) >> v = View2d(z)

To add more meaningful labels to the axes and data do,

>> v = View2d(z, coords={‘x’:x, ‘y’:y}, dims=(‘x’, ‘y’), name=’wow’)

or use a DataArray,

>> d = DataArray(z, coords={‘x’:x, ‘y’:y}, dims=(‘x’, ‘y’), name=’wow’) >> v = View2d(d)

Note that the coordinates should be 1D. Initializing a view with regular grid, 2D coordinates will result in an attempt to slice them appropriately. This is done for consistency with some matplotlib 2D plotting routines, but is not recommended.

>> v = View2d(z, coords=dict(x=x, y=y), dims=(‘x’, ‘y’))

If you have irregularly distributed 2D data, it is recomended that you first interpolate it to a 2D grid in whatever way is most applicable. If you do not, initializing a view will result in an attempt to linearly interpolate to a automatically chosen grid.

>> x = np.random.rand(1000) >> y = np.random.rand(1000) * 2 >> z = np.exp(-x**2 - y**2) >> v = View2d(z, coords=dict(x=x, y=y), dims=(‘x’, ‘y’))

The same applies for 2D collections of irregular points and values.

>> x = x.reshape((50, 20)) >> y = y.reshape((50, 20)) >> z = z.reshape((50, 20)) >> v = View2d(z, coords=[(‘x’, x), (‘y’, y)], dims=(‘x’, ‘y’))

Parameters
  • data – DataArray or array-like 2D or 3D data values to be viewed.

  • coords – dict-like Dictionary of Coordinate objects that label values along each dimension.

  • dims – tuple Dimension names associated with this array.

  • name – string Label used in legend. Empty or begining with ‘_’ produces no legend label.

  • dim – string, DataArray Dimension plotted on x-axis. If DataArray, must have same dims as data.

  • axes – Axes instance The axes plotting is done in.

  • quiet – bool Suppress printed messages.

  • use_uband – bool Use uband for 1D slice plots instead of uerrorbar.

  • contour_levels – int or np.ndarray Number of or specific levels used to draw black contour lines over the 2D image.

  • imag_options – dict Key word arguments passed to the DataArray plot method (modified pcolormesh).

  • plot_options – dict Key word arguments passed to plot or uerrorbar. Color will be determined by cmap variable.

  • **indexers – dict Dictionary with keys given by dimension names and values given by arrays of coordinate index values.

use_uband(use=True)[source]

Toggle use of uband instead of uerrorbar in 1D slice plots

key_navigate_cuts(key_press)[source]
set_data(data=None, **kw)[source]

Set base data and axes, as well as displayed values. Call with no arguments re-sets to interactively displayed values (use after der(), or int()).

Parameters

data (DataArray) – 2D array of data.

Returns

The main 2d matplotlib.collections.QuadMesh from pcolormesh

der(axis=0)[source]

Set 2D values to derivative of data along specified axis.

Parameters

axis (int) – Axis along which derivative is taken (0=horizontal, 1=vertical).

Returns

The main 2d matplotlib.collections.QuadMesh from pcolormesh

int(axis=0)[source]

Set 2D values to derivative of data along specified axis.

Parameters

axis (int) – Axis along which integration is taken (0=horizontal, 1=vertical).

Returns

The main 2d matplotlib.collections.QuadMesh from pcolormesh

toggle_log()[source]

Toggle log/linear scaling of data in 2D plot.

vslice(xmin, xmax, std=False, draw=True, force=False, **kw)[source]

Plot line collection of x slices.

Parameters
  • xmin (float) – Lower bound of slices dispayed in line plot.

  • xmax (float) – Upper bound of slices dispayed in line plot.

  • std (bool) – Display mean and standard deviation instead of individual slices.

  • draw (bool) – Redraw the figure canvas

  • Force – Re-slice even if arguments are identical to last slice.

  • **kw – Extra key words are passed to the 1D plot() function

Returns

Possibly modified (xmin,xmax,std)

Return type

tuple

hslice(ymin, ymax, std=False, draw=True, force=False, **kw)[source]

Plot line collection of y slices.

Parameters
  • ymin (float) – Lower bound of slices dispayed in line plot.

  • ymax (float) – Upper bound of slices dispayed in line plot.

  • std (bool) – Display mean and standard deviation instead of individual slices.

  • draw (bool) – Redraw the figure canvas

  • Force – Re-slice even if arguments are identical to last slice.

  • **kw – Extra key words are passed to the 1D plot() function

Returns

Possibly modified (ymin,ymax,std)

Return type

tuple

line_select_callback(eclick, erelease=None)[source]

Call vslice and hslice for the range of x and y spanned by the rectangle between mouse press and release. Called by RectangleSelector.

Parameters
  • eclick (matplotlib Event) – Matplotlib mouse click.

  • erelease (matplotlib Event) – Matplotlib mouse release.

Returns

None

toggle_selector(event)[source]

Connected to key press events to turn on (a) or off (q) selector.

Parameters

event (matplotlib event) – key press event.

Returns

None

Link all actions in this view to the given view.

Unlink actions in this view from controlling the given view.

class omfit_classes.utils_plot.View3d(data, coords=None, dims=None, name=None, axes=None, use_uband=False, quiet=False, contour_levels=0, imag_options={}, plot_options={'ls': '-', 'marker': ''}, **indexers)[source]

Bases: omfit_classes.utils_plot.View2d

View 3D data by scrolling through 3rd dimension in a View2d plot.

Parameters
  • data – DataArray or array-like 2D or 3D data values to be viewed.

  • coords – dict-like Dictionary of Coordinate objects that label values along each dimension.

  • dims – tuple Dimension names associated with this array.

  • name – string Label used in legend. Empty or begining with ‘_’ produces no legend label.

  • dim – string, DataArray Dimension plotted on x-axis. If DataArray, must have same dims as data.

  • axes – Axes instance The axes plotting is done in.

  • quiet – bool Suppress printed messages.

  • use_uband – bool Use uband for 1D slice plots instead of uerrorbar.

  • contour_levels – int or np.ndarray Number of or specific levels used to draw black contour lines over the 2D image.

  • imag_options – dict Key word arguments passed to the DataArray plot method (modified pcolormesh).

  • plot_options – dict Key word arguments passed to plot or uerrorbar/uband. Color will be determined by cmap variable.

  • **indexers – dict Dictionary with keys given by dimension names and values given by arrays of coordinate index values.

set_3d(x2, index=False, draw=True)[source]

Set third dimension of view to value nearest to slice.

Parameters
  • x2 (float or int) – Slice in third dimension. Type depends on value of index.

  • index (bool) – Set True if x2 is the integer index

  • draw (bool) – If True, redraw the canvas

Returns

None

class omfit_classes.utils_plot.DragPoints(yArray, xArray=None, eyArray=None, exArray=None, editY=True, editX=True, exDelta=1, eyDelta=1, sorted=False, continuous=True, showOriginal=False, func=None, fargs=[], cyArray=None, cxArray=None, onMove=None, resize=False, show_hints=True, ax=None, **kwargs)[source]

Bases: object

This class is used to define matplotlib draggable arrays

Parameters
  • yArray – location in the OMFIT tree for the y array

  • xArray – location in the OMFIT tree for the x array

  • eyArray – location in the OMFIT tree for the x error array

  • exArray – location in the OMFIT tree for the y error array

  • editY – allow dragging of points along Y axis

  • editX – allow dragging of points along X axis

  • exDelta – increments of exArray

  • eyDelta – increments of eyArray

  • editX – allow dragging of points along X axis

  • sorted – keep points sorted in x

  • continuous – draw continuously even while user is dragging (may be good to disable if takes long time)

  • showOriginal – show original points

  • func

    a function with signature like func(x,y,motion,fargs=[]). where: * x : x coordinate of control points

    • y : y coordinate of control points

    • motion : True if the user has dragged a point

    This function must return x_, y_, x, y where:

    • x_ : interpolating points between x coordinate of control points

    • y_ : interpolating points between y coordinate of control points

    • x : x coordinate of control points

    • y : y coordinate of control points

  • fargs – arguments to the function

  • cyArray – location in the OMFIT tree for the y interpolation array

  • cxArray – location in the OMFIT tree for the x interpolation array

  • onMove – function handle to call whenever control point is moved

  • resize – boolean to whether update axes size when drag point gets to an edge for the figure

  • show_hints – bool. Show a cornernote with tips for interaction.

  • ax – Axes. Axes in which to plot.

All other key word arguments passed to the matplotlib plot function.

Returns

epsilon = 5
draw_callback(event=None)[source]
get_ind_under_point(event)[source]
button_press_callback(event)[source]
key_press_callback(event=None)[source]
button_release_callback(event)[source]
motion_notify_callback(event)[source]
omfit_classes.utils_plot.editProfile(yi, xi=None, n=None, showOriginal=True, func='spline', onMove=None)[source]

This function opens an interactive figure for convenient editing of profiles via spline

Parameters
  • yi – string whose eval yields the y data

  • xi – string whose eval yields the x data

  • n – number of control points

  • showOriginal – plot original data

  • func – interpolation function used to interpolate between control points ‘linear’, ‘spline’, ‘pchip’, ‘circular’

  • onMove – function to call when moving of control points occurs

  • resize – boolean to whether update axes size when drag point gets to an edge for the figure

Returns

DragPoints object

omfit_classes.utils_plot.cornernote(text='', root=None, device=None, shot=None, time=None, ax=None, fontsize='small', clean=True, remove=False, remove_specific=False)[source]

Write text at the bottom right corner of a figure

Parameters
  • text – text to appear in the bottom left corner

  • root

    • if ‘’ append nothing

    • if None append shot/time as from OMFIT[‘MainSettings’][‘EXPERIMENT’]

    • if OMFITmodule append shot/time as from root[‘SETTINGS’][‘EXPERIMENT’]

  • device – override device string (does not print device at all if empty string)

  • shot – override shot string (does not print shot at all if empty string)

  • time – override time string (does not print time at all if empty string)

  • ax – axis to plot on

  • fontsize – str or float. Sets font size of the Axes annotate method.

  • clean – delete existing cornernote(s) from current axes before drawing a new cornernote

  • remove – delete existing cornernote(s) and return before drawing any new ones

  • remove_specific – delete existing cornernote(s) from current axes only if text matches the text that would be printed by the current call to cornernote() (such as identical shot, time, etc.)

Returns

Matplotlib annotate object

class omfit_classes.utils_plot.axdline(slope=1, intercept=0, *args, **kwargs)[source]

Bases: matplotlib.lines.Line2D

Draw a line based on its slope and y-intercept. Additional arguments are passed to the <matplotlib.lines.Line2D> constructor.

From stackoverflow anser by ali_m: http://stackoverflow.com/a/14348481/6605826 Originally named ABLine2D

Create a .Line2D instance with x and y data in sequences of xdata, ydata.

Additional keyword arguments are .Line2D properties:

Properties:

agg_filter: a filter function, which takes a (m, n, 3) float array and a dpi value, and returns a (m, n, 3) array alpha: float or None animated: bool antialiased or aa: bool clip_box: .Bbox clip_on: bool clip_path: Patch or (Path, Transform) or None color or c: color contains: unknown dash_capstyle: {‘butt’, ‘round’, ‘projecting’} dash_joinstyle: {‘miter’, ‘round’, ‘bevel’} dashes: sequence of floats (on/off ink in points) or (None, None) data: (2, N) array or two 1D arrays drawstyle or ds: {‘default’, ‘steps’, ‘steps-pre’, ‘steps-mid’, ‘steps-post’}, default: ‘default’ figure: .Figure fillstyle: {‘full’, ‘left’, ‘right’, ‘bottom’, ‘top’, ‘none’} gid: str in_layout: bool label: object linestyle or ls: {‘-‘, ‘–’, ‘-.’, ‘:’, ‘’, (offset, on-off-seq), …} linewidth or lw: float marker: marker style string, ~.path.Path or ~.markers.MarkerStyle markeredgecolor or mec: color markeredgewidth or mew: float markerfacecolor or mfc: color markerfacecoloralt or mfcalt: color markersize or ms: float markevery: None or int or (int, int) or slice or List[int] or float or (float, float) or List[bool] path_effects: .AbstractPathEffect picker: unknown pickradius: float rasterized: bool or None sketch_params: (scale: float, length: float, randomness: float) snap: bool or None solid_capstyle: {‘butt’, ‘round’, ‘projecting’} solid_joinstyle: {‘miter’, ‘round’, ‘bevel’} transform: matplotlib.transforms.Transform url: str visible: bool xdata: 1D array ydata: 1D array zorder: float

See set_linestyle() for a description of the line styles, set_marker() for a description of the markers, and set_drawstyle() for a description of the draw styles.

omfit_classes.utils_plot.square_subplots(nplots, ncol_max=inf, flip=False, sparse_column=True, just_numbers=False, identify=False, fig=None, **kw)[source]

Creates a set of subplots in an approximate square, with a few empty subplots if needed

Parameters
  • nplots – int Number of subplots desired

  • ncol_max – int Maximum number of columns to allow

  • flip – bool True: Puts row 0 at the bottom, so every plot on the bottom row can accept an X axis label False: Normal plot numbering with row 0 at the top. The bottom row may be sparsely populated.

  • sparse_column

    bool Controls the arrangement of empty subplots. True: the last column is sparse. That is, all the empty plots will be in the last column. There will be at most

    one plot missing from the last row, and potentially several from the last column. The advantage is this provides plenty of X axes on the bottom row to accept labels. To get natural numbering of flattened subplots, transpose before flattening: axs.T.flatten(), or just use the 1D axsf array that’s returned.

    False: the last row is sparse. All the empty plots will be in the last row. The last column will be missing at

    most one plot, but the last row may be missing several. This arrangement goes more smoothly with the numbering of axes after flattening.

  • just_numbers – bool Don’t create any axes, but instead just return the number of rows, columns, and empty subplots in the array.

  • identify – bool For debugging: write the number (as flattened) and [row, col] coordinates of each subplot on the plot itself. These go in the center, in black. In the top left corner in red is the naive flattened count, which will appear on empty plots as well to show how wrong it is. In the bottom right corner in blue is the proper flattened count based on axsf.

  • fig – Figure instance [optional]

  • **kw – keywords passed to pyplot.subplots when creating axes (like sharex, etc.)

Returns

(axs, axsf) or (nr, nc, on, empty) axs: 2d array of Axes instances. It is flipped vertically relative to normal axes output by pyplot.subplots,

so the 0th row is the bottom. This is so the bottom row will be populated and can receive x axis labels.

axsf: 1d array of Axes instances, leaving out the empty ones (they might not be in order nicely) empty: int: number of empty cells in axs.

The first empty, if there is one, is [-1, -1] (top right), then [-1, -2] (top row, 2nd from the right), etc.

nr: int: number of rows nc: int: number of columns on: 2d bool array: flags indicating which axes should be on (True) and which should be hidden/off (False)

utils_fit

omfit_classes.utils_fit.autoknot(x, y, x0, evaluate=False, minDist=None, minKnotSpacing=0, s=3, w=None, allKnots=False, userFunc=None, *args, **kwargs)[source]

This function returns the optimal location of the inner-knots for a nth degree spline interpolation of y=f(x)

Parameters
  • x – input x array

  • y – input y array

  • x0 – initial knots distribution (list) or number of knots (integer)

  • s – order of the spline

  • w – input weights array

  • allKnots – returns all knots or only the central ones exluding the extremes

  • userFunc – autoknot with user defined function with signature y0=userFunc(x,y)(x0)

  • minDist – a number between >0 and infinity (though usually <1), which sets the minimum distance between knots. If small knots will be allowed to be close to one another, if large knots will be equispaced. Use None to automatically determine this parameter based on: 0.01*len(knots) If minDist is a string than it will be evaluated (the knots locations in the string can be accessed as knots).

  • minKnotSpacing – a number in x input units that denotes the minimal inter-knot space that autoknot should aim for. It shows up as an addition term in the cost function and is heavily punished against. If x is too large it will force autoknow to output evenly spaced knots. Default to 0. (ie no limit)

Returns

x0 optimal location of inner knots to spline interpolate y=f(x) f1=interpolate.LSQUnivariateSpline(x,y,x0,k=s,w=w)

class omfit_classes.utils_fit.knotted_fit_base(x, y, yerr)[source]

Bases: object

The base class for the types of fits that have free knots and locations

Does basic checking for x,y,yerr then stores them in self.x, self.y, self.yerr such that x is monotonically increasing

kw_vars = ['min_slope', 'monotonic', 'min_dist', 'first_knot', 'knots', 'fixed_knots', 'fit_SOL', 'outliers']
kw_defaults = [None, False, 0, None, 3, False, False, 3]
fit_knot_range(knots_range, **kw)[source]

Try a range of number of knots Stop the loop if the current number of knots did no better than the previous best

Parameters
  • knots_range – A tuple, (min_num_knots, max_num_knots) passed to range

  • **kw – The keywords passed to fit_single

get_tag(**kw)[source]

Get the tag for the settings given by **kw

Parameters

**kw – The Fit Keywords documented in the __init__ method

fit_single(**keyw)[source]

Perform a single fit for the given **keyw

Parameters

**keyw – The Fit Keywords documented in the __init__ method

Returns

The lmfit.MinimizerResult instance for the fit

The fit is also stored in self.fits[self.get_tag(**keyw)]

restore_orig_data()[source]

Restore the original data

get_zk(params)[source]

Return the array of knot values from the lmfit.Parameters object

get_xk(params)[source]

Return the array of knot locations from the lmfit.Parameters object

valid_xk(params)[source]
Parameters

params – An lmfit.Paramaters object to check

Returns

True if np.all(min_dist< xk_i-xk_(i-1)))

residual(params)[source]

Return the array of np.sqrt(((ymodel-y)/yerr)**2) given the lmfit.Parameters instance

Developers note: For the default data set tested, choosing the sqrt of the square did better than the signed residual

get_param_unc(lmfit_out)[source]

Get parameters with correlated uncertainties

Parameters

lmfit_outlmfit.MinimizerResult instance

Returns

tuple of xk, zk, ybc, xbc

plot(**kw)[source]

Plot all fits calculated so far, each in its own tab of a FigureNotebook, where the tab is labeled by the shortened tag of the tag of the fit

Parameters

**kw – Dictionary passed to self.plot_individual_fit

Returns

The FigureNotebook instance created

short_tag(tag)[source]

Return a shortened version of the tag

plot_individual_fit(tag, ax=None, x=array([0.0, 0.0011, 0.0022, ..., 1.0978, 1.0989, 1.1]))[source]

Plot a single fit, characterized by tag

Parameters
  • tag – The tag of the fit that is to be plotted, must be in self.fits.keys()

  • ax – The axes to plot into (one is created if None)

  • x – The x values to use for plotting the fitted curve

has_large_xerrorbars(lmfit_out)[source]
Parameters

lmfit_out – A lmfit.MinimizerResult object

Returns

True if the errorbars of the knot locations are larger than the distance between knots

has_large_errorbars(lmfit_out, verbose=False)[source]
Parameters

lmfit_out – A lmfit.MinimizerResult object

Returns

True if any of the following are False:

  1. the errorbars of the knot locations are smaller than the distance between the knots

  2. the errorbars in the fit at the data locations is not larger than the range in data

  3. the errorbars in the fit at the data locations is not larger than the fit value at that location, if the data are all of one sign

get_best_fit(verbose=False, allow_no_errorbar=None)[source]

Figure out which is the best fit so far

The best fit is characterized as being the fit with the lowest reduced chi^2 that is valid. The definition of valid is

  1. the knots are in order

  2. the knots are at least min_dist apart

  3. the errorbars on the fit parameters were able to be determined

  4. the errorbars of the knot locations are smaller than the distance between the knots

  5. the errorbars in the fit at the data locations is not larger than the range in data

  6. the errorbars in the fit at the data locations is not larger than the fit value at that location

Parameters
  • verbose – If True, print the tag and reduced chi2 of all fits

  • allow_no_errorbar – If True, if there is no valid fit found with errorbars, return the best fit without errorbars

Returns

A tuple of (best_tag, best_reduced_chi2)

plot_best_fit(**kw)[source]

A convenience function for plotting the best fit

class omfit_classes.utils_fit.fitSL(x, y, yerr, knots=3, min_dist=0, first_knot=None, fixed_knots=False, fit_SOL=False, monotonic=False, min_slope=None, outliers=3, plot_best=False, allow_no_errorbar=False)[source]

Bases: omfit_classes.utils_fit.knotted_fit_base

Fit a profile of data using integrated scale lengths, ideally obtaining uncertainties in the fitting parameters.

Due to the nature of integrating scale lengths, this fitter is only good for data > 0 or data < 0.

Examples

>>> pkl = OMFITpickle(OMFITsrc+'/../samples/data_pickled.pkl')
>>> x = pkl['x'][0,:]
>>> y = pkl['y'][0,:]
>>> yerr = pkl['e'][0,:]
>>> fit = fitSL(x,y,yerr,fixed_knots=True,knots=-7,plot_best=True)

Along the way of obtaining the fit with the desired parameters, other intermediate fits may be obtained. These are stored in the fits attribute (a dict), the value of whose keys provide an indication of how the fit was obtained, relative to the starting fit. For instance, to provide a variable knot fit, a fixed knot (equally spaced) fit is performed first. Also an initial fit is necessary to know if there are any outliers, and then the outliers can be detected. The get_best_fit method is useful for determing which of all of the fits is the best, meaning the valid fit with the lowest reduced chi^2. Here valid means

  1. the knots are in order

  2. the knots are at least min_dist apart

  3. the errorbars on the fit parameters were able to be determined

  4. the errorbars of the knot locations are smaller than the distance between the knots

Note that 1) and 2) should be satisfied by using lmfit Parameter constraints, but it doesn’t hurt to double check :-)

Developer note: If the fitter is always failing to find the errorbars due to tolerance problems, there are some tolerance keywords that can be passed to lmfit.minimize: xtol, ftol, gtol that could be exposed.

Initialize the fitSL object, including calculating the first fit(s)

Parameters
  • x – The x values of the data

  • y – The values of the data

  • yerr – The errors of the data

Fit Keywords:

Parameters
  • knots

    • Positive integer: Use this number of knots as default (>=3)

    • Negative integer: Invoke the fit_knot_range method for the range (3,abs(knots))

    • list-like: Use this list as the starting point for the knot locations

  • min_dist – The minimum distance between knot locations * min_dist > 0 (faster) enforced by construction * min_dist < 0 (much slower) enforced with lmfit

  • first_knot – The first knot can be constrained to be above first_knot. The default is above min(x)+min_dist (The zeroth knot is at 0.)

  • fixed_knots – If True, do not allow the knot locations to change

  • fit_SOL – If True, include data points with x>1

  • monotonic – If True, only allow positive scale lengths

  • min_slope – Constrain the scale lengths to be above min_slope

  • outliers – Do an initial fit, then throw out any points a factor of outliers standard deviations away from the fit

Convenience Keywords:

Parameters
  • plot_best – Plot the best fit

  • allow_no_errorbar – If True, get_best_fit will return the best fit without errorbars if no valid fit with errorbars exists

build_params(**keyw)[source]

Build the lmfit.Parameters object needed for the fitting

Parameters

**keyw – The Fit Keywords documented in the __init__ method

Returns

The lmfit.Parameters translation of the settings given by **keyw

model(params, x, lmfit_out=None)[source]

Return the model integrated scale length curve at x

Parameters
  • params – The lmfit.Parameters object

  • x – evaluate model at x

  • lmfit_outlmfit.MinimizerResult instance to use for getting uncertainties in the curve

plot(showZ=True, x=array([0.0, 0.01, 0.02, 0.03, 0.04, 0.05, 0.06, 0.07, 0.08, 0.09, 0.1, 0.11, 0.12, 0.13, 0.14, 0.15, 0.16, 0.17, 0.18, 0.19, 0.2, 0.21, 0.22, 0.23, 0.24, 0.25, 0.26, 0.27, 0.28, 0.29, 0.3, 0.31, 0.32, 0.33, 0.34, 0.35, 0.36, 0.37, 0.38, 0.39, 0.4, 0.41, 0.42, 0.43, 0.44, 0.45, 0.46, 0.47, 0.48, 0.49, 0.5, 0.51, 0.52, 0.53, 0.54, 0.55, 0.56, 0.57, 0.58, 0.59, 0.6, 0.61, 0.62, 0.63, 0.64, 0.65, 0.66, 0.67, 0.68, 0.69, 0.7, 0.71, 0.72, 0.73, 0.74, 0.75, 0.76, 0.77, 0.78, 0.79, 0.8, 0.81, 0.82, 0.83, 0.84, 0.85, 0.86, 0.87, 0.88, 0.89, 0.9, 0.91, 0.92, 0.93, 0.94, 0.95, 0.96, 0.97, 0.98, 0.99, 1.0, 1.01, 1.02, 1.03, 1.04, 1.05, 1.06, 1.07, 1.08, 1.09, 1.1]))[source]

Plot all fits calculated so far, each in its own tab of a FigureNotebook, where the tab is labeled by the shortened tag of the tag of the fit

Parameters
  • showZ – Overplot the values of the inverse scale lengths in red

  • x – The x values to use for plotting the fitted curve

Returns

The FigureNotebook instance created

plot_individual_fit(tag, ax=None, showZ=True, x=array([0.0, 0.0011, 0.0022, ..., 1.0978, 1.0989, 1.1]))[source]

Plot a single fit, characterized by tag

Parameters
  • tag – The tag of the fit that is to be plotted, must be in self.fits.keys()

  • ax – The axes to plot into (one is created if None)

  • showZ – Overplot the values of the inverse scale lengths in red

  • x – The x values to use for plotting the fitted curve

class omfit_classes.utils_fit.fitSpline(x, y, yerr, knots=3, min_dist=0, first_knot=None, fixed_knots=False, fit_SOL=False, monotonic=False, min_slope=None, outliers=3, plot_best=False, allow_no_errorbar=False)[source]

Bases: omfit_classes.utils_fit.knotted_fit_base

Fit a spline to some data; return the fit with uncertainties

Does basic checking for x,y,yerr then stores them in self.x, self.y, self.yerr such that x is monotonically increasing

build_params(**keyw)[source]

Build the lmfit.Parameters object needed for the fitting

Parameters

**keyw – The Fit Keywords documented in the __init__ method

Returns

The lmfit.Parameters translation of the settings given by **keyw

model(params, x, lmfit_out=None)[source]

Return the model spline curve at x

Parameters
  • params – The lmfit.Parameters object

  • x – evaluate model at x

  • lmfit_outlmfit.MinimizerResult instance to use for getting uncertainties in the curve

omfit_classes.utils_fit.xy_outliers(x, y, cutoff=1.2, return_valid=False)[source]

This function returns the index of the outlier x,y data useful to run before doing a fit of experimental data to remove outliers. This function works assuming that the first and the last samples of the x/y data set are valid data points (i.e. not outliers).

Parameters
  • x – x data (e.g. rho)

  • y – y data (e.g. ne)

  • cutoff – sensitivity of the cutoff (smaller numbers -> more sensitive [min=1])

  • return_valid – if False returns the index of the outliers, if True returns the index of the valid data

Returns

index of outliers or valid data depending on return_valid switch

class omfit_classes.utils_fit.fitGP(xx, yy, ey, noise_ub=2.0, random_starts=20, zero_value_outside=True, ntanh=1, verbose=False)[source]

Bases: object

x: array or array of arrays

Independent variable data points.

y: array or array of arrays

Dependent variable data points.

e: array or array of arrays

Uncertainty in the dependent variable data points.

noise_ub: float, optional

Upper bound on a multiplicative factor that will be optimized to infer the most probable systematic underestimation of uncertainties. Note that this factor is applied over the entire data profile, although diagnostic uncertainties are expected to be heteroschedastic. Default is 2 (giving significant freedom to the optimizer).

random_starts: int, optional

Number of random starts for the optimization of the hyperparameters of the GP. Each random starts begins sampling the posterior distribution in a different way. The optimization that gives the largest posterior probability is chosen. It is recommended to increase this value if the fit results difficult. If the regression fails, it might be necessary to vary the constraints given in the _fit method of the class GPfit2 below, which has been kept rather general for common usage. Default is 20.

zero_value_outside: bool, optional

Set to True if the profile to be evaluated is expected to go to zero beyond the LCFS, e.g. for electron temperature and density; note that this option does NOT force the value to be 0 at the LCFS, but only attempts to constrain the fit to stabilize to 0 well beyond rho=1. Profiles like those of omega_tor_12C6 and T_12C6 are experimentally observed not to go to 0 at the LCFS, so this option should be set to False for these. Default is True.

ntanh: integer, optional

Set to 2 if an internal transport barrier is expected. Default is 1 (no ITB expected). This parameter has NOT been tested recently.

verbose: bool, optional

If set to True, outputs messages from non-linear kernel optimization. Default is False.

(object) fit: call at points at which the profile is to be evaluated, e.g. if locations are stored in

an array ``xo’’, call fo = fit(xo). For an example, see 7_fit.py in OMFITprofiles.

plot(profile=None, ngauss=1)[source]

Function to conveniently plot the input data and the result of the fit. Inputs: ———– Profile: int, optional

Profile to evaluate if more than one has been computed and included in the gp object. To call the nth profile, set profile=n. If None, it will return an array of arrays.

ngauss: int, optional

Number of shaded standard deviations

None

class omfit_classes.utils_fit.fitCH(x, y, yerr, m=18)[source]

Bases: object

Fitting of kinetic profiles by Chebyshev polynomials Adapted from MATLAB function by A. Marinoni <marinoni@fusion.gat.com>

Parameters
  • x – radial coordinate

  • y – data

  • yerr – data uncertainties

  • m – Polynomial degree

plot()[source]

Plotting of the raw data and fit with uncertainties

Returns

None

class omfit_classes.utils_fit.fitLG(x, y, e, d, ng=100, sm=1, nmax=None)[source]

Bases: object

This class provides linear fitting of experimental profiles, with gaussian bluring for smoothing.

This procedure was inspired by discussions with David Eldon about the Weighted Average of Interpolations to a Common base (WAIC) technique that he describes in his thesis. However the implementation here is quite a bit different, in that instead of using a weighted average the median profile is taken, which allows for robust rejection of outliers. In this implementation the profiles smoothing is obtained by radially perturbing the measurements based on the farthest distance to their neighboring point.

Parameters
  • x – x values of the experimental data

  • y – experimental data

  • e – uncertainties of the experimental data

  • d – data time identifier

  • ng – number of gaussian instances

  • sm – smoothing

  • nmax – take no more than nmax data points

plot(variations=True)[source]
omfit_classes.utils_fit.mtanh(c, x, y=None, e=1.0, a2_plus_a3=None)[source]

Modified tanh function

>> if len(c)==6: >> y=a0*(a1+tanh((a2-x)/a3))+a4*(a5-x)*(x<a5)

a0: scale a1: offset in y a2: shift in x a3: width a4: slope of linear a5: when linear part takes off

>> if len(c)==5: >> y=a0*(a1+tanh((a2-x)/a3))+a4*(a2-a3-x)*(x<a2-a3)

a0: scale a1: offset in y a2: shift in x a3: width a4: slope of linear

Parameters
  • c – array of coefficients [a0,a1,a2,a3,a4,(a5)]

  • x – x data to fit

  • y – y data to fit

  • e – y error of the data to fit

  • a2_plus_a3 – force sum of a2 and a3 to be some value NOTE: In this case a3 should be removed from input vector c

Returns

cost, or evaluates y if y==None

omfit_classes.utils_fit.mtanh_gauss_model(x, bheight, bsol, bpos, bwidth, bslope, aheight, awidth, aexp)[source]

Modified hyperbolic tangent function for fitting pedestal with gaussian function for the fitting of the core. Stefanikova, E., et al., RewSciInst, 87 (11), Nov 2016 This function is design to fit H-mode density and temeprature profiles as a function of psi_n.

omfit_classes.utils_fit.tanh_model(x, a1, a2, a3, c)[source]
omfit_classes.utils_fit.tanh_poly_model(x, a1, a2, a3, c, p2, p3, p4)[source]
class omfit_classes.utils_fit.fit_base(x, y, yerr)[source]

Bases: object

get_param_unc(lmfit_out)[source]
property uparams
property params
valid_fit(model_out)[source]
class omfit_classes.utils_fit.fit_mtanh_gauss(x, y, yerr, **kw)[source]

Bases: omfit_classes.utils_fit.fit_base

class omfit_classes.utils_fit.fit_tanh(x, y, yerr, **kw)[source]

Bases: omfit_classes.utils_fit.fit_base

class omfit_classes.utils_fit.fit_tanh_poly(x, y, yerr, **kw)[source]

Bases: omfit_classes.utils_fit.fit_base

omfit_classes.utils_fit.tanh_through_2_points(x0, y0, x=None)[source]

Find tanh passing through two points x0=[x_0,x_1] and y0=[y_0,y_1]

y=a1*tanh((a2-x)/a3)+c

Parameters
  • x0 – iterable of two values (x coords)

  • y0 – iterable of two values (y coords)

  • x – array of coordinates where to evaluate tanh. If None function will return fit parameters

Returns

if x is None then return tanh coefficients (a1,a2,a3,c), otherwise returns y(x) points

omfit_classes.utils_fit.toq_profiles(psin, width, core, ped, edge, expin, expout)[source]

This is a direct fortran->python translation from the TOQ source

Parameters
  • psin – psi grid

  • width – width of tanh

  • core – core value

  • ped – pedestal value

  • edge – separatrix value

  • expin – inner exponent (EPED1: 1.1 for n_e, and 1.2 for T_e)

  • expout – outner exponent (EPED1: 1.1 for n_e, and 1.4 for T_e)

Returns

profile

omfit_classes.utils_fit.mtanh_polyexp(x, params)[source]

given lmfit Parameters object, generate mtanh_poly

omfit_classes.utils_fit.mtanh_polyexp_res(params, x_data, y_data, y_err)[source]
class omfit_classes.utils_fit.GASpline(xdata, ydata, ydataerr, xufit, y0=None, y0p=0.0, y1=None, y1p=None, sol=True, solc=False, numknot=3, autoknot=True, knotloc=None, knotbias=0, sepval_min=None, sepval_max=None, scrapewidth_min=None, scrapewidth_max=None, method='leastsq', maxiter=2000, monte=0, verbose=False, doPlot=False)[source]

Bases: object

Python based replacement for GAprofiles IDL spline routine “spl_mod” * Code accepts irregularly spaced (x,y,e) data and returns fit on regularly spaced grid * Numerical spline procedure based on Numerical Recipes Sec. 3.3 equations 3.3.2, 3.3.4, 3.3.5, 3.3.7 * Auto-knotting uses LMFIT minimization with chosen scheme * Boundary conditions enforced with matrix elements

The logic of this implementation is as follows: * The defualt is to auto-knot, which uses least-squares minimization to choose the knot locations. * If auto-knotting, then there are options of the guess of the knot locations or option to bias the knots. Else manual knots are used and LMFIT is not called * If the knot guess is present, then it is used else there are two options. Else we use the knot guess. * If the knot bias is None or >-1 then the knot guess is uniformly distributed using linspace. Else we use linspace with a knot bias.

For the edge data, the logic is as follows: * We can have auto/manual knots, free/fixed boundary value, and fit/ignore edge data. * When we fit edge data, then that edge data places a constraint on boundary value.

When monte-carlo is used, the return value is an unumpy uncertainties array that contains the mean and standard-deviation of the monte-carlo trials.

>>> x = np.linspace(0,1,21)
>>> y = np.sin(2*np.pi*x)
>>> e = np.repeat(0.1,len(x))
>>> xo = np.linspace(0,1,101)
>>> fit_obj = GASpline(x,y,e,xo,numknot=4,doPlot=True,monte=20)
>>> uerrorbar(x,uarray(y,e))
>>> pyplot.plot(xo,nominal_values(fit_obj(xo)))
>>> uband(xo,fit_obj(xo))
design_matrix(xcore, knotloc, bcs)[source]

Design Matrix for cubic spline interpolation Numerical Recipes Sec. 3.3 :param xcore: rho values for data on [0,1] :param knotloc: knot locations on [0,1] :param bcs: Dictionary of boundary conditions :return: design matrix for cubic interpolating spline

get_knotvals(xcore, ycore, wgt, d, geeinvc, fx, b, knotloc, bcs)[source]

Get the spline y-values at the knot locations that best fit the data :param xdata: x values of measured data :param ydata: values of measured data :param wgt: weight of measured data :param knotloc: location of spline knots [0, …, 1] :param bcs: dictionary of boundary conditions :param d, geeinvc, fx, b: Return values from design_matrix :return: values of the cubic interpolating spline at knot locations that best match the data

get_spl(x, y, w, bcs, k)[source]
Parameters
  • xdata – rho values

  • ydata – data values

  • wgt – data weight (1/uncertainty)

  • knotloc – location of knots

  • bcs – boundary conditions

Returns

spline fit on rho values

do_fit()[source]
plot()[source]
class omfit_classes.utils_fit.MtanhPolyExpFit(x_data, y_data, y_err, pow_core=2, method='leastsq', verbose=False, onAxis_gradzero=False, onAxis_value=None, fitEdge=True, edge_value=None, maxiter=None, blend_width_min=0.001, edge_width_min=0.001, edge_width_max=None, sym_guess=0.975, sym_min=0.9, sym_max=1.0, pos_edge_exp=False)[source]

Bases: object

Generalized fitter derived from B. Grierson tools fits core with pow_core polynomial C(x), and edge with offset exponential of form E(x) = offset + A*np.exp(-(x - xsym)/edge_width) blends functions together about x=x_sym with tanh-like behavior y_fit = (C(x)*np.exp(z) + E(x)*np.exp(-z))/(np.exp(z) + np.exp(-z)) where z = (xsym - x)/blend_width

Parameters
  • method – minimization method to use

  • verbose – turns on details of set flags

  • onAxis_gradzero – turn on to force C’(0) = 0 (effectively y’(0) for xsym/blend_width >> 1)

  • onAxis_value – set to force y(0) = onAxis_value

  • fitEdge – set = False to require E(x) = offset, A=edge_width=0

  • edge_value – set to force y(x=1) = edge_value

  • maxiter – controls maximum # of iterations

  • blend_width_min – minimum value for the core edge blending

  • edge_width_min – minimum value for the edge

  • sym_guess – guess for the x location of the pedestal symmetry point

  • sym_min – constraint for minimum x location for symmetry point

  • sym_max – constraint for maximum x location for symmetry point

Pos_edge_exp

force exponential to be positively valued so that exponential will have a negative slope in the SOL

Methods

__call__(x)

Evaluate mtanh_polyexp at x, propagating correlated uncertainties in the fit parameters.

class omfit_classes.utils_fit.UncertainRBF(x, d, e, centers=None, function='multiquadric', epsilon=None, norm=None)[source]

Bases: object

A class for radial basis function fitting of n-dimensional uncertain scattered data

Parameters:

Parameters
  • *args – arrays x, y, z, …, d, e where x, y, z, … are the coordinates of the nodes d is the array of values at the nodes, and e is the standard deviation error of the values at the nodes

  • centers – None the RBFs are centered on the input data points (can be very expensive for large number of nodes points) -N: N nodes randomly distributed in the domain N: N*N nodes uniformly distributed in the domain np.array(N,X): user-defined array with X coordinates of the N nodes

  • epsilon – float Adjustable constant for gaussian - defaults to approximate average distance between nodes

  • function – ‘multiquadric’: np.sqrt((r / self.epsilon) ** 2 + 1) #<— default ‘inverse’: 1.0 / np.sqrt((r / self.epsilon) ** 2 + 1) ‘gaussian’: np.exp(-(r**2 / self.epsilon)) ‘linear’: r ‘cubic’: r ** 3 ‘quintic’: r ** 5 ‘thin_plate’: r ** 2 * np.log(r)

  • norm – default “distance” is the euclidean norm (2-norm)

Examples

>>> x=np.linspace(0,1,21)
>>> y=np.linspace(0,1,21)
>>> e=x*0+.5
>>> e[abs(x-0.5)<0.1]=8
>>> y[abs(x-0.5)<0.1]=10
>>>
>>> x1 = np.linspace(0,1,100)
>>> y1 = UncertainRBF(x, y, e, centers=None, epsilon=1)(x1)
>>> y0 = UncertainRBF(x, y, e*0+1, centers=None, epsilon=1)(x1)
>>>
>>> pyplot.subplot(2,2,1)
>>> errorbar(x,y,e,ls='',marker='.',label='raw 1D data')
>>> uband(x1,y1,label='1D RBF w/ uncertainty')
>>> pyplot.plot(x1,nominal_values(y0),label='1D RBF w/o uncertainty')
>>> pyplot.title('1D')
>>> legend