# ___________________________________________________________________________ # # 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. # ___________________________________________________________________________ from pyutilib.misc import flatten_tuple from pyomo.common import DeveloperError from pyomo.core.base.set import SetOf, _SetDataBase from pyomo.core.base.component import Component, ComponentData from pyomo.core.base.indexed_component import ( IndexedComponent, UnindexedComponent_set ) from pyomo.core.base.indexed_component_slice import ( IndexedComponent_slice, _IndexedComponent_slice_iter ) import six from six import iteritems, advance_iterator if six.PY3: from collections.abc import MutableMapping as collections_MutableMapping from collections.abc import Set as collections_Set else: from collections import MutableMapping as collections_MutableMapping from collections import Set as collections_Set _NotSpecified = object() class _fill_in_known_wildcards(object): """Variant of "six.advance_iterator" that substitutes wildcard values This object is initialized with a tuple of index values. Calling the resulting object on a :py:class:`_slice_generator` will "advance" the iterator, substituting values from the tuple into the slice wildcards (":" indices), and returning the resulting object. The motivation for implementing this as an iterator is so that we can re-use all the logic from :py:meth:`_IndexedComponent_slice_iter.__next__` when looking up specific indices within the slice. Parameters ---------- wildcard_values : tuple of index values a tuple containing index values to substitute into the slice wildcards look_in_index : :py:class:`bool` [optional] If True, the iterator will also look for matches using the components' underlying index_set() in addition to the (sparse) indices matched by the components' __contains__() method. [default: False] get_if_not_present : :py:class:`bool` [optional] If True, the iterator will attempt to retrieve data objects (through getitem) for indexes that match the underlying component index_set() but do not appear in the (sparse) indices matched by __contains__. get_if_not_present implies look_in_index. [default: False] """ def __init__(self, wildcard_values, look_in_index=False, get_if_not_present=False): self.base_key = wildcard_values self.key = list(wildcard_values) self.known_slices = set() self.look_in_index = look_in_index or get_if_not_present self.get_if_not_present = get_if_not_present def __call__(self, _slice): """Advance the specified slice generator, substituting wildcard values This advances the passed :py:class:`_slice_generator ` by substituting values from the `wildcard_values` list for any wildcard slices ("`:`"). Parameters ---------- _slice : pyomo.core.base.indexed_component_slice._slice_generator the slice to advance """ if _slice in self.known_slices: raise StopIteration() self.known_slices.add(_slice) if _slice.ellipsis is None: idx_count = _slice.explicit_index_count elif not _slice.component.is_indexed(): idx_count = 1 else: idx_count = _slice.component.index_set().dimen if idx_count is None: raise SliceEllipsisLookupError( "Cannot lookup elements in a _ReferenceDict when the " "underlying slice object contains ellipsis over a jagged " "(dimen=None) Set") try: idx = tuple( _slice.fixed[i] if i in _slice.fixed else self.key.pop(0) for i in range(idx_count)) except IndexError: raise KeyError( "Insufficient values for slice of indexed component '%s' " "(found evaluating slice index %s)" % (_slice.component.name, self.base_key)) if idx in _slice.component: _slice.last_index = idx return _slice.component[idx] elif len(idx) == 1 and idx[0] in _slice.component: _slice.last_index = idx return _slice.component[idx[0]] elif self.look_in_index: if idx in _slice.component.index_set(): _slice.last_index = idx return _slice.component[idx] if self.get_if_not_present \ else None elif len(idx) == 1 and idx[0] in _slice.component.index_set(): _slice.last_index = idx return _slice.component[idx[0]] if self.get_if_not_present \ else None raise KeyError( "Index %s is not valid for indexed component '%s' " "(found evaluating slice index %s)" % (idx, _slice.component.name, self.base_key)) def check_complete(self): if self.key: raise KeyError("Extra (unused) values for slice index %s" % ( self.base_key, )) class SliceEllipsisLookupError(Exception): pass class _ReferenceDict(collections_MutableMapping): """A dict-like object whose values are defined by a slice. This implements a dict-like object whose keys and values are defined by a component slice (:py:class:`IndexedComponent_slice`). The intent behind this object is to replace the normal ``_data`` :py:class:`dict` in :py:class:`IndexedComponent` containers to create "reference" components. Parameters ---------- component_slice : :py:class:`IndexedComponent_slice` The slice object that defines the "members" of this mutable mapping. """ def __init__(self, component_slice): self._slice = component_slice def __contains__(self, key): try: advance_iterator(self._get_iter(self._slice, key)) return True except (StopIteration, KeyError): return False except SliceEllipsisLookupError: if type(key) is tuple and len(key) == 1: key = key[0] # Brute force (linear time) lookup _iter = iter(self._slice) for item in _iter: if _iter.get_last_index_wildcards() == key: return True return False def __getitem__(self, key): try: return advance_iterator( self._get_iter(self._slice, key, get_if_not_present=True) ) except StopIteration: raise KeyError("KeyError: %s" % (key,)) except SliceEllipsisLookupError: if type(key) is tuple and len(key) == 1: key = key[0] # Brute force (linear time) lookup _iter = iter(self._slice) for item in _iter: if _iter.get_last_index_wildcards() == key: return item raise KeyError("KeyError: %s" % (key,)) def __setitem__(self, key, val): tmp = self._slice.duplicate() op = tmp._call_stack[-1][0] if op == IndexedComponent_slice.get_item: tmp._call_stack[-1] = ( IndexedComponent_slice.set_item, tmp._call_stack[-1][1], val ) elif op == IndexedComponent_slice.slice_info: tmp._call_stack[-1] = ( IndexedComponent_slice.set_item, tmp._call_stack[-1][1], val ) elif op == IndexedComponent_slice.get_attribute: tmp._call_stack[-1] = ( IndexedComponent_slice.set_attribute, tmp._call_stack[-1][1], val ) else: raise DeveloperError( "Unexpected slice _call_stack operation: %s" % op) try: advance_iterator(self._get_iter(tmp, key, get_if_not_present=True)) except StopIteration: pass def __delitem__(self, key): tmp = self._slice.duplicate() op = tmp._call_stack[-1][0] if op == IndexedComponent_slice.get_item: # If the last attribute of the slice gets an item, # change it to delete the item tmp._call_stack[-1] = ( IndexedComponent_slice.del_item, tmp._call_stack[-1][1] ) elif op == IndexedComponent_slice.slice_info: assert len(tmp._call_stack) == 1 _iter = self._get_iter(tmp, key) try: advance_iterator(_iter) del _iter._iter_stack[0].component[_iter.get_last_index()] return except StopIteration: raise KeyError("KeyError: %s" % (key,)) elif op == IndexedComponent_slice.get_attribute: # If the last attribute of the slice retrieves an attribute, # change it to delete the attribute tmp._call_stack[-1] = ( IndexedComponent_slice.del_attribute, tmp._call_stack[-1][1] ) else: raise DeveloperError( "Unexpected slice _call_stack operation: %s" % op) try: advance_iterator(self._get_iter(tmp, key)) except StopIteration: pass def __iter__(self): return self._slice.wildcard_keys() def __len__(self): return sum(1 for i in self._slice) def iteritems(self): """Return the wildcard, value tuples for this ReferenceDict This method is necessary because the default implementation iterates over the keys and looks the values up in the dictionary. Unfortunately some slices have structures that make looking up components by the wildcard keys very expensive (linear time; e.g., the use of elipses with jagged sets). By implementing this method without using lookups, general methods that iterate over everything (like component.pprint()) will still be linear and not quadratic time. """ return self._slice.wildcard_items() def itervalues(self): """Return the values for this ReferenceDict This method is necessary because the default implementation iterates over the keys and looks the values up in the dictionary. Unfortunately some slices have structures that make looking up components by the wildcard keys very expensive (linear time; e.g., the use of elipses with jagged sets). By implementing this method without using lookups, general methods that iterate over everything (like component.pprint()) will still be linear and not quadratic time. """ return iter(self._slice) def _get_iter(self, _slice, key, get_if_not_present=False): if key.__class__ not in (tuple, list): key = (key,) return _IndexedComponent_slice_iter( _slice, _fill_in_known_wildcards(flatten_tuple(key), get_if_not_present=get_if_not_present) ) if six.PY3: _ReferenceDict.items = _ReferenceDict.iteritems _ReferenceDict.values = _ReferenceDict.itervalues class _ReferenceSet(collections_Set): """A set-like object whose values are defined by a slice. This implements a dict-like object whose members are defined by a component slice (:py:class:`IndexedComponent_slice`). :py:class:`_ReferenceSet` differs from the :py:class:`_ReferenceDict` above in that it looks in the underlying component ``index_set()`` for values that match the slice, and not just the (sparse) indices defined by the slice. Parameters ---------- component_slice : :py:class:`IndexedComponent_slice` The slice object that defines the "members" of this set """ def __init__(self, component_slice): self._slice = component_slice def __contains__(self, key): try: advance_iterator(self._get_iter(self._slice, key)) return True except (StopIteration, KeyError): return False except SliceEllipsisLookupError: if type(key) is tuple and len(key) == 1: key = key[0] # Brute force (linear time) lookup _iter = iter(self._slice) for item in _iter: if _iter.get_last_index_wildcards() == key: return True return False def __iter__(self): return self._slice.index_wildcard_keys() def __len__(self): return sum(1 for _ in self) def _get_iter(self, _slice, key): if key.__class__ not in (tuple, list): key = (key,) return _IndexedComponent_slice_iter( _slice, _fill_in_known_wildcards(flatten_tuple(key), look_in_index=True), iter_over_index=True ) def _identify_wildcard_sets(iter_stack, index): # if we have already decided that there isn't a comon index for the # slices, there is nothing more we can do. Bail. if index is None: return index # Walk the iter_stack that led to the current item and try to # identify the component wildcard sets tmp = [None]*len(iter_stack) for i, level in enumerate(iter_stack): if level is not None: offset = 0 wildcard_sets = {} for j,s in enumerate(level.component.index_set().subsets()): if s is UnindexedComponent_set: wildcard_sets[j] = s offset += 1 continue if s.dimen is None: return None wild = sum( 1 for k in range(s.dimen) if k+offset not in level.fixed ) if wild == s.dimen: wildcard_sets[j] = s elif wild != 0: # This subset is "touched" by an explicit slice, but # the whole set is not (i.e. there is a fixed # component to this subset). Therefore, as we # cannot extract that subset, we quit. return None offset += s.dimen # I believe that this check is not necessary: weirdnesses # with elipsis should get caught by the check for s.dimen # above. # #if offset != level.explicit_index_count: # return None tmp[i] = wildcard_sets if not index: return tmp # Any of the following would preclude identifying common sets. # However, I can't see a way to actually create any of these # situations (i.e., I can't test them). Assertions are left in for # defensive programming. assert len(index) == len(tmp) for i, level in enumerate(tmp): assert (index[i] is None) == (level is None) # No slices at this level in the slice if level is None: continue # if there are a differing number of subsets if len(index[i]) != len(level): return None # if any subset differs if any(index[i].get(j,None) is not _set for j,_set in iteritems(level)): return None return index def Reference(reference, ctype=_NotSpecified): """Creates a component that references other components ``Reference`` generates a *reference component*; that is, an indexed component that does not contain data, but instead references data stored in other components as defined by a component slice. The ctype parameter sets the :py:meth:`Component.type` of the resulting indexed component. If the ctype parameter is not set and all data identified by the slice (at construction time) share a common :py:meth:`Component.type`, then that type is assumed. If either the ctype parameter is ``None`` or the data has more than one ctype, the resulting indexed component will have a ctype of :py:class:`IndexedComponent`. If the indices associated with wildcards in the component slice all refer to the same :py:class:`Set` objects for all data identifed by the slice, then the resulting indexed component will be indexed by the product of those sets. However, if all data do not share common set objects, or only a subset of indices in a multidimentional set appear as wildcards, then the resulting indexed component will be indexed by a :py:class:`SetOf` containing a :py:class:`_ReferenceSet` for the slice. Parameters ---------- reference : :py:class:`IndexedComponent_slice` component slice that defines the data to include in the Reference component ctype : :py:class:`type` [optional] the type used to create the resulting indexed component. If not specified, the data's ctype will be used (if all data share a common ctype). If multiple data ctypes are found or type is ``None``, then :py:class:`IndexedComponent` will be used. Examples -------- .. doctest:: >>> from pyomo.environ import * >>> m = ConcreteModel() >>> @m.Block([1,2],[3,4]) ... def b(b,i,j): ... b.x = Var(bounds=(i,j)) ... >>> m.r1 = Reference(m.b[:,:].x) >>> m.r1.pprint() r1 : Size=4, Index=r1_index Key : Lower : Value : Upper : Fixed : Stale : Domain (1, 3) : 1 : None : 3 : False : True : Reals (1, 4) : 1 : None : 4 : False : True : Reals (2, 3) : 2 : None : 3 : False : True : Reals (2, 4) : 2 : None : 4 : False : True : Reals Reference components may also refer to subsets of the original data: .. doctest:: >>> m.r2 = Reference(m.b[:,3].x) >>> m.r2.pprint() r2 : Size=2, Index=b_index_0 Key : Lower : Value : Upper : Fixed : Stale : Domain 1 : 1 : None : 3 : False : True : Reals 2 : 2 : None : 3 : False : True : Reals Reference components may have wildcards at multiple levels of the model hierarchy: .. doctest:: >>> from pyomo.environ import * >>> m = ConcreteModel() >>> @m.Block([1,2]) ... def b(b,i): ... b.x = Var([3,4], bounds=(i,None)) ... >>> m.r3 = Reference(m.b[:].x[:]) >>> m.r3.pprint() r3 : Size=4, Index=r3_index Key : Lower : Value : Upper : Fixed : Stale : Domain (1, 3) : 1 : None : None : False : True : Reals (1, 4) : 1 : None : None : False : True : Reals (2, 3) : 2 : None : None : False : True : Reals (2, 4) : 2 : None : None : False : True : Reals The resulting reference component may be used just like any other component. Changes to the stored data will be reflected in the original objects: .. doctest:: >>> m.r3[1,4] = 10 >>> m.b[1].x.pprint() x : Size=2, Index=b[1].x_index Key : Lower : Value : Upper : Fixed : Stale : Domain 3 : 1 : None : None : False : True : Reals 4 : 1 : 10 : None : False : False : Reals """ if isinstance(reference, IndexedComponent_slice): pass elif isinstance(reference, Component): reference = reference[...] else: raise TypeError( "First argument to Reference constructors must be a " "component or component slice (received %s)" % (type(reference).__name__,)) _data = _ReferenceDict(reference) _iter = iter(reference) if ctype is _NotSpecified: ctypes = set() else: # If the caller specified a ctype, then we will prepopulate the # list to improve our chances of avoiding a scan of the entire # Reference ctypes = set((1,2)) index = [] for obj in _iter: ctypes.add(obj.ctype) if not isinstance(obj, ComponentData): # This object is not a ComponentData (likely it is a pure # IndexedComponent container). As the Reference will treat # it as if it *were* a ComponentData, we will skip ctype # identification and return a base IndexedComponent, thereby # preventing strange exceptions in the writers and with # things like pprint(). Of course, all of this logic is # skipped if the User knows better and forced a ctype on us. ctypes.add(0) if index is not None: index = _identify_wildcard_sets(_iter._iter_stack, index) # Note that we want to walk the entire slice, unless we can # prove that BOTH there aren't common indexing sets AND there is # more than one ctype. elif len(ctypes) > 1: break if not index: index = SetOf(_ReferenceSet(reference)) else: wildcards = sum((sorted(iteritems(lvl)) for lvl in index if lvl is not None), []) index = wildcards[0][1] if not isinstance(index, _SetDataBase): index = SetOf(index) for lvl, idx in wildcards[1:]: if not isinstance(idx, _SetDataBase): idx = SetOf(idx) index = index * idx if ctype is _NotSpecified: if len(ctypes) == 1: ctype = ctypes.pop() else: ctype = IndexedComponent elif ctype is None: ctype = IndexedComponent obj = ctype(index, ctype=ctype) obj._constructed = True obj._data = _data return obj