# ___________________________________________________________________________
#
# Pyomo: Python Optimization Modeling Objects
# Copyright 2017 National Technology and Engineering Solutions of Sandia, LLC
# Under the terms of Contract DE-NA0003525 with National Technology and
# Engineering Solutions of Sandia, LLC, the U.S. Government retains certain
# rights in this software.
# This software is distributed under the 3-clause BSD License.
# ___________________________________________________________________________
import pyutilib.th as unittest
from pyomo.core.base.indexed_component import normalize_index
from pyomo.core.base.indexed_component_slice import IndexedComponent_slice
from pyomo.core.base.global_set import UnindexedComponent_set
import pyomo.environ as pyo
import pyomo.dae as dae
from pyomo.common.collections import ComponentSet, ComponentMap
from pyomo.util.slices import (
get_component_call_stack,
slice_component_along_sets,
replace_indices,
get_location_set_map,
)
class TestGetComponentCallStack(unittest.TestCase):
get_attribute = IndexedComponent_slice.get_attribute
get_item = IndexedComponent_slice.get_item
def assertSameStack(self, stack1, stack2):
for (call1, arg1), (call2, arg2) in zip(stack1, stack2):
self.assertIs(call1, call2)
self.assertEqual(arg1, arg2)
def assertCorrectStack(self, comp, pred_stack, context=None):
act_stack = get_component_call_stack(comp, context=None)
self.assertSameStack(pred_stack, act_stack)
def model(self):
m = pyo.ConcreteModel()
m.s1 = pyo.Set(initialize=[1,2,3])
m.s2 = pyo.Set(initialize=[('a',1),('b',2)])
m.v0 = pyo.Var()
m.v1 = pyo.Var(m.s1)
m.v2 = pyo.Var(m.s2)
@m.Block()
def b(b):
@b.Block(m.s1)
def b1(b):
b.v0 = pyo.Var()
b.v1 = pyo.Var(m.s1)
@b.Block()
def b(b):
b.v0 = pyo.Var()
b.v2 = pyo.Var(m.s2)
@b.Block(m.s1, m.s2)
def b2(b,i1,i2,i3=None):
# Optional i3 is to allow for the case
# when normalize_index.flatten is False
b.v0 = pyo.Var()
b.v1 = pyo.Var(m.s1)
@b.Block()
def b(b):
b.v0 = pyo.Var()
b.v2 = pyo.Var(m.s2)
return m
def test_no_context(self):
m = self.model()
comp = m.v1[1]
pred_stack = [
(self.get_item, 1),
(self.get_attribute, 'v1'),
]
self.assertCorrectStack(comp, pred_stack)
comp = m.v1
pred_stack = [
(self.get_attribute, 'v1'),
]
self.assertCorrectStack(comp, pred_stack)
comp = m.b.b1[1].b.b2
pred_stack = [
(self.get_attribute, 'b2'),
(self.get_attribute, 'b'),
(self.get_item, 1),
(self.get_attribute, 'b1'),
(self.get_attribute, 'b'),
]
self.assertCorrectStack(comp, pred_stack)
comp = m.b.b1[1].b.b2[1,'a',1]
pred_stack = [
(self.get_item, (1,'a',1)),
(self.get_attribute, 'b2'),
(self.get_attribute, 'b'),
(self.get_item, 1),
(self.get_attribute, 'b1'),
(self.get_attribute, 'b'),
]
self.assertCorrectStack(comp, pred_stack)
normalize_index.flatten = False
comp = m.b.b1[1].b.b2[1,('a',1)]
# NOTE: This fails without the parenthesis around ('a',1)
# as comp.index() is None. The actual stack just starts one
# layer higher...
# It is confusing that comp can even be generated in this case;
# shouldn't (1,'a',1) be an invalid index?
pred_stack = [
(self.get_item, (1,('a',1))),
(self.get_attribute, 'b2'),
(self.get_attribute, 'b'),
(self.get_item, 1),
(self.get_attribute, 'b1'),
(self.get_attribute, 'b'),
]
self.assertCorrectStack(comp, pred_stack)
normalize_index.flatten = True
def test_from_block(self):
m = self.model()
comp = m.v0
pred_stack = [
(self.get_attribute, 'v0'),
]
self.assertCorrectStack(comp, pred_stack, context=m)
comp = m.b.b1[2].b.b2[1,'a',1]
pred_stack = [
(self.get_item, (1,'a',1)),
(self.get_attribute, 'b2'),
(self.get_attribute, 'b'),
(self.get_item, 2),
(self.get_attribute, 'b1'),
(self.get_attribute, 'b'),
]
self.assertCorrectStack(comp, pred_stack, context=m)
comp = m.b.b1[2].b.b2[1,'a',1].b.v2['b',2]
pred_stack = [
(self.get_item, ('b',2)),
(self.get_attribute, 'v2'),
(self.get_attribute, 'b'),
(self.get_item, (1,'a',1)),
(self.get_attribute, 'b2'),
(self.get_attribute, 'b'),
(self.get_item, 2),
]
self.assertCorrectStack(comp, pred_stack, context=m.b.b1)
comp = m.b.b1[2].b.b2
pred_stack = [
(self.get_attribute, 'b2'),
(self.get_attribute, 'b'),
(self.get_item, 2),
]
self.assertCorrectStack(comp, pred_stack, context=m.b.b1)
comp = m.b.b1[2]
pred_stack = [
(self.get_item, 2),
]
self.assertCorrectStack(comp, pred_stack, context=m.b.b1)
comp = m.b.b1
act_stack = get_component_call_stack(comp, context=m.b.b1)
self.assertEqual(len(act_stack), 0)
def test_from_blockdata(self):
m = self.model()
context = m.b.b1[3].b.b2[2,'b',2]
comp = m.b.b1[3].b.b2[2,'b',2].b
pred_stack = [
(self.get_attribute, 'b'),
]
self.assertCorrectStack(comp, pred_stack, context=context)
comp = m.b.b1[3].b.b2[2,'b',2].b.v2['a',1]
pred_stack = [
(self.get_item, ('a',1)),
(self.get_attribute, 'v2'),
(self.get_attribute, 'b'),
]
self.assertCorrectStack(comp, pred_stack, context=context)
context = m.b.b1[3]
comp = m.b.b1[3]
act_stack = get_component_call_stack(comp, context=context)
self.assertEqual(len(act_stack), 0)
class TestGetLocationAndReplacement(unittest.TestCase):
def model(self):
m = pyo.ConcreteModel()
m.time = pyo.Set(initialize=[1,2,3])
m.space = dae.ContinuousSet(initialize=[0,2])
m.comp = pyo.Set(initialize=['a','b'])
m.d_2 = pyo.Set(initialize=[('a',1),('b',2)])
m.d_none = pyo.Set(initialize=[('c',1,10), ('d',3)], dimen=None)
m.b0 = pyo.Block()
m.b1 = pyo.Block(m.time)
m.b2 = pyo.Block(m.time, m.space)
m.b0d2 = pyo.Block(m.d_2)
m.b1d2 = pyo.Block(m.time, m.d_2)
m.b2d2 = pyo.Block(m.time, m.d_2, m.space)
m.b3d2 = pyo.Block(m.time, m.d_2, m.space, m.time)
m.b0dn = pyo.Block(m.d_none)
m.b1dn = pyo.Block(m.time, m.d_none)
m.b1dnd2 = pyo.Block(m.time, m.d_none, m.d_2)
m.dnd2b1 = pyo.Block(m.d_none, m.d_2, m.time)
m.b2dn = pyo.Block(m.time, m.space, m.d_none)
m.b2dnd2 = pyo.Block(m.time, m.d_none, m.d_2, m.space)
m.b3dn = pyo.Block(m.time, m.d_2, m.d_none, m.space, m.d_2)
m.dn2 = pyo.Block(m.time, m.d_none, m.d_2, m.d_none, m.d_2)
return m
def assertSameMap(self, pred, act):
self.assertEqual(len(pred), len(act))
for k, v in pred.items():
self.assertIn(k, act)
self.assertIs(pred[k], act[k])
def test_one_to_one(self):
m = self.model()
index_set = m.b0.index_set()
index = None
pred_map = {0: UnindexedComponent_set}
location_set_map = get_location_set_map(index, index_set)
self.assertEqual(pred_map, location_set_map)
index_set = m.b1.index_set()
index = 1
pred_map = {
0: m.time,
}
location_set_map = get_location_set_map(index, index_set)
self.assertSameMap(pred_map, location_set_map)
index_set = m.b2.index_set()
index = (1,2)
pred_map = {
0: m.time,
1: m.space,
}
location_set_map = get_location_set_map(index, index_set)
self.assertSameMap(pred_map, location_set_map)
def test_1d_replacement(self):
m = self.model()
# Unindexed
index_set = m.b0.index_set()
index = None
location_set_map = get_location_set_map(index, index_set)
sets = ComponentSet()
pred_index = (None,)
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(pred_index, new_index)
# One one-d set
index_set = m.b1.index_set()
index = 1
location_set_map = get_location_set_map(index, index_set)
sets = ComponentSet()
pred_index = (1,)
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(pred_index, new_index)
sets = ComponentSet((m.time,))
pred_index = (slice(None),)
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(pred_index, new_index)
# Two one-d sets
index_set = m.b2.index_set()
index = (1,2)
location_set_map = get_location_set_map(index, index_set)
sets = ComponentSet()
pred_index = (1,2)
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(pred_index, new_index)
sets = ComponentSet((m.space,))
pred_index = (1, slice(None))
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(pred_index, new_index)
sets = ComponentSet((m.space, m.time))
pred_index = (slice(None), slice(None))
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(pred_index, new_index)
def test_multi_dim(self):
m = self.model()
index_set = m.b0d2.index_set()
index = ('a',1)
pred_map = {
0: m.d_2,
1: m.d_2,
}
location_set_map = get_location_set_map(index, index_set)
self.assertSameMap(pred_map, location_set_map)
normalize_index.flatten = False
# A single multi-dimensional index will behave the same when
# normalize_index.flatten is False. This is probably unexpected
# Should we just fail in this case?
# Note that the index (('a',1),) will not work.
index_set = m.b0d2.index_set()
index = ('a',1)
pred_map = {
0: m.d_2,
1: m.d_2,
}
with self.assertRaises(ValueError) as cm:
location_set_map = get_location_set_map(index, index_set)
self.assertIn('normalize_index.flatten', str(cm.exception))
normalize_index.flatten = True
index_set = m.b1d2.index_set()
index = (1, 'a', 1)
pred_map = {
0: m.time,
1: m.d_2,
2: m.d_2,
}
location_set_map = get_location_set_map(index, index_set)
self.assertSameMap(pred_map, location_set_map)
index_set = m.b2d2.index_set()
index = (1, 'a', 1, 2)
pred_map = {
0: m.time,
1: m.d_2,
2: m.d_2,
3: m.space,
}
location_set_map = get_location_set_map(index, index_set)
self.assertSameMap(pred_map, location_set_map)
index_set = m.b3d2.index_set()
index = (1, 'a', 1, 2, 1)
pred_map = {
0: m.time,
1: m.d_2,
2: m.d_2,
3: m.space,
4: m.time,
}
location_set_map = get_location_set_map(index, index_set)
self.assertSameMap(pred_map, location_set_map)
def test_multi_dim_replacement(self):
m = self.model()
# One two-dimen set
index_set = m.b0d2.index_set()
index = ('a',1)
location_set_map = get_location_set_map(index, index_set)
sets = ComponentSet()
pred_index = ('a', 1)
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
sets = ComponentSet((m.d_2,))
pred_index = (slice(None), slice(None))
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
# One two-dimen, one one-dimen
index_set = m.b1d2.index_set()
index = (1, 'a', 1)
location_set_map = get_location_set_map(index, index_set)
sets = ComponentSet()
pred_index = (1, 'a', 1)
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
sets = ComponentSet((m.d_2,))
pred_index = (1, slice(None), slice(None))
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
# One two-dimen, two one-dimen
index_set = m.b2d2.index_set()
index = (1, 'a', 1, 2)
location_set_map = get_location_set_map(index, index_set)
sets = ComponentSet((m.d_2, m.time))
pred_index = (slice(None), slice(None), slice(None), 2)
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
sets = ComponentSet((m.space,))
pred_index = (1, 'a', 1, slice(None))
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
# One two-dimen, three one-dimen
index_set = m.b3d2.index_set()
index = (1, 'a', 1, 2, 1)
location_set_map = get_location_set_map(index, index_set)
sets = ComponentSet((m.time,))
pred_index = (slice(None), 'a', 1, 2, slice(None))
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
def test_dimen_none(self):
m = self.model()
index_set = m.b0dn.index_set()
index = ('c',1,10)
pred_map = {
0: m.d_none,
1: m.d_none,
2: m.d_none,
}
location_set_map = get_location_set_map(index, index_set)
self.assertSameMap(pred_map, location_set_map)
index_set = m.b1dn.index_set()
index = (1, 'c', 1, 10)
pred_map = {
0: m.time,
1: m.d_none,
2: m.d_none,
3: m.d_none,
}
location_set_map = get_location_set_map(index, index_set)
self.assertSameMap(pred_map, location_set_map)
index_set = m.b1dnd2.index_set()
index = (1, 'c', 1, 10, 'a', 1)
pred_map = {
0: m.time,
1: m.d_none,
2: m.d_none,
3: m.d_none,
4: m.d_2,
5: m.d_2,
}
location_set_map = get_location_set_map(index, index_set)
self.assertSameMap(pred_map, location_set_map)
index_set = m.b2dn.index_set()
index = (1, 0, 'd', 3)
pred_map = {
0: m.time,
1: m.space,
2: m.d_none,
3: m.d_none,
}
location_set_map = get_location_set_map(index, index_set)
self.assertSameMap(pred_map, location_set_map)
index_set = m.b2dnd2.index_set()
index = (1, 'c', 1, 10, 'b', 2, 0)
pred_map = {
0: m.time,
1: m.d_none,
2: m.d_none,
3: m.d_none,
4: m.d_2,
5: m.d_2,
6: m.space,
}
location_set_map = get_location_set_map(index, index_set)
self.assertSameMap(pred_map, location_set_map)
index_set = m.dnd2b1.index_set()
index = ('d', 3, 'b', 2, 1)
pred_map = {
0: m.d_none,
1: m.d_none,
2: m.d_2,
3: m.d_2,
4: m.time,
}
location_set_map = get_location_set_map(index, index_set)
self.assertSameMap(pred_map, location_set_map)
index_set = m.b3dn.index_set()
index = (1, 'a', 1, 'd', 3, 0, 'b', 2)
pred_map = {
0: m.time,
1: m.d_2,
2: m.d_2,
3: m.d_none,
4: m.d_none,
5: m.space,
6: m.d_2,
7: m.d_2,
}
location_set_map = get_location_set_map(index, index_set)
self.assertSameMap(pred_map, location_set_map)
index_set = m.dn2.index_set()
index = (1, 'c', 1, 10, 'b', 2, 'd', 3, 'a', 1)
with self.assertRaises(RuntimeError) as cm:
location_set_map = get_location_set_map(index, index_set)
self.assertIn('multiple sets of dimen==None', str(cm.exception))
def test_dimen_none_replacement(self):
m = self.model()
# Single set of dimen None
index_set = m.b0dn.index_set()
index = ('c',1,10)
location_set_map = get_location_set_map(index, index_set)
sets = ComponentSet((m.d_none,))
pred_index = (Ellipsis,)
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
sets = ComponentSet()
pred_index = ('c', 1, 10)
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
# One None-dimen, one one-dimen
index_set = m.b1dn.index_set()
index = (1, 'c', 1, 10)
location_set_map = get_location_set_map(index, index_set)
sets = ComponentSet((m.d_none,))
pred_index = (1, Ellipsis)
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
sets = ComponentSet((m.d_none, m.time))
pred_index = (slice(None), Ellipsis)
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
# One None-dimen, one one-dimen, one two-dimen
index_set = m.b1dnd2.index_set()
index = (1, 'c', 1, 10, 'a', 1)
location_set_map = get_location_set_map(index, index_set)
sets = ComponentSet((m.d_none, m.time, m.d_2))
pred_index = (slice(None), Ellipsis, slice(None), slice(None))
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
sets = ComponentSet((m.time, m.d_2))
pred_index = (slice(None), 'c', 1, 10, slice(None), slice(None))
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
sets = ComponentSet((m.d_2,))
pred_index = (1, 'c', 1, 10, slice(None), slice(None))
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
# One None-dimen, one two-dimen
index_set = m.b2dn.index_set()
index = (1, 0, 'd', 3)
location_set_map = get_location_set_map(index, index_set)
sets = ComponentSet()
pred_index = (1, 0, 'd', 3)
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
sets = ComponentSet((m.d_none,))
pred_index = (1, 0, Ellipsis)
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
# One None-dimen, one two-dimen, two one-dimen
index_set = m.b2dnd2.index_set()
index = (1, 'c', 1, 10, 'b', 2, 0)
location_set_map = get_location_set_map(index, index_set)
sets = ComponentSet((m.d_none, m.space))
pred_index = (1, Ellipsis, 'b', 2, slice(None))
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
# None-dimen set comes first
index_set = m.dnd2b1.index_set()
index = ('d', 3, 'b', 2, 1)
location_set_map = get_location_set_map(index, index_set)
sets = ComponentSet((m.d_none, m.time))
pred_index = (Ellipsis, 'b', 2, slice(None))
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
# None-dimen, plus 2-d set shows up twice
index_set = m.b3dn.index_set()
index = (1, 'a', 1, 'd', 3, 0, 'b', 2)
location_set_map = get_location_set_map(index, index_set)
sets = ComponentSet((m.d_none, m.d_2))
pred_index = (
1,
slice(None),
slice(None),
Ellipsis,
0,
slice(None),
slice(None),
)
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
def test_replacement_normalize_flatten_false(self):
m = self.model()
# index of m.b2d2
index = (1, 0, ('a', 1))
location_set_map = {
0: m.time,
1: m.space,
2: m.d_2,
}
sets = ComponentSet((m.d_2,))
pred_index = (1, 0, slice(None))
# If you can assemble the proper location_set_map
# (get_location_set_map will not work), this function
# should work when normalize_index.flatten is False.
normalize_index.flatten = False
new_index = replace_indices(index, location_set_map, sets)
self.assertEqual(new_index, pred_index)
normalize_index.flatten = True
class SliceComponentDataAlongSets(unittest.TestCase):
def model(self):
m = pyo.ConcreteModel()
m.time = pyo.Set(initialize=[1,2,3])
m.space = dae.ContinuousSet(initialize=[0,2])
m.comp = pyo.Set(initialize=['a','b'])
m.d_2 = pyo.Set(initialize=[('a',1),('b',2)])
m.d_none = pyo.Set(initialize=[('c',1,10), ('d',3)], dimen=None)
@m.Block()
def b(b):
b.v0 = pyo.Var()
b.v1 = pyo.Var(m.time)
b.v2 = pyo.Var(m.time, m.space)
@b.Block(m.time, m.space)
def b2(b2):
b2.v0 = pyo.Var()
b2.v1 = pyo.Var(m.comp)
b2.v2 = pyo.Var(m.time, m.comp)
b2.vn = pyo.Var(m.time, m.d_none, m.d_2)
@b.Block(m.d_none, m.d_2)
def bn(bn):
bn.v0 = pyo.Var()
bn.v2 = pyo.Var(m.time, m.space)
bn.v3 = pyo.Var(m.time, m.space, m.time)
bn.vn = pyo.Var(m.time, m.d_none, m.d_2)
return m
def test_no_context(self):
m = self.model()
comp = m.b.v0
sets = ComponentSet((m.time, m.space))
_slice = slice_component_along_sets(comp, sets)
# Use `assertIs` when the "slice" is actually just a component.
self.assertIs(_slice, m.b.v0)
comp = m.b.v1[1]
sets = ComponentSet((m.time, m.space))
_slice = slice_component_along_sets(comp, sets)
self.assertEqual(_slice, m.b.v1[:])
comp = m.b.v2[1,0]
sets = ComponentSet((m.time, m.space))
_slice = slice_component_along_sets(comp, sets)
self.assertEqual(_slice, m.b.v2[:,:])
comp = m.b.b2[1,0].v1['a']
sets = ComponentSet((m.time, m.space))
_slice = slice_component_along_sets(comp, sets)
self.assertEqual(_slice, m.b.b2[:,:].v1['a'])
comp = m.b.b2[1,0].v1
sets = ComponentSet((m.time, m.space))
_slice = slice_component_along_sets(comp, sets)
self.assertEqual(_slice, m.b.b2[:,:].v1)
comp = m.b.b2[1,0].v2[1,'a']
sets = ComponentSet((m.time,))
_slice = slice_component_along_sets(comp, sets)
self.assertEqual(_slice, m.b.b2[:,0].v2[:,'a'])
comp = m.b.bn['c',1,10,'a',1].v3[1,0,1]
sets = ComponentSet((m.time,))
_slice = slice_component_along_sets(comp, sets)
self.assertEqual(_slice, m.b.bn['c',1,10,'a',1].v3[:,0,:])
def test_context(self):
m = self.model()
comp = m.b.v2[1,0]
context = m.b
sets = ComponentSet((m.time,))
_slice = slice_component_along_sets(comp, sets)
# Context makes no difference in resulting slice here.
self.assertEqual(_slice, m.b.v2[:,0])
comp = m.b.b2[1,0].v2[1,'a']
context = m.b.b2[1,0]
sets = ComponentSet((m.time,))
_slice = slice_component_along_sets(comp, sets, context=context)
# Here, context does make a difference
self.assertEqual(_slice, m.b.b2[1,0].v2[:,'a'])
comp = m.b.b2[1,0].v1['a']
context = m.b.b2[1,0]
sets = ComponentSet((m.time,))
_slice = slice_component_along_sets(comp, sets, context=context)
self.assertIs(_slice, m.b.b2[1,0].v1['a'])
sets = ComponentSet((m.comp,))
_slice = slice_component_along_sets(comp, sets, context=context)
self.assertEqual(_slice, m.b.b2[1,0].v1[:])
context = m.b.b2
sets = ComponentSet((m.time, m.comp))
_slice = slice_component_along_sets(comp, sets, context=context)
self.assertEqual(_slice, m.b.b2[:,0].v1[:])
def test_none_dimen(self):
m = self.model()
comp = m.b.b2[1,0].vn
sets = ComponentSet((m.d_none,))
_slice = slice_component_along_sets(comp, sets)
self.assertIs(_slice, m.b.b2[1,0].vn)
comp = m.b.b2[1,0].vn[1,'d',3,'a',1]
sets = ComponentSet((m.d_2, m.time))
_slice = slice_component_along_sets(comp, sets)
self.assertEqual(_slice, m.b.b2[:,0].vn[:,'d',3,:,:])
sets = ComponentSet((m.d_none, m.d_2))
_slice = slice_component_along_sets(comp, sets)
self.assertEqual(_slice, m.b.b2[1,0].vn[1,...,:,:])
comp = m.b.bn['c',1,10,'a',1].v3[1,0,1]
sets = ComponentSet((m.d_none, m.time))
_slice = slice_component_along_sets(comp, sets)
self.assertEqual(_slice, m.b.bn[...,'a',1].v3[:,0,:])
comp = m.b.bn['c',1,10,'a',1].vn[1,'d',3,'b',2]
sets = ComponentSet((m.d_none,))
_slice = slice_component_along_sets(comp, sets)
self.assertEqual(_slice, m.b.bn[...,'a',1].vn[1,...,'b',2])
sets = ComponentSet((m.d_none, m.d_2))
_slice = slice_component_along_sets(comp, sets)
self.assertEqual(_slice, m.b.bn[...,:,:].vn[1,...,:,:])
comp = m.b.bn['c',1,10,'a',1].vn[1,'d',3,'b',2]
context = m.b.bn['c',1,10,'a',1]
sets = ComponentSet((m.d_none,))
_slice = slice_component_along_sets(comp, sets, context=context)
self.assertEqual(_slice, m.b.bn['c',1,10,'a',1].vn[1,...,'b',2])
sets = ComponentSet((m.d_none, m.d_2))
_slice = slice_component_along_sets(comp, sets, context=context)
self.assertEqual(_slice, m.b.bn['c',1,10,'a',1].vn[1,...,:,:])
context = m.b.bn
sets = ComponentSet((m.d_none,))
_slice = slice_component_along_sets(comp, sets, context=context)
self.assertEqual(_slice, m.b.bn[...,'a',1].vn[1,...,'b',2])
if __name__ == '__main__':
unittest.main()