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

SCM Repository

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

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

Parent Directory Parent Directory | Revision Log Revision Log


Revision 3969 - (download) (annotate)
Wed Jun 15 01:37:57 2016 UTC (2 years, 11 months ago) by cchiw
File size: 30782 byte(s)
fixed error in cleanIndex and ein-to-scalar
(* mk-operators.sml
 *
 * Functions to create the various Ein operators.
 *
 * This code is part of the Diderot Project (http://diderot-language.cs.uchicago.edu)
 *
 * COPYRIGHT (c) 2015 The University of Chicago
 * All rights reserved.
 *)

structure MkOperators : sig

    type dim = int
    type shape = dim list
    type ids = Ein.index_id list

    val addRR : Ein.ein
    val addTT : shape -> Ein.ein
    val addTF : dim * shape -> Ein.ein
    val addFF : dim * shape -> Ein.ein

    val subRR : Ein.ein
    val subTT : shape -> Ein.ein
    val subTF : dim * shape -> Ein.ein
    val subFT : dim * shape -> Ein.ein
    val subFF : dim * shape -> Ein.ein

    val mulRT : shape -> Ein.ein
    val mulRR : Ein.ein
    val mulRF : dim * shape -> Ein.ein
    val mulST : dim * shape -> Ein.ein
    val mulSS : dim -> Ein.ein
    val mulSF : dim * shape -> Ein.ein

    val divTR : shape -> Ein.ein
    val divRR : Ein.ein
    val divFR : dim * shape -> Ein.ein
    val divSS : dim -> Ein.ein
    val divFS : dim * shape -> Ein.ein

    val negTT : shape -> Ein.ein
    val negFF : dim * shape -> Ein.ein

    val cross2TT : Ein.ein
    val cross3TT : Ein.ein
    val cross2FF : Ein.ein
    val cross3FF : Ein.ein

    val outerTT : shape * shape -> Ein.ein
    val outerFF : dim * shape * shape -> Ein.ein
    val outerTF : dim * shape * shape -> Ein.ein
    val outerFT : dim * shape * shape -> Ein.ein

    val innerTT : shape * ids -> Ein.ein
    val innerFF : shape * dim * ids -> Ein.ein
    val innerFT : shape * dim * ids -> Ein.ein
    val innerTF : shape * dim * ids -> Ein.ein

    val colonTT : shape * ids -> Ein.ein
    val colonFF : dim * shape * ids -> Ein.ein
    val colonFT : dim * shape * ids -> Ein.ein
    val colonTF : dim * shape * ids -> Ein.ein

    val normT : shape -> Ein.ein
    val normF : dim * shape -> Ein.ein

    val normalizeTT : shape -> Ein.ein
    val normalizeFF : dim * shape -> Ein.ein

    val traceT : dim -> Ein.ein
    val traceF : dim * shape -> Ein.ein

    val transposeT : shape -> Ein.ein
(* QUESTION: should these be index_kind? *)
    val transposeF : dim * Ein.index_id * Ein.index_id -> Ein.ein

    val det2T : Ein.ein
    val det3T : Ein.ein
    val det2F : Ein.ein
    val det3F_full : Ein.ein
    val det3FA : Ein.ein
    val det3F : Ein.ein

    val expF : dim -> Ein.ein
    val expT : Ein.ein

    val powF  : dim * int -> Ein.ein

    val sqrtR : Ein.ein
    val sqrtF : dim -> Ein.ein
    val cosR  : Ein.ein
    val cosF  : dim -> Ein.ein
    val acosR : Ein.ein
    val acosF : dim -> Ein.ein
    val sinR  : Ein.ein
    val sinF  : dim -> Ein.ein
    val asinR : Ein.ein
    val asinF : dim -> Ein.ein
    val tanR  : Ein.ein
    val tanF  : dim -> Ein.ein
    val atanR : Ein.ein
    val atanF : dim -> Ein.ein

    val modulate : dim -> Ein.ein

    val identity : dim -> Ein.ein
    val zeros : shape -> Ein.ein

    val slice : shape * bool list * int list * shape -> Ein.ein

    val conv : dim * shape -> Ein.ein
    val probe : shape * dim -> Ein.ein

    val curl2d : Ein.ein
    val curl3d : Ein.ein
    val grad : shape -> Ein.ein
    val dotimes : dim * shape -> Ein.ein (* ?? *)
    val divergence : dim * shape -> Ein.ein

  end = struct

    structure E = Ein

    type dim = int
    type shape = dim list
    type ids = int list

  (* controls whether tensor operations should be *)
    val canSubst = true

  (* A constructor function for tensor variables that (by default) can be substituted for;
   * this behavior is controlled by the canSubst flag.
   *)
    fun mkTEN alpha = E.TEN(canSubst, alpha)

  (* a constructor function for tensor parameters that should never be substituted for.  *)
    fun mkNoSubstTEN alpha = E.TEN(false, alpha)

    fun specialize (alpha, inc) = List.mapi (fn (i, _) => E.V(i + inc)) alpha

    fun sumIds (n, inc, alpha) = let
        val vs = List.tabulate(n, fn v => E.V(v+inc))
        in
          ListPair.map (fn(v, i) => (v, 0, i-1)) (vs, alpha)
        end

    fun sumIds2 (n, i) = List.tabulate(n, fn v => (E.V v, 0, i))

  (******************************* Addition *****************************************)

  (* Adding tensors : < X{\alpha} + Y_{\alpha}>_{\alpha} *)
    fun addTT alpha = let
          val expindex = specialize(alpha, 0)
          in
            E.EIN{
                params = [mkTEN alpha, mkTEN alpha],
                index = alpha,
                body = E.Opn(E.Add, [E.Tensor(0, expindex), E.Tensor(1, expindex)])
              }
          end

    val addRR = addTT []

  (* Tensor and Fields *)
    fun addTF (dim, shape) =let
          val expindex = specialize(shape, 0)
          in
            E.EIN{
                params = [mkTEN shape, E.FLD dim],
                index = shape,
                body = E.Opn(E.Add, [E.Lift(E.Tensor(0, expindex)), E.Field(1, expindex)])
              }
          end

  (* Adding Fields : < F{\alpha} + G_{\alpha}>_{\alpha} *)
    fun addFF (dim, shape) =let
          val expindex = specialize(shape, 0)
          in
            E.EIN{
                params = [E.FLD dim, E.FLD dim],
                index = shape,
                body = E.Opn(E.Add, [E.Field(0, expindex), E.Field(1, expindex)])
              }
          end

  (********************************* Subtraction **************************************)

    fun subTT alpha = let
          val expindex = specialize(alpha, 0)
          in
            E.EIN{
                params = [mkTEN alpha, mkTEN alpha],
                index = alpha,
                body = E.Op2(E.Sub, E.Tensor(0, expindex), E.Tensor(1, expindex))
              }
          end

    val subRR = subTT []

    fun subTF (dim, shape) = let
          val expindex = specialize(shape, 0)
          in
            E.EIN{
                params = [mkTEN shape, E.FLD dim],
                index = shape,
                body = E.Opn(E.Add,
                  [E.Lift(E.Tensor(0, expindex)), E.Op1(E.Neg, E.Field(1, expindex))])
              }
          end

    fun subFT (dim, shape) = let
          val expindex = specialize(shape, 0)
          in
            E.EIN{
                params = [mkTEN shape, E.FLD dim],
                index = shape,
                body = E.Op2(E.Sub, E.Field(1, expindex), E.Lift(E.Tensor(0, expindex)))
              }
          end

    fun subFF (dim, shape) = let
          val expindex = specialize(shape, 0)
          in
            E.EIN{
                params = [E.FLD dim, E.FLD dim],
                index = shape,
                body = E.Op2(E.Sub, E.Field(0, expindex), E.Field(1, expindex))
              }
          end

  (********************************** Multiplication *************************************)

  (* scalar times tensor product: <s * T_{\alpha}>_{\alpha} *)
    fun mulRT alpha = let
          val expindex = specialize(alpha, 0)
          in
            E.EIN{
                params = [mkTEN [], mkTEN alpha],
                index = alpha,
                body = E.Opn(E.Prod, [E.Tensor(0, []),  E.Tensor(1, expindex)])
              }
          end

    val mulRR = mulRT []

    fun mulRF (dim, shape) =let
         val expindex = specialize(shape, 0)
         in
            E.EIN{
             params = [mkTEN [], E.FLD dim],
             index = shape,
             body = E.Opn(E.Prod, [E.Lift( E.Tensor(0, [])), E.Field(1,expindex)])
             }
         end

    fun mulST (dim, shape) =let
         val expindex = specialize(shape, 0)
         in
            E.EIN{
            params = [mkTEN shape, E.FLD dim],
            index = shape,
            body = E.Opn(E.Prod, [E.Lift( E.Tensor(0,expindex)), E.Field(1, [])])
            }
         end

    fun mulSS dim = E.EIN{
             params = [E.FLD dim, E.FLD dim],
             index = [],
             body = E.Opn(E.Prod, [E.Field(0, []), E.Field(1, [])])
         }

    fun mulSF(dim, shape) =let
         val  expindex= specialize(shape, 0)
         in E.EIN{
             params = [E.FLD dim, E.FLD dim],
             index = shape,
             body = E.Opn(E.Prod, [E.Field(0, []), E.Field(1, expindex)])
            }
         end

  (************************************ Division ************************************)

    fun divTR alpha = let
          val expindex = specialize(alpha, 0)
          in
            E.EIN{
                params = [mkTEN alpha, mkTEN []],
                index = alpha,
                body = E.Op2(E.Div, E.Tensor(0, expindex), E.Tensor(1, []))
              }
          end

    val divRR = divTR []

    fun divFR (dim, shape) = let
          val expindex = specialize(shape, 0)
          in
            E.EIN{
                params = [E.FLD dim, mkTEN []],
                index = shape,
                body = E.Op2(E.Div, E.Field(0, expindex), E.Lift(E.Tensor(1, [])))
              }
          end

    fun divSS dim = E.EIN{
            params = [E.FLD dim, E.FLD dim],
            index = [],
            body = E.Op2(E.Div, E.Field(0, []), E.Field(1, []))
          }

    fun divFS(dim, shape) = let
          val expindex = specialize(shape, 0)
          in
            E.EIN{
                params = [E.FLD dim, E.FLD dim],
                index = shape,
                body = E.Opn(E.Prod, [E.Field(0, expindex), E.Op2(E.Div, E.Const 1, E.Field(1, []))])
              }
          end

  (************************************* Negation **********************************)

    fun negTT alpha = let
          val expindex= specialize(alpha, 0)
            (*changed tensor lift variable here *)
          in
            E.EIN {
                params = [mkTEN alpha], index = alpha,
                body = E.Op1(E.Neg, E.Tensor(0, expindex))
              }
          end

    fun negFF (dim, shape) = let
          val expindex = specialize(shape, 0)
          in
            E.EIN{
                params = [E.FLD dim], index = shape,
                body = E.Op1(E.Neg, E.Field(0, expindex))
              }
          end

  (****************************** cross product ***********************************)

  (* 2-d cross product Eps_{ij}U_i V_j *)
    val cross2TT = E.EIN{
             params = [mkTEN [2], mkTEN [2]],
             index = [],
             body = E.Sum([(E.V 0, 0, 1), (E.V 1, 0, 1)],
             E.Opn(E.Prod, [E.Eps2(0, 1), E.Tensor(0, [E.V 0]), E.Tensor(1, [E.V 1])]))
          }

  (* crossProduct is on 3D vectors ..vec3 t8=t0 × t1; *)
    val cross3TT = E.EIN{
             params = [mkTEN [3], mkTEN [3]],
             index = [3],
             body = E.Sum([(E.V 1, 0, 2), (E.V 2, 0, 2)],
                E.Opn(E.Prod, [
                    E.Epsilon(0, 1, 2),
                    E.Tensor(0, [E.V 1]),
                    E.Tensor(1, [E.V 2])
                  ]))
          }

  (* Field Cross Product *)
    val cross2FF = E.EIN{
            params = [E.FLD(2), E.FLD(2)], index = [],
            body = E.Sum([(E.V 0, 0, 1), (E.V 1, 0, 1)],
                E.Opn(E.Prod, [E.Eps2(0, 1), E.Field(0, [E.V 0]), E.Field(1, [E.V 1])]))
          }

  (* Field Cross Product *)
    val cross3FF = E.EIN{
            params = [E.FLD(3), E.FLD(3)], index= [3],
            body = E.Sum([(E.V 1, 0, 2), (E.V 2, 0, 2)],
                E.Opn(E.Prod, [E.Epsilon(0, 1, 2),
                E.Field(0, [E.V 1]), E.Field(1, [E.V 2])]))
          }

  (******************** outer product ********************************)

  fun outerTT (alpha, beta) = let
        val expIdxA = specialize (alpha, 0)
        val expIdxB = specialize (beta, length alpha)
        in
          E.EIN{
              params = [mkTEN alpha, mkTEN beta],
              index = alpha@beta,
              body = E.Opn(E.Prod, [E.Tensor(0, expIdxA), E.Tensor(1, expIdxB)])
            }
        end

    (*Assumes same dimension vector field *)
    fun outerFF (dim, alpha, beta) =let
        val expIdxA = specialize (alpha, 0)
        val expIdxB = specialize (beta, length alpha)
        in
          E.EIN{
              params = [E.FLD dim, E.FLD dim],
              index = alpha@beta,
              body = E.Opn(E.Prod, [E.Field(0, expIdxA), E.Field(1, expIdxB)])
            }
        end

    fun outerTF (dim, alpha, beta) =let
        val expIdxA = specialize (alpha, 0)
        val expIdxB = specialize (beta, length alpha)
        in
          E.EIN{
              params = [mkTEN alpha, E.FLD dim],
              index = alpha@beta,
              body = E.Opn(E.Prod, [E.Lift(E.Tensor(0, expIdxA)), E.Field(1, expIdxB)])
            }
        end

    fun outerFT (dim, alpha, beta) =let
        val expIdxA = specialize(alpha, 0)
        val expIdxB = specialize(beta, length alpha)
        in
          E.EIN{
              params = [E.FLD dim, mkTEN alpha],
              index = alpha@beta,
              body = E.Opn(E.Prod, [E.Field(0, expIdxA), E.Lift(E.Tensor(1, expIdxB))])
            }
        end

  (*************************** inner product **********************************)
    (* generic inner product: <T_{\alpha i} * T_{i \beta}>_{\alpha \beta} *)
    fun innerTT (shape1, i::beta) = let
          val alpha = List.take(shape1, length shape1 - 1)
          val expindexA = specialize(alpha, 0)
          val expindexB = specialize(beta, length alpha)
          val s' = E.V(length alpha + length beta)
          val s'' = [(s', 0, i-1)]
          in
            E.EIN{
                params = [mkTEN shape1, mkTEN(i :: beta)],
                index = alpha@beta,
                body = E.Sum(s'', E.Opn(E.Prod, [
                    E.Tensor(0, expindexA@[s']),   (* T_{\alpha i} *)
                    E.Tensor(1, [s']@expindexB )   (* T'_{i \beta} *)
                  ]))
              }
          end
     | innerTT _ = raise Fail "Wrong shape for inner product"

  (* generic inner product: <T_{\alpha i} * T_{i \beta}>_{\alpha \beta} *)
    fun innerFF (shape1, dim, i::beta) = let
          val alpha = List.take(shape1, length shape1 - 1)
          val expindexA = specialize(alpha, 0)
          val expindexB = specialize(beta, length alpha)
          val sid = E.V(length alpha + length beta)
          in
            E.EIN{
                params = [E.FLD dim, E.FLD dim],
                index = alpha @ beta,
                body = E.Sum([(sid, 0, i-1)],
                  E.Opn(E.Prod, [
                      E.Field(0, expindexA @ [sid]),   (* F_{\alpha i} *)
                      E.Field(1, [sid] @ expindexB)    (* F'_{i \beta} *)
                    ]))
              }
          end
      | innerFF _ = raise Fail "Wrong shape for innerProductField"

    fun innerFT (shape1, dim, i::beta) = let
          val alpha = List.take(shape1, length shape1-1)
          val expindexA = specialize(alpha, 0)
          val expindexB = specialize(beta, length alpha)
          val sid = E.V(length alpha + length beta)
          in
            E.EIN{
                params = [E.FLD dim, mkTEN(i::beta)],
                index = alpha @ beta,
                body = E.Sum([(sid, 0, i-1)],
                  E.Opn(E.Prod, [
                      E.Field(0, expindexA @ [sid]),           (* F_{\alpha i} *)
                      E.Lift(E.Tensor(1, [sid] @ expindexB ))  (* F'_{i \beta} *)
                    ]))
              }
          end
      | innerFT _ = raise Fail "Wrong shape for innerProductFieldTensor"

    fun innerTF (shape1, dim, i::beta) = let
          val alpha = List.take(shape1, length shape1 - 1)
          val expindexA = specialize(alpha, 0)
          val expindexB = specialize(beta, length alpha)
          val sid = E.V(length(alpha) + length beta)
          in
            E.EIN{
                params = [mkTEN shape1, E.FLD dim],
                index = alpha @ beta,
                body = E.Sum([(sid, 0, i-1)],
                  E.Opn(E.Prod, [
                      E.Lift(E.Tensor(0, expindexA @ [sid])),   (* F_{\alpha i} *)
                      E.Field(1, [sid] @ expindexB)             (* F'_{i \beta} *)
                    ]))
              }
          end
      | innerTF _ = raise Fail "Wrong shape for innerProductTensorField"

  (*************************** colon product **********************************)

  (* <T_{\alpha i j} * B{i j \beta }>_\alpha \beta *)
    fun colonTT (shape1, i::j::beta) = let
          val lenAlpha = length shape1 - 2
          val alpha = List.take(shape1, lenAlpha)
          val expindexA = specialize(alpha, 0)
          val expindexB = specialize(beta, lenAlpha)
          val sumi = lenAlpha + length beta
          val s' = [E.V sumi, E.V(sumi+1)]
          val sx = [(E.V sumi, 0, i-1), (E.V(sumi+1), 0, j-1)]
          in
            E.EIN{
                params = [mkTEN shape1, mkTEN(i::j::beta)],
                index = alpha@beta,
                body = E.Sum(sx,
                  E.Opn(E.Prod, [E.Tensor(0, expindexA@s'), E.Tensor(1, s'@expindexB)]))
              }
          end

  (* <F_{\alpha i j} * G_{i j \beta }>_\alpha \beta *)
    fun colonFF (dim, shape1, i::j::beta) = let
          val lenAlpha = length shape1 - 2
          val alpha = List.take(shape1, lenAlpha)
          val expindexA = specialize(alpha, 0)
          val expindexB = specialize(beta, lenAlpha)
          val sumi = lenAlpha + length beta
          val s' = [E.V sumi, E.V(sumi+1)]
          val sx = [(E.V sumi, 0, i-1), (E.V(sumi+1), 0, j-1)]
          in
            E.EIN{
                params = [E.FLD dim, E.FLD dim],
                index = alpha@beta,
                body = E.Sum(sx,
                  E.Opn(E.Prod, [E.Field(0, expindexA@s'), E.Field(1, s'@expindexB)]))
              }
          end


  (* <F_{\alpha i j} * T_{i j \beta }>_\alpha \beta *)
    fun colonFT (dim, shape1, i::j::beta) = let
          val lenAlpha = length shape1 - 2
          val alpha = List.take(shape1, lenAlpha)
          val expindexA = specialize(alpha, 0)
          val expindexB = specialize(beta, lenAlpha)
          val sumi = lenAlpha + length beta
          val s' = [E.V sumi, E.V(sumi+1)]
          val sx = [(E.V sumi, 0, i-1), (E.V(sumi+1), 0, j-1)]
          in
            E.EIN{
                params = [E.FLD dim, mkTEN shape1],
                index = alpha@beta,
                body = E.Sum(sx,
                  E.Opn(E.Prod, [E.Field(0, expindexA@s'), E.Lift(E.Tensor(1, s'@expindexB))]))
              }
          end

  (* <T_{\alpha i j} * G{i j \beta }>_\alpha \beta *)
    fun colonTF (dim, shape1, i::j::beta) = let
          val lenAlpha = length shape1 - 2
          val alpha = List.take(shape1, lenAlpha)
          val expindexA = specialize(alpha, 0)
          val expindexB = specialize(beta, lenAlpha)
          val sumi = lenAlpha + length beta
          val s' = [E.V sumi, E.V(sumi+1)]
          val sx = [(E.V sumi, 0, i-1), (E.V(sumi+1), 0, j-1)]
          in
            E.EIN{
                params = [mkTEN(i::j::beta), E.FLD dim],
                index = alpha@beta,
                body = E.Sum(sx,
                  E.Opn(E.Prod, [E.Lift(E.Tensor(0, expindexA@s')), E.Field(1, s'@expindexB)]))
              }
          end

  (******************** Norm  ********************************)

    fun normT alpha = let
          val expIdx = specialize(alpha, 0)
          val sx = sumIds(length alpha, 0, alpha)
          in
            E.EIN{
                params = [mkTEN alpha],
                index = [],
                body = E.Op1(E.Sqrt,
                  E.Sum(sx, E.Opn(E.Prod, [E.Tensor(0, expIdx), E.Tensor(0, expIdx)])))
              }
          end

    fun normF (dim, alpha) =let
          val expIdx = specialize(alpha, 0)
          val sx = sumIds(length alpha, 0, alpha)
          in
            E.EIN{
                params = [E.FLD dim],
                index = [],
                body = E.Op1(E.Sqrt,
                  E.Sum(sx, E.Opn(E.Prod, [E.Field(0, expIdx), E.Field(0, expIdx)])))
              }
          end

    fun normalizeTT alpha = let
          val expindex = specialize(alpha, 0)
          val len = length alpha
          val expindexDot = specialize(alpha, len)
          val sx = sumIds(len, len, alpha)
          val f = E.Tensor(0, expindex)
          val g = E.Tensor(1, expindexDot)
          in
            E.EIN{
                params = [mkTEN alpha, mkTEN alpha],
                index = alpha,
                body = E.Opn(E.Prod, [
                    f, E.Op2(E.Div, E.Const 1, E.Op1(E.Sqrt, E.Sum(sx, E.Opn(E.Prod, [g, g]))))
                  ])
              }
          end

    fun normalizeFF (dim, alpha as i::_) = let
          val expindex = specialize(alpha, 0)
          val len = length alpha
          val expindexDot = specialize(alpha, len)
          val sx = sumIds(len, len, alpha)
          val f = E.Field(0, expindex)
          val g = E.Field(1, expindexDot)
          in
            E.EIN{
                params = [E.FLD dim, E.FLD dim],
                index = alpha,
                body = E.Opn(E.Prod, [
                    f, E.Op2(E.Div, E.Const 1, E.Op1(E.Sqrt, E.Sum(sx, E.Opn(E.Prod, [g, g]))))
                  ])
              }
          end

  (************************* trace *************************)

  (* Trace: <M_{i, i}>  This one Sx represents both i's*)
    fun traceT dim = E.EIN{
            params = [mkTEN [dim, dim]], index = [],
            body = E.Sum([(E.V 0, 0, dim-1)], E.Tensor(0, [E.V 0, E.V 0]))
          }

  (* Trace: <Sigma_i F_{\alpha i, i}>  This one Sx represents both i's *)
    fun traceF (dim, alpha) = let
          val expindex = specialize(alpha, 0)
          val s = E.V(length alpha)
          in
            E.EIN{
                params = [E.FLD dim],
                index = alpha,
                body = E.Sum([(s, 0, dim-1)], E.Field(0, expindex@[s, s]))
              }
          end

  (************************* tranpose *************************)

    fun transposeT alpha = E.EIN{
            params = [mkTEN alpha],
            index = List.rev alpha,
            body = E.Tensor(0, [E.V 1, E.V 0])
          }

  (* Transpose Field F_{ji} *)
    fun transposeF (dim, i, j) = E.EIN{
            params = [E.FLD dim],
            index = [i, j],
            body = E.Field(0, [E.V 1, E.V 0])
          }

  (************************* determinant *************************)

    val det2T = E.EIN{
            params = [mkNoSubstTEN [2, 2]],
            index = [],
            body = E.Op2(E.Sub,
              E.Opn(E.Prod, [E.Tensor(0, [E.C 0, E.C 0]), E.Tensor(0, [E.C 1, E.C 1])]),
              E.Opn(E.Prod, [E.Tensor(0, [E.C 0, E.C 1]), E.Tensor(0, [E.C 1, E.C 0])]))
          }

    val det3T = let
          val a = E.Tensor(0, [E.C 0, E.C 0])
          val b = E.Tensor(0, [E.C 0, E.C 1])
          val c = E.Tensor(0, [E.C 0, E.C 2])
          val d = E.Tensor(0, [E.C 1, E.C 0])
          val e = E.Tensor(0, [E.C 1, E.C 1])
          val f = E.Tensor(0, [E.C 1, E.C 2])
          val g = E.Tensor(0, [E.C 2, E.C 0])
          val h = E.Tensor(0, [E.C 2, E.C 1])
          val i = E.Tensor(0, [E.C 2, E.C 2])
          in
            E.EIN{
                params = [mkNoSubstTEN [3, 3]],
                index = [],
                body = E.Op2(E.Sub,
                  E.Opn(E.Add, [
                      E.Opn(E.Prod, [a, e, i]), E.Opn(E.Prod, [b, f, g]), E.Opn(E.Prod, [c, d, h])
                    ]),
                  E.Opn(E.Add, [
                      E.Opn(E.Prod, [c, e, g]), E.Opn(E.Prod, [b, d, i]), E.Opn(E.Prod, [a, f, h])
                    ]))
              }
          end

    val det2F = E.EIN{
            params = [E.FLD 2],
            index = [],
            body = E.Op2(E.Sub,
              E.Opn(E.Prod, [E.Field(0, [E.C 0, E.C 0]), E.Field(0, [E.C 1, E.C 1])]),
              E.Opn(E.Prod, [E.Field(0, [E.C 0, E.C 1]), E.Field(0, [E.C 1, E.C 0])]))
          }

    val det3F_full = let
          val a = E.Field(0, [E.C 0, E.C 0])
          val b = E.Field(0, [E.C 0, E.C 1])
          val c = E.Field(0, [E.C 0, E.C 2])
          val d = E.Field(0, [E.C 1, E.C 0])
          val e = E.Field(0, [E.C 1, E.C 1])
          val f = E.Field(0, [E.C 1, E.C 2])
          val g = E.Field(0, [E.C 2, E.C 0])
          val h = E.Field(0, [E.C 2, E.C 1])
          val i = E.Field(0, [E.C 2, E.C 2])
          in
            E.EIN{
                params = [E.FLD 3],
                index = [],
                body = E.Op2(E.Sub,
                  E.Opn(E.Add, [
                      E.Opn(E.Prod, [a, e, i]), E.Opn(E.Prod, [b, f, g]), E.Opn(E.Prod, [c, d, h])
                    ]),
                  E.Opn(E.Add, [
                      E.Opn(E.Prod, [c, e, g]), E.Opn(E.Prod, [b, d, i]), E.Opn(E.Prod, [a, f, h])
                    ]))
              }
          end

    val det3FA = E.EIN{
            params = [E.FLD 3],
            index = [],
            body = E.Sum([(E.V 0, 0, 2), (E.V 1, 0, 2), (E.V 2, 0, 2)],
              E.Opn(E.Prod, [
                  E.Epsilon(0, 1, 2),
                  E.Field(0, [E.C 0, E.V 0]),
                  E.Field(0, [E.C 1, E.V 1]),
                  E.Field(0, [E.C 2, E.V 2])
                ]))
          }

    val det3F = E.EIN{
            params = [E.FLD 3],
            index = [],
            body = E.Sum([(E.V 0, 0, 2)],
              E.Opn(E.Prod, [
                  E.Field(0, [E.C 0, E.V 0]),
                  E.Sum([(E.V 1, 0, 2)],
                    E.Opn(E.Prod, [
                        E.Field(0, [E.C 1, E.V 1]),
                        E.Sum([(E.V 2, 0, 2)],
                          E.Opn(E.Prod, [E.Epsilon(0, 1, 2), E.Field(0, [E.C 2, E.V 2])]))
                      ]))
                ]))
          }

  (************************* Exponential **************************)
    fun expF dim = E.EIN{params = [E.FLD dim], index = [], body = E.Op1(E.Exp, E.Field(0, []))}
    val expT = E.EIN{params = [mkNoSubstTEN []], index = [], body = E.Op1(E.Exp, E.Tensor(0, []))}

  (************************* Lifted single-argument math functions *************************)
    local
      fun tensorFn rator = E.EIN{
              params = [mkNoSubstTEN []],
              index = [],
              body = E.Op1(rator, E.Tensor(0, []))
            }
      fun liftFn rator dim = E.EIN{
              params = [E.FLD dim],
              index = [],
              body = E.Op1(rator, E.Field(0, []))
            }
    in
    fun powF (dim, n) = E.EIN{
            params = [E.FLD dim],
            index = [], body = E.Op1(E.PowInt n, E.Field(0, []))
          }
    val sqrtR = tensorFn E.Sqrt
    val sqrtF = liftFn E.Sqrt
    val cosR  = tensorFn E.Cosine
    val cosF  = liftFn E.Cosine
    val acosR = tensorFn E.ArcCosine
    val acosF = liftFn E.ArcCosine
    val sinR  = tensorFn E.Sine
    val sinF  = liftFn E.Sine
    val asinR = tensorFn E.ArcSine
    val asinF = liftFn E.ArcSine
    val tanR  = tensorFn E.Tangent
    val tanF  = liftFn E.Tangent
    val atanR = tensorFn E.ArcTangent
    val atanF = liftFn E.ArcTangent
    end (* local *)

  (************************* other tensor ops *************************)
    fun modulate dim = E.EIN{
            params = [mkTEN [dim], mkTEN [dim]],
            index = [dim],
            body = E.Opn(E.Prod, [E.Tensor(0, [E.V 0]), E.Tensor(1, [E.V 0])])
          }

    fun identity dim = E.EIN{
            params = [], index = [dim, dim], body = E.Delta(E.V(0), E.V(1))
          }

    fun zeros shape = E.EIN{
            params = [], index = shape, body = E.Const 0
          }

(* QUESTION: why do we need the const list?  The indices are implicit in the position of
 * of the mask element!  Likewise, the result type can be determined from the argTy and
 * mask.
 *)
    fun slice (argTy, mask, const, rstTy) = let
          fun iter ([], _, cnt) = []
            | iter (true::es, c::cs, cnt) = E.C c :: iter(es, cs, cnt)
            | iter (false::es, cs, cnt) = E.V cnt :: iter(es, cs, cnt+1)
          val ix = iter(mask, const, 0)
          in
            E.EIN{
                params = [mkTEN argTy],
                index = rstTy,
                body = E.Tensor(0, ix)
              }
          end

  (******************** other field ops  ********************************)

  (* FLD here is bounded to image field, and dimension of h *)
    fun conv (dim, shape) =let
          val expindex = specialize(shape, 0)
          in
            E.EIN{
                params = [E.IMG(dim, shape), E.KRN],
                index = shape,
                body = E.Conv(0,expindex, 1, [])
              }
          end

  (* Probe: <F(x)>_{\alpha}   *)
    fun probe (alpha, dim) = let
          val expindex = specialize(alpha, 0)
          in
            E.EIN{
                params = [E.FLD dim, mkNoSubstTEN []], index = alpha,
                body = E.Probe(E.Field(0, expindex), E.Tensor(1, []))
              }
          end

  (***************************** derivative ****************************)

  (* \EinExp{\sum_{ij}\mathcal{E}_{ij} \frac{F_j}{\partial x_i} *)
    val curl2d = E.EIN{
            params = [E.FLD 2],
            index = [],
            body = E.Sum([(E.V 0, 0, 1), (E.V 1, 0, 1)],
              E.Opn(E.Prod, [
                  E.Eps2(0, 1),
                  E.Apply(E.Partial[E.V 0], E.Field(0, [E.V 1]))
                ]))
          }

    val curl3d = E.EIN{
            params = [mkTEN [3]],
            index = [3],
            body = E.Sum([(E.V 1, 0, 2), (E.V 2, 0, 2)],
              E.Opn(E.Prod, [
                  E.Epsilon(0, 1, 2),
                  E.Apply(E.Partial[E.V 1], E.Field(0, [E.V 2]))
                ]))
          }

  (*< d F /  d_i>_i  *)
    fun grad (alpha as a::_) = let
          val expindex = specialize(alpha, 0)
          in
            E.EIN{
                params = [E.FLD a],
                index = alpha,
                body = E.Apply(E.Partial expindex, E.Field(0, []))
              }
          end

  (*< Sigma d F_alpha /  d x_i>ALpha  i CHANGE HERE *)
    fun dotimes (dim, alpha) = let
          val n = length alpha
(*
          fun expIndex(n, inc) = List.tabulate(n, (fn(x)=>E.V (x+inc)))
          val i' = expIndex(n, 0)
*)
          val i' = List.tabulate (n, fn x => E.V x)
          in
            E.EIN{
                params = [E.FLD dim], index = alpha@[dim],
                body = E.Apply(E.Partial[E.V n], E.Field(0, i'))
              }
          end

  (*  <d F_i /d_i> *)
    fun divergence (dim, alpha) = let
          val expindex = specialize(alpha, 0)
          val sumI = length alpha
          val sumIndex = E.V sumI
          val sumIndexL = [sumIndex]
          val S = expindex@sumIndexL
          in
            E.EIN{
                params = [E.FLD dim],
                index = alpha,
                body = E.Sum([(sumIndex, 0, dim-1)], E.Apply(E.Partial sumIndexL, E.Field(0, S)))
              }
          end

  end (* mkOperators *)

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