From 50b32bfa8f4a1f345684e14fd0211c7f161faf63 Mon Sep 17 00:00:00 2001 From: Nicolas Date: Wed, 3 Sep 2025 15:52:38 +0200 Subject: [PATCH] Numpy-like array slicing for fluids and elements added; get_n() function added for fluids and elements --- src/gaspype/_main.py | 80 ++++++++++++++++++++++++++++++++++++------- src/gaspype/typing.py | 4 +++ tests/test_slicing.py | 42 +++++++++++------------ 3 files changed, 92 insertions(+), 34 deletions(-) diff --git a/src/gaspype/_main.py b/src/gaspype/_main.py index 36f3575..1bd0afe 100644 --- a/src/gaspype/_main.py +++ b/src/gaspype/_main.py @@ -7,7 +7,7 @@ from gaspype._phys_data import atomic_weights, db_reader import re import pkgutil from .constants import R, epsy, p0 -from .typing import FloatArray, NDFloat, Shape +from .typing import FloatArray, NDFloat, Shape, ArrayIndices T = TypeVar('T', 'fluid', 'elements') @@ -483,6 +483,28 @@ class fluid: assert set(species) <= set(self.fs.species), f'Species {", ".join([s for s in species if s not in self.fs.species])} is/are not part of the fluid system' return self.array_fractions[..., [self.fs.species.index(k) for k in species]] + def get_n(self, species: str | list[str] | None = None) -> FloatArray: + """Get molar amount of fluid species + + Args: + species: A single species name, a list of species names or None for + returning the amount of all species + + Returns: + Returns an array of floats with the molar amount of the species. + If the a single species name is provided the return float array has + the same dimensions as the fluid type. If a list or None is provided + the return array has an additional dimension for the species. + """ + if not species: + return self.array_composition + elif isinstance(species, str): + assert species in self.fs.species, f'Species {species} is not part of the fluid system' + return self.array_composition[..., self.fs.species.index(species)] + else: + assert set(species) <= set(self.fs.species), f'Species {", ".join([s for s in species if s not in self.fs.species])} is/are not part of the fluid system' + return self.array_composition[..., [self.fs.species.index(k) for k in species]] + def __add__(self, other: T) -> T: return array_operation(self, other, np.add) @@ -510,16 +532,21 @@ class fluid: # def __array__(self) -> FloatArray: # return self.array_composition - def __getitem__(self, key: str | int | list[str] | list[int] | slice) -> FloatArray: + @overload + def __getitem__(self, key: str) -> FloatArray: + pass + + @overload + def __getitem__(self, key: ArrayIndices) -> 'fluid': + pass + + def __getitem__(self, key: str | ArrayIndices) -> Any: if isinstance(key, str): assert key in self.fs.species, f'Species {key} is not part of the fluid system' return self.array_composition[..., self.fs.species.index(key)] - elif isinstance(key, (slice, int)): - return self.array_composition[..., key] else: - mset = set(self.fs.species) | set(range(len(self.fs.species))) - assert set(key) <= mset, f'Species {", ".join([str(s) for s in key if s not in mset])} is/are not part of the fluid system' - return self.array_composition[..., [self.fs.species.index(k) if isinstance(k, str) else k for k in key]] + key_tuple = key if isinstance(key, tuple) else (key,) + return fluid(self.array_composition[(*key_tuple, slice(None))], self.fs) def __iter__(self) -> Iterator[dict[str, float]]: assert len(self.shape) < 2, 'Cannot iterate over species with more than one dimension' @@ -614,6 +641,28 @@ class elements: """ return np.sum(self.array_elemental_composition * self.fs.array_atomic_mass, axis=-1, dtype=NDFloat) + def get_n(self, elemental_species: str | list[str] | None = None) -> FloatArray: + """Get molar amount of elements + + Args: + elemental_species: A single element name, a list of element names or None for + returning the amount of all element + + Returns: + Returns an array of floats with the molar amount of the elements. + If the a single element name is provided the return float array has + the same dimensions as the fluid type. If a list or None is provided + the return array has an additional dimension for the elements. + """ + if not elemental_species: + return self.array_elemental_composition + elif isinstance(elemental_species, str): + assert elemental_species in self.fs.elements, f'Element {elemental_species} is not part of the fluid system' + return self.array_elemental_composition[..., self.fs.elements.index(elemental_species)] + else: + assert set(elemental_species) <= set(self.fs.elements), f'Elements {", ".join([s for s in elemental_species if s not in self.fs.elements])} is/are not part of the fluid system' + return self.array_elemental_composition[..., [self.fs.elements.index(k) for k in elemental_species]] + def __add__(self, other: 'fluid | elements') -> 'elements': return array_operation(self, other, np.add) @@ -639,16 +688,21 @@ class elements: def __array__(self) -> FloatArray: return self.array_elemental_composition - def __getitem__(self, key: str | int | list[str] | list[int] | slice) -> FloatArray: + @overload + def __getitem__(self, key: str) -> FloatArray: + pass + + @overload + def __getitem__(self, key: ArrayIndices) -> 'elements': + pass + + def __getitem__(self, key: str | ArrayIndices) -> Any: if isinstance(key, str): assert key in self.fs.elements, f'Element {key} is not part of the fluid system' return self.array_elemental_composition[..., self.fs.elements.index(key)] - elif isinstance(key, (slice, int)): - return self.array_elemental_composition[..., key] else: - mset = set(self.fs.elements) | set(range(len(self.fs.elements))) - assert set(key) <= mset, f'Elements {", ".join([str(s) for s in key if s not in mset])} is/are not part of the fluid system' - return self.array_elemental_composition[..., [self.fs.elements.index(k) if isinstance(k, str) else k for k in key]] + key_tuple = key if isinstance(key, tuple) else (key,) + return elements(self.array_elemental_composition[(*key_tuple, slice(None))], self.fs) def __iter__(self) -> Iterator[dict[str, float]]: assert len(self.shape) < 2, 'Cannot iterate over elements with more than one dimension' diff --git a/src/gaspype/typing.py b/src/gaspype/typing.py index ad935c4..fdf963f 100644 --- a/src/gaspype/typing.py +++ b/src/gaspype/typing.py @@ -1,6 +1,10 @@ from numpy import float64 from numpy.typing import NDArray +from typing import Sequence +from types import EllipsisType Shape = tuple[int, ...] NDFloat = float64 FloatArray = NDArray[NDFloat] +ArrayIndex = int | slice | None | EllipsisType | Sequence[int] +ArrayIndices = ArrayIndex | tuple[ArrayIndex, ...] diff --git a/tests/test_slicing.py b/tests/test_slicing.py index d4a6633..ed41e36 100644 --- a/tests/test_slicing.py +++ b/tests/test_slicing.py @@ -12,28 +12,28 @@ def test_str_index(): assert el['C'].shape == (2, 3, 4) -def test_str_list_index(): - assert fl[['CO2', 'H2', 'CO']].shape == (2, 3, 4, 3) - assert el[['C', 'H', 'O']].shape == (2, 3, 4, 3) +def test_single_axis_int_index(): + assert fl[0].shape == (3, 4) + assert fl[1].shape == (3, 4) + assert el[1].shape == (3, 4) + assert el[0].shape == (3, 4) -def test_int_list_index(): - assert fl[[1, 2, 0, 5]].shape == (2, 3, 4, 4) - assert el[[1, 2, 0, 3]].shape == (2, 3, 4, 4) +def test_single_axis_int_list(): + assert fl[:, [0, 1]].shape == (2, 2, 4) + assert el[:, [0, 1]].shape == (2, 2, 4) -def test_mixed_list_index(): - assert el[[1, 'H', 0, 'O']].shape == (2, 3, 4, 4) - - -def test_int_index(): - assert fl[5].shape == (2, 3, 4) - assert el[-1].shape == (2, 3, 4) - - -def test_slice_index(): - assert fl[0:3].shape == (2, 3, 4, 3) - assert fl[:].shape == (2, 3, 4, 6) - - assert el[0:3].shape == (2, 3, 4, 3) - assert el[:].shape == (2, 3, 4, 4) +def test_multi_axis_int_index(): + assert fl[0, 1].shape == (4,) + assert fl[0, 1, 2].shape == tuple() + assert fl[0, 2].shape == (4,) + assert fl[:, 2, :].shape == (2, 4) + assert fl[0, [1, 2]].shape == (2, 4) + assert fl[..., 0].shape == (2, 3) + assert el[0, 1].shape == (4,) + assert el[0, 1, 2].shape == tuple() + assert el[0, 2].shape == (4,) + assert el[:, 2, :].shape == (2, 4) + assert el[0, [1, 2]].shape == (2, 4) + assert el[..., 0].shape == (2, 3)