diff --git a/Raytrace/PlotRays.py b/Raytrace/PlotRays.py new file mode 100644 index 0000000..bb7e01a --- /dev/null +++ b/Raytrace/PlotRays.py @@ -0,0 +1,59 @@ +from Raytrace.TriangleMesh import TriangleMesh +from Raytrace.SideScan import ScanReading +import numpy as np +try: + import plotly.graph_objs as go # type: ignore +except ModuleNotFoundError: + print('plotly not installed, please run "pip install plotly"') + exit() + +def plot_mesh(mesh:TriangleMesh, **kwargs) -> go.Mesh3d: + x, y, z = mesh.triangles.reshape((-1,3)).swapaxes(0, 1) + i, j, k = [list(range(i,x.shape[0],3)) for i in range(3)] + return go.Mesh3d(x=x, y=y, z=z, i=i, j=j, k=k, **kwargs) + +def plot_rays(reading:ScanReading, **kwargs) -> go.Scatter3d: + origin = reading.origins[reading.finite] + inters = reading.intersections[reading.finite] + + x, y, z = np.stack((origin, inters, origin)).swapaxes(0,1).reshape((-1,3)).swapaxes(0,1) + return go.Scatter3d(x=x, y=y, z=z, **kwargs) + +def plot_reading(reading:ScanReading) -> go.Figure: + return go.Figure(data = [ + go.Scatter({ + 'name': 'Smoothed', + 'showlegend': True, + 'mode': 'lines', + 'x': np.arange(reading.result_reselution)/reading.result_reselution*(reading.max_dist-reading.min_dist)+reading.min_dist, + 'y': reading.result, + }), + go.Scatter({ + 'mode': 'markers', + 'name': 'Raw', + 'showlegend': True, + 'x': reading.distances[reading.finite], + 'y': np.zeros(reading.intersection_count), + }) + ]) +def plot_intersections(reading:ScanReading, **kwargs) -> go.Scatter3d: + inters = reading.intersections[reading.finite] + + x, y, z = inters.reshape((-1,3)).swapaxes(0,1) + return go.Scatter3d(x=x, y=y, z=z, **kwargs) + +def plot_intersections_list(readings:list[ScanReading], **kwargs) -> list[go.Scatter3d]: + return [plot_intersections(reading, **kwargs) for reading in readings] + +def plot_readings_heatmap(readings) -> go.Figure: + return go.Figure(data = [ + go.Heatmap({ + 'y': np.arange(readings[0].min_dist, readings[0].max_dist, (readings[0].max_dist - readings[0].min_dist) / readings[0].result_reselution), + 'z': np.array([reading.result for reading in readings]).swapaxes(0,1), + }) + ], + layout = { + 'xaxis_title' : 'Reading Index', + 'yaxis_title' : 'Distance (units)', + } + ) \ No newline at end of file