Velocity gauge

quantum_systems/basis_set.py

@@ -46,6 +46,7 @@ def __init__(
         self._u = None
         self._s = None
         self._position = None
+        self._momentum = None
 
         self._sigma_x = None
         self._sigma_y = None
@@ -130,6 +131,19 @@ def position(self, position):
     def dipole_moment(self):
         return self.particle_charge * self.position
 
+    @property
+    def momentum(self):
+        return self._momentum
+
+    @momentum.setter
+    def momentum(self, momentum):
+        assert len(momentum) == self.dim
+
+        for i in range(self.dim):
+            assert all(self.check_axis_lengths(momentum[i], self.l))
+
+        self._momentum = momentum
+
     @property
     def spin_x(self):
         return self._spin_x
@@ -272,6 +286,7 @@ def change_module(self, np):
             ("_spf", self.spf),
             ("_bra_spf", self.bra_spf),
             ("_position", self.position),
+            ("_momentum", self.momentum),
             ("_spin_x", self.spin_x),
             ("_spin_y", self.spin_y),
             ("_spin_z", self.spin_z),
@@ -293,6 +308,7 @@ def cast_to_complex(self):
             ("_spf", self.spf),
             ("_bra_spf", self.bra_spf),
             ("_position", self.position),
+            ("_momentum", self.momentum),
             ("_spin_x", self.spin_x),
             ("_spin_y", self.spin_y),
             ("_spin_z", self.spin_z),
@@ -377,6 +393,18 @@ def _change_basis_position_elements(self, C, C_tilde):
 
         self.position = self.np.asarray(position)
 
+    def _change_basis_momentum_elements(self, C, C_tilde):
+        momentum = []
+
+        for i in range(self.momentum.shape[0]):
+            momentum.append(
+                self.transform_one_body_elements(
+                    self.momentum[i], C, np=self.np, C_tilde=C_tilde
+                )
+            )
+
+        self.momentum = self.np.asarray(momentum)
+
     def _change_basis_spf(self, C, C_tilde):
         self.bra_spf = self.transform_bra_spf(self.bra_spf, C_tilde, self.np)
 
@@ -429,6 +457,9 @@ def change_basis(self, C, C_tilde=None):
         if self.position is not None:
             self._change_basis_position_elements(C, C_tilde)
 
+        if self.momentum is not None:
+            self._change_basis_momentum_elements(C, C_tilde)
+
         if self.spf is not None:
             self._change_basis_spf(C, C_tilde)
 
@@ -582,6 +613,14 @@ def change_to_general_orbital_basis(
 
             self.position = self.np.array(position)
 
+        if not self.momentum is None:
+            momentum = [
+                self.add_spin_one_body(self.momentum[i], np=self.np)
+                for i in range(len(self.momentum))
+            ]
+
+            self.momentum = self.np.array(momentum)
+
         if not self.spf is None:
             self.spf = self.add_spin_spf(self.spf, self.np)
 

quantum_systems/system.py

@@ -108,6 +108,11 @@ def position(self):
         """Getter retuning the position matrix."""
         return self._basis_set.position
 
+    @property
+    def momentum(self):
+        """Getter returning the momentum matrix."""
+        return self._basis_set.momentum
+
     @property
     def dipole_moment(self):
         """Getter returning the dipole moment elements."""

quantum_systems/time_evolution_operators/operator.py

@@ -106,10 +106,27 @@ class DipoleFieldInteraction(TimeEvolutionOperator):
         axis as a function of ``current_time``. Default is ``None`` which gives
         a constant polarization along the first axis of the dipole moment
         integrals.
+    gauge : str
+        String specifying the gauge choice. 'length' (Default) or 'velocity'.
+    quadratic_term : bool
+        Specifying whether to include quadratic vector potential term. Only
+        relevant for velocity gauge.
     """
 
-    def __init__(self, electric_field_strength, polarization_vector=None):
-        self._electric_field_strength = electric_field_strength
+    def __init__(
+        self,
+        field_strength,
+        polarization_vector=None,
+        gauge="length",
+        quadratic_term=True,
+    ):
+        assert gauge in [
+            "length",
+            "velocity",
+        ], "gauge must be either length or velocity."
+        self._length_gauge = True if gauge == "length" else False
+        self._quadratic_term = quadratic_term
+        self._field_strength = field_strength
         self._polarization = polarization_vector
 
     @property
@@ -119,9 +136,9 @@ def is_one_body_operator(self):
     def h_t(self, current_time):
         np = self._system.np
 
-        if not callable(self._electric_field_strength):
-            tmp = self._electric_field_strength
-            self._electric_field_strength = lambda t: tmp
+        if not callable(self._field_strength):
+            tmp = self._field_strength
+            self._field_strength = lambda t: tmp
 
         if self._polarization is None:
             # Set default polarization along x-axis
@@ -132,13 +149,26 @@ def h_t(self, current_time):
             tmp = self._polarization
             self._polarization = lambda t: tmp
 
-        return self._system.h - self._electric_field_strength(
-            current_time
-        ) * np.tensordot(
-            self._polarization(current_time),
-            self._system.dipole_moment,
-            axes=(0, 0),
-        )
+        if self._length_gauge:
+            H_t = -self._field_strength(current_time) * np.tensordot(
+                self._polarization(current_time),
+                self._system.dipole_moment,
+                axes=(0, 0),
+            )
+        else:
+            H_t = self._field_strength(current_time) * np.tensordot(
+                self._polarization(current_time),
+                self._system.momentum,
+                axes=(0, 0),
+            )
+            if self._quadratic_term:
+                H_t += (
+                    0.5
+                    * self._field_strength(current_time) ** 2
+                    * np.eye(self._system.l)
+                )
+
+        return self._system.h + H_t
 
 
 class AdiabaticSwitching(TimeEvolutionOperator):