diff --git a/src/correctionlib_gradients/_category_with_grad.py b/src/correctionlib_gradients/_category_with_grad.py new file mode 100644 index 0000000..ac6f0a7 --- /dev/null +++ b/src/correctionlib_gradients/_category_with_grad.py @@ -0,0 +1,116 @@ +from typing import Union, Dict, List, Any +from dataclasses import dataclass + +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 + + +@dataclass +class CategoryWithGrad: + """A JAX-friendly representation of a Category correction.""" + 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 + + @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 = {} + + # 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 + + if isinstance(item.value, float): + content[item.key] = item.value + elif isinstance(item.value, (schema.Formula, schema.FormulaRef)): + 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) + 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)}") + + return CategoryWithGrad( + var=category.input, + content=content, + input_vars=[v.name for v in inputs], + 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.""" + 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: + 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) + orig_key = self._idx_to_key[idx] + value = self.content[orig_key] + + 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, 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 + 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, *other_inputs) + else: + # Vectorize the function using jax.vmap for array inputs + vectorized_handler = jax.vmap(_handle_single_input) + return vectorized_handler(x, *other_inputs) + + def __call__(self, inputs: Dict[str, Value]) -> Value: + """Alias for evaluate().""" + return self.evaluate(inputs) diff --git a/src/correctionlib_gradients/_correction_with_gradient.py b/src/correctionlib_gradients/_correction_with_gradient.py index 5280238..9c482ed 100644 --- a/src/correctionlib_gradients/_correction_with_gradient.py +++ b/src/correctionlib_gradients/_correction_with_gradient.py @@ -1,10 +1,12 @@ # SPDX-FileCopyrightText: 2023-present Enrico Guiraud # # SPDX-License-Identifier: BSD-3-Clause -import correctionlib.schemav2 as schema +from typing import Union + 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 @@ -18,10 +20,12 @@ 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 + ) 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 +37,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/_correctiondag.py b/src/correctionlib_gradients/_correctiondag.py index 0cea5b0..6a27de4 100644 --- a/src/correctionlib_gradients/_correctiondag.py +++ b/src/correctionlib_gradients/_correctiondag.py @@ -8,8 +8,9 @@ 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 +DAGNode: TypeAlias = float | SplineWithGrad | FormulaDAG | CompoundBinning | CategoryWithGrad class CorrectionDAG: @@ -49,6 +50,8 @@ 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 _: msg = f"Correction '{c.name}' contains the unsupported operation type '{type(c.data).__name__}'" raise ValueError(msg) @@ -68,3 +71,5 @@ 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) 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 diff --git a/tests/test_correction_with_gradient.py b/tests/test_correction_with_gradient.py index 9c8677c..7fea2bd 100644 --- a/tests/test_correction_with_gradient.py +++ b/tests/test_correction_with_gradient.py @@ -180,16 +180,72 @@ flow="clamp", ), ), - # this type of correction is unsupported - "categorical": schemav2.Correction( - name="categorical", + # can be differentiated w.r.t. x, but not c + "int-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", + ), + ), + ] + ), + ), + # same as above 'int-categorical-with-formula' but with str indexing + "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 @@ -234,11 +290,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 +569,53 @@ 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["int-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) + + +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")