Source code for rhodent.calculators.energy

from __future__ import annotations

from collections.abc import Sequence
import numpy as np
from typing import Any, Generator

from .base import BaseObservableCalculator
from ..density_matrices.base import WorkMetadata
from ..density_matrices.time import ConvolutionDensityMatrixMetadata
from ..density_matrices.distributed.pulse import PulsePerturbation
from ..utils import ResultKeys, Result, broaden_n2e


[docs] class EnergyCalculator(BaseObservableCalculator): r""" Calculate energy contributions in the time domain. The total energy can be written .. math:: E_\text{tot}(t) = E^{(0)}_\text{tot} + \sum_{ia}^\text{eh} E_{ia}(t) + E_\text{pulse}(t). The contributions to the total energy are .. math:: E_{ia} = \frac{1}{2} \left[ p_{ia}\dot{q}_{ia} - q_{ia} \dot{p}_{ia} - v_{ia} q_{ia} \right], the contributions to the Hartree energy are .. math:: E_{ia}^\text{c} = -\frac{1}{2} \left[ \omega_{ia} q_{ia}^2 - q_{ia} \dot{p}_{ia} - v_{ia} q_{ia} \right], and the rate of energy change is .. math:: \dot{E}_{ia} = \frac{1}{2} \left[ p_{ia}\ddot{q}_{ia} - q_{ia} \ddot{p}_{ia} - v_{ia} \dot{q}_{ia} - \dot{v}_{ia} q_{ia} \right]. The matrix element :math:`v_{ia}` can be computed from the dipole matrix element .. math:: \boldsymbol{\mu}_{ia} = \int \psi^{(0)}_i(\boldsymbol{r}) \boldsymbol{r} \psi^{(0)}_a(\boldsymbol{r}) \mathrm{d}\boldsymbol{r} projected on the direction of the perturbation :math:`\hat{\boldsymbol{e}}`, the occupation number difference :math:`f_{ia}` and the pulse amplitude :math:`v_\text{pulse}(t)` .. math:: v_{ia} = \sqrt{2 f_{ia}} \boldsymbol{\mu}_{ia} \cdot\hat{\boldsymbol{e}} v_\text{pulse}. Parameters ---------- response Response object. voronoi Voronoi weights object. energies_occ Energy grid in units of eV for occupied levels (holes). energies_unocc Energy grid in units of eV for unoccupied levels (electrons). sigma Gaussian broadening width for energy grid in units of eV. times Compute energies in the time domain, for these times (or as close to them as possible). In units of as. pulses Compute energies in the time domain, in response to these pulses. If none, then no pulse convolution is performed. """
[docs] def get_result_keys(self, yield_total_E_ia: bool = False, yield_proj_E_ia: bool = False, yield_total_E_ou: bool = False, yield_total_dists: bool = False, direction: int | Sequence[int] = 2, ) -> ResultKeys: r""" Get the keys that each result will contain, and dimensions thereof. Parameters ---------- yield_total_E_ia The results should include the contributions in the electron-hole basis to the total energy :math:`E_{ia}` and Coulomb energy :math:`E_{ia}^\text{c}` yield_proj_E_ia The results should include the contributions in the electron-hole basis to the total energy projected on the occupied and unoccupied Voronoi weights :math:`E_{ia} w_i` and :math:`E_{ia} w_a`. yield_total_E_ou The results should include the contributions the total energy broadened on the occupied and unoccupied energy grids :math:`\sum_{ia} E_{ia}\delta(\varepsilon_\text{occ}-\varepsilon_{i}) \delta(\varepsilon_\text{unocc}-\varepsilon_{a})` and yield_total_dists The results should include the contributions the total energy and Coulomb energy broadened by electronic transition energy onto the unoccupied energies grid :math:`\sum_{ia} E_{ia} \delta(\varepsilon-\omega_{ia})` and :math:`\sum_{ia} E_{ia}^\text{C} \delta(\varepsilon-\omega_{ia})` direction Direction :math:`\hat{\boldsymbol{e}}` of the polarization of the pulse. Integer 0, 1 or 2 to specify x, y or z, or the direction vector specified as a list of three values. Default: polarization along z. """ nI = self.nproj imin, imax, amin, amax = self.ksd.ialims() ni = int(imax - imin + 1) na = int(amax - amin + 1) no = len(self.energies_occ) nu = len(self.energies_unocc) assert direction in [0, 1, 2] or (isinstance(direction, Sequence) and len(direction) == 3) resultkeys = ResultKeys('dm', 'total', 'total_coulomb', 'Epulse') if yield_total_E_ia: resultkeys.add_key('E_ia', (ni, na)) resultkeys.add_key('Ec_ia', (ni, na)) if yield_total_dists: resultkeys.add_key('E_transition_u', nu) resultkeys.add_key('Ec_transition_u', nu) if yield_total_E_ou: resultkeys.add_key('E_ou', (no, nu)) resultkeys.add_key('Ec_ou', (no, nu)) resultkeys.add_key('total_proj_II', (nI, nI)) resultkeys.add_key('total_coulomb_proj_II', (nI, nI)) if yield_proj_E_ia: resultkeys.add_key('E_occ_proj_Iia', (nI, ni, na)) resultkeys.add_key('E_unocc_proj_Iia', (nI, ni, na)) return resultkeys
[docs] def icalculate(self, yield_total_E_ia: bool = False, yield_proj_E_ia: bool = False, yield_total_E_ou: bool = False, yield_total_dists: bool = False, direction: int | Sequence[int] = 2, ) -> Generator[tuple[WorkMetadata, Result], None, None]: """ Iteratively calculate energies. Parameters ---------- yield_total_E_ia The results should include the contributions in the electron-hole basis to the total energy :math:`E_{ia}` and Coulomb energy :math:`E_{ia}^\\text{c}` yield_proj_E_ia The results should include the contributions in the electron-hole basis to the total energy projected on the occupied and unoccupied Voronoi weights :math:`E_{ia} w_i` and :math:`E_{ia} w_a`. yield_total_E_ou The results should include the contributions the total energy broadened on the occupied and unoccupied energy grids :math:`\\sum_{ia} E_{ia}\\delta(\\varepsilon_\\text{occ}-\\varepsilon_{i}) \\delta(\\varepsilon_\\text{unocc}-\\varepsilon_{a})` and yield_total_dists The results should include the contributions the total energy and Coulomb energy broadened by electronic transition energy onto the unoccupied energies grid :math:`\\sum_{ia} E_{ia} \\delta(\\varepsilon-\\omega_{ia})` and :math:`\\sum_{ia} E_{ia}^\\text{C} \\delta(\\varepsilon-\\omega_{ia})` direction Direction :math:`\\hat{\\boldsymbol{e}}` of the polarization of the pulse. Integer 0, 1 or 2 to specify x, y or z, or the direction vector specified as a list of three values. Default: polarization along z. Yields ------ Tuple (work, result) on the root rank of the calculation communicator. \ Does not yield on non-root ranks of the calculation communicator. work An object representing the metadata (time or pulse) for the work done. result Object containg the calculation results for this time and pulse. """ from gpaw.tddft.units import as_to_au, au_to_eV, au_to_eA include_energy_dists = (yield_total_dists or yield_total_E_ou) if include_energy_dists: assert self.sigma is not None assert all([order in self.density_matrices.derivative_order_s for order in [0, 1]]) assert 'Re' in self.density_matrices.reim and 'Im' in self.density_matrices.reim assert direction in [0, 1, 2] or (isinstance(direction, Sequence) and len(direction) == 3) if direction in [0, 1, 2]: direction_v = np.zeros(3) direction_v[direction] = 1 else: direction_v = np.array(direction) direction_v /= np.linalg.norm(direction_v) dm_p = direction_v @ self.ksd.dm_vp v0_p = dm_p * np.sqrt(2 * self.ksd.f_p) w_p = self.ksd.w_p # List all keys that this method computes # This is necessary to safely send and receive data across ranks resultkeys = self.get_result_keys(yield_total_dists=yield_total_dists, yield_total_E_ia=yield_total_E_ia, yield_proj_E_ia=yield_proj_E_ia, yield_total_E_ou=yield_total_E_ou, ) self._read_weights_diagonal() # Iterate over the pulses and times for work, dm in self.density_matrices: dm.rho_p # Read these now, so non calc_comm root ranks can continue dm.drho_p if dm.rank > 0: continue self.log.start('calculate') assert isinstance(work, ConvolutionDensityMatrixMetadata) assert isinstance(work.pulse, PulsePerturbation) pulsestr = work.pulse.pulse.strength(work.time * as_to_au) dipmom = - 2 * dm_p @ dm.rho_p.real Epulse = dipmom * pulsestr * au_to_eV # Calculate v_ia v_p = v0_p * pulsestr E_p = -v_p * dm.Q_p E_p -= dm.Q_p * dm.dP_p Ec_p = E_p.copy() E_p += dm.P_p * dm.dQ_p E_p *= 0.5 * au_to_eV Ec_p -= w_p * dm.Q_p ** 2 Ec_p *= 0.5 * au_to_eV if yield_total_E_ia or yield_total_E_ou: E_ia = self.ksd.M_ia_from_M_p(E_p) Ec_ia = self.ksd.M_ia_from_M_p(Ec_p) result = Result() if yield_total_E_ia: result['E_ia'] = E_ia result['Ec_ia'] = Ec_ia result['dm'] = dipmom * au_to_eA result['Epulse'] = Epulse result['total'] = np.sum(E_p) result['total_coulomb'] = np.sum(Ec_p) # (Optional) Broaden transitions by transition energy if yield_total_dists: assert self.sigma is not None result['E_transition_u'] = broaden_n2e(E_p, w_p * au_to_eV, self.energies_unocc, self.sigma) result['Ec_transition_u'] = broaden_n2e(Ec_p, w_p * au_to_eV, self.energies_unocc, self.sigma) # (Optional) Compute energy contribution matrix if yield_total_E_ou: result['E_ou'] = self.broaden_ia2ou(E_ia) result['Ec_ou'] = self.broaden_ia2ou(Ec_ia) # Initialize the remaining empty arrays result.create_all_empty(resultkeys) # Iterate over projections for iI, weight_n in enumerate(self._iterate_weights_diagonal): assert weight_n is not None weight_i = weight_n[self.flti] weight_a = weight_n[self.flta] for iI2, weight2_n in enumerate(self._iterate_weights_diagonal): assert weight2_n is not None weight2_a = weight2_n[self.flta] result['total_proj_II'][iI, iI2] = np.einsum( 'ia,i,a->', E_ia, weight_i, weight2_a, optimize=True) result['total_coulomb_proj_II'][iI, iI2] = np.einsum( 'ia,i,a->', Ec_ia, weight_i, weight2_a, optimize=True) if yield_proj_E_ia: E_occ_proj_ia = E_ia * weight_i[:, None] E_unocc_proj_ia = E_ia * weight_a[None, :] result['E_occ_proj_Iia'][iI] = E_occ_proj_ia result['E_unocc_proj_Iia'][iI] = E_unocc_proj_ia self.log_parallel(f'Calculated and broadened energies in {self.log.elapsed("calculate"):.2f}s ' f'for {work.desc}', flush=True) yield work, result if self.calc_comm.rank == 0: self.log_parallel('Finished calculating energies', flush=True)
@property def _need_derivatives_real_imag(self) -> tuple[list[int], bool, bool]: return ([0, 1], True, True) @property def _voronoi_shapes(self) -> dict[str, tuple[int, ...]]: nI = self.voronoi.nproj if nI == 0: return {} Nn = self.voronoi.nn # Diagonal weights only return {'diagonal': (nI, Nn)}
[docs] def calculate_and_write(self, out_fname: str, write_extra: dict[str, Any] = dict(), save_matrix: bool = False, save_dist: bool = False, only_one_pulse: bool = True): """ Calculate energy contributions. Energies are saved in a numpy archive if the file extension is ``.npz`` or in an ULM file if the file extension is ``.ulm``. Parameters ---------- out_fname File name of the resulting data file. write_extra Dictionary of extra key-value pairs to write to the data file. save_matrix Whether the electron-hole map matrix should be computed and saved. save_dist Whether the transition energy distributions should be computed and saved. only_one_pulse If False, group arrays by pulse. """ from ..writers.energy import EnergyWriter from ..writers.writer import TimeResultsCollector, PulseConvolutionResultsCollector out_fname = str(out_fname) if out_fname[-4:] == '.npz': calc_kwargs = dict(yield_total_E_ou=save_matrix, yield_total_dists=save_dist) cls = TimeResultsCollector if only_one_pulse else PulseConvolutionResultsCollector writer = EnergyWriter(cls(self, calc_kwargs), only_one_pulse=only_one_pulse) writer.calculate_and_save_npz(out_fname, write_extra=write_extra) elif out_fname[-4:] == '.ulm': assert only_one_pulse calc_kwargs = dict(yield_total_E_ou=save_matrix, yield_total_dists=save_dist) writer = EnergyWriter(TimeResultsCollector(self, calc_kwargs, exclude=['E_ou']), only_one_pulse=True) writer.calculate_and_save_ulm(out_fname, write_extra=write_extra) else: raise ValueError(f'output-file must have ending .npz or .ulm, is {out_fname}')