Fix find_idx

.github/workflows/pythonapp.yml

@@ -16,6 +16,14 @@ jobs:
         channels: conda-forge,defaults
         channel-priority: true
         activate-environment: anaconda-client-env
+    - shell: bash -el {0}
+      name: Configure conda to use default solver
+      run: |
+        conda config --set solver classic
+    - shell: bash -el {0}
+      name: Install mamba
+      run: |
+        conda install -n base -c conda-forge mamba
     - shell: bash -el {0}
       name: Install dependencies
       run: |

andes/core/model/modeldata.py

@@ -11,6 +11,7 @@
 from andes.core.param import (BaseParam, DataParam, IdxParam, NumParam,
                               TimerParam)
 from andes.shared import pd
+from andes.utils.func import list_flatten
 
 logger = logging.getLogger(__name__)
 
@@ -277,7 +278,8 @@ def find_param(self, prop):
 
         return out
 
-    def find_idx(self, keys, values, allow_none=False, default=False):
+    def find_idx(self, keys, values, allow_none=False, default=False,
+                 no_flatten=False):
         """
         Find `idx` of devices whose values match the given pattern.
 
@@ -292,11 +294,30 @@ def find_idx(self, keys, values, allow_none=False, default=False):
             Allow key, value to be not found. Used by groups.
         default : bool
             Default idx to return if not found (missing)
+        no_flatten : bool
+            If True, return the non-flattened list of idxes. Otherwise, flatten the list.
 
         Returns
         -------
         list
             indices of devices
+
+        Examples
+        --------
+        >>> # Use example case of IEEE 14-bus system with PVD1
+        >>> ss = andes.load(andes.get_case('ieee14/ieee14_pvd1.xlsx'))
+
+        >>> # To find the idx of `PVD1` with `name` of 'PVD1_1' and 'PVD1_2'
+        >>> ss.find_idx(keys='name', values=['PVD1_1', 'PVD1_2'])
+        [1, 2]
+
+        >>> # To find the idx of `PVD1` with `gammap` equals to 0.1
+        >>> ss.find_idx(keys='gammap', values=0.1)
+        [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]
+
+        >>> # To find the idx of `PVD1` with `gammap` equals to 0.1 and `name` of 'PVD1_1'
+        >>> ss.find_idx(keys=['gammap', 'name'], values=[[0.1], ['PVD1_1']])
+        [1]
         """
         if isinstance(keys, str):
             keys = (keys,)
@@ -316,21 +337,33 @@ def find_idx(self, keys, values, allow_none=False, default=False):
             if len(keys) != len(values):
                 raise ValueError("keys and values must have the same length")
 
+            # check if all elements in values have the same length
+            iterator = iter(values)
+            try:
+                first_length = len(next(iterator))
+                ok_len = all(len(element) == first_length for element in iterator)
+            except StopIteration:  # empty iterable
+                ok_len = True
+            if not ok_len:
+                raise ValueError("All elements in values must have the same length")
+
         v_attrs = [self.__dict__[key].v for key in keys]
 
         idxes = []
         for v_search in zip(*values):
-            v_idx = None
+            v_idx_list = []
             for pos, v_attr in enumerate(zip(*v_attrs)):
                 if all([i == j for i, j in zip(v_search, v_attr)]):
-                    v_idx = self.idx.v[pos]
-                    break
-            if v_idx is None:
+                    v_idx_list.append(self.idx.v[pos])
+            if not v_idx_list:
                 if allow_none is False:
                     raise IndexError(f'{list(keys)}={v_search} not found in {self.class_name}')
                 else:
-                    v_idx = default
+                    v_idx_list.append(default)
 
-            idxes.append(v_idx)
+            idxes.append(v_idx_list)
 
-        return idxes
+        if not no_flatten:
+            return list_flatten(idxes)
+        else:
+            return idxes

andes/models/group.py

@@ -243,31 +243,41 @@ def set(self, src: str, idx, attr, value):
 
         return True
 
-    def find_idx(self, keys, values, allow_none=False, default=None):
+    def find_idx(self, keys, values, allow_none=False, default=None,
+                 no_flatten=False):
         """
         Find indices of devices that satisfy the given `key=value` condition.
 
         This method iterates over all models in this group.
         """
-        indices_found = []
-        # `indices_found` contains found indices returned from all models of this group
+        idx_mdls = []
         for model in self.models.values():
-            indices_found.append(model.find_idx(keys, values, allow_none=True, default=default))
-
-        out = []
-        for idx, idx_found in enumerate(zip(*indices_found)):
-            if not allow_none:
-                if idx_found.count(None) == len(idx_found):
-                    missing_values = [item[idx] for item in values]
-                    raise IndexError(f'{list(keys)} = {missing_values} not found in {self.class_name}')
-
-            real_idx = default
-            for item in idx_found:
-                if item is not None:
-                    real_idx = item
-                    break
-            out.append(real_idx)
-        return out
+            idx_mdls.append(model.find_idx(keys, values, allow_none=True, default=default,
+                                           no_flatten=True))
+
+        # `indices_found` contains found indices returned from all models of this group
+        # NOTE: if the idx returned to [default] across all models, it means there is
+        # no such idx in this group. If so, return default or raise an key error.
+        indices_found = []
+        uid_missing = []
+        for uid, col in enumerate(zip(*idx_mdls)):
+            if all(item == [default] for item in col):
+                if allow_none:
+                    indices_found.append([default])
+                else:
+                    uid_missing.append(uid)
+            else:
+                col_filter = [item for item in col if item != [default]]
+                indices_found.append(list_flatten(col_filter))
+
+        if uid_missing:
+            miss_str = f'{keys}={[v[u] for v in values for u in uid_missing]}'
+            raise IndexError(f'Group <{self.class_name}> does not contain device with {miss_str}')
+
+        if not no_flatten:
+            return list_flatten(indices_found)
+        else:
+            return indices_found
 
     def _check_src(self, src: str):
         """

andes/models/misc/output.py

@@ -50,9 +50,9 @@ def in1d(self, addr, v_code):
         """
 
         if v_code == 'x':
-            return np.in1d(self.xidx, addr)
+            return np.isin(self.xidx, addr)
         if v_code == 'y':
-            return np.in1d(self.yidx, addr)
+            return np.isin(self.yidx, addr)
 
         raise NotImplementedError("v_code <%s> not recognized" % v_code)
 

docs/source/release-notes.rst

@@ -19,6 +19,7 @@ v1.9.3 (2024-04-XX)
 - Adjust `BusFreq.Tw.default` to 0.1.
 - Add parameter from_csv=None in TDS.run() to allow loading data from CSV files at TDS begining.
 - Fix `TDS.init()` and `TDS._csv_step()` to fit loading from CSV when `Output` exists.
+- Fix `ModelData.find_idx()` to return all matches
 
 v1.9.2 (2024-03-25)
 -------------------

tests/test_group.py

@@ -71,6 +71,7 @@ def test_group_access(self):
                                 [6, 7, 8, 1])
 
         # --- find_idx ---
+        # same Model
         self.assertListEqual(ss.DG.find_idx('name', ['PVD1_1', 'PVD1_2']),
                              ss.PVD1.find_idx('name', ['PVD1_1', 'PVD1_2']),
                              )
@@ -82,6 +83,22 @@ def test_group_access(self):
                                               [('PVD1_1', 'PVD1_2'),
                                                (1.0, 1.0)]))
 
+        # cross Model, given results
+        self.assertListEqual(ss.StaticGen.find_idx(keys='bus',
+                                                   values=[1, 2, 3, 4]),
+                             [1, 2, 3, 6])
+
+        self.assertListEqual(ss.StaticGen.find_idx(keys='bus',
+                                                   values=[1, 2, 3, 4],
+                                                   no_flatten=True),
+                             [[1], [2], [3], [6]])
+
+        self.assertListEqual(ss.StaticGen.find_idx(keys='bus',
+                                                   values=[1, 2, 3, 4, 2024],
+                                                   allow_none=True,
+                                                   default=2011),
+                             [1, 2, 3, 6, 2011])
+
         # --- get_field ---
         ff = ss.DG.get_field('f', list(ss.DG._idx2model.keys()), 'v_code')
         self.assertTrue(any([item == 'y' for item in ff]))

tests/test_model_set.py

@@ -54,3 +54,52 @@ def test_model_set(self):
         ss.GENROU.set("M", np.array(["GENROU_4"]), "v", 6.0)
         np.testing.assert_equal(ss.GENROU.M.v[3], 6.0)
         self.assertEqual(ss.TDS.Teye[omega_addr[3], omega_addr[3]], 6.0)
+
+    def test_find_idx(self):
+        ss = andes.load(andes.get_case('ieee14/ieee14_pvd1.xlsx'))
+        mdl = ss.PVD1
+
+        # multiple values
+        self.assertListEqual(mdl.find_idx(keys='name', values=['PVD1_1', 'PVD1_2'],
+                                          allow_none=False, default=False, no_flatten=False),
+                             [1, 2])
+        # multiple values, no flatten
+        self.assertListEqual(mdl.find_idx(keys='name', values=['PVD1_1', 'PVD1_2'],
+                                          allow_none=False, default=False, no_flatten=True),
+                             [[1], [2]])
+        # non-existing value
+        self.assertListEqual(mdl.find_idx(keys='name', values=['PVD1_999'],
+                                          allow_none=True, default=False, no_flatten=False),
+                             [False])
+
+        # non-existing value is not allowed
+        with self.assertRaises(IndexError):
+            mdl.find_idx(keys='name', values=['PVD1_999'],
+                         allow_none=False, default=False, no_flatten=False)
+
+        # multiple keys
+        self.assertListEqual(mdl.find_idx(keys=['gammap', 'name'],
+                                          values=[[0.1, 0.1], ['PVD1_1', 'PVD1_2']]),
+                             [1, 2])
+
+        # multiple keys, with non-existing values
+        self.assertListEqual(mdl.find_idx(keys=['gammap', 'name'],
+                                          values=[[0.1, 0.1], ['PVD1_1', 'PVD1_999']],
+                                          allow_none=True, default='CURENT'),
+                             [1, 'CURENT'])
+
+        # multiple keys, with non-existing values not allowed
+        with self.assertRaises(IndexError):
+            mdl.find_idx(keys=['gammap', 'name'],
+                         values=[[0.1, 0.1], ['PVD1_1', 'PVD1_999']],
+                         allow_none=False, default=999)
+
+        # multiple keys, values are not iterable
+        with self.assertRaises(ValueError):
+            mdl.find_idx(keys=['gammap', 'name'],
+                         values=[0.1, 0.1])
+
+        # multiple keys, items length are inconsistent in values
+        with self.assertRaises(ValueError):
+            mdl.find_idx(keys=['gammap', 'name'],
+                         values=[[0.1, 0.1], ['PVD1_1']])