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

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