Merge pull request #871 from zerothi/viz-tbt-ados

allowed atoms in tbtnc, and enabled PDOS of ADOS in viz

CHANGELOG.md

@@ -9,6 +9,7 @@ we hit release version 1.0.0.
 ## [0.15.3] - YYYY-MM-DD
 
 ### Added
+- added ADOS extraction of PDOS data in `sisl.viz`
 - enabled submodule access without imports:
 
       import sisl

src/sisl/io/tbtrans/_cdf.py

@@ -54,7 +54,13 @@ def read_lattice(self) -> Lattice:
         return lattice
 
     def read_geometry(self, *args, **kwargs) -> Geometry:
-        """Returns `Geometry` object from this file"""
+        """Returns `Geometry` object from this file
+
+        Parameters
+        ----------
+        atoms :
+            atoms used instead of random species
+        """
         lattice = self.read_lattice()
 
         xyz = _a.arrayd(np.copy(self.xa))
@@ -64,23 +70,23 @@ def read_geometry(self, *args, **kwargs) -> Geometry:
         lasto = _a.arrayi(np.copy(self.lasto) + 1)
         nos = np.diff(lasto, prepend=0)
 
-        if "atom" in kwargs:
-            # The user "knows" which atoms are present
-            atms = kwargs["atom"]
+        atoms = kwargs.get("atoms", kwargs.get("atom"))
+
+        if atoms is not None:
             # Check that all atoms have the correct number of orbitals.
             # Otherwise we will correct them
-            for i in range(len(atms)):
-                if atms[i].no != nos[i]:
-                    atms[i] = Atom(atms[i].Z, [-1] * nos[i], tag=atms[i].tag)
+            for i in range(len(atoms)):
+                if atoms[i].no != nos[i]:
+                    atoms[i] = Atom(atoms[i].Z, [-1] * nos[i], tag=atoms[i].tag)
 
         else:
             # Default to Hydrogen atom with nos[ia] orbitals
             # This may be counterintuitive but there is no storage of the
             # actual species
-            atms = [Atom("H", [-1] * o) for o in nos]
+            atoms = [Atom("H", [-1] * o) for o in nos]
 
         # Create and return geometry object
-        geom = Geometry(xyz, atms, lattice=lattice)
+        geom = Geometry(xyz, atoms, lattice=lattice)
 
         return geom
 

src/sisl/viz/data/pdos.py

@@ -255,9 +255,7 @@ def from_fdf(
             Extra arguments to be passed to the PDOSData constructor, which depends
             on the source requested.
 
-            Except for the hamiltonian source, no extra arguments are needed (and they
-            won't be used). See PDOSData.from_hamiltonian for the extra arguments accepted
-            by the hamiltonian data constructor.
+            One should check ``PDOSData.from_*`` for details for each PDOS retrieval.
         """
         if source == "pdos":
             sile = FileDataSIESTA(fdf=fdf, cls=pdosSileSiesta)
@@ -272,7 +270,7 @@ def from_fdf(
 
             geometry = fdf.read_geometry(output=True)
 
-            return cls.new(sile, geometry=geometry)
+            return cls.new(sile, geometry=geometry, **kwargs)
         elif source == "wfsx":
             sile = FileDataSIESTA(fdf=fdf, cls=wfsxSileSiesta)
 
@@ -304,7 +302,10 @@ def from_siesta_pdos(cls, pdos_file: pdosSileSiesta):
     @new.register
     @classmethod
     def from_tbtrans(
-        cls, tbt_nc: tbtncSileTBtrans, geometry: Union[Geometry, None] = None
+        cls,
+        tbt_nc: tbtncSileTBtrans,
+        geometry: Union[Geometry, None] = None,
+        elec: Union[int, str, None] = None,
     ):
         """Reads the PDOS from a *.TBT.nc file coming from a TBtrans run.
 
@@ -316,18 +317,27 @@ def from_tbtrans(
             Full geometry of the system (including scattering and electrode regions).
             Right now only used to get the basis of each atom, which is not
             stored in the TBT.nc file.
+        elec:
+            which electrode to get the PDOS from. Can be None for the Green function,
+            otherwise the specified. Otherwise it is the index/name of an electrode
+            so that one gets the ADOS from that electrode.
         """
-        PDOS = tbt_nc.DOS(sum=False).T
+        if elec is None:
+            elec = "Gf"
+            PDOS = tbt_nc.DOS(sum=False).T
+        else:
+            PDOS = tbt_nc.ADOS(elec, sum=False).T
+            elec = f"A{elec}"
         E = tbt_nc.E
 
         read_geometry_kwargs = {}
         if geometry is not None:
-            read_geometry_kwargs["atom"] = geometry.atoms
+            read_geometry_kwargs["atoms"] = geometry.atoms
 
         # Read the geometry from the TBT.nc file and get only the device part
         geometry = tbt_nc.read_geometry(**read_geometry_kwargs).sub(tbt_nc.a_dev)
 
-        return cls.new(PDOS, geometry, E)
+        return cls.new(PDOS, geometry, E, extra_attrs={"elec": elec})
 
     @new.register
     @classmethod