diff --git a/andes/models/base.py b/andes/models/base.py index 586399a70..c48e43020 100644 --- a/andes/models/base.py +++ b/andes/models/base.py @@ -424,15 +424,21 @@ def get_field(self, field, idx=None, astype=None): # ================================================================= uid = self.get_uid(idx) - if not astype: - astype = type(self.__dict__[field]) - ret = matrix(self.__dict__[field])[uid] + field_data = self.__dict__[field] - if isinstance(idx, (float, int, str)): - return ret - else: - return astype(ret) + if isinstance(field_data, matrix): + ret = field_data[uid] + elif isinstance(field_data, list): + if isinstance(idx, (float, int, str)): + ret = field_data[uid] + else: + ret = [field_data[x] for x in uid] + + if astype is not None: + ret = astype(ret) + + return ret def _alloc(self): """ diff --git a/andes/models/zone.py b/andes/models/zone.py index 3bc0c60d7..8d360c4a5 100644 --- a/andes/models/zone.py +++ b/andes/models/zone.py @@ -208,6 +208,7 @@ class Area(Zone): def __init__(self, system, name): super().__init__(system, name) self._name = 'Area' + self._group = 'Collection' self._init() def setup(self): @@ -221,6 +222,7 @@ class Region(Zone): def __init__(self, system, name): super().__init__(system, name) self._name = 'Region' + self._group = 'Balancing' self._params.extend(['Ptol', 'slack']) self._descr.update({ 'Ptol': 'Total transfer capacity', diff --git a/andes/system.py b/andes/system.py index cd912d6e2..3bca53ca8 100644 --- a/andes/system.py +++ b/andes/system.py @@ -625,6 +625,7 @@ def __init__(self, system, name): self.name = name self.all_models = [] self._idx_model = {} + self._idx = [] def register_model(self, model): """ @@ -650,9 +651,14 @@ def register_element(self, model, idx): if idx is None: idx = model + '_' + str(len(self._idx_model)) - assert idx not in self._idx_model.values() + # TODO: `in` a list test is slow for large lists. Consider `bisect` + if idx in self._idx: + raise IndexError("Model {} idx {} already exist in model {}". + format(model, idx, self._idx_model[idx])) self._idx_model[idx] = model + self._idx.append(idx) + return idx def get_field(self, field, idx):