From cf3da70c224855887effcb9b53fa50e29861542c Mon Sep 17 00:00:00 2001
From: odow <o.dowson@gmail.com>
Date: Thu, 6 Mar 2025 15:16:39 +1300
Subject: [PATCH] [Nonlinear.SymbolicAD] add two more simplification cases

---
 src/Nonlinear/SymbolicAD/SymbolicAD.jl |  8 ++++++++
 test/Nonlinear/SymbolicAD.jl           | 19 +++++++++----------
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/src/Nonlinear/SymbolicAD/SymbolicAD.jl b/src/Nonlinear/SymbolicAD/SymbolicAD.jl
index 8a94707400..203da4ad76 100644
--- a/src/Nonlinear/SymbolicAD/SymbolicAD.jl
+++ b/src/Nonlinear/SymbolicAD/SymbolicAD.jl
@@ -77,6 +77,12 @@ function simplify!(f::MOI.ScalarAffineFunction{T}) where {T}
     if isempty(f.terms)
         return f.constant
     end
+    if iszero(f.constant) && length(f.terms) == 1
+        term = only(f.terms)
+        if isone(term.coefficient)
+            return term.variable
+        end
+    end
     return f
 end
 
@@ -1589,6 +1595,8 @@ function _add_to_affine!(
             end
         end
         return _add_to_affine!(ret, something(y, one(T)), convert(T, scale))
+    elseif f.head == :/ && f.args[2] isa Real
+        return _add_to_affine!(ret, f.args[1], convert(T, scale / f.args[2]))
     end
     return  # An unsupported f.head
 end
diff --git a/test/Nonlinear/SymbolicAD.jl b/test/Nonlinear/SymbolicAD.jl
index d94b941623..c6f557f2f1 100644
--- a/test/Nonlinear/SymbolicAD.jl
+++ b/test/Nonlinear/SymbolicAD.jl
@@ -199,22 +199,21 @@ function test_simplify_ScalarAffineFunction()
     x = MOI.VariableIndex(1)
     @test SymbolicAD.simplify(1.0 * x + 1.0) ≈ 1.0 * x + 1.0
     @test SymbolicAD.simplify(1.0 * x + 2.0 * x + 1.0) ≈ 3.0 * x + 1.0
+    @test SymbolicAD.simplify(1.0 * x) == x
+    @test SymbolicAD.simplify(2.0 * x) ≈ 2.0 * x
     return
 end
 
 function test_simplify_ScalarQuadraticFunction()
     x = MOI.VariableIndex(1)
-    f = MOI.ScalarQuadraticFunction(
-        MOI.ScalarQuadraticTerm{Float64}[],
-        [MOI.ScalarAffineTerm{Float64}(1.0, x)],
-        1.0,
-    )
+    terms = MOI.ScalarQuadraticTerm{Float64}[]
+    f = MOI.ScalarQuadraticFunction(terms, [MOI.ScalarAffineTerm(1.0, x)], 1.0)
     @test SymbolicAD.simplify(f) ≈ 1.0 * x + 1.0
-    f = MOI.ScalarQuadraticFunction(
-        MOI.ScalarQuadraticTerm{Float64}[],
-        MOI.ScalarAffineTerm{Float64}[],
-        2.0,
-    )
+    f = MOI.ScalarQuadraticFunction(terms, [MOI.ScalarAffineTerm(2.0, x)], 0.0)
+    @test SymbolicAD.simplify(f) ≈ 2.0 * x + 0.0
+    f = MOI.ScalarQuadraticFunction(terms, [MOI.ScalarAffineTerm(1.0, x)], 0.0)
+    @test SymbolicAD.simplify(f) == x
+    f = MOI.ScalarQuadraticFunction(terms MOI.ScalarAffineTerm{Float64}[], 2.0)
     @test SymbolicAD.simplify(f) === 2.0
     @test SymbolicAD.simplify(1.0 * x * x + 1.0) ≈ 1.0 * x * x + 1.0
     g = 1.0 * x * x + 2.0 * x * x + 1.0