Skeleton for coherent/lindbladian light interaction

src/qimpy/transport/material/ab_initio/__init__.py

@@ -1,6 +1,7 @@
-__all__ = ["PackedHermitian", "RelaxationTime", "Lindblad", "AbInitio"]
+__all__ = ["PackedHermitian", "RelaxationTime", "Lindblad", "Light", "AbInitio"]
 
 from ._packed_hermitian import PackedHermitian
 from ._relaxation_time import RelaxationTime
 from ._lindblad import Lindblad
+from ._light import Light
 from ._ab_initio import AbInitio

src/qimpy/transport/material/ab_initio/_ab_initio.py

@@ -10,7 +10,7 @@
 from qimpy.io import Checkpoint, CheckpointPath, Unit
 from qimpy.mpi import ProcessGrid
 from .. import Material, fermi
-from . import PackedHermitian, RelaxationTime, Lindblad
+from . import PackedHermitian, RelaxationTime, Lindblad, Light
 
 
 class DynamicsTerm(Protocol):
@@ -31,6 +31,7 @@ class AbInitio(Material):
     S: Optional[torch.Tensor]  #: Spin matrix elements
     L: Optional[torch.Tensor]  #: Angular momentum matrix elements
     scattering: DynamicsTerm  #: scattering functional
+    light: Light  #: coherent-light interaction
     dynamics_terms: list[DynamicsTerm]  #: all active drho/dt contributions
 
     def __init__(
@@ -45,6 +46,7 @@ def __init__(
         ),
         relaxation_time: Optional[Union[RelaxationTime, dict]] = None,
         lindblad: Optional[Union[Lindblad, dict]] = None,
+        light: Optional[Union[Light, dict]] = None,
         process_grid: ProcessGrid,
         checkpoint_in: CheckpointPath = CheckpointPath(),
     ):
@@ -121,6 +123,10 @@ def __init__(
             if not isinstance(self.scattering, RelaxationTime) or self.scattering:
                 self.dynamics_terms.append(self.scattering)
 
+            if light is not None:
+                self.add_child("light", Light, light, checkpoint_in, ab_initio=self)
+                self.dynamics_terms.append(self.light)
+
     def read_scalars(self, data_file: Checkpoint, name: str) -> torch.Tensor:
         """Read quantities that don't transform with rotations from data_file."""
         dset = data_file[name]

src/qimpy/transport/material/ab_initio/_light.py

@@ -0,0 +1,75 @@
+from __future__ import annotations
+from typing import Optional
+
+import torch
+
+from qimpy import rc, TreeNode
+from qimpy.io import CheckpointPath, InvalidInputException
+from qimpy.transport import material
+
+
+class Light(TreeNode):
+    ab_initio: material.ab_initio.AbInitio
+    coherent: bool  #: Whether term is Coherent or Lindbladian
+    gauge: str  #: Gauge: one of velocity or length
+    A0: torch.Tensor  #: Vector potential amplitude
+    omega: float  #: light frequency
+    t0: float  #: center of Gaussian pulse, if sigma is non-zero
+    sigma: float  #: width of Gaussian pulse in time, if non-zero
+
+    def __init__(
+        self,
+        *,
+        ab_initio: material.ab_initio.AbInitio,
+        coherent: bool,
+        gauge: str = "velocity",
+        A0: Optional[list[complex]] = None,
+        E0: Optional[list[complex]] = None,
+        omega: float,
+        t0: float = 0.0,
+        sigma: float = 0.0,
+        checkpoint_in: CheckpointPath = CheckpointPath(),
+    ) -> None:
+        """
+        Initialize coherent light interaction.
+
+        Parameters
+        ----------
+        coherent
+            :yaml:`Switch between coherent and Lindbladian implementations.`
+        gauge
+            :yaml:`Switch between 'velocity' or 'length' gauge.`
+        A0
+            :yaml:`Vector potential amplitude.`
+            TODO: specify details about differences in CW vs pulse mode.
+            Exactly one of A0 or E0 must be specified.
+        E0:
+            :yaml:`Electric-field amplitude.`
+            Exactly one of A0 or E0 must be specified.
+        omega
+            :yaml:`Angular frequency / photon energy of the light.`
+        t0
+            :yaml:`Center of Gaussian pulse, used only if sigma is non-zero.`
+        sigma
+            :yaml:`Time width of Gaussian pulse, if non-zero.`
+        """
+        super().__init__()
+        self.coherent = coherent
+        self.ab_initio = ab_initio
+        self.gauge = gauge
+
+        # Get amplitude from A0 or E0:
+        if (A0 is None) == (E0 is None):
+            raise InvalidInputException("Exactly one of A0 and E0 must be specified")
+        if A0 is not None:
+            self.A0 = torch.tensor(A0, device=rc.device)
+        if E0 is not None:
+            self.A0 = torch.tensor(E0, device=rc.device) / omega
+        if self.A0.shape[-1] == 2:  # handle complex tensors
+            self.A0 = torch.view_as_complex(self.A0)
+        self.omega = omega
+        self.t0 = t0
+        self.sigma = sigma
+
+    def rho_dot(self, rho: torch.Tensor, t: float) -> torch.Tensor:
+        raise NotImplementedError

src/qimpy/transport/material/ab_initio/_relaxation_time.py

@@ -43,4 +43,4 @@ def __bool__(self) -> bool:
     def rho_dot(self, rho: torch.Tensor, t: float) -> torch.Tensor:
         if not self:
             return torch.zeros_like(rho)
-        return NotImplemented
+        raise NotImplementedError