Extending Linearization

Guide for developers who want to add new linearization techniques or customize existing ones.

Overview

The linearization module is designed to be extensible. You can:

1. Add new linearization techniques for different term types

2. Customize existing techniques with new formulations

3. Add pre-built function approximations to LXNonLinearFunctions

4. Create custom configuration options for your techniques

Adding a New Linearization Technique

Example: Linearizing Absolute Value Squared (abs(x)²)

Let’s create a new linearizer for absolute value squared terms.

Step 1: Define the Linearizer Class

from typing import List
from lumix.core.variables import LXVariable
from lumix.core.constraints import LXConstraint
from lumix.core.expressions import LXLinearExpression
from lumix.linearization.config import LXLinearizerConfig

class LXAbsSquaredLinearizer:
    """
    Linearize absolute value squared: z = |x|²

    Uses the reformulation:
    - t = |x| (linearized via standard absolute value)
    - z = t² (linearized via piecewise-linear or McCormick if bounded)
    """

    def __init__(self, config: LXLinearizerConfig):
        """
        Initialize the linearizer.

        Args:
            config: Linearization configuration
        """
        self.config = config
        self.auxiliary_vars: List[LXVariable] = []
        self.auxiliary_constraints: List[LXConstraint] = []
        self._aux_counter = 0

    def linearize_abs_squared(
        self,
        var: LXVariable,
        coefficient: float = 1.0
    ) -> LXVariable:
        """
        Linearize |x|² term.

        Args:
            var: Input variable
            coefficient: Multiplier

        Returns:
            Auxiliary variable representing |x|²

        Example:
            >>> linearizer = LXAbsSquaredLinearizer(config)
            >>> z = linearizer.linearize_abs_squared(x_var)
            >>> # z represents |x|²
        """
        # Step 1: Create auxiliary variable for |x|
        abs_name = f"aux_abs_{var.name}_{self._aux_counter}"
        self._aux_counter += 1

        t = (
            LXVariable[str, float](abs_name)
            .continuous()
            .bounds(lower=0, upper=None)  # |x| ≥ 0
            .indexed_by(lambda x: x)
            .from_data([abs_name])
        )
        self.auxiliary_vars.append(t)

        # Step 2: Add constraints for |x|
        # t ≥ x
        self.auxiliary_constraints.append(
            LXConstraint(f"{abs_name}_ge_x")
            .expression(
                LXLinearExpression()
                .add_term(t, 1.0)
                .add_term(var, -1.0)
            )
            .ge()
            .rhs(0)
        )

        # t ≥ -x
        self.auxiliary_constraints.append(
            LXConstraint(f"{abs_name}_ge_neg_x")
            .expression(
                LXLinearExpression()
                .add_term(t, 1.0)
                .add_term(var, 1.0)
            )
            .ge()
            .rhs(0)
        )

        # Step 3: Create auxiliary variable for t²
        # Use piecewise-linear approximation
        squared_name = f"aux_squared_{abs_name}_{self._aux_counter}"
        self._aux_counter += 1

        # Determine bounds for t²
        t_upper = None
        if var.upper_bound is not None and var.lower_bound is not None:
            t_upper = max(abs(var.lower_bound), abs(var.upper_bound))
            z_upper = t_upper ** 2
        else:
            z_upper = None

        z = (
            LXVariable[str, float](squared_name)
            .continuous()
            .bounds(lower=0, upper=z_upper)
            .indexed_by(lambda x: x)
            .from_data([squared_name])
        )
        self.auxiliary_vars.append(z)

        # Step 4: Add piecewise-linear constraints for z = t²
        # (This is simplified - in practice, use LXPiecewiseLinearizer)
        # For demonstration, we'll assume these are added elsewhere

        return z

Step 2: Integrate into Main Engine

# In lumix/linearization/engine.py

from .techniques.abs_squared import LXAbsSquaredLinearizer

class LXLinearizer:
    def __init__(self, model, solver_capability, config=None):
        # ... existing code ...
        self._abs_squared_linearizer = LXAbsSquaredLinearizer(self.config)

    def _linearize_expression(self, expr):
        # ... existing code ...

        # Handle absolute value squared terms
        elif isinstance(term, LXAbsSquaredTerm):
            aux_var = self._abs_squared_linearizer.linearize_abs_squared(
                term.var,
                term.coefficient
            )
            linear_expr = linear_expr + LXLinearExpression().add_term(
                aux_var, 1.0
            )

    def linearize_model(self):
        # ... existing code ...

        # Collect auxiliary elements from abs_squared linearizer
        for aux_var in self._abs_squared_linearizer.auxiliary_vars:
            if aux_var not in self.auxiliary_vars:
                linearized.add_variable(aux_var)
                self.auxiliary_vars.append(aux_var)

        for aux_constraint in self._abs_squared_linearizer.auxiliary_constraints:
            if aux_constraint not in self.auxiliary_constraints:
                linearized.add_constraint(aux_constraint)
                self.auxiliary_constraints.append(aux_constraint)

        # ... existing code ...

Step 3: Add Configuration Options

# In lumix/linearization/config.py

@dataclass
class LXLinearizerConfig:
    # ... existing fields ...

    # Absolute value squared settings
    abs_squared_pwl_segments: int = 25  # For t² approximation
    abs_squared_use_mccormick: bool = False  # Alternative formulation

Step 4: Add Tests

# tests/linearization/test_abs_squared.py

def test_abs_squared_linearization():
    """Test linearization of |x|² term."""
    config = LXLinearizerConfig(abs_squared_pwl_segments=30)
    linearizer = LXAbsSquaredLinearizer(config)

    # Create variable
    x = LXVariable[Model, float]("x").bounds(lower=-10, upper=10)

    # Linearize
    z = linearizer.linearize_abs_squared(x)

    # Verify
    assert z is not None
    assert len(linearizer.auxiliary_vars) >= 2  # t and z
    assert len(linearizer.auxiliary_constraints) >= 2  # |x| constraints

def test_abs_squared_accuracy():
    """Test approximation accuracy."""
    import numpy as np

    config = LXLinearizerConfig(abs_squared_pwl_segments=50)
    linearizer = LXAbsSquaredLinearizer(config)

    x_test = np.linspace(-10, 10, 100)
    max_error = 0

    for x_val in x_test:
        true_value = abs(x_val) ** 2
        # Evaluate linearized approximation
        approx_value = evaluate_linearization(linearizer, x_val)
        error = abs(true_value - approx_value) / (true_value + 1e-10)
        max_error = max(max_error, error)

    assert max_error < 0.02, f"Error too large: {max_error}"

Adding Pre-built Functions

Example: Hyperbolic Tangent (tanh)

Add to lumix/linearization/functions.py:

class LXNonLinearFunctions:
    # ... existing methods ...

    @staticmethod
    def tanh(
        var: LXVariable,
        linearizer: LXPiecewiseLinearizer,
        segments: int = 40
    ) -> LXVariable:
        """
        Hyperbolic tangent function: tanh(x) = (e^x - e^(-x)) / (e^x + e^(-x))

        Use Cases:
            - Activation functions in neural networks
            - Saturation curves
            - Signal processing

        Args:
            var: Input variable
            linearizer: Piecewise linearizer instance
            segments: Number of segments (default: 40 for smooth approximation)

        Returns:
            Output variable representing tanh(var) ∈ [-1, 1]

        Example:
            >>> # Activation function
            >>> activation = LXNonLinearFunctions.tanh(
            ...     net_input,
            ...     linearizer,
            ...     segments=50
            ... )
        """
        return linearizer.approximate_function(
            lambda x: math.tanh(x),
            var,
            num_segments=segments,
            adaptive=True  # tanh curves sharply around x=0
        )

    @staticmethod
    def relu(
        var: LXVariable,
        linearizer: LXPiecewiseLinearizer,
        segments: int = 2  # ReLU is piecewise linear with 2 segments
    ) -> LXVariable:
        """
        Rectified Linear Unit: ReLU(x) = max(0, x)

        Use Cases:
            - Neural network activation
            - Non-negative constraints with smooth approximation

        Args:
            var: Input variable
            linearizer: Piecewise linearizer instance
            segments: Number of segments (default: 2, exact for ReLU)

        Returns:
            Output variable representing max(0, var)

        Example:
            >>> # Non-negative activation
            >>> output = LXNonLinearFunctions.relu(
            ...     weighted_sum,
            ...     linearizer
            ... )
        """
        def relu_func(x: float) -> float:
            return max(0, x)

        return linearizer.approximate_function(
            relu_func,
            var,
            num_segments=segments,
            adaptive=False  # ReLU is piecewise linear, uniform is fine
        )

Customizing Existing Techniques

Example: Custom McCormick with Additional Constraints

Extend the bilinear linearizer to add custom bound tightening:

from lumix.linearization.techniques.bilinear import LXBilinearLinearizer

class LXTightMcCormickLinearizer(LXBilinearLinearizer):
    """
    Enhanced McCormick linearizer with additional bound tightening.
    """

    def _mccormick_envelope(self, x, y, coeff):
        """
        Override to add custom bound tightening before McCormick.
        """
        # Step 1: Apply custom bound tightening
        x_tight = self._tighten_bounds(x)
        y_tight = self._tighten_bounds(y)

        # Step 2: Call parent McCormick with tightened bounds
        z = super()._mccormick_envelope(x_tight, y_tight, coeff)

        # Step 3: Add custom strengthening constraints
        self._add_strengthening_constraints(z, x_tight, y_tight)

        return z

    def _tighten_bounds(self, var: LXVariable) -> LXVariable:
        """
        Custom bound tightening logic.

        Args:
            var: Variable to tighten

        Returns:
            Variable with tightened bounds
        """
        # Implement custom bound tightening
        # This could use constraint propagation, domain reduction, etc.
        ...

    def _add_strengthening_constraints(self, z, x, y):
        """
        Add custom constraints to strengthen McCormick relaxation.

        Args:
            z: Product variable
            x: First variable
            y: Second variable
        """
        # Add custom constraints
        # For example, additional cuts based on problem structure
        ...

Creating Custom Formulations

Example: Custom PWL Formulation Using Convex Combination

from lumix.linearization.techniques.piecewise import LXPiecewiseLinearizer

class LXCustomPWLLinearizer(LXPiecewiseLinearizer):
    """
    Custom piecewise-linear linearizer with specialized formulation.
    """

    def _custom_convex_formulation(
        self,
        var: LXVariable,
        breakpoints: List[float],
        values: List[float]
    ) -> LXVariable:
        """
        Custom convex combination formulation.

        Similar to SOS2 but with additional constraints for specific
        problem structures.

        Args:
            var: Input variable
            breakpoints: Breakpoint x-coordinates
            values: Function values at breakpoints

        Returns:
            Output variable
        """
        n = len(breakpoints)

        # Lambda variables
        lambda_vars = []
        for i in range(n):
            lambda_name = f"lambda_custom_{var.name}_{i}"
            lambda_var = (
                LXVariable[str, float](lambda_name)
                .continuous()
                .bounds(lower=0, upper=1)
                .indexed_by(lambda x: x)
                .from_data([lambda_name])
            )
            lambda_vars.append(lambda_var)
        self.auxiliary_vars.extend(lambda_vars)

        # Output variable
        output_name = f"pwl_custom_{var.name}_{self._aux_counter}"
        self._aux_counter += 1
        output = (
            LXVariable[str, float](output_name)
            .continuous()
            .bounds(lower=min(values), upper=max(values))
            .indexed_by(lambda x: x)
            .from_data([output_name])
        )
        self.auxiliary_vars.append(output)

        # Standard convexity constraint
        convex_expr = LXLinearExpression()
        for lv in lambda_vars:
            convex_expr.add_term(lv, 1.0)
        self.auxiliary_constraints.append(
            LXConstraint(f"custom_convex_{output_name}")
            .expression(convex_expr)
            .eq()
            .rhs(1.0)
        )

        # Custom adjacency constraints (instead of SOS2)
        # Force at most 2 adjacent lambdas to be positive
        for i in range(n - 2):
            # λ[i] + λ[i+1] + λ[i+2] ≤ 1 would be too restrictive
            # Instead, use binary variables to select active segment
            pass  # Implement custom logic

        # x and y definitions
        # ... similar to SOS2 formulation ...

        return output

    def approximate_function(self, func, var, **kwargs):
        """
        Override to use custom formulation.
        """
        # Generate breakpoints
        breakpoints = ...
        values = ...

        # Use custom formulation
        return self._custom_convex_formulation(var, breakpoints, values)

Testing Custom Extensions

Unit Tests

import pytest
from lumix.linearization.config import LXLinearizerConfig

class TestCustomLinearizer:
    @pytest.fixture
    def linearizer(self):
        config = LXLinearizerConfig()
        return LXCustomLinearizer(config)

    def test_creation(self, linearizer):
        """Test linearizer can be created."""
        assert linearizer is not None
        assert linearizer.auxiliary_vars == []

    def test_linearize_custom_term(self, linearizer):
        """Test custom linearization."""
        x = LXVariable[Model, float]("x").bounds(0, 100)
        result = linearizer.linearize_custom_term(x)

        assert result is not None
        assert len(linearizer.auxiliary_vars) > 0

    def test_accuracy(self, linearizer):
        """Test approximation accuracy."""
        # Implement accuracy validation
        pass

Integration Tests

def test_custom_linearizer_in_model():
    """Test custom linearizer in full model."""
    # Build model
    model = build_test_model()

    # Configure with custom linearizer
    config = LXLinearizerConfig()
    # ... configure custom settings ...

    # Linearize
    linearizer = LXLinearizer(model, solver_cap, config)
    linearized = linearizer.linearize_model()

    # Verify
    assert linearized.name.endswith("_linearized")

    # Solve
    solution = optimizer.solve(linearized)
    assert solution.is_optimal()

Best Practices

1. Follow Naming Conventions

   # Good
   class LXYourFeatureLinearizer:
       def linearize_your_feature(self, term):
           aux_name = f"aux_your_feature_{var.name}_{self._aux_counter}"
   
   # Avoid
   class MyCustomThing:
       def do_it(self, x):
           temp = f"tmp_{x}"
   

2. Maintain Auxiliary Element Lists

   # Always append to auxiliary lists
   self.auxiliary_vars.append(new_var)
   self.auxiliary_constraints.append(new_constraint)
   

3. Document Thoroughly

   def linearize_custom(self, term):
       """
       Linearize custom term using XYZ method.
   
       Mathematical Formulation:
           Given: ...
           Creates: ...
           Constraints: ...
   
       Args:
           term: Custom term to linearize
   
       Returns:
           Auxiliary variable representing linearized term
   
       Example:
           >>> linearizer = LXCustomLinearizer(config)
           >>> z = linearizer.linearize_custom(term)
       """
   

4. Handle Edge Cases

   def linearize_term(self, term):
       # Validate inputs
       if term.var.lower_bound is None:
           raise ValueError("Variable must have lower bound")
   
       # Handle zero coefficient
       if abs(term.coefficient) < self.config.tolerance:
           return None  # Skip linearization
   
       # Continue with linearization
       ...
   

5. Add Configuration Options

   @dataclass
   class LXLinearizerConfig:
       # Add settings for your technique
       custom_method: str = "default"
       custom_precision: float = 1e-6
       custom_use_advanced: bool = True
   

Documentation

Document your extension in the appropriate places:

1. Docstrings: Add Google-style docstrings to all classes and methods

2. User Guide: Add usage examples to user guide

3. API Reference: Ensure autodoc picks up your classes

4. Development Guide: Document architecture and design decisions

Example Documentation Structure:

Custom Linearization Technique
===============================

Overview
--------

Description of your technique...

Mathematical Background
-----------------------

Formulation details...

Usage
-----

.. code-block:: python

   from lumix.linearization.techniques import LXCustomLinearizer

   linearizer = LXCustomLinearizer(config)
   result = linearizer.linearize_custom(term)

See Also
--------

- :doc:`/api/linearization/index` - API reference

See Also

• Linearization Architecture - Architecture overview

• Design Decisions - Design rationale

• Linearization Concepts - User guide

• Linearization Module API - API reference