Extending Nonlinear Module¶

Guide for adding custom nonlinear term types to LumiX.

Overview¶

You can extend the nonlinear module by:

Adding New Term Types: Define new dataclasses for additional nonlinear operations
Custom Linearization: Implement linearization methods for your terms
Integration: Register terms with the linearization engine

Adding New Term Types¶

Step-by-Step Guide¶

Example: Adding Cubic Terms¶

Let’s add support for cubic terms (x³):

Step 1: Define the dataclass in src/lumix/nonlinear/terms.py:

@dataclass
class LXCubicTerm:
    """Cubic term: coefficient * x³.

    Represents the cube of a variable, linearized using piecewise-linear
    approximation or auxiliary variables.

    Attributes:
        var: The variable to cube.
        coefficient: Coefficient to multiply the cubic term by (default: 1.0).

    Example:
        Cubic cost function::

            volume = LXVariable[Container, float]("volume").continuous().bounds(0, 10)
            cubic_cost = LXCubicTerm(var=volume, coefficient=2.5)
    """

    var: LXVariable
    coefficient: float = 1.0

Step 2: Export from __init__.py:

from .terms import (
    LXAbsoluteTerm,
    LXBilinearTerm,
    LXCubicTerm,  # Add new term
    LXIndicatorTerm,
    LXMinMaxTerm,
    LXPiecewiseLinearTerm,
)

__all__ = [
    "LXAbsoluteTerm",
    "LXBilinearTerm",
    "LXCubicTerm",  # Add to exports
    "LXIndicatorTerm",
    "LXMinMaxTerm",
    "LXPiecewiseLinearTerm",
]

Step 3: Add linearization in src/lumix/linearization/engine.py:

def linearize_cubic(self, term: LXCubicTerm) -> ...:
    """Linearize x³ using piecewise-linear approximation."""
    # Get variable bounds
    x_min = term.var.lower_bound or 0
    x_max = term.var.upper_bound or 1

    # Create piecewise-linear approximation
    def cubic_func(x):
        return term.coefficient * (x ** 3)

    # Use existing piecewise linearization
    piecewise_term = LXPiecewiseLinearTerm(
        var=term.var,
        func=cubic_func,
        num_segments=20,
        x_min=x_min,
        x_max=x_max,
        adaptive=True
    )

    return self.linearize_piecewise(piecewise_term)

Step 4: Add tests:

def test_cubic_term():
    var = LXVariable[Item, float]("x").continuous().bounds(0, 10)
    term = LXCubicTerm(var=var, coefficient=2.0)

    assert term.var == var
    assert term.coefficient == 2.0

def test_cubic_linearization():
    var = LXVariable[Item, float]("x").continuous().bounds(0, 10)
    term = LXCubicTerm(var=var, coefficient=1.0)

    linearizer = LXLinearizer()
    result = linearizer.linearize_cubic(term)
    # Verify linearization constraints

More Examples¶

Logical AND Term¶

@dataclass
class LXAndTerm:
    """Logical AND of multiple binary variables.

    z = x₁ AND x₂ AND ... AND xₙ

    Attributes:
        vars: List of binary variables to AND together.
    """
    vars: List[LXVariable]

# Linearization (in engine):
def linearize_and(self, term: LXAndTerm):
    """z = 1 iff all vars = 1."""
    z = self._create_aux_var("and", var_type=BINARY)

    # z <= xᵢ for all i
    for var in term.vars:
        self._add_constraint(z <= var)

    # z >= sum(xᵢ) - n + 1
    self._add_constraint(z >= sum(term.vars) - len(term.vars) + 1)

    return z

Absolute Difference¶

@dataclass
class LXAbsDifferenceTerm:
    """Absolute difference between two variables: |x - y|.

    Attributes:
        var1: First variable.
        var2: Second variable.
        coefficient: Coefficient (default: 1.0).
    """
    var1: LXVariable
    var2: LXVariable
    coefficient: float = 1.0

# Linearization:
def linearize_abs_difference(self, term: LXAbsDifferenceTerm):
    """Linearize |x - y|."""
    # Create auxiliary for difference
    diff = self._create_aux_var("diff")
    self._add_constraint(diff == term.var1 - term.var2)

    # Use absolute value linearization
    abs_term = LXAbsoluteTerm(var=diff, coefficient=term.coefficient)
    return self.linearize_absolute(abs_term)

Custom Linearization Methods¶

Overriding Default Linearization¶

Create custom linearizer subclass:

from lumix.linearization import LXLinearizer, LXLinearizerConfig
from lumix.nonlinear import LXBilinearTerm

class TightBoundsLinearizer(LXLinearizer):
    """Linearizer with improved bound computation."""

    def linearize_bilinear(self, term: LXBilinearTerm):
        """Custom McCormick with tighter bounds."""
        # Compute tighter bounds using problem structure
        tight_bounds = self._compute_tight_bounds(term)

        # Apply custom McCormick
        return self._mccormick_with_bounds(term, tight_bounds)

    def _compute_tight_bounds(self, term):
        # Custom logic to compute tighter bounds
        pass

Alternative Formulations¶

Provide multiple linearization strategies:

class FlexibleLinearizer(LXLinearizer):
    def __init__(self, config: LXLinearizerConfig):
        super().__init__(config)
        self.bilinear_method = config.bilinear_method

    def linearize_bilinear(self, term: LXBilinearTerm):
        if self.bilinear_method == "mccormick":
            return self._mccormick(term)
        elif self.bilinear_method == "logarithmic":
            return self._logarithmic_formulation(term)
        else:
            return super().linearize_bilinear(term)

Integration with Model Builder¶

Automatic Detection¶

# In linearization engine
TERM_HANDLERS = {
    LXAbsoluteTerm: linearize_absolute,
    LXBilinearTerm: linearize_bilinear,
    LXCubicTerm: linearize_cubic,  # New handler
    # ...
}

def linearize_term(self, term):
    """Dispatch to appropriate handler."""
    handler = TERM_HANDLERS.get(type(term))
    if handler is None:
        raise ValueError(f"No handler for {type(term)}")
    return handler(self, term)

Testing Custom Terms¶

Unit Tests¶

def test_custom_term_creation():
    var = LXVariable[Data, float]("x").continuous()
    term = LXCubicTerm(var=var, coefficient=2.0)
    assert isinstance(term, LXCubicTerm)
    assert term.coefficient == 2.0

Integration Tests¶

def test_custom_term_in_model():
    var = LXVariable[Data, float]("x").continuous().bounds(0, 10)
    term = LXCubicTerm(var=var)

    model = LXModel("test")
    model.add_variable(var)
    # Add term to objective or constraint

    linearizer = CustomLinearizer()
    linear_model = linearizer.linearize(model)

    # Verify linearization
    assert len(linear_model.constraints) > 0

Best Practices¶

Term Design¶

Immutable dataclasses: Use @dataclass with no methods
Type annotations: Full typing for all attributes
Default values: Provide sensible defaults where appropriate
Documentation: Comprehensive docstrings with examples

Linearization Implementation¶

Validate inputs: Check variable bounds before linearization
Efficient formulations: Minimize auxiliary variables and constraints
Numerical stability: Avoid large M values
Error handling: Raise informative errors for invalid inputs

Example: Validation¶

def linearize_cubic(self, term: LXCubicTerm):
    """Linearize x³."""
    # Validate bounds
    if term.var.lower_bound is None or term.var.upper_bound is None:
        raise ValueError(
            f"Variable '{term.var.name}' must have finite bounds "
            f"for cubic linearization"
        )

    # Validate coefficient
    if term.coefficient == 0:
        raise ValueError("Coefficient cannot be zero")

    # Proceed with linearization
    ...

Documentation¶

Document your custom terms:

@dataclass
class LXCustomTerm:
    """One-line summary.

    Detailed description of the term, including:
    - Mathematical formulation
    - Use cases
    - Linearization approach

    Attributes:
        var: Description of variable.
        param: Description of parameter.

    Example:
        Basic usage::

            var = LXVariable[Data, float]("x").continuous()
            term = LXCustomTerm(var=var, param=1.0)

    Note:
        Important notes about usage, requirements, or limitations.
    """
    var: LXVariable
    param: float = 1.0

Contributing to LumiX¶

If you develop useful custom terms, consider contributing them to LumiX:

Fork the repository
Add your term following the guidelines above
Include comprehensive tests
Update documentation
Submit a pull request