Feature/generalize conditioning

docs/examples/conditional_luminosity_function.rst

@@ -387,10 +387,12 @@ The model therefore changes both in amplitude and in shape.
        cmap_range=(0.0, 0.2),
    )
 
-   lf = ConditionalLuminosityFunction.schechter(
+   lf = ConditionalLuminosityFunction.double_schechter(
        phi_star=lambda z: 1.2e-3 * (1.0 + z) ** 0.5,
        m_star=lambda z: -20.3 - 0.5 * (z - 0.1),
        alpha=lambda z: -1.15 - 0.10 * z,
+       beta=lambda z: -0.45 - 0.05 * z,
+       m_transition=lambda z: -19.0 - 0.3 * (z - 0.1),
    )
 
    fig, ax = plt.subplots(figsize=(7.0, 5.0))

src/lfkit/api/conditional_luminosity_function.py

@@ -21,10 +21,10 @@
 class ConditionalLuminosityFunction(LuminosityFunction):
     """User-facing wrapper for conditional luminosity function models.
 
-    A conditional luminosity function evaluates ``Phi(M | x)``, where ``M`` is
-    absolute magnitude and ``x`` is an external conditioning variable such as
-    redshift, halo mass, environment, richness, stellar mass, or another
-    model-specific quantity.
+    A conditional luminosity function evaluates ``Phi(M | x_1, x_2, ...)``,
+    where ``M`` is absolute magnitude and the ``x_i`` are external conditioning
+    variables such as redshift, halo mass, environment, richness, stellar mass,
+    or other model-specific quantities.
 
     Instances can be created either with the generic constructor or with
     automatically generated model constructors.
@@ -33,34 +33,39 @@ class ConditionalLuminosityFunction(LuminosityFunction):
     def phi(
         self,
         absolute_mag: FloatInput,
-        condition: FloatInput | None = None,
+        *conditions: FloatInput,
     ) -> FloatArray:
         """Evaluate the conditional luminosity function.
 
         Args:
             absolute_mag: Absolute magnitude value or array.
-            condition: Conditioning variable value or array. The meaning of this
-                variable depends on the selected conditional luminosity function model.
+            *conditions: One or more conditioning variable values or arrays. The
+                meaning of each variable depends on the selected conditional
+                luminosity function model.
 
         Returns:
             Conditional luminosity function evaluated at ``absolute_mag`` and
-            ``condition``.
+            the supplied conditioning variables.
 
         Raises:
-            ValueError: If ``condition`` is not provided or the model is not registered
-                as a conditional luminosity function.
+            ValueError: If no conditioning variables are provided, or if the
+                model is not registered as a conditional luminosity function.
         """
         model_spec = get_conditional_lf_model(self.model)
 
-        if condition is None:
+        if not conditions:
             raise ValueError(
-                f"condition is required for conditional luminosity function "
-                f"model '{self.model}'."
+                f"At least one conditioning variable is required for conditional "
+                f"luminosity function model '{self.model}'."
             )
 
+        condition_arrays = tuple(
+            np.asarray(condition_value, dtype=float) for condition_value in conditions
+        )
+
         return model_spec.function(
             np.asarray(absolute_mag, dtype=float),
-            np.asarray(condition, dtype=float),
+            *condition_arrays,
             **self.parameters_dict,
         )
 
@@ -125,7 +130,8 @@ def constructor(
 
         Examples:
             >>> clf = ConditionalLuminosityFunction.{model_name}(...)
-            >>> phi = clf.phi(absolute_mag=-20.0, condition=0.5)
+            >>> phi = clf.phi(-20.0, 0.5)
+            >>> phi = clf.phi(-20.0, halo_mass, redshift)
         """
 
     return constructor
@@ -138,8 +144,8 @@ def _parameters_from_signature(
 ) -> dict[str, Any]:
     """Build stored parameters from a function signature and user values.
 
-    Independent variables such as ``absolute_mag`` and ``condition`` are not stored
-    as model parameters. They are supplied later when calling
+    Independent variables such as ``absolute_mag`` and conditioning variables are
+    not stored as model parameters. They are supplied later when calling
     :meth:`ConditionalLuminosityFunction.phi`.
 
     Args:
@@ -156,8 +162,20 @@ def _parameters_from_signature(
     """
     payload: dict[str, Any] = {}
 
+    independent_names = {
+        "absolute_mag",
+        "condition",
+        "conditions",
+        "z",
+        "redshift",
+        "x",
+        "halo_mass",
+        "environment",
+        "galaxy_type",
+    }
+
     for name, parameter in signature.parameters.items():
-        if name in {"absolute_mag", "condition", "z", "redshift", "x"}:
+        if name in independent_names:
             continue
 
         if parameter.kind in {

src/lfkit/luminosity_functions/conditional_integrals.py

@@ -1,12 +1,12 @@
 """Conditional luminosity function integration utilities.
 
 This module provides numerical helpers for conditional luminosity functions of
-the form ``Phi(M | x)``, where ``M`` is absolute magnitude and ``x`` is an
-external conditioning variable.
+the form ``Phi(M | x_1, x_2, ...)``, where ``M`` is absolute magnitude and the
+``x_i`` are external conditioning variables.
 
-The conditioning variable is intentionally generic. It may represent halo mass,
-environment, galaxy type, richness, stellar mass, or any other quantity. This
-module does not implement HOD or halo-model machinery.
+The conditioning variables are intentionally generic. They may represent halo
+mass, redshift, environment, galaxy type, richness, stellar mass, or any other
+quantities. This module does not implement HOD or halo-model machinery.
 
 The goal is to support conditional luminosity function evaluation and
 integration while keeping halo model calculations outside LFKit.
@@ -15,6 +15,7 @@
 from __future__ import annotations
 
 from collections.abc import Callable
+from typing import Any
 
 import numpy as np
 
@@ -28,31 +29,46 @@
 ]
 
 
+def _validate_conditions(conditions: tuple[FloatInput, ...]) -> tuple[FloatArray, ...]:
+    """Return validated conditioning-variable arrays."""
+    if not conditions:
+        raise ValueError("At least one conditioning variable is required.")
+
+    return tuple(
+        validate_array(condition, name=f"condition_{i}")
+        for i, condition in enumerate(conditions)
+    )
+
+
 def evaluate_conditional_luminosity_function(
     absolute_mag: FloatInput,
-    condition: FloatInput,
-    conditional_lf: Callable[[FloatArray, FloatArray], FloatArray],
+    *conditions: FloatInput,
+    conditional_lf: Callable[..., Any],
 ) -> FloatArray:
     """Evaluate a conditional luminosity function.
 
     Args:
         absolute_mag: Absolute magnitude values.
-        condition: Values of the conditioning variable.
-        conditional_lf: Callable returning ``Phi(M | x)`` for absolute
-            magnitude ``M`` and condition ``x``.
+        *conditions: Values of one or more conditioning variables.
+        conditional_lf: Callable returning ``Phi(M | x_1, x_2, ...)`` for
+            absolute magnitude ``M`` and the supplied conditioning variables.
 
     Returns:
         Conditional luminosity function values evaluated at the requested
         absolute magnitudes and conditioning values.
 
     Raises:
-        ValueError: If the inputs contain non-finite values, or if the evaluated
-            conditional luminosity function contains non-finite or negative values.
+        ValueError: If no conditioning variables are supplied, if the inputs
+            contain non-finite values, or if the evaluated conditional luminosity
+            function contains non-finite or negative values.
     """
     absolute_mag_arr = validate_array(absolute_mag, name="absolute_mag")
-    condition_arr = validate_array(condition, name="condition")
+    condition_arrays = _validate_conditions(conditions)
 
-    phi = np.asarray(conditional_lf(absolute_mag_arr, condition_arr), dtype=float)
+    phi = np.asarray(
+        conditional_lf(absolute_mag_arr, *condition_arrays),
+        dtype=float,
+    )
 
     if not np.all(np.isfinite(phi)):
         raise ValueError("conditional_lf returned NaN or infinite values.")
@@ -67,32 +83,32 @@ def evaluate_conditional_luminosity_function(
 
 def integrate_conditional_luminosity_function(
     absolute_mag: FloatInput,
-    condition: FloatInput,
-    conditional_lf: Callable[[FloatArray, FloatArray], FloatArray],
-    *,
+    *conditions: FloatInput,
+    conditional_lf: Callable[..., Any],
     axis: int = -1,
 ) -> FloatArray:
     """Integrate a conditional luminosity function over absolute magnitude.
 
     Args:
         absolute_mag: Absolute magnitude grid.
-        condition: Values of the conditioning variable.
-        conditional_lf: Callable returning ``Phi(M | x)`` for absolute
-            magnitude ``M`` and condition ``x``.
+        *conditions: Values of one or more conditioning variables.
+        conditional_lf: Callable returning ``Phi(M | x_1, x_2, ...)`` for
+            absolute magnitude ``M`` and the supplied conditioning variables.
         axis: Axis corresponding to the absolute magnitude grid.
 
     Returns:
         Conditional luminosity function integrated over absolute magnitude.
 
     Raises:
-        ValueError: If the inputs contain non-finite values, or if the evaluated
-            conditional luminosity function contains non-finite or negative values.
+        ValueError: If no conditioning variables are supplied, if the inputs
+            contain non-finite values, or if the evaluated conditional luminosity
+            function contains non-finite or negative values.
     """
     absolute_mag_arr = validate_array(absolute_mag, name="absolute_mag")
 
     phi = evaluate_conditional_luminosity_function(
-        absolute_mag=absolute_mag_arr,
-        condition=condition,
+        absolute_mag_arr,
+        *conditions,
         conditional_lf=conditional_lf,
     )
 
@@ -104,42 +120,45 @@ def integrate_conditional_luminosity_function(
 
 def integrate_weighted_conditional_luminosity_function(
     absolute_mag: FloatInput,
-    condition: FloatInput,
-    conditional_lf: Callable[[FloatArray, FloatArray], FloatArray],
-    weight: Callable[[FloatArray, FloatArray], FloatArray],
-    *,
+    *conditions: FloatInput,
+    conditional_lf: Callable[..., Any],
+    weight: Callable[..., Any],
     axis: int = -1,
 ) -> FloatArray:
     """Integrate a weighted conditional luminosity function.
 
     Args:
         absolute_mag: Absolute magnitude grid.
-        condition: Values of the conditioning variable.
-        conditional_lf: Callable returning ``Phi(M | x)`` for absolute
-            magnitude ``M`` and condition ``x``.
-        weight: Callable returning weights ``w(M, x)`` for absolute magnitude
-            ``M`` and condition ``x``.
+        *conditions: Values of one or more conditioning variables.
+        conditional_lf: Callable returning ``Phi(M | x_1, x_2, ...)`` for
+            absolute magnitude ``M`` and the supplied conditioning variables.
+        weight: Callable returning weights ``w(M, x_1, x_2, ...)`` for absolute
+            magnitude ``M`` and the supplied conditioning variables.
         axis: Axis corresponding to the absolute magnitude grid.
 
     Returns:
         Weighted conditional luminosity function integrated over absolute
         magnitude.
 
     Raises:
-        ValueError: If the inputs contain non-finite values, if the evaluated
-            conditional luminosity function contains non-finite or negative values,
-            or if the weights contain non-finite values.
+        ValueError: If no conditioning variables are supplied, if the inputs
+            contain non-finite values, if the evaluated conditional luminosity
+            function contains non-finite or negative values, or if the weights
+            contain non-finite values.
     """
     absolute_mag_arr = validate_array(absolute_mag, name="absolute_mag")
-    condition_arr = validate_array(condition, name="condition")
+    condition_arrays = _validate_conditions(conditions)
 
     phi = evaluate_conditional_luminosity_function(
-        absolute_mag=absolute_mag_arr,
-        condition=condition_arr,
+        absolute_mag_arr,
+        *condition_arrays,
         conditional_lf=conditional_lf,
     )
 
-    weight_arr = np.asarray(weight(absolute_mag_arr, condition_arr), dtype=float)
+    weight_arr = np.asarray(
+        weight(absolute_mag_arr, *condition_arrays),
+        dtype=float,
+    )
 
     if not np.all(np.isfinite(weight_arr)):
         raise ValueError("weight returned NaN or infinite values.")

src/lfkit/luminosity_functions/conditional_models.py

@@ -3,10 +3,10 @@
 This module provides generic conditional wrappers around LFKit luminosity
 function models.
 
-A conditional luminosity function has the form ``Phi(M | x)``, where ``M`` is
-absolute magnitude and ``x`` is an external conditioning variable. Callable
-model parameters are evaluated at ``x`` before the wrapped luminosity function
-is evaluated.
+A conditional luminosity function has the form ``Phi(M | x_1, x_2, ...)``,
+where ``M`` is absolute magnitude and the ``x_i`` are external conditioning
+variables. Callable model parameters are evaluated at the supplied conditioning
+variables before the wrapped luminosity function is evaluated.
 """
 
 from __future__ import annotations
@@ -27,42 +27,64 @@ def conditionalize_lf_model(
 ) -> Callable[..., FloatArray]:
     """Return a conditional version of a luminosity function model.
 
-    Callable keyword arguments are interpreted as parameter models and evaluated as
-    functions of ``condition``. Non-callable keyword arguments are passed through
-    unchanged.
+    Callable keyword arguments are interpreted as parameter models and evaluated
+    as functions of the supplied conditioning variables. Non-callable keyword
+    arguments are passed through unchanged.
 
     Args:
         lf_model: Luminosity function model to wrap.
 
     Returns:
         Conditional luminosity function model with signature
-        ``conditional_model(absolute_mag, condition, **kwargs)``.
+        ``conditional_model(absolute_mag, *conditions, **kwargs)``.
     """
 
     @wraps(lf_model)
     def conditional_model(
         absolute_mag: FloatInput,
-        condition: FloatInput,
+        *conditions: FloatInput,
         **kwargs: Any,
     ) -> FloatArray:
-        condition_arr = validate_array(condition, name="condition")
+        absolute_mag_arr = validate_array(absolute_mag, name="absolute_mag")
+        condition_arrays = _validate_conditions(conditions)
 
         evaluated_kwargs: dict[str, Any] = {}
         for name, value in kwargs.items():
             if callable(value):
                 evaluated_kwargs[name] = validate_array(
-                    value(condition_arr),
+                    value(*condition_arrays),
                     name=name,
                 )
             else:
                 evaluated_kwargs[name] = value
 
-        phi = lf_model(absolute_mag, **evaluated_kwargs)
+        phi = lf_model(absolute_mag_arr, **evaluated_kwargs)
         return _validate_lf_output(phi, name=lf_model.__name__)
 
     return conditional_model
 
 
+def _validate_conditions(conditions: tuple[FloatInput, ...]) -> tuple[FloatArray, ...]:
+    """Return validated conditioning variable arrays.
+
+    Args:
+        conditions: Conditioning variable values.
+
+    Returns:
+        Validated conditioning variable arrays.
+
+    Raises:
+        ValueError: If no conditioning variables are supplied.
+    """
+    if not conditions:
+        raise ValueError("At least one conditioning variable is required.")
+
+    return tuple(
+        validate_array(condition, name=f"condition_{i}")
+        for i, condition in enumerate(conditions)
+    )
+
+
 def _conditional_model_name(name: str) -> str:
     """Return the generated conditional wrapper name for a luminosity function model.