diff --git a/src/gaspype/_main.py b/src/gaspype/_main.py index 9c83251..0ba80e8 100644 --- a/src/gaspype/_main.py +++ b/src/gaspype/_main.py @@ -525,7 +525,9 @@ class fluid: return self.array_composition[..., [self.fs.species.index(k) if isinstance(k, str) else k for k in key]] def __iter__(self) -> Iterator[dict[str, float]]: - return iter({s: c for s, c in zip(self.fs.species, spa)} for spa in np.array(self.array_composition, ndmin=2)) + assert len(self.shape) < 2, 'Cannot iterate over species with more than one dimension' + aec = self.array_composition.reshape(-1, len(self.fs.species)) + return iter({s: c for s, c in zip(self.fs.species, aec[i, :])} for i in range(aec.shape[0])) def __repr__(self) -> str: if len(self.array_fractions.shape) == 1: @@ -652,7 +654,9 @@ class elements: return self.array_elemental_composition[..., [self.fs.elements.index(k) if isinstance(k, str) else k for k in key]] def __iter__(self) -> Iterator[dict[str, float]]: - return iter({s: c for s, c in zip(self.fs.elements, spa)} for spa in np.array(self.array_elemental_composition, ndmin=2)) + assert len(self.shape) < 2, 'Cannot iterate over elements with more than one dimension' + aec = self.array_elemental_composition.reshape(-1, len(self.fs.elements)) + return iter({s: c for s, c in zip(self.fs.elements, aec[i, :])} for i in range(aec.shape[0])) def __repr__(self) -> str: if len(self.array_elemental_composition.shape) == 1: diff --git a/tests/test_to_pandas.py b/tests/test_to_pandas.py index f9d646e..ed3d122 100644 --- a/tests/test_to_pandas.py +++ b/tests/test_to_pandas.py @@ -21,3 +21,13 @@ def test_elements(): df = pd.DataFrame(list(gp.elements(fl * np.array([1, 2, 3, 4])))) assert df.shape == (4, 2) + +def test_iter(): + fl = gp.fluid({'O2': 1, 'H2': 2, 'H2O': 3}) + + fl2 = fl * np.array([1, 2, 3, 4]) + for i, f in enumerate(fl2): + if i == 1: + assert f == {'O2': np.float64(2.0), 'H2': np.float64(4.0), 'H2O': np.float64(6.0)} + if i == 3: + assert f == {'O2': np.float64(4.0), 'H2': np.float64(8.0), 'H2O': np.float64(12.0)}