# 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/.
from __future__ import annotations
import numpy as np
from sisl import Atom, Geometry, Grid, Lattice, SislError
from sisl._internal import set_module
# Import sile objects
from sisl.io.sile import *
from sisl.messages import deprecate_argument
from sisl.unit import unit_convert
from ._help import header_to_dict
__all__ = ["cubeSile"]
@set_module("sisl.io")
class cubeSile(Sile):
"""CUBE file object
By default the cube file is written using Bohr units.
one can define the units by passing a respective unit argument.
Note that the grid data is assumed unit-less and thus no conversion
will be done for this data, only atomic coordinates and lattice vectors.
"""
[docs]
@sile_fh_open()
@deprecate_argument("sc", "lattice", "use lattice= instead of sc=", "0.15", "0.16")
def write_lattice(
self,
lattice: Lattice,
fmt: str = "15.10e",
size=None,
origin=None,
unit: str = "Bohr",
*args,
**kwargs,
):
"""Writes `Lattice` object attached to this grid
Parameters
----------
lattice :
lattice to be written
fmt :
floating point format for stored values
size : (3, ), optional
shape of the stored grid (``[1, 1, 1]``)
origin : (3, ), optional
origin of the cell (``[0, 0, 0]``)
unit:
what length unit should the cube file data be written in
"""
sile_raise_write(self)
Ang2unit = unit_convert("Ang", unit)
# Write header
self._write("\n")
self._write(f"sisl-version=1 unit={unit}\n")
if size is None:
size = np.ones([3], np.int32)
if origin is None:
origin = lattice.origin[:]
_fmt = "{:d} {:15.10e} {:15.10e} {:15.10e}\n"
# Add #-of atoms and origin
self._write(_fmt.format(1, *(origin * Ang2unit)))
# Write the cell and voxels
for ix in range(3):
dcell = lattice.cell[ix, :] / size[ix] * Ang2unit
self._write(_fmt.format(size[ix], *dcell))
self._write("1 0. 0. 0. 0.\n")
[docs]
@sile_fh_open()
def write_geometry(
self,
geometry: Geometry,
fmt: str = "15.10e",
size=None,
origin=None,
unit: str = "Bohr",
*args,
**kwargs,
):
"""Writes `Geometry` object attached to this grid
Parameters
----------
geometry :
geometry to be written
fmt :
floating point format for stored values
size : (3, ), optional
shape of the stored grid (``[1, 1, 1]``)
origin : (3, ), optional
origin of the cell (``[0, 0, 0]``)
unit:
what length unit should the cube file data be written in
"""
sile_raise_write(self)
Ang2unit = unit_convert("Ang", unit)
# Write header
self._write("\n")
self._write(f"sisl-version=1 unit={unit}\n")
if size is None:
size = np.ones([3], np.int32)
if origin is None:
origin = geometry.origin[:]
_fmt = "{:d} {:15.10e} {:15.10e} {:15.10e}\n"
valid_Z = (geometry.atoms.Z > 0).nonzero()[0]
geometry = geometry.sub(valid_Z)
# Add #-of atoms and origin
self._write(_fmt.format(len(geometry), *(origin * Ang2unit)))
# Write the cell and voxels
for ix in range(3):
dcell = geometry.cell[ix, :] / size[ix] * Ang2unit
self._write(_fmt.format(size[ix], *dcell))
tmp = " {:" + fmt + "}"
_fmt = "{:d} 0.0" + tmp + tmp + tmp + "\n"
for ia in geometry:
self._write(
_fmt.format(geometry.atoms[ia].Z, *geometry.xyz[ia, :] * Ang2unit)
)
[docs]
@sile_fh_open()
def write_grid(
self,
grid: Grid,
fmt: str = ".5e",
imag: bool = False,
unit: str = "Bohr",
*args,
**kwargs,
):
"""Write `Grid` to the contained file
Parameters
----------
grid :
the grid to be written in the CUBE file
fmt :
format used for precision output
imag :
write only imaginary part of the grid, default to only writing the
real part.
unit:
what length unit should the cube file data be written in.
The grid data is assumed to be unit-less, this unit only refers
to the lattice vectors and atomic coordinates.
buffersize : int, optional
size of the buffer while writing the data, (6144)
"""
# Check that we can write to the file
sile_raise_write(self)
if grid.geometry is None:
self.write_lattice(
grid.lattice, size=grid.shape, unit=unit, *args, **kwargs
)
else:
self.write_geometry(
grid.geometry, size=grid.shape, unit=unit, *args, **kwargs
)
buffersize = kwargs.get("buffersize", min(6144, grid.grid.size))
buffersize += buffersize % 6 # ensure multiple of 6
# A CUBE file contains grid-points aligned like this:
# for x
# for y
# for z
# write...
_fmt1 = "{:" + fmt + "} "
_fmt6 = (_fmt1 * 6)[:-1] + "\n"
__fmt = _fmt6 * (buffersize // 6)
if imag:
for z in np.nditer(
np.asarray(grid.grid.imag, order="C").reshape(-1),
flags=["external_loop", "buffered"],
op_flags=[["readonly"]],
order="C",
buffersize=buffersize,
):
if z.shape[0] != buffersize:
s = z.shape[0]
__fmt = _fmt6 * (s // 6) + _fmt1 * (s % 6) + "\n"
self._write(__fmt.format(*z.tolist()))
else:
for z in np.nditer(
np.asarray(grid.grid.real, order="C").reshape(-1),
flags=["external_loop", "buffered"],
op_flags=[["readonly"]],
order="C",
buffersize=buffersize,
):
if z.shape[0] != buffersize:
s = z.shape[0]
__fmt = _fmt6 * (s // 6) + _fmt1 * (s % 6) + "\n"
self._write(__fmt.format(*z.tolist()))
# Add a finishing line to ensure empty ending
self._write("\n")
def _r_header_dict(self):
"""Reads the header of the file"""
self.fh.seek(0)
self.readline()
header = header_to_dict(self.readline())
header["unit"] = unit_convert(header.get("unit", "Bohr"), "Ang")
return header
[docs]
def read_basis(self) -> Atoms:
"""Reads the `Atoms` object from the CUBE file"""
return self.read_geometry().atoms
[docs]
@sile_fh_open()
def read_lattice(self, ret_na: bool = False) -> Lattice:
"""Returns `Lattice` object from the CUBE file
Parameters
----------
ret_na : bool, optional
whether to also return the number of atoms in the geometry
Returns
-------
lattice: Lattice
the lattice object
na : int
number of atoms (only if `ret_na`)
"""
unit2Ang = self._r_header_dict()["unit"]
origin = self.readline().split() # origin
na = int(origin[0])
origin = np.fromiter(map(float, origin[1:]), np.float64)
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 * unit2Ang
origin = origin * unit2Ang
if ret_na:
return Lattice(cell, origin=origin), na
return Lattice(cell, origin=origin)
[docs]
@sile_fh_open()
def read_geometry(self) -> Geometry:
"""Returns `Geometry` object from the CUBE file"""
unit2Ang = self._r_header_dict()["unit"]
lattice, na = self.read_lattice(ret_na=True)
if na == 0:
return None
# Start reading the geometry
xyz = np.empty([na, 3], np.float64)
atom = []
for ia in range(na):
tmp = self.readline().split()
atom.append(Atom(int(tmp[0])))
xyz[ia, 0] = float(tmp[2])
xyz[ia, 1] = float(tmp[3])
xyz[ia, 2] = float(tmp[4])
return Geometry(xyz * unit2Ang, atom, lattice=lattice)
[docs]
@sile_fh_open()
def read_grid(self, imag=None) -> Grid:
"""Returns `Grid` object from the CUBE file
Parameters
----------
imag : str or Sile or Grid
the imaginary part of the grid. If the geometries does not match
an error will be raised.
"""
if not imag is None:
if not isinstance(imag, Grid):
imag = Grid.read(imag)
geom = self.read_geometry()
if geom is None:
self.fh.seek(0)
lattice = self.read_lattice()
else:
lattice = geom.lattice
# read headers (and seek to start)
self._r_header_dict()
na = int(self.readline().split()[0])
ngrid = [0] * 3
for i in (0, 1, 2):
tmp = self.readline().split()
ngrid[i] = int(tmp[0])
# Read past the atoms
for i in range(na):
self.readline()
if geom is None:
grid = Grid(ngrid, dtype=np.float64, lattice=lattice)
else:
grid = Grid(ngrid, dtype=np.float64, geometry=geom)
grid.grid.shape = (-1,)
# TODO check performance of this
# We are currently doing this to enable reading
# 1-column data and 6-column data.
lines = [item for sublist in self.fh.readlines() for item in sublist.split()]
grid.grid[:] = np.array(lines).astype(grid.dtype, copy=False)
grid.grid.shape = ngrid
if imag is None:
return grid
# We are expecting an imaginary part
if not grid.geometry.equal(imag.geometry):
raise SislError(
f"{self!s} and its imaginary part does not have the same "
"geometry. Hence a combined complex Grid cannot be formed."
)
if grid != imag:
raise SislError(
f"{self!s} and its imaginary part does not have the same "
"shape. Hence a combined complex Grid cannot be formed."
)
# Now we have a complex grid
grid.grid = grid.grid + 1j * imag.grid
return grid
add_sile("cube", cubeSile, case=False, gzip=True)