Generating data for CSV stage or function

 1"""This script shows how to generate a yaml config containing a CSV stage and
 2a normal stage with a CSV function. It also shows how to generate the CSV
 3data that can be used by the stage and function.
 4
 5It generates three files:
 6- csv_experiment_config.yml containing the yaml config to be imported
 7- experiment_stage.csv containing the csv data that needs to be selected in the
 8  ``csv_stage`` stage's settings in the GUI
 9- experiment_function.csv containing the csv data that needs to be selected in
10  the csv_func_stage stage's function settings in the GUI
11
12Either the csv_stage or the csv_func_stage can then be run and the values
13will be read from the appropriate csv file.
14
15See the docs for the CSVStage and CSVFunc for details on the format of the CSV
16files.
17"""
18from ceed.shape import CeedPaintCanvasBehavior, CeedPaintCircle, \
19    CeedPaintEllipse
20from ceed.function import FunctionFactoryBase, register_all_functions
21from ceed.stage import StageFactoryBase, register_all_stages
22from ceed.storage.controller import CeedDataWriterBase
23import csv
24import math
25
26# create shape/function/stage factory
27shape_factory = CeedPaintCanvasBehavior()
28function_factory = FunctionFactoryBase()
29stage_factory = StageFactoryBase(
30    function_factory=function_factory, shape_factory=shape_factory)
31
32# now we need to register all the builtin functions and stage plugins
33register_all_functions(function_factory)
34register_all_stages(stage_factory)
35
36# create the shapes and name them
37ellipse = CeedPaintEllipse.create_shape(
38    center=(250, 450), radius_x=200, radius_y=400, name='ellipse')
39circle = CeedPaintCircle.create_shape(
40    center=(700, 300), radius=200, name='circle')
41
42# add the shapes to factory
43shape_factory.add_shape(ellipse)
44shape_factory.add_shape(circle)
45
46
47def create_csv_stage():
48    # create a csv stage that will have two shapes and accept a csv file
49    # that is generated below
50    csv_stage = stage_factory.get('CSVStage')(
51        function_factory=function_factory, stage_factory=stage_factory,
52        shape_factory=shape_factory, name='csv_stage', color_r=True,
53        color_g=True, color_b=True)
54    stage_factory.add_stage(csv_stage)
55
56    # generate the csv file for the stage, ellipse is going to be red, circle
57    # yellow
58    with open('experiment_stage.csv', 'w', newline='') as fp:
59        writer = csv.writer(fp)
60        writer.writerow(['ellipse_red', 'circle_red', 'circle_green'])
61        for i in range(1000):
62            val = .5 * math.sin(i / 120 * 2 * math.pi) + .5
63            writer.writerow([val, val, val])
64
65
66def create_csv_function_stage():
67    # create a stage that will have two shapes and one csv function. The csv
68    # function will accept a csv file that is generated below
69    func = function_factory.get('CSVFunc')(
70        function_factory=function_factory, name='csv')
71
72    stage = stage_factory.get('CeedStage')(
73        function_factory=function_factory, stage_factory=stage_factory,
74        shape_factory=shape_factory, name='csv_func_stage', color_b=True)
75    stage_factory.add_stage(stage)
76    stage.add_func(func)
77    stage.add_shape(circle)
78    stage.add_shape(ellipse)
79
80    # generate the csv file for the function; ellipse and circle are going to
81    # be red
82    with open('experiment_function.csv', 'w', newline='') as fp:
83        writer = csv.writer(fp)
84        for i in range(1000):
85            writer.writerow([.5 * math.sin(i / 120 * 2 * math.pi) + .5])
86
87
88# create the stages/functions and csv files
89create_csv_stage()
90create_csv_function_stage()
91
92# save it to disk for import later
93CeedDataWriterBase.save_config_to_yaml(
94    'csv_experiment_config.yml', shape_factory=shape_factory,
95    function_factory=function_factory, stage_factory=stage_factory,
96    overwrite=True)