From 744b043461e67bb2b1ea8455afa70f56cd168087 Mon Sep 17 00:00:00 2001
From: "Alinson S. Xavier" <git@axavier.org>
Date: Thu, 18 Sep 2025 12:08:58 -0500
Subject: [PATCH] Implement _add_pwl_constraints

---
 src/model/jumpext.jl           |  75 ++++++++++++++++-
 test/src/RELOGT.jl             |   2 +
 test/src/model/jumpext_test.jl | 144 +++++++++++++++++++++++++++++++++
 3 files changed, 220 insertions(+), 1 deletion(-)
 create mode 100644 test/src/model/jumpext_test.jl

diff --git a/src/model/jumpext.jl b/src/model/jumpext.jl
index f69c93f..43150a2 100644
--- a/src/model/jumpext.jl
+++ b/src/model/jumpext.jl
@@ -14,7 +14,7 @@ function fix(x::Float64, v::Float64; force)
     return abs(x - v) < 1e-6 || error("Value mismatch: $x != $v")
 end
 
-function set_name(x::Number, n::String)
+function set_name(::Number, ::String)
     # nop
 end
 
@@ -45,3 +45,76 @@ function _set_names!(dict::Dict)
         end
     end
 end
+
+
+"""
+    _add_pwl_constraints(model, xvar, yvars, xpts, ypts)
+
+Add piecewise-linear constraints to a JuMP model for multiple y variables.
+
+Creates constraints y_i = f_i(x) where each f_i is a piecewise-linear function 
+defined by the breakpoints (xpts, ypts[:, i]).
+
+# Arguments
+- `model`: JuMP model
+- `xvar`: The x variable (JuMP variable)
+- `yvars`: Vector of y variables (JuMP variables)  
+- `xpts`: Vector of x values for breakpoints (must be in non-decreasing order)
+- `ypts`: Matrix of y values where ypts[i, j] is the y value for the j-th variable 
+         at the i-th breakpoint
+
+# Example
+```julia
+@variable(model, y1)
+@variable(model, y2)
+ypts_matrix = [1.5 2.0; 0.0 1.5; 3.0 0.5]  # 3 breakpoints, 2 y variables
+_add_pwl_constraints(model, x, [y1, y2], [0.0, 1.0, 2.0], ypts_matrix, name="multiPWL")
+```
+"""
+function _add_pwl_constraints(model, xvar, yvars, xpts, ypts)
+    # Input validation
+    ypts isa AbstractMatrix || throw(ArgumentError("ypts must be a matrix"))
+    length(xpts) == size(ypts, 1) ||
+        throw(ArgumentError("xpts length must match number of rows in ypts"))
+    length(yvars) == size(ypts, 2) ||
+        throw(ArgumentError("Number of y variables must match number of columns in ypts"))
+    length(xpts) >= 1 || throw(ArgumentError("At least one breakpoint is required"))
+
+    # Check that xpts is increasing
+    for i = 2:length(xpts)
+        xpts[i] > xpts[i-1] || throw(ArgumentError("xpts must be in increasing order"))
+    end
+
+    n_points = length(xpts)
+    n_yvars = length(yvars)
+
+    if n_points == 1
+        # Single point case: y_j = ypts[1,j], x = xpts[1]
+        @constraint(model, xvar == xpts[1])
+        for j = 1:n_yvars
+            @constraint(model, yvars[j] == ypts[1, j])
+        end
+
+    elseif n_points == 2
+        # Two points case: single linear segment for each y variable
+        x1, x2 = xpts[1], xpts[2]
+
+        # Linear relationship for each y variable: y_j = y1_j + slope_j * (x-x1)
+        for j = 1:n_yvars
+            y1, y2 = ypts[1, j], ypts[2, j]
+            slope = (y2 - y1) / (x2 - x1)
+            @constraint(model, yvars[j] == y1 + slope * (xvar - x1))
+        end
+    else
+        # Multiple segments case (3+ points): use SOS2 formulation
+        λ = @variable(model, [1:n_points], lower_bound = 0, upper_bound = 1)
+        @constraint(model, λ in SOS2())
+        @constraint(model, sum(λ) == 1)
+        @constraint(model, xvar == sum(xpts[i] * λ[i] for i = 1:n_points))
+        for j = 1:n_yvars
+            @constraint(model, yvars[j] == sum(ypts[i, j] * λ[i] for i = 1:n_points))
+        end
+    end
+
+    return
+end
diff --git a/test/src/RELOGT.jl b/test/src/RELOGT.jl
index f1ac80b..c4e241a 100644
--- a/test/src/RELOGT.jl
+++ b/test/src/RELOGT.jl
@@ -7,6 +7,7 @@ using JuliaFormatter
 include("instance/parse_test.jl")
 include("model/build_test.jl")
 include("model/dist_test.jl")
+include("model/jumpext_test.jl")
 include("reports_test.jl")
 include("../fixtures/boat_example.jl")
 
@@ -23,6 +24,7 @@ function runtests()
         model_build_test()
         model_dist_test()
         report_tests()
+        jumpext_test()
     end
     return
 end
diff --git a/test/src/model/jumpext_test.jl b/test/src/model/jumpext_test.jl
new file mode 100644
index 0000000..5f085c3
--- /dev/null
+++ b/test/src/model/jumpext_test.jl
@@ -0,0 +1,144 @@
+# RELOG: Reverse Logistics Optimization
+# Copyright (C) 2020, UChicago Argonne, LLC. All rights reserved.
+# Released under the modified BSD license. See COPYING.md for more details.
+
+using RELOG
+using JuMP
+using HiGHS
+using Test
+
+function jumpext_test()
+    jumpext_pwl_single_point()
+    jumpext_pwl_two_points()
+    jumpext_pwl_multiple_points()
+    jumpext_pwl_input_validation()
+    return
+end
+
+function jumpext_pwl_single_point()
+    model = Model(HiGHS.Optimizer)
+    set_silent(model)
+    @variable(model, x)
+    @variable(model, y1)
+    @variable(model, y2)
+    xpts = [5.0]
+    ypts = [10.0 20.0]
+    RELOG._add_pwl_constraints(model, x, [y1, y2], xpts, ypts)
+    optimize!(model)
+    @test is_solved_and_feasible(model)
+    @test value(x) ≈ 5.0 atol = 1e-6
+    @test value(y1) ≈ 10.0 atol = 1e-6
+    @test value(y2) ≈ 20.0 atol = 1e-6
+    return
+end
+
+function jumpext_pwl_two_points()
+    model = Model(HiGHS.Optimizer)
+    set_silent(model)
+    @variable(model, x)
+    @variable(model, y1)
+    @variable(model, y2)
+    xpts = [0.0, 2.0]
+    ypts = [0.0 10.0; 4.0 6.0]
+    RELOG._add_pwl_constraints(model, x, [y1, y2], xpts, ypts)
+
+    # Test at x = 1
+    JuMP.fix(x, 1.0)
+    optimize!(model)
+    @test is_solved_and_feasible(model)
+    @test value(y1) ≈ 2.0 atol = 1e-6
+    @test value(y2) ≈ 8.0 atol = 1e-6
+
+    # Test at x = 2
+    JuMP.fix(x, 2.0)
+    optimize!(model)
+    @test is_solved_and_feasible(model)
+    @test value(y1) ≈ 4.0 atol = 1e-6
+    @test value(y2) ≈ 6.0 atol = 1e-6
+    return
+end
+
+function jumpext_pwl_multiple_points()
+    model = Model(HiGHS.Optimizer)
+    set_silent(model)
+    @variable(model, x)
+    @variable(model, y1)
+    @variable(model, y2)
+    xpts = [0.0, 1.0, 2.0]
+    ypts = [0.0 5.0; 2.0 3.0; 1.0 4.0]
+    RELOG._add_pwl_constraints(model, x, [y1, y2], xpts, ypts)
+
+    # Test at x = 0.5
+    JuMP.fix(x, 0.5)
+    optimize!(model)
+    @test is_solved_and_feasible(model)
+    @test value(y1) ≈ 1.0 atol = 1e-6
+    @test value(y2) ≈ 4.0 atol = 1e-6
+
+    # Test at x = 1
+    JuMP.fix(x, 1.0)
+    optimize!(model)
+    @test is_solved_and_feasible(model)
+    @test value(y1) ≈ 2.0 atol = 1e-6
+    @test value(y2) ≈ 3.0 atol = 1e-6
+
+    # Test at x = 1.5
+    JuMP.fix(x, 1.5)
+    optimize!(model)
+    @test is_solved_and_feasible(model)
+    @test value(y1) ≈ 1.5 atol = 1e-6
+    @test value(y2) ≈ 3.5 atol = 1e-6
+    return
+end
+
+function jumpext_pwl_input_validation()
+    model = Model(HiGHS.Optimizer)
+    @variable(model, x)
+    @variable(model, y)
+
+    # Test non-matrix ypts
+    @test_throws ArgumentError RELOG._add_pwl_constraints(model, x, [y], [1.0], [1.0])
+
+    # Test mismatched dimensions
+    @test_throws ArgumentError RELOG._add_pwl_constraints(
+        model,
+        x,
+        [y],
+        [1.0, 2.0],
+        [1.0 2.0],
+    )
+    @test_throws ArgumentError RELOG._add_pwl_constraints(
+        model,
+        x,
+        [y],
+        [1.0],
+        [1.0 2.0; 3.0 4.0],
+    )
+
+    # Test empty breakpoints
+    @test_throws ArgumentError RELOG._add_pwl_constraints(
+        model,
+        x,
+        [y],
+        Float64[],
+        Matrix{Float64}(undef, 0, 1),
+    )
+
+    # Test non-increasing x points
+    @test_throws ArgumentError RELOG._add_pwl_constraints(
+        model,
+        x,
+        [y],
+        [2.0, 1.0],
+        [1.0; 2.0],
+    )
+    @test_throws ArgumentError RELOG._add_pwl_constraints(
+        model,
+        x,
+        [y],
+        [1.0, 1.0],
+        [1.0; 2.0],
+    )
+
+    return
+end