From 6e1102d4ff56d5dc193c3030dda20c7f3a32496d Mon Sep 17 00:00:00 2001
From: Joshua Pulsipher <33848311+pulsipher@users.noreply.github.com>
Date: Thu, 20 Jun 2024 10:00:19 -0400
Subject: [PATCH] Support Expression Restrictions (#348)

* Support expression restrictions

* fix missing end
---
 docs/src/manual/expression.md |  1 +
 src/expressions.jl            | 64 +++++++++++++++++++++++++++++++++++
 test/expressions.jl           | 28 +++++++++++++++
 3 files changed, 93 insertions(+)

diff --git a/docs/src/manual/expression.md b/docs/src/manual/expression.md
index 0c5f5c860..1ffc24067 100644
--- a/docs/src/manual/expression.md
+++ b/docs/src/manual/expression.md
@@ -54,6 +54,7 @@ add_operators_to_jump
 ## Expression Methods
 ```@docs
 parameter_refs(::Union{JuMP.GenericAffExpr, JuMP.GenericQuadExpr, JuMP.NonlinearExpr})
+restrict(::JuMP.AbstractJuMPScalar)
 map_expression
 map_expression_to_ast
 ```
diff --git a/src/expressions.jl b/src/expressions.jl
index d65df8549..010fbf4c8 100644
--- a/src/expressions.jl
+++ b/src/expressions.jl
@@ -812,6 +812,70 @@ function map_expression_to_ast(
     return ex
 end
 
+################################################################################
+#                            EXPRESSION RESTICTIONS
+################################################################################
+"""
+    restrict(expr::JuMP.AbstractJuMPScalar, supps...)::JuMP.AbstractJuMPScalar
+
+Restrict an infinite expression `expr` to be enforced over infinite parameter 
+supports `supps`. This is limited to expressions only contain infinite variables 
+with the same kind of infinite parameter dependencies. Note that more conveniently 
+the expression can be treated as a function for the syntax `expr(supps...)`. 
+
+**Example**
+```julia-repl
+julia> ex = @expression(model, 3y - 2)
+3 y(t) - 2
+
+julia> restrict(ex, 0)
+3 y(0) - 2
+
+julia> ex(0)
+3 y(0) - 2
+```
+"""
+function restrict(expr::JuMP.AbstractJuMPScalar, supps...)
+    # check to make sure all variables depend on the same infinite parameters
+    pref_tuples = Set()
+    _interrogate_variables(expr) do v
+        if v.index_type <: InfiniteParameterIndex
+            error("Restrictions on expressions with infinite parameters are not supported.")
+        elseif !(v.index_type <: Union{FiniteVariableIndex, PointVariableIndex, FiniteParameterIndex})
+            push!(pref_tuples, raw_parameter_refs(v))
+        end
+    end
+    if length(pref_tuples) > 1
+        error("Unable to restrict expression `$expr` with supports `$supps`. " *
+              "Not all the variables use the same infinite parameters in the " *
+              "same format.")
+    end
+    # restrict the expression using supps and return
+    return map_expression(expr) do v
+        if isempty(_object_numbers(v))
+            return v
+        else
+            return restrict(v, supps...)
+        end
+    end
+end
+
+## Make expressions callable for restrictions
+# AffExprs
+function (aff::JuMP.GenericAffExpr{Float64, GeneralVariableRef})(supps...) 
+    return restrict(aff, supps...)
+end
+
+# QuadExprs
+function (quad::JuMP.GenericQuadExpr{Float64, GeneralVariableRef})(supps...) 
+    return restrict(quad, supps...)
+end
+
+# NonlinearExprs
+function (nl::JuMP.GenericNonlinearExpr{GeneralVariableRef})(supps...) 
+    return restrict(nl, supps...)
+end
+
 ################################################################################
 #                             COEFFICIENT METHODS
 ################################################################################
diff --git a/test/expressions.jl b/test/expressions.jl
index d7bc97f6f..719257bd5 100644
--- a/test/expressions.jl
+++ b/test/expressions.jl
@@ -702,6 +702,34 @@ end
     end
 end
 
+# Test restrictions
+@testset "restrict" begin 
+    # setup model
+    m = InfiniteModel()
+    @infinite_parameter(m, t in [0, 1])
+    @infinite_parameter(m, x in [-1, 1])
+    @variable(m, y, Infinite(t, x))
+    @variable(m, q, Infinite(x, t))
+    @variable(m, w, Infinite(t))
+    @variable(m, z)
+    # test AffExpr
+    @testset "AffExpr" begin
+        @test (2z + y + 42)(0, -1) == 2z + y(0, -1) + 42
+        @test (2z + y + 42)(0, x) == 2z + y(0, x) + 42
+        @test_throws ErrorException (y + q)(0, 0)
+        @test_throws ErrorException (y + t)(0, x)
+    end
+    # test QuadExpr 
+    @testset "QuadExpr" begin
+        @test (z^2 + 3y - 2)(0, -1) == z^2 + 3y(0, -1) - 2
+        @test (w^2 - 2)(0) == w(0)^2 - 2
+    end
+    # test GenericNonlinearExpr
+    @testset "GenericNonlinearExpr" begin
+        @test isequal_canonical((sin(y) * z)(t, -1), sin(y(t, -1)) * z)
+    end
+end
+
 # Test _set_variable_coefficient!
 @testset "_set_variable_coefficient!" begin
     # initialize model and references