Python string_utils.str_aligned函数代码示例

def get_magnetizations(mydir, ion_list):
    data = []
    max_row = 0
    for (parent, subdirs, files) in os.walk(mydir):
        for f in files:
            if re.match("OUTCAR*", f):
                try:
                    row = []
                    fullpath = os.path.join(parent, f)
                    outcar = Outcar(fullpath)
                    mags = outcar.magnetization
                    mags = [m["tot"] for m in mags]
                    all_ions = xrange(len(mags))
                    row.append(fullpath.lstrip("./"))
                    if ion_list:
                        all_ions = ion_list
                    for ion in all_ions:
                        row.append(str(mags[ion]))
                    data.append(row)
                    if len(all_ions) > max_row:
                        max_row = len(all_ions)
                except:
                    pass

    for d in data:
        if len(d) < max_row + 1:
            d.extend([""] * (max_row + 1 - len(d)))
    headers = ["Filename"]
    for i in xrange(max_row):
        headers.append(str(i))
    print str_aligned(data, headers)

def get_energies(rootdir, reanalyze, verbose, pretty, detailed, sort):
    if verbose:
        FORMAT = "%(relativeCreated)d msecs : %(message)s"
        logging.basicConfig(level=logging.INFO, format=FORMAT)

    if not detailed:
        drone = SimpleVaspToComputedEntryDrone(inc_structure=True)
    else:
        drone = VaspToComputedEntryDrone(inc_structure=True,
                                         data=["filename",
                                               "initial_structure"])

    ncpus = multiprocessing.cpu_count()
    logging.info("Detected {} cpus".format(ncpus))
    queen = BorgQueen(drone, number_of_drones=ncpus)
    if os.path.exists(save_file) and not reanalyze:
        msg = "Using previously assimilated data from {}.".format(save_file) \
            + " Use -f to force re-analysis."
        queen.load_data(save_file)
    else:
        if ncpus > 1:
            queen.parallel_assimilate(rootdir)
        else:
            queen.serial_assimilate(rootdir)
        msg = "Analysis results saved to {} for faster ".format(save_file) + \
              "subsequent loading."
        queen.save_data(save_file)

    entries = queen.get_data()
    if sort == "energy_per_atom":
        entries = sorted(entries, key=lambda x: x.energy_per_atom)
    elif sort == "filename":
        entries = sorted(entries, key=lambda x: x.data["filename"])

    all_data = []
    for e in entries:
        if not detailed:
            delta_vol = "{:.2f}".format(e.data["delta_volume"] * 100)
        else:
            delta_vol = e.structure.volume / \
                e.data["initial_structure"].volume - 1
            delta_vol = "{:.2f}".format(delta_vol * 100)
        all_data.append((e.data["filename"].replace("./", ""),
                         re.sub("\s+", "", e.composition.formula),
                         "{:.5f}".format(e.energy),
                         "{:.5f}".format(e.energy_per_atom),
                         delta_vol))
    if len(all_data) > 0:
        headers = ("Directory", "Formula", "Energy", "E/Atom", "% vol chg")
        if pretty:
            from prettytable import PrettyTable
            t = PrettyTable(headers)
            t.set_field_align("Directory", "l")
            map(t.add_row, all_data)
            print t
        else:
            print str_aligned(all_data, headers)
        print msg
    else:
        print "No valid vasp run found."

    def test_str_aligned_delimited(self):
        data = [["a", "bb"], ["ccc", "dddd"]]
        ans = """  a     bb
ccc   dddd"""
        self.assertEqual(str_aligned(data), ans)
        self.assertEqual(str_aligned(data, header=["X", "Y"]),
                         '  X      Y\n----------\n  a     bb\nccc   dddd')
        self.assertEqual(str_delimited(data), 'a\tbb\nccc\tdddd')

    def get_string(self, sort_keys=True, pretty=True):
        """
        Returns a string representation of the Feff_tag file.  The reason why
        this method is different from the __str__ method is to provide options
        for pretty printing.

        Args:
            sort_keys: Set to True to sort the Feff parameters alphabetically.
                Defaults to False.
            pretty: Set to True for pretty aligned output, False for no.

        Returns:
            String representation of FeffTags.
        """
        keys = self.keys()
        if sort_keys:
            keys = sorted(keys)
        lines = []
        for k in keys:
            if isinstance(self[k], list):
                lines.append([k, " ".join([str(i) for i in self[k]])])
            else:
                lines.append([k, self[k]])

        if pretty:
            return str_aligned(lines)
        else:
            return str_delimited(lines, None, "  ")

    def get_string(self, pretty=False):
        """
        Returns a string representation of self. The reason why this
        method is different from the __str__ method is to provide options for pretty printing.

        Args:
            pretty:
                Set to True for pretty aligned output.
        """
        lines = []
        app = lines.append

        # Write the Abinit objects first.
        for obj in self.abiobjects:
            app([80*"#", ""])
            app(["#", "%s" % obj.__class__.__name__])
            app([80*"#", ""])
            for (k, v) in obj.to_abivars().items():
                app(self._format_kv(k, v))

        # Extra variables.
        if self.extra_abivars:
            app([80*"#", ""])
            app(["#", "Extra_Abivars"])
            app([80*"#", ""])
            for (k, v) in self.extra_abivars.items():
                app(self._format_kv(k, v))

        if pretty:
            return str_aligned(lines, header=None)
        else:
            return str_delimited(lines, header=None, delimiter=5*" ")

    def get_string(self, sort_keys=False, pretty=False):
        """
        Returns a string representation of the INCAR.  The reason why this
        method is different from the __str__ method is to provide options for
        pretty printing.

        Args:
            sort_keys:
                Set to True to sort the INCAR parameters alphabetically.
                Defaults to False.
            pretty:
                Set to True for pretty aligned output. Defaults to False.
        """
        keys = self.keys()
        if sort_keys:
            keys = sorted(keys)
        lines = []
        for k in keys:
            if k == "MAGMOM" and isinstance(self[k], list):
                value = []
                for m, g in itertools.groupby(self[k]):
                    value.append("{}*{}".format(len(tuple(g)), m))
                lines.append([k, " ".join(value)])
            elif isinstance(self[k], list):
                lines.append([k, " ".join([str(i) for i in self[k]])])
            else:
                lines.append([k, self[k]])

        if pretty:
            return str_aligned(lines)
        else:
            return str_delimited(lines, None, " = ")

def get_energies(rootdir, reanalyze):
    FORMAT = "%(relativeCreated)d msecs : %(message)s"
    logging.basicConfig(level=logging.INFO, format=FORMAT)
    drone = VaspToComputedEntryDrone(inc_structure=True, data=['filename'])
    ncpus = multiprocessing.cpu_count()
    logging.info('Detected {} cpus'.format(ncpus))
    queen = BorgQueen(drone, number_of_drones=ncpus)
    if os.path.exists('vasp_analyzer_data.gz') and not reanalyze:
        logging.info('Using previously assimilated data file vasp_analyzer_data.gz. Use -f to force re-analysis')
        queen.load_data('vasp_analyzer_data.gz')
    else:
        queen.parallel_assimilate(rootdir)
        queen.save_data('vasp_analyzer_data.gz')
    entries = queen.get_assimilated_data()
    entries = sorted(entries, key=lambda x:x.data['filename'])
    all_data = [(e.data['filename'].replace("./", ""), e.composition.formula, "{:.5f}".format(e.energy), "{:.5f}".format(e.energy_per_atom), "{:.2f}".format(e.structure.volume)) for e in entries]
    print str_aligned(all_data, ("Directory", "Formula", "Energy", "E/Atom", "Vol"))

def get_energies(rootdir, reanalyze, verbose, pretty):
    if verbose:
        FORMAT = "%(relativeCreated)d msecs : %(message)s"
        logging.basicConfig(level=logging.INFO, format=FORMAT)
    drone = GaussianToComputedEntryDrone(inc_structure=True,
                                         parameters=['filename'])
    ncpus = multiprocessing.cpu_count()
    logging.info('Detected {} cpus'.format(ncpus))
    queen = BorgQueen(drone, number_of_drones=ncpus)
    if os.path.exists(save_file) and not reanalyze:
        msg = 'Using previously assimilated data from {}. ' + \
              'Use -f to force re-analysis'.format(save_file)
        queen.load_data(save_file)
    else:
        queen.parallel_assimilate(rootdir)
        msg = 'Results saved to {} for faster reloading.'.format(save_file)
        queen.save_data(save_file)

    entries = queen.get_data()
    entries = sorted(entries, key=lambda x: x.parameters['filename'])
    all_data = [(e.parameters['filename'].replace("./", ""),
                 re.sub("\s+", "", e.composition.formula),
                 "{}".format(e.parameters['charge']),
                 "{}".format(e.parameters['spin_mult']),
                 "{:.5f}".format(e.energy), "{:.5f}".format(e.energy_per_atom),
                 ) for e in entries]
    headers = ("Directory", "Formula", "Charge", "Spin Mult.", "Energy",
               "E/Atom")
    if pretty:
        from prettytable import PrettyTable
        t = PrettyTable(headers)
        t.set_field_align("Directory", "l")
        map(t.add_row, all_data)
        print t
    else:
        print str_aligned(all_data, headers)
    print msg

    def get_string(self, radius=10.):
        """
        Returns a string representation of atomic shell coordinates to be used
        in the feff.inp file.

        Args:
            radius:
                Maximum atomic shell radius to include in atoms list

        Returns:
            String representation of Atomic Coordinate Shells.
        """
        nopts = len(self.struct.species)

        ptatoms = []
        for i in range(0, nopts):
            ptatoms.append(self.struct.species[i].symbol)

        index = ptatoms.index(self.central_atom)
        pt = self.struct.cart_coords[index]
        sphere = Structure.get_sites_in_sphere(self.struct, pt, radius)
        xshift = pt[0]
        yshift = pt[1]
        zshift = pt[2]
        end = len(sphere)
        row = []

        for i in range(end):
            atom = sphere[i][0]
            atm = re.sub(r"[^aA-zZ]+", "", atom.species_string)
            ipot = self.pot_dict[atm]
            x = atom.coords[0] - xshift
            y = atom.coords[1] - yshift
            z = atom.coords[2] - zshift
            distance = sphere[i][1]
            row.append(["{:f}".format(x), "{:f}".format(y), "{:f}".format(z),
                        ipot, atm, "{:f}".format(distance), i])

        row_sorted = sorted(row, key=itemgetter(5))
        row_sorted[0][3] = 0

        for i in range(end):
            row_sorted[i][6] = i

        row_sorted = str_aligned(row_sorted, ["*       x", "y", "z", "ipot",
                                              "Atom", "Distance", "Number"])
        atom_list = row_sorted.replace("--", "**")

        return ''.join(["ATOMS\n", atom_list, "\nEND\n"])

    def __str__(self):
        """
        Returns a string representation of potential parameters to be used in
        the feff.inp file,
        determined from structure object.

                The lines are arranged as follows:

          ipot   Z   element   lmax1   lmax2   stoichometry   spinph

        Returns:
            String representation of Atomic Coordinate Shells.
        """

        noelements = len(self.struct.composition.items())
        nopts = len(self.struct.species)

        ptatoms = []
        for i in range(0, nopts):
            ptatoms.append(self.struct.species[i].symbol)

        index = ptatoms.index(self.central_atom)
        center = self.struct.species[index]
        cs = center.symbol
        cz = center.Z
        ipotrow = [[0, cz, cs, -1, -1, .0001, 0]]
        for i in range(0, noelements):
            center = self.struct.composition.items()[i][0]
            cs = center.symbol
            cz = center.Z
            ipot = self.pot_dict[cs]
            stoic = self.struct.composition.items()[i][1]
            ipotrow.append([ipot, cz, cs, -1, -1, stoic, 0])
        ipot_sorted = sorted(ipotrow, key=itemgetter(0))
        ipotrow = str_aligned(ipot_sorted, ["*ipot", "Z", "tag", "lmax1",
                                            "lmax2", "xnatph(stoichometry)",
                                            "spinph"])
        ipotlist = ipotrow.replace("--", "**")
        ipotlist = ''.join(["POTENTIALS\n", ipotlist])

        return ipotlist

    def get_string(self, pretty=False):
        """
        Returns a string representation of self. The reason why this
        method is different from the __str__ method is to provide options for pretty printing.

        Args:
            pretty: Set to True for pretty aligned output.
        """
        lines = []
        app = lines.append

        # extra_abivars can contain variables that are already defined 
        # in the object. In this case, the value in extra_abivars is used
        # TODO: Should find a more elegant way to avoid collission between objects
        # and extra_abivars
        extra_abivars = self.extra_abivars.copy()

        # Write the Abinit objects first.
        for obj in self.abiobjects:
            #print(obj)
            app([80*"#", ""])
            app(["#", "%s" % obj.__class__.__name__])
            app([80*"#", ""])
            for (k, v) in obj.to_abivars().items():
                v = extra_abivars.pop(k, v)
                app(self._format_kv(k, v))

        # Extra variables.
        if self.extra_abivars:
            app([80*"#", ""])
            app(["#", "Extra_Abivars"])
            app([80*"#", ""])
            for (k, v) in extra_abivars.items():
                app(self._format_kv(k, v))

        #lines = self._cut_lines(lines)

        if pretty:
            return str_aligned(lines, header=None)
        else:
            return str_delimited(lines, header=None, delimiter=5*" ")

Last updated: Oct 26 2011''')
parser.add_argument('incar_file', metavar='filename', type=str, nargs=2,
                    help='files to process')

args = parser.parse_args()

filepath1 = args.incar_file[0]
filepath2 = args.incar_file[1]
incar1 = Incar.from_file(filepath1)
incar2 = Incar.from_file(filepath2)


def format_lists(v):
    if isinstance(v, (tuple, list)):
        return " ".join(["%d*%.2f" % (len(tuple(group)), i)
                         for (i, group) in itertools.groupby(v)])
    return v

d = incar1.diff(incar2)
output = [['SAME PARAMS', '', '']]
output.append(['---------------', '', ''])
output.extend([(k, format_lists(d['Same'][k]), format_lists(d['Same'][k]))
               for k in sorted(d['Same'].keys()) if k != "SYSTEM"])
output.append(['', '', ''])
output.append(['DIFFERENT PARAMS', '', ''])
output.append(['----------------', '', ''])
output.extend([(k, format_lists(d['Different'][k]['INCAR1']),
                format_lists(d['Different'][k]['INCAR2']))
               for k in sorted(d['Different'].keys()) if k != "SYSTEM"])
print str_aligned(output, ['', filepath1, filepath2])