Merge pull request #48 from Schoyen/custom-one-body-operator

Custom one-body operator
HyQD · Sep 21, 2020 · 41af6da · 41af6da
2 parents f3b31fb + 254a620
commit 41af6da
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 1 deletion.
diff --git a/quantum_systems/time_evolution_operators/__init__.py b/quantum_systems/time_evolution_operators/__init__.py
@@ -1 +1,6 @@
-from .operator import TimeEvolutionOperator, LaserField, AdiabaticSwitching
+from .operator import (
+    TimeEvolutionOperator,
+    LaserField,
+    AdiabaticSwitching,
+    CustomOneBodyOperator,
+)
diff --git a/quantum_systems/time_evolution_operators/operator.py b/quantum_systems/time_evolution_operators/operator.py
@@ -123,3 +123,22 @@ def u_t(self, current_time):
         np = self._system.np
 
         return self._switching_function(current_time) * self._system.u
+
+
+class CustomOneBodyOperator(TimeEvolutionOperator):
+    def __init__(self, weight, operator):
+        self._weight = weight
+        self._operator = operator
+
+    @property
+    def is_one_body_operator(self):
+        return True
+
+    def h_t(self, current_time):
+        np = self._system.np
+
+        if not callable(self._weight):
+            tmp = self._weight
+            self._weight = lambda t: tmp
+
+        return self._system.h + self._weight(current_time) * self._operator