A recursive tensor data type

Tensor packages in Haskell

There are several tensor packages for Haskell. The Data.Tensor module is intended for projective geometry for 3D modelling using OpenGL, and is restricted to four dimensions. By contrast, the hTensor package is intended for tensor calculations and is quite general. If you want to do tensor computations, you should probably use it. It supports named indices, co- and contravariant indices and the Einstein summation convention. It came out just after I started writing my own module, and I might not have written it if I had known.

If, on the other hand, you want to know how to use recursive Haskell data types for something other than parsing, read on. They are a perfect match for describing tensors, and many functions can be expressed very elegantly using folds over the data type.

The data type

After reading the headline, you will probably already guess the data type I am about to declare:

data Tensor n = Scalar n | Vector [Tensor n]

By that definition, a tensor is either a scalar, with constructor Scalar, or a list of sub-tensors, via the constructor Vector. This nomenclature is mathematically dubious, as even order-one sub-tensors are (AFAIK) never called vectors, to say nothing about higher-order ones. But as we have to pick a name for the second constructor, Vector is probably the least misleading choice.

A more practical objection to the data type is that the tree structure it defines covers more than just tensors. All subtensors of a tensor at the same level of recursion have the same size (dimension), while the Tensor type allows more general tree structures. This is solved — practically if not formally — by providing functions for the creation of tensors, and a function to verify the tensor property for users who want to use the constructors.

Having defined a recursive data type, it remains to define functions that operate on it. Many of them can be defined by iteration over the data type, which in functional programming is called a fold. Taking a leaf out of the book of formal languages (such as this one, chapters 6 and 7), we define the data type of a "semantic function" over the Tensor data type:

type TensorAlgebra n r = (n -> r, [r] -> r)

This is a tuple of two functions, the first of which processes a Scalar, and the second of which combines the results of processing sub-tensors into the result of a higher level. The operation is applied to a Tensor by the following function:

tenFold :: TensorAlgebra n r -> Tensor n -> r
tenFold (fs, _) (Scalar s) = fs s
tenFold f@(_, fv) (Vector v) = fv $ map (tenFold f) v

The many uses of folds

The tenFold function defined in the previous section has a multitude of applications. Perhaps the simplest is scalar multiplication:

scalarmul :: Num n => n -> Tensor n -> Tensor n
scalarmul n = tenFold (Scalar . (n*), Vector)

A further fold gets us the product of two tensors without contracting:

externalprod :: Num n => Tensor n -> Tensor n -> Tensor n
externalprod t1 t2 = tenFold (flip scalarmul t2, Vector) t1

This function can for instance be used in linear algebra to generate a projection matrix by multiplying a normalised vector by itself. It has nothing to do with the antisymmetric exterior product.

More generally, one can define a map function that makes Tensor an instance of the Functor class:

tenMap :: (a -> b) -> Tensor a -> Tensor b
tenMap f = tenFold (Scalar . f, Vector)

Operations that take a Tensor to a different data type have a combining function (second part of the TensorAlgebra) that is not Vector. For instance, one could add up all the elements of a tensor using the following function:

crosssum :: Num n => Tensor n -> n
crosssum = tenFold (id, sum)

A more complicated example is turning a tensor into a list of its dimensions and a flat list of its elements. For the purpose of the latter, the top level of the tree corresponds to the most slowly changing index. Stripped of the error handling, the function is:

tensor2list :: Tensor n -> ([Int], [n])
tensor2list = tenFold ((\x -> ([], [x])), (\l -> (length l : fst (head l), concat $ map snd l)))

A Scalar has no dimension, so the first function only puts the tensor element into a list in the second half of its result. The combining function is faced with a list of sub-results that (for a correct tensor) all have the same dimensions. So it picks the first and prepends it by the dimension at its level, the length of the sub-tensor list. The lists of elements of the sub-tensors are simply concatenated.

Almost as many uses for zips

Just as tenFold applies an operation to one tensor, one can define a function that applies an operation to two equally-sized tensors. In analogy to the corresponding function for lists, I called it vecZipWith. Without the error handling, its definition is the following:

tenZipWith :: TensorAlgebra (a, b) c -> Tensor a -> Tensor b -> c
tenZipWith (fs, _) (Scalar a) (Scalar b) = fs (a, b)
tenZipWith f@(_, fv) (Vector a) (Vector b) = fv $ zipWith (tenZipWith f) a b

This makes it easy to define the element-wise sum of two tensors, or a general element-wise binary operation:

tenAdd :: Num n => Tensor n -> Tensor n -> Tensor n
tenAdd = tenZipWith (Scalar . uncurry (+), Vector)
tenBinOp :: (a -> b -> c) -> Tensor a -> Tensor b -> Tensor c
tenBinOp f = tenZipWith (Scalar . uncurry f, Vector)

As the simpler examples for the fold, these functions keep the tensor as it is. Since the second function is the Vector constructor, the result has the same structure (meaning the same dimensions) as the two arguments. As for folds, the usefulness of the zip does not end there. Consider for example the full contraction of two equally-sized tensors by corresponding indices:

fullcontract :: Num n => Tensor n -> Tensor n -> n
fullcontract = tenZipWith (uncurry (*), sum)

The zip functional can also be used to recursively compare two equally-sizes tensors, in order to make Tensor an instance of the Eq class:

instance (Eq n) => Eq (Tensor n)
  where (==) a b = tenDims a == tenDims b && tenZipWith (uncurry (==), and) a b
        (/=) a b = tenDims a /= tenDims b || tenZipWith (uncurry (/=), or) a b

tenDims is a small function which returns a list of the tensor's dimensions. The equality of the dimension has to be checked first because that is one of the assumptions necessary for the use of tenZipWith.

Now it gets interesting: index rearrangements and general contractions

Operations on Tensors that can be use a fold or a zip can be expressed briefly and elegantly, as you have seen above. Functions that do something to specific indices but not others are more complex. (This may be the reason why the hTensor package uses a flat list as the underlying data structure — but where is the fun in that?)

Acting on a certain index corresponds to applying a function to all sub-tensors at a specific level in the tree structure. A fold is not suitable for that, as it always recurses over the whole tree. One needs a functional that descends only one level, or a given number of levels.

vecMap :: (Tensor a -> Tensor b) -> Tensor a -> Tensor b
vecMap f (Vector l) = Vector $ map f l
nVecMap :: Int -> (Tensor a -> Tensor b) -> Tensor a -> Tensor b
nVecMap n = (!! n) . iterate vecMap

Note that the type signature of the single-level map and a map for a given number of levels (nVecMap n for fixed n) have the same type signature. The recursive definition of Tensor is what causes this, and what allows declaring nVecMap as an iterate in the first place. It also gives rise to a certain conceptual difficulty, namely that one cannot read off from the types whether one is descending into the tree. The expressions t and Vector [t] have the same type, so one has to keep track of the depth mentally without help from the formalism.

Using vecMap, one can compute the trace of a matrix:

trace :: Num n => Tensor n -> n
trace m
  | vnull m = 0
  | otherwise = (vhead $ vhead m) + (trace $ vecMap vtail $ vtail m)

vnull, vhead and vtail are the same as the corresponding list functions, just for Vectors. Any error checking, such as making sure the tensor is indeed a square matrix, has again been left out.

Rearranging indices requires a second element besides vecMap. One has to turn the tree structure "inside out" in order to exchange indices. The simplest example is transposing a matrix:

transpose :: Tensor n -> Tensor n
transpose v | vnull $ vhead v = Vector []
            | otherwise = vcons (fst headtail) (transpose $ snd headtail)
            where headtail = vecUnzipWith vheadtail v
vheadtail :: Tensor n -> (Tensor n, Tensor n)
vheadtail (Vector (vh:vt)) = (vh, Vector vt)
vecUnzipWith :: (Tensor a -> (Tensor b, Tensor c)) -> Tensor a -> (Tensor b, Tensor c)
vecUnzipWith f (Vector l) = (Vector $ fst fl, Vector $ snd fl)
        where fl = unzip $ map f l

The transpose function recursively pulls out the first elements of all sub-tensors as a Vector and prepends them to the result tensor. (The vcons function is the Vector equivalent of the (:) operator.) vheadtail separates the head and tail of sub-tensors. The actual "pulling out" of inner indices is performed by vecUnzipWith, which turns the resulting Vector of pairs into a pair of Vectors. The first half of the pair is the first column of the matrix if it was previously in row-major format.

As you can see from the type signature of vecUnzipWith, it can be iterated. This allows pulling out indices by more than one level. Conversely, indices can be pushed to the back by using the self-explanatory functional vecZipWith. vecZipWith can also be iterated in order to push indices back by more than one level:

pushindex :: Int -> Tensor n -> Tensor n
pushindex n v@(Vector (vh:vt))
        | null vt   = nVecMap n (\x -> Vector [x]) vh
        | otherwise = nVecZipWith n vcons vh (pushindex n $ Vector vt)

Now that we can manipulate indices, we have all the prerequisites for defining a general contraction of two tensors. The idea is to push the indices to be contracted to the back in each of the tensors, then doing a full contraction of the sub-tensors.

contract :: Num n => [(Int,Int)] -> Tensor n -> Tensor n -> Tensor n
contract inds x y
  = nVecMap (ox-li) (\xs -> nVecMap (oy-li) (Scalar . fullcontract xs) y') x'
  where
    inds' = unzip $ sortind $ checkind inds
    x' = pushall (fst inds') x
    y' = pushall (snd inds') y
    ox = order x
    oy = order y
    li = length inds

x' and y' are the argument tensors that have their affected indices shifted to the end. The first nVecMap function simply bypasses the unaffected indices of x' and makes the outer structure of the result tensor that of the remaining part of x (that is, the first indices of the result are the remaining indices of x). For each sub-tensor of x' at that level, the unaffected indices of y' are also skipped, and the sub-tensor xs is contracted with each sub-tensor of y'.

The details are rather more tricky. When the affected indices of x and y are shifted to the end, they have to end up in corresponding order. That order is computed by the sortind function, which returns a pair of index lists (for x and for y). The pushall function shifts indices to the back in the given order, taking into account that the number of levels they have to be shifted depends on how other indices have been shifted previously. If you are interested in these details, however, here is the full source code for you:

Download

If I have made you curious and you want to have a closer look at my recursive tensor implementation, or if you — heaven forbid — actually want to use it, you can get it here:

Download the Tensor module

It is documented for Haddock, so you can generate HTML documentation (see also below) with:

haddock -h tensor.hs

or, if you want the documentation of internal functions as well:

haddock -h --ignore-all-exports tensor.hs

You probably best run these commands in a newly created directory (or use the -o option, as they create between five and ten files.

Haddock documentation of Tensor

TOS / Impressum