From fb1b40503ad24d6e3786bf2bc5cf1029fd0dba51 Mon Sep 17 00:00:00 2001 From: Cal Wing <20716204+calw20@users.noreply.github.com> Date: Tue, 31 Oct 2023 23:05:24 +1000 Subject: [PATCH] So things to main and update makeGraph --- main.py | 34 ++++++++++--------- makeGraph.py | 95 +++++++++++++++++++++++++++++++++++++++++++--------- 2 files changed, 99 insertions(+), 30 deletions(-) diff --git a/main.py b/main.py index 9253fcf..c944c67 100644 --- a/main.py +++ b/main.py @@ -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_Z (in m/s^2) - the specific force in the body-fixed z-direction IMU_TIME_HEADER = ["Time [s]"] -IMU_WBE_HEADERS = ["WBE_1 [rad/s]", "WBE_2 [rad/s]", "WBE_3 [rad/s]"] -IMU_FSP_HEADERS = ["FSP_X [m/s^2]", "FSP_Y [m/s^2]", "FSP_Z [m/s^2]"] +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]"] # Specific Force X, Y , Z IMU_DATA_HEADER = IMU_TIME_HEADER + IMU_WBE_HEADERS + IMU_FSP_HEADERS def importIMUData(mission, imu): # 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 INIT_EULER_ANGLES = (0, 0, 0) -def translate2NED(angles, euler_angles): +def attidude_rate_2NED(angles, euler_angles): phi, theta, psi = euler_angles p, q, r = angles @@ -113,14 +113,14 @@ def translate2NED(angles, euler_angles): return np.matmul(transMat, angleMat) 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(theta), cos(theta)*sin(phi), cos(theta)*cos(phi) ] ]) - return forceMat + return R41 # Calculate Mission Translation & Force Data 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") translatedData = pd.DataFrame(columns=TRANS_DATA_HEADER) for i in tqdm(range(len(missionWBEData))): - dataPoint = missionWBEData[IMU_WBE_HEADERS].iloc[i] - trans = translate2NED(dataPoint.values, INIT_EULER_ANGLES).flatten() + dataPoint = missionWBEData[IMU_WBE_HEADERS].iloc[i] # Get the time point data + trans = attidude_rate_2NED(dataPoint.values, INIT_EULER_ANGLES).flatten() # Translate to Net forces = getNEDForces(trans) 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] 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[1]: trans[0], TRANS_DATA_HEADER[2]: trans[1], @@ -143,23 +143,27 @@ def calculateTranslatedData(missionWBEData) -> pd.DataFrame: TRANS_DATA_HEADER[5]: forceSum, TRANS_DATA_HEADER[6]: forceCumSum, } - translatedData.loc[i] = dataRow - + return translatedData if __name__ == '__main__': missionWBEData = m1_IMUData[0][IMU_TIME_HEADER + IMU_WBE_HEADERS] 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 fig = plt.figure() ax = plt.axes(projection='3d') plots = translatedData[TRANS_DATA_HEADER[6]].values - x = [p[0] for p in plots] - y = [p[1] for p in plots] - z = [p[2] for p in plots] + #x = [p[0] for p in plots] + #y = [p[1] for p in plots] + #z = [p[2] for p in plots] ax.plot3D(x, y, z) - + ax.plot3D(x+5, y-2, z+1) + plt.show() diff --git a/makeGraph.py b/makeGraph.py index 1a843c8..c0085e1 100644 --- a/makeGraph.py +++ b/makeGraph.py @@ -6,7 +6,7 @@ #### 2023 - Added UQ Default Colours to MatPlotLib __author__ = "Cal Wing" -__version__ = "0.1.10" +__version__ = "0.1.12" from collections.abc import Iterator import numpy as np @@ -15,6 +15,7 @@ import matplotlib.pyplot as plt import matplotlib.colors as colors from mpl_toolkits.axes_grid1 import make_axes_locatable from cycler import cycler +from tqdm import tqdm # Define the UQ Colours UQ_COLOURS_DICT = { @@ -62,6 +63,10 @@ class ColourValue(str): def hsv(self) -> np.ndarray: 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 class ColourList(object): @@ -243,7 +248,7 @@ def makeGraph(graphData, showPlot=True, doProgramBlock=True, figSavePath=None, h "xLabel", "yLabel", "title", "axis", "grid", "xPos", "yPos", "xLabelPos", "yLabelPos", "xTickPos", "yTickPos", "xScale", "yScale", "xTickMap", "yTickMap", "plots", "xLim", "yLim", "ledgLoc", "y2Label", - "ticklabel" + "ticklabel", "xInvert", "yInvert" ] #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"]): 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"]])): ax2 = ax1.twinx() else: ax2 = None + additional_legends = [] + # Duct Tape ax = ax1 - + #Draw many plots as needed # Also provide functions for drawing other types of lines 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 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: xData = pData["x"] @@ -286,22 +300,43 @@ def makeGraph(graphData, showPlot=True, doProgramBlock=True, figSavePath=None, h ax = ax2 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": ax.scatter(xData, yData, marker=getSafeValue("marker"), label=getSafeValue("label"), color=getSafeColour, zorder=getSafeValue("zorder", 2), - **optArgs ) + **optArgs) 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": - 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": 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": 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": 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) 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 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 "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 "xInvert" in axGraphData and axGraphData["xInvert"]: ax.invert_xaxis() + if "yInvert" in axGraphData and axGraphData["yInvert"]: ax.invert_yaxis() if "xLim" in axGraphData: xLimit = () if type(axGraphData["xLim"]) in [int, float]: @@ -405,21 +465,24 @@ def makeGraph(graphData, showPlot=True, doProgramBlock=True, figSavePath=None, h yLimit = axGraphData["yLim"] ax.set_ylim(yLimit) 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) 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) if "plots" in axGraphData and bool(sum([("label" in pData) for pData in axGraphData["plots"]])): 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() if ax2: 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: - ax1.legend(lines1, labels1, loc=locPoint) + ax1.legend(lines1 + add_lines, labels1 + add_labels, loc=locPoint) if "ticklabel" in axGraphData: 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() - 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: 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.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