Skip to content

Commit

Permalink
Support expression restrictions
Browse files Browse the repository at this point in the history
  • Loading branch information
pulsipher committed Jun 19, 2024
1 parent 20c1162 commit 87b0506
Show file tree
Hide file tree
Showing 3 changed files with 92 additions and 0 deletions.
1 change: 1 addition & 0 deletions docs/src/manual/expression.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
```
Expand Down
64 changes: 64 additions & 0 deletions src/expressions.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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 warning on line 838 in src/expressions.jl

View check run for this annotation

Codecov / codecov/patch

src/expressions.jl#L838

Added line #L838 was not covered by tests
# 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))

Check warning on line 845 in src/expressions.jl

View check run for this annotation

Codecov / codecov/patch

src/expressions.jl#L840-L845

Added lines #L840 - L845 were not covered by tests
end
end
if length(pref_tuples) > 1
error("Unable to restrict expression `$expr` with supports `$supps`. " *

Check warning on line 849 in src/expressions.jl

View check run for this annotation

Codecov / codecov/patch

src/expressions.jl#L848-L849

Added lines #L848 - L849 were not covered by tests
"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

Check warning on line 856 in src/expressions.jl

View check run for this annotation

Codecov / codecov/patch

src/expressions.jl#L854-L856

Added lines #L854 - L856 were not covered by tests
else
return restrict(v, supps...)

Check warning on line 858 in src/expressions.jl

View check run for this annotation

Codecov / codecov/patch

src/expressions.jl#L858

Added line #L858 was not covered by tests
end
end
end

## Make expressions callable for restrictions
# AffExprs
function (aff::JuMP.GenericAffExpr{Float64, GeneralVariableRef})(supps...)
return restrict(aff, supps...)

Check warning on line 866 in src/expressions.jl

View check run for this annotation

Codecov / codecov/patch

src/expressions.jl#L865-L866

Added lines #L865 - L866 were not covered by tests
end

# QuadExprs
function (quad::JuMP.GenericQuadExpr{Float64, GeneralVariableRef})(supps...)
return restrict(quad, supps...)

Check warning on line 871 in src/expressions.jl

View check run for this annotation

Codecov / codecov/patch

src/expressions.jl#L870-L871

Added lines #L870 - L871 were not covered by tests
end

# NonlinearExprs
function (nl::JuMP.GenericNonlinearExpr{GeneralVariableRef})(supps...)
return restrict(nl, supps...)

Check warning on line 876 in src/expressions.jl

View check run for this annotation

Codecov / codecov/patch

src/expressions.jl#L875-L876

Added lines #L875 - L876 were not covered by tests
end

################################################################################
# COEFFICIENT METHODS
################################################################################
Expand Down
27 changes: 27 additions & 0 deletions test/expressions.jl
Original file line number Diff line number Diff line change
Expand Up @@ -702,6 +702,33 @@ 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
# 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
Expand Down

0 comments on commit 87b0506

Please sign in to comment.