def testRestraint(self): """Test the Restraint class.""" p1 = Parameter("p1", 1) p2 = Parameter("p2", 2) factory = EquationFactory() factory.registerArgument("p1", p1) factory.registerArgument("p2", p2) # Restrain 1 < p1 + p2 < 5 eq = equationFromString("p1 + p2", factory) r = Restraint(eq, 1, 5) # This should have no penalty p1.setValue(1) p2.setValue(1) self.assertEqual(0, r.penalty()) # Make p1 + p2 = 0 # This should have a penalty of 1*(1 - 0)**2 = 1 p1.setValue(-1) p2.setValue(1) self.assertEqual(1, r.penalty()) # Make p1 + p2 = 8 # This should have a penalty of 1*(8 - 5)**2 = 9 p1.setValue(4) p2.setValue(4) self.assertEqual(9, r.penalty()) # Set the chi^2 to get a dynamic penalty r.scaled = True self.assertEqual(13.5, r.penalty(1.5)) # Make a really large number to check the upper bound import numpy r.ub = numpy.inf p1.setValue(1e100) self.assertEqual(0, r.penalty()) return
def restrain(self, res, lb=-inf, ub=inf, sig=1, scaled=False, ns={}): """Restrain an expression to specified bounds res -- An equation string or Parameter to restrain. lb -- The lower bound on the restraint evaluation (default -inf). ub -- The lower bound on the restraint evaluation (default inf). sig -- The uncertainty on the bounds (default 1). scaled -- A flag indicating if the restraint is scaled (multiplied) by the unrestrained point-average chi^2 (chi^2/numpoints) (default False). ns -- A dictionary of Parameters, indexed by name, that are used in the equation string, but not part of the RecipeOrganizer (default {}). The penalty is calculated as (max(0, lb - val, val - ub)/sig)**2 and val is the value of the calculated equation. This is multipled by the average chi^2 if scaled is True. Raises ValueError if ns uses a name that is already used for a Parameter. Raises ValueError if res depends on a Parameter that is not part of the RecipeOrganizer and that is not defined in ns. Returns the Restraint object for use with the 'unrestrain' method. """ if isinstance(res, basestring): eqstr = res eq = equationFromString(res, self._eqfactory, ns) else: eq = Equation(root=res) eqstr = res.name # Make and store the restraint res = Restraint(eq, lb, ub, sig, scaled) res.eqstr = eqstr self.addRestraint(res) return res