diff --git a/src/krotov/objectives.py b/src/krotov/objectives.py index 5ec100f7..f45b7e71 100644 --- a/src/krotov/objectives.py +++ b/src/krotov/objectives.py @@ -16,12 +16,6 @@ import numpy as np import qutip -from packaging.version import parse as parse_version - -if parse_version(qutip.__version__) < parse_version("5"): - is_qutip5 = False -else: - is_qutip5 = True from qutip.solver import Options as QutipSolverOptions from qutip.solver import Result as QutipSolverResult @@ -382,15 +376,13 @@ def propagate( e_ops = [] if args is None: args = {} - if is_qutip5: - result_options = { - "store_final_state": False, - "store_states": None, - "normalize_output": False, - } - result = QutipSolverResult(e_ops=e_ops, options=result_options) - else: - result = QutipSolverResult() + result_options = { + "store_final_state": False, + "store_states": None, + "normalize_output": False, + } + result = QutipSolverResult(e_ops=e_ops, options=result_options) + try: result.solver = propagator.__name__ except AttributeError: @@ -401,21 +393,7 @@ def propagate( state = rho0 if state is None: state = self.initial_state - if is_qutip5: - result.add(tlist[0], state) - else: - result.times = np.array(tlist) - result.states = [] - result.expect = [] - result.num_expect = len(e_ops) - result.num_collapse = len(c_ops) - for _ in e_ops: - result.expect.append([]) - if len(e_ops) == 0: - result.states.append(state) - else: - for i, oper in enumerate(e_ops): - result.expect[i].append(expect(oper, state)) + result.add(tlist[0], state) controls = extract_controls([self]) pulses_mapping = extract_controls_mapping([self], controls) mapping = pulses_mapping[0] # "first objective" (dummy structure) @@ -437,16 +415,7 @@ def propagate( c_ops_at_t, initialize=True, # initialize=(time_index == 0) ) - if is_qutip5: - result.add(tlist[time_index + 1], state) - else: - if len(e_ops) == 0: - result.states.append(state) - else: - for i, oper in enumerate(e_ops): - result.expect[i].append(expect(oper, state)) - if not is_qutip5: - result.expect = [np.array(a) for a in result.expect] + result.add(tlist[time_index + 1], state) return result @classmethod diff --git a/src/krotov/optimize.py b/src/krotov/optimize.py index d6ac35df..ae53f331 100644 --- a/src/krotov/optimize.py +++ b/src/krotov/optimize.py @@ -6,8 +6,7 @@ import numpy as np import threadpoolctl -from qutip import Qobj -from qutip import serial_map +from qutip import Qobj, serial_map from .conversions import ( control_onto_interval, @@ -451,8 +450,8 @@ def optimize_pulses( if second_order: σ = sigma(tlist[time_index] + 0.5 * dt) # pulse update - for (i_pulse, guess_pulse) in enumerate(guess_pulses): - for (i_obj, objective) in enumerate(objectives): + for i_pulse, guess_pulse in enumerate(guess_pulses): + for i_obj, objective in enumerate(objectives): χ = backward_states[i_obj][time_index] μ = mu( objectives, @@ -718,7 +717,7 @@ def _restore_from_previous_result(result, objectives, tlist, store_all_pulses): "When continuing from a previous Result, the number of " "objectives must be the same" ) - for (a, b) in zip(objectives, result.objectives): + for a, b in zip(objectives, result.objectives): if a != b: raise ValueError( "When continuing from a previous Result, the objectives must " diff --git a/src/krotov/parallelization.py b/src/krotov/parallelization.py index 15995f75..da8653b5 100644 --- a/src/krotov/parallelization.py +++ b/src/krotov/parallelization.py @@ -96,6 +96,7 @@ .. _ipyparallel: https://ipyparallel.readthedocs.io/en/latest/ """ + import contextlib import multiprocessing from concurrent.futures import ProcessPoolExecutor @@ -407,10 +408,10 @@ def __call__(self, data): time_index = self.time_index pulse_vals = self.pulse_vals i_state = self.i_state - for (pulse, pulse_val) in zip(pulses, pulse_vals): + for pulse, pulse_val in zip(pulses, pulse_vals): pulse[time_index] = pulse_val - for (pulse, pulse_val) in zip(pulses, pulse_vals): + for pulse, pulse_val in zip(pulses, pulse_vals): pulse[time_index] = pulse_val state = states[i_state] mapping = pulses_mapping[i_state] @@ -584,7 +585,7 @@ def _pmfw_forward_prop_step(pulse_vals, time_index): _threadpool_limits = threadpool_limits i_state = _pmfw_data['state_index'] pulses = _pmfw_data['pulses'] - for (pulse, pulse_val) in zip(pulses, pulse_vals): + for pulse, pulse_val in zip(pulses, pulse_vals): pulse[time_index] = pulse_val mapping = _pmfw_data['pulses_mapping'][i_state] obj = _pmfw_data['objective'] diff --git a/src/krotov/propagators.py b/src/krotov/propagators.py index ea42da31..4139f8d4 100644 --- a/src/krotov/propagators.py +++ b/src/krotov/propagators.py @@ -70,15 +70,7 @@ import scipy import threadpoolctl from packaging.version import parse as parse_version - -if parse_version(qutip.__version__) < parse_version("5"): - is_qutip5 = False - from qutip.cy.spconvert import dense2D_to_fastcsr_fmode - from qutip.cy.spmatfuncs import spmvpy_csr - from qutip.superoperator import mat2vec, vec2mat -else: - is_qutip5 = True - from qutip.core.superoperator import stack_columns, unstack_columns +from qutip.core.superoperator import stack_columns, unstack_columns __all__ = ['expm', 'Propagator', 'DensityMatrixODEPropagator'] @@ -259,21 +251,12 @@ def __call__( self._t += dt self._r.integrate(self._t) self._y = self._r.y - if is_qutip5: - return qutip.Qobj( - unstack_columns(self._y), - dims=state.dims, - isherm=True, - dtype="csr", - ) - else: - return qutip.Qobj( - dense2D_to_fastcsr_fmode( - vec2mat(self._y), state.shape[0], state.shape[1] - ), - dims=state.dims, - isherm=True, - ) + return qutip.Qobj( + unstack_columns(self._y), + dims=state.dims, + isherm=True, + dtype="csr", + ) @staticmethod def _rhs(t, rho, L_list): @@ -282,17 +265,11 @@ def _rhs(t, rho, L_list): out = np.zeros(rho.shape[0], dtype=complex) L = L_list[0][0] L_coeff = L_list[0][1] - if is_qutip5: - out = L_coeff * L @ rho - else: - spmvpy_csr(L.data, L.indices, L.indptr, rho, L_coeff, out) + out = L_coeff * L @ rho for n in range(1, len(L_list)): L = L_list[n][0] L_coeff = L_list[n][1] - if is_qutip5: - out = L_coeff * L @ rho - else: - spmvpy_csr(L.data, L.indices, L.indptr, rho, L_coeff, out) + out = L_coeff * L @ rho return out def _initialize(self, L, rho, dt, c_ops, backwards): @@ -327,12 +304,7 @@ def _initialize_data(self, L, rho, dt, c_ops, backwards): self._control_indices = control_indices if rho.type == 'oper': - if is_qutip5: - self._y = unstack_columns(rho.full()).ravel( - 'F' - ) # initial state - else: - self._y = mat2vec(rho.full()).ravel('F') # initial state + self._y = unstack_columns(rho.full()).ravel('F') # initial state else: raise ValueError("rho must be a density matrix") diff --git a/tests/test_dump_result.py b/tests/test_dump_result.py index b4cc708c..dae81e16 100644 --- a/tests/test_dump_result.py +++ b/tests/test_dump_result.py @@ -6,13 +6,8 @@ import krotov import pytest -import qutip -from packaging.version import parse as parse_version -if parse_version(qutip.__version__) < parse_version("5"): - oct_result_name = "oct_result_qutip4.dump" -else: - oct_result_name = "oct_result_qutip5.dump" +oct_result_name = "oct_result.dump" def incl_range(a, b, step=1): diff --git a/tests/test_dump_result/oct_result_qutip5.dump b/tests/test_dump_result/oct_result.dump similarity index 100% rename from tests/test_dump_result/oct_result_qutip5.dump rename to tests/test_dump_result/oct_result.dump diff --git a/tests/test_dump_result/oct_result_qutip4.dump b/tests/test_dump_result/oct_result_qutip4.dump deleted file mode 100644 index 1aedd140..00000000 Binary files a/tests/test_dump_result/oct_result_qutip4.dump and /dev/null differ