from __future__ import annotations
from abc import ABC, abstractmethod
from typing import Any, Union
from numbers import Number
import numpy as np
from numpy.typing import NDArray
from gpaw.lcaotddft.laser import Laser, create_laser
def create_perturbation(perturbation: PerturbationLike):
if isinstance(perturbation, Perturbation):
return perturbation
if perturbation is None:
return NoPerturbation()
if isinstance(perturbation, Laser):
return PulsePerturbation(perturbation)
assert isinstance(perturbation, dict)
if perturbation['name'] == 'none':
return NoPerturbation
if perturbation['name'] == 'deltakick':
return DeltaKick(strength=perturbation['strength'])
pulse = create_laser(perturbation)
return PulsePerturbation(pulse)
[docs]
class Perturbation(ABC):
""" Perturbation that density matrices are a response to """
def timestep(self,
times: NDArray[np.float64]):
dt = times[1] - times[0]
assert np.allclose(times[1:] - dt, times[:-1]), 'Variable time step'
return dt
[docs]
@abstractmethod
def fourier(self,
times: NDArray[np.float64],
padnt: int) -> NDArray[np.complex128]:
"""
Fourier transform of perturbation
Parameters
----------
times
Time grid in atomic units
padnt
Length of zero-padded time grid
"""
raise NotImplementedError
@abstractmethod
def todict(self) -> dict[str, Any]:
raise NotImplementedError
def __eq__(self, other) -> bool:
""" Equal if dicts are identical (up to numerical tolerance)
"""
try:
d1 = self.todict()
d2 = other.todict()
except AttributeError:
return False
if d1.keys() != d2.keys():
return False
for key in d1.keys():
if isinstance(d1[key], Number) and isinstance(d2[key], Number):
if not np.isclose(d1[key], d2[key]):
return False
else:
if not d1[key] == d2[key]:
return False
return True
PerturbationLike = Union[Perturbation, Laser, dict, None]
[docs]
class NoPerturbation(Perturbation):
""" No perturbation
Used to indicate that we do not know the perturbation,
and that it should not matter.
"""
def __init__(self):
pass
[docs]
def fourier(self,
times: NDArray[np.float64],
padnt: int) -> NDArray[np.complex128]:
raise RuntimeError('Not possible for no perturbation')
def __str__(self) -> str:
return 'No perturbation'
def todict(self) -> dict[str, Any]:
return {'name': 'none'}
[docs]
class DeltaKick(Perturbation):
""" Delta-kick perturbation
Parameters
----------
strength
Strength of the perturbation
"""
def __init__(self,
strength: float):
self.strength = strength
[docs]
def fourier(self,
times: NDArray[np.float64],
padnt: int) -> NDArray[np.complex128]:
# The strength is specified in the frequency domain, hence no multiplication by timestep
norm = self.strength
return np.array([norm])
def todict(self) -> dict[str, Any]:
return {'name': 'deltakick', 'strength': self.strength}
def __str__(self) -> str:
return f'Delta-kick perturbation (strength {self.strength:.1e})'
[docs]
class PulsePerturbation(Perturbation):
""" Perturbation as a time-dependent function
Parameters
----------
pulse
Object representing the pulse
"""
def __init__(self,
pulse: Laser | dict):
self.pulse = create_laser(pulse)
[docs]
def fourier(self,
times: NDArray[np.float64],
padnt: int) -> NDArray[np.complex128]:
pulse_t = self.pulse.strength(times)
return np.fft.rfft(pulse_t, n=padnt) * self.timestep(times)
def todict(self) -> dict[str, Any]:
try:
return self.pulse.todict()
except AttributeError:
return {'name': self.pulse.__class__.__name__}
def __str__(self) -> str:
lines: list[str] = []
width = 50
for key, value in self.todict().items():
line = f'{key}: {value}'
if len(lines) == 0:
lines.append(line)
continue
if len(lines[-1]) + len(line) + 2 < width:
lines[-1] = lines[-1] + ', ' + line
else:
lines.append(line)
return '\n'.join(lines)