5 files changed, 94 insertions, 4 deletions

diff --git a/lib/Data/Packed/Internal/Vector.hs b/lib/Data/Packed/Internal/Vector.hs
index ba68909..70bbae2 100644
--- a/lib/Data/Packed/Internal/Vector.hs
+++ b/lib/Data/Packed/Internal/Vector.hs
@@ -22,6 +22,7 @@ module Data.Packed.Internal.Vector (
     foldVector, foldVectorG, foldLoop, foldVectorWithIndex,
     createVector, vec,
     asComplex, asReal, float2DoubleV, double2FloatV,
+    stepF, stepD,
     fwriteVector, freadVector, fprintfVector, fscanfVector,
     cloneVector,
     unsafeToForeignPtr,
@@ -292,6 +293,23 @@ double2FloatV v = unsafePerformIO $ do
 foreign import ccall "float2double" c_float2double:: TFV
 foreign import ccall "double2float" c_double2float:: TVF
 
+---------------------------------------------------------------
+
+stepF :: Vector Float -> Vector Float
+stepF v = unsafePerformIO $ do
+    r <- createVector (dim v)
+    app2 c_stepF vec v vec r "stepF"
+    return r
+
+stepD :: Vector Double -> Vector Double
+stepD v = unsafePerformIO $ do
+    r <- createVector (dim v)
+    app2 c_stepD vec v vec r "stepD"
+    return r
+
+foreign import ccall "stepF" c_stepF :: TFF
+foreign import ccall "stepD" c_stepD :: TVV
+
 ----------------------------------------------------------------
 
 cloneVector :: Storable t => Vector t -> IO (Vector t)
diff --git a/lib/Numeric/ContainerBoot.hs b/lib/Numeric/ContainerBoot.hs
index 992a501..5a8e243 100644
--- a/lib/Numeric/ContainerBoot.hs
+++ b/lib/Numeric/ContainerBoot.hs
@@ -45,6 +45,7 @@ module Numeric.ContainerBoot (
 ) where
 
 import Data.Packed
+import Data.Packed.ST as ST
 import Numeric.Conversion
 import Data.Packed.Internal
 import Numeric.GSL.Vector
@@ -120,6 +121,12 @@ class (Complexable c, Fractional e, Element e) => Container c e where
     sumElements :: c e -> e
     -- | the product of elements (faster than using @fold@)
     prodElements :: c e -> e
+    -- | map (if x_i>0 then 1.0 else 0.0)
+    step :: RealFloat e => c e -> c e
+    -- | find index of elements which satisfy a predicate
+    find :: (e -> Bool) -> c e -> [IndexOf c]
+    -- | create a structure from an association list
+    assoc :: IndexOf c -> e -> [(IndexOf c, e)] -> c e
 
 --------------------------------------------------------------------------
 
@@ -145,6 +152,9 @@ instance Container Vector Float where
     maxElement  = toScalarF Max
     sumElements  = sumF
     prodElements = prodF
+    step = stepF
+    find = findV
+    assoc = assocV
 
 instance Container Vector Double where
     scale = vectorMapValR Scale
@@ -168,6 +178,9 @@ instance Container Vector Double where
     maxElement  = toScalarR Max
     sumElements  = sumR
     prodElements = prodR
+    step = stepD
+    find = findV
+    assoc = assocV
 
 instance Container Vector (Complex Double) where
     scale = vectorMapValC Scale
@@ -191,6 +204,9 @@ instance Container Vector (Complex Double) where
     maxElement  = ap (@>) maxIndex
     sumElements  = sumC
     prodElements = prodC
+    step = undefined -- cannot match
+    find = findV
+    assoc = assocV
 
 instance Container Vector (Complex Float) where
     scale = vectorMapValQ Scale
@@ -214,6 +230,9 @@ instance Container Vector (Complex Float) where
     maxElement  = ap (@>) maxIndex
     sumElements  = sumQ
     prodElements = prodQ
+    step = undefined -- cannot match
+    find = findV
+    assoc = assocV
 
 ---------------------------------------------------------------
 
@@ -243,6 +262,9 @@ instance (Container Vector a) => Container Matrix a where
     maxElement = ap (@@>) maxIndex
     sumElements = sumElements . flatten
     prodElements = prodElements . flatten
+    step = liftMatrix step
+    find = findM
+    assoc = assocM
 
 ----------------------------------------------------
 
@@ -580,3 +602,21 @@ diag v = diagRect 0 v n n where n = dim v
 -- | creates the identity matrix of given dimension
 ident :: (Num a, Element a) => Int -> Matrix a
 ident n = diag (constantD 1 n)
+
+--------------------------------------------------------
+
+findV p x = foldVectorWithIndex g [] x where
+    g k z l = if p z then k:l else l
+
+findM p x = map ((`divMod` cols x)) $ findV p (flatten x)
+
+assocV n z xs = ST.runSTVector $ do
+        v <- ST.newVector z n
+        mapM_ (\(k,x) -> ST.writeVector v k x) xs
+        return v
+
+assocM (r,c) z xs = ST.runSTMatrix $ do
+        m <- ST.newMatrix z r c
+        mapM_ (\((i,j),x) -> ST.writeMatrix m i j x) xs
+        return m
+
diff --git a/lib/Numeric/LinearAlgebra/LAPACK/lapack-aux.c b/lib/Numeric/LinearAlgebra/LAPACK/lapack-aux.c
index e8bbbdb..ae437d2 100644
--- a/lib/Numeric/LinearAlgebra/LAPACK/lapack-aux.c
+++ b/lib/Numeric/LinearAlgebra/LAPACK/lapack-aux.c
@@ -1247,3 +1247,23 @@ int conjugateC(KCVEC(x),CVEC(t)) {
     OK
 }
 
+//////////////////// step /////////////////////////
+
+int stepF(FVEC(x),FVEC(y)) {
+    DEBUGMSG("stepF")
+    int k;
+    for(k=0;k<xn;k++) {
+        yp[k]=xp[k]>0;
+    }
+    OK
+}
+
+int stepD(DVEC(x),DVEC(y)) {
+    DEBUGMSG("stepD")
+    int k;
+    for(k=0;k<xn;k++) {
+        yp[k]=xp[k]>0;
+    }
+    OK
+}
+
diff --git a/lib/Numeric/LinearAlgebra/LAPACK/lapack-aux.h b/lib/Numeric/LinearAlgebra/LAPACK/lapack-aux.h
index 0543f7a..6207a59 100644
--- a/lib/Numeric/LinearAlgebra/LAPACK/lapack-aux.h
+++ b/lib/Numeric/LinearAlgebra/LAPACK/lapack-aux.h
@@ -87,6 +87,9 @@ int double2float(DVEC(x),FVEC(y));
 int conjugateQ(KQVEC(x),QVEC(t));
 int conjugateC(KCVEC(x),CVEC(t));
 
+int stepF(FVEC(x),FVEC(y));
+int stepD(DVEC(x),DVEC(y));
+
 int svd_l_R(KDMAT(x),DMAT(u),DVEC(s),DMAT(v));
 int svd_l_Rdd(KDMAT(x),DMAT(u),DVEC(s),DMAT(v));
 int svd_l_C(KCMAT(a),CMAT(u),DVEC(s),CMAT(v));
diff --git a/lib/Numeric/LinearAlgebra/Tests.hs b/lib/Numeric/LinearAlgebra/Tests.hs
index 68e77f1..181cfbf 100644
--- a/lib/Numeric/LinearAlgebra/Tests.hs
+++ b/lib/Numeric/LinearAlgebra/Tests.hs
@@ -342,8 +342,8 @@ lift_maybe m = MaybeT $ do
 
 -- | apply a test to successive elements of a vector, evaluates to true iff test passes for all pairs
 --successive_ :: Storable a => (a -> a -> Bool) -> Vector a -> Bool
-successive_ t v = maybe False (\_ -> True) $ evalState (runMaybeT (mapVectorM_ step (subVector 1 (dim v - 1) v))) (v @> 0)
-   where step e = do
+successive_ t v = maybe False (\_ -> True) $ evalState (runMaybeT (mapVectorM_ stp (subVector 1 (dim v - 1) v))) (v @> 0)
+   where stp e  = do
                   ep <- lift_maybe $ state_get
                   if t e ep
                      then lift_maybe $ state_put e
@@ -351,8 +351,8 @@ successive_ t v = maybe False (\_ -> True) $ evalState (runMaybeT (mapVectorM_ s
 
 -- | operate on successive elements of a vector and return the resulting vector, whose length 1 less than that of the input
 --successive :: (Storable a, Storable b) => (a -> a -> b) -> Vector a -> Vector b
-successive f v = evalState (mapVectorM step (subVector 1 (dim v - 1) v)) (v @> 0)
-   where step e = do
+successive f v = evalState (mapVectorM stp (subVector 1 (dim v - 1) v)) (v @> 0)
+   where stp  e = do
                   ep <- state_get
                   state_put e
                   return $ f ep e
@@ -365,6 +365,14 @@ succTest = utest "successive" $
 
 ---------------------------------------------------------------------
 
+findAssocTest = utest "findAssoc" ok
+  where
+    ok = m1 == m2
+    m1 = assoc (6,6) 7 $ zip (find (>0) (ident 5 :: Matrix Float)) [10 ..] :: Matrix Double
+    m2 = diagRect 7 (fromList[10..14]) 6 6 :: Matrix Double
+
+---------------------------------------------------------------------
+
 
 -- | All tests must pass with a maximum dimension of about 20
 --  (some tests may fail with bigger sizes due to precision loss).
@@ -533,6 +541,7 @@ runTests n = do
         , sumprodTest
         , chainTest
         , succTest
+        , findAssocTest
         ]
     return ()