Merge pull request #142 from bbean23/141-add-arrow-marker-type

141 add arrow marker type

opencsp/common/lib/render/test/test_View3d.py

@@ -0,0 +1,112 @@
+import unittest
+
+import opencsp.common.lib.tool.file_tools as ft
+
+import numpy as np
+from PIL import Image
+
+import opencsp.common.lib.render.figure_management as fm
+import opencsp.common.lib.render.view_spec as vs
+import opencsp.common.lib.render_control.RenderControlAxis as rca
+import opencsp.common.lib.render_control.RenderControlFigure as rcfg
+import opencsp.common.lib.render_control.RenderControlFigureRecord as rcfr
+import opencsp.common.lib.render_control.RenderControlPointSeq as rcps
+import opencsp.common.lib.tool.file_tools as ft
+
+
+class test_View3d(unittest.TestCase):
+    @classmethod
+    def setUpClass(cls) -> None:
+        path, name, _ = ft.path_components(__file__)
+        cls.in_dir = ft.join(path, 'data/input', name.split('test_')[-1])
+        cls.out_dir = ft.join(path, 'data/output', name.split('test_')[-1])
+        ft.create_directories_if_necessary(cls.out_dir)
+        ft.delete_files_in_directory(cls.out_dir, '*')
+        return super().setUpClass()
+
+    def setUp(self) -> None:
+        self.test_name = self.id().split('.')[-1]
+
+    def setup_figure(self) -> rcfr.RenderControlFigureRecord:
+        # clear existing figures
+        fm.reset_figure_management()
+
+        # setup the new figure
+        axis_control = rca.meters(grid=False)
+        figure_control = rcfg.RenderControlFigure()
+        view_spec_2d = vs.view_spec_xy()
+        fig_record = fm.setup_figure(
+            figure_control,
+            axis_control,
+            view_spec_2d,
+            title=self.test_name,
+            code_tag=f"{__file__}.{self.test_name}()",
+            equal=False,
+        )
+
+        return fig_record
+
+    def test_plot_arrows(self):
+        fig_record = self.setup_figure()
+
+        # draw
+        square_corners = [(0, 0), (1, 0), (1, 1), (0, 1)]
+        arrow_style = rcps.RenderControlPointSeq(marker='arrow', markersize=0.1)
+        fig_record.view.draw_pq_list(square_corners, close=True, style=arrow_style)
+        fig_record.view.show(equal=True, block=False)
+        actual = fig_record.to_array()
+        fig_record.close()
+
+        # load and compare
+        expected = np.array(Image.open(ft.join(self.in_dir, f"{self.test_name}.png")))
+        np.testing.assert_array_equal(expected, actual)
+
+        # save
+        img = Image.fromarray(actual)
+        img.save(ft.join(self.out_dir, f"{self.test_name}.png"))
+
+    def test_draw_xyz(self):
+        """Verify that the various accepted input arguments produce the same output."""
+        style = rcps.RenderControlPointSeq('None', marker='.')
+
+        # with single points
+        fig_record = self.setup_figure()
+        fig_record.view.draw_xyz((1, 1), style=style)
+        fig_record.view.draw_xyz((2, 2), style=style)
+        fig_record.view.draw_xyz((2, 1), style=style)
+        fig_record.view.draw_xyz((1, 2), style=style)
+        fig_record.view.show(equal=True)
+        image_single_points = fig_record.to_array()
+        fig_record.close()
+
+        # with lists
+        x = [1, 2, 2, 1]
+        y = [1, 1, 2, 2]
+        fig_record = self.setup_figure()
+        fig_record.view.draw_xyz((x, y), style=style)
+        fig_record.view.show(equal=True)
+        image_xy_lists = fig_record.to_array()
+        fig_record.close()
+
+        # with a numpy array
+        arr = np.array([[1, 1], [1, 2], [2, 2], [2, 1]])
+        fig_record = self.setup_figure()
+        fig_record.view.draw_xyz(arr, style=style)
+        fig_record.view.show(equal=True)
+        image_arr = fig_record.to_array()
+        fig_record.close()
+
+        # save the output
+        Image.fromarray(image_single_points).save(ft.join(self.out_dir, f"{self.test_name}_single_points.png"))
+        Image.fromarray(image_xy_lists).save(ft.join(self.out_dir, f"{self.test_name}_xy_lists.png"))
+        Image.fromarray(image_arr).save(ft.join(self.out_dir, f"{self.test_name}_arr.png"))
+
+        # load and compare
+        expected = np.array(Image.open(ft.join(self.in_dir, f"{self.test_name}.png")))
+        np.testing.assert_array_equal(expected, image_single_points)
+        np.testing.assert_array_equal(expected, image_xy_lists)
+        np.testing.assert_array_equal(expected, image_arr)
+
+
+if __name__ == '__main__':
+    unittest.main()

opencsp/common/lib/render/test/test_figure_management.py

@@ -31,6 +31,9 @@ def setUpClass(cls) -> None:
     def setUp(self) -> None:
         self.test_name = self.id().split('.')[-1]
 
+    def setUp(self) -> None:
+        self.test_name = self.id().split('.')[-1]
+
     def tearDown(self):
         # Make sure we release all matplotlib resources.
         plt.close('all')

opencsp/common/lib/render_control/RenderControlFigureRecord.py

@@ -4,10 +4,11 @@
 """
 
 import os
+
 import matplotlib.pyplot as plt
 from matplotlib.figure import Figure
+import numpy as np
 import PIL.Image as PilImage
-from typing import Optional
 
 from opencsp.common.lib.render.View3d import View3d
 import opencsp.common.lib.tool.file_tools as ft
@@ -112,6 +113,17 @@ def print_comments(self):
         for comment_line in self.comments:
             lt.info(comment_line)
 
+    def to_array(self):
+        return self.figure_to_array(self.figure)
+
+    @staticmethod
+    def figure_to_array(figure: Figure):
+        # Force matplotlib to render to it's internal buffer
+        figure.canvas.draw()
+
+        # Convert the buffer to a numpy array
+        return np.asarray(figure.canvas.buffer_rgba())
+
     def save(self, output_dir: str, output_file_body: str = None, format: str = None, dpi=600, close_after_save=True):
         """Saves this figure record to an image file.
 
@@ -160,7 +172,7 @@ def save(self, output_dir: str, output_file_body: str = None, format: str = None
                 #     print('Skipping save of existing figure: ' + output_figure_dir_body_ext)
                 # else:
                 output_figure_dir_body_ext = output_figure_dir_body + '.' + format
-                lt.info('Saving figure: ' + output_figure_dir_body_ext)
+                lt.info('In RenderControlFigureRecord.save(), saving figure: ' + output_figure_dir_body_ext)
                 plt.savefig(output_figure_dir_body_ext, format=format, dpi=dpi)
         finally:
             # Close figure if desired.

opencsp/common/lib/render_control/RenderControlPointSeq.py

@@ -3,6 +3,8 @@
 
 """
 
+import opencsp.common.lib.render.color as clr
+
 
 class RenderControlPointSeq:
     """
@@ -37,52 +39,97 @@ class RenderControlPointSeq:
 
         Markers
         -------
-        '.' 	point marker
-        ',' 	pixel marker
-        'o' 	circle marker
-        'v' 	triangle_down marker
-        '^' 	triangle_up marker
-        '<' 	triangle_left marker
-        '>' 	triangle_right marker
-        '1' 	tri_down marker
-        '2' 	tri_up marker
-        '3' 	tri_left marker
-        '4' 	tri_right marker
-        '8' 	octagon marker
-        's' 	square marker
-        'p' 	pentagon marker
-        'P' 	plus (filled) marker
-        '*' 	star marker
-        'h' 	hexagon1 marker
-        'H' 	hexagon2 marker
-        '+' 	plus marker
-        'x' 	x marker
-        'X' 	x (filled) marker
-        'D' 	diamond marker
-        'd' 	thin_diamond marker
-        '|' 	vline marker
-        '_' 	hline marker
-        'None'  no marker
+        '.' 	   point marker
+        ',' 	   pixel marker
+        'o' 	   circle marker
+        'v' 	   triangle_down marker
+        '^' 	   triangle_up marker
+        '<' 	   triangle_left marker
+        '>' 	   triangle_right marker
+        '1' 	   tri_down marker (three lines from the center to points on 30, 150, and 270 degrees)
+        '2' 	   tri_up marker (three lines from the center to points on 90, 210, and 330 degrees)
+        '3' 	   tri_left marker (three lines from the center to points on 60, 180, and 300 degrees)
+        '4' 	   tri_right marker (three lines from the center to points on 0, 120, and 240 degrees)
+        '8' 	   octagon marker
+        's' 	   square marker
+        'p' 	   pentagon marker
+        'P' 	   plus (filled) marker
+        '*' 	   star marker
+        'h' 	   hexagon1 marker
+        'H' 	   hexagon2 marker
+        '+' 	   plus marker
+        'x' 	   x marker
+        'X' 	   x (filled) marker
+        'D' 	   diamond marker
+        'd' 	   thin_diamond marker
+        '|' 	   vline marker
+        '_' 	   hline marker
+        'None'     no marker
+        '$\u266B$' two quarter notes
+        'arrow'    draws an arrow at the end of every line
+
+        For more markers, see:
+         https://matplotlib.org/stable/gallery/lines_bars_and_markers/marker_reference.html
 
     """
 
     def __init__(
         self,  # See above for details:
         linestyle='-',  # '-', '--', '-.', ':', '' or 'None'
-        linewidth=1,  # float
-        color='b',  #    bgrcmykw
-        marker='x',  #    .,ov^<>12348sp*hH+xXDd|_ or None
-        markersize=6,  # float
-        markeredgecolor=None,  # Defaults to color above if not set.
+        linewidth: float = 1,  # float
+        color: str | clr.Color = 'b',  # line color
+        marker='x',  # .,ov^<>12348sp*hH+xXDd|_ or None
+        markersize: float = 6,  # float
+        markeredgecolor: str | clr.Color = None,  # Defaults to color above if not set.
         markeredgewidth=None,  # Defaults to linewidth if not set.
-        markerfacecolor=None,  # Defaults to color above if not set.
-        vector_color='b',  # Used if points are in a vector field.
-        vector_linewidth=1,  # Used if points are in a vector field.
-        vector_scale=1.0,  # Facter to grow/srhink vector length, for points in a vector field.
+        markerfacecolor: str | clr.Color = None,  # Defaults to color above if not set.
+        vector_color: str | clr.Color = 'b',  # Used if points are in a vector field.
+        vector_linewidth: float = 1,  # Used if points are in a vector field.
+        vector_scale: float = 1.0,  # Facter to grow/srhink vector length, for points in a vector field.
     ):
-
+        """
+        linestyle : str
+            Determines how lines are drawn. One of '-', '--', '-.', ':', '' or
+            'None'. Default is '-' (solid line).
+        linewidth : float
+            Width of lines in the number of pixels. Default is 1.
+        color : str | Color
+            The primary color use for everything that doesn't have a color
+            specified. If a Color object then the rgb() value of the Color
+            object will be used. Default is 'b'.
+        marker : str | None, optional
+            The style of marker to use. See the class description for more
+            information. Default is point '.'.
+        markersize : float, optional
+            Size of the marker, in pixels. Default is 6.
+        markeredgecolor : str | Color | None, optional
+            The color of the marker edges. Default is 'color'.
+        markeredgewidth : float | None, optional
+            Width of the marker edge in pixels. Defaults is 'linewidth'.
+        markerfacecolor : str | Color | None, optional
+            The color of the marker faces. Default is 'color'.
+        vector_color : str | Color | None, optional
+            The color for vectors. Only applies to points in a vector field.
+            Default is 'b'.
+        vector_linewidth : float, optional
+            The line width for vectors, in pixels. Only applies to points in a
+            vector field. default is 1.
+        vector_scale : float, optional
+            Facter to grow/srhink vector length. Only applies to points in a
+            vector field. Default is 1.
+        """
         super(RenderControlPointSeq, self).__init__()
 
+        # Convert color values
+        if isinstance(color, clr.Color):
+            color = color.rgb()
+        if isinstance(markeredgecolor, clr.Color):
+            markeredgecolor = markeredgecolor.rgb()
+        if isinstance(markerfacecolor, clr.Color):
+            markerfacecolor = markerfacecolor.rgb()
+        if isinstance(vector_color, clr.Color):
+            vector_color = vector_color.rgb()
+
         # Set defaults.
         if markeredgecolor == None:
             markeredgecolor = color
@@ -115,11 +162,11 @@ def set_color(self, color):
 # COMMON CASES
 
 
-def default(marker='o', color='b', linewidth=1, markersize=8):
+def default(marker='.', color='b', linewidth=1, markersize=8):
     """
     What to draw if no particular preference is expressed.
     """
-    return RenderControlPointSeq(linestyle='-', linewidth=1, color=color, marker='.', markersize=markersize)
+    return RenderControlPointSeq(linestyle='-', linewidth=linewidth, color=color, marker=marker, markersize=markersize)
 
 
 def outline(color='k', linewidth=1):