So things to main and update makeGraph

This commit is contained in:
Cal Wing 2023-10-31 23:05:24 +10:00
parent 17cfa82b4b
commit fb1b40503a
2 changed files with 99 additions and 30 deletions

34
main.py
View File

@ -75,8 +75,8 @@ def cacheData(dataFilePath: str, calcFunction: callable = lambda x: x, args: tup
# - FSP_Y (in m/s^2) - the specific force in the body-fixed y-direction # - FSP_Y (in m/s^2) - the specific force in the body-fixed y-direction
# - FSP_Z (in m/s^2) - the specific force in the body-fixed z-direction # - FSP_Z (in m/s^2) - the specific force in the body-fixed z-direction
IMU_TIME_HEADER = ["Time [s]"] IMU_TIME_HEADER = ["Time [s]"]
IMU_WBE_HEADERS = ["WBE_1 [rad/s]", "WBE_2 [rad/s]", "WBE_3 [rad/s]"] IMU_WBE_HEADERS = ["WBE_1 [rad/s]", "WBE_2 [rad/s]", "WBE_3 [rad/s]"] # Roll, Pitch, Yaw - Rate
IMU_FSP_HEADERS = ["FSP_X [m/s^2]", "FSP_Y [m/s^2]", "FSP_Z [m/s^2]"] IMU_FSP_HEADERS = ["FSP_X [m/s^2]", "FSP_Y [m/s^2]", "FSP_Z [m/s^2]"] # Specific Force X, Y , Z
IMU_DATA_HEADER = IMU_TIME_HEADER + IMU_WBE_HEADERS + IMU_FSP_HEADERS IMU_DATA_HEADER = IMU_TIME_HEADER + IMU_WBE_HEADERS + IMU_FSP_HEADERS
def importIMUData(mission, imu): def importIMUData(mission, imu):
# If IMU is not a string, then convert based on bool eval where "H" == True # If IMU is not a string, then convert based on bool eval where "H" == True
@ -96,7 +96,7 @@ m2_IMUData = importIMUData(2, 0), importIMUData(2, 1)
# NED Translation & Force Functions # NED Translation & Force Functions
INIT_EULER_ANGLES = (0, 0, 0) INIT_EULER_ANGLES = (0, 0, 0)
def translate2NED(angles, euler_angles): def attidude_rate_2NED(angles, euler_angles):
phi, theta, psi = euler_angles phi, theta, psi = euler_angles
p, q, r = angles p, q, r = angles
@ -113,14 +113,14 @@ def translate2NED(angles, euler_angles):
return np.matmul(transMat, angleMat) return np.matmul(transMat, angleMat)
def getNEDForces(NEDPos): def getNEDForces(NEDPos):
phi, theta, psi = NEDPos phi, theta, psi = NEDPos # Roll, Pitch, Yaw
forceMat = np.array([ [cos(psi)*cos(theta), cos(psi)*sin(theta)*sin(phi)-sin(psi)*cos(phi), cos(psi)*sin(theta)*cos(phi)+sin(psi)*sin(phi)], R41 = np.array([ [cos(psi)*cos(theta), cos(psi)*sin(theta)*sin(phi)-sin(psi)*cos(phi), cos(psi)*sin(theta)*cos(phi)+sin(psi)*sin(phi)],
[sin(psi)*cos(theta), sin(psi)*sin(theta)*sin(phi)+cos(psi)*cos(phi), sin(psi)*sin(theta)*cos(phi)-cos(psi)*sin(phi)], [sin(psi)*cos(theta), sin(psi)*sin(theta)*sin(phi)+cos(psi)*cos(phi), sin(psi)*sin(theta)*cos(phi)-cos(psi)*sin(phi)],
[-sin(theta), cos(theta)*sin(phi), cos(theta)*cos(phi) ] [-sin(theta), cos(theta)*sin(phi), cos(theta)*cos(phi) ]
]) ])
return forceMat return R41
# Calculate Mission Translation & Force Data # Calculate Mission Translation & Force Data
TRANS_DATA_HEADER = IMU_TIME_HEADER + IMU_WBE_HEADERS + ["Forces", "ForceSum", "ForceCumSum"] TRANS_DATA_HEADER = IMU_TIME_HEADER + IMU_WBE_HEADERS + ["Forces", "ForceSum", "ForceCumSum"]
@ -128,13 +128,13 @@ def calculateTranslatedData(missionWBEData) -> pd.DataFrame:
print("Translating Motion & Calculating Resulting Forces") print("Translating Motion & Calculating Resulting Forces")
translatedData = pd.DataFrame(columns=TRANS_DATA_HEADER) translatedData = pd.DataFrame(columns=TRANS_DATA_HEADER)
for i in tqdm(range(len(missionWBEData))): for i in tqdm(range(len(missionWBEData))):
dataPoint = missionWBEData[IMU_WBE_HEADERS].iloc[i] dataPoint = missionWBEData[IMU_WBE_HEADERS].iloc[i] # Get the time point data
trans = translate2NED(dataPoint.values, INIT_EULER_ANGLES).flatten() trans = attidude_rate_2NED(dataPoint.values, INIT_EULER_ANGLES).flatten() # Translate to Net
forces = getNEDForces(trans) forces = getNEDForces(trans)
forceSum = np.array([np.sum(forces[:,0].flatten()), np.sum(forces[:,1].flatten()), np.sum(forces[:,2].flatten())]) forceSum = np.array([np.sum(forces[:,0].flatten()), np.sum(forces[:,1].flatten()), np.sum(forces[:,2].flatten())])
privForce = translatedData[TRANS_DATA_HEADER[5]].iloc[i-1] if i > 0 else [0, 0, 0] privForce = translatedData[TRANS_DATA_HEADER[5]].iloc[i-1] if i > 0 else [0, 0, 0]
forceCumSum = np.array([ forceSum[ii] + value for ii, value in enumerate(privForce) ]) forceCumSum = np.array([ forceSum[ii] + value for ii, value in enumerate(privForce) ])
dataRow = { translatedData.loc[i] = {
TRANS_DATA_HEADER[0]: missionWBEData[IMU_TIME_HEADER].iloc[i], TRANS_DATA_HEADER[0]: missionWBEData[IMU_TIME_HEADER].iloc[i],
TRANS_DATA_HEADER[1]: trans[0], TRANS_DATA_HEADER[1]: trans[0],
TRANS_DATA_HEADER[2]: trans[1], TRANS_DATA_HEADER[2]: trans[1],
@ -143,23 +143,27 @@ def calculateTranslatedData(missionWBEData) -> pd.DataFrame:
TRANS_DATA_HEADER[5]: forceSum, TRANS_DATA_HEADER[5]: forceSum,
TRANS_DATA_HEADER[6]: forceCumSum, TRANS_DATA_HEADER[6]: forceCumSum,
} }
translatedData.loc[i] = dataRow
return translatedData return translatedData
if __name__ == '__main__': if __name__ == '__main__':
missionWBEData = m1_IMUData[0][IMU_TIME_HEADER + IMU_WBE_HEADERS] missionWBEData = m1_IMUData[0][IMU_TIME_HEADER + IMU_WBE_HEADERS]
translatedData = cacheData("./tmp/m1_transData.dfz", calculateTranslatedData, (missionWBEData,)) translatedData = cacheData("./tmp/m1_transData.dfz", calculateTranslatedData, (missionWBEData,))
x = np.cumsum(translatedData[IMU_WBE_HEADERS[0]])
y = np.cumsum(translatedData[IMU_WBE_HEADERS[1]])
z = np.cumsum(translatedData[IMU_WBE_HEADERS[2]])
import matplotlib.pyplot as plt import matplotlib.pyplot as plt
fig = plt.figure() fig = plt.figure()
ax = plt.axes(projection='3d') ax = plt.axes(projection='3d')
plots = translatedData[TRANS_DATA_HEADER[6]].values plots = translatedData[TRANS_DATA_HEADER[6]].values
x = [p[0] for p in plots] #x = [p[0] for p in plots]
y = [p[1] for p in plots] #y = [p[1] for p in plots]
z = [p[2] for p in plots] #z = [p[2] for p in plots]
ax.plot3D(x, y, z) ax.plot3D(x, y, z)
ax.plot3D(x+5, y-2, z+1)
plt.show() plt.show()

View File

@ -6,7 +6,7 @@
#### 2023 - Added UQ Default Colours to MatPlotLib #### 2023 - Added UQ Default Colours to MatPlotLib
__author__ = "Cal Wing" __author__ = "Cal Wing"
__version__ = "0.1.10" __version__ = "0.1.12"
from collections.abc import Iterator from collections.abc import Iterator
import numpy as np import numpy as np
@ -15,6 +15,7 @@ import matplotlib.pyplot as plt
import matplotlib.colors as colors import matplotlib.colors as colors
from mpl_toolkits.axes_grid1 import make_axes_locatable from mpl_toolkits.axes_grid1 import make_axes_locatable
from cycler import cycler from cycler import cycler
from tqdm import tqdm
# Define the UQ Colours # Define the UQ Colours
UQ_COLOURS_DICT = { UQ_COLOURS_DICT = {
@ -62,6 +63,10 @@ class ColourValue(str):
def hsv(self) -> np.ndarray: def hsv(self) -> np.ndarray:
return colors.rgb_to_hsv(self.rgb()) return colors.rgb_to_hsv(self.rgb())
def alpha_adj(self, alpha):
r, g, b, a = self.rgba()
return (r, g, b, alpha)
# Define the UQ Colours in a nicer object # Define the UQ Colours in a nicer object
class ColourList(object): class ColourList(object):
@ -243,7 +248,7 @@ def makeGraph(graphData, showPlot=True, doProgramBlock=True, figSavePath=None, h
"xLabel", "yLabel", "title", "axis", "grid", "xPos", "yPos", "xLabel", "yLabel", "title", "axis", "grid", "xPos", "yPos",
"xLabelPos", "yLabelPos", "xTickPos", "yTickPos", "xScale", "yScale", "xLabelPos", "yLabelPos", "xTickPos", "yTickPos", "xScale", "yScale",
"xTickMap", "yTickMap", "plots", "xLim", "yLim", "ledgLoc", "y2Label", "xTickMap", "yTickMap", "plots", "xLim", "yLim", "ledgLoc", "y2Label",
"ticklabel" "ticklabel", "xInvert", "yInvert"
] ]
#Feel like this could be optimized #Feel like this could be optimized
@ -258,15 +263,20 @@ def makeGraph(graphData, showPlot=True, doProgramBlock=True, figSavePath=None, h
for i, axGraphData in enumerate(graphData["subPlots"]): for i, axGraphData in enumerate(graphData["subPlots"]):
ax1 = flatAxes[i] ax1 = flatAxes[i]
if "skip" in axGraphData and axGraphData["skip"]:
ax1.set_axis_off()
continue
if bool(sum([("y2" in pData) for pData in axGraphData["plots"]])): if bool(sum([("y2" in pData) for pData in axGraphData["plots"]])):
ax2 = ax1.twinx() ax2 = ax1.twinx()
else: else:
ax2 = None ax2 = None
additional_legends = []
# Duct Tape # Duct Tape
ax = ax1 ax = ax1
#Draw many plots as needed #Draw many plots as needed
# Also provide functions for drawing other types of lines # Also provide functions for drawing other types of lines
if "plots" in axGraphData: if "plots" in axGraphData:
@ -276,6 +286,10 @@ def makeGraph(graphData, showPlot=True, doProgramBlock=True, figSavePath=None, h
getSafeColour = getSafeValue("colour") or getSafeValue("color") #Frigen American Spelling getSafeColour = getSafeValue("colour") or getSafeValue("color") #Frigen American Spelling
optArgs = getSafeValue("args", {}) #Allow for other args to be passed in optArgs = getSafeValue("args", {}) #Allow for other args to be passed in
currentLine = None
if "skip" in pData and pData["skip"]: continue
if "x" in pData: if "x" in pData:
xData = pData["x"] xData = pData["x"]
@ -286,22 +300,43 @@ def makeGraph(graphData, showPlot=True, doProgramBlock=True, figSavePath=None, h
ax = ax2 ax = ax2
if "type" not in pData or pData["type"] == "plot": if "type" not in pData or pData["type"] == "plot":
ax.plot(xData, yData, label=getSafeValue("label"), color=getSafeColour, **optArgs) currentLine = ax.plot(xData, yData, label=getSafeValue("label"), color=getSafeColour, **optArgs)
currentLineColour = currentLine.get_color()
if "maxPoint" in pData:
labelText = pData["maxPoint"] if type(pData["maxPoint"]) == str else "Maximum Point ({0:.2f})"
x, y = np.array(xData), np.array(yData)
maxPoint = np.max(x), y[np.argmax(x)]
ax.scatter(maxPoint[0], maxPoint[1],
marker=getSafeValue("maxMarker", "o"), label=labelText.format(*maxPoint),
color=currentLineColour, zorder=getSafeValue("maxZorder", 2))
if "minPoint" in pData:
labelText = pData["minPoint"] if type(pData["minPoint"]) == str else "Minimum Point ({0:.2f})"
x, y = np.array(xData), np.array(yData)
minPoint = np.min(x), y[np.argmin(x)]
ax.scatter(minPoint[0], minPoint[1],
marker=getSafeValue("minMarker", "*" if "maxPoint" in pData else "o"),
label=labelText.format(*minPoint), color=currentLineColour,
zorder=getSafeValue("minxZorder", 2))
elif pData["type"] == "point": elif pData["type"] == "point":
ax.scatter(xData, yData, ax.scatter(xData, yData,
marker=getSafeValue("marker"), label=getSafeValue("label"), marker=getSafeValue("marker"), label=getSafeValue("label"),
color=getSafeColour, zorder=getSafeValue("zorder", 2), color=getSafeColour, zorder=getSafeValue("zorder", 2),
**optArgs ) **optArgs)
elif pData["type"] == "hLine": elif pData["type"] == "hLine":
ax.hlines(yData, *xData, label=getSafeValue("label"), color=getSafeColour, **optArgs) currentLine = ax.hlines(yData, *xData, label=getSafeValue("label"), color=getSafeColour, **optArgs)
elif pData["type"] == "vLine": elif pData["type"] == "vLine":
ax.vlines(xData, *yData, label=getSafeValue("label"), color=getSafeColour, **optArgs) currentLine = ax.vlines(xData, *yData, label=getSafeValue("label"), color=getSafeColour, **optArgs)
elif pData["type"] == "axvLine": elif pData["type"] == "axvLine":
if "y" not in pData: yData = (0, 1) #Span the whole graph if "y" not in pData: yData = (0, 1) #Span the whole graph
ax.axvline(xData, *yData, label=getSafeValue("label"), color=getSafeColour, **optArgs) currentLine = ax.axvline(xData, *yData, label=getSafeValue("label"), color=getSafeColour, **optArgs)
elif pData["type"] == "axhLine": elif pData["type"] == "axhLine":
if "x" not in pData: xData = (0, 1) #Span the whole graph if "x" not in pData: xData = (0, 1) #Span the whole graph
ax.axhline(yData, *xData, label=getSafeValue("label"), color=getSafeColour, **optArgs) currentLine = ax.axhline(yData, *xData, label=getSafeValue("label"), color=getSafeColour, **optArgs)
elif pData["type"] == "scatter": elif pData["type"] == "scatter":
ax.scatter(xData, yData, marker=getSafeValue("marker"), label=getSafeValue("label"), color=getSafeColour, **optArgs) ax.scatter(xData, yData, marker=getSafeValue("marker"), label=getSafeValue("label"), color=getSafeColour, **optArgs)
@ -368,7 +403,30 @@ def makeGraph(graphData, showPlot=True, doProgramBlock=True, figSavePath=None, h
align = getSafeValue("align", align) align = getSafeValue("align", align)
ax.text(getSafeValue("x", 0.05), getSafeValue("y", 0.95), pData["text"], transform=ax.transAxes, fontsize=getSafeValue("fontsize", None), va=align[0], ha=align[1], bbox=props) ax.text(getSafeValue("x", 0.05), getSafeValue("y", 0.95), pData["text"], transform=ax.transAxes, fontsize=getSafeValue("fontsize", None), va=align[0], ha=align[1], bbox=props)
elif pData["type"] == "addLabel":
additional_legends.append(pData["line"])
if "fillAlpha" in pData or ("fill" in pData and ("type" not in pData or pData["type"] in ["plot"])):
currentLine = ax.get_lines()[-1] if currentLine is None else currentLine
line_colour = colors.to_rgba(currentLine.get_color())
fillData = pData["fill"] if type(pData["fill"]) == dict and "fill" in pData else {}
if "fillAlpha" not in fillData: fillData["fillAlpha"] = pData["fillAlpha"] if "fillAlpha" in pData else 0.07
line_colour = (line_colour[0], line_colour[1], line_colour[2], fillData["fillAlpha"])
if "color" not in fillData: fillData["color"] = line_colour
if "y2" not in fillData: fillData["y2"] = 0
if "where" not in fillData: fillData["where"] = None
if "interpolate" not in fillData: fillData["interpolate"] = False
if "step" not in fillData: fillData["step"] = None
if "data" not in fillData: fillData["data"] = None
if "args" not in fillData: fillData["args"] = {}
ax.fill_between(xData, yData, fillData["y2"], color=fillData["color"],
where=fillData["where"], interpolate=fillData["interpolate"],
step=fillData["step"], data=fillData["data"], **fillData["args"]
)
#Set extra options as needed #Set extra options as needed
ax = ax1 ax = ax1
@ -390,6 +448,8 @@ def makeGraph(graphData, showPlot=True, doProgramBlock=True, figSavePath=None, h
if "yTickPos" in axGraphData: ax.yaxis.set_ticks_position(axGraphData["yTickPos"]) # Add the ability to move the y axis ticks if "yTickPos" in axGraphData: ax.yaxis.set_ticks_position(axGraphData["yTickPos"]) # Add the ability to move the y axis ticks
if "xScale" in axGraphData: ax.set_xscale(axGraphData["xScale"]) #Add x axis scaling if needed if "xScale" in axGraphData: ax.set_xscale(axGraphData["xScale"]) #Add x axis scaling if needed
if "yScale" in axGraphData: ax.set_yscale(axGraphData["yScale"]) #Add y axis scaling if needed if "yScale" in axGraphData: ax.set_yscale(axGraphData["yScale"]) #Add y axis scaling if needed
if "xInvert" in axGraphData and axGraphData["xInvert"]: ax.invert_xaxis()
if "yInvert" in axGraphData and axGraphData["yInvert"]: ax.invert_yaxis()
if "xLim" in axGraphData: if "xLim" in axGraphData:
xLimit = () xLimit = ()
if type(axGraphData["xLim"]) in [int, float]: if type(axGraphData["xLim"]) in [int, float]:
@ -405,21 +465,24 @@ def makeGraph(graphData, showPlot=True, doProgramBlock=True, figSavePath=None, h
yLimit = axGraphData["yLim"] yLimit = axGraphData["yLim"]
ax.set_ylim(yLimit) ax.set_ylim(yLimit)
if "xTickMap" in axGraphData: #Allow for the mapping / transformation of the xAxis Ticks if "xTickMap" in axGraphData: #Allow for the mapping / transformation of the xAxis Ticks
xTicks = matplotlib.ticker.FuncFormatter(lambda x, pos: '{0:g}'.format(axGraphData["xTickMap"](x))) formatter = axGraphData["xTickMap"]
xTicks = matplotlib.ticker.FuncFormatter(formatter)
ax.xaxis.set_major_formatter(xTicks) ax.xaxis.set_major_formatter(xTicks)
if "yTickMap" in axGraphData: #Allow for the mapping / transformation of the yAxis Ticks if "yTickMap" in axGraphData: #Allow for the mapping / transformation of the yAxis Ticks
yTicks = matplotlib.ticker.FuncFormatter(lambda y, pos: '{0:g}'.format(axGraphData["yTickMap"](y))) formatter = axGraphData["yTickMap"]
yTicks = matplotlib.ticker.FuncFormatter(formatter)
ax.yaxis.set_major_formatter(yTicks) ax.yaxis.set_major_formatter(yTicks)
if "plots" in axGraphData and bool(sum([("label" in pData) for pData in axGraphData["plots"]])): if "plots" in axGraphData and bool(sum([("label" in pData) for pData in axGraphData["plots"]])):
locPoint = axGraphData["ledgLoc"] if "ledgLoc" in axGraphData else None locPoint = axGraphData["ledgLoc"] if "ledgLoc" in axGraphData else None
add_lines, add_labels = additional_legends, [line.get_label() for line in additional_legends]
lines1, labels1 = ax1.get_legend_handles_labels() lines1, labels1 = ax1.get_legend_handles_labels()
if ax2: if ax2:
lines2, labels2 = ax2.get_legend_handles_labels() lines2, labels2 = ax2.get_legend_handles_labels()
ax2.legend(lines1 + lines2, labels1 + labels2, loc=locPoint) ax2.legend(lines1 + lines2 + add_lines, labels1 + labels2 + add_labels, loc=locPoint)
else: else:
ax1.legend(lines1, labels1, loc=locPoint) ax1.legend(lines1 + add_lines, labels1 + add_labels, loc=locPoint)
if "ticklabel" in axGraphData: if "ticklabel" in axGraphData:
style = axGraphData["ticklabel"]["style"] if "style" in axGraphData["ticklabel"] else "" style = axGraphData["ticklabel"]["style"] if "style" in axGraphData["ticklabel"] else ""
@ -434,12 +497,14 @@ def makeGraph(graphData, showPlot=True, doProgramBlock=True, figSavePath=None, h
ax.set_axis_off() ax.set_axis_off()
if "title" in graphData and not "figTitle" in graphData: fig.canvas.manager.set_window_title(graphData["title"].replace("\n", " ")) #Set the figure title correctly if "title" in graphData or "winTitle" in graphData:
title = graphData["winTitle"] if "winTitle" in graphData else graphData["title"]
fig.canvas.manager.set_window_title(title.replace("\n", " ")) #Set the figure title correctly
if "figTitle" in graphData: if "figTitle" in graphData:
getSafeValue = lambda key: graphData[key] if key in graphData else None #Only return the key-value if present in graphData getSafeValue = lambda key: graphData[key] if key in graphData else None #Only return the key-value if present in graphData
fig.suptitle(graphData["figTitle"], fontsize=getSafeValue("figTitleFontSize")) fig.suptitle(graphData["figTitle"], fontsize=getSafeValue("figTitleFontSize"))
fig.canvas.manager.set_window_title(graphData["figTitle"].replace("\n", " ")) if "title" not in graphData: fig.canvas.manager.set_window_title(graphData["figTitle"].replace("\n", " "))
fig.tight_layout() #Fix labels being cut off sometimes fig.tight_layout() #Fix labels being cut off sometimes