# coding: utf-8
"""Tools for writing Xcrysden files."""
from __future__ import annotations
import numpy as np
from pymatgen.core.units import Energy, EnergyArray
from abipy.tools.numtools import transpose_last3dims, add_periodic_replicas
__all__ = [
"xsf_write_structure_and_data_to_path",
"xsf_write_structure",
"xsf_write_data",
"bxsf_write",
]
[docs]
def xsf_write_structure(file, structures: list) -> None:
"""
Write the crystalline structure in the Xcrysden format (XSF)
Args:
file: file-like object.
structures: :class:`Structure` or list of :class:`Structure` objects.
"""
animation = True
if not isinstance(structures, (list, tuple)):
structures = [structures]
animation = False
fwrite = file.write
if animation:
# axsf file.
fwrite('ANIMSTEPS %s\n' % len(structures))
fwrite('CRYSTAL\n')
for n, struct in enumerate(structures):
cell = struct.lattice.matrix
fwrite('# Primitive lattice vectors in Angstrom\n')
fwrite('PRIMVEC %d\n' % (n + 1))
for i in range(3):
fwrite(' %.14f %.14f %.14f\n' % tuple(cell[i]))
cart_coords = struct.cart_coords
atomic_numbers = struct.atomic_numbers
# TODO
cart_forces = None
#if "cartesian_forces" in structure.site_properties:
#cart_forces = ArrayWithUnit().to("Ha ang^-1")
fwrite("# Cartesian coordinates in Angstrom.\n")
fwrite('PRIMCOORD %d\n' % (n + 1))
fwrite(' %d 1\n' % len(cart_coords))
for a in range(len(cart_coords)):
fwrite(' %2d' % atomic_numbers[a])
fwrite(' %20.14f %20.14f %20.14f' % tuple(cart_coords[a]))
if cart_forces is None:
fwrite('\n')
else:
fwrite(' %20.14f %20.14f %20.14f\n' % tuple(cart_forces[a]))
[docs]
def xsf_write_structure_and_data_to_path(filepath, structure, datar, **kwargs) -> None:
"""Simplified interface to write structure and data to filepath in XSF format."""
with open(filepath, mode="wt") as fh:
xsf_write_structure(fh, structure)
xsf_write_data(fh, structure, datar, **kwargs)
#def xsf_write_structure_and_multidata(filepath, structure, multi_datar, tags=None, **kwargs) -> None:
# """
# """
# multi_datar = np.array(multi_datar, ndmin=4)
# if tags is not None:
# if len(tags) != len(multi_datar):
# raise ValueError(f"tags and multi_datar should have same length but {len(tags)} != {len(multi_datar)}")
#
# with open(filepath, mode="wt") as fh:
# xsf_write_structure(fh, structure)
# for i, datar in enumerate(multi_datar):
# tag = tags[i] if tags is not None else f"data_{i}"
# xsf_write_data(fh, structure, multi_datar[i], tag=tag, **kwargs)
[docs]
def xsf_write_data(file, structure, data, add_replicas=True, cplx_mode=None,
idname="data", tag="_UNKNOWN") -> None:
#idname="data", tag="_grid") -> None:
"""
Write data in the Xcrysden format (XSF)
Args:
file: file-like object.
structure: :class:`Structure` object.
data: array-like object in C-order, i.e data[nx, ny, nz] or data[ngrids, nx, ny, nz]
add_replicas: If True, data is padded with redundant data points.
in order to have a periodic 3D array of shape: (nx+1, ny+1, nz+1).
cplx_mode: string defining the data to print when data is a complex array.
Possible choices are (case-insensitive):
- "re" for real part.
- "im" for imaginary part.
- "abs" for the absolute value
"""
fwrite = file.write
# Check this one
if add_replicas:
data = add_periodic_replicas(data)
if np.iscomplexobj(data):
if cplx_mode is None:
raise TypeError("cplx_mode must be specified when data is a complex array.")
cplx_mode = cplx_mode.lower()
if cplx_mode == "re":
data = data.real
elif cplx_mode == "im":
data = data.imag
elif cplx_mode == "abs":
data = np.abs(data)
else:
raise ValueError(f"Wrong value for {cplx_mode=}")
shape, ndim = data.shape, data.ndim
if ndim == 3:
ngrids = 1
data = np.asarray([data])
elif ndim == 4:
ngrids = shape[0]
else:
raise ValueError(f"{ndim=} is not supported")
# Xcrysden uses Fortran-order.
# Transpose (...,x,y,z) --> (...,z,y,x) to speed up the write below.
fdata = transpose_last3dims(data)
fgrid = fdata.shape[-3:]
cell = structure.lattice.matrix
origin = np.zeros(3)
fwrite('BEGIN_BLOCK_DATAGRID_3D\n')
fwrite(f' {idname}\n')
for dg in range(ngrids):
if ngrids != 1:
fwrite(f" BEGIN_DATAGRID_3D{tag}#{dg+1}" + "\n")
else:
fwrite(f" BEGIN_DATAGRID_3D{tag}" + "\n")
fwrite('%d %d %d\n' % shape[-3:])
fwrite('%f %f %f\n' % tuple(origin))
for i in range(3):
fwrite('%f %f %f\n' % tuple(cell[i]))
for z in range(fgrid[0]):
for y in range(fgrid[1]):
slice_x = fdata[dg,z,y]
fwrite(' '.join(['%f' % d for d in slice_x]))
fwrite('\n')
fwrite('\n')
fwrite(' END_DATAGRID_3D\n')
fwrite('END_BLOCK_DATAGRID_3D\n')
[docs]
def bxsf_write(file, structure, nsppol, nband, ngkpt, ucdata_sbk, fermie, unit="eV") -> None:
"""
Write band structure data in the Xcrysden format (XSF)
Args:
file: file-like object.
structure: :class:`Structure` object.
nsppol: Number of spins.
nband: Number of bands.
ngkpt: Number of divisions of the full k-mesh.
ucdata_sbk: Array [nsppol, nband, ngkpt[0], ngkpt[1], ngkpt[2]] with energies
in the unit cell mesh in unit `unit`.
fermie: Fermi energy.
unit=Unit of input `ucdata_sbk` and `fermie`. Energies will be converted to Hartree before writing.
.. note::
#. The k-points must span the reciprocal unit cell, not the Brillouin zone.
#. The mesh must be closed and centered on Gamma.
#. Energies are written in row-major (i.e. C) order.
# Energies are in Hartree.
See also http://www.xcrysden.org/doc/XSF.html
"""
# Xscryden uses Ha for energies.
ucdata_sbk = EnergyArray(ucdata_sbk, unit).to("Ha")
fermie = Energy(fermie, unit).to("Ha")
ucdata_sbk = np.reshape(ucdata_sbk, (nsppol, nband, np.prod(ngkpt)))
close_it = False
if not hasattr(file, "write"):
file = open(str(file), mode="wt")
close_it = True
fw = file.write
# Write the header.
fw('BEGIN_INFO\n')
fw('# Band-XCRYSDEN-Structure-File for Visualization of Fermi Surface generated by the ABINIT package\n')
fw('# NOTE: the first band is relative to spin-up electrons,\n')
fw('# the second band to spin-down electrons (if any) and so on ...\n#\n')
fw('# Launch as: xcrysden --bxsf\n#\n')
fw(' Fermi Energy: %f\n' % fermie)
fw('END_INFO\n\n')
fw('BEGIN_BLOCK_BANDGRID_3D\n')
fw(' band_energies\n')
fw(' BEGIN_BANDGRID_3D\n')
fw(str(nsppol * nband) + "\n") # Number of bands written.
fw("%d %d %d\n" % tuple(ngkpt)) # Number of division in the full BZ mesh.
fw("0 0 0\n") # NB: Unshifted meshes are not supported.
# Reciprocal lattice vectors in Ang^{-1}
gcell = structure.lattice.reciprocal_lattice.matrix
for i in range(3):
fw('%f %f %f\n' % tuple(gcell[i]))
# Write energies on the full mesh for all spins and bands.
idx = 0
for band in range(nband):
for spin in range(nsppol):
idx += 1
enes = ucdata_sbk[spin, band, :]
fw(" BAND: %d\n" % idx)
fw("\n".join("%.18e" % v for v in enes))
fw("\n")
fw(' END_BANDGRID_3D\n')
fw('END_BLOCK_BANDGRID_3D\n')
file.flush()
if close_it:
file.close()