From fa0d2b5ebce0733cf78c2fd109876fc9d97a2c62 Mon Sep 17 00:00:00 2001 From: ligerlac Date: Thu, 13 Feb 2025 20:08:56 +0100 Subject: [PATCH 01/11] add support for Category corrections --- .../_category_with_grad.py | 165 ++++++++++++++++++ src/correctionlib_gradients/_correctiondag.py | 9 + src/correctionlib_gradients/_formuladag.py | 2 + 3 files changed, 176 insertions(+) create mode 100644 src/correctionlib_gradients/_category_with_grad.py diff --git a/src/correctionlib_gradients/_category_with_grad.py b/src/correctionlib_gradients/_category_with_grad.py new file mode 100644 index 0000000..ea78d70 --- /dev/null +++ b/src/correctionlib_gradients/_category_with_grad.py @@ -0,0 +1,165 @@ +from typing import Union, Dict, List +import jax +import jax.numpy as jnp +import correctionlib.schemav2 as schema +from correctionlib_gradients._formuladag import FormulaDAG +from correctionlib_gradients._utils import get_result_size +from correctionlib_gradients._typedefs import Value +from dataclasses import dataclass + +from correctionlib_gradients._typedefs import Value + + +@dataclass +class CategoryWithGrad: + """A JAX-friendly representation of a Category correction.""" + var: str # The category input variable (e.g. 'q') + content: Dict[int, Union[float, 'FormulaDAG']] # Map from category to value/formula + input_vars: List[str] # All input variables needed (e.g. ['phi', 'q']) + default: Union[float, None] + + @staticmethod + def from_category(category: schema.Category, inputs: List[schema.Variable], generic_formulas: Dict[str, schema.Formula] = None) -> "CategoryWithGrad": + content = {} + all_vars = {v.name for v in inputs} + + # Process each key-value pair in the category + for item in category.content: + if isinstance(item.value, float): + content[item.key] = item.value + elif isinstance(item.value, (schema.Formula, schema.FormulaRef)): + # If it's a formula, convert it to FormulaDAG + if isinstance(item.value, schema.FormulaRef) and generic_formulas: + formula = generic_formulas[item.value.ref] + else: + formula = item.value + content[item.key] = FormulaDAG(formula, inputs) + else: + raise ValueError(f"Unsupported content type in Category: {type(item.value)}") + + return CategoryWithGrad( + var=category.input, + content=content, + input_vars=[v.name for v in inputs], + default=getattr(category, 'default', None) + ) + + def evaluate(self, inputs: Dict[str, Value]) -> Value: + """Evaluate the category correction for the given inputs.""" + x = inputs[self.var] # Get the category selector value + + def _handle_single_input(xi: Value, *args: Value) -> Value: + # Convert to int for category lookup + key = int(xi) + + if key in self.content: + value = self.content[key] + if isinstance(value, float): + return value + else: # FormulaDAG + # Create input dict for formula evaluation + input_dict = {name: arg for name, arg in zip(self.input_vars, args)} + return value.evaluate(input_dict) + elif self.default is not None: + return self.default + else: + raise ValueError(f"Category key '{key}' not found and no default value specified") + + # Get all required inputs as arrays + input_arrays = [inputs[name] for name in self.input_vars] + + # Handle both scalar and array inputs + if jnp.isscalar(x) or (isinstance(x, jax.Array) and x.ndim == 0): + return _handle_single_input(x, *input_arrays) + else: + # Vectorize the function using jax.vmap for array inputs + vectorized_handler = jax.vmap(_handle_single_input) + return vectorized_handler(x, *input_arrays) + + def __call__(self, inputs: Dict[str, Value]) -> Value: + """Alias for evaluate().""" + return self.evaluate(inputs) + + +# class CategoryWithGrad: +# """A JAX-friendly representation of a Category correction.""" + +# def __init__(self, category: schema.Category, inputs: list[schema.Variable], generic_formulas: Dict[str, schema.Formula] = None): +# self.var = category.input +# self.inputs = inputs +# self.content: Dict[str, Union[float, 'FormulaDAG']] = {} + +# print(category.content) +# for cont in category.content: +# key, value = cont.key, cont.value +# if isinstance(value, float): +# self.content[key] = value +# elif isinstance(value, (schema.Formula, schema.FormulaRef)): +# # If it's a formula, convert it to FormulaDAG +# # You'll need to handle FormulaRef appropriately with generic_formulas +# if isinstance(value, schema.FormulaRef) and generic_formulas: +# formula = generic_formulas[value.ref] +# else: +# formula = value +# self.content[key] = FormulaDAG(formula, inputs) # You might need to adjust the inputs here +# else: +# raise ValueError(f"Unsupported content type in Category: {type(value)}") + +# # Handle default value if present +# self.default = category.default if hasattr(category, 'default') else None + + +# def __call__(self, x: Value) -> Value: +# """Evaluate the category correction for the given input.""" + +# print(f"x = {x}") + +# def _handle_single_input(xi): +# # Convert input to string since categories use string keys +# # key = str(int(xi)) # Convert to int first to handle both float and int inputs +# key = int(xi) + +# if key in self.content: +# value = self.content[key] +# if isinstance(value, float): +# return value +# else: # FormulaDAG +# input_dict = {"phi": jnp.array([1.7]), "q": jnp.array([1])} +# return value.evaluate(input_dict) +# elif self.default is not None: +# return self.default +# else: +# raise ValueError(f"Category key '{key}' not found and no default value specified") + +# # Handle both scalar and array inputs +# if jnp.isscalar(x) or (isinstance(x, jax.Array) and x.ndim == 0): +# return _handle_single_input(x) +# else: +# # Vectorize the function using jax.vmap for array inputs +# vectorized_handler = jax.vmap(_handle_single_input) +# return vectorized_handler(x) + +# # def __call__(self, v: Value) -> Value: +# # """Evaluate the category correction for the given input.""" +# # result_size = get_result_size({'x': x}) + +# # # Create a function to handle single input +# # def _handle_single_input(xi): +# # # Convert input to string since categories use string keys +# # key = str(xi) + +# # if key in self.content: +# # value = self.content[key] +# # if isinstance(value, float): +# # return value +# # else: # FormulaDAG +# # return value.evaluate({}) # You might need to pass appropriate inputs +# # elif self.default is not None: +# # return self.default +# # else: +# # raise ValueError(f"Category key '{key}' not found and no default value specified") + +# # # Vectorize the function using jax.vmap +# # vectorized_handler = jax.vmap(_handle_single_input) + +# # return vectorized_handler(x) diff --git a/src/correctionlib_gradients/_correctiondag.py b/src/correctionlib_gradients/_correctiondag.py index 0cea5b0..573a220 100644 --- a/src/correctionlib_gradients/_correctiondag.py +++ b/src/correctionlib_gradients/_correctiondag.py @@ -8,6 +8,7 @@ from correctionlib_gradients._compound_binning import CompoundBinning from correctionlib_gradients._formuladag import FormulaDAG from correctionlib_gradients._spline_with_grad import SplineWithGrad +from correctionlib_gradients._category_with_grad import CategoryWithGrad DAGNode: TypeAlias = float | SplineWithGrad | FormulaDAG | CompoundBinning @@ -49,6 +50,10 @@ def __init__(self, c: schema.Correction): raise ValueError(msg) case schema.Formula() as f: self.node = FormulaDAG(f, c.inputs) + case schema.Category() as category: + self.node = CategoryWithGrad.from_category(category, c.inputs, c.generic_formulas) + # case schema.Category() as category: + # self.node = CategoryWithGrad(category, c.inputs, c.generic_formulas) case _: msg = f"Correction '{c.name}' contains the unsupported operation type '{type(c.data).__name__}'" raise ValueError(msg) @@ -68,3 +73,7 @@ def evaluate(self, inputs: dict[str, jax.Array]) -> jax.Array: return f.evaluate(inputs) case CompoundBinning() as cb: return cb.evaluate(inputs) + case CategoryWithGrad() as cat: + return cat.evaluate(inputs) + # case CategoryWithGrad() as cat: + # return cat(inputs[cat.var]) diff --git a/src/correctionlib_gradients/_formuladag.py b/src/correctionlib_gradients/_formuladag.py index 45130cb..9369ec6 100644 --- a/src/correctionlib_gradients/_formuladag.py +++ b/src/correctionlib_gradients/_formuladag.py @@ -74,6 +74,8 @@ class Op: class FormulaDAG: def __init__(self, f: schema.Formula, inputs: list[schema.Variable]): + print(f"f.json() = {f.json()}") + print(f"inputs = {inputs}") cpp_formula = Formula.from_string(f.json(), [CPPVariable.from_string(v.json()) for v in inputs]) self.input_names = [v.name for v in inputs] self.node: FormulaNode = self._make_node(cpp_formula.ast) From 6d582a9b96ce84f1b170b8c78a3d7f53e0e4b748 Mon Sep 17 00:00:00 2001 From: ligerlac Date: Fri, 14 Feb 2025 10:37:21 +0100 Subject: [PATCH 02/11] wrote tests for new category support --- src/correctionlib_gradients/_formuladag.py | 2 - tests/test_correction_with_gradient.py | 66 ++++++++++++++++++---- 2 files changed, 54 insertions(+), 14 deletions(-) diff --git a/src/correctionlib_gradients/_formuladag.py b/src/correctionlib_gradients/_formuladag.py index 9369ec6..45130cb 100644 --- a/src/correctionlib_gradients/_formuladag.py +++ b/src/correctionlib_gradients/_formuladag.py @@ -74,8 +74,6 @@ class Op: class FormulaDAG: def __init__(self, f: schema.Formula, inputs: list[schema.Variable]): - print(f"f.json() = {f.json()}") - print(f"inputs = {inputs}") cpp_formula = Formula.from_string(f.json(), [CPPVariable.from_string(v.json()) for v in inputs]) self.input_names = [v.name for v in inputs] self.node: FormulaNode = self._make_node(cpp_formula.ast) diff --git a/tests/test_correction_with_gradient.py b/tests/test_correction_with_gradient.py index 9c8677c..5fd96af 100644 --- a/tests/test_correction_with_gradient.py +++ b/tests/test_correction_with_gradient.py @@ -180,16 +180,38 @@ flow="clamp", ), ), - # this type of correction is unsupported - "categorical": schemav2.Correction( - name="categorical", + # can be differentiated w.r.t. x, but not c + "categorical-with-formula": schemav2.Correction( + name="categorical-with-formula", version=2, - inputs=[schemav2.Variable(name="c", type="int")], - output=schemav2.Variable(name="weight", type="real"), + inputs=[ + schemav2.Variable(name="x", type="real"), + schemav2.Variable(name="c", type="int"), + ], + output=schemav2.Variable(name="a scale", type="real"), data=schemav2.Category( nodetype="category", - input="x", - content=[schemav2.CategoryItem(key=0, value=1.234)], + input="c", + content=[ + schemav2.CategoryItem( + key=1, + value=schemav2.Formula( + nodetype="formula", + variables=["x"], + parser="TFormula", + expression="42*x", + ), + ), + schemav2.CategoryItem( + key=-1, + value=schemav2.Formula( + nodetype="formula", + variables=["x"], + parser="TFormula", + expression="1337*x", + ), + ), + ] ), ), # this type of correction is unsupported @@ -234,11 +256,6 @@ def test_missing_input(): cg.evaluate() -def test_unsupported_correction(): - with pytest.raises(ValueError, match="Correction 'categorical' contains the unsupported operation type 'Category'"): - CorrectionWithGradient(schemas["categorical"]) - - def test_unsupported_flow_type(): with pytest.raises( ValueError, @@ -518,3 +535,28 @@ def test_compound_binning_with_formularef(): value, grad = jax.value_and_grad(cg.evaluate)(0.5) assert math.isclose(value, 0.2) assert math.isclose(grad, 0.2) + + +def test_categorical_with_formula(): + cg = CorrectionWithGradient(schemas["categorical-with-formula"]) + + value = cg.evaluate(0.5, 1) + assert math.isclose(value, 21.0) + + value, grad = jax.value_and_grad(cg.evaluate, argnums=0)(1.0, 1) + assert math.isclose(value, 42.0) + assert math.isclose(grad, 42.0) + + value = cg.evaluate(1.0, -1) + assert math.isclose(value, 1337.0) + + value, grad = jax.value_and_grad(cg.evaluate, argnums=0)(1.0, -1) + assert math.isclose(value, 1337.0) + assert math.isclose(grad, 1337.0) + + # differenttiating w.r.t. to the category key should not work + with pytest.raises(TypeError): + value, grad = jax.value_and_grad(cg.evaluate, argnums=1)(1.0, 1) + + with pytest.raises(TypeError): + value, grad = jax.value_and_grad(cg.evaluate, argnums=1)(1.0, -1) From 9d5cddc01195ae9ed784d4f4a3698231f9d3fed0 Mon Sep 17 00:00:00 2001 From: ligerlac Date: Fri, 14 Feb 2025 10:40:25 +0100 Subject: [PATCH 03/11] prettify --- .../_category_with_grad.py | 85 ------------------- src/correctionlib_gradients/_correctiondag.py | 4 - 2 files changed, 89 deletions(-) diff --git a/src/correctionlib_gradients/_category_with_grad.py b/src/correctionlib_gradients/_category_with_grad.py index ea78d70..b84ae98 100644 --- a/src/correctionlib_gradients/_category_with_grad.py +++ b/src/correctionlib_gradients/_category_with_grad.py @@ -12,7 +12,6 @@ @dataclass class CategoryWithGrad: - """A JAX-friendly representation of a Category correction.""" var: str # The category input variable (e.g. 'q') content: Dict[int, Union[float, 'FormulaDAG']] # Map from category to value/formula input_vars: List[str] # All input variables needed (e.g. ['phi', 'q']) @@ -79,87 +78,3 @@ def _handle_single_input(xi: Value, *args: Value) -> Value: def __call__(self, inputs: Dict[str, Value]) -> Value: """Alias for evaluate().""" return self.evaluate(inputs) - - -# class CategoryWithGrad: -# """A JAX-friendly representation of a Category correction.""" - -# def __init__(self, category: schema.Category, inputs: list[schema.Variable], generic_formulas: Dict[str, schema.Formula] = None): -# self.var = category.input -# self.inputs = inputs -# self.content: Dict[str, Union[float, 'FormulaDAG']] = {} - -# print(category.content) -# for cont in category.content: -# key, value = cont.key, cont.value -# if isinstance(value, float): -# self.content[key] = value -# elif isinstance(value, (schema.Formula, schema.FormulaRef)): -# # If it's a formula, convert it to FormulaDAG -# # You'll need to handle FormulaRef appropriately with generic_formulas -# if isinstance(value, schema.FormulaRef) and generic_formulas: -# formula = generic_formulas[value.ref] -# else: -# formula = value -# self.content[key] = FormulaDAG(formula, inputs) # You might need to adjust the inputs here -# else: -# raise ValueError(f"Unsupported content type in Category: {type(value)}") - -# # Handle default value if present -# self.default = category.default if hasattr(category, 'default') else None - - -# def __call__(self, x: Value) -> Value: -# """Evaluate the category correction for the given input.""" - -# print(f"x = {x}") - -# def _handle_single_input(xi): -# # Convert input to string since categories use string keys -# # key = str(int(xi)) # Convert to int first to handle both float and int inputs -# key = int(xi) - -# if key in self.content: -# value = self.content[key] -# if isinstance(value, float): -# return value -# else: # FormulaDAG -# input_dict = {"phi": jnp.array([1.7]), "q": jnp.array([1])} -# return value.evaluate(input_dict) -# elif self.default is not None: -# return self.default -# else: -# raise ValueError(f"Category key '{key}' not found and no default value specified") - -# # Handle both scalar and array inputs -# if jnp.isscalar(x) or (isinstance(x, jax.Array) and x.ndim == 0): -# return _handle_single_input(x) -# else: -# # Vectorize the function using jax.vmap for array inputs -# vectorized_handler = jax.vmap(_handle_single_input) -# return vectorized_handler(x) - -# # def __call__(self, v: Value) -> Value: -# # """Evaluate the category correction for the given input.""" -# # result_size = get_result_size({'x': x}) - -# # # Create a function to handle single input -# # def _handle_single_input(xi): -# # # Convert input to string since categories use string keys -# # key = str(xi) - -# # if key in self.content: -# # value = self.content[key] -# # if isinstance(value, float): -# # return value -# # else: # FormulaDAG -# # return value.evaluate({}) # You might need to pass appropriate inputs -# # elif self.default is not None: -# # return self.default -# # else: -# # raise ValueError(f"Category key '{key}' not found and no default value specified") - -# # # Vectorize the function using jax.vmap -# # vectorized_handler = jax.vmap(_handle_single_input) - -# # return vectorized_handler(x) diff --git a/src/correctionlib_gradients/_correctiondag.py b/src/correctionlib_gradients/_correctiondag.py index 573a220..3b2f9b0 100644 --- a/src/correctionlib_gradients/_correctiondag.py +++ b/src/correctionlib_gradients/_correctiondag.py @@ -52,8 +52,6 @@ def __init__(self, c: schema.Correction): self.node = FormulaDAG(f, c.inputs) case schema.Category() as category: self.node = CategoryWithGrad.from_category(category, c.inputs, c.generic_formulas) - # case schema.Category() as category: - # self.node = CategoryWithGrad(category, c.inputs, c.generic_formulas) case _: msg = f"Correction '{c.name}' contains the unsupported operation type '{type(c.data).__name__}'" raise ValueError(msg) @@ -75,5 +73,3 @@ def evaluate(self, inputs: dict[str, jax.Array]) -> jax.Array: return cb.evaluate(inputs) case CategoryWithGrad() as cat: return cat.evaluate(inputs) - # case CategoryWithGrad() as cat: - # return cat(inputs[cat.var]) From 0a533d940e6baaaeec8eed8d97b188140cd60c39 Mon Sep 17 00:00:00 2001 From: ligerlac Date: Fri, 14 Feb 2025 10:41:18 +0100 Subject: [PATCH 04/11] add CategoryWithGrad to node types --- src/correctionlib_gradients/_correctiondag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/correctionlib_gradients/_correctiondag.py b/src/correctionlib_gradients/_correctiondag.py index 3b2f9b0..6a27de4 100644 --- a/src/correctionlib_gradients/_correctiondag.py +++ b/src/correctionlib_gradients/_correctiondag.py @@ -10,7 +10,7 @@ from correctionlib_gradients._spline_with_grad import SplineWithGrad from correctionlib_gradients._category_with_grad import CategoryWithGrad -DAGNode: TypeAlias = float | SplineWithGrad | FormulaDAG | CompoundBinning +DAGNode: TypeAlias = float | SplineWithGrad | FormulaDAG | CompoundBinning | CategoryWithGrad class CorrectionDAG: From 5c5e4c08dfcb7d170c72676b28eb6b462e80366b Mon Sep 17 00:00:00 2001 From: ligerlac Date: Fri, 14 Feb 2025 10:43:55 +0100 Subject: [PATCH 05/11] prettify --- src/correctionlib_gradients/_category_with_grad.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/correctionlib_gradients/_category_with_grad.py b/src/correctionlib_gradients/_category_with_grad.py index b84ae98..0e7e45a 100644 --- a/src/correctionlib_gradients/_category_with_grad.py +++ b/src/correctionlib_gradients/_category_with_grad.py @@ -12,15 +12,14 @@ @dataclass class CategoryWithGrad: - var: str # The category input variable (e.g. 'q') + var: str # The category input variable name content: Dict[int, Union[float, 'FormulaDAG']] # Map from category to value/formula - input_vars: List[str] # All input variables needed (e.g. ['phi', 'q']) + input_vars: List[str] # All input variable names needed default: Union[float, None] @staticmethod def from_category(category: schema.Category, inputs: List[schema.Variable], generic_formulas: Dict[str, schema.Formula] = None) -> "CategoryWithGrad": content = {} - all_vars = {v.name for v in inputs} # Process each key-value pair in the category for item in category.content: From 2e3b3d9808bffbcb5c8a6a704d46fc5d75b70312 Mon Sep 17 00:00:00 2001 From: ligerlac Date: Fri, 14 Feb 2025 12:03:14 +0100 Subject: [PATCH 06/11] add support for string inputs --- my-script.py | 116 ++++++++++++++++++ .../_category_with_grad.py | 85 +++++++++---- .../_correction_with_gradient.py | 44 +++++-- src/correctionlib_gradients/_utils.py | 11 +- 4 files changed, 214 insertions(+), 42 deletions(-) create mode 100644 my-script.py diff --git a/my-script.py b/my-script.py new file mode 100644 index 0000000..ca68b91 --- /dev/null +++ b/my-script.py @@ -0,0 +1,116 @@ +import math + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +from correctionlib import schemav2 + +from correctionlib_gradients import CorrectionWithGradient + + +schemas = { + "categorical-with-formula": schemav2.Correction( + name="categorical-with-formula", + version=2, + inputs=[ + schemav2.Variable(name="x", type="real"), + schemav2.Variable(name="c", type="int"), + ], + output=schemav2.Variable(name="a scale", type="real"), + data=schemav2.Category( + nodetype="category", + input="c", + content=[ + schemav2.CategoryItem( + key=1, + value=schemav2.Formula( + nodetype="formula", + variables=["x"], + parser="TFormula", + expression="42*x", + ), + ), + schemav2.CategoryItem( + key=-1, + value=schemav2.Formula( + nodetype="formula", + variables=["x"], + parser="TFormula", + expression="1337*x", + ), + ), + ] + ), + ), + # can be differentiated w.r.t. x, but not c + "str-categorical-with-formula": schemav2.Correction( + name="str-categorical-with-formula", + version=2, + inputs=[ + schemav2.Variable(name="x", type="real"), + schemav2.Variable(name="c", type="string"), + ], + output=schemav2.Variable(name="a scale", type="real"), + data=schemav2.Category( + nodetype="category", + input="c", + content=[ + schemav2.CategoryItem( + key="up", + value=schemav2.Formula( + nodetype="formula", + variables=["x"], + parser="TFormula", + expression="42*x", + ), + ), + schemav2.CategoryItem( + key="down", + value=schemav2.Formula( + nodetype="formula", + variables=["x"], + parser="TFormula", + expression="1337*x", + ), + ), + ] + ), + ), + # this type of correction is unsupported + "compound-binning-with-categorical": schemav2.Correction( + name="compound binning with categorical", + version=2, + inputs=[schemav2.Variable(name="x", type="real")], + output=schemav2.Variable(name="weight", type="real"), + data=schemav2.Binning( + nodetype="binning", + input="x", + edges=schemav2.UniformBinning(n=1, low=0.0, high=1.0), + content=[ + schemav2.Category(nodetype="category", input="x", content=[schemav2.CategoryItem(key=0, value=1.234)]) + ], + flow="clamp", + ), + ), +} + +cg = CorrectionWithGradient(schemas["categorical-with-formula"]) +value = cg.evaluate(0.5, 1) +assert math.isclose(value, 21.0) + +value = cg.evaluate(0.5, -1) +assert math.isclose(value, 668.5) + + +cg = CorrectionWithGradient(schemas["str-categorical-with-formula"]) +value = cg.evaluate(0.5, "up") +assert math.isclose(value, 21.0) + +value = cg.evaluate(0.5, "down") +assert math.isclose(value, 668.5) + + +# cg = CorrectionWithGradient(schemas["str-categorical-with-formula"]) +# value = cg.evaluate(0.5, "up") +# assert math.isclose(value, 21.0) diff --git a/src/correctionlib_gradients/_category_with_grad.py b/src/correctionlib_gradients/_category_with_grad.py index 0e7e45a..3bcfe73 100644 --- a/src/correctionlib_gradients/_category_with_grad.py +++ b/src/correctionlib_gradients/_category_with_grad.py @@ -1,4 +1,4 @@ -from typing import Union, Dict, List +from typing import Union, Dict, List, Any import jax import jax.numpy as jnp import correctionlib.schemav2 as schema @@ -9,20 +9,36 @@ from correctionlib_gradients._typedefs import Value +from typing import Union, Dict, List, Any +import jax +import jax.numpy as jnp +import correctionlib.schemav2 as schema +from correctionlib_gradients._utils import get_result_size +from correctionlib_gradients._typedefs import Value +from dataclasses import dataclass @dataclass class CategoryWithGrad: - var: str # The category input variable name - content: Dict[int, Union[float, 'FormulaDAG']] # Map from category to value/formula - input_vars: List[str] # All input variable names needed + """A JAX-friendly representation of a Category correction.""" + var: str # The category input variable + content: Dict[Any, Union[float, 'FormulaDAG']] # Map from original keys to value/formula + input_vars: List[str] # All input variables needed default: Union[float, None] + _key_to_idx: Dict[Any, int] # Map from external keys (str/int) to internal indices + _idx_to_key: Dict[int, Any] # Map from internal indices to external keys @staticmethod def from_category(category: schema.Category, inputs: List[schema.Variable], generic_formulas: Dict[str, schema.Formula] = None) -> "CategoryWithGrad": content = {} + key_to_idx = {} + idx_to_key = {} - # Process each key-value pair in the category - for item in category.content: + # Create mapping between external keys and internal indices + # But keep content using original keys + for idx, item in enumerate(category.content): + key_to_idx[item.key] = idx + idx_to_key[idx] = item.key + if isinstance(item.value, float): content[item.key] = item.value elif isinstance(item.value, (schema.Formula, schema.FormulaRef)): @@ -39,41 +55,56 @@ def from_category(category: schema.Category, inputs: List[schema.Variable], gene var=category.input, content=content, input_vars=[v.name for v in inputs], - default=getattr(category, 'default', None) + default=getattr(category, 'default', None), + _key_to_idx=key_to_idx, + _idx_to_key=idx_to_key ) def evaluate(self, inputs: Dict[str, Value]) -> Value: """Evaluate the category correction for the given inputs.""" - x = inputs[self.var] # Get the category selector value + orig_x = inputs[self.var] # Get the category selector value + + # Convert JAX array to Python value if needed + if isinstance(orig_x, (jax.Array, jnp.ndarray)): + lookup_key = orig_x.item() + else: + lookup_key = orig_x + + if lookup_key not in self._key_to_idx: + raise ValueError(f"Category key '{lookup_key}' not found") + + # Convert to internal index for JAX compatibility + x = jnp.array(self._key_to_idx[lookup_key]) + + print(f"lookup_key: {lookup_key}") + print(f"self._key_to_idx: {self._key_to_idx}") + print(f"self._idx_to_key: {self._idx_to_key}") + print(f"orig_x: {orig_x}, x: {x}") def _handle_single_input(xi: Value, *args: Value) -> Value: - # Convert to int for category lookup - key = int(xi) + # Convert back to original key for content lookup + idx = int(xi) + orig_key = self._idx_to_key[idx] + value = self.content[orig_key] - if key in self.content: - value = self.content[key] - if isinstance(value, float): - return value - else: # FormulaDAG - # Create input dict for formula evaluation - input_dict = {name: arg for name, arg in zip(self.input_vars, args)} - return value.evaluate(input_dict) - elif self.default is not None: - return self.default - else: - raise ValueError(f"Category key '{key}' not found and no default value specified") + if isinstance(value, float): + return value + else: # FormulaDAG + # Create input dict for formula evaluation + input_dict = {name: arg for name, arg in zip(self.input_vars, args)} + return value.evaluate(input_dict) - # Get all required inputs as arrays - input_arrays = [inputs[name] for name in self.input_vars] + # Get all required inputs as arrays except the category variable + other_inputs = [inputs[name] for name in self.input_vars if name != self.var] # Handle both scalar and array inputs if jnp.isscalar(x) or (isinstance(x, jax.Array) and x.ndim == 0): - return _handle_single_input(x, *input_arrays) + return _handle_single_input(x, *other_inputs) else: # Vectorize the function using jax.vmap for array inputs vectorized_handler = jax.vmap(_handle_single_input) - return vectorized_handler(x, *input_arrays) + return vectorized_handler(x, *other_inputs) def __call__(self, inputs: Dict[str, Value]) -> Value: """Alias for evaluate().""" - return self.evaluate(inputs) + return self.evaluate(inputs) \ No newline at end of file diff --git a/src/correctionlib_gradients/_correction_with_gradient.py b/src/correctionlib_gradients/_correction_with_gradient.py index 5280238..041ee1e 100644 --- a/src/correctionlib_gradients/_correction_with_gradient.py +++ b/src/correctionlib_gradients/_correction_with_gradient.py @@ -1,6 +1,7 @@ # SPDX-FileCopyrightText: 2023-present Enrico Guiraud # # SPDX-License-Identifier: BSD-3-Clause +from typing import TypeAlias, Union import correctionlib.schemav2 as schema import jax import jax.numpy as jnp @@ -18,10 +19,14 @@ def __init__(self, c: schema.Correction): def evaluate(self, *inputs: Value) -> jax.Array: self._check_num_inputs(inputs) - inputs_as_jax = tuple(jnp.array(i) for i in inputs) + # inputs_as_jax = tuple(jnp.array(i) for i in inputs) + + inputs_as_jax = tuple( + i if isinstance(i, str) else jnp.array(i) + for i in inputs + ) self._check_input_types(inputs_as_jax) input_names = (v.name for v in self._input_vars) - input_dict = dict(zip(input_names, inputs_as_jax)) return self._dag.evaluate(input_dict) @@ -33,14 +38,31 @@ def _check_num_inputs(self, inputs: tuple[Value, ...]) -> None: ) raise ValueError(msg) - def _check_input_types(self, inputs: tuple[jax.Array, ...]) -> None: + def _check_input_types(self, inputs: tuple[Union[jax.Array, str], ...]) -> None: for i, v in enumerate(inputs): - in_type = v.dtype expected_type_str = self._input_vars[i].type - expected_type = {"real": np.floating, "int": np.integer}[expected_type_str] - if not np.issubdtype(in_type, expected_type): - msg = ( - f"Variable '{self._input_vars[i].name}' has type {in_type}" - f" instead of the expected {expected_type.__name__}" - ) - raise ValueError(msg) + + if expected_type_str == "string": + if not isinstance(v, str): + msg = ( + f"Variable '{self._input_vars[i].name}' should be a string " + f"but got {type(v).__name__}" + ) + raise ValueError(msg) + else: + # For numeric types, check the dtype + if isinstance(v, str): + msg = ( + f"Variable '{self._input_vars[i].name}' should be numeric " + f"but got a string" + ) + raise ValueError(msg) + + in_type = v.dtype + expected_type = {"real": np.floating, "int": np.integer}[expected_type_str] + if not np.issubdtype(in_type, expected_type): + msg = ( + f"Variable '{self._input_vars[i].name}' has type {in_type}" + f" instead of the expected {expected_type.__name__}" + ) + raise ValueError(msg) diff --git a/src/correctionlib_gradients/_utils.py b/src/correctionlib_gradients/_utils.py index b88c18c..3e8c25d 100644 --- a/src/correctionlib_gradients/_utils.py +++ b/src/correctionlib_gradients/_utils.py @@ -1,18 +1,21 @@ # SPDX-FileCopyrightText: 2023-present Enrico Guiraud # # SPDX-License-Identifier: BSD-3-Clause +from typing import Union import jax -def get_result_size(inputs: dict[str, jax.Array]) -> int: +def get_result_size(inputs: dict[str, Union[jax.Array, str]]) -> int: """Calculate what size the result of a DAG evaluation should have. - The size is equal to the one, common size (shape[0], or number or rows) of all - the non-scalar inputs we require, or 0 if all inputs are scalar. + the non-scalar numeric inputs we require, or 0 if all inputs are scalar or strings. An error is thrown in case the shapes of two non-scalar inputs differ. """ result_shape: tuple[int, ...] = () for value in inputs.values(): + # Skip string inputs when determining result size + if isinstance(value, str): + continue if result_shape == (): result_shape = value.shape elif value.shape != result_shape: @@ -21,4 +24,4 @@ def get_result_size(inputs: dict[str, jax.Array]) -> int: if result_shape != (): return result_shape[0] else: - return 0 + return 0 \ No newline at end of file From ed79d2dfec403763ad78736dae40b1f2d963697c Mon Sep 17 00:00:00 2001 From: ligerlac Date: Fri, 14 Feb 2025 12:08:11 +0100 Subject: [PATCH 07/11] prettify --- my-script.py | 116 ------------------ .../_category_with_grad.py | 23 +--- .../_correction_with_gradient.py | 5 +- 3 files changed, 7 insertions(+), 137 deletions(-) delete mode 100644 my-script.py diff --git a/my-script.py b/my-script.py deleted file mode 100644 index ca68b91..0000000 --- a/my-script.py +++ /dev/null @@ -1,116 +0,0 @@ -import math - -import jax -import jax.numpy as jnp -import numpy as np -import pytest -from correctionlib import schemav2 - -from correctionlib_gradients import CorrectionWithGradient - - -schemas = { - "categorical-with-formula": schemav2.Correction( - name="categorical-with-formula", - version=2, - inputs=[ - schemav2.Variable(name="x", type="real"), - schemav2.Variable(name="c", type="int"), - ], - output=schemav2.Variable(name="a scale", type="real"), - data=schemav2.Category( - nodetype="category", - input="c", - content=[ - schemav2.CategoryItem( - key=1, - value=schemav2.Formula( - nodetype="formula", - variables=["x"], - parser="TFormula", - expression="42*x", - ), - ), - schemav2.CategoryItem( - key=-1, - value=schemav2.Formula( - nodetype="formula", - variables=["x"], - parser="TFormula", - expression="1337*x", - ), - ), - ] - ), - ), - # can be differentiated w.r.t. x, but not c - "str-categorical-with-formula": schemav2.Correction( - name="str-categorical-with-formula", - version=2, - inputs=[ - schemav2.Variable(name="x", type="real"), - schemav2.Variable(name="c", type="string"), - ], - output=schemav2.Variable(name="a scale", type="real"), - data=schemav2.Category( - nodetype="category", - input="c", - content=[ - schemav2.CategoryItem( - key="up", - value=schemav2.Formula( - nodetype="formula", - variables=["x"], - parser="TFormula", - expression="42*x", - ), - ), - schemav2.CategoryItem( - key="down", - value=schemav2.Formula( - nodetype="formula", - variables=["x"], - parser="TFormula", - expression="1337*x", - ), - ), - ] - ), - ), - # this type of correction is unsupported - "compound-binning-with-categorical": schemav2.Correction( - name="compound binning with categorical", - version=2, - inputs=[schemav2.Variable(name="x", type="real")], - output=schemav2.Variable(name="weight", type="real"), - data=schemav2.Binning( - nodetype="binning", - input="x", - edges=schemav2.UniformBinning(n=1, low=0.0, high=1.0), - content=[ - schemav2.Category(nodetype="category", input="x", content=[schemav2.CategoryItem(key=0, value=1.234)]) - ], - flow="clamp", - ), - ), -} - -cg = CorrectionWithGradient(schemas["categorical-with-formula"]) -value = cg.evaluate(0.5, 1) -assert math.isclose(value, 21.0) - -value = cg.evaluate(0.5, -1) -assert math.isclose(value, 668.5) - - -cg = CorrectionWithGradient(schemas["str-categorical-with-formula"]) -value = cg.evaluate(0.5, "up") -assert math.isclose(value, 21.0) - -value = cg.evaluate(0.5, "down") -assert math.isclose(value, 668.5) - - -# cg = CorrectionWithGradient(schemas["str-categorical-with-formula"]) -# value = cg.evaluate(0.5, "up") -# assert math.isclose(value, 21.0) diff --git a/src/correctionlib_gradients/_category_with_grad.py b/src/correctionlib_gradients/_category_with_grad.py index 3bcfe73..f8bb5cf 100644 --- a/src/correctionlib_gradients/_category_with_grad.py +++ b/src/correctionlib_gradients/_category_with_grad.py @@ -1,28 +1,21 @@ from typing import Union, Dict, List, Any -import jax -import jax.numpy as jnp -import correctionlib.schemav2 as schema -from correctionlib_gradients._formuladag import FormulaDAG -from correctionlib_gradients._utils import get_result_size -from correctionlib_gradients._typedefs import Value from dataclasses import dataclass -from correctionlib_gradients._typedefs import Value - -from typing import Union, Dict, List, Any import jax import jax.numpy as jnp import correctionlib.schemav2 as schema + +from correctionlib_gradients._formuladag import FormulaDAG from correctionlib_gradients._utils import get_result_size from correctionlib_gradients._typedefs import Value -from dataclasses import dataclass + @dataclass class CategoryWithGrad: """A JAX-friendly representation of a Category correction.""" - var: str # The category input variable + var: str # The category input variable name content: Dict[Any, Union[float, 'FormulaDAG']] # Map from original keys to value/formula - input_vars: List[str] # All input variables needed + input_vars: List[str] default: Union[float, None] _key_to_idx: Dict[Any, int] # Map from external keys (str/int) to internal indices _idx_to_key: Dict[int, Any] # Map from internal indices to external keys @@ -42,7 +35,6 @@ def from_category(category: schema.Category, inputs: List[schema.Variable], gene if isinstance(item.value, float): content[item.key] = item.value elif isinstance(item.value, (schema.Formula, schema.FormulaRef)): - # If it's a formula, convert it to FormulaDAG if isinstance(item.value, schema.FormulaRef) and generic_formulas: formula = generic_formulas[item.value.ref] else: @@ -76,11 +68,6 @@ def evaluate(self, inputs: Dict[str, Value]) -> Value: # Convert to internal index for JAX compatibility x = jnp.array(self._key_to_idx[lookup_key]) - print(f"lookup_key: {lookup_key}") - print(f"self._key_to_idx: {self._key_to_idx}") - print(f"self._idx_to_key: {self._idx_to_key}") - print(f"orig_x: {orig_x}, x: {x}") - def _handle_single_input(xi: Value, *args: Value) -> Value: # Convert back to original key for content lookup idx = int(xi) diff --git a/src/correctionlib_gradients/_correction_with_gradient.py b/src/correctionlib_gradients/_correction_with_gradient.py index 041ee1e..30a9817 100644 --- a/src/correctionlib_gradients/_correction_with_gradient.py +++ b/src/correctionlib_gradients/_correction_with_gradient.py @@ -2,10 +2,11 @@ # # SPDX-License-Identifier: BSD-3-Clause from typing import TypeAlias, Union -import correctionlib.schemav2 as schema + import jax import jax.numpy as jnp import numpy as np +import correctionlib.schemav2 as schema from correctionlib_gradients._correctiondag import CorrectionDAG from correctionlib_gradients._typedefs import Value @@ -19,8 +20,6 @@ def __init__(self, c: schema.Correction): def evaluate(self, *inputs: Value) -> jax.Array: self._check_num_inputs(inputs) - # inputs_as_jax = tuple(jnp.array(i) for i in inputs) - inputs_as_jax = tuple( i if isinstance(i, str) else jnp.array(i) for i in inputs From bd70d4807a686528793b6de859e8e27cfdb03f20 Mon Sep 17 00:00:00 2001 From: ligerlac Date: Fri, 14 Feb 2025 12:11:02 +0100 Subject: [PATCH 08/11] prettify --- src/correctionlib_gradients/_correction_with_gradient.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/correctionlib_gradients/_correction_with_gradient.py b/src/correctionlib_gradients/_correction_with_gradient.py index 30a9817..9c482ed 100644 --- a/src/correctionlib_gradients/_correction_with_gradient.py +++ b/src/correctionlib_gradients/_correction_with_gradient.py @@ -1,7 +1,7 @@ # SPDX-FileCopyrightText: 2023-present Enrico Guiraud # # SPDX-License-Identifier: BSD-3-Clause -from typing import TypeAlias, Union +from typing import Union import jax import jax.numpy as jnp From 83c975962b64ddb802fcecba38031fd4934c71e8 Mon Sep 17 00:00:00 2001 From: ligerlac Date: Fri, 14 Feb 2025 12:13:42 +0100 Subject: [PATCH 09/11] test string categorical --- tests/test_correction_with_gradient.py | 63 +++++++++++++++++++++++++- 1 file changed, 61 insertions(+), 2 deletions(-) diff --git a/tests/test_correction_with_gradient.py b/tests/test_correction_with_gradient.py index 5fd96af..7e76024 100644 --- a/tests/test_correction_with_gradient.py +++ b/tests/test_correction_with_gradient.py @@ -181,7 +181,7 @@ ), ), # can be differentiated w.r.t. x, but not c - "categorical-with-formula": schemav2.Correction( + "int-categorical-with-formula": schemav2.Correction( name="categorical-with-formula", version=2, inputs=[ @@ -214,6 +214,40 @@ ] ), ), + # can be differentiated w.r.t. x, but not c + "str-categorical-with-formula": schemav2.Correction( + name="categorical-with-formula", + version=2, + inputs=[ + schemav2.Variable(name="x", type="real"), + schemav2.Variable(name="c", type="string"), + ], + output=schemav2.Variable(name="a scale", type="real"), + data=schemav2.Category( + nodetype="category", + input="c", + content=[ + schemav2.CategoryItem( + key="up", + value=schemav2.Formula( + nodetype="formula", + variables=["x"], + parser="TFormula", + expression="42*x", + ), + ), + schemav2.CategoryItem( + key="down", + value=schemav2.Formula( + nodetype="formula", + variables=["x"], + parser="TFormula", + expression="1337*x", + ), + ), + ] + ), + ), # this type of correction is unsupported "compound-binning-with-categorical": schemav2.Correction( name="compound binning with categorical", @@ -538,7 +572,7 @@ def test_compound_binning_with_formularef(): def test_categorical_with_formula(): - cg = CorrectionWithGradient(schemas["categorical-with-formula"]) + cg = CorrectionWithGradient(schemas["int-categorical-with-formula"]) value = cg.evaluate(0.5, 1) assert math.isclose(value, 21.0) @@ -560,3 +594,28 @@ def test_categorical_with_formula(): with pytest.raises(TypeError): value, grad = jax.value_and_grad(cg.evaluate, argnums=1)(1.0, -1) + + +def test_str_categorical_with_formula(): + cg = CorrectionWithGradient(schemas["str-categorical-with-formula"]) + + value = cg.evaluate(0.5, "up") + assert math.isclose(value, 21.0) + + value, grad = jax.value_and_grad(cg.evaluate, argnums=0)(1.0, "up") + assert math.isclose(value, 42.0) + assert math.isclose(grad, 42.0) + + value = cg.evaluate(1.0, "down") + assert math.isclose(value, 1337.0) + + value, grad = jax.value_and_grad(cg.evaluate, argnums=0)(1.0, "down") + assert math.isclose(value, 1337.0) + assert math.isclose(grad, 1337.0) + + # differenttiating w.r.t. to the category key should not work + with pytest.raises(TypeError): + value, grad = jax.value_and_grad(cg.evaluate, argnums=1)(1.0, "up") + + with pytest.raises(TypeError): + value, grad = jax.value_and_grad(cg.evaluate, argnums=1)(1.0, "down") From 224f019f8a5e167fb8ede4e275bfd9787d5731d4 Mon Sep 17 00:00:00 2001 From: ligerlac Date: Fri, 14 Feb 2025 12:15:10 +0100 Subject: [PATCH 10/11] prettify --- tests/test_correction_with_gradient.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_correction_with_gradient.py b/tests/test_correction_with_gradient.py index 7e76024..7fea2bd 100644 --- a/tests/test_correction_with_gradient.py +++ b/tests/test_correction_with_gradient.py @@ -214,7 +214,7 @@ ] ), ), - # can be differentiated w.r.t. x, but not c + # same as above 'int-categorical-with-formula' but with str indexing "str-categorical-with-formula": schemav2.Correction( name="categorical-with-formula", version=2, From 2b50092ff5259af76bb4e1b9ad0ed4083d622e6f Mon Sep 17 00:00:00 2001 From: ligerlac Date: Fri, 14 Feb 2025 13:46:11 +0100 Subject: [PATCH 11/11] support nested categories --- .../_category_with_grad.py | 41 ++++++++++++++----- 1 file changed, 30 insertions(+), 11 deletions(-) diff --git a/src/correctionlib_gradients/_category_with_grad.py b/src/correctionlib_gradients/_category_with_grad.py index f8bb5cf..ac6f0a7 100644 --- a/src/correctionlib_gradients/_category_with_grad.py +++ b/src/correctionlib_gradients/_category_with_grad.py @@ -13,9 +13,9 @@ @dataclass class CategoryWithGrad: """A JAX-friendly representation of a Category correction.""" - var: str # The category input variable name - content: Dict[Any, Union[float, 'FormulaDAG']] # Map from original keys to value/formula - input_vars: List[str] + var: str # The category input variable + content: Dict[Any, Union[float, 'FormulaDAG', 'CategoryWithGrad']] # Map from original keys to value/formula/category + input_vars: List[str] # All input variables needed default: Union[float, None] _key_to_idx: Dict[Any, int] # Map from external keys (str/int) to internal indices _idx_to_key: Dict[int, Any] # Map from internal indices to external keys @@ -26,9 +26,12 @@ def from_category(category: schema.Category, inputs: List[schema.Variable], gene key_to_idx = {} idx_to_key = {} - # Create mapping between external keys and internal indices - # But keep content using original keys - for idx, item in enumerate(category.content): + # Sort the keys to ensure consistent ordering + sorted_items = sorted(category.content, key=lambda x: x.key) + + for item in sorted_items: + # Use the original key directly as the index + idx = len(key_to_idx) key_to_idx[item.key] = idx idx_to_key[idx] = item.key @@ -40,6 +43,9 @@ def from_category(category: schema.Category, inputs: List[schema.Variable], gene else: formula = item.value content[item.key] = FormulaDAG(formula, inputs) + elif isinstance(item.value, schema.Category): + # Recursively handle nested categories + content[item.key] = CategoryWithGrad.from_category(item.value, inputs, generic_formulas) else: raise ValueError(f"Unsupported content type in Category: {type(item.value)}") @@ -63,11 +69,13 @@ def evaluate(self, inputs: Dict[str, Value]) -> Value: lookup_key = orig_x if lookup_key not in self._key_to_idx: - raise ValueError(f"Category key '{lookup_key}' not found") + if self.default is not None: + return jnp.array(self.default) + raise ValueError(f"Category key '{lookup_key}' not found and no default specified") # Convert to internal index for JAX compatibility x = jnp.array(self._key_to_idx[lookup_key]) - + def _handle_single_input(xi: Value, *args: Value) -> Value: # Convert back to original key for content lookup idx = int(xi) @@ -76,9 +84,20 @@ def _handle_single_input(xi: Value, *args: Value) -> Value: if isinstance(value, float): return value + elif isinstance(value, CategoryWithGrad): + # For nested categories, ensure we pass through only the inputs they need + needed_inputs = {k: v for k, v in inputs.items() if k in value.input_vars} + return value.evaluate(needed_inputs) else: # FormulaDAG - # Create input dict for formula evaluation - input_dict = {name: arg for name, arg in zip(self.input_vars, args)} + # Create input dict for formula evaluation, checking that we have all needed variables + input_dict = {} + # We need all variables that the formula depends on + for name, arg in zip(self.input_vars, args): + input_dict[name] = arg + # Also include any other variables from the inputs that the formula might need + for name in inputs: + if name not in input_dict: + input_dict[name] = inputs[name] return value.evaluate(input_dict) # Get all required inputs as arrays except the category variable @@ -94,4 +113,4 @@ def _handle_single_input(xi: Value, *args: Value) -> Value: def __call__(self, inputs: Dict[str, Value]) -> Value: """Alias for evaluate().""" - return self.evaluate(inputs) \ No newline at end of file + return self.evaluate(inputs)