Add support for category corrections with gradients

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)

src/correctionlib_gradients/_correction_with_gradient.py

@@ -1,10 +1,12 @@
 # SPDX-FileCopyrightText: 2023-present Enrico Guiraud <enrico.guiraud@pm.me>
 #
 # 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)

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)

src/correctionlib_gradients/_utils.py

@@ -1,18 +1,21 @@
 # SPDX-FileCopyrightText: 2023-present Enrico Guiraud <enrico.guiraud@pm.me>
 #
 # 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

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")