diff --git a/maxent.cabal b/maxent.cabal
--- a/maxent.cabal
+++ b/maxent.cabal
@@ -10,7 +10,7 @@
 -- PVP summary:      +-+------- breaking API changes
 --                   | | +----- non-breaking API additions
 --                   | | | +--- code changes with no API change
-version:             0.1.0.1
+version:             0.2.0.0
 
 -- A short (one-line) description of the package.
 synopsis:            Compute Maximum Entropy Distributions
diff --git a/src/MaxEnt.hs b/src/MaxEnt.hs
--- a/src/MaxEnt.hs
+++ b/src/MaxEnt.hs
@@ -31,12 +31,14 @@
 -- 
 module MaxEnt (
     Constraint,
+    ExpectationFunction,
     constraint,
     average,
     variance,
     maxent
 ) where
 import MaxEnt.Internal (Constraint,
+                        ExpectationFunction,
                         constraint,
                         average,
                         variance,
diff --git a/src/MaxEnt/Internal.hs b/src/MaxEnt/Internal.hs
--- a/src/MaxEnt/Internal.hs
+++ b/src/MaxEnt/Internal.hs
@@ -9,18 +9,34 @@
 sumWith :: Num c => (a -> b -> c) -> [a] -> [b] -> c 
 sumWith f xs = sum . zipWith f xs
 
-pOfK :: Floating a => [a] -> [a -> a] -> [a] -> Int -> a
-pOfK values fs ls k = exp (negate . sumWith (\l f -> l * f (values !! k)) ls $ fs) / 
+pOfK :: Floating a => [a] -> [ExpectationFunction a] -> [a] -> Int -> a
+pOfK values fs ls k = exp (negate . sumWith (\l f -> l * f k (values !! k)) ls $ fs) / 
     partitionFunc values fs ls 
 
-probs :: Floating b => [b] -> [b -> b] -> [b] -> [b]    
+probs :: Floating b 
+      => [b] 
+      -> [ExpectationFunction b] 
+      -> [b] 
+      -> [b]    
 probs values fs ls = map (pOfK values fs ls) [0..length values - 1] 
 
-partitionFunc :: Floating a => [a] -> [a -> a] -> [a] -> a
-partitionFunc values fs ls = sum $ [ exp ((-l) * f x) | x <- values, (f, l) <- zip fs ls]
+partitionFunc :: Floating a 
+              => [a] 
+              -> [ExpectationFunction a]
+              -> [a] 
+              -> a
+partitionFunc values fs ls = sum $ [ exp ((-l) * f i x) | 
+                                (i, x) <- zip [0..] values, 
+                                (f, l) <- zip fs ls]
 
-objectiveFunc :: Floating a => [a] -> [a -> a] -> [a] -> [a] -> a
-objectiveFunc values fs moments ls = log (partitionFunc values fs ls) + sumWith (*) ls moments
+objectiveFunc :: Floating a 
+              => [a] 
+              -> [ExpectationFunction a] 
+              -> [a] 
+              -> [a] 
+              -> a
+objectiveFunc values fs moments ls = log (partitionFunc values fs ls) 
+                                   + sumWith (*) ls moments
 
 toFunction :: (forall a. Floating a => [a] -> a) -> Function Simple
 toFunction f = VFunction (f . U.toList)
@@ -31,20 +47,34 @@
 toDoubleF :: (forall a. Floating a => [a] -> a) -> [Double] -> Double
 toDoubleF f x = f x 
 
--- | Constraint type. Think of this as f and c in sum pi (f x) = c
-type Constraint a = (a -> a, a)
+-- | Constraint type. A function and the constant it equals.
+-- 
+--   Think of it as the pair @(f, c)@ in the constraint 
+--
+-- @
+--     &#931; p&#8336; f(a, x&#8336;) = c
+-- @
+--
+--  such that we are summing over all values and @a@ is the index.
+--
+--  For example, for a variance constraint the @f@ would be @(\\_ x -> x*x)@ and @c@ would be the variance.
+type Constraint a = (ExpectationFunction a, a)
 
+-- | A function that takes an index and value and returns a value.
+--   See 'average' and 'variance' for examples.
+type ExpectationFunction a = (Int -> a -> a)
+
 -- make a constraint from function and constant
-constraint :: Floating a => (a -> a) -> a -> Constraint a
+constraint :: Floating a => ExpectationFunction a -> a -> Constraint a
 constraint = (,)
 
 -- The average constraint
 average :: Floating a => a -> Constraint a
-average m = constraint id m
+average m = constraint (const id) m
 
 -- The variance constraint
 variance :: Floating a => a -> Constraint a
-variance sigma = constraint (^(2 :: Int)) sigma
+variance sigma = constraint (const (^(2 :: Int))) sigma
 
 -- | The main entry point for computing discrete maximum entropy distributions.
 --   
@@ -57,18 +87,18 @@
     values :: Floating a => [a]
     values = fst params
     
-    constraints :: Floating a => [(a -> a, a)]
+    constraints :: Floating a => [(ExpectationFunction a, a)]
     constraints = snd params
     
-    fsmoments :: Floating a => ([a -> a], [a])
+    fsmoments :: Floating a => ([ExpectationFunction a], [a])
     fsmoments = unzip constraints 
     
-    fs :: [Double -> Double]
+    fs :: [Int -> Double -> Double]
     fs = fst fsmoments
     
     -- hmm maybe there is a better way to get rid of the defaulting
     guess = U.fromList $ replicate 
-        (length (constraints :: [(Double -> Double, Double)])) (1.0 :: Double) 
+        (length fs) (1.0 :: Double) 
     
     result = case unsafePerformIO (optimize defaultParameters 0.00001 guess 
                         (toFunction obj)