Source code for gpaw.lcaotddft.magneticmomentwriter

import json
import re
from typing import Tuple

import numpy as np
from ase import Atoms
from ase.units import Bohr
from ase.utils import IOContext
from gpaw.fd_operators import Gradient
from gpaw.lcaotddft.densitymatrix import DensityMatrix
from import TDDFTObserver
from gpaw.typing import Vector
from import coordinates

def calculate_magnetic_moment_on_grid(wfs, grad_v, r_vG, dM_vaii, *,
    """Calculate magnetic moment on grid.

        Wave functions object
        List of gradient operators
        Grid point coordinates
        Atomic PAW corrections for magnetic moment
        If true, do not add atomic corrections

    Magnetic moment vector
    gd =
    mode = wfs.mode
    bd =
    kpt_u = wfs.kpt_u

    rxnabla_v = np.zeros(3, dtype=complex)
    if mode == 'lcao':
        psit_G = gd.empty(dtype=complex)
    nabla_psit_vG = gd.empty(3, dtype=complex)
    for kpt in kpt_u:
        for n, f in enumerate(kpt.f_n):
            if mode == 'lcao':
                psit_G[:] = 0.0
                wfs.basis_functions.lcao_to_grid(kpt.C_nM[n], psit_G, kpt.q)
                psit_G = kpt.psit_nG[n]

            for v in range(3):
                grad_v[v].apply(psit_G, nabla_psit_vG[v], kpt.phase_cd)

            # rxnabla   = <psi1| r x nabla |psi2>
            # rxnabla_x = <psi1| r_y nabla_z - r_z nabla_y |psi2>
            # rxnabla_y = <psi1| r_z nabla_x - r_x nabla_z |psi2>
            # rxnabla_z = <psi1| r_x nabla_y - r_y nabla_x |psi2>
            for v in range(3):
                v1 = (v + 1) % 3
                v2 = (v + 2) % 3
                rnabla_psit_G = (r_vG[v1] * nabla_psit_vG[v2]
                                 - r_vG[v2] * nabla_psit_vG[v1])
                rxnabla_v[v] += f * gd.integrate(psit_G.conj() * rnabla_psit_G)

    if not only_pseudo:
        paw_rxnabla_v = np.zeros(3, dtype=complex)
        for kpt in kpt_u:
            for v in range(3):
                for a, P_ni in kpt.P_ani.items():
                    paw_rxnabla_v[v] += np.einsum('n,ni,ij,nj',
                                                  kpt.f_n, P_ni.conj(),
                                                  dM_vaii[v][a], P_ni,
        rxnabla_v += paw_rxnabla_v

    return -0.5 * rxnabla_v.imag

def calculate_magnetic_moment_atomic_corrections(R_av, setups, partition):
    """Calculate atomic PAW augmentation corrections for magnetic moment.

        Atom positions
        PAW setups object
        Atom partition object

    Atomic correction matrices
    # augmentation contributions to magnetic moment
    # <psi1| r x nabla |psi2> = <psi1| (r - Ra + Ra) x nabla |psi2>
    #                         = <psi1| (r - Ra) x nabla |psi2>
    #                             + Ra x <psi1| nabla |psi2>

    def shape(a):
        ni = setups[a].ni
        return ni, ni

    dM_vaii = []
    for _ in range(3):
        dM_aii = partition.arraydict(shapes=shape, dtype=complex)

    for a in partition.my_indices:
        Ra_v = R_av[a]
        rxnabla_iiv = setups[a].rxnabla_iiv
        nabla_iiv = setups[a].nabla_iiv

        # rxnabla   = <psi1| (r - Ra) x nabla |psi2>
        # Rxnabla   = Ra x <psi1| nabla |psi2>
        # Rxnabla_x = Ra_y nabla_z - Ra_z nabla_y
        # Rxnabla_y = Ra_z nabla_x - Ra_x nabla_z
        # Rxnabla_z = Ra_x nabla_y - Ra_y nabla_x
        for v in range(3):
            v1 = (v + 1) % 3
            v2 = (v + 2) % 3
            Rxnabla_ii = (Ra_v[v1] * nabla_iiv[:, :, v2]
                          - Ra_v[v2] * nabla_iiv[:, :, v1])
            dM_vaii[v][a][:] = Rxnabla_ii + rxnabla_iiv[:, :, v]

    return dM_vaii

def calculate_magnetic_moment_matrix(kpt_u, bfs, correction, r_vG, dM_vaii, *,
    """Calculate magnetic moment matrix in LCAO basis.

        Basis functions object
        Correction object
        Grid point coordinates
        Atomic PAW corrections for magnetic moment
        If true, do not add PAW corrections

    Magnetic moment matrix
    Mstart = correction.Mstart
    Mstop = correction.Mstop
    mynao = Mstop - Mstart
    nao = bfs.Mmax

    assert bfs.Mstart == Mstart
    assert bfs.Mstop == Mstop

    M_vmM = np.zeros((3, mynao, nao), dtype=complex)
    rnabla_vmM = np.empty((3, mynao, nao), dtype=complex)

    for v in range(3):
        v1 = (v + 1) % 3
        v2 = (v + 2) % 3
        rnabla_vmM[:] = 0.0
        bfs.calculate_potential_matrix_derivative(r_vG[v], rnabla_vmM, 0)
        M_vmM[v1] += rnabla_vmM[v2]
        M_vmM[v2] -= rnabla_vmM[v1]

    if not only_pseudo:
        for kpt in kpt_u:
            assert kpt.k == 0
            for v in range(3):
                correction.calculate(kpt.q, dM_vaii[v], M_vmM[v],
                                     Mstart, Mstop)

    # The matrix should be real
    assert np.max(np.absolute(M_vmM.imag)) == 0.0
    M_vmM = M_vmM.real.copy()
    return -0.5 * M_vmM

def calculate_magnetic_moment_in_lcao(ksl, rho_mm, M_vmm):
    """Calculate magnetic moment in LCAO.

        Kohn-Sham Layouts object
        Density matrix in LCAO basis
        Magnetic moment matrix in LCAO basis

    Magnetic moment vector
    assert M_vmm.dtype == float
    mm_v = np.sum(rho_mm.imag * M_vmm, axis=(1, 2))
    if ksl.using_blacs:
    return mm_v

def get_origin_coordinates(atoms: Atoms,
                           origin: str,
                           origin_shift: Vector) -> np.ndarray:
    """Get origin coordinates.

        Atoms object
        See :class:`~gpaw.tddft.MagneticMomentWriter`
        See :class:`~gpaw.tddft.MagneticMomentWriter`

    Origin coordinates in atomic units
    if origin == 'COM':
        origin_v = atoms.get_center_of_mass()
    elif origin == 'COC':
        origin_v = 0.5 * atoms.get_cell().sum(0)
    elif origin == 'zero':
        origin_v = np.zeros(3, dtype=float)
        raise ValueError('unknown origin')
    origin_v += np.asarray(origin_shift, dtype=float)
    return origin_v / Bohr

def parse_header(line: str) -> Tuple[str, int, dict]:
    """Parse header line.

    Example header line (keyword arguments as json):

        NameOfWriter[version=1](**{"arg1": "abc", ...})

        Header line

        Keyword arguments

        Line cannot be parsed
    regexp = r"^(?P<name>\w+)\[version=(?P<ver>\d+)\]\(\*\*(?P<args>.*)\)$"
    m = re.match(regexp, line)
    if m is None:
        raise ValueError('unable parse header')
    name ='name')
    version = int('ver'))
        kwargs = json.loads('args'))
    except json.decoder.JSONDecodeError:
        raise ValueError('unable parse keyword arguments')
    return name, version, kwargs

[docs]class MagneticMomentWriter(TDDFTObserver): """Observer for writing time-dependent magnetic moment data. The data is written in atomic units. The observer attaches to the TDDFT calculator during creation. Parameters ---------- paw TDDFT calculator filename File for writing magnetic moment data origin Origin of the coordinate system used in calculation. Possible values: ``'COM'``: center of mass (default), ``'COC'``: center of cell, ``'zero'``: (0, 0, 0) origin_shift Vector in Å shifting the origin from the position defined by *origin*. dmat Density matrix object. Define this for LCAO calculations to avoid expensive recalculations of the density matrix. calculate_on_grid Parameter for testing. In LCAO mode, if true, calculation is performed on real-space grid. In fd mode, calculation is always performed on real-space grid and this parameter is neglected. only_pseudo Parameter for testing. If true, PAW corrections are neglected. interval Update interval. Value of 1 corresponds to evaluating and writing data after every propagation step. """ version = 5 def __init__(self, paw, filename: str, *, origin: str = None, origin_shift: Vector = None, dmat: DensityMatrix = None, calculate_on_grid: bool = None, only_pseudo: bool = None, interval: int = 1): TDDFTObserver.__init__(self, paw, interval) self.ioctx = IOContext() mode = paw.wfs.mode assert mode in ['fd', 'lcao'], f'unknown mode: {mode}' if paw.niter == 0: if origin is None: origin = 'COM' if origin_shift is None: origin_shift = [0., 0., 0.] if calculate_on_grid is None: calculate_on_grid = mode == 'fd' if only_pseudo is None: only_pseudo = False _kwargs = dict(origin=origin, origin_shift=origin_shift, calculate_on_grid=calculate_on_grid, only_pseudo=only_pseudo) # Initialize self.fd = self.ioctx.openfile(filename,, mode='w') self._write_header(paw, _kwargs) else: if origin is not None: raise ValueError('Do not set origin in restart') if origin_shift is not None: raise ValueError('Do not set origin_shift in restart') if calculate_on_grid is not None: raise ValueError('Do not set calculate_on_grid in restart') if only_pseudo is not None: raise ValueError('Do not set only_pseudo in restart') # Read and continue _kwargs = self._read_header(filename) origin = _kwargs['origin'] # type: ignore origin_shift = _kwargs['origin_shift'] # type: ignore calculate_on_grid = _kwargs['calculate_on_grid'] # type: ignore only_pseudo = _kwargs['only_pseudo'] # type: ignore self.fd = self.ioctx.openfile(filename,, mode='a') atoms = paw.atoms gd = self.timer = paw.timer assert isinstance(origin, str) assert isinstance(origin_shift, list) origin_v = get_origin_coordinates(atoms, origin, origin_shift) R_av = atoms.positions / Bohr - origin_v[np.newaxis, :] r_vG, _ = coordinates(gd, origin=origin_v) dM_vaii = calculate_magnetic_moment_atomic_corrections( R_av, paw.setups, paw.hamiltonian.dH_asp.partition) self.calculate_on_grid = calculate_on_grid if self.calculate_on_grid: self.only_pseudo = only_pseudo self.r_vG = r_vG self.dM_vaii = dM_vaii grad_v = [] for v in range(3): grad_v.append(Gradient(gd, v, dtype=complex, n=2)) self.grad_v = grad_v else: M_vmM = calculate_magnetic_moment_matrix( paw.wfs.kpt_u, paw.wfs.basis_functions, paw.wfs.atomic_correction, r_vG, dM_vaii, only_pseudo=only_pseudo) # TODO: All observers recalculate density matrix # unless dmat is given. # Calculator itself could have a density matrix object to avoid # this expensive recalculation in a clean way. if dmat is None: self.dmat = DensityMatrix(paw) else: self.dmat = dmat ksl = paw.wfs.ksl if ksl.using_blacs: self.M_vmm = ksl.distribute_overlap_matrix(M_vmM) else: gd.comm.sum(M_vmM) self.M_vmm = M_vmM def _write(self, line): self.fd.write(line) self.fd.flush() def _write_header(self, paw, kwargs): origin_v = get_origin_coordinates( paw.atoms, kwargs['origin'], kwargs['origin_shift']) lines = [f'{self.__class__.__name__}[version={self.version}]' f'(**{json.dumps(kwargs)})', 'origin_v = [%.6f, %.6f, %.6f] Å' % tuple(origin_v * Bohr)] self._write('# ' + '\n# '.join(lines) + '\n') self._write(f'# {"time":>15} {"mmx":>17} {"mmy":>22} {"mmz":>22}\n') def _read_header(self, filename): with open(filename, encoding='utf-8') as fd: line = fd.readline() try: name, version, kwargs = parse_header(line[2:]) except ValueError as e: raise ValueError(f'File {filename} cannot be parsed: {e}') if name != self.__class__.__name__: raise ValueError(f'File {filename} is not ' f'for {self.__class__.__name__}') if version != self.version: raise ValueError(f'File {filename} is not ' f'of version {self.version}') return kwargs def _write_init(self, paw): time = paw.time line = '# Start; Time = %.8lf\n' % time self._write(line) def _write_kick(self, paw): time = paw.time kick = paw.kick_strength line = '# Kick = [%22.12le, %22.12le, %22.12le]; ' % tuple(kick) line += 'Time = %.8lf\n' % time self._write(line) def _calculate_mm(self, paw): if self.calculate_on_grid: self.timer.start('Calculate magnetic moment on grid') mm_v = calculate_magnetic_moment_on_grid( paw.wfs, self.grad_v, self.r_vG, self.dM_vaii, only_pseudo=self.only_pseudo) self.timer.stop('Calculate magnetic moment on grid') else: self.timer.start('Calculate magnetic moment in LCAO') u = 0 rho_mm = self.dmat.get_density_matrix((paw.niter, paw.action))[u] mm_v = calculate_magnetic_moment_in_lcao( paw.wfs.ksl, rho_mm, self.M_vmm) self.timer.stop('Calculate magnetic moment in LCAO') assert mm_v.shape == (3,) assert mm_v.dtype == float return mm_v def _write_mm(self, paw): time = paw.time mm_v = self._calculate_mm(paw) line = ('%20.8lf %22.12le %22.12le %22.12le\n' % (time, mm_v[0], mm_v[1], mm_v[2])) self._write(line) def _update(self, paw): if paw.action == 'init': self._write_init(paw) elif paw.action == 'kick': self._write_kick(paw) self._write_mm(paw) def __del__(self): self.ioctx.close()