Skip to content

Commit

Permalink
Merge pull request #167 from FarBo/adding_ReIm
Browse files Browse the repository at this point in the history 
Added functionality of choosing between Re/Im and Re in 2D complex plots
  • Loading branch information
@jenshnielsen
jenshnielsen committed Jan 8, 2021
2 parents 07f67b8 + ea22e76 commit 91e2d49
Showing 1 changed file with 145 additions and 52 deletions.
197 changes: 145 additions & 52 deletions plottr/plot/mpl.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,6 @@ class ComplexRepresentation(Enum):
#: magnitude and phase
magAndPhase = auto()


def determinePlotDataType(data: Optional[DataDictBase]) -> PlotDataType:
"""
Analyze input data and determine most likely :class:`PlotDataType`.
Expand Down Expand Up @@ -614,7 +613,7 @@ class _AutoPlotToolBar(QtWidgets.QToolBar):
plotTypeSelected = Signal(PlotType)

#: signal emitted when the complex data option has been changed
complexPolarSelected = Signal(bool)
complexPolarSelected = Signal(ComplexRepresentation)

def __init__(self, name: str, parent: Optional[QtWidgets.QWidget] = None):
"""Constructor for :class:`AutoPlotToolBar`"""
Expand Down Expand Up @@ -656,9 +655,20 @@ def __init__(self, name: str, parent: Optional[QtWidgets.QWidget] = None):
# other options
self.addSeparator()

self.plotComplexPolar = self.addAction('Mag/Phase')
self.plotComplexPolar.setCheckable(True)
self.plotComplexPolar.triggered.connect(self._trigger_complex_mag_phase)
self.plotReal = self.addAction('Real')
self.plotReal.setCheckable(True)
self.plotReal.triggered.connect(
lambda: self.selectComplexType(ComplexRepresentation.real))

self.plotComplexReIm = self.addAction('Real/Imag')
self.plotComplexReIm.setCheckable(True)
self.plotComplexReIm.triggered.connect(
lambda: self.selectComplexType(ComplexRepresentation.realAndImag))

self.plotMagPhase = self.addAction('Mag/Phase')
self.plotMagPhase.setCheckable(True)
self.plotMagPhase.triggered.connect(
lambda: self.selectComplexType(ComplexRepresentation.magAndPhase))

self.plotTypeActions = OrderedDict({
PlotType.multitraces: self.plotasMultiTraces,
Expand All @@ -668,11 +678,18 @@ def __init__(self, name: str, parent: Optional[QtWidgets.QWidget] = None):
PlotType.scatter2d: self.plotasScatter2d,
})

self.ComplexActions = OrderedDict({
ComplexRepresentation.real: self.plotReal,
ComplexRepresentation.realAndImag: self.plotComplexReIm,
ComplexRepresentation.magAndPhase: self.plotMagPhase
})

self._currentPlotType = PlotType.empty
self._currentlyAllowedPlotTypes: Tuple[PlotType, ...] = ()

def _trigger_complex_mag_phase(self, enable: bool) -> None:
self.complexPolarSelected.emit(enable)
self._currentComplex = ComplexRepresentation.realAndImag
self.ComplexActions[self._currentComplex].setChecked(True)
self._currentlyAllowedComplexTypes: Tuple[ComplexRepresentation, ...] = ()

def selectPlotType(self, plotType: PlotType) -> None:
"""makes sure that the selected `plotType` is active (checked), all
Expand Down Expand Up @@ -724,6 +741,54 @@ def setAllowedPlotTypes(self, *args: PlotType) -> None:

self._currentlyAllowedPlotTypes = args

def selectComplexType(self, comp: ComplexRepresentation) -> None:
"""makes sure that the selected `comp` is active (checked), all
others are not active.
This method should be used to catch a trigger from the UI.
If the active plot type has been changed by using this method,
we emit `complexPolarSelected`.
"""
# deselect all other types
for k, v in self.ComplexActions.items():
if k is not comp and v is not None:
v.setChecked(False)

# don't want un-toggling - can only be done by selecting another type
self.ComplexActions[comp].setChecked(True)

if comp is not self._currentPlotType:
self._currentComplex = comp
self.complexPolarSelected.emit(self._currentComplex)

def setAllowedComplexTypes(self, *complexOptions: ComplexRepresentation) -> None:
"""Disable all choices that are not allowed.
If the current selection is now disabled, instead select the first
enabled one.
"""

if complexOptions == self._currentlyAllowedComplexTypes:
return

for k, v in self.ComplexActions.items():
if k not in complexOptions:
v.setChecked(False)
v.setEnabled(False)
else:
v.setEnabled(True)

if self._currentComplex not in complexOptions:
self._currentComplex = ComplexRepresentation.realAndImag
for k, v in self.ComplexActions.items():
if k in complexOptions:
v.setChecked(True)
self._currentComplex = k
break

self.complexPolarSelected.emit(self._currentComplex)

self._currentlyAllowedComplexTypes = complexOptions

class AutoPlot(_MPLPlotWidget):
"""A widget for plotting with matplotlib.
Expand Down Expand Up @@ -766,8 +831,10 @@ def __init__(self, parent: Optional[QtWidgets.QWidget] = None):

self.plotDataType = PlotDataType.unknown
self.plotType = PlotType.empty
self.complexRepresentation = ComplexRepresentation.real
self.complexPreference = ComplexRepresentation.realAndImag

# The default complex behavior is set here. This choice is according to
# the first item in ComplexActions in class _AutoPlotToolBar
self.complexRepresentation = ComplexRepresentation.realAndImag

self.dataType: Optional[Type[DataDictBase]] = None
self.dataStructure: Optional[DataDictBase] = None
Expand Down Expand Up @@ -852,6 +919,7 @@ def setData(self, data: Optional[DataDictBase]) -> None:
self.plotDataType = determinePlotDataType(data)

self._processPlotTypeOptions()
self._processComplexTypeOptions()
self._plotData(adjustSize=True)

def _processPlotTypeOptions(self) -> None:
Expand All @@ -875,20 +943,31 @@ def _processPlotTypeOptions(self) -> None:
else:
self.plotOptionsToolBar.setAllowedPlotTypes()

def _processComplexTypeOptions(self) -> None:
"""Given data is complex or not, define what complex options to be selected."""
if self.data is not None:
if self.dataIsComplex() == True:
self.plotOptionsToolBar.setAllowedComplexTypes(
ComplexRepresentation.real,
ComplexRepresentation.realAndImag,
ComplexRepresentation.magAndPhase,
)
else:
self.plotOptionsToolBar.setAllowedComplexTypes(
ComplexRepresentation.real
)

@Slot(PlotType)
def _plotTypeFromToolBar(self, plotType: PlotType) -> None:
if plotType is not self.plotType:
self.plotType = plotType
self._plotData(adjustSize=True)

@Slot(bool)
def _complexPreferenceFromToolBar(self, magPhasePreferred: bool) -> None:
if magPhasePreferred:
self.complexPreference = ComplexRepresentation.magAndPhase
else:
self.complexPreference = ComplexRepresentation.realAndImag

self._plotData(adjustSize=True)
@Slot(ComplexRepresentation)
def _complexPreferenceFromToolBar(self, complexRepresentation: ComplexRepresentation) -> None:
if complexRepresentation is not self.complexRepresentation:
self.complexRepresentation = complexRepresentation
self._plotData(adjustSize=True)

def _makeAxes(self, nAxes: int) -> List[Axes]:
"""Create a grid of axes.
Expand All @@ -911,8 +990,6 @@ def _plotData(self, adjustSize: bool = False) -> None:

if not self.dataIsComplex():
self.complexRepresentation = ComplexRepresentation.real
else:
self.complexRepresentation = self.complexPreference

if self.plotType is PlotType.multitraces:
logger.debug(f"Plotting lines in a single panel")
Expand Down Expand Up @@ -949,7 +1026,8 @@ def _plot1dSinglepanel(self) -> None:
depvals = [self.data.data_vals(d) for d in depnames]

# count the number of panels we need.
if self.complexRepresentation is ComplexRepresentation.magAndPhase:
if self.complexRepresentation is ComplexRepresentation.magAndPhase or\
self.complexRepresentation is ComplexRepresentation.realAndImag:
nAxes = 2
else:
nAxes = 1
Expand All @@ -968,15 +1046,17 @@ def _plot1dSinglepanel(self) -> None:
if isinstance(yvals, np.ma.MaskedArray):
yvals = yvals.filled(np.nan)

if self.complexRepresentation in [ComplexRepresentation.real,
ComplexRepresentation.realAndImag]:
plot1dTrace(axes[0], xvals, np.asanyarray(yvals),
axLabels=(self.data.label(xname), ylbl),
curveLabel=self.data.label(yname),
addLegend=(yname == depnames[-1]))

elif self.complexRepresentation is ComplexRepresentation.magAndPhase:
if self.dataIsComplex(yname):
if self.dataIsComplex(yname):
if self.complexRepresentation is ComplexRepresentation.realAndImag:
plot1dTrace(axes[0], xvals, np.asanyarray(yvals).real,
axLabels=(self.data.label(xname), ylbl),
curveLabel=f"Re({self.data.label(yname)})",
addLegend=(yname == depnames[-1]))
plot1dTrace(axes[1], xvals, np.asanyarray(yvals).imag,
axLabels=(self.data.label(xname), ylbl),
curveLabel=f"Im({yname})",
addLegend=(yname == depnames[-1]))
elif self.complexRepresentation is ComplexRepresentation.magAndPhase:
plot1dTrace(axes[0], xvals, np.real(np.abs(yvals)),
axLabels=(self.data.label(xname), ylbl),
curveLabel=f"Abs({self.data.label(yname)})",
Expand All @@ -985,11 +1065,16 @@ def _plot1dSinglepanel(self) -> None:
axLabels=(self.data.label(xname), phlbl),
curveLabel=f"Arg({yname})",
addLegend=(yname == depnames[-1]))
else:
plot1dTrace(axes[0], xvals, np.asanyarray(yvals),
elif self.complexRepresentation is ComplexRepresentation.real:
plot1dTrace(axes[0], xvals, np.asanyarray(yvals).real,
axLabels=(self.data.label(xname), ylbl),
curveLabel=self.data.label(yname),
curveLabel=f"Re({self.data.label(yname)})",
addLegend=(yname == depnames[-1]))
else:
plot1dTrace(axes[0], xvals, np.asanyarray(yvals),
axLabels=(self.data.label(xname), ylbl),
curveLabel=self.data.label(yname),
addLegend=(yname == depnames[-1]))

def _plot1dSeparatePanels(self) -> None:
assert self.data is not None
Expand All @@ -998,8 +1083,7 @@ def _plot1dSeparatePanels(self) -> None:
depnames = self.data.dependents()
depvals = [self.data.data_vals(d) for d in depnames]

if self.complexRepresentation in [ComplexRepresentation.real,
ComplexRepresentation.realAndImag]:
if self.complexRepresentation is ComplexRepresentation.real:
nAxes = len(depnames)
else:
nAxes = 0
Expand All @@ -1019,27 +1103,36 @@ def _plot1dSeparatePanels(self) -> None:
if isinstance(yvals, np.ma.MaskedArray):
yvals = yvals.filled(np.nan)

if self.complexRepresentation in [ComplexRepresentation.real,
ComplexRepresentation.realAndImag]:
plot1dTrace(axes[iax], xvals, np.asanyarray(yvals),
axLabels=(self.data.label(xname), self.data.label(yname)),
addLegend=self.dataIsComplex(yname))
iax += 1

elif self.complexRepresentation is ComplexRepresentation.magAndPhase:
if self.dataIsComplex(yname):
plot1dTrace(axes[iax], xvals, np.real(np.abs(yvals)),
if self.dataIsComplex(yname):
if self.complexRepresentation is ComplexRepresentation.realAndImag:
plot1dTrace(axes[iax], xvals, np.real(yvals),
axLabels=(self.data.label(xname),
f"Abs({self.data.label(yname)})"))
plot1dTrace(axes[iax+1], xvals, np.angle(yvals),
f"Re({self.data.label(yname)})"))
plot1dTrace(axes[iax+1], xvals, np.imag(yvals),
axLabels=(self.data.label(xname),
f"Arg({yname})"))
f"Im({self.data.label(yname)})"))
iax += 2
else:
plot1dTrace(axes[iax], xvals, np.asanyarray(yvals),
axLabels=(self.data.label(xname),
self.data.label(yname)))
iax += 1
elif self.complexRepresentation is ComplexRepresentation.magAndPhase:
if self.dataIsComplex(yname):
plot1dTrace(axes[iax], xvals, np.real(np.abs(yvals)),
axLabels=(self.data.label(xname),
f"Abs({self.data.label(yname)})"))
plot1dTrace(axes[iax+1], xvals, np.angle(yvals),
axLabels=(self.data.label(xname),
f"Arg({yname})"))
iax += 2
elif self.complexRepresentation is ComplexRepresentation.real:
if self.dataIsComplex(yname):
plot1dTrace(axes[iax], xvals, np.asanyarray(yvals).real,
axLabels=(self.data.label(xname),
f"Re({self.data.label(yname)})"),
)
iax += 1
else:
plot1dTrace(axes[iax], xvals, np.asanyarray(yvals),
axLabels=(self.data.label(xname),
self.data.label(yname)))
iax += 1

def _colorplot2d(self) -> None:
assert self.data is not None
Expand Down

0 comments on commit 91e2d49

Please sign in to comment.