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

SCM Repository

[diderot] View of /branches/charisee/src/compiler/ein/mkoperators.sml
ViewVC logotype

View of /branches/charisee/src/compiler/ein/mkoperators.sml

Parent Directory Parent Directory | Revision Log Revision Log


Revision 2870 - (download) (annotate)
Wed Feb 25 21:47:43 2015 UTC (4 years, 7 months ago) by cchiw
File size: 18764 byte(s)
added sqrt,pow, and examples
(* creates EIN operators 
 *
 * COPYRIGHT (c) 2012 The Diderot Project (http://diderot-language.cs.uchicago.edu)
 * All rights reserved.
 *)

structure mkOperators = struct

    local

    structure E = Ein
    structure P=Printer
    in
 
    fun specialize(alpha,inc)=  List.tabulate(length(alpha), (fn(x)=>E.V (x+inc)))
    fun sumIds(n,i)=List.tabulate(n, (fn v=>(E.V v, 0, i)))
 
    (* Adding tensors : < X{\alpha} + Y_{\alpha}>_{\alpha} *)
    fun addTen alpha =let
        val expindex= specialize(alpha,0)
        in 
            E.EIN{
                params = [E.TEN(1,alpha), E.TEN(1,alpha)],
                index = alpha,
                body = E.Add[E.Tensor(0, expindex), E.Tensor(1, expindex)]
            }
        end
        
    fun createVec dim=
        E.EIN{params = [E.TEN(1,[dim])], index = [dim], body = E.Tensor(0, [E.V 0])}

     val zero=
        E.EIN{params = [], index = [], body = E.Const(0)}
        
     fun subTen alpha=let
        val expindex= specialize(alpha,0)
        in
            E.EIN{
                params = [E.TEN(1,alpha), E.TEN(1,alpha)],
                index = alpha,
                body = E.Sub(E.Tensor(0, expindex), E.Tensor(1, expindex))
            }
        end
    
    fun divTen alpha =let
        val expindex= specialize(alpha,0)
        in
            E.EIN{
                params = [E.TEN(1,alpha), E.TEN(1,[])],
                index = alpha,
                body = E.Div(E.Tensor(0, expindex), E.Tensor(1,[]))
            }
        end
        
    (* Trace: <M_{i, i}>  This one Sx represents both i's*)
    fun trace dim = E.EIN{
            params = [E.TEN(1,[dim,dim])],
            index = [],
            body = E.Sum([(E.V 0,0,dim-1)],E.Tensor(0, [E.V 0, E.V 0]))
        }
        
    fun negTen alpha=let
        val  expindex= specialize(alpha,0)
        in
            E.EIN{
                params = [E.TEN(1,alpha)],
                index = alpha,
                body = E.Neg(E.Tensor(0, expindex))
            }
        end
        
    (* scalar times tensor product: <s * T_{\alpha}>_{\alpha} *)
    fun scaleTen alpha = let
        val expindex= specialize(alpha,0)
        in
            E.EIN{
                params = [E.TEN(1,[]), E.TEN(1,alpha)],
                index = alpha,
                body = E.Prod[ E.Tensor(0, []),  E.Tensor(1, expindex)]
            }
        end 

    fun dot i =
        E.EIN{
                params = [E.TEN(1,[i]) ,E.TEN(1,[i])],
                index = [],
                body = E.Sum([(E.V 0,0,i-1)], E.Prod[E.Tensor(0, [E.V 0]), E.Tensor(1, [E.V 0])])
            }


    (* generic inner product: <T_{\alpha i} * T_{i \beta}>_{\alpha \beta} *)
    fun innerProduct(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 = [E.TEN(1,shape1) ,E.TEN(1,i::beta)],
                index = alpha@beta,
                body = E.Sum(s'', E.Prod[
                    E.Tensor(0, expindexA@[s']),   (* T_{\alpha i} *)
                    E.Tensor(1, [s']@expindexB )  (* T'_{i \beta} *)])
            }
        end
        | innerProduct _ = raise Fail "Wrong shape for inner product"

    (*<T_{\alpha i j} * B{i j \beta }>_\alpha \beta*)
    fun doubleDot(shape1,i::j::beta) = let
        val alpha= List.take(shape1,length(shape1)-2)
        val expindexA= specialize(alpha,0)
        val expindexB= specialize(beta,(length(alpha)))
        val sumi=length(alpha)+ length(beta)
        val s'=[E.V sumi,E.V(sumi+1)]
        val s''=[(E.V sumi,0,i-1),(E.V(sumi+1),0,j-1)]
        in
            E.EIN{
                params = [E.TEN(1,shape1),E.TEN(1,i::j::beta)],
                index = alpha@beta,
                body = E.Sum(s'',E.Prod[ E.Tensor(0, expindexA@s'), E.Tensor(1,s'@expindexB)])
            }
        end
    | doubleDot _ = raise Fail "Wrong shape for double dot "
        
    (*Vector Examples : <T_i * T_j>_ij..t0⊗t1*)
    fun outerProduct(dimA,dimB) =
        E.EIN{
            params = [E.TEN(1,[dimA]), E.TEN(1,[dimB])],
            index= [dimA,dimB],
            body= E.Prod[E.Tensor(0, [E.V 0]), E.Tensor(1, [E.V 1])]
        }
       
        
    fun transpose alpha =E.EIN{
            params = [E.TEN(1,alpha)],
            index= List.rev alpha,
            body= E.Tensor(0, [E.V 1,E.V 0])
      }
   
    fun modulate dim =E.EIN{
            params = [E.TEN(1,[dim]), E.TEN(1,[dim])],
            index = [dim],
            body = E.Prod[E.Tensor(0, [E.V 0]), E.Tensor(1, [E.V 0])]
        }
        
    (*get norm, but without the sqrt
    * implemented as a summation over a modulate
    *)
    fun norm alpha =let
        val i=List.hd alpha
        val expindex= specialize(alpha,0)
        val sx= sumIds(length(alpha),i-1)
        in E.EIN{
            params = [E.TEN(1,alpha), E.TEN(1,alpha)],
            index = [],
            body = E.Sum(sx,
                E.Prod[E.Tensor(0, expindex), E.Tensor(1, expindex)])
            }
        end
        
    fun norm2 dim =let
        in E.EIN{
        params = [ E.TEN(1,[]),E.TEN(1,[dim])],
        index = [dim],
        body =E.Prod[ E.Div(E.Const 1,E.Tensor(0, [])),E.Tensor(1, [E.V 0])]
        }
        end
        
        
    (*crossProduct is on 3D vectors ..vec3 t8=t0 × t1; *)
    val crossProduct = E.EIN{
            params = [E.TEN(1,[3]), E.TEN(1,[3])],
            index= [3],
            body=E.Sum([(E. V 1,0,2),(E.V 2,0,2)],
                E.Prod[ E.Epsilon(0, 1, 2), E.Tensor(0, [E.V 1]),  E.Tensor(1, [E.V 2 ]) ])
        }
        
    (*2-d cross product Eps_{ij}U_i V_j*)
    val crossProduct2 = E.EIN{
            params = [E.TEN(1,[2]), E.TEN(1,[2])],
            index= [],
            body=E.Sum([(E. V 0,0,1),(E.V 1,0,1)],
                E.Prod[ E.Eps2(0, 1), E.Tensor(0, [E.V 0]),  E.Tensor(1, [E.V 1])])
        }

        
    (* Identiy: <\delta_{i j}>_{i j}  *)
    fun identity dim =E.EIN{
            params = [],index = [dim,dim], body = E.Delta(E.V(0), E.V(1))
        }
            
    (*Tensor and Fields*)
    fun addTenField dim = E.EIN{
            params = [E.TEN(1,[]),E.FLD(dim)],
            index = [],
            body = E.Add[E.Lift(E.Tensor(0, [])),E.Field(1, [])]
        }
        
    fun subTenField dim = E.EIN{
            params = [E.TEN(1,[]),E.FLD(dim)],
            index = [],
            body = E.Add[E.Lift(E.Tensor(0, [])),E.Neg(E.Field(1, []))]
        }
    
    fun subFieldTen dim = E.EIN{
            params = [E.TEN(1,[]),E.FLD(dim)],
            index = [],
            body = E.Sub(E.Field(1, []),E.Lift(E.Tensor(0, [])))
        }
    
    (* mkField functions*)
    (*Adding Fields : < F{\alpha} + G_{\alpha}>_{\alpha} *)
    fun addField(dim,shape) =let
        val expindex= specialize(shape,0)
        in E.EIN{
            params = [E.FLD(dim),E.FLD(dim)],
            index = shape,
            body = E.Add[E.Field(0, expindex),E.Field(1, expindex)]
        }
        end
        
    fun subField(dim,shape) =let
        val expindex= specialize(shape,0)
        in E.EIN{
            params = [E.FLD(dim),E.FLD(dim)],
            index = shape,
            body = E.Sub(E.Field(0, expindex),E.Field(1, expindex))
        }
        end
    
    fun scaleField(dim,shape) =let
        val  expindex= specialize(shape,0)
        in E.EIN{
            params = [E.TEN(1,[]),E.FLD(dim)],
            index = shape,
            body = E.Prod[E.Lift( E.Tensor(0,[])), E.Field(1,expindex)]
        }
        end
    
    fun divideField(dim,shape) = let
        val expindex= specialize(shape,0)
        in E.EIN{
            params = [E.FLD(dim),E.TEN(1,[])],
            index = shape,
            body = E.Div(E.Field(0, expindex),E.Lift(  E.Tensor(1, [])))
        }
        end
         
    fun negField(dim,shape) = let
        val expindex = specialize(shape,0)
        in E.EIN{
            params = [E.FLD(dim)],
            index = shape,
            body = E.Neg(E.Field(0, expindex))
        }
        end

    (*< d F /  d_i>_i  *)
    fun grad alpha=let
        val a=List.hd(alpha)
        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)
        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
        
    (*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),E.TEN(0,[])],
            index= alpha,
            body= E.Probe(E.Field(0, expindex), E.Tensor(1,[]))
        }
        end

    (*(F_y/dx - F_x/dy )k*)
    (*val curl2d=E.EIN{
        params = [E.FLD 2],
        index = [],
        body = E.Sub(E.Apply(E.Partial([E.C 0]), E.Field(0,[E.C 1])),
                     E.Apply(E.Partial([E.C 1]), E.Field(0,[E.C 0])))
    }*)
        
    (*\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.Prod[E.Eps2(0, 1),E.Apply( E.Partial([E.V 0]), E.Field(0,[E.V 1]))])
        }
        
    val curl3d=E.EIN{
            params = [E.TEN(1,[3])],
            index = [3],
            body = E.Sum([(E.V 1,0,2), (E.V 2,0,2)],E.Prod[E.Epsilon(0, 1, 2),
                E.Apply( E.Partial([E.V 1]), E.Field(0,[E.V 2]))])
        }
    
    (*Scalars*)
    (* Add Scalars*)
    val addScalar = E.EIN{
            params = [E.TEN(1,[]), E.TEN(1,[])] ,
            index = [],
            body = E.Add[ E.Tensor(0, []), E.Tensor(1, [])]
        }
        
    (* Subtract Scalars*)
    val subScalar = E.EIN{
            params = [E.TEN(1,[]), E.TEN(1,[])],
            index = [],
            body = E.Sub( E.Tensor(0, []), E.Tensor(1, []))
        }

    (* Divide Scalars*)
    val divScalar = E.EIN{
            params = [E.TEN(1,[]), E.TEN(1,[])],
            index = [],
            body = E.Div( E.Tensor(0, []), E.Tensor(1, []))
        }
        
    (* Product Scalars*)
    val prodScalar = E.EIN{
            params =[E.TEN(1,[]), E.TEN(1,[])],
            index = [],
            body = E.Prod[ E.Tensor(0, []), E.Tensor(1, [])]
        }
        
    (*Transform M_ij x_j+T*)
    fun transform(i, j) = E.EIN{
            params = [E.TEN(1,[i,j]), E.TEN(1,[j]), E.TEN(1,[j])],
            index = [i],
            body = E.Add[E.Sum([(E.V 1, 0,j-1)],E.Prod[E.Tensor(0, [E.V 0, E.V 1]),
                E.Tensor(1, [E.V 1])]), E.Tensor(2,[E.V 0])]
        }
 
    fun transformA(i, j) = E.EIN{
            params = [E.TEN(1,[i,j]), E.TEN(1,[j])],
            index = [i],
            body = E.Sum([(E.V 1, 0,j-1)],E.Prod[E.Tensor(0, [E.V 0, E.V 1]), E.Tensor(1, [E.V 1])])
        }

    fun transformB i = E.EIN{
            params = [E.TEN(1,[i]), E.TEN(1,[i])],
            index = [i],
            body = E.Add[E.Tensor(0, [E.V 0]), E.Tensor(1,[E.V 0])]
        }

        
    (*multiply scalar fields *)
    fun mulFieldss dim = E.EIN{
            params = [E.FLD(dim),E.FLD(dim)],
            index = [],
            body = E.Prod[E.Field(0, []),E.Field(1, [])]
        }
        
    fun mulFieldsf(dim,shape) =let
        val  expindex= specialize(shape,0)
        in E.EIN{
            params = [E.FLD(dim),E.FLD(dim)],
            index = shape,
            body = E.Prod[E.Field(0, []),E.Field(1, expindex)]
            }
        end
        
    fun divFieldss dim = E.EIN{
            params = [E.FLD(dim),E.FLD(dim)],
            index = [],
            body = E.Div(E.Field(0, []),E.Field(1, []))
        }
        
    fun divFieldfs(dim,shape) = let
        val  expindex= specialize(shape,0)
        in  E.EIN{
            params = [E.FLD(dim),E.FLD(dim)],
            index = shape,
            body = E.Prod[E.Field(0, expindex),E.Div(E.Const 1,E.Field(1, []))]
            }
        end 
        
    (*Assumes same dimension vector field *)
    fun outerField dim =
        E.EIN{
            params = [E.FLD(dim),E.FLD(dim)],
            index = [dim, dim],
            body = E.Prod[E.Field(0, [E.V 0]),E.Field(1, [E.V 1])]
        }
        
        fun fs x=Int.toString(x)
        fun f x=fs(length(x))
        
     (* generic inner product: <T_{\alpha i} * T_{i \beta}>_{\alpha \beta} *)
    fun innerProductField(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 s'=E.V(length(alpha)+ length(beta))
        val s''=[(s',0,i-1)]
        in E.EIN{
            params = [E.FLD(dim) ,E.FLD(dim)],              (* T and T' *)
            index = alpha@beta,(* \alpha \beta, i *)
            body = E.Sum(s'', E.Prod[
                E.Field(0, expindexA@[s']),   (* F_{\alpha i} *)
                E.Field(1, [s']@expindexB )  (* F'_{i \beta} *)
            ])}
        end
    | innerProductField _ = raise Fail "Wrong shape for innerProductField"
            
    (*Field Cross Product*)
    val crossProductField2d = E.EIN{
            params = [E.FLD(2), E.FLD(2)],
            index= [],
            body=E.Sum([(E. V 0,0,1),(E.V 1,0,1)],
            E.Prod[ E.Eps2(0, 1), E.Field(0, [E.V 0]),  E.Field(1, [E.V 1]) ])
         }
         
        (*Field Cross Product*)
    val crossProductField3d = 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.Prod[ E.Epsilon(0, 1, 2), E.Field(0, [E.V 1]),  E.Field(1, [E.V 2 ]) ])
         }
         
    (* Trace: <Sigma_i F_{\alpha i, i}>  This one Sx represents both i's*)
    fun traceField(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
    
    (*Transpose Field F_{ji}*)
    fun transposeField(dim,i,j) =E.EIN{
            params = [E.FLD(dim)], index= [i,j],
            body= E.Field(0, [E.V 1,E.V 0])
        }
        
    (*<F_{\alpha i j} * G{i j \beta }>_\alpha \beta*)
    fun doubleDotField(shape1,i::j::beta) = let
        val alpha= List.take(shape1,length(shape1)-2)
        val expindexA= specialize(alpha,0)
        val expindexB= specialize(beta,(length(alpha)))
        val sumi=length(alpha)+ length(beta)
        val s'=[E.V sumi,E.V(sumi+1)]
        val s''=[(E.V sumi,0,i-1),(E.V(sumi+1),0,j-1)]
        in
            E.EIN{
                params = [E.TEN(1,shape1),E.TEN(1,i::j::beta)],
                index = alpha@beta,
                body = E.Sum(s'',E.Prod[E.Field(0, expindexA@s'), E.Field(1,s'@expindexB)])
            }
        end
    | doubleDotField _ = raise Fail "Wrong shape for double dot "
        
        
    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 = [E.TEN(1,argTy)],
                index = rstTy,
                body = E.Tensor(0, ix)
            }
        end

        
        
    fun magnitudeTenVec(i::beta) = let
            val shape1=i::beta
            val n=length(shape1)
            val alpha= List.take(shape1,n-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 = [E.TEN(1,shape1),E.TEN(1,shape1)],
            index = alpha@beta,
            body = E.Sqrt(E.Sum(s'', E.Prod[
                E.Tensor(0, expindexA@[s']),
                E.Tensor(0, [s']@expindexB )
                ]))
            }

        end 
    fun magnitudeFldVec(dim,[]) =
        E.EIN{
            params = [E.FLD(dim)],             
            index = [],
            body =  E.Field(0, [])
            }        
        | magnitudeFldVec(dim,i::beta) = let
        val shape1=i::beta
        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 = [E.FLD(dim) ,E.FLD(dim)],              (* T and T' *)
            index = alpha@beta,(* \alpha \beta, i *)
            body = E.Sqrt(E.Sum(s'', E.Prod[
                E.Field(0, expindexA@[s']),   (* F_{\alpha i} *)
                E.Field(1, [s']@expindexB )  (* F'_{i \beta} *)
                ]))
            }
        end
        
        fun normalizeFldVec(dim,[]) =raise Fail"normalize of a scalar"
          | normalizeFldVec(dim,[i]) = let    
            val sx=[(E.V 0,0,i-1)]
            val f=E.Field(0, [E.V 0])
            in E.EIN{
                params = [E.FLD(dim) ,E.FLD(dim)],
                index = [dim],
                body = E.Prod[f,
                    E.Div(E.Const 1,
                        E.Sqrt(E.Sum(sx, E.Prod[f,f]))
                    )]
            }
            end
        
        
        fun sqrt dim=
            E.EIN{
                params = [E.FLD(dim)],
                index = [],
                body = E.Sqrt(E.Field(0, []))
            }

    

  end; (* local *)

    end (* local *)

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