def binning(flow):
cset = wrap(
schema.Correction(
name="test",
version=2,
inputs=[schema.Variable(name="x", type="real")],
output=schema.Variable(name="a scale", type="real"),
data=schema.Binning(
nodetype="binning",
input="x",
edges=[0.0, 1.0, 3.0],
content=[1.0, 2.0],
flow=flow,
),
))
return cset["test"]
def test_formularef():
cset = wrap(
schema.Correction(
name="reftest",
version=2,
inputs=[
schema.Variable(name="x", type="real"),
],
output=schema.Variable(name="a scale", type="real"),
generic_formulas=[
schema.Formula(
nodetype="formula",
expression="[0] + [1]*x",
parser="TFormula",
variables=["x"],
),
],
data=schema.Binning(
nodetype="binning",
input="x",
edges=[0, 1, 2, 3],
content=[
schema.FormulaRef(
nodetype="formularef", index=0, parameters=[0.1, 0.2]
),
schema.FormulaRef(
nodetype="formularef", index=0, parameters=[1.1, -0.2]
),
schema.FormulaRef(
nodetype="formularef", index=0, parameters=[3.1, 0.5]
),
],
flow="error",
),
)
)
corr = cset["reftest"]
assert corr.evaluate(0.5) == 0.1 + 0.2 * 0.5
assert corr.evaluate(1.5) == 1.1 + -0.2 * 1.5
assert corr.evaluate(2.5) == 3.1 + 0.5 * 2.5
def multibinning(flow):
cset = wrap(
schema.Correction(
name="test",
version=2,
inputs=[
schema.Variable(name="x", type="real"),
schema.Variable(name="y", type="real"),
],
output=schema.Variable(name="a scale", type="real"),
data=schema.MultiBinning(
nodetype="multibinning",
inputs=["x", "y"],
edges=[
[0.0, 1.0, 3.0],
[10.0, 20.0, 30.0, 40.0],
],
content=[float(i) for i in range(2 * 3)],
flow=flow,
),
))
return cset["test"]
def make_cat(items, default):
cset = wrap(
schema.Correction(
name="test",
version=2,
inputs=[
schema.Variable(
name="cat",
type="string" if isinstance(next(iter(items)), str) else "int",
)
],
output=schema.Variable(name="a scale", type="real"),
data=schema.Category(
nodetype="category",
input="cat",
content=[
{"key": key, "value": value} for key, value in items.items()
],
default=default,
),
)
)
return cset["test"]
def test_highlevel():
cset = correctionlib.CorrectionSet(
model.CorrectionSet(
schema_version=model.VERSION,
corrections=[
model.Correction(
name="test corr",
version=2,
inputs=[
model.Variable(name="a", type="real"),
model.Variable(name="b", type="real"),
],
output=model.Variable(name="a scale", type="real"),
data=1.234,
)
],
))
assert set(cset) == {"test corr"}
sf = cset["test corr"]
assert sf.version == 2
assert sf.description == ""
with pytest.raises(RuntimeError):
sf.evaluate(0, 1.2, 35.0, 0.01)
assert sf.evaluate(1.0, 1.0) == 1.234
numpy.testing.assert_array_equal(
sf.evaluate(numpy.ones((3, 4)), 1.0),
numpy.full((3, 4), 1.234),
)
numpy.testing.assert_array_equal(
sf.evaluate(numpy.ones((3, 4)), numpy.ones(4)),
numpy.full((3, 4), 1.234),
)
sf2 = pickle.loads(pickle.dumps(sf))
assert sf2.evaluate(1.0, 1.0) == 1.234
def test_transform():
cset = wrap(
schema.Correction(
name="test",
version=2,
inputs=[
schema.Variable(name="torewrite", type="real"),
],
output=schema.Variable(name="a scale", type="real"),
data=schema.Transform(
nodetype="transform",
input="torewrite",
rule=0.1,
content=schema.Formula(
nodetype="formula",
expression="x",
parser="TFormula",
variables=["torewrite"],
),
),
))
corr = cset["test"]
assert corr.evaluate(0.5) == 0.1
assert corr.evaluate(1.5) == 0.1
cset = wrap(
schema.Correction(
name="test",
version=2,
inputs=[
schema.Variable(name="torewrite", type="int"),
],
output=schema.Variable(name="a scale", type="real"),
data=schema.Transform(
nodetype="transform",
input="torewrite",
rule=schema.Category(
nodetype="category",
input="torewrite",
content=[
{
"key": 0,
"value": 0
},
{
"key": 1,
"value": 4
},
{
"key": 2,
"value": 0
},
{
"key": 9,
"value": 3.000001
},
{
"key": 10,
"value": 2.999999
},
],
),
content=schema.Category(
nodetype="category",
input="torewrite",
content=[
{
"key": 0,
"value": 0.0
},
{
"key": 3,
"value": 0.1
},
{
"key": 4,
"value": 0.2
},
],
),
),
))
corr = cset["test"]
assert corr.evaluate(0) == 0.0
assert corr.evaluate(1) == 0.2
assert corr.evaluate(2) == 0.0
with pytest.raises(IndexError):
corr.evaluate(3)
assert corr.evaluate(9) == 0.1
assert corr.evaluate(10) == 0.1
def test_evaluator_v1():
with pytest.raises(RuntimeError):
cset = core.CorrectionSet.from_string("{")
with pytest.raises(RuntimeError):
cset = core.CorrectionSet.from_string("{}")
with pytest.raises(RuntimeError):
cset = core.CorrectionSet.from_string('{"schema_version": "blah"}')
cset = wrap(
schema.Correction(
name="test corr",
version=2,
inputs=[],
output=schema.Variable(name="a scale", type="real"),
data=1.234,
))
assert set(cset) == {"test corr"}
sf = cset["test corr"]
assert sf.version == 2
assert sf.description == ""
with pytest.raises(RuntimeError):
sf.evaluate(0, 1.2, 35.0, 0.01)
assert sf.evaluate() == 1.234
cset = wrap(
schema.Correction(
name="test corr",
version=2,
inputs=[
schema.Variable(name="pt", type="real"),
schema.Variable(name="syst", type="string"),
],
output=schema.Variable(name="a scale", type="real"),
data=schema.Binning.parse_obj({
"nodetype":
"binning",
"input":
"pt",
"edges": [0, 20, 40],
"flow":
"error",
"content": [
schema.Category.parse_obj({
"nodetype":
"category",
"input":
"syst",
"content": [
{
"key": "blah",
"value": 1.1
},
{
"key": "blah2",
"value": 2.2
},
],
}),
schema.Category.parse_obj({
"nodetype":
"category",
"input":
"syst",
"content": [
{
"key": "blah2",
"value": 1.3
},
{
"key": "blah3",
"value": {
"nodetype": "formula",
"expression": "0.25*x + exp([0])",
"parser": "TFormula",
"variables": ["pt"],
"parameters": [3.1],
},
},
],
}),
],
}),
))
assert set(cset) == {"test corr"}
sf = cset["test corr"]
assert sf.version == 2
assert sf.description == ""
with pytest.raises(RuntimeError):
# too many inputs
sf.evaluate(0, 1.2, 35.0, 0.01)
with pytest.raises(RuntimeError):
# not enough inputs
sf.evaluate(1.2)
with pytest.raises(RuntimeError):
# wrong type
sf.evaluate(5)
with pytest.raises(RuntimeError):
# wrong type
sf.evaluate("asdf")
assert sf.evaluate(12.0, "blah") == 1.1
# Do we need pytest.approx? Maybe not
assert sf.evaluate(31.0, "blah3") == 0.25 * 31.0 + math.exp(3.1)
def test_core_vectorized():
cset = wrap(
schema.Correction(
name="test",
version=1,
inputs=[
schema.Variable(name="a", type="real"),
schema.Variable(name="b", type="int"),
schema.Variable(name="c", type="string"),
],
output=schema.Variable(name="a scale", type="real"),
data={
"nodetype":
"category",
"input":
"b",
"content": [{
"key": 1,
"value": {
"nodetype": "formula",
"expression": "x",
"parser": "TFormula",
"variables": ["a"],
},
}],
"default":
-99.0,
},
))
corr = cset["test"]
assert corr.evaluate(0.3, 1, "") == 0.3
assert corr.evalv(0.3, 1, "") == 0.3
numpy.testing.assert_array_equal(
corr.evalv(numpy.full(10, 0.3), 1, ""),
numpy.full(10, 0.3),
)
numpy.testing.assert_array_equal(
corr.evalv(0.3, numpy.full(10, 1), ""),
numpy.full(10, 0.3),
)
numpy.testing.assert_array_equal(
corr.evalv(numpy.full(10, 0.3), numpy.full(10, 1), ""),
numpy.full(10, 0.3),
)
with pytest.raises(ValueError):
corr.evalv(numpy.full(5, 0.3), numpy.full(10, 1), "")
with pytest.raises(ValueError):
corr.evalv(numpy.full((10, 2), 0.3), 1, "")
with pytest.raises(ValueError):
corr.evalv(0.3)
with pytest.raises(ValueError):
corr.evalv(0.3, 1, 1, 1)
with pytest.raises(ValueError):
corr.evalv(0.3, 1, numpy.full(10, "asdf"))
a = numpy.linspace(-3, 3, 100)
b = numpy.arange(100) % 3
numpy.testing.assert_array_equal(
corr.evalv(a, b, ""),
numpy.where(b == 1, a, -99.0),
)