354 lines
12 KiB
Python
354 lines
12 KiB
Python
# -*- coding: utf-8 -*-
|
|
from ..Node import Node
|
|
from pyqtgraph.Qt import QtGui, QtCore
|
|
import numpy as np
|
|
from .common import *
|
|
from pyqtgraph.SRTTransform import SRTTransform
|
|
from pyqtgraph.Point import Point
|
|
from pyqtgraph.widgets.TreeWidget import TreeWidget
|
|
from pyqtgraph.graphicsItems.LinearRegionItem import LinearRegionItem
|
|
|
|
from . import functions
|
|
|
|
try:
|
|
import metaarray
|
|
HAVE_METAARRAY = True
|
|
except:
|
|
HAVE_METAARRAY = False
|
|
|
|
class ColumnSelectNode(Node):
|
|
"""Select named columns from a record array or MetaArray."""
|
|
nodeName = "ColumnSelect"
|
|
def __init__(self, name):
|
|
Node.__init__(self, name, terminals={'In': {'io': 'in'}})
|
|
self.columns = set()
|
|
self.columnList = QtGui.QListWidget()
|
|
self.axis = 0
|
|
self.columnList.itemChanged.connect(self.itemChanged)
|
|
|
|
def process(self, In, display=True):
|
|
if display:
|
|
self.updateList(In)
|
|
|
|
out = {}
|
|
if HAVE_METAARRAY and isinstance(In, metaarray.MetaArray):
|
|
for c in self.columns:
|
|
out[c] = In[self.axis:c]
|
|
elif isinstance(In, np.ndarray) and In.dtype.fields is not None:
|
|
for c in self.columns:
|
|
out[c] = In[c]
|
|
else:
|
|
self.In.setValueAcceptable(False)
|
|
raise Exception("Input must be MetaArray or ndarray with named fields")
|
|
|
|
return out
|
|
|
|
def ctrlWidget(self):
|
|
return self.columnList
|
|
|
|
def updateList(self, data):
|
|
if HAVE_METAARRAY and isinstance(data, metaarray.MetaArray):
|
|
cols = data.listColumns()
|
|
for ax in cols: ## find first axis with columns
|
|
if len(cols[ax]) > 0:
|
|
self.axis = ax
|
|
cols = set(cols[ax])
|
|
break
|
|
else:
|
|
cols = list(data.dtype.fields.keys())
|
|
|
|
rem = set()
|
|
for c in self.columns:
|
|
if c not in cols:
|
|
self.removeTerminal(c)
|
|
rem.add(c)
|
|
self.columns -= rem
|
|
|
|
self.columnList.blockSignals(True)
|
|
self.columnList.clear()
|
|
for c in cols:
|
|
item = QtGui.QListWidgetItem(c)
|
|
item.setFlags(QtCore.Qt.ItemIsEnabled|QtCore.Qt.ItemIsUserCheckable)
|
|
if c in self.columns:
|
|
item.setCheckState(QtCore.Qt.Checked)
|
|
else:
|
|
item.setCheckState(QtCore.Qt.Unchecked)
|
|
self.columnList.addItem(item)
|
|
self.columnList.blockSignals(False)
|
|
|
|
|
|
def itemChanged(self, item):
|
|
col = str(item.text())
|
|
if item.checkState() == QtCore.Qt.Checked:
|
|
if col not in self.columns:
|
|
self.columns.add(col)
|
|
self.addOutput(col)
|
|
else:
|
|
if col in self.columns:
|
|
self.columns.remove(col)
|
|
self.removeTerminal(col)
|
|
self.update()
|
|
|
|
def saveState(self):
|
|
state = Node.saveState(self)
|
|
state['columns'] = list(self.columns)
|
|
return state
|
|
|
|
def restoreState(self, state):
|
|
Node.restoreState(self, state)
|
|
self.columns = set(state.get('columns', []))
|
|
for c in self.columns:
|
|
self.addOutput(c)
|
|
|
|
|
|
|
|
class RegionSelectNode(CtrlNode):
|
|
"""Returns a slice from a 1-D array. Connect the 'widget' output to a plot to display a region-selection widget."""
|
|
nodeName = "RegionSelect"
|
|
uiTemplate = [
|
|
('start', 'spin', {'value': 0, 'step': 0.1}),
|
|
('stop', 'spin', {'value': 0.1, 'step': 0.1}),
|
|
('display', 'check', {'value': True}),
|
|
('movable', 'check', {'value': True}),
|
|
]
|
|
|
|
def __init__(self, name):
|
|
self.items = {}
|
|
CtrlNode.__init__(self, name, terminals={
|
|
'data': {'io': 'in'},
|
|
'selected': {'io': 'out'},
|
|
'region': {'io': 'out'},
|
|
'widget': {'io': 'out', 'multi': True}
|
|
})
|
|
self.ctrls['display'].toggled.connect(self.displayToggled)
|
|
self.ctrls['movable'].toggled.connect(self.movableToggled)
|
|
|
|
def displayToggled(self, b):
|
|
for item in self.items.values():
|
|
item.setVisible(b)
|
|
|
|
def movableToggled(self, b):
|
|
for item in self.items.values():
|
|
item.setMovable(b)
|
|
|
|
|
|
def process(self, data=None, display=True):
|
|
#print "process.."
|
|
s = self.stateGroup.state()
|
|
region = [s['start'], s['stop']]
|
|
|
|
if display:
|
|
conn = self['widget'].connections()
|
|
for c in conn:
|
|
plot = c.node().getPlot()
|
|
if plot is None:
|
|
continue
|
|
if c in self.items:
|
|
item = self.items[c]
|
|
item.setRegion(region)
|
|
#print " set rgn:", c, region
|
|
#item.setXVals(events)
|
|
else:
|
|
item = LinearRegionItem(values=region)
|
|
self.items[c] = item
|
|
#item.connect(item, QtCore.SIGNAL('regionChanged'), self.rgnChanged)
|
|
item.sigRegionChanged.connect(self.rgnChanged)
|
|
item.setVisible(s['display'])
|
|
item.setMovable(s['movable'])
|
|
#print " new rgn:", c, region
|
|
#self.items[c].setYRange([0., 0.2], relative=True)
|
|
|
|
if self.selected.isConnected():
|
|
if data is None:
|
|
sliced = None
|
|
elif isinstance(data, MetaArray):
|
|
sliced = data[0:s['start']:s['stop']]
|
|
else:
|
|
mask = (data['time'] >= s['start']) * (data['time'] < s['stop'])
|
|
sliced = data[mask]
|
|
else:
|
|
sliced = None
|
|
|
|
return {'selected': sliced, 'widget': self.items, 'region': region}
|
|
|
|
|
|
def rgnChanged(self, item):
|
|
region = item.getRegion()
|
|
self.stateGroup.setState({'start': region[0], 'stop': region[1]})
|
|
self.update()
|
|
|
|
|
|
class EvalNode(Node):
|
|
"""Return the output of a string evaluated/executed by the python interpreter.
|
|
The string may be either an expression or a python script, and inputs are accessed as the name of the terminal.
|
|
For expressions, a single value may be evaluated for a single output, or a dict for multiple outputs.
|
|
For a script, the text will be executed as the body of a function."""
|
|
nodeName = 'PythonEval'
|
|
|
|
def __init__(self, name):
|
|
Node.__init__(self, name,
|
|
terminals = {
|
|
'input': {'io': 'in', 'renamable': True},
|
|
'output': {'io': 'out', 'renamable': True},
|
|
},
|
|
allowAddInput=True, allowAddOutput=True)
|
|
|
|
self.ui = QtGui.QWidget()
|
|
self.layout = QtGui.QGridLayout()
|
|
#self.addInBtn = QtGui.QPushButton('+Input')
|
|
#self.addOutBtn = QtGui.QPushButton('+Output')
|
|
self.text = QtGui.QTextEdit()
|
|
self.text.setTabStopWidth(30)
|
|
self.text.setPlainText("# Access inputs as args['input_name']\nreturn {'output': None} ## one key per output terminal")
|
|
#self.layout.addWidget(self.addInBtn, 0, 0)
|
|
#self.layout.addWidget(self.addOutBtn, 0, 1)
|
|
self.layout.addWidget(self.text, 1, 0, 1, 2)
|
|
self.ui.setLayout(self.layout)
|
|
|
|
#QtCore.QObject.connect(self.addInBtn, QtCore.SIGNAL('clicked()'), self.addInput)
|
|
#self.addInBtn.clicked.connect(self.addInput)
|
|
#QtCore.QObject.connect(self.addOutBtn, QtCore.SIGNAL('clicked()'), self.addOutput)
|
|
#self.addOutBtn.clicked.connect(self.addOutput)
|
|
self.text.focusOutEvent = self.focusOutEvent
|
|
self.lastText = None
|
|
|
|
def ctrlWidget(self):
|
|
return self.ui
|
|
|
|
#def addInput(self):
|
|
#Node.addInput(self, 'input', renamable=True)
|
|
|
|
#def addOutput(self):
|
|
#Node.addOutput(self, 'output', renamable=True)
|
|
|
|
def focusOutEvent(self, ev):
|
|
text = str(self.text.toPlainText())
|
|
if text != self.lastText:
|
|
self.lastText = text
|
|
print("eval node update")
|
|
self.update()
|
|
return QtGui.QTextEdit.focusOutEvent(self.text, ev)
|
|
|
|
def process(self, display=True, **args):
|
|
l = locals()
|
|
l.update(args)
|
|
## try eval first, then exec
|
|
try:
|
|
text = str(self.text.toPlainText()).replace('\n', ' ')
|
|
output = eval(text, globals(), l)
|
|
except SyntaxError:
|
|
fn = "def fn(**args):\n"
|
|
run = "\noutput=fn(**args)\n"
|
|
text = fn + "\n".join([" "+l for l in str(self.text.toPlainText()).split('\n')]) + run
|
|
exec(text)
|
|
return output
|
|
|
|
def saveState(self):
|
|
state = Node.saveState(self)
|
|
state['text'] = str(self.text.toPlainText())
|
|
state['terminals'] = self.saveTerminals()
|
|
return state
|
|
|
|
def restoreState(self, state):
|
|
Node.restoreState(self, state)
|
|
self.text.clear()
|
|
self.text.insertPlainText(state['text'])
|
|
self.restoreTerminals(state['terminals'])
|
|
self.update()
|
|
|
|
class ColumnJoinNode(Node):
|
|
"""Concatenates record arrays and/or adds new columns"""
|
|
nodeName = 'ColumnJoin'
|
|
|
|
def __init__(self, name):
|
|
Node.__init__(self, name, terminals = {
|
|
'output': {'io': 'out'},
|
|
})
|
|
|
|
#self.items = []
|
|
|
|
self.ui = QtGui.QWidget()
|
|
self.layout = QtGui.QGridLayout()
|
|
self.ui.setLayout(self.layout)
|
|
|
|
self.tree = TreeWidget()
|
|
self.addInBtn = QtGui.QPushButton('+ Input')
|
|
self.remInBtn = QtGui.QPushButton('- Input')
|
|
|
|
self.layout.addWidget(self.tree, 0, 0, 1, 2)
|
|
self.layout.addWidget(self.addInBtn, 1, 0)
|
|
self.layout.addWidget(self.remInBtn, 1, 1)
|
|
|
|
self.addInBtn.clicked.connect(self.addInput)
|
|
self.remInBtn.clicked.connect(self.remInput)
|
|
self.tree.sigItemMoved.connect(self.update)
|
|
|
|
def ctrlWidget(self):
|
|
return self.ui
|
|
|
|
def addInput(self):
|
|
#print "ColumnJoinNode.addInput called."
|
|
term = Node.addInput(self, 'input', renamable=True)
|
|
#print "Node.addInput returned. term:", term
|
|
item = QtGui.QTreeWidgetItem([term.name()])
|
|
item.term = term
|
|
term.joinItem = item
|
|
#self.items.append((term, item))
|
|
self.tree.addTopLevelItem(item)
|
|
|
|
def remInput(self):
|
|
sel = self.tree.currentItem()
|
|
term = sel.term
|
|
term.joinItem = None
|
|
sel.term = None
|
|
self.tree.removeTopLevelItem(sel)
|
|
self.removeTerminal(term)
|
|
self.update()
|
|
|
|
def process(self, display=True, **args):
|
|
order = self.order()
|
|
vals = []
|
|
for name in order:
|
|
if name not in args:
|
|
continue
|
|
val = args[name]
|
|
if isinstance(val, np.ndarray) and len(val.dtype) > 0:
|
|
vals.append(val)
|
|
else:
|
|
vals.append((name, None, val))
|
|
return {'output': functions.concatenateColumns(vals)}
|
|
|
|
def order(self):
|
|
return [str(self.tree.topLevelItem(i).text(0)) for i in range(self.tree.topLevelItemCount())]
|
|
|
|
def saveState(self):
|
|
state = Node.saveState(self)
|
|
state['order'] = self.order()
|
|
return state
|
|
|
|
def restoreState(self, state):
|
|
Node.restoreState(self, state)
|
|
inputs = [inp.name() for inp in self.inputs()]
|
|
for name in inputs:
|
|
if name not in state['order']:
|
|
self.removeTerminal(name)
|
|
for name in state['order']:
|
|
if name not in inputs:
|
|
Node.addInput(self, name, renamable=True)
|
|
|
|
self.tree.clear()
|
|
for name in state['order']:
|
|
term = self[name]
|
|
item = QtGui.QTreeWidgetItem([name])
|
|
item.term = term
|
|
term.joinItem = item
|
|
#self.items.append((term, item))
|
|
self.tree.addTopLevelItem(item)
|
|
|
|
def terminalRenamed(self, term, oldName):
|
|
Node.terminalRenamed(self, term, oldName)
|
|
item = term.joinItem
|
|
item.setText(0, term.name())
|
|
self.update()
|
|
|
|
|