def test_compare(solvername, x0, **kwds):
  exec "my = solvers.%s" % solvername
  exec "sp = %s" % solvername
  maxiter = kwds.get('maxiter', None)
  maxfun = kwds.get('maxfun', None)
  my_x = my(rosen, x0, disp=0, full_output=True, **kwds)
# itermon = kwds.pop('itermon',None)
  sp_x = sp(rosen, x0, disp=0, full_output=True, **kwds)
  # similar bestSolution and bestEnergy
# print 'my:', my_x[0:2]
# print 'sp:', sp_x[0:2]
  if my_x[3] == sp_x[-2]: # mystic can stop at iter=0, scipy can't
    assert almostEqual(my_x[0], sp_x[0])
    assert almostEqual(my_x[1], sp_x[1])
  # print (iters, fcalls) and [maxiter, maxfun]
# print my_x[2:4], (sp_x[-3],sp_x[-2]), [maxiter, maxfun]
  # test same number of iters and fcalls
  if maxiter and maxfun is not None:
    assert my_x[2] == sp_x[-3]
    assert my_x[3] == sp_x[-2]
#   # test fcalls <= maxfun
#   assert my_x[3] <= maxfun
  if maxiter is not None:
    # test iters <= maxiter
    assert my_x[2] <= maxiter
  return 

def test_impose_reweighted_variance():

  x0 = [1,2,3,4,5]
  w0 = [3,1,1,1,1]
  v = 1.0

  w = impose_reweighted_variance(v, x0, w0)
  assert almostEqual(variance(x0,w), v)
  assert almostEqual(mean(x0,w0), mean(x0,w))

def test_generate_constraint():

  constraints = """
  spread([x0, x1, x2]) = 10.0
  mean([x0, x1, x2]) = 5.0"""

  from mystic.math.measures import mean, spread
  solv = generate_solvers(constraints)
  assert almostEqual(mean(solv[0]([1,2,3])), 5.0)
  assert almostEqual(spread(solv[1]([1,2,3])), 10.0)

  constraint = generate_constraint(solv)
  assert almostEqual(constraint([1,2,3]), [0.0,5.0,10.0], 1e-10)

def test_with_mean_spread():

  from mystic.math.measures import mean, spread, impose_mean, impose_spread

  @with_spread(50.0)
  @with_mean(5.0)
  def constrained_squared(x):
    return [i**2 for i in x]

  from numpy import array
  x = array([1,2,3,4,5])
  y = impose_spread(50.0, impose_mean(5.0,[i**2 for i in x]))
  assert almostEqual(mean(y), 5.0, tol=1e-15)
  assert almostEqual(spread(y), 50.0, tol=1e-15)
  assert constrained_squared(x) == y

 def factory(x, *args, **kwds):
     # apply decorated constraints function
     x = constraints(x, *args, **kwds)
     # constrain x such that mean(x) == target
     if not almostEqual(mean(x), target):
         x = impose_mean(target, x)#, weights=weights)
     return x

def test_constrain():

  from mystic.math.measures import mean, spread
  from mystic.math.measures import impose_mean, impose_spread
  def mean_constraint(x, mean=0.0):
    return impose_mean(mean, x)

  def range_constraint(x, spread=1.0):
    return impose_spread(spread, x)

  @inner(inner=range_constraint, kwds={'spread':5.0})
  @inner(inner=mean_constraint, kwds={'mean':5.0})
  def constraints(x):
    return x

  def cost(x):
    return abs(sum(x) - 5.0)

  from mystic.solvers import fmin_powell
  from numpy import array
  x = array([1,2,3,4,5])
  y = fmin_powell(cost, x, constraints=constraints, disp=False)

  assert mean(y) == 5.0
  assert spread(y) == 5.0
  assert almostEqual(cost(y), 4*(5.0))

 def constraints(x):
     # constrain the last x_i to be the same value as the first x_i
     x[-1] = x[0]
     # constrain x such that mean(x) == target
     if not almostEqual(mean(x), target):
         x = impose_mean(target, x)
     return x

def test_impose_reweighted_mean():

  x0 = [1,2,3,4,5]
  w0 = [3,1,1,1,1]
  m = 3.5

  w = impose_reweighted_mean(m, x0, w0)
  assert almostEqual(mean(x0,w), m)

def test_solve_constraint():

  from mystic.math.measures import mean
  @with_mean(1.0)
  def constraint(x):
    x[-1] = x[0]
    return x

  x = solve(constraint, guess=[2,3,1])

  assert almostEqual(mean(x), 1.0, tol=1e-15)
  assert x[-1] == x[0]
  assert issolution(constraint, x)

 def constraints(rv):
   c = product_measure().load(rv, npts)
   # NOTE: bounds wi in [0,1] enforced by filtering
   # impose norm on each discrete measure
   for measure in c:
     if not almostEqual(float(measure.mass), 1.0, tol=atol, rel=rtol):
       measure.normalize()
   # impose expectation on product measure
   ##################### begin function-specific #####################
   E = float(c.expect(model))
   if not (E <= float(target[0] + error[0])) \
   or not (float(target[0] - error[0]) <= E):
     c.set_expect(target[0], model, (x_lb,x_ub), tol=error[0])
   ###################### end function-specific ######################
   # extract weights and positions
   return c.flatten()

def test_numpy_penalty():

  constraints = """
  mean([x0, x1, x2]) = 5.0
  x0 = x1 + x2"""

  ineq,eq = generate_conditions(constraints)
  assert eq[0]([7,5,3]) == 0.0
  assert eq[1]([7,4,3]) == 0.0

  penalty = generate_penalty((ineq,eq))
  assert penalty([9.0,5,4.0]) == 100.0
  assert penalty([7.5,4,3.5]) == 0.0

  constraint = as_constraint(penalty, solver='fmin')
  assert almostEqual(penalty(constraint([3,4,5])), 0.0, 1e-10)

def test_generate_penalty():

  constraints = """
  x0**2 = 2.5*x3 - a
  exp(x2/x0) >= b"""

  ineq,eq = generate_conditions(constraints, nvars=4, locals={'a':5.0, 'b':7.0})
  assert ineq[0]([4,0,0,1,0]) == 6.0
  assert eq[0]([4,0,0,1,0]) == 18.5

  penalty = generate_penalty((ineq,eq))
  assert penalty([1,0,2,2.4]) == 0.0
  assert penalty([1,0,0,2.4]) == 7200.0
  assert penalty([1,0,2,2.8]) == 100.0

  constraint = as_constraint(penalty, nvars=4, solver='fmin')
  assert almostEqual(penalty(constraint([1,0,0,2.4])), 0.0, 1e-10)

def impose_reweighted_variance(v, samples, weights=None, solver=None):
    """impose a variance on a list of points by reweighting weights"""
    ndim = len(samples)
    if weights is None:
        weights = [1.0/ndim] * ndim
    if solver is None or solver == 'fmin':
        from mystic.solvers import fmin as solver
    elif solver == 'fmin_powell':
        from mystic.solvers import fmin_powell as solver
    elif solver == 'diffev':
        from mystic.solvers import diffev as solver
    elif solver == 'diffev2':
        from mystic.solvers import diffev2 as solver
    norm = sum(weights)
    m = mean(samples, weights)

    inequality = ""
    equality = ""; equality2 = ""; equality3 = ""
    for i in range(ndim):
        inequality += "x%s >= 0.0\n" % (i) # positive
        equality += "x%s + " % (i)         # normalized
        equality2 += "%s * x%s + " % (float(samples[i]),(i)) # mean
        equality3 += "x%s*(%s-%s)**2 + " % ((i),float(samples[i]),m) # var

    equality += "0.0 = %s\n" % float(norm)
    equality += equality2 + "0.0 = %s*%s\n" % (float(norm),m)
    equality += equality3 + "0.0 = %s*%s\n" % (float(norm),v)

    penalties = generate_penalty(generate_conditions(inequality))
    constrain = generate_constraint(generate_solvers(solve(equality)))

    def cost(x): return sum(x)

    results = solver(cost, weights, constraints=constrain, \
                     penalty=penalties, disp=False, full_output=True)
    wts = list(results[0])
    _norm = results[1] # should have _norm == norm
    warn = results[4]  # nonzero if didn't converge

    #XXX: better to fail immediately if xlo < m < xhi... or the below?
    if warn or not almostEqual(_norm, norm):
        print "Warning: could not impose mean through reweighting"
        return None #impose_variance(v, samples, weights), weights

    return wts #samples, wts  # "mean-preserving"

  def cone_mesh(length):
    """ construct a conical mesh for a given length of cone """
    L1,L2,L3 = slope
    radius = length / L3 #XXX: * 0.5
    r0 = ZERO

    if almostEqual(radius, r0, tol=r0): radius = r0
    r = np.linspace(radius,radius,6) 
    r[0]= np.zeros(r[0].shape) 
    r[1] *= r0/radius
    r[5] *= r0/radius
    r[3]= np.zeros(r[3].shape) 

    p = np.linspace(0,2*np.pi,50) 
    R,P = np.meshgrid(r,p) 
    X,Y = L1 * R*np.cos(P), L2 * R*np.sin(P) 

    tmp=list() 
    for i in range(np.size(p)): 
      tmp.append([0,0,length,length,length,0]) # len = size(r)
    Z = np.array(tmp) 
    return X,Z,Y

from mystic.symbolic import generate_constraint, generate_solvers, solve
from mystic.symbolic import generate_penalty, generate_conditions

equations = """
x0**2 - x1 + 1.0 <= 0.0
1.0 - x0 + (x1 - 4)**2 <= 0.0
"""
#cf = generate_constraint(generate_solvers(solve(equations))) #XXX: inequalities
pf = generate_penalty(generate_conditions(equations), k=1e12)

from mystic.constraints import as_constraint

cf = as_constraint(pf)



if __name__ == '__main__':

    from mystic.solvers import buckshot
    from mystic.math import almostEqual

    result = buckshot(objective, 2, 40, bounds=bounds, penalty=pf, disp=False, full_output=True)

    assert almostEqual(result[0], xs, tol=1e-2)
    assert almostEqual(result[1], ys, rel=1e-2)



# EOF

from mystic.symbolic import generate_conditions, generate_penalty
pf = generate_penalty(generate_conditions(equations))
from mystic.symbolic import generate_constraint, generate_solvers, simplify
cf = generate_constraint(generate_solvers(simplify(equations)))

# inverted objective, used in solving for the maximum
_objective = lambda x: -objective(x)


if __name__ == '__main__':

  from mystic.solvers import diffev2, fmin_powell
  from mystic.math import almostEqual

  result = diffev2(objective, x0=bounds, bounds=bounds, constraint=cf, penalty=pf, npop=40, disp=False, full_output=True)
  assert almostEqual(result[0], xs, rel=1e-2)
  assert almostEqual(result[1], ys, rel=1e-2)

  result = fmin_powell(objective, x0=[0.0,0.0], bounds=bounds, constraint=cf, penalty=pf, disp=False, full_output=True)
  assert almostEqual(result[0], xs, rel=1e-2)
  assert almostEqual(result[1], ys, rel=1e-2)

  # alternately, solving for the maximum
  result = diffev2(_objective, x0=bounds, bounds=bounds, constraint=cf, penalty=pf, npop=40, disp=False, full_output=True)
  assert almostEqual( result[0], _xs, rel=1e-2)
  assert almostEqual(-result[1], _ys, rel=1e-2)

  result = fmin_powell(_objective, x0=[0,0], bounds=bounds, constraint=cf, penalty=pf, npop=40, disp=False, full_output=True)
  assert almostEqual( result[0], _xs, rel=1e-2)
  assert almostEqual(-result[1], _ys, rel=1e-2)

solver = as_constraint(penalty)
#solver = discrete(range(11))(solver)  #XXX: MOD = range(11) instead of LARGE
#FIXME: constrain to 'int' with discrete is very fragile!  required #MODs

def constraint(x):
    from numpy import round
    return round(solver(x))

# better is to constrain to integers, penalize otherwise
from mystic.constraints import integers

@integers()
def round(x):
  return x


if __name__ == '__main__':

    from mystic.solvers import diffev2
    from mystic.math import almostEqual

    result = diffev2(objective, x0=bounds, bounds=bounds, penalty=penalty, constraints=round, npop=30, gtol=50, disp=True, full_output=True)

    print(result[0])
    assert almostEqual(result[0], xs, tol=1e-8) #XXX: fails b/c rel & zero?
    assert almostEqual(result[1], ys, tol=1e-4)


# EOF

def test_griewangk():
    """Test Griewangk's function, which has many local minima.

Testing Griewangk:
Expected: x=[0.]*10 and f=0

Using DifferentialEvolutionSolver:
Solution:  [  8.87516194e-09   7.26058147e-09   1.02076001e-08   1.54219038e-08
  -1.54328461e-08   2.34589663e-08   2.02809360e-08  -1.36385836e-08
   1.38670373e-08   1.59668900e-08]
f value:  0.0
Iterations:  4120
Function evaluations:  205669
Time elapsed:  34.4936850071  seconds

Using DifferentialEvolutionSolver2:
Solution:  [ -2.02709316e-09   3.22017968e-09   1.55275472e-08   5.26739541e-09
  -2.18490470e-08   3.73725584e-09  -1.02315312e-09   1.24680355e-08
  -9.47898116e-09   2.22243557e-08]
f value:  0.0
Iterations:  4011
Function evaluations:  200215
Time elapsed:  32.8412370682  seconds
"""

    print "Testing Griewangk:"
    print "Expected: x=[0.]*10 and f=0"
    from mystic.models import griewangk as costfunc
    ndim = 10
    lb = [-400.]*ndim
    ub = [400.]*ndim
    maxiter = 10000
    seed = 123 # Re-seed for each solver to have them all start at same x0
    
    # DifferentialEvolutionSolver
    print "\nUsing DifferentialEvolutionSolver:"
    npop = 50
    random_seed(seed)
    from mystic.solvers import DifferentialEvolutionSolver
    from mystic.termination import ChangeOverGeneration as COG
    from mystic.termination import CandidateRelativeTolerance as CRT
    from mystic.termination import VTR
    from mystic.strategy import Rand1Bin, Best1Bin, Rand1Exp
    esow = Monitor()
    ssow = Monitor() 
    solver = DifferentialEvolutionSolver(ndim, npop)
    solver.SetRandomInitialPoints(lb, ub)
    solver.SetStrictRanges(lb, ub)
    solver.SetEvaluationLimits(generations=maxiter)
    solver.SetEvaluationMonitor(esow)
    solver.SetGenerationMonitor(ssow)
    solver.enable_signal_handler()
    #term = COG(1e-10)
    #term = CRT()
    term = VTR(0.)
    time1 = time.time() # Is this an ok way of timing?
    solver.Solve(costfunc, term, strategy=Rand1Exp, \
                 CrossProbability=0.3, ScalingFactor=1.0)
    sol = solver.Solution()
    time_elapsed = time.time() - time1
    fx = solver.bestEnergy
    print "Solution: ", sol
    print "f value: ", fx
    print "Iterations: ", solver.generations
    print "Function evaluations: ", len(esow.x)
    print "Time elapsed: ", time_elapsed, " seconds"
    assert almostEqual(fx, 0.0, tol=3e-3)

    # DifferentialEvolutionSolver2
    print "\nUsing DifferentialEvolutionSolver2:"
    npop = 50
    random_seed(seed)
    from mystic.solvers import DifferentialEvolutionSolver2
    from mystic.termination import ChangeOverGeneration as COG
    from mystic.termination import CandidateRelativeTolerance as CRT
    from mystic.termination import VTR
    from mystic.strategy import Rand1Bin, Best1Bin, Rand1Exp
    esow = Monitor()
    ssow = Monitor() 
    solver = DifferentialEvolutionSolver2(ndim, npop)
    solver.SetRandomInitialPoints(lb, ub)
    solver.SetStrictRanges(lb, ub)
    solver.SetEvaluationLimits(generations=maxiter)
    solver.SetEvaluationMonitor(esow)
    solver.SetGenerationMonitor(ssow)
    #term = COG(1e-10)
    #term = CRT()
    term = VTR(0.)
    time1 = time.time() # Is this an ok way of timing?
    solver.Solve(costfunc, term, strategy=Rand1Exp, \
                 CrossProbability=0.3, ScalingFactor=1.0)
    sol = solver.Solution()
    time_elapsed = time.time() - time1
    fx = solver.bestEnergy
    print "Solution: ", sol
    print "f value: ", fx
    print "Iterations: ", solver.generations
    print "Function evaluations: ", len(esow.x)
    print "Time elapsed: ", time_elapsed, " seconds"
    assert almostEqual(fx, 0.0, tol=3e-3)

def test_rosenbrock():
    """Test the 2-dimensional Rosenbrock function.

Testing 2-D Rosenbrock:
Expected: x=[1., 1.] and f=0

Using DifferentialEvolutionSolver:
Solution:  [ 1.00000037  1.0000007 ]
f value:  2.29478683682e-13
Iterations:  99
Function evaluations:  3996
Time elapsed:  0.582273006439  seconds

Using DifferentialEvolutionSolver2:
Solution:  [ 0.99999999  0.99999999]
f value:  3.84824937598e-15
Iterations:  100
Function evaluations:  4040
Time elapsed:  0.577210903168  seconds

Using NelderMeadSimplexSolver:
Solution:  [ 0.99999921  1.00000171]
f value:  1.08732211477e-09
Iterations:  70
Function evaluations:  130
Time elapsed:  0.0190329551697  seconds

Using PowellDirectionalSolver:
Solution:  [ 1.  1.]
f value:  0.0
Iterations:  28
Function evaluations:  859
Time elapsed:  0.113857030869  seconds
"""

    print "Testing 2-D Rosenbrock:"
    print "Expected: x=[1., 1.] and f=0"
    from mystic.models import rosen as costfunc
    ndim = 2
    lb = [-5.]*ndim
    ub = [5.]*ndim
    x0 = [2., 3.]
    maxiter = 10000
    
    # DifferentialEvolutionSolver
    print "\nUsing DifferentialEvolutionSolver:"
    npop = 40
    from mystic.solvers import DifferentialEvolutionSolver
    from mystic.termination import ChangeOverGeneration as COG
    from mystic.strategy import Rand1Bin
    esow = Monitor()
    ssow = Monitor() 
    solver = DifferentialEvolutionSolver(ndim, npop)
    solver.SetInitialPoints(x0)
    solver.SetStrictRanges(lb, ub)
    solver.SetEvaluationLimits(generations=maxiter)
    solver.SetEvaluationMonitor(esow)
    solver.SetGenerationMonitor(ssow)
    term = COG(1e-10)
    time1 = time.time() # Is this an ok way of timing?
    solver.Solve(costfunc, term, strategy=Rand1Bin)
    sol = solver.Solution()
    time_elapsed = time.time() - time1
    fx = solver.bestEnergy
    print "Solution: ", sol
    print "f value: ", fx
    print "Iterations: ", solver.generations
    print "Function evaluations: ", len(esow.x)
    print "Time elapsed: ", time_elapsed, " seconds"
    assert almostEqual(fx, 2.29478683682e-13, tol=3e-3)

    # DifferentialEvolutionSolver2
    print "\nUsing DifferentialEvolutionSolver2:"
    npop = 40
    from mystic.solvers import DifferentialEvolutionSolver2
    from mystic.termination import ChangeOverGeneration as COG
    from mystic.strategy import Rand1Bin
    esow = Monitor()
    ssow = Monitor() 
    solver = DifferentialEvolutionSolver2(ndim, npop)
    solver.SetInitialPoints(x0)
    solver.SetStrictRanges(lb, ub)
    solver.SetEvaluationLimits(generations=maxiter)
    solver.SetEvaluationMonitor(esow)
    solver.SetGenerationMonitor(ssow)
    term = COG(1e-10)
    time1 = time.time() # Is this an ok way of timing?
    solver.Solve(costfunc, term, strategy=Rand1Bin)
    sol = solver.Solution()
    time_elapsed = time.time() - time1
    fx = solver.bestEnergy
    print "Solution: ", sol
    print "f value: ", fx
    print "Iterations: ", solver.generations
    print "Function evaluations: ", len(esow.x)
    print "Time elapsed: ", time_elapsed, " seconds"
    assert almostEqual(fx, 3.84824937598e-15, tol=3e-3)

    # NelderMeadSimplexSolver
    print "\nUsing NelderMeadSimplexSolver:"
#.........这里部分代码省略.........

pens = ms.generate_penalty(ms.generate_conditions(eqns), k=1e3)
bounds = [(0., None), (0., 4.)]

# get the objective
def objective(x):
  x = np.asarray(x)
  return x[0]**2 + 4*x[1]**2 - 32*x[1] + 64

x0 = np.random.rand(2)

# compare against the exact minimum
xs = np.array([2., 3.])
ys = objective(xs)


sol = my.fmin_powell(objective, x0, constraint=cons, penalty=pens, disp=False,
                     bounds=bounds, gtol=3, ftol=1e-6, full_output=True)

assert mm.almostEqual(sol[0], xs, tol=1e-2)
assert mm.almostEqual(sol[1], ys, tol=1e-2)


sol = my.diffev(objective, bounds, constraint=cons, penalty=pens, disp=False,
                bounds=bounds, npop=10, gtol=100, ftol=1e-6, full_output=True)

assert mm.almostEqual(sol[0], xs, tol=1e-2)
assert mm.almostEqual(sol[1], ys, tol=1e-2)


# EOF