1#lang typed/racket/base
2
3(require typed/racket/class racket/match racket/list math/flonum racket/math
4         (only-in typed/pict pict)
5         plot/utils
6         "../common/type-doc.rkt"
7         "../common/utils.rkt")
8
9(provide (all-defined-out))
10
11;; ===================================================================================================
12;; Surfaces of constant value (isosurfaces)
13
14(: isosurface3d-render-proc (-> 3D-Sampler Real Positive-Integer
15                                Plot-Color Plot-Brush-Style
16                                Plot-Color Nonnegative-Real Plot-Pen-Style
17                                Nonnegative-Real
18                                3D-Render-Proc))
19(define ((isosurface3d-render-proc
20          f d samples color style line-color line-width line-style alpha)
21         area)
22  (match-define (vector x-ivl y-ivl z-ivl) (send area get-bounds-rect))
23  (match-define (ivl x-min x-max) x-ivl)
24  (match-define (ivl y-min y-max) y-ivl)
25  (match-define (ivl z-min z-max) z-ivl)
26  (define num (animated-samples samples))
27  (define sample (f (vector x-ivl y-ivl z-ivl) (vector num num num)))
28  (match-define (3d-sample xs ys zs dsss d-min d-max) sample)
29
30  (send area put-alpha alpha)
31  (send area put-brush color style)
32  (send area put-pen line-color line-width line-style)
33  (for-3d-sample
34   (xa xb ya yb za zb d1 d2 d3 d4 d5 d6 d7 d8) sample
35   (for ([vs  (in-list (heights->cube-polys xa xb ya yb za zb d d1 d2 d3 d4 d5 d6 d7 d8))])
36     (send area put-polygon vs)))
37  (void))
38
39(:: isosurface3d
40    (->* [(-> Real Real Real Real) Real]
41         [(U Real #f) (U Real #f)
42          (U Real #f) (U Real #f)
43          (U Real #f) (U Real #f)
44          #:samples Positive-Integer
45          #:color Plot-Color
46          #:style Plot-Brush-Style
47          #:line-color Plot-Color
48          #:line-width Nonnegative-Real
49          #:line-style Plot-Pen-Style
50          #:alpha Nonnegative-Real
51          #:label (U String pict #f)]
52         renderer3d))
53(define (isosurface3d f d [x-min #f] [x-max #f] [y-min #f] [y-max #f] [z-min #f] [z-max #f]
54                      #:samples [samples (plot3d-samples)]
55                      #:color [color (surface-color)]
56                      #:style [style (surface-style)]
57                      #:line-color [line-color (surface-line-color)]
58                      #:line-width [line-width (surface-line-width)]
59                      #:line-style [line-style (surface-line-style)]
60                      #:alpha [alpha (surface-alpha)]
61                      #:label [label #f])
62  (define fail/pos (make-raise-argument-error 'isosurface3d f d x-min x-max y-min y-max z-min z-max))
63  (define fail/kw (make-raise-keyword-error 'isosurface3d))
64  (cond
65    [(not (rational? d))  (fail/pos "rational?" 1)]
66    [(and x-min (not (rational? x-min)))  (fail/pos "#f or rational" 2)]
67    [(and x-max (not (rational? x-max)))  (fail/pos "#f or rational" 3)]
68    [(and y-min (not (rational? y-min)))  (fail/pos "#f or rational" 4)]
69    [(and y-max (not (rational? y-max)))  (fail/pos "#f or rational" 5)]
70    [(and z-min (not (rational? z-min)))  (fail/pos "#f or rational" 6)]
71    [(and z-max (not (rational? z-max)))  (fail/pos "#f or rational" 7)]
72    [(< samples 2)  (fail/kw "Integer >= 2" '#:samples samples)]
73    [(not (rational? line-width))  (fail/kw "rational?" '#:line-width line-width)]
74    [(or (> alpha 1) (not (rational? alpha)))  (fail/kw "real in [0,1]" '#:alpha alpha)]
75    [else
76     (define x-ivl (ivl x-min x-max))
77     (define y-ivl (ivl y-min y-max))
78     (define z-ivl (ivl z-min z-max))
79     (define g (3d-function->sampler f (vector x-ivl y-ivl z-ivl)))
80     (renderer3d (vector x-ivl y-ivl z-ivl) #f default-ticks-fun
81                 (and label (λ (_) (rectangle-legend-entry
82                                    label color style line-color line-width line-style)))
83                 (isosurface3d-render-proc
84                  g d samples color style line-color line-width line-style alpha))]))
85
86;; ===================================================================================================
87;; Nested isosurfaces
88
89(: make-isosurfaces3d-label-and-renderer
90   (-> 3D-Sampler (U Real #f) (U Real #f) Contour-Levels Positive-Integer
91       (Plot-Colors (Listof Real)) (Plot-Brush-Styles (Listof Real))
92       (Plot-Colors (Listof Real)) (Pen-Widths (Listof Real)) (Plot-Pen-Styles (Listof Real))
93       (Alphas (Listof Real))
94       (U String pict #f)
95       (Values (U #f (-> Rect (Treeof legend-entry)))
96               3D-Render-Proc)))
97(define (make-isosurfaces3d-label-and-renderer f rd-min rd-max levels samples colors styles
98                                    line-colors line-widths line-styles alphas label)
99  ;; g is a 3D-sampler, which is a memoized proc. Recalculation here should be cheap.
100  ;; if we memoize here based on rect, smooth interactions are disabled (because of the call to `plot-z-ticks`)
101  (define (calculate-ds/labels [rect : Rect]) : (Values 3d-sample (Listof Real) (Listof String))
102    (match-define (vector x-ivl y-ivl z-ivl) rect)
103    (match-define (ivl x-min x-max) x-ivl)
104    (match-define (ivl y-min y-max) y-ivl)
105    (match-define (ivl z-min z-max) z-ivl)
106    (define num (animated-samples samples))
107    (define sample (f (vector x-ivl y-ivl z-ivl) (vector num num num)))
108    (match-define (3d-sample xs ys zs dsss fd-min fd-max) sample)
109
110    (define d-min (if rd-min rd-min fd-min))
111    (define d-max (if rd-max rd-max fd-max))
112
113    (match-define (list (tick #{ds : (Listof Real)}
114                              #{_ : (Listof Bolean)}
115                              #{labels : (Listof String)})
116                        ...)
117      (cond [(and d-min d-max)  (contour-ticks (plot-d-ticks) d-min d-max levels #f)]
118            [else  empty]))
119
120    (values sample ds labels))
121
122  (define label-proc
123    (and label
124         (λ ([rect : Rect])
125           (define-values (_ ds labels) (calculate-ds/labels rect))
126           (cond
127             [(empty? ds) empty]
128             [else (rectangle-legend-entries
129                    label ds colors styles line-colors line-widths line-styles)]))))
130
131  (: render-proc 3D-Render-Proc)
132  (define (render-proc area)
133    (define-values (sample ds _) (calculate-ds/labels (send area get-bounds-rect)))
134
135    (let* ([colors  (generate-list colors ds)]
136           [styles  (generate-list styles ds)]
137           [alphas  (generate-list alphas ds)]
138           [line-colors  (generate-list line-colors ds)]
139           [line-widths  (generate-list line-widths ds)]
140           [line-styles  (generate-list line-styles ds)])
141      (for ([d      (in-list ds)]
142            [color  (in-cycle* colors)]
143            [style  (in-cycle* styles)]
144            [alpha : Nonnegative-Real  (in-cycle* alphas)]
145            [line-color  (in-cycle* line-colors)]
146            [line-width : Nonnegative-Real  (in-cycle* line-widths)]
147            [line-style  (in-cycle* line-styles)])
148        (send area put-alpha alpha)
149        (send area put-brush color style)
150        (send area put-pen line-color line-width line-style)
151        (for-3d-sample
152         (xa xb ya yb za zb d1 d2 d3 d4 d5 d6 d7 d8) sample
153         (for ([vs  (in-list (heights->cube-polys xa xb ya yb za zb d d1 d2 d3 d4 d5 d6 d7 d8))])
154           (send area put-polygon vs))))))
155
156  (values label-proc render-proc))
157
158(:: isosurfaces3d
159    (->* [(-> Real Real Real Real)]
160         [(U Real #f) (U Real #f)
161          (U Real #f) (U Real #f)
162          (U Real #f) (U Real #f)
163          #:d-min (U Real #f) #:d-max (U Real #f)
164          #:samples Positive-Integer
165          #:levels Contour-Levels
166          #:colors (Plot-Colors (Listof Real))
167          #:styles (Plot-Brush-Styles (Listof Real))
168          #:line-colors (Plot-Colors (Listof Real))
169          #:line-widths (Pen-Widths (Listof Real))
170          #:line-styles (Plot-Pen-Styles (Listof Real))
171          #:alphas (Alphas (Listof Real))
172          #:label (U String pict #f)]
173         renderer3d))
174(define (isosurfaces3d f [x-min #f] [x-max #f] [y-min #f] [y-max #f] [z-min #f] [z-max #f]
175                       #:d-min [d-min #f] #:d-max [d-max #f]
176                       #:samples [samples (plot3d-samples)]
177                       #:levels [levels (isosurface-levels)]
178                       #:colors [colors (isosurface-colors)]
179                       #:styles [styles (isosurface-styles)]
180                       #:line-colors [line-colors (isosurface-line-colors)]
181                       #:line-widths [line-widths (isosurface-line-widths)]
182                       #:line-styles [line-styles (isosurface-line-styles)]
183                       #:alphas [alphas (isosurface-alphas)]
184                       #:label [label #f])
185  (define fail/pos (make-raise-argument-error 'isosurfaces3d f x-min x-max y-min y-max z-min z-max))
186  (define fail/kw (make-raise-keyword-error 'isosurfaces3d))
187  (cond
188    [(and x-min (not (rational? x-min)))  (fail/pos "#f or rational" 1)]
189    [(and x-max (not (rational? x-max)))  (fail/pos "#f or rational" 2)]
190    [(and y-min (not (rational? y-min)))  (fail/pos "#f or rational" 3)]
191    [(and y-max (not (rational? y-max)))  (fail/pos "#f or rational" 4)]
192    [(and z-min (not (rational? z-min)))  (fail/pos "#f or rational" 5)]
193    [(and z-max (not (rational? z-max)))  (fail/pos "#f or rational" 6)]
194    [(and d-min (not (rational? d-min)))  (fail/kw "#f or rational" '#:d-min d-min)]
195    [(and d-max (not (rational? d-max)))  (fail/kw "#f or rational" '#:d-max d-max)]
196    [(< samples 2)  (fail/kw "Integer >= 2" '#:samples samples)]
197    [else
198     (define x-ivl (ivl x-min x-max))
199     (define y-ivl (ivl y-min y-max))
200     (define z-ivl (ivl z-min z-max))
201     (define g (3d-function->sampler f (vector x-ivl y-ivl z-ivl)))
202     (define-values (label-proc render-proc)
203       (make-isosurfaces3d-label-and-renderer g d-min d-max levels samples colors styles
204                                              line-colors line-widths line-styles alphas
205                                              label))
206     (renderer3d (vector x-ivl y-ivl z-ivl) #f default-ticks-fun
207                 label-proc render-proc)]))
208
209;; ===================================================================================================
210
211(: polar3d-render-proc (-> (-> Real Real Real Real) 3D-Sampler Positive-Integer
212                           Plot-Color Plot-Brush-Style
213                           Plot-Color Nonnegative-Real Plot-Pen-Style
214                           Nonnegative-Real
215                           3D-Render-Proc))
216(define ((polar3d-render-proc f g samples color style line-color line-width line-style alpha)
217         area)
218  (match-define (vector x-ivl y-ivl z-ivl) (send area get-bounds-rect))
219  (match-define (ivl x-min x-max) x-ivl)
220  (match-define (ivl y-min y-max) y-ivl)
221  (match-define (ivl z-min z-max) z-ivl)
222  (define num (animated-samples samples))
223  (define sample (g (vector x-ivl y-ivl z-ivl) (vector num num num)))
224  (match-define (3d-sample xs ys zs dsss d-min d-max) sample)
225
226  (: draw-cube (-> Real Real Real Real Real Real Real Real Real Real Real Real Real Real Real Void))
227  (define (draw-cube xa xb ya yb za zb d d1 d2 d3 d4 d5 d6 d7 d8)
228    (for ([vs  (in-list (heights->cube-polys xa xb ya yb za zb d d1 d2 d3 d4 d5 d6 d7 d8))])
229      (send area put-polygon vs)))
230
231  (send area put-alpha alpha)
232  (send area put-brush color style)
233  (send area put-pen line-color line-width line-style)
234  (define d 0)
235  (for-3d-sample
236   (xa xb ya yb za zb d1 d2 d3 d4 d5 d6 d7 d8) sample
237   (cond [(and (xb . > . 0) (ya . < . 0) (yb . > . 0))
238          (let* ([yb  -0.00001]
239                 [d3  (f xb yb za)]
240                 [d4  (f xa yb za)]
241                 [d7  (f xb yb zb)]
242                 [d8  (f xa yb zb)])
243            (draw-cube xa xb ya yb za zb d d1 d2 d3 d4 d5 d6 d7 d8))
244          (let* ([ya  0.00001]
245                 [d1  (f xa ya za)]
246                 [d2  (f xb ya za)]
247                 [d5  (f xa ya zb)]
248                 [d6  (f xb ya zb)])
249            (draw-cube xa xb ya yb za zb d d1 d2 d3 d4 d5 d6 d7 d8))]
250         [else
251          (draw-cube xa xb ya yb za zb d d1 d2 d3 d4 d5 d6 d7 d8)]))
252  (void))
253
254(define 2pi (* 2 pi))
255
256(: flmodulo (-> Flonum Flonum Flonum))
257(define (flmodulo x y)
258  (fl- x (fl* y (flfloor (fl/ x y)))))
259
260(: 2d-polar->3d-function (-> (-> Real Real Real) (-> Real Real Real Real)))
261(define ((2d-polar->3d-function f) x y z)
262  (let ([x  (fl x)]
263        [y  (fl y)]
264        [z  (fl z)])
265    (define-values (θ ρ)
266      (cond [(and (fl= x 0.0) (fl= y 0.0))  (values 0.0 0.0)]
267            [else  (values (flmodulo (atan y x) 2pi)
268                           (flatan (fl/ z (fldistance x y))))]))
269    (fl- (fl (f θ ρ)) (fldistance x y z))))
270
271(:: polar3d
272    (->* [(-> Real Real Real)]
273         [#:x-min (U Real #f) #:x-max (U Real #f)
274          #:y-min (U Real #f) #:y-max (U Real #f)
275          #:z-min (U Real #f) #:z-max (U Real #f)
276          #:samples Positive-Integer
277          #:color Plot-Color
278          #:style Plot-Brush-Style
279          #:line-color Plot-Color
280          #:line-width Nonnegative-Real
281          #:line-style Plot-Pen-Style
282          #:alpha Nonnegative-Real
283          #:label (U String pict #f)]
284         renderer3d))
285(define (polar3d f
286                 #:x-min [x-min #f] #:x-max [x-max #f]
287                 #:y-min [y-min #f] #:y-max [y-max #f]
288                 #:z-min [z-min #f] #:z-max [z-max #f]
289                 #:samples [samples (plot3d-samples)]
290                 #:color [color (surface-color)]
291                 #:style [style (surface-style)]
292                 #:line-color [line-color (surface-line-color)]
293                 #:line-width [line-width (surface-line-width)]
294                 #:line-style [line-style (surface-line-style)]
295                 #:alpha [alpha (surface-alpha)]
296                 #:label [label #f])
297  (define fail/kw (make-raise-keyword-error 'polar3d))
298  (cond
299    [(and x-min (not (rational? x-min)))  (fail/kw "#f or rational" '#:x-min x-min)]
300    [(and x-max (not (rational? x-max)))  (fail/kw "#f or rational" '#:x-max x-max)]
301    [(and y-min (not (rational? y-min)))  (fail/kw "#f or rational" '#:x-min y-min)]
302    [(and y-max (not (rational? y-max)))  (fail/kw "#f or rational" '#:x-max y-max)]
303    [(and z-min (not (rational? z-min)))  (fail/kw "#f or rational" '#:x-min z-min)]
304    [(and z-max (not (rational? z-max)))  (fail/kw "#f or rational" '#:x-max z-max)]
305    [(< samples 2)  (fail/kw "Integer >= 2" '#:samples samples)]
306    [(not (rational? line-width))  (fail/kw "rational?" '#:line-width line-width)]
307    [(or (> alpha 1) (not (rational? alpha)))  (fail/kw "real in [0,1]" '#:alpha alpha)]
308    [else
309     (define vs
310       (for*/list : (Listof (Vectorof Real))
311         ([θ  (in-list (linear-seq 0.0 2pi (* 4 samples)))]
312          [ϕ  (in-list (linear-seq (* -1/2 pi) (* 1/2 pi) (* 2 samples)))])
313         (3d-polar->3d-cartesian θ ϕ (f θ ϕ))))
314     (define rvs (filter vrational? vs))
315     (cond [(empty? rvs)  empty-renderer3d]
316           [else
317            (match-define (list (vector #{rxs : (Listof Real)}
318                                        #{rys : (Listof Real)}
319                                        #{rzs : (Listof Real)})
320                                ...)
321              rvs)
322            (let ([x-min  (if x-min x-min (apply min* rxs))]
323                  [x-max  (if x-max x-max (apply max* rxs))]
324                  [y-min  (if y-min y-min (apply min* rys))]
325                  [y-max  (if y-max y-max (apply max* rys))]
326                  [z-min  (if z-min z-min (apply min* rzs))]
327                  [z-max  (if z-max z-max (apply max* rzs))])
328              (define x-ivl (ivl x-min x-max))
329              (define y-ivl (ivl y-min y-max))
330              (define z-ivl (ivl z-min z-max))
331              (define new-f (2d-polar->3d-function f))
332              (define g (3d-function->sampler new-f (vector x-ivl y-ivl z-ivl)))
333              (renderer3d (vector x-ivl y-ivl z-ivl) #f
334                          default-ticks-fun
335                          (and label (λ (_) (rectangle-legend-entry
336                                             label color style line-color line-width line-style)))
337                          (polar3d-render-proc new-f g samples color style
338                                               line-color line-width line-style alpha)))])]))
339