Exemplo n.º 1
0
    def ArgumentParser(self, p=None, *args, **kwargs):
        """ Returns the arguments that is available for this Sile """

        # We limit the import to occur here
        import argparse

        namespace = default_namespace(_tbtse=self, _geometry=self.geom)

        class Info(argparse.Action):
            """ Action to print information contained in the TBT.SE.nc file, helpful before performing actions """

            def __call__(self, parser, ns, value, option_string=None):
                # First short-hand the file
                print(ns._tbtse.info(value))

        p.add_argument('--info', '-i', action=Info, nargs='?', metavar='ELEC',
                       help='Print out what information is contained in the TBT.SE.nc file, optionally only for one of the electrodes.')

        return p, namespace
Exemplo n.º 2
0
    def ArgumentParser(self, p=None, *args, **kwargs):
        """ Returns the arguments that is available for this Sile """

        # We limit the import to occur here
        import argparse
        import warnings

        comment = 'Fermi-level shifted to 0'
        with warnings.catch_warnings(record=True) as w:
            # Cause all warnings to always be triggered.
            warnings.simplefilter("always")
            geometry, E, PDOS = self.read_data()

            if len(w) > 0:
                if issubclass(w[-1].category, SislWarning):
                    comment = 'Fermi-level unknown'

        def _sum_filter(PDOS):
            if PDOS.ndim == 2:
                # non-polarized
                return PDOS
            elif PDOS.ndim == 3:
                # polarized
                return PDOS.sum(0)
            return PDOS[0]

        namespace = default_namespace(_geometry=geometry,
                                      _E=E,
                                      _PDOS=PDOS,
                                      _Erng=None,
                                      _PDOS_filter=_sum_filter,
                                      _data=[],
                                      _data_header=[])

        def ensure_E(func):
            """ This decorater ensures that E is the first element in the _data container """
            def assign_E(self, *args, **kwargs):
                ns = args[1]
                if len(ns._data) == 0:
                    # We immediately extract the energies
                    ns._data.append(ns._E[ns._Erng].flatten())
                    ns._data_header.append('Energy[eV]')
                return func(self, *args, **kwargs)

            return assign_E

        class ERange(argparse.Action):
            def __call__(self, parser, ns, value, option_string=None):
                E = ns._E
                Emap = strmap(float, value, E.min(), E.max())

                def Eindex(e):
                    return np.abs(E - e).argmin()

                # Convert to actual indices
                E = []
                for begin, end in Emap:
                    if begin is None and end is None:
                        ns._Erng = None
                        return
                    elif begin is None:
                        E.append(range(Eindex(end) + 1))
                    elif end is None:
                        E.append(range(Eindex(begin), len(E)))
                    else:
                        E.append(range(Eindex(begin), Eindex(end) + 1))
                # Issuing unique also sorts the entries
                ns._Erng = np.unique(arrayi(E).flatten())

        p.add_argument(
            '--energy',
            '-E',
            action=ERange,
            help=
            """Denote the sub-section of energies that are extracted: "-1:0,1:2" [eV]
                       
                       This flag takes effect on all energy-resolved quantities and is reset whenever --plot or --out is called"""
        )

        if PDOS.ndim == 3:
            # Add a spin-action
            class Spin(argparse.Action):
                def __call__(self, parser, ns, value, option_string=None):
                    if value.lower() in ["up", "u"]:

                        def _filter(PDOS):
                            return PDOS[0]
                    elif value.lower() in ["down", "dn", "dw", "d"]:

                        def _filter(PDOS):
                            return PDOS[1]
                    elif value.lower() in ["sum", "+"]:

                        def _filter(PDOS):
                            return PDOS.sum(0)

                    ns._PDOS_filter = _filter

            p.add_argument(
                '--spin',
                '-S',
                action=Spin,
                nargs=1,
                help="Which spin-component to store, up/u, down/d or sum/+")

        elif PDOS.ndim == 3:
            # Add a spin-action
            class Spin(argparse.Action):
                def __call__(self, parser, ns, value, option_string=None):
                    value = value.lower()
                    if value in ["sum", "+"]:

                        def _filter(PDOS):
                            return PDOS[0]
                    else:
                        # the stuff must be a range of directions
                        # so simply put it in
                        idx = list(map(direction, value))

                        def _filter(PDOS):
                            return PDOS[idx].sum(0)

                    ns._PDOS_filter = _filter

            p.add_argument(
                '--spin',
                '-S',
                action=Spin,
                nargs=1,
                help=
                "Which spin-component to store, sum/+, x, y, z or a sum of either of the directions xy, zx etc."
            )

        def parse_atom_range(geom, value):
            value = ",".join(  # ensure only single commas (no space between them)
                "".join(  # ensure no empty whitespaces
                    ",".join(  # join different lines with a comma
                        value.splitlines()).split()).split(","))

            # Sadly many shell interpreters does not
            # allow simple [] because they are expansion tokens
            # in the shell.
            # We bypass this by allowing *, [, {
            # * will "only" fail if files are named accordingly, else
            # it will be passed as-is.
            #       {    [    *
            sep = ['c', 'b', '*']
            failed = True
            while failed and len(sep) > 0:
                try:
                    ranges = lstranges(
                        strmap(int, value, 0, len(geom), sep.pop()))
                    failed = False
                except:
                    pass
            if failed:
                print(value)
                raise ValueError("Could not parse the atomic/orbital ranges")

            # we have only a subset of the orbitals
            orbs = []
            no = 0
            for atoms in ranges:
                if isinstance(atoms, list):
                    # Get atoms and orbitals
                    ob = geom.a2o(atoms[0] - 1, True)
                    # We normalize for the total number of orbitals
                    # on the requested atoms.
                    # In this way the user can compare directly the DOS
                    # for same atoms with different sets of orbitals and the
                    # total will add up.
                    no += len(ob)
                    ob = ob[asarrayi(atoms[1]) - 1]
                else:
                    ob = geom.a2o(atoms - 1, True)
                    no += len(ob)
                orbs.append(ob)

            if len(orbs) == 0:
                print('Available atoms:')
                print(f'  1-{len(geometry)}')
                print('Input atoms:')
                print('  ', value)
                raise ValueError(
                    'Atomic/Orbital requests are not fully included in the device region.'
                )

            # Add one to make the c-index equivalent to the f-index
            return np.concatenate(orbs).flatten()

        # Try and add the atomic specification
        class AtomRange(argparse.Action):
            @collect_action
            @ensure_E
            def __call__(self, parser, ns, value, option_string=None):
                orbs = parse_atom_range(ns._geometry, value)
                ns._data.append(ns._PDOS_filter(ns._PDOS)[orbs].sum(0))
                ns._data_header.append(f"PDOS[1/eV]{value}")

        p.add_argument(
            '--atom',
            '-a',
            type=str,
            action=AtomRange,
            help=
            """Limit orbital resolved PDOS to a sub-set of atoms/orbitals: "1-2[3,4]" will yield the 1st and 2nd atom and their 3rd and fourth orbital. Multiple comma-separated specifications are allowed. Note that some shells does not allow [] as text-input (due to expansion), {, [ or * are allowed orbital delimiters. Each invocation will create a new column/line in output"""
        )

        class Out(argparse.Action):
            @run_actions
            def __call__(self, parser, ns, value, option_string=None):
                out = value[0]

                try:
                    # We figure out if the user wants to write
                    # to a geometry
                    obj = get_sile(out, mode='w')
                    if hasattr(obj, 'write_geometry'):
                        with obj as fh:
                            fh.write_geometry(ns._geometry)
                        return
                    raise NotImplementedError
                except:
                    pass

                if len(ns._data) == 0:
                    orbs = parse_atom_range(ns._geometry, f"1-{len(geometry)}")
                    ns._data.append(ns._E)
                    ns._data.append(ns._PDOS_filter(ns._PDOS)[orbs].sum(0))
                    ns._data_header.append("DOS[1/eV]")

                from sisl.io import tableSile
                tableSile(out, mode='w').write(*ns._data,
                                               comment=comment,
                                               header=ns._data_header)
                # Clean all data
                ns._data = []
                ns._data_header = []
                ns._PDOS_filter = _sum_filter
                ns._Erng = None

        p.add_argument(
            '--out',
            '-o',
            nargs=1,
            action=Out,
            help=
            'Store currently collected PDOS (at its current invocation) to the out file.'
        )

        class Plot(argparse.Action):
            @run_actions
            def __call__(self, parser, ns, value, option_string=None):

                if len(ns._data) == 0:
                    orbs = parse_atom_range(ns._geometry, f"1-{len(geometry)}")
                    ns._data.append(ns._E)
                    ns._data.append(ns._PDOS_filter(ns._PDOS)[orbs].sum(0))
                    ns._data_header.append("DOS[1/eV]")

                from matplotlib import pyplot as plt
                plt.figure()

                def _get_header(header):
                    header = header.split(']', 1)[1]
                    if len(header) == 0:
                        return "DOS"
                    return header

                for i in range(1, len(ns._data)):
                    plt.plot(ns._data[0],
                             ns._data[i],
                             label=_get_header(ns._data_header[i]))

                plt.ylabel('DOS [1/eV]')
                if 'unknown' in comment:
                    plt.xlabel('E [eV]')
                else:
                    plt.xlabel('E - E_F [eV]')

                plt.legend(loc=8, ncol=3, bbox_to_anchor=(0.5, 1.0))
                if value is None:
                    plt.show()
                else:
                    plt.savefig(value)

                # Clean all data
                ns._data = []
                ns._data_header = []
                ns._PDOS_filter = _sum_filter
                ns._Erng = None

        p.add_argument(
            '--plot',
            '-p',
            action=Plot,
            nargs='?',
            metavar='FILE',
            help=
            'Plot the currently collected information (at its current invocation).'
        )

        return p, namespace
Exemplo n.º 3
0
Arquivo: pdos.py Projeto: juijan/sisl
    def ArgumentParser(self, p=None, *args, **kwargs):
        """ Returns the arguments that is available for this Sile """

        # We limit the import to occur here
        import argparse
        import warnings

        comment = 'Fermi-level shifted to 0'
        with warnings.catch_warnings(record=True) as w:
            # Cause all warnings to always be triggered.
            warnings.simplefilter("always")
            geometry, E, PDOS = self.read_data()

            if len(w) > 0:
                if issubclass(w[-1].category, SislWarning):
                    comment = 'Fermi-level unknown'

        def norm(geom, orbitals=None, norm='none'):
            r""" Normalization factor depending on the input

            The normalization can be performed in one of the below methods.
            In the following :math:`N` refers to the normalization constant
            that is to be used (i.e. the divisor):

            ``'none'``
               :math:`N=1`
            ``'all'``
               :math:`N` equals the number of orbitals in the total geometry
            ``'atom'``
               :math:`N` equals the total number of orbitals in the selected
               atoms. If `orbitals` is an argument a conversion of `orbitals` to the equivalent
               unique atoms is performed, and subsequently the total number of orbitals on the
               atoms is used. This makes it possible to compare the fraction of orbital DOS easier.
            ``'orbital'``
               :math:`N` is the sum of selected orbitals, if `atoms` is specified, this
               is equivalent to the 'atom' option.

            Parameters
            ----------
            orbitals : array_like of int or bool, optional
               only return for a given set of orbitals (default to all)
            norm : {'none', 'atom', 'orbital', 'all'}
               how the normalization of the summed DOS is performed (see `norm` routine)
            """
            # Cast to lower
            norm = norm.lower()
            if norm == 'none':
                NORM = 1
            elif norm in ['all', 'atom', 'orbital']:
                NORM = geom.no
            else:
                raise ValueError(
                    f"norm error on norm keyword in when requesting normalization!"
                )

            # If the user requests all orbitals
            if orbitals is None:
                return NORM

            # Now figure out what to do
            # Get pivoting indices to average over
            if norm == 'orbital':
                NORM = len(orbitals)
            elif norm == 'atom':
                a = np.unique(geom.o2a(orbitals))
                # Now sum the orbitals per atom
                NORM = geom.orbitals[a].sum()
            return NORM

        def _sum_filter(PDOS):
            """ Default sum is the total DOS, no projection on directions """
            if PDOS.ndim == 2:
                # non-polarized
                return PDOS
            elif PDOS.shape[0] == 2:
                # polarized
                return PDOS.sum(0)
            return PDOS[0]

        namespace = default_namespace(
            _geometry=geometry,
            _E=E,
            _PDOS=PDOS,
            # The energy range of all data
            _Erng=None,
            _norm="none",
            _PDOS_filter_name='total',
            _PDOS_filter=_sum_filter,
            _data=[],
            _data_description=[],
            _data_header=[])

        def ensure_E(func):
            """ This decorater ensures that E is the first element in the _data container """
            def assign_E(self, *args, **kwargs):
                ns = args[1]
                if len(ns._data) == 0:
                    # We immediately extract the energies
                    ns._data.append(ns._E[ns._Erng].flatten())
                    ns._data_header.append('Energy[eV]')
                return func(self, *args, **kwargs)

            return assign_E

        class ERange(argparse.Action):
            def __call__(self, parser, ns, value, option_string=None):
                E = ns._E
                Emap = strmap(float, value, E.min(), E.max())

                def Eindex(e):
                    return np.abs(E - e).argmin()

                # Convert to actual indices
                E = []
                for begin, end in Emap:
                    if begin is None and end is None:
                        ns._Erng = None
                        return
                    elif begin is None:
                        E.append(range(Eindex(end) + 1))
                    elif end is None:
                        E.append(range(Eindex(begin), len(E)))
                    else:
                        E.append(range(Eindex(begin), Eindex(end) + 1))
                # Issuing unique also sorts the entries
                ns._Erng = np.unique(arrayi(E).flatten())

        p.add_argument(
            '--energy',
            '-E',
            action=ERange,
            help=
            """Denote the sub-section of energies that are extracted: "-1:0,1:2" [eV]
                       
                       This flag takes effect on all energy-resolved quantities and is reset whenever --plot or --out is called"""
        )

        # The normalization method
        class NormAction(argparse.Action):
            @collect_action
            def __call__(self, parser, ns, value, option_string=None):
                ns._norm = value

        p.add_argument(
            '--norm',
            '-N',
            action=NormAction,
            default='atom',
            choices=['none', 'atom', 'orbital', 'all'],
            help=
            """Specify the normalization method; "none") no normalization, "atom") total orbitals in selected atoms,
                       "orbital") selected orbitals or "all") all orbitals.

                       Will only take effect on subsequent --atom ranges.

                       This flag is reset whenever --plot or --out is called"""
        )

        if PDOS.ndim == 2:
            # no spin is possible
            pass
        elif PDOS.shape[0] == 2:
            # Add a spin-action
            class Spin(argparse.Action):
                @collect_action
                def __call__(self, parser, ns, value, option_string=None):
                    value = value[0].lower()
                    if value in ("up", "u"):
                        name = "up"

                        def _filter(PDOS):
                            return PDOS[0]
                    elif value in ("down", "dn", "dw", "d"):
                        name = "down"

                        def _filter(PDOS):
                            return PDOS[1]
                    elif value in ("sum", "+", "total"):
                        name = "total"

                        def _filter(PDOS):
                            return PDOS.sum(0)
                    else:
                        raise ValueError(
                            f"Wrong argument for --spin [up, down, sum], found {value}"
                        )
                    ns._PDOS_filter_name = name
                    ns._PDOS_filter = _filter

            p.add_argument(
                '--spin',
                '-S',
                action=Spin,
                nargs=1,
                help=
                "Which spin-component to store, up/u, down/d or sum/+/total")

        elif PDOS.shape[0] == 4:
            # Add a spin-action
            class Spin(argparse.Action):
                @collect_action
                def __call__(self, parser, ns, value, option_string=None):
                    value = value[0].lower()
                    if value in ("sum", "+", "total"):
                        name = "total"

                        def _filter(PDOS):
                            return PDOS[0]
                    else:
                        # the stuff must be a range of directions
                        # so simply put it in
                        idx = list(map(direction, value))
                        name = value

                        def _filter(PDOS):
                            return PDOS[idx].sum(0)

                    ns._PDOS_filter_name = name
                    ns._PDOS_filter = _filter

            p.add_argument(
                '--spin',
                '-S',
                action=Spin,
                nargs=1,
                help=
                "Which spin-component to store, sum/+/total, x, y, z or a sum of either of the directions xy, zx etc."
            )

        def parse_atom_range(geom, value):
            if value.lower() in ("all", ":"):
                return np.arange(geom.no), "all"

            value = ",".join(  # ensure only single commas (no space between them)
                "".join(  # ensure no empty whitespaces
                    ",".join(  # join different lines with a comma
                        value.splitlines()).split()).split(","))

            # Sadly many shell interpreters does not
            # allow simple [] because they are expansion tokens
            # in the shell.
            # We bypass this by allowing *, [, {
            # * will "only" fail if files are named accordingly, else
            # it will be passed as-is.
            #       {    [    *
            sep = ['c', 'b', '*']
            failed = True
            while failed and len(sep) > 0:
                try:
                    ranges = lstranges(
                        strmap(int, value, 0, len(geom), sep.pop()))
                    failed = False
                except Exception:
                    pass
            if failed:
                print(value)
                raise ValueError("Could not parse the atomic/orbital ranges")

            # we have only a subset of the orbitals
            orbs = []
            no = 0
            for atoms in ranges:
                if isinstance(atoms, list):
                    # Get atoms and orbitals
                    ob = geom.a2o(atoms[0] - 1, True)
                    # We normalize for the total number of orbitals
                    # on the requested atoms.
                    # In this way the user can compare directly the DOS
                    # for same atoms with different sets of orbitals and the
                    # total will add up.
                    no += len(ob)
                    ob = ob[asarrayi(atoms[1]) - 1]
                else:
                    ob = geom.a2o(atoms - 1, True)
                    no += len(ob)
                orbs.append(ob)

            if len(orbs) == 0:
                print('Available atoms:')
                print(f'  1-{len(geometry)}')
                print('Input atoms:')
                print('  ', value)
                raise ValueError(
                    'Atomic/Orbital requests are not fully included in the device region.'
                )

            # Add one to make the c-index equivalent to the f-index
            return np.concatenate(orbs).flatten(), value

        # Try and add the atomic specification
        class AtomRange(argparse.Action):
            @collect_action
            @ensure_E
            def __call__(self, parser, ns, value, option_string=None):
                # get which orbitals to extract
                orbs, value = parse_atom_range(ns._geometry, value)
                # calculate the normalization
                scale = norm(ns._geometry, orbs, ns._norm)
                # Calculate PDOS on the selected atoms with the norm
                ns._data.append(ns._PDOS_filter(ns._PDOS)[orbs].sum(0) / scale)
                index = len(ns._data)
                if value == "all":
                    DOS = "DOS"
                else:
                    DOS = "PDOS"

                if ns._PDOS_filter_name is not None:
                    ns._data_header.append(
                        f"{DOS}[spin={ns._PDOS_filter_name}:{value}][1/eV]")
                    ns._data_description.append(
                        f"Column {index} is the sum of spin={ns._PDOS_filter_name} on atoms[orbs] {value} with normalization 1/{scale}"
                    )
                else:
                    ns._data_header.append(f"{DOS}[{value}][1/eV]")
                    ns._data_description.append(
                        f"Column {index} is the total PDOS on atoms[orbs] {value} with normalization 1/{scale}"
                    )

        p.add_argument(
            '--atom',
            '-a',
            type=str,
            action=AtomRange,
            help=
            """Limit orbital resolved PDOS to a sub-set of atoms/orbitals: "1-2[3,4]" will yield the 1st and 2nd atom and their 3rd and fourth orbital. Multiple comma-separated specifications are allowed. Note that some shells does not allow [] as text-input (due to expansion), {, [ or * are allowed orbital delimiters.

Multiple options will create a new column/line in output, the --norm and --E should be before any of these arguments"""
        )

        class Out(argparse.Action):
            @run_actions
            def __call__(self, parser, ns, value, option_string=None):
                out = value[0]

                try:
                    # We figure out if the user wants to write
                    # to a geometry
                    obj = get_sile(out, mode='w')
                    if hasattr(obj, 'write_geometry'):
                        with obj as fh:
                            fh.write_geometry(ns._geometry)
                        return
                    raise NotImplementedError
                except Exception:
                    pass

                if len(ns._data) == 0:
                    ns._data.append(ns._E)
                    ns._data_header.append('Energy[eV]')
                    ns._data.append(ns._PDOS_filter(ns._PDOS).sum(0))
                    if ns._PDOS_filter_name is not None:
                        ns._data_header.append(
                            f"DOS[spin={ns._PDOS_filter_name}][1/eV]")
                    else:
                        ns._data_header.append("DOS[1/eV]")

                from sisl.io import tableSile
                tableSile(out, mode='w').write(*ns._data,
                                               comment=[comment] +
                                               ns._data_description,
                                               header=ns._data_header)
                # Clean all data
                ns._norm = "none"
                ns._data = []
                ns._data_header = []
                ns._data_description = []

                ns._PDOS_filter_name = None
                ns._PDOS_filter = _sum_filter
                ns._Erng = None

        p.add_argument(
            '--out',
            '-o',
            nargs=1,
            action=Out,
            help=
            'Store currently collected PDOS (at its current invocation) to the out file.'
        )

        class Plot(argparse.Action):
            @run_actions
            def __call__(self, parser, ns, value, option_string=None):

                if len(ns._data) == 0:
                    ns._data.append(ns._E)
                    ns._data_header.append('Energy[eV]')
                    ns._data.append(ns._PDOS_filter(ns._PDOS).sum(0))
                    if ns._PDOS_filter_name is not None:
                        ns._data_header.append(
                            f"DOS[spin={ns._PDOS_filter_name}][1/eV]")
                    else:
                        ns._data_header.append("DOS[1/eV]")

                from matplotlib import pyplot as plt
                plt.figure()

                def _get_header(header):
                    header = (header.replace("PDOS",
                                             "").replace("DOS", "").replace(
                                                 "[1/eV]", ""))
                    if len(header) == 0:
                        return "Total"
                    if header.startswith("["):
                        header = header[1:]
                    if header.endswith("]"):
                        header = header[:-1]
                    return header

                kwargs = {}
                if len(ns._data) > 2:
                    kwargs['alpha'] = 0.6
                for i in range(1, len(ns._data)):
                    plt.plot(ns._data[0],
                             ns._data[i],
                             label=_get_header(ns._data_header[i]),
                             **kwargs)

                plt.ylabel('DOS [1/eV]')
                if 'unknown' in comment:
                    plt.xlabel('E [eV]')
                else:
                    plt.xlabel('E - E_F [eV]')

                plt.legend(loc=8, ncol=3, bbox_to_anchor=(0.5, 1.0))
                if value is None:
                    plt.show()
                else:
                    plt.savefig(value)

                # Clean all data
                ns._norm = "none"
                ns._data = []
                ns._data_header = []
                ns._data_description = []

                ns._PDOS_filter_name = None
                ns._PDOS_filter = _sum_filter
                ns._Erng = None

        p.add_argument(
            '--plot',
            '-p',
            action=Plot,
            nargs='?',
            metavar='FILE',
            help=
            'Plot the currently collected information (at its current invocation).'
        )

        return p, namespace