from __future__ import annotations
import numpy as np
from numpy.typing import NDArray
from ..density_matrices.frequency import FrequencyDensityMatrices
from ..density_matrices.base import WorkMetadata
from ..density_matrices.time import ConvolutionDensityMatrices
from ..utils import Result, get_gaussian_pulse_values
from .writer import ResultsCollector, Writer
[docs]
class DipoleWriter(Writer):
""" Calculate dipole moment contributions, optionally broadened onto
an energy grid as a transition contribution map
Parameters
----------
collector
ResultsCollector object
"""
def __init__(self,
collector: ResultsCollector):
super().__init__(collector)
if isinstance(self.density_matrices, ConvolutionDensityMatrices):
self._ulm_tag = 'Time TCM'
assert len(self.density_matrices.pulses) == 1, 'Only one pulse allowed'
else:
assert isinstance(self.density_matrices, FrequencyDensityMatrices)
self._ulm_tag = 'TCM'
@property
def common_arrays(self) -> dict[str, NDArray[np.float64] | int | float]:
common = super().common_arrays
if self.calc.sigma is not None:
# There is an energy grid
common['sigma'] = self.calc.sigma
common['energy_o'] = np.array(self.calc.energies_occ)
common['energy_u'] = np.array(self.calc.energies_unocc)
if isinstance(self.density_matrices, ConvolutionDensityMatrices):
common['time_t'] = self.density_matrices.times * 1e-3
else:
assert isinstance(self.density_matrices, FrequencyDensityMatrices)
common['freq_w'] = self.density_matrices.frequencies
common['scale_w'] = 4 * common['freq_w'] / np.pi
if isinstance(self.density_matrices, ConvolutionDensityMatrices):
# If pulse is Gaussian pulse, then get dictionary with 'pulsefreq' and 'pulsefwhm'
common.update(**get_gaussian_pulse_values(self.density_matrices.pulses[0]))
return common
[docs]
def fill_ulm(self,
writer,
work: WorkMetadata,
result: Result):
if self.collector.calc_kwargs.get('yield_total_ou', False):
writer.fill(result['dm_ouv'])
[docs]
def write_empty_arrays_ulm(self, writer):
if not self.collector.calc_kwargs.get('yield_total_ou', False):
return
shape_ou = (len(self.calc.energies_occ), len(self.calc.energies_unocc))
if isinstance(self.density_matrices, ConvolutionDensityMatrices):
Nt = len(self.density_matrices.times)
writer.add_array('dm_touv', (Nt, ) + shape_ou + (3, ), dtype=float)
else:
assert isinstance(self.density_matrices, FrequencyDensityMatrices)
Nw = len(self.density_matrices.frequencies)
writer.add_array('dm_wouv', (Nw, ) + shape_ou + (3, ), dtype=float)