Home My Page Projects Code Snippets Project Openings diderot
Summary Activity Tracker Tasks SCM

SCM Repository

[diderot] Annotation of /branches/vis15/src/compiler/ein/mk-operators.sml
ViewVC logotype

Annotation of /branches/vis15/src/compiler/ein/mk-operators.sml

Parent Directory Parent Directory | Revision Log Revision Log


Revision 3702 - (view) (download)

1 : jhr 3477 (* mk-operators.sml
2 : jhr 3476 *
3 : jhr 3477 * Functions to create the various Ein operators.
4 :     *
5 : jhr 3476 * This code is part of the Diderot Project (http://diderot-language.cs.uchicago.edu)
6 :     *
7 :     * COPYRIGHT (c) 2015 The University of Chicago
8 :     * All rights reserved.
9 :     *)
10 :    
11 :     structure MkOperators : sig
12 :    
13 : jhr 3477 type dim = int
14 :     type shape = dim list
15 :     type ids = Ein.index_id list
16 : jhr 3476
17 : jhr 3477 val addRR : Ein.ein
18 :     val addTT : shape -> Ein.ein
19 :     val addTF : dim * shape -> Ein.ein
20 :     val addFF : dim * shape -> Ein.ein
21 :    
22 :     val subRR : Ein.ein
23 :     val subTT : shape -> Ein.ein
24 :     val subTF : dim * shape -> Ein.ein
25 : jhr 3478 val subFT : dim * shape -> Ein.ein
26 : jhr 3477 val subFF : dim * shape -> Ein.ein
27 :    
28 :     val mulRT : shape -> Ein.ein
29 :     val mulRR : Ein.ein
30 :     val mulRF : dim * shape -> Ein.ein
31 :     val mulST : dim * shape -> Ein.ein
32 :     val mulSS : dim -> Ein.ein
33 :     val mulSF : dim * shape -> Ein.ein
34 :    
35 :     val divTR : shape -> Ein.ein
36 :     val divRR : Ein.ein
37 :     val divFR : dim * shape -> Ein.ein
38 :     val divSS : dim -> Ein.ein
39 :     val divFS : dim * shape -> Ein.ein
40 :    
41 :     val negTT : shape -> Ein.ein
42 :     val negFF : dim * shape -> Ein.ein
43 :    
44 :     val cross2TT : Ein.ein
45 :     val cross3TT : Ein.ein
46 :     val cross2FF : Ein.ein
47 :     val cross3FF : Ein.ein
48 :    
49 : cchiw 3495 val outerTT : shape * shape -> Ein.ein
50 :     val outerFF : dim * shape * shape -> Ein.ein
51 : cchiw 3496 val outerTF : dim * shape * shape -> Ein.ein
52 :     val outerFT : dim * shape * shape -> Ein.ein
53 : jhr 3477
54 :     val innerTT : shape * ids -> Ein.ein
55 :     val innerFF : shape * dim * ids -> Ein.ein
56 :     val innerFT : shape * dim * ids -> Ein.ein
57 :     val innerTF : shape * dim * ids -> Ein.ein
58 :    
59 :     val colonTT : shape * ids -> Ein.ein
60 :     val colonFF : dim * shape * ids -> Ein.ein
61 :     val colonFT : dim * shape * ids -> Ein.ein
62 :     val colonTF : dim * shape * ids -> Ein.ein
63 :    
64 : cchiw 3494 val normT : shape -> Ein.ein
65 :     val normF : dim * shape -> Ein.ein
66 : jhr 3477
67 :     val normalizeTT : shape -> Ein.ein
68 :     val normalizeFF : dim * shape -> Ein.ein
69 :    
70 :     val traceT : dim -> Ein.ein
71 :     val traceF : dim * shape -> Ein.ein
72 :    
73 :     val transposeT : shape -> Ein.ein
74 :     (* QUESTION: should these be index_kind? *)
75 :     val transposeF : dim * Ein.index_id * Ein.index_id -> Ein.ein
76 :    
77 : jhr 3478 val det2T : Ein.ein
78 :     val det3T : Ein.ein
79 :     val det2F : Ein.ein
80 :     val det3F_full : Ein.ein
81 :     val det3FA : Ein.ein
82 :     val det3F : Ein.ein
83 : jhr 3477
84 :     val expF : dim -> Ein.ein
85 :     val expT : Ein.ein
86 :    
87 : jhr 3478 val powF : dim * int -> Ein.ein
88 :    
89 : jhr 3498 val sqrtR : Ein.ein
90 : jhr 3477 val sqrtF : dim -> Ein.ein
91 : jhr 3498 val cosR : Ein.ein
92 : jhr 3477 val cosF : dim -> Ein.ein
93 : jhr 3498 val acosR : Ein.ein
94 : jhr 3477 val acosF : dim -> Ein.ein
95 : jhr 3498 val sinR : Ein.ein
96 : jhr 3477 val sinF : dim -> Ein.ein
97 : jhr 3498 val asinR : Ein.ein
98 : jhr 3477 val asinF : dim -> Ein.ein
99 : jhr 3498 val tanR : Ein.ein
100 : jhr 3477 val tanF : dim -> Ein.ein
101 : jhr 3498 val atanR : Ein.ein
102 : jhr 3477 val atanF : dim -> Ein.ein
103 :    
104 :     val modulate : dim -> Ein.ein
105 :    
106 :     val identity : dim -> Ein.ein
107 : jhr 3478 val zeros : shape -> Ein.ein
108 : jhr 3477
109 :     val slice : shape * bool list * int list * shape -> Ein.ein
110 :    
111 :     val conv : dim * shape -> Ein.ein
112 :     val probe : shape * dim -> Ein.ein
113 :    
114 :     val curl2d : Ein.ein
115 :     val curl3d : Ein.ein
116 :     val grad : shape -> Ein.ein
117 :     val dotimes : dim * shape -> Ein.ein (* ?? *)
118 :     val divergence : dim * shape -> Ein.ein
119 :    
120 : jhr 3476 end = struct
121 :    
122 :     structure E = Ein
123 :    
124 : jhr 3477 type dim = int
125 :     type shape = dim list
126 :     type ids = int list
127 : jhr 3701
128 :     (* controls whether tensor operations should be *)
129 :     val canSubst = true
130 :    
131 :     (* A constructor function for tensor variables that (by default) can be substituted for;
132 :     * this behavior is controlled by the canSubst flag.
133 :     *)
134 :     fun mkTEN alpha = E.TEN(canSubst, alpha)
135 :    
136 :     (* a constructor function for tensor parameters that should never be substituted for. *)
137 :     fun mkNoSubstTEN alpha = E.TEN(false, alpha)
138 :    
139 : jhr 3476 fun specialize (alpha, inc) = List.mapi (fn (i, _) => E.V(i + inc)) alpha
140 :    
141 :     fun sumIds (n, inc, alpha) = let
142 :     val vs = List.tabulate(n, fn v => E.V(v+inc))
143 :     in
144 : jhr 3498 ListPair.map (fn(v, i) => (v, 0, i-1)) (vs, alpha)
145 :     end
146 : jhr 3476
147 :     fun sumIds2 (n, i) = List.tabulate(n, fn v => (E.V v, 0, i))
148 :    
149 :     (******************************* Addition *****************************************)
150 :    
151 :     (* Adding tensors : < X{\alpha} + Y_{\alpha}>_{\alpha} *)
152 :     fun addTT alpha = let
153 : jhr 3498 val expindex = specialize(alpha, 0)
154 : jhr 3702 in
155 : jhr 3476 E.EIN{
156 : jhr 3701 params = [mkTEN alpha, mkTEN alpha],
157 : jhr 3476 index = alpha,
158 :     body = E.Opn(E.Add, [E.Tensor(0, expindex), E.Tensor(1, expindex)])
159 : jhr 3498 }
160 :     end
161 : jhr 3476
162 :     val addRR = addTT []
163 :    
164 : jhr 3477 (* Tensor and Fields *)
165 :     fun addTF (dim, shape) =let
166 : jhr 3498 val expindex = specialize(shape, 0)
167 :     in
168 :     E.EIN{
169 : jhr 3701 params = [mkTEN shape, E.FLD dim],
170 : jhr 3498 index = shape,
171 :     body = E.Opn(E.Add, [E.Lift(E.Tensor(0, expindex)), E.Field(1, expindex)])
172 :     }
173 :     end
174 : jhr 3476
175 : jhr 3477 (* Adding Fields : < F{\alpha} + G_{\alpha}>_{\alpha} *)
176 :     fun addFF (dim, shape) =let
177 : jhr 3498 val expindex = specialize(shape, 0)
178 :     in
179 :     E.EIN{
180 :     params = [E.FLD dim, E.FLD dim],
181 :     index = shape,
182 :     body = E.Opn(E.Add, [E.Field(0, expindex), E.Field(1, expindex)])
183 :     }
184 :     end
185 : jhr 3476
186 :     (********************************* Subtraction **************************************)
187 :    
188 :     fun subTT alpha = let
189 : jhr 3498 val expindex = specialize(alpha, 0)
190 :     in
191 : jhr 3476 E.EIN{
192 : jhr 3701 params = [mkTEN alpha, mkTEN alpha],
193 : jhr 3498 index = alpha,
194 :     body = E.Op2(E.Sub, E.Tensor(0, expindex), E.Tensor(1, expindex))
195 :     }
196 :     end
197 : jhr 3476
198 :     val subRR = subTT []
199 :    
200 :     fun subTF (dim, shape) = let
201 : jhr 3498 val expindex = specialize(shape, 0)
202 :     in
203 :     E.EIN{
204 : jhr 3701 params = [mkTEN shape, E.FLD dim],
205 : jhr 3498 index = shape,
206 :     body = E.Opn(E.Add,
207 :     [E.Lift(E.Tensor(0, expindex)), E.Op1(E.Neg, E.Field(1, expindex))])
208 :     }
209 :     end
210 : jhr 3476
211 :     fun subFT (dim, shape) = let
212 : jhr 3498 val expindex = specialize(shape, 0)
213 :     in
214 :     E.EIN{
215 : jhr 3701 params = [mkTEN shape, E.FLD dim],
216 : jhr 3498 index = shape,
217 :     body = E.Op2(E.Sub, E.Field(1, expindex), E.Lift(E.Tensor(0, expindex)))
218 :     }
219 :     end
220 : jhr 3476
221 :     fun subFF (dim, shape) = let
222 : jhr 3498 val expindex = specialize(shape, 0)
223 :     in
224 :     E.EIN{
225 :     params = [E.FLD dim, E.FLD dim],
226 :     index = shape,
227 :     body = E.Op2(E.Sub, E.Field(0, expindex), E.Field(1, expindex))
228 :     }
229 :     end
230 : jhr 3476
231 :     (********************************** Multiplication *************************************)
232 :    
233 :     (* scalar times tensor product: <s * T_{\alpha}>_{\alpha} *)
234 :     fun mulRT alpha = let
235 : jhr 3498 val expindex = specialize(alpha, 0)
236 :     in
237 : jhr 3476 E.EIN{
238 : jhr 3701 params = [mkTEN [], mkTEN alpha],
239 : jhr 3476 index = alpha,
240 : jhr 3477 body = E.Opn(E.Prod, [E.Tensor(0, []), E.Tensor(1, expindex)])
241 : jhr 3498 }
242 :     end
243 : jhr 3476
244 :     val mulRR = mulRT []
245 :    
246 : jhr 3477 fun mulRF (dim, shape) =let
247 : jhr 3476 val expindex = specialize(shape, 0)
248 :     in
249 :     E.EIN{
250 : jhr 3701 params = [mkTEN [], E.FLD dim],
251 : jhr 3476 index = shape,
252 : jhr 3477 body = E.Opn(E.Prod, [E.Lift( E.Tensor(0, [])), E.Field(1,expindex)])
253 : jhr 3476 }
254 :     end
255 :    
256 :     fun mulST (dim, shape) =let
257 :     val expindex = specialize(shape, 0)
258 :     in
259 :     E.EIN{
260 : jhr 3701 params = [mkTEN shape, E.FLD dim],
261 : jhr 3476 index = shape,
262 : jhr 3477 body = E.Opn(E.Prod, [E.Lift( E.Tensor(0,expindex)), E.Field(1, [])])
263 : jhr 3476 }
264 :     end
265 :    
266 :     fun mulSS dim = E.EIN{
267 : jhr 3477 params = [E.FLD dim, E.FLD dim],
268 : jhr 3476 index = [],
269 : jhr 3477 body = E.Opn(E.Prod, [E.Field(0, []), E.Field(1, [])])
270 : jhr 3476 }
271 :    
272 : jhr 3477 fun mulSF(dim, shape) =let
273 :     val expindex= specialize(shape, 0)
274 : jhr 3476 in E.EIN{
275 : jhr 3477 params = [E.FLD dim, E.FLD dim],
276 : jhr 3476 index = shape,
277 : jhr 3477 body = E.Opn(E.Prod, [E.Field(0, []), E.Field(1, expindex)])
278 : jhr 3476 }
279 :     end
280 :    
281 :     (************************************ Division ************************************)
282 :    
283 : jhr 3477 fun divTR alpha = let
284 : jhr 3498 val expindex = specialize(alpha, 0)
285 :     in
286 : jhr 3476 E.EIN{
287 : jhr 3701 params = [mkTEN alpha, mkTEN []],
288 : jhr 3476 index = alpha,
289 : jhr 3477 body = E.Op2(E.Div, E.Tensor(0, expindex), E.Tensor(1, []))
290 : jhr 3498 }
291 :     end
292 : jhr 3476
293 : jhr 3477 val divRR = divTR []
294 : jhr 3476
295 : jhr 3477 fun divFR (dim, shape) = let
296 : jhr 3498 val expindex = specialize(shape, 0)
297 :     in
298 :     E.EIN{
299 : jhr 3701 params = [E.FLD dim, mkTEN []],
300 : jhr 3498 index = shape,
301 :     body = E.Op2(E.Div, E.Field(0, expindex), E.Lift(E.Tensor(1, [])))
302 :     }
303 :     end
304 : jhr 3477
305 : jhr 3476 fun divSS dim = E.EIN{
306 : jhr 3498 params = [E.FLD dim, E.FLD dim],
307 :     index = [],
308 :     body = E.Op2(E.Div, E.Field(0, []), E.Field(1, []))
309 :     }
310 : jhr 3476
311 : jhr 3477 fun divFS(dim, shape) = let
312 : jhr 3498 val expindex = specialize(shape, 0)
313 :     in
314 :     E.EIN{
315 :     params = [E.FLD dim, E.FLD dim],
316 :     index = shape,
317 : jhr 3512 body = E.Opn(E.Prod, [E.Field(0, expindex), E.Op2(E.Div, E.Const 1, E.Field(1, []))])
318 : jhr 3498 }
319 :     end
320 : jhr 3476
321 : jhr 3477 (************************************* Negation **********************************)
322 :    
323 : jhr 3476 fun negTT alpha = let
324 : jhr 3498 val expindex= specialize(alpha, 0)
325 : jhr 3476 (*changed tensor lift variable here *)
326 : jhr 3477 in
327 :     E.EIN {
328 : jhr 3701 params = [mkTEN alpha], index = alpha,
329 : jhr 3498 body = E.Op1(E.Neg, E.Tensor(0, expindex))
330 : jhr 3702 }
331 : jhr 3498 end
332 : jhr 3476
333 : jhr 3477 fun negFF (dim, shape) = let
334 : jhr 3498 val expindex = specialize(shape, 0)
335 :     in
336 :     E.EIN{
337 :     params = [E.FLD dim], index = shape,
338 :     body = E.Op1(E.Neg, E.Field(0, expindex))
339 :     }
340 : jhr 3477 end
341 : jhr 3476
342 : jhr 3477 (****************************** cross product ***********************************)
343 : jhr 3476
344 : jhr 3477 (* 2-d cross product Eps_{ij}U_i V_j *)
345 : jhr 3476 val cross2TT = E.EIN{
346 : jhr 3701 params = [mkTEN [2], mkTEN [2]],
347 : jhr 3477 index = [],
348 :     body = E.Sum([(E. V 0, 0, 1), (E.V 1, 0, 1)],
349 : jhr 3512 E.Opn(E.Prod, [E.Eps2(0, 1), E.Tensor(0, [E.V 0]), E.Tensor(1, [E.V 1])]))
350 : jhr 3498 }
351 : jhr 3476
352 : jhr 3477 (* crossProduct is on 3D vectors ..vec3 t8=t0 × t1; *)
353 :     val cross3TT = E.EIN{
354 : jhr 3701 params = [mkTEN [3], mkTEN [3]],
355 : jhr 3477 index = [3],
356 :     body = E.Sum([(E.V 1, 0, 2), (E.V 2, 0, 2)],
357 : jhr 3498 E.Opn(E.Prod, [
358 : jhr 3512 E.Epsilon(0, 1, 2),
359 : jhr 3498 E.Tensor(0, [E.V 1]),
360 :     E.Tensor(1, [E.V 2])
361 :     ]))
362 :     }
363 : jhr 3477
364 :     (* Field Cross Product *)
365 : jhr 3476 val cross2FF = E.EIN{
366 : jhr 3477 params = [E.FLD(2), E.FLD(2)], index = [],
367 :     body = E.Sum([(E.V 0, 0, 1), (E.V 1, 0, 1)],
368 : jhr 3512 E.Opn(E.Prod, [E.Eps2(0, 1), E.Field(0, [E.V 0]), E.Field(1, [E.V 1])]))
369 : jhr 3498 }
370 : jhr 3476
371 : jhr 3477 (* Field Cross Product *)
372 : jhr 3476 val cross3FF = E.EIN{
373 : jhr 3701 params = [E.FLD(3), E.FLD(3)], index= [3],
374 : jhr 3477 body = E.Sum([(E.V 1, 0, 2), (E.V 2, 0, 2)],
375 : jhr 3512 E.Opn(E.Prod, [E.Epsilon(0, 1, 2),
376 : jhr 3477 E.Field(0, [E.V 1]), E.Field(1, [E.V 2])]))
377 : jhr 3498 }
378 : jhr 3476
379 : jhr 3477 (******************** outer product ********************************)
380 : jhr 3498
381 :     fun outerTT (alpha, beta) = let
382 :     val expIdxA = specialize (alpha, 0)
383 :     val expIdxB = specialize (beta, length alpha)
384 :     in
385 :     E.EIN{
386 : jhr 3701 params = [mkTEN alpha, mkTEN beta],
387 : jhr 3498 index = alpha@beta,
388 :     body = E.Opn(E.Prod, [E.Tensor(0, expIdxA), E.Tensor(1, expIdxB)])
389 :     }
390 : cchiw 3495 end
391 : jhr 3476
392 : cchiw 3495 (*Assumes same dimension vector field *)
393 : jhr 3498 fun outerFF (dim, alpha, beta) =let
394 :     val expIdxA = specialize (alpha, 0)
395 :     val expIdxB = specialize (beta, length alpha)
396 :     in
397 :     E.EIN{
398 :     params = [E.FLD dim, E.FLD dim],
399 :     index = alpha@beta,
400 :     body = E.Opn(E.Prod, [E.Field(0, expIdxA), E.Field(1, expIdxB)])
401 :     }
402 :     end
403 : jhr 3477
404 : jhr 3498 fun outerTF (dim, alpha, beta) =let
405 :     val expIdxA = specialize (alpha, 0)
406 :     val expIdxB = specialize (beta, length alpha)
407 :     in
408 :     E.EIN{
409 : jhr 3701 params = [mkTEN alpha, E.FLD dim],
410 : jhr 3498 index = alpha@beta,
411 :     body = E.Opn(E.Prod, [E.Lift(E.Tensor(0, expIdxA)), E.Field(1, expIdxB)])
412 :     }
413 :     end
414 : jhr 3702
415 : jhr 3498 fun outerFT (dim, alpha, beta) =let
416 :     val expIdxA = specialize(alpha, 0)
417 :     val expIdxB = specialize(beta, length alpha)
418 :     in
419 :     E.EIN{
420 : jhr 3701 params = [E.FLD dim, mkTEN alpha],
421 : jhr 3498 index = alpha@beta,
422 :     body = E.Opn(E.Prod, [E.Field(0, expIdxA), E.Lift(E.Tensor(1, expIdxB))])
423 :     }
424 :     end
425 : jhr 3477
426 :     (*************************** inner product **********************************)
427 : jhr 3476 (* generic inner product: <T_{\alpha i} * T_{i \beta}>_{\alpha \beta} *)
428 : jhr 3477 fun innerTT (shape1, i::beta) = let
429 : jhr 3498 val alpha = List.take(shape1, length shape1 - 1)
430 :     val expindexA = specialize(alpha, 0)
431 :     val expindexB = specialize(beta, length alpha)
432 :     val s' = E.V(length alpha + length beta)
433 :     val s'' = [(s', 0, i-1)]
434 :     in
435 : jhr 3476 E.EIN{
436 : jhr 3702 params = [mkTEN shape1, mkTEN(i :: beta)],
437 : jhr 3476 index = alpha@beta,
438 : jhr 3477 body = E.Sum(s'', E.Opn(E.Prod, [
439 : jhr 3498 E.Tensor(0, expindexA@[s']), (* T_{\alpha i} *)
440 :     E.Tensor(1, [s']@expindexB ) (* T'_{i \beta} *)
441 :     ]))
442 :     }
443 :     end
444 : jhr 3476 | innerTT _ = raise Fail "Wrong shape for inner product"
445 :    
446 : jhr 3477 (* generic inner product: <T_{\alpha i} * T_{i \beta}>_{\alpha \beta} *)
447 :     fun innerFF (shape1, dim, i::beta) = let
448 : jhr 3498 val alpha = List.take(shape1, length shape1 - 1)
449 :     val expindexA = specialize(alpha, 0)
450 :     val expindexB = specialize(beta, length alpha)
451 :     val sid = E.V(length alpha + length beta)
452 :     in
453 :     E.EIN{
454 :     params = [E.FLD dim, E.FLD dim],
455 : jhr 3702 index = alpha @ beta,
456 : jhr 3498 body = E.Sum([(sid, 0, i-1)],
457 :     E.Opn(E.Prod, [
458 : jhr 3702 E.Field(0, expindexA @ [sid]), (* F_{\alpha i} *)
459 :     E.Field(1, [sid] @ expindexB) (* F'_{i \beta} *)
460 : jhr 3498 ]))
461 :     }
462 :     end
463 : jhr 3477 | innerFF _ = raise Fail "Wrong shape for innerProductField"
464 : jhr 3476
465 : jhr 3477 fun innerFT (shape1, dim, i::beta) = let
466 : jhr 3498 val alpha = List.take(shape1, length shape1-1)
467 :     val expindexA = specialize(alpha, 0)
468 :     val expindexB = specialize(beta, length alpha)
469 :     val sid = E.V(length alpha + length beta)
470 :     in
471 :     E.EIN{
472 : jhr 3702 params = [E.FLD dim, mkTEN(i::beta)],
473 :     index = alpha @ beta,
474 : jhr 3498 body = E.Sum([(sid, 0, i-1)],
475 :     E.Opn(E.Prod, [
476 : jhr 3702 E.Field(0, expindexA @ [sid]), (* F_{\alpha i} *)
477 :     E.Lift(E.Tensor(1, [sid] @ expindexB )) (* F'_{i \beta} *)
478 : jhr 3498 ]))
479 :     }
480 :     end
481 : jhr 3477 | innerFT _ = raise Fail "Wrong shape for innerProductFieldTensor"
482 : jhr 3476
483 : jhr 3477 fun innerTF (shape1, dim, i::beta) = let
484 : jhr 3498 val alpha = List.take(shape1, length shape1 - 1)
485 :     val expindexA = specialize(alpha, 0)
486 :     val expindexB = specialize(beta, length alpha)
487 :     val sid = E.V(length(alpha) + length beta)
488 :     in
489 :     E.EIN{
490 : jhr 3702 params = [mkTEN shape1, E.FLD dim],
491 :     index = alpha @ beta,
492 : jhr 3498 body = E.Sum([(sid, 0, i-1)],
493 :     E.Opn(E.Prod, [
494 : jhr 3702 E.Lift(E.Tensor(0, expindexA @ [sid])), (* F_{\alpha i} *)
495 :     E.Field(1, [sid] @ expindexB) (* F'_{i \beta} *)
496 : jhr 3498 ]))
497 :     }
498 :     end
499 : jhr 3477 | innerTF _ = raise Fail "Wrong shape for innerProductTensorField"
500 : jhr 3476
501 : jhr 3477 (*************************** colon product **********************************)
502 : jhr 3476
503 : jhr 3477 (* <T_{\alpha i j} * B{i j \beta }>_\alpha \beta *)
504 :     fun colonTT (shape1, i::j::beta) = let
505 : jhr 3498 val lenAlpha = length shape1 - 2
506 :     val alpha = List.take(shape1, lenAlpha)
507 :     val expindexA = specialize(alpha, 0)
508 :     val expindexB = specialize(beta, lenAlpha)
509 :     val sumi = lenAlpha + length beta
510 :     val s' = [E.V sumi, E.V(sumi+1)]
511 :     val sx = [(E.V sumi, 0, i-1), (E.V(sumi+1), 0, j-1)]
512 :     in
513 : jhr 3476 E.EIN{
514 : jhr 3702 params = [mkTEN shape1, mkTEN(i::j::beta)],
515 : jhr 3476 index = alpha@beta,
516 : jhr 3477 body = E.Sum(sx,
517 : jhr 3498 E.Opn(E.Prod, [E.Tensor(0, expindexA@s'), E.Tensor(1, s'@expindexB)]))
518 :     }
519 :     end
520 : jhr 3476
521 : jhr 3477 (* <F_{\alpha i j} * G_{i j \beta }>_\alpha \beta *)
522 :     fun colonFF (dim, shape1, i::j::beta) = let
523 : jhr 3498 val lenAlpha = length shape1 - 2
524 :     val alpha = List.take(shape1, lenAlpha)
525 :     val expindexA = specialize(alpha, 0)
526 :     val expindexB = specialize(beta, lenAlpha)
527 :     val sumi = lenAlpha + length beta
528 :     val s' = [E.V sumi, E.V(sumi+1)]
529 :     val sx = [(E.V sumi, 0, i-1), (E.V(sumi+1), 0, j-1)]
530 :     in
531 : jhr 3476 E.EIN{
532 : jhr 3477 params = [E.FLD dim, E.FLD dim],
533 : jhr 3476 index = alpha@beta,
534 : jhr 3477 body = E.Sum(sx,
535 : jhr 3498 E.Opn(E.Prod, [E.Field(0, expindexA@s'), E.Field(1, s'@expindexB)]))
536 :     }
537 :     end
538 : jhr 3476
539 :    
540 : jhr 3477 (* <F_{\alpha i j} * T_{i j \beta }>_\alpha \beta *)
541 :     fun colonFT (dim, shape1, i::j::beta) = let
542 : jhr 3498 val lenAlpha = length shape1 - 2
543 :     val alpha = List.take(shape1, lenAlpha)
544 :     val expindexA = specialize(alpha, 0)
545 :     val expindexB = specialize(beta, lenAlpha)
546 :     val sumi = lenAlpha + length beta
547 :     val s' = [E.V sumi, E.V(sumi+1)]
548 :     val sx = [(E.V sumi, 0, i-1), (E.V(sumi+1), 0, j-1)]
549 :     in
550 : jhr 3476 E.EIN{
551 : jhr 3701 params = [E.FLD dim, mkTEN shape1],
552 : jhr 3498 index = alpha@beta,
553 :     body = E.Sum(sx,
554 :     E.Opn(E.Prod, [E.Field(0, expindexA@s'), E.Lift(E.Tensor(1, s'@expindexB))]))
555 :     }
556 :     end
557 : jhr 3476
558 : jhr 3477 (* <T_{\alpha i j} * G{i j \beta }>_\alpha \beta *)
559 :     fun colonTF (dim, shape1, i::j::beta) = let
560 : jhr 3498 val lenAlpha = length shape1 - 2
561 :     val alpha = List.take(shape1, lenAlpha)
562 :     val expindexA = specialize(alpha, 0)
563 :     val expindexB = specialize(beta, lenAlpha)
564 :     val sumi = lenAlpha + length beta
565 :     val s' = [E.V sumi, E.V(sumi+1)]
566 :     val sx = [(E.V sumi, 0, i-1), (E.V(sumi+1), 0, j-1)]
567 :     in
568 : jhr 3476 E.EIN{
569 : jhr 3702 params = [mkTEN(i::j::beta), E.FLD dim],
570 : jhr 3498 index = alpha@beta,
571 :     body = E.Sum(sx,
572 :     E.Opn(E.Prod, [E.Lift(E.Tensor(0, expindexA@s')), E.Field(1, s'@expindexB)]))
573 :     }
574 :     end
575 : jhr 3476
576 : jhr 3477 (******************** Norm ********************************)
577 : jhr 3476
578 : jhr 3498 fun normT alpha = let
579 :     val expIdx = specialize(alpha, 0)
580 :     val sx = sumIds(length alpha, 0, alpha)
581 :     in
582 : jhr 3476 E.EIN{
583 : jhr 3701 params = [mkTEN alpha],
584 : jhr 3498 index = [],
585 :     body = E.Op1(E.Sqrt,
586 :     E.Sum(sx, E.Opn(E.Prod, [E.Tensor(0, expIdx), E.Tensor(0, expIdx)])))
587 :     }
588 :     end
589 : jhr 3702
590 : jhr 3498 fun normF (dim, alpha) =let
591 :     val expIdx = specialize(alpha, 0)
592 :     val sx = sumIds(length alpha, 0, alpha)
593 :     in
594 :     E.EIN{
595 :     params = [E.FLD dim],
596 :     index = [],
597 :     body = E.Op1(E.Sqrt,
598 :     E.Sum(sx, E.Opn(E.Prod, [E.Field(0, expIdx), E.Field(0, expIdx)])))
599 :     }
600 :     end
601 : jhr 3476
602 : cchiw 3494 fun normalizeTT alpha = let
603 : jhr 3498 val expindex = specialize(alpha, 0)
604 :     val len = length alpha
605 :     val expindexDot = specialize(alpha, len)
606 :     val sx = sumIds(len, len, alpha)
607 :     val f = E.Tensor(0, expindex)
608 :     val g = E.Tensor(1, expindexDot)
609 :     in
610 : jhr 3476 E.EIN{
611 : jhr 3701 params = [mkTEN alpha, mkTEN alpha],
612 : jhr 3498 index = alpha,
613 :     body = E.Opn(E.Prod, [
614 : jhr 3512 f, E.Op2(E.Div, E.Const 1, E.Op1(E.Sqrt, E.Sum(sx, E.Opn(E.Prod, [g, g]))))
615 : jhr 3498 ])
616 :     }
617 :     end
618 : jhr 3476
619 : cchiw 3494 fun normalizeFF (dim, alpha as i::_) = let
620 : jhr 3498 val expindex = specialize(alpha, 0)
621 :     val len = length alpha
622 :     val expindexDot = specialize(alpha, len)
623 :     val sx = sumIds(len, len, alpha)
624 :     val f = E.Field(0, expindex)
625 :     val g = E.Field(1, expindexDot)
626 :     in
627 : jhr 3476 E.EIN{
628 : jhr 3477 params = [E.FLD dim, E.FLD dim],
629 : cchiw 3545 index = alpha,
630 : jhr 3477 body = E.Opn(E.Prod, [
631 : jhr 3512 f, E.Op2(E.Div, E.Const 1, E.Op1(E.Sqrt, E.Sum(sx, E.Opn(E.Prod, [g, g]))))
632 : jhr 3498 ])
633 :     }
634 :     end
635 : jhr 3476
636 : jhr 3477 (************************* trace *************************)
637 :    
638 :     (* Trace: <M_{i, i}> This one Sx represents both i's*)
639 : jhr 3476 fun traceT dim = E.EIN{
640 : jhr 3701 params = [mkTEN [dim, dim]], index = [],
641 : jhr 3498 body = E.Sum([(E.V 0, 0, dim-1)], E.Tensor(0, [E.V 0, E.V 0]))
642 :     }
643 : jhr 3476
644 : jhr 3477 (* Trace: <Sigma_i F_{\alpha i, i}> This one Sx represents both i's *)
645 :     fun traceF (dim, alpha) = let
646 : jhr 3498 val expindex = specialize(alpha, 0)
647 :     val s = E.V(length alpha)
648 :     in
649 : jhr 3476 E.EIN{
650 : jhr 3477 params = [E.FLD dim],
651 : jhr 3476 index = alpha,
652 : jhr 3477 body = E.Sum([(s, 0, dim-1)], E.Field(0, expindex@[s, s]))
653 : jhr 3498 }
654 :     end
655 : jhr 3476
656 : jhr 3477 (************************* tranpose *************************)
657 : jhr 3476
658 : jhr 3477 fun transposeT alpha = E.EIN{
659 : jhr 3701 params = [mkTEN alpha],
660 : jhr 3498 index = List.rev alpha,
661 :     body = E.Tensor(0, [E.V 1, E.V 0])
662 :     }
663 : jhr 3476
664 : jhr 3477 (* Transpose Field F_{ji} *)
665 :     fun transposeF (dim, i, j) = E.EIN{
666 : jhr 3498 params = [E.FLD dim],
667 :     index = [i, j],
668 :     body = E.Field(0, [E.V 1, E.V 0])
669 :     }
670 : jhr 3476
671 : jhr 3477 (************************* determinant *************************)
672 : jhr 3476
673 : jhr 3478 val det2T = E.EIN{
674 : jhr 3701 params = [mkNoSubstTEN [2, 2]],
675 : jhr 3498 index = [],
676 :     body = E.Op2(E.Sub,
677 :     E.Opn(E.Prod, [E.Tensor(0, [E.C 0, E.C 0]), E.Tensor(0, [E.C 1, E.C 1])]),
678 :     E.Opn(E.Prod, [E.Tensor(0, [E.C 0, E.C 1]), E.Tensor(0, [E.C 1, E.C 0])]))
679 :     }
680 : jhr 3476
681 : jhr 3478 val det3T = let
682 : jhr 3498 val a = E.Tensor(0, [E.C 0, E.C 0])
683 :     val b = E.Tensor(0, [E.C 0, E.C 1])
684 :     val c = E.Tensor(0, [E.C 0, E.C 2])
685 :     val d = E.Tensor(0, [E.C 1, E.C 0])
686 :     val e = E.Tensor(0, [E.C 1, E.C 1])
687 :     val f = E.Tensor(0, [E.C 1, E.C 2])
688 :     val g = E.Tensor(0, [E.C 2, E.C 0])
689 :     val h = E.Tensor(0, [E.C 2, E.C 1])
690 :     val i = E.Tensor(0, [E.C 2, E.C 2])
691 :     in
692 :     E.EIN{
693 : jhr 3701 params = [mkNoSubstTEN [3, 3]],
694 : jhr 3498 index = [],
695 :     body = E.Op2(E.Sub,
696 :     E.Opn(E.Add, [
697 :     E.Opn(E.Prod, [a, e, i]), E.Opn(E.Prod, [b, f, g]), E.Opn(E.Prod, [c, d, h])
698 :     ]),
699 :     E.Opn(E.Add, [
700 :     E.Opn(E.Prod, [c, e, g]), E.Opn(E.Prod, [b, d, i]), E.Opn(E.Prod, [a, f, h])
701 :     ]))
702 :     }
703 :     end
704 : jhr 3477
705 : jhr 3478 val det2F = E.EIN{
706 : jhr 3498 params = [E.FLD 2],
707 :     index = [],
708 :     body = E.Op2(E.Sub,
709 :     E.Opn(E.Prod, [E.Field(0, [E.C 0, E.C 0]), E.Field(0, [E.C 1, E.C 1])]),
710 :     E.Opn(E.Prod, [E.Field(0, [E.C 0, E.C 1]), E.Field(0, [E.C 1, E.C 0])]))
711 :     }
712 : jhr 3477
713 : jhr 3478 val det3F_full = let
714 : jhr 3498 val a = E.Field(0, [E.C 0, E.C 0])
715 :     val b = E.Field(0, [E.C 0, E.C 1])
716 :     val c = E.Field(0, [E.C 0, E.C 2])
717 :     val d = E.Field(0, [E.C 1, E.C 0])
718 :     val e = E.Field(0, [E.C 1, E.C 1])
719 :     val f = E.Field(0, [E.C 1, E.C 2])
720 :     val g = E.Field(0, [E.C 2, E.C 0])
721 :     val h = E.Field(0, [E.C 2, E.C 1])
722 :     val i = E.Field(0, [E.C 2, E.C 2])
723 :     in
724 :     E.EIN{
725 :     params = [E.FLD 3],
726 :     index = [],
727 :     body = E.Op2(E.Sub,
728 :     E.Opn(E.Add, [
729 :     E.Opn(E.Prod, [a, e, i]), E.Opn(E.Prod, [b, f, g]), E.Opn(E.Prod, [c, d, h])
730 :     ]),
731 :     E.Opn(E.Add, [
732 :     E.Opn(E.Prod, [c, e, g]), E.Opn(E.Prod, [b, d, i]), E.Opn(E.Prod, [a, f, h])
733 :     ]))
734 :     }
735 :     end
736 : jhr 3476
737 : jhr 3478 val det3FA = E.EIN{
738 : jhr 3498 params = [E.FLD 3],
739 :     index = [],
740 :     body = E.Sum([(E.V 0, 0, 2), (E.V 1, 0, 2), (E.V 2, 0, 2)],
741 :     E.Opn(E.Prod, [
742 : jhr 3512 E.Epsilon(0, 1, 2),
743 : jhr 3498 E.Field(0, [E.C 0, E. V 0]),
744 :     E.Field(0, [E.C 1, E. V 1]),
745 :     E.Field(0, [E.C 2, E. V 2])
746 :     ]))
747 :     }
748 : jhr 3476
749 : jhr 3478 val det3F = E.EIN{
750 : jhr 3498 params = [E.FLD 3],
751 :     index = [],
752 :     body = E.Sum([(E.V 0, 0, 2)],
753 :     E.Opn(E.Prod, [
754 :     E.Field(0, [E.C 0, E. V 0]),
755 :     E.Sum([(E.V 1, 0, 2)],
756 :     E.Opn(E.Prod, [
757 :     E.Field(0, [E.C 1, E. V 1]),
758 :     E.Sum([(E.V 2, 0, 2)],
759 : jhr 3512 E.Opn(E.Prod, [E.Epsilon(0, 1, 2), E.Field(0, [E.C 2, E. V 2])]))
760 : jhr 3498 ]))
761 :     ]))
762 :     }
763 : jhr 3476
764 : jhr 3477 (************************* Exponential **************************)
765 :     fun expF dim = E.EIN{params = [E.FLD dim], index = [], body = E.Op1(E.Exp, E.Field(0, []))}
766 : jhr 3701 val expT = E.EIN{params = [mkNoSubstTEN []], index = [], body = E.Op1(E.Exp, E.Tensor(0, []))}
767 : jhr 3476
768 : jhr 3477 (************************* Lifted single-argument math functions *************************)
769 :     local
770 : jhr 3498 fun tensorFn rator = E.EIN{
771 : jhr 3701 params = [mkNoSubstTEN []],
772 : jhr 3498 index = [],
773 :     body = E.Op1(rator, E.Field(0, []))
774 :     }
775 :     fun liftFn rator dim = E.EIN{
776 :     params = [E.FLD dim],
777 :     index = [],
778 :     body = E.Op1(rator, E.Field(0, []))
779 :     }
780 : jhr 3477 in
781 : jhr 3478 fun powF (dim, n) = E.EIN{
782 : jhr 3498 params = [E.FLD dim],
783 :     index = [], body = E.Op1(E.PowInt n, E.Field(0, []))
784 :     }
785 :     val sqrtR = tensorFn E.Sqrt
786 : jhr 3477 val sqrtF = liftFn E.Sqrt
787 : jhr 3498 val cosR = tensorFn E.Cosine
788 :     val cosF = liftFn E.Cosine
789 :     val acosR = tensorFn E.ArcCosine
790 : jhr 3477 val acosF = liftFn E.ArcCosine
791 : jhr 3498 val sinR = tensorFn E.Sine
792 :     val sinF = liftFn E.Sine
793 :     val asinR = tensorFn E.ArcSine
794 : jhr 3477 val asinF = liftFn E.ArcSine
795 : jhr 3498 val tanR = tensorFn E.Tangent
796 :     val tanF = liftFn E.Tangent
797 :     val atanR = tensorFn E.ArcTangent
798 : jhr 3477 val atanF = liftFn E.ArcTangent
799 :     end (* local *)
800 : jhr 3476
801 : jhr 3477 (************************* other tensor ops *************************)
802 :     fun modulate dim = E.EIN{
803 : jhr 3701 params = [mkTEN [dim], mkTEN [dim]],
804 : jhr 3498 index = [dim],
805 :     body = E.Opn(E.Prod, [E.Tensor(0, [E.V 0]), E.Tensor(1, [E.V 0])])
806 :     }
807 : jhr 3476
808 : jhr 3477 fun identity dim = E.EIN{
809 : jhr 3512 params = [], index = [dim, dim], body = E.Delta(E.V(0), E.V(1))
810 : jhr 3498 }
811 : jhr 3476
812 : jhr 3478 fun zeros shape = E.EIN{
813 : jhr 3512 params = [], index = shape, body = E.Const 0
814 : jhr 3498 }
815 : jhr 3478
816 : jhr 3477 (* QUESTION: why do we need the const list? The indices are implicit in the position of
817 :     * of the mask element! Likewise, the result type can be determined from the argTy and
818 :     * mask.
819 :     *)
820 :     fun slice (argTy, mask, const, rstTy) = let
821 : jhr 3498 fun iter ([], _, cnt) = []
822 :     | iter (true::es, c::cs, cnt) = E.C c :: iter(es, cs, cnt)
823 :     | iter (false::es, cs, cnt) = E.V cnt :: iter(es, cs, cnt+1)
824 :     val ix = iter(mask, const, 0)
825 :     in
826 :     E.EIN{
827 : jhr 3701 params = [mkTEN argTy],
828 : jhr 3498 index = rstTy,
829 :     body = E.Tensor(0, ix)
830 :     }
831 :     end
832 : jhr 3476
833 : jhr 3477 (******************** other field ops ********************************)
834 : jhr 3476
835 : jhr 3477 (* FLD here is bounded to image field, and dimension of h *)
836 :     fun conv (dim, shape) =let
837 : jhr 3498 val expindex = specialize(shape, 0)
838 :     in
839 :     E.EIN{
840 : jhr 3645 params = [E.IMG(dim, shape), E.KRN],
841 : jhr 3498 index = shape,
842 :     body = E.Conv(0,expindex, 1, [])
843 :     }
844 :     end
845 : jhr 3476
846 : jhr 3477 (* Probe: <F(x)>_{\alpha} *)
847 :     fun probe (alpha, dim) = let
848 : jhr 3498 val expindex = specialize(alpha, 0)
849 :     in
850 :     E.EIN{
851 : jhr 3701 params = [E.FLD dim, mkNoSubstTEN []], index = alpha,
852 : jhr 3498 body = E.Probe(E.Field(0, expindex), E.Tensor(1, []))
853 :     }
854 :     end
855 : jhr 3476
856 : jhr 3477 (***************************** derivative ****************************)
857 : jhr 3476
858 : jhr 3477 (* \EinExp{\sum_{ij}\mathcal{E}_{ij} \frac{F_j}{\partial x_i} *)
859 :     val curl2d = E.EIN{
860 : jhr 3498 params = [E.FLD 2],
861 :     index = [],
862 :     body = E.Sum([(E.V 0, 0, 1), (E.V 1, 0, 1)],
863 :     E.Opn(E.Prod, [
864 : jhr 3512 E.Eps2(0, 1),
865 : jhr 3498 E.Apply(E.Partial[E.V 0], E.Field(0, [E.V 1]))
866 :     ]))
867 :     }
868 : jhr 3476
869 : jhr 3477 val curl3d = E.EIN{
870 : jhr 3701 params = [mkTEN [3]],
871 : jhr 3498 index = [3],
872 :     body = E.Sum([(E.V 1, 0, 2), (E.V 2, 0, 2)],
873 :     E.Opn(E.Prod, [
874 : jhr 3512 E.Epsilon(0, 1, 2),
875 : jhr 3498 E.Apply(E.Partial[E.V 1], E.Field(0, [E.V 2]))
876 :     ]))
877 :     }
878 : jhr 3476
879 : jhr 3477 (*< d F / d_i>_i *)
880 :     fun grad (alpha as a::_) = let
881 : jhr 3498 val expindex = specialize(alpha, 0)
882 :     in
883 :     E.EIN{
884 :     params = [E.FLD a],
885 :     index = alpha,
886 :     body = E.Apply(E.Partial expindex, E.Field(0, []))
887 :     }
888 :     end
889 : jhr 3477
890 :     (*< Sigma d F_alpha / d x_i>ALpha i CHANGE HERE *)
891 :     fun dotimes (dim, alpha) = let
892 : jhr 3498 val n = length alpha
893 : jhr 3477 (*
894 : jhr 3498 fun expIndex(n, inc) = List.tabulate(n, (fn(x)=>E.V (x+inc)))
895 :     val i' = expIndex(n, 0)
896 : jhr 3477 *)
897 : jhr 3498 val i' = List.tabulate (n, fn x => E.V x)
898 :     in
899 :     E.EIN{
900 :     params = [E.FLD dim], index = alpha@[dim],
901 :     body = E.Apply(E.Partial[E.V n], E.Field(0, i'))
902 :     }
903 :     end
904 : jhr 3477
905 :     (* <d F_i /d_i> *)
906 :     fun divergence (dim, alpha) = let
907 : jhr 3498 val expindex = specialize(alpha, 0)
908 :     val sumI = length alpha
909 :     val sumIndex = E.V sumI
910 :     val sumIndexL = [sumIndex]
911 :     val S = expindex@sumIndexL
912 :     in
913 :     E.EIN{
914 :     params = [E.FLD dim],
915 :     index = alpha,
916 :     body = E.Sum([(sumIndex, 0, dim-1)], E.Apply(E.Partial sumIndexL, E.Field(0, S)))
917 :     }
918 :     end
919 : jhr 3477
920 : jhr 3476 end (* mkOperators *)

root@smlnj-gforge.cs.uchicago.edu
ViewVC Help
Powered by ViewVC 1.0.0