Source code for sisl.io.table

# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.
import re

import numpy as np

import sisl._array as _a
from sisl._internal import set_module

from .sile import Sile, add_sile, sile_fh_open, sile_raise_write

__all__ = ["tableSile", "TableSile"]


@set_module("sisl.io")
class tableSile(Sile):
    """ASCII tabular formatted data

    A generic table data which will easily accommodate the most common write-outs of data

    Examples
    --------

    >>> a = np.arange(6).reshape(3, 2)
    >>> tbl = tableSile('test.dat', 'w')
    >>> tbl.write_data(a)
    >>> print(''.join(open('test.dat').readlines())) # doctest: +NORMALIZE_WHITESPACE
    0.00000e+00	2.00000e+00	4.00000e+00
    1.00000e+00	3.00000e+00	5.00000e+00
    <BLANKLINE>

    >>> a = np.arange(6).reshape(3, 2)
    >>> tbl = tableSile('test.dat', 'w')
    >>> tbl.write_data(a, comment='Hello')
    >>> print(''.join(open('test.dat').readlines())) # doctest: +NORMALIZE_WHITESPACE
    # Hello
    0.00000e+00	2.00000e+00	4.00000e+00
    1.00000e+00	3.00000e+00	5.00000e+00
    <BLANKLINE>

    >>> a = np.arange(6).reshape(3, 2)
    >>> tbl = tableSile('test.dat', 'w')
    >>> tbl.write_data(a, header=['x', 'y'])
    >>> print(''.join(open('test.dat').readlines())) # doctest: +NORMALIZE_WHITESPACE
    #x	y
    0.00000e+00	2.00000e+00	4.00000e+00
    1.00000e+00	3.00000e+00	5.00000e+00
    <BLANKLINE>

    >>> a = np.arange(6).reshape(3, 2)
    >>> tbl = tableSile('test.dat', 'w')
    >>> tbl.write_data(a.T)
    >>> print(''.join(open('test.dat').readlines())) # doctest: +NORMALIZE_WHITESPACE
    0.00000e+00	1.00000e+00
    2.00000e+00	3.00000e+00
    4.00000e+00	5.00000e+00
    <BLANKLINE>

    >>> x = np.arange(4)
    >>> y = np.arange(8).reshape(2, 4) + 4
    >>> tbl = tableSile('test.dat', 'w')
    >>> tbl.write_data(x, y)
    >>> print(''.join(open('test.dat').readlines())) # doctest: +NORMALIZE_WHITESPACE
    0.00000e+00	4.00000e+00	8.00000e+00
    1.00000e+00	5.00000e+00	9.00000e+00
    2.00000e+00	6.00000e+00	1.00000e+01
    3.00000e+00	7.00000e+00	1.10000e+01
    <BLANKLINE>

    """

    def _setup(self, *args, **kwargs):
        """Setup the `tableSile` after initialization"""
        super()._setup(*args, **kwargs)
        self._comment = ["#"]

[docs] @sile_fh_open() def write_data(self, *args, **kwargs): """Write tabular data to the file with optional header. Parameters ---------- args : array_like or list of array_like the different columns in the tabular data. This may either be several 1D data, or 2D arrays. Internally the data is stacked via `numpy.vstack` and for any dimension higher than 2 it gets separated by two newline characters (like gnuplot acceptable data). fmt : str, optional The formatting string, defaults to ``.5e``. fmts : str, optional The formatting string (for all columns), defaults to ``fmt * len(args)``. `fmts` has precedence over `fmt`. newline : str, optional Defaults to ``\\n``. delimiter : str, optional Defaults to ``\\t``. comment : str or list of str, optional A pre-header text at the top of the file. This comment is automatically prepended a ``#`` at the start of new lines, for lists each item corresponds to a newline which is automatically appended. header : str or list of str, optional A header for the data (this will be prepended a comment if not present). footer : str or list of str, optional A footer written at the end of the file. Examples -------- >>> tbl = tableSile('test.dat', 'w') >>> tbl.write_data(range(2), range(1, 3), comment='A comment', header=['index', 'value']) >>> print(''.join(open('test.dat').readlines())) # doctest: +NORMALIZE_WHITESPACE # A comment #index value 0.00000e+00 1.00000e+00 1.00000e+00 2.00000e+00 <BLANKLINE> """ sile_raise_write(self) fmt = kwargs.get("fmt", ".5e") newline = kwargs.get("newline", "\n") delimiter = kwargs.get("delimiter", "\t") _com = self._comment[0] def comment_newline(line, prefix=""): """Converts a list of str arguments into nicely formatted commented and newlined output""" nonlocal _com line = map(lambda s: s.strip(), line.strip().split(newline)) # always append a newline line = ( newline.join( [s if s.startswith(_com) else f"{_com}{prefix}{s}" for s in line] ) + newline ) return line comment = kwargs.get("comment", None) if comment is None: comment = "" elif isinstance(comment, str): comment = comment_newline(comment, " ") else: comment = comment_newline(newline.join(comment), " ") header = kwargs.get("header", None) if header is None: header = "" elif isinstance(header, str): header = comment_newline(header) else: header = comment_newline(delimiter.join(header)) # Finalize output header = comment + header footer = kwargs.get("footer", None) if footer is None: footer = "" elif isinstance(footer, str): pass else: footer = newline.join(footer) # Now we are ready to write if len(header) > 0: self._write(header) # Create a unified data table # This is rather difficult because we have to guess the # input size if len(args) == 0: # This may be used in cases where one wishes to accummulate content # in a file. return elif len(args) == 1: if isinstance(args[0], np.ndarray): # Ensure that when we change the shape we are doing it on a view dat = args[0].view() else: # Probably a tuple/list passed dat = np.vstack(args[0]) else: dat = np.vstack(args) _fmt = "{:" + fmt + "}" # Reshape such that it becomes easy ndim = dat.ndim if ndim > 2: dat.shape = (-1, dat.shape[-2], dat.shape[-1]) elif ndim == 1: dat.shape = (1, -1) if ndim > 2: _fmt = kwargs.get( "fmts", (_fmt + delimiter) * (dat.shape[1] - 1) + _fmt + newline ) for i in range(dat.shape[0]): for j in range(dat.shape[2]): self._write(_fmt.format(*dat[i, :, j])) self._write(newline * 2) else: _fmt = kwargs.get( "fmts", (_fmt + delimiter) * (dat.shape[0] - 1) + _fmt + newline ) for i in range(dat.shape[1]): self._write(_fmt.format(*dat[:, i])) if len(footer) > 0: self._write(newline * 2 + footer)
[docs] @sile_fh_open() def read_data(self, *args, **kwargs): """Read tabular data from the file. Parameters ---------- columns : list of int, optional only return the indices of the columns that are provided delimiter : str, optional the delimiter used in the file, will automatically try to guess if not specified ret_comment : bool, optional also return the comments at the top of the file (if queried) ret_header : bool, optional also return the header information (if queried) comment : str, optional lines starting with this are discarded as comments """ # Override the comment in the file self._comment = [kwargs.get("comment", self._comment[0])] # Skip to next line comment = [] header = "" # Also read comments line = self.readline(True) while line.startswith(self._comment[0] + " "): comment.append(line) line = self.readline(True) if line.startswith(self._comment[0]): header = line line = self.readline() # Now we are ready to read the data dat = [[]] # First we need to figure out the separator: len_sep = 0 sep = kwargs.get("delimiter", "") if len(sep) == 0: for cur_sep in ["\t", " ", ","]: s = line.split(cur_sep) if len(s) > len_sep: len_sep = len(s) sep = cur_sep if len(sep) == 0: raise ValueError( self.__class__.__name__ + ".read_data could not determine " "column separator..." ) empty = re.compile(r"\s*\n") while len(line) > 0: # If we start a line by a comment, or a newline # then we have a new data set if empty.match(line) is not None: if len(dat[-1]) > 0: dat[-1] = _a.asarrayd(dat[-1]) dat.append([]) else: line = [l for l in line.split(sep) if len(l) > 0] dat[-1].append(_a.fromiterd(map(float, line))) line = self.readline() if len(dat[-1]) > 0: dat[-1] = _a.asarrayd(dat[-1]) # Ensure we have no false positives dat = _a.asarrayd([d for d in dat if len(d) > 0]) if dat.shape[0] == 1: s = list(dat.shape) s.pop(0) dat.shape = tuple(s) if dat.ndim == 2: if dat.shape[1] == 1: # surely a 1D data dat.shape = (-1,) # For 2D data we need to transpose because the data is # read row wise, but stored column wise if dat.ndim > 1: dat = np.swapaxes(dat, -2, -1) ret_comment = kwargs.get("ret_comment", False) ret_header = kwargs.get("ret_header", False) if ret_comment: if ret_header: return dat, comment, header return dat, comment elif ret_header: return dat, header return dat
# Specify the default write function _write_default = write_data TableSile = tableSile add_sile("table", tableSile, case=False, gzip=True) add_sile("dat", tableSile, case=False, gzip=True)