1010
1111from mathics .core .evaluation import Evaluation
1212from mathics .core .symbols import strip_context
13+ from mathics .core .systemsymbols import SymbolNone , SymbolRGBColor
1314from mathics .timing import Timer
1415
1516from .plot_compile import plot_compile
@@ -26,55 +27,107 @@ def eval_Plot3D(
2627 # pull out plot options
2728 _ , xmin , xmax = plot_options .ranges [0 ]
2829 _ , ymin , ymax = plot_options .ranges [1 ]
29- nx , ny = plot_options .plotpoints
3030 names = [strip_context (str (range [0 ])) for range in plot_options .ranges ]
3131
32- # compute (nx, ny) grids of xs and ys for corresponding vertexes
33- xs = np .linspace (xmin , xmax , nx )
34- ys = np .linspace (ymin , ymax , ny )
35- xs , ys = np .meshgrid (xs , ys )
32+ # Mesh option
33+ nmesh = 20
34+ if plot_options .mesh is SymbolNone :
35+ nmesh = 0
36+
37+ # color-blind friendly palette from https://davidmathlogic.com/colorblind
38+ palette = [
39+ (255 , 176 , 0 ), # orange
40+ (100 , 143 , 255 ), # blue
41+ (220 , 38 , 127 ), # red
42+ (50 , 150 , 140 ), # green
43+ (120 , 94 , 240 ), # purple
44+ (254 , 97 , 0 ), # dark orange
45+ (0 , 114 , 178 ), # dark blue
46+ ]
47+
48+ # compile the functions
49+ with Timer ("compile" ):
50+ compiled_functions = [
51+ plot_compile (evaluation , function , names )
52+ for function in plot_options .functions
53+ ]
54+
55+ def compute_over_grid (nx , ny ):
56+ """
57+ For each function, computes an (nx*ny, 3) array of coordinates (xyzs),
58+ and an (nx, ny) array of indices (inxs) into xyzs representing
59+ the index in xyzs of the corresponding position in the grid.
60+ Returns an iterator over (xyzs,inxs) pairs, one for each function.
61+
62+ This is used for computing the full grid of quads representing the
63+ surface defined by each function, and also for computing a sparse
64+ grid used to display a mesh of lines on the surface.
65+ """
66+
67+ # compute (nx, ny) grids of xs and ys for corresponding vertexes
68+ xs = np .linspace (xmin , xmax , nx )
69+ ys = np .linspace (ymin , ymax , ny )
70+ xs , ys = np .meshgrid (xs , ys )
71+
72+ # (nx,ny) array of numbers from 0 to n-1 that are
73+ # indexes into xyzs array for corresponding vertex
74+ # +1 because these will be used as WL indexes, which are 1-based
75+ inxs = np .arange (math .prod (xs .shape )).reshape (xs .shape ) + 1
76+
77+ for function in compiled_functions :
78+ # compute zs from xs and ys using compiled function
79+ with Timer ("compute zs" ):
80+ zs = function (** {str (names [0 ]): xs , str (names [1 ]): ys })
81+
82+ # sometimes expr gets compiled into something that returns a complex
83+ # even though the imaginary part is 0
84+ # TODO: check that imag is all 0?
85+ # TODO: needed this for Hypergeometric - look into that
86+ # assert np.all(np.isreal(zs)), "array contains complex values"
87+ zs = np .real (zs )
88+
89+ # if it's a constant, make it a full array
90+ if isinstance (zs , (float , int , complex )):
91+ zs = np .full (xs .shape , zs )
3692
37- for function in plot_options .functions :
38- with Timer ("compile" ):
39- function = plot_compile (evaluation , function , names )
40-
41- # compute zs from xs and ys using compiled function
42- with Timer ("compute zs" ):
43- zs = function (** {str (names [0 ]): xs , str (names [1 ]): ys })
44-
45- # sometimes expr gets compiled into something that returns a complex
46- # even though the imaginary part is 0
47- # TODO: check that imag is all 0?
48- # TODO: needed this for Hypergeometric - look into that
49- # assert np.all(np.isreal(zs)), "array contains complex values"
50- zs = np .real (zs )
51-
52- # if it's a constant, make it a full array
53- if isinstance (zs , (float , int , complex )):
54- zs = np .full (xs .shape , zs )
55-
56- with Timer ("stack" ):
5793 # (nx*ny, 3) array of points, to be indexed by quads
5894 xyzs = np .stack ([xs , ys , zs ]).transpose (1 , 2 , 0 ).reshape (- 1 , 3 )
5995
60- # (nx,ny) array of numbers from 0 to n-1 that are
61- # indexes into xyzs array for corresponding vertex
62- inxs = np .arange (math .prod (xs .shape )).reshape (xs .shape )
63-
64- # shift inxs array four different ways and stack to form
65- # (4, nx-1, ny-1) array of quads represented as indexes into xyzs array
66- quads = np .stack (
67- [inxs [:- 1 , :- 1 ], inxs [:- 1 , 1 :], inxs [1 :, 1 :], inxs [1 :, :- 1 ]]
68- )
69-
70- # transpose and flatten to ((nx-1)*(ny-1), 4) array, suitable for use in GraphicsComplex
71- quads = quads .T .reshape (- 1 , 4 )
72-
73- # ugh - indexes in Polygon are 1-based
74- quads += 1
75-
76- # add a GraphicsComplex for this function
77- graphics .add_complex (xyzs , lines = None , polys = quads )
96+ yield xyzs , inxs
97+
98+ # generate the quads and emit a GraphicsComplex containing them
99+ for i , (xyzs , inxs ) in enumerate (compute_over_grid (* plot_options .plotpoints )):
100+ # shift inxs array four different ways and stack to form
101+ # (4, nx-1, ny-1) array of quads represented as indexes into xyzs array
102+ quads = np .stack ([inxs [:- 1 , :- 1 ], inxs [:- 1 , 1 :], inxs [1 :, 1 :], inxs [1 :, :- 1 ]])
103+
104+ # transpose and flatten to ((nx-1)*(ny-1), 4) array, suitable for use in GraphicsComplex
105+ quads = quads .T .reshape (- 1 , 4 )
106+
107+ # choose a color
108+ rgb = palette [i % len (palette )]
109+ rgb = [c / 255.0 for c in rgb ]
110+ # graphics.add_color(SymbolRGBColor, rgb)
111+ graphics .add_directives ([SymbolRGBColor , * rgb ])
112+
113+ # add a GraphicsComplex displaying a surface for this function
114+ graphics .add_complex (xyzs , lines = None , polys = quads )
115+
116+ # if requested by the Mesh attribute create a mesh of lines covering the surfaces
117+ if nmesh :
118+ # meshes are black for now
119+ graphics .add_directives ([SymbolRGBColor , 0 , 0 , 0 ])
120+
121+ with Timer ("Mesh" ):
122+ nx , ny = plot_options .plotpoints
123+ # Do nmesh lines in each direction, each line formed
124+ # from one row or one column of the inxs array.
125+ # Each mesh line has high res (nx or ny) so it follows
126+ # the contours of the surface.
127+ for xyzs , inxs in compute_over_grid (nx , nmesh ):
128+ graphics .add_complex (xyzs , lines = inxs , polys = None )
129+ for xyzs , inxs in compute_over_grid (nmesh , ny ):
130+ graphics .add_complex (xyzs , lines = inxs .T , polys = None )
78131
79132 return graphics
80133
0 commit comments