diff --git a/quantum_systems/time_evolution_operators/__init__.py b/quantum_systems/time_evolution_operators/__init__.py index 39bdbb2..0dcbce9 100644 --- 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 index 21a95ef..28ee11f 100644 --- 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