From cba22ebb29bb5826f84e58fc91d2f3d7762cf578 Mon Sep 17 00:00:00 2001 From: Nick Papior Date: Thu, 5 Dec 2024 20:10:56 +0100 Subject: [PATCH] enabled ellipsis here and there (#863) * enabled ellipsis here and there Should be improved so None fully gets removed. * Enabled ... for atoms, orbitals arguments This seems a bit more natural than None. Signed-off-by: Nick Papior --------- Signed-off-by: Nick Papior --- CHANGELOG.md | 6 ++++++ src/sisl/_core/geometry.py | 5 +++++ src/sisl/_core/tests/test_geometry.py | 5 +++++ src/sisl/io/siesta/stdout.py | 25 +++++++++++++++---------- src/sisl/typing/_indices.py | 8 ++++++++ 5 files changed, 39 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1a0aa3965a..0562510e70 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,12 @@ we hit release version 1.0.0. sisl.geom.graphene - added Nambu spin configuration, this is still experimental +- enabled `...` when extracting slices of MD steps in siesta output + files. + Here it is the same as `:`. But it also allows + inline arguments: `read_scf(imd=...)` where `imd=:` is not + allowed, partly fixes #835 +- enabled `...` for `atoms=` arguments. Selects all atoms. ### Fixed - `projection` arguments of several functions has been streamlined diff --git a/src/sisl/_core/geometry.py b/src/sisl/_core/geometry.py index 18254644a1..6c090c137e 100644 --- a/src/sisl/_core/geometry.py +++ b/src/sisl/_core/geometry.py @@ -366,9 +366,12 @@ def _sanitize_atoms(self, atoms) -> ndarray: - name -> self._names[name] - `Atom` -> self.atoms.index(atom) - range/list/ndarray -> ndarray + - `...` -> ndarray """ if atoms is None: return np.arange(self.na) + elif atoms is Ellipsis: + return np.arange(self.na) atoms = _a.asarray(atoms) if atoms.size == 0: return _a.asarrayl([]) @@ -452,6 +455,8 @@ def _sanitize_orbs(self, orbitals) -> ndarray: """ if orbitals is None: return np.arange(self.no) + elif orbitals is Ellipsis: + return np.arange(self.no) orbitals = _a.asarray(orbitals) if orbitals.size == 0: return _a.asarrayl([]) diff --git a/src/sisl/_core/tests/test_geometry.py b/src/sisl/_core/tests/test_geometry.py index 081cba296a..9354d2b056 100644 --- a/src/sisl/_core/tests/test_geometry.py +++ b/src/sisl/_core/tests/test_geometry.py @@ -1654,6 +1654,11 @@ def test_geometry_sort_int(): assert np.all(np.diff(bi.fxyz[ix, 1] * bi.lattice.length[i]) <= atol) +def test_geometry_ellipsis(): + gr = sisl_geom.graphene() + assert np.allclose(gr.axyz(...), gr.axyz(None)) + + def test_geometry_sort_atom(): bi = sisl_geom.bilayer().tile(2, 0).repeat(2, 1) diff --git a/src/sisl/io/siesta/stdout.py b/src/sisl/io/siesta/stdout.py index 89b73a71f4..27e8852d47 100644 --- a/src/sisl/io/siesta/stdout.py +++ b/src/sisl/io/siesta/stdout.py @@ -830,7 +830,7 @@ def read_data(self, *args, **kwargs) -> Any: def read_scf( self, key: str = "scf", - iscf: Optional[int] = -1, + iscf: Optional[Union[int, Ellipsis]] = -1, as_dataframe: bool = False, ret_header: bool = False, ): @@ -842,7 +842,7 @@ def read_scf( parse SCF information from Siesta SCF or TranSiesta SCF iscf : which SCF cycle should be stored. If ``-1`` only the final SCF step is stored, - for None *all* SCF cycles are returned. When `iscf` values queried are not found they + for `...`/`None` *all* SCF cycles are returned. When `iscf` values queried are not found they will be truncated to the nearest SCF step. as_dataframe: whether the information should be returned as a `pandas.DataFrame`. The advantage of this @@ -856,7 +856,9 @@ def read_scf( # These are the properties that are written in SIESTA scf props = ["iscf", "Eharris", "E_KS", "FreeEng", "dDmax", "Ef", "dHmax"] - if not iscf is None: + if iscf is Ellipsis: + iscf = None + elif iscf is not None: if iscf == 0: raise ValueError( f"{self.__class__.__name__}.read_scf requires iscf argument to *not* be 0!" @@ -1086,8 +1088,8 @@ def construct_data(d, data): def read_charge( self, name: Literal["voronoi", "hirshfeld", "mulliken", "mulliken:<5.2"], - iscf=Opt.ANY, - imd=Opt.ANY, + iscf: Union[Opt, int, Ellipsis] = Opt.ANY, + imd: Union[Opt, int, Ellipsis] = Opt.ANY, key_scf: str = "scf", as_dataframe: bool = False, ): @@ -1129,15 +1131,15 @@ def read_charge( ---------- name: the name of the charges that you want to read - iscf: int or Opt, optional + iscf: int or Opt or `...`, optional index (0-based) of the scf iteration you want the charges for. - If the enum specifier `Opt.ANY` or `Opt.ALL` are used, then + If the enum specifier `Opt.ANY` or `Opt.ALL`/`...` are used, then the returned quantities depend on what is present. If ``None/Opt.NONE`` it will not return any SCF charges. If both `imd` and `iscf` are ``None`` then only the final charges will be returned. imd: int or Opt, optional index (0-based) of the md step you want the charges for. - If the enum specifier `Opt.ANY` or `Opt.ALL` are used, then + If the enum specifier `Opt.ANY` or `Opt.ALL`/`...` are used, then the returned quantities depend on what is present. If ``None/Opt.NONE`` it will not return any MD charges. If both `imd` and `iscf` are ``None`` then only the final charges will be returned. @@ -1570,9 +1572,9 @@ def try_parse_int(s): md_scf_charge = pd.concat( [ pd.concat( - iscf, keys=pd.RangeIndex(1, len(iscf) + 1, name="iscf") + iscf_, keys=pd.RangeIndex(1, len(iscf_) + 1, name="iscf") ) - for iscf in md_scf_charge + for iscf_ in md_scf_charge ], keys=pd.RangeIndex(1, len(md_scf_charge) + 1, name="imd"), ) @@ -1607,6 +1609,9 @@ def _p(flag, found): flag : corrected flag """ + if flag is Ellipsis: + flag = Opt.ALL + if isinstance(flag, Opt): # correct flag depending on what `found` is # If the values have been found we diff --git a/src/sisl/typing/_indices.py b/src/sisl/typing/_indices.py index cfe914ce37..e32954347f 100644 --- a/src/sisl/typing/_indices.py +++ b/src/sisl/typing/_indices.py @@ -94,3 +94,11 @@ None, ] """Indexing orbitals via various construct methods""" + +try: + from types import EllipsisType + + AtomsIndex = Union[AtomsIndex, EllipsisType] + OrbitalsIndex = Union[OrbitalsIndex, EllipsisType] +except ImportError: + pass