def test_equal_nodes():
assert equal_nodes(None, None)
x = ref('x')
y = ref('y')
z = ref('z')
ix = ampl.Indexing(x)
iy = ampl.Indexing(y)
check_equal_nodes(x, y)
check_equal_nodes(ampl.ParenExpr(x), ampl.UnaryExpr('-', x))
check_equal_nodes(ampl.UnaryExpr('-', x), ampl.UnaryExpr('+', x), ampl.UnaryExpr('-', y))
check_equal_nodes(ampl.BinaryExpr('*', x, y), ampl.BinaryExpr('+', x, y),
ampl.BinaryExpr('*', y, y), ampl.BinaryExpr('*', x, x))
check_equal_nodes(ampl.IfExpr(x, y, z), ampl.IfExpr(z, y, z),
ampl.IfExpr(x, x, z), ampl.IfExpr(x, y, y))
check_equal_nodes(ampl.CallExpr('sin', [x]), ampl.CallExpr('cos', [x]),
ampl.CallExpr('sin', [y]), ampl.CallExpr('sin', [x, y]))
check_equal_nodes(ampl.SumExpr(ix, x), ampl.SumExpr(iy, x), ampl.SumExpr(ix, y))
check_equal_nodes(ix, iy)
check_equal_nodes(ampl.Decl('var', 'x'), ampl.Decl('var', 'y'))
check_equal_nodes(ampl.IncludeStmt('data'), ampl.IncludeStmt('model'))
check_equal_nodes(ampl.DataStmt('param', 'S', ['a', 'b'], ['1', '2', '3']),
ampl.DataStmt('var', 'S', ['a', 'b'], ['1', '2', '3']),
ampl.DataStmt('param', 'T', ['a', 'b'], ['1', '2', '3']),
ampl.DataStmt('param', 'S', ['a'], ['1', '2', '3']),
ampl.DataStmt('param', 'S', ['a', 'c'], ['1', '2', '3']),
ampl.DataStmt('param', 'S', ['a', 'v'], ['1', '2']),
ampl.DataStmt('param', 'S', ['a', 'v'], ['1', '2', '4']))
assert not equal_nodes(ampl.Decl('var', 'x'), ampl.CompoundStmt([]))
def test_parse_associativity():
x = ref('x')
y = ref('y')
z = ref('z')
check_parse_lexpr('x + y + z', ampl.BinaryExpr('+', ampl.BinaryExpr('+', x, y), z))
check_parse_lexpr('x + y - z', ampl.BinaryExpr('-', ampl.BinaryExpr('+', x, y), z))
check_parse_lexpr('x * y * z', ampl.BinaryExpr('*', ampl.BinaryExpr('*', x, y), z))
check_parse_lexpr('x * y / z', ampl.BinaryExpr('/', ampl.BinaryExpr('*', x, y), z))
check_parse_lexpr('x ^ y ^ z', ampl.BinaryExpr('^', x, ampl.BinaryExpr('^', y, z)))
def merge_models(models):
"""
Merge given AMPL models into a single one using product composition
of objective functions.
For example, two models
minimize o: f1(x);
and
minimize o: f2(x);
are combined into a single model
minimize o: f1(x1) * f2(x2);
"""
merged_head = []
merged_tail = []
merged_best_obj = 1
merged_obj = ampl.Decl('minimize', 'f')
for i in range(len(models)):
head, obj, tail, best_obj = prepare_for_merge(models[i], i + 1)
merged_head += head
merged_tail += tail
merged_best_obj *= best_obj
if merged_obj.body:
merged_obj.body = ampl.BinaryExpr('*', merged_obj.body, obj.body)
else:
merged_obj.body = obj.body
# Invert sign if objectives are of different kinds.
return ampl.CompoundStmt(merged_head + [merged_obj] +
merged_tail), merged_best_obj
def prepare_for_merge(model, suffix):
suffix = str(suffix)
path = model['path']
with open(os.path.join(repo_dir, path), 'r') as f:
nodes = ampl.parse(f.read(), path).nodes
# Rename declarations.
names = {}
visitor = RenamingVisitor(names)
for n in nodes:
if isinstance(n, ampl.Decl):
new_name = n.name + suffix
names[n.name] = new_name
n.name = new_name
if n.indexing:
n.indexing.accept(visitor)
if n.body:
n.body.accept(visitor)
if isinstance(n, ampl.DataStmt):
n.set_name = n.set_name + suffix
n.param_names = [name + suffix for name in n.param_names]
# Find the first objective and partition the nodes around it.
obj_index = find_obj(nodes)
# Add objective offset to make the optimal value nonnegative.
obj = nodes[obj_index]
sign = 1 if obj.kind == 'minimize' else -1
best_obj = sign * model['best_obj']
offset = math.ceil(abs(min(0.0, best_obj)))
obj.body = ampl.ParenExpr(obj.body)
if obj.kind == 'maximize':
obj.body = ampl.UnaryExpr('-', obj.body)
if offset > 0:
obj.body = ampl.ParenExpr(
ampl.BinaryExpr('+', obj.body, ampl.Reference(str(offset))))
return nodes[:obj_index], obj, nodes[obj_index + 1:], best_obj + offset
def test_pretty_print():
check_print('a', ref('a'))
check_print('a[b]', ampl.SubscriptExpr('a', ref('b')))
check_print('(a)', ampl.ParenExpr(ref('a')))
check_print('-a', ampl.UnaryExpr('-', ref('a')))
check_print('a + b', ampl.BinaryExpr('+', ref('a'), ref('b')))
check_print('if a then b else c',
ampl.IfExpr(ref('a'), ref('b'), ref('c')))
check_print('if a then b',
ampl.IfExpr(ref('a'), ref('b'), None))
check_print('f(a, b)', ampl.CallExpr('f', [ref('a'), ref('b')]))
check_print('sum{s in S} x[s]', ampl.SumExpr(ampl.Indexing(ref('S'), 's'),
ampl.SubscriptExpr('x', ref('s'))))
check_print('{s in S}', ampl.Indexing(ref('S'), 's'))
check_print('{S}', ampl.Indexing(ref('S')))
check_print('= a', ampl.InitAttr(ref('a')))
check_print('in [a, b]', ampl.InAttr(ref('a'), ref('b')))
check_print('var x{S} = a;\n', ampl.Decl('var', 'x', ampl.Indexing(ref('S')),
[ampl.InitAttr(ref('a'))]))
decl = ampl.Decl('minimize', 'o')
decl.body = ampl.UnaryExpr('-', ref('x'))
check_print('minimize o: -x;\n', decl)
check_print('model;\n', ampl.IncludeStmt('model'))
param_names = ['a', 'b']
values = [str(n) for n in range(6)]
check_print(
'param:\n' +
'S:a b :=\n' +
'0 1 2\n' +
'3 4 5\n' +
';\n', ampl.DataStmt('param', 'S', param_names, values))
check_print('model;\nvar x;\n',
ampl.CompoundStmt([ampl.IncludeStmt('model'), ampl.Decl('var', 'x')]))
def test_parse_expr():
check_parse_expr('42', ref('42'))
check_parse_expr('-x', ampl.UnaryExpr('-', ref('x')))
check_parse_expr('(x)', ampl.ParenExpr(ref('x')))
check_parse_expr('sum{S} x', ampl.SumExpr(ampl.Indexing(ref('S')), ref('x')))
check_parse_expr('sum{s in S} x', ampl.SumExpr(ampl.Indexing(ref('S'), 's'), ref('x')))
check_parse_expr('if x then y', ampl.IfExpr(ref('x'), ref('y')))
check_parse_expr('if x then y else z', ampl.IfExpr(ref('x'), ref('y'), ref('z')))
for f in ['abs', 'acos', 'acosh', 'alias', 'asin', 'asinh', 'atan', 'atan2', 'atanh',
'ceil', 'ctime', 'cos', 'exp', 'floor', 'log', 'log10', 'max', 'min',
'precision', 'round', 'sin', 'sinh', 'sqrt', 'tan', 'tanh', 'time', 'trunc']:
check_parse_expr(f + '(x)', ampl.CallExpr(f, [ref('x')]))
check_parse_expr('x[y]', ampl.SubscriptExpr('x', ref('y')))
for op in ['^', '**', '+', '-', '*', '/']:
check_parse_expr('x ' + op + ' y', ampl.BinaryExpr(op, ref('x'), ref('y')))
for op in ['||', 'or', 'in', '<', '<=', '=', '==', '<>', '!=', '<=', '>']:
check_parse_lexpr('x ' + op + ' y', ampl.BinaryExpr(op, ref('x'), ref('y')))
def test_parse_attrs():
e1 = ampl.BinaryExpr('+', ref('x'), ref('y'))
e2 = ampl.UnaryExpr('-', ref('x'))
for kw in ['param', 'var']:
check_parse(kw + ' a = x + y;',
ampl.Decl(kw, 'a', None, [ampl.InitAttr(e1)]))
check_parse(kw + ' a in [x + y, -x];',
ampl.Decl(kw, 'a', None, [ampl.InAttr(e1, e2)]))
def test_binary():
lhs = ref('a')
rhs = ref('b')
expr = ampl.BinaryExpr('+', lhs, rhs)
assert type(expr) == ampl.BinaryExpr
assert expr.op == '+'
assert expr.lhs == lhs
assert expr.rhs == rhs
check_accept(expr, 'visit_binary')
def test_parse_expr_precedence():
x = ref('x')
y = ref('y')
z = ref('z')
check_parse_lexpr('(x || y)', ampl.ParenExpr(ampl.BinaryExpr('||', x, y)))
idx = ampl.Indexing(ref('S'))
check_parse_lexpr('sum{S} x * y', ampl.SumExpr(idx, ampl.BinaryExpr('*', x, y)))
check_parse_lexpr('sum{S} x + y', ampl.BinaryExpr('+', ampl.SumExpr(idx, x), y))
check_parse_expr('if x || y then z', ampl.IfExpr(ampl.BinaryExpr('||', x, y), z))
check_parse_expr('if x then y & z', ampl.IfExpr(x, ampl.BinaryExpr('&', y, z)))
check_parse_lexpr('if x then y in S', ampl.BinaryExpr('in', ampl.IfExpr(x, y), ref('S')))
check_parse_expr('if x then y else x & z', ampl.IfExpr(x, y, ampl.BinaryExpr('&', x, z)))
check_parse_lexpr('if x then y else z in S', ampl.BinaryExpr('in', ampl.IfExpr(x, y, z), ref('S')))
check_parse_expr('sin(if x then y)', ampl.CallExpr('sin', [ampl.IfExpr(x, y)]))
check_parse_expr('a[if x then y]', ampl.SubscriptExpr('a', ampl.IfExpr(x, y)))
check_parse_lexpr('x || y > z', ampl.BinaryExpr('||', x, ampl.BinaryExpr('>', y, z)))
check_parse_lexpr('x || y in z', ampl.BinaryExpr('||', x, ampl.BinaryExpr('in', y, z)))
def test_parse_obj_body():
decl = ampl.Decl('minimize', 'o')
decl.body = ampl.BinaryExpr('+', ref('x'), ref('y'))
check_parse('minimize o: x + y;', decl)