<commit_before>from __future__ import print_function

import sys
import numpy as np

# Import sile objects
from sisl.io.sile import *

# Import the geometry object
from sisl import Geometry, Atom, SuperCell, Grid
from sisl.unit import unit_convert

__all__ = ['CUBESile']

if sys.version_info >= (3, 0):
    _w = ['wb', 'ab']
else:
    _w = ['w', 'a']

Ang2Bohr = unit_convert('Ang', 'Bohr')


class CUBESile(Sile):
    """ CUBE file object """

    @Sile_fh_open
    def write_geometry(self, geom, size=None,
            fmt='15.10e', origo=None,
            *args, **kwargs):
        """ Writes the `Geometry` object attached to this grid """
        sile_raise_write(self)

        # Write header
        self._write('\n')
        self._write('sisl --- CUBE file\n')

        if size is None:
            size = np.ones([3], np.int32)
        if origo is None:
            origo = np.zeros([3], np.float64)

        _fmt = '{:d} {:15.10e} {:15.10e} {:15.10e}\n'

        # Add #-of atoms and origo
        self._write(_fmt.format(len(geom), *origo * Ang2Bohr))

        # Write the cell and voxels
        dcell = np.empty([3, 3], np.float64)
        for ix in range(3):
            dcell[ix, :] = geom.cell[ix, :] / size[ix] * Ang2Bohr
        self._write(_fmt.format(size[0], *dcell[0, :]))
        self._write(_fmt.format(size[1], *dcell[1, :]))
        self._write(_fmt.format(size[2], *dcell[2, :]))

        tmp = ' {:' + fmt + '}'
        _fmt = '{:d} 0.0' + tmp + tmp + tmp + '\n'
        for ia in geom:
            self._write(_fmt.format(geom.atom[ia].Z, *geom.xyz[ia, :] * Ang2Bohr))

    @Sile_fh_open
    def write_grid(self, grid, fmt='%.5e', *args, **kwargs):
        """ Writes the geometry to the contained file """
        # Check that we can write to the file
        sile_raise_write(self, _w)

        # Write the geometry
        self.write_geometry(grid.geom, size=grid.grid.shape, *args, **kwargs)

        g_size = np.copy(grid.grid.shape)

        grid.grid.shape = (-1,)

        # Write the grid
        np.savetxt(self.fh, grid.grid[:], fmt)

        # Add a finishing line to ensure empty ending
        self._write('\n')

        grid.grid.shape = g_size

    @Sile_fh_open
    def read_supercell(self, na=False):
        """ Returns `SuperCell` object from the CUBE file

        If `na=True` it will return a tuple (na,SuperCell)
        """

        self.readline()  # header 1
        self.readline()  # header 2
        tmp = self.readline().split()  # origo
        na = int(tmp[0])

        cell = np.empty([3, 3], np.float64)
        for i in [0, 1, 2]:
            tmp = self.readline().split()
            s = int(tmp[0])
            tmp = tmp[1:]
            for j in [0, 1, 2]:
                cell[i, j] = float(tmp[j]) * s

        cell = cell / Ang2Bohr
        if na:
            return na, SuperCell(cell)
        return SuperCell(cell)<commit_msg>bug: again, trying to fix savetxt<commit_after>from __future__ import print_function

import numpy as np

# Import sile objects
from sisl.io.sile import *

# Import the geometry object
from sisl import Geometry, Atom, SuperCell, Grid
from sisl.unit import unit_convert

__all__ = ['CUBESile']

Ang2Bohr = unit_convert('Ang', 'Bohr')


class CUBESile(Sile):
    """ CUBE file object """

    def _open(self):
        if self.file.endswith('gz'):
            self.fh = gzip.open(self.file)
        else:
            # Default to buffer 1MB
            self.fh = open(self.file, self._mode, buffering=1024**2)
        self._line = 0

    @Sile_fh_open
    def write_geometry(self, geom, size=None,
            fmt='15.10e', origo=None,
            *args, **kwargs):
        """ Writes the `Geometry` object attached to this grid """
        sile_raise_write(self)

        # Write header
        self._write('\n')
        self._write('sisl --- CUBE file\n')

        if size is None:
            size = np.ones([3], np.int32)
        if origo is None:
            origo = np.zeros([3], np.float64)

        _fmt = '{:d} {:15.10e} {:15.10e} {:15.10e}\n'

        # Add #-of atoms and origo
        self._write(_fmt.format(len(geom), *origo * Ang2Bohr))

        # Write the cell and voxels
        dcell = np.empty([3, 3], np.float64)
        for ix in range(3):
            dcell[ix, :] = geom.cell[ix, :] / size[ix] * Ang2Bohr
        self._write(_fmt.format(size[0], *dcell[0, :]))
        self._write(_fmt.format(size[1], *dcell[1, :]))
        self._write(_fmt.format(size[2], *dcell[2, :]))

        tmp = ' {:' + fmt + '}'
        _fmt = '{:d} 0.0' + tmp + tmp + tmp + '\n'
        for ia in geom:
            self._write(_fmt.format(geom.atom[ia].Z, *geom.xyz[ia, :] * Ang2Bohr))

    @Sile_fh_open
    def write_grid(self, grid, fmt='%.5e', *args, **kwargs):
        """ Writes the geometry to the contained file """
        # Check that we can write to the file
        sile_raise_write(self)

        # Write the geometry
        self.write_geometry(grid.geom, size=grid.grid.shape, *args, **kwargs)

        # A CUBE file contains grid-points aligned like this:
        # for x
        #   for y
        #     for z
        #       write...
        cast = np.vectorize((fmt + '\n').__mod__)
        for z in np.nditer(np.asarray(grid.grid, order='C'), flags=['external_loop', 'buffered'],
                           op_flags=[['readonly']], order='C', buffersize=grid.shape[2]):
            self._write(cast(z))

        # Add a finishing line to ensure empty ending
        self._write('\n')

    @Sile_fh_open
    def read_supercell(self, na=False):
        """ Returns `SuperCell` object from the CUBE file

        If `na=True` it will return a tuple (na,SuperCell)
        """

        self.readline()  # header 1
        self.readline()  # header 2
        tmp = self.readline().split()  # origo
        na = int(tmp[0])

        cell = np.empty([3, 3], np.float64)
        for i in [0, 1, 2]:
            tmp = self.readline().split()
            s = int(tmp[0])
            tmp = tmp[1:]
            for j in [0, 1, 2]:
                cell[i, j] = float(tmp[j]) * s

        cell = cell / Ang2Bohr
        if na:
            return na, SuperCell(cell)
        return SuperCell(cell)
{'fraction_non_alphanumeric': 0.14695918696563962, 'fraction_numerical': 0.019841909985481528, 'mean_word_length': 2.011170471102477, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3573727', 'n_tokens_mistral': 2225, 'n_tokens_neox': 2095, 'n_words': 544}