Source code for abipy.tools.tensors

# coding: utf-8
"""
This modules provides subclasses of pymatgen tensor objects.
"""
import numpy as np
import pandas as pd

from pymatgen.core.tensors import Tensor, SquareTensor
from pymatgen.analysis.elasticity.elastic import ElasticTensor  # noqa: F401
from pymatgen.analysis.elasticity.stress import Stress as pmg_Stress
from pymatgen.analysis.piezo import PiezoTensor # noqa: F401
from abipy.iotools import ETSF_Reader


class _Tensor33(object):

    def _repr_html_(self):
        """Integration with jupyter notebooks."""
        return self.get_dataframe()._repr_html_()

    def get_dataframe(self, tol=1e-3, cmode=None):
        """
        Return |pandas-Dataframe| with tensor elements set to zero below `tol`.

        Args:
            cmode: "real" or "imag" to include only the real/imaginary part.
        """
        tensor = self.zeroed(tol=tol)
        if cmode == "real": tensor = tensor.real
        if cmode == "imag": tensor = tensor.imag

        return pd.DataFrame({"x": tensor[:,0], "y": tensor[:,1], "z": tensor[:,2]}, index=["x", "y", "z"])

    def get_voigt_dataframe(self, tol=1e-3):
        """
        Return |pandas-DataFrame| with Voigt indices as colums.
        Elements below tol are set to zero.

        Useful to analyze the converge of individual elements.
        """
        tensor = self.zeroed(tol=tol)
        columns = ["xx", "yy", "zz", "yz", "xz", "xy"]
        d = {k: v for k, v in zip(columns, tensor.voigt)}
        return pd.DataFrame(d, index=[0], columns=columns)


[docs]class Stress(pmg_Stress, _Tensor33): """ Stress tensor. rank2 symmetric tensor with shape [3, 3]. .. rubric:: Inheritance Diagram .. inheritance-diagram:: Stress """
[docs]class DielectricTensor(SquareTensor, _Tensor33): """ Subclass of |pmg-Tensor| describing a dielectric tensor. rank2 symmetric tensor with shape [3, 3]. .. rubric:: Inheritance Diagram .. inheritance-diagram:: DielectricTensor """
[docs] def reflectivity(self, n1=1, tol=1e-6): """ If the tensor is diagonal (with off diagonal elements smaller than tol) returns the three components of the reflectivity :math:`|n1 - n2| / | n1 + n2 |` """ d = np.diag(self) if np.max(np.abs(self - np.diag(d))) > tol: raise ValueError("The tensor is not diagonal.") n2 = np.sqrt(d) return np.abs((n1 - n2) / (n1 + n2)) ** 2
[docs]class ZstarTensor(SquareTensor, _Tensor33): """ Born effective charge tensor (for a single atom). .. rubric:: Inheritance Diagram .. inheritance-diagram:: ZstarTensor """
[docs]class NLOpticalSusceptibilityTensor(Tensor): """ Subclass of |pmg-Tensor| containing the non-linear optical susceptibility tensor. .. rubric:: Inheritance Diagram .. inheritance-diagram:: NLOpticalSusceptibilityTensor """
[docs] @classmethod def from_file(cls, filepath): """ Creates the tensor from an anaddb.nc netcdf file containing ``dchide``. This requires to run anaddb with ``tnlflag`` > 0 """ with ETSF_Reader(filepath) as reader: try: return cls(reader.read_value("dchide")) except Exception as exc: import traceback msg = traceback.format_exc() msg += ("Error while trying to read from file.\n" "Verify that nlflag > 0 in anaddb\n") raise ValueError(msg)