Nonlinear Module Architecture

Deep dive into the nonlinear module’s architecture and design patterns.

Design Philosophy

The nonlinear module follows a declarative, type-safe approach to nonlinear modeling:

1. Separation of Concerns: Term definitions are separate from linearization logic

2. Dataclass Pattern: All terms are immutable dataclasses

3. Metadata Carriers: Terms carry information about structure, not implementation

4. Late Binding: Linearization happens during model building, not at term creation

Architecture Overview

        classDiagram
    class LXAbsoluteTerm {
        +var: LXVariable
        +coefficient: float
    }

    class LXMinMaxTerm {
        +vars: List[LXVariable]
        +operation: Literal["min", "max"]
        +coefficients: List[float]
    }

    class LXBilinearTerm {
        +var1: LXVariable
        +var2: LXVariable
        +coefficient: float
    }

    class LXIndicatorTerm {
        +binary_var: LXVariable
        +condition: bool
        +linear_expr: LXLinearExpression
        +sense: Literal["<=", ">=", "=="]
        +rhs: float
    }

    class LXPiecewiseLinearTerm {
        +var: LXVariable
        +func: Callable
        +num_segments: int
        +x_min: Optional[float]
        +x_max: Optional[float]
        +adaptive: bool
        +method: Literal["sos2", "incremental", "logarithmic"]
    }

    class LXLinearizer {
        +linearize_absolute(term)
        +linearize_minmax(term)
        +linearize_bilinear(term)
        +linearize_indicator(term)
        +linearize_piecewise(term)
    }

    LXLinearizer --> LXAbsoluteTerm
    LXLinearizer --> LXMinMaxTerm
    LXLinearizer --> LXBilinearTerm
    LXLinearizer --> LXIndicatorTerm
    LXLinearizer --> LXPiecewiseLinearTerm
    

Term Design Patterns

Immutable Dataclasses

All terms are frozen dataclasses:

from dataclasses import dataclass

@dataclass
class LXAbsoluteTerm:
    """Immutable representation of |x|."""
    var: LXVariable
    coefficient: float = 1.0

Benefits:

• Thread-safe

• Hashable (can be used in sets/dicts)

• Clear intent (terms are data, not behavior)

• Easy to serialize

Metadata Carriers

Terms carry what to linearize, not how:

# User creates term (metadata)
bilinear = LXBilinearTerm(var1=x, var2=y, coefficient=1.0)

# Linearizer decides how based on variable types
if x.var_type == BINARY and y.var_type == BINARY:
    linearizer.use_and_logic(bilinear)
elif x.var_type == BINARY and y.var_type == CONTINUOUS:
    linearizer.use_big_m(bilinear)
else:
    linearizer.use_mccormick(bilinear)

Type Safety

Full type annotations for compile-time checking:

from typing import Callable, Literal, Optional

@dataclass
class LXPiecewiseLinearTerm:
    var: LXVariable
    func: Callable[[float], float]  # Function signature
    num_segments: int = 20
    method: Literal["sos2", "incremental", "logarithmic"] = "sos2"
    # Literal types enforce valid values

Integration with Linearization

Linearization Workflow

        sequenceDiagram
    participant User
    participant Term
    participant Model
    participant Linearizer
    participant Solver

    User->>Term: Create LXBilinearTerm(x, y)
    User->>Model: Add constraints with term
    Model->>Linearizer: linearize_model()
    Linearizer->>Term: Inspect variable types
    Linearizer->>Linearizer: Select linearization method
    Linearizer->>Model: Add auxiliary variables
    Linearizer->>Model: Add linearization constraints
    Model->>Solver: Build solver-specific model
    

Key Point: Linearization is deferred until model building.

Linearization Methods

Each term type has a dedicated linearization method:

class LXLinearizer:
    def linearize_absolute(self, term: LXAbsoluteTerm) -> ...:
        """Linearize |x| using auxiliary variable."""
        aux_var = self._create_aux_var(f"abs_{term.var.name}")
        self._add_constraint(aux_var >= term.var)
        self._add_constraint(aux_var >= -term.var)
        return aux_var

    def linearize_bilinear(self, term: LXBilinearTerm) -> ...:
        """Linearize x*y based on variable types."""
        if self._is_binary_times_binary(term):
            return self._linearize_binary_and(term)
        elif self._is_binary_times_continuous(term):
            return self._linearize_big_m(term)
        else:
            return self._linearize_mccormick(term)

Module Structure

File Organization

src/lumix/nonlinear/
├── __init__.py      # Module exports and documentation
└── terms.py         # All term dataclass definitions

Design Decision: Single terms.py file keeps related definitions together.

Dependencies

        graph TD
    A[nonlinear] --> B[core.variables]
    A --> C[core.expressions]
    D[linearization] --> A
    E[solvers] --> D

    style A fill:#e1f5ff
    style D fill:#fff4e1
    

Dependency Direction:

• Nonlinear module depends only on core (variables, expressions)

• Linearization module depends on nonlinear (consumes terms)

• Solvers depend on linearization (receives linearized models)

This ensures clean separation of concerns.

Extension Points

Adding New Term Types

To add a new nonlinear term:

1. Define Dataclass in terms.py:

@dataclass
class LXQuadraticTerm:
    """Quadratic term: a*x^2 + b*x + c."""
    var: LXVariable
    a: float
    b: float = 0.0
    c: float = 0.0

2. Export from __init__.py:

from .terms import LXQuadraticTerm

__all__ = [..., "LXQuadraticTerm"]

3. Add Linearization in lumix/linearization/engine.py:

def linearize_quadratic(self, term: LXQuadraticTerm):
    # Linearization logic here
    pass

4. Add Tests:

def test_quadratic_term():
    term = LXQuadraticTerm(var=x, a=2.0, b=1.0, c=0.5)
    assert term.a == 2.0

See Extending Nonlinear Module for detailed guide.

Custom Linearization Methods

Subclass LXLinearizer to customize linearization:

from lumix.linearization import LXLinearizer

class CustomLinearizer(LXLinearizer):
    def linearize_bilinear(self, term: LXBilinearTerm):
        # Use custom McCormick with tighter bounds
        return self._custom_mccormick(term)

Performance Considerations

Memory Footprint

Terms are lightweight dataclasses:

import sys

term = LXAbsoluteTerm(var=x, coefficient=1.0)
print(sys.getsizeof(term))  # ~64 bytes

Implication: Creating thousands of terms is cheap.

Linearization Cost

Term Type	Aux Vars	Constraints	Notes
Absolute	1	2	Very efficient
Min/Max	1	n (inputs)	Scales with inputs
Bilinear (Bin×Bin)	1	3	Exact, efficient
Bilinear (Bin×Cont)	1	4	Exact, Big-M
Bilinear (Cont×Cont)	1	4	Relaxation
Indicator	0	1	Big-M
Piecewise (SOS2)	n+1	n	Best with SOS2 support

Testing Strategy

Unit Tests

Test term creation and properties:

def test_absolute_term_creation():
    var = LXVariable[Product, float]("x").continuous()
    term = LXAbsoluteTerm(var=var, coefficient=2.0)
    assert term.coefficient == 2.0
    assert term.var == var

Integration Tests

Test with linearization engine:

def test_bilinear_linearization():
    x = LXVariable("x").binary()
    y = LXVariable("y").continuous().bounds(0, 10)
    term = LXBilinearTerm(var1=x, var2=y)

    linearizer = LXLinearizer()
    result = linearizer.linearize_bilinear(term)
    assert len(result.constraints) == 4  # Big-M

Type Tests

Use mypy for type checking:

mypy src/lumix/nonlinear

Future Directions

Planned Extensions

• Higher-order terms: Cubic, polynomial

• Multi-variate terms: f(x, y, z) with multiple inputs

• Logical terms: AND, OR, NOT combinations

• Cardinality constraints: AtMost, AtLeast, Exactly

Design Considerations

For new terms:

1. Keep terms as pure data (no methods beyond dataclass)

2. Maintain immutability

3. Full type annotations

4. Document linearization method in docstring

5. Provide usage examples

Next Steps

• Extending Nonlinear Module - How to add new nonlinear terms

• Linearization Architecture - Linearization engine design

• Design Decisions - Overall design philosophy