Changeset 626 for branches

Timestamp:

12/31/12 09:59:13 (3 years ago)

Author:

mmckerns

Message:

discrete decorator can be applied to given indicies

File:

• branches/decorate/wrapper.py (modified) (4 diffs)

branches/decorate/wrapper.py

@@ -41 +41 @@
 
 
-from numpy import asarray, any, inf, vectorize
+from numpy import asarray, any, inf, vectorize, choose
 def bounded(min=None, max=None):
     """impose an infinite barrier box constraint on a function
@@ -119 +119 @@
 
 #from random import sample, choice
-def discrete(samples):
+def discrete(samples, index=None):
     """impose a discrete set of input values for the selected function
 
@@ -126 +126 @@
 >>> @discrete([1.0, 2.0])
 ... def identity(x):
-...   return x
+...     return x
 
 >>> identity([0.123, 1.789, 4.000])
-[1.0, 2.0, 2.0]""" #FIXME: should return same type as input
+[1.0, 2.0, 2.0]
+
+>>> @discrete([1,3,5,7], index=(0,3)])
+... def squared(x):
+....    return [i**2 for i in x]
+
+>>> squared([0,2,4,6,8,10])
+[1, 4, 16, 25, 64, 100]"""
     samples = [asarray(samples)]
     samples[0].sort()
@@ -157 +164 @@
     def dec(f):
         def func(x, *args, **kwds):
+            xtype = type(x)
             arglo, arghi = argnear(x)
             xp = near(x, samples[0][arglo], samples[0][arghi])
+            # create a choice array from given indices
+            mask = xp*False #FIXME: better ways to do this
+            try: mask[sorted(index)] = True #FIXME: better ways to do this
+            except IndexError: pass
+            xp = xtype(choose(mask, (x,xp)))
             return f(xp, *args, **kwds)
         func.samples = _points