Add tests for the new model class

src/qutip_qip/device/processor.py

@@ -512,6 +512,9 @@ def save_coeff(self, file_name, inctime=True):
         """
         self._is_pulses_valid()
         coeffs = np.array(self.get_full_coeffs())
+
+        if not all([isinstance(pulse.label, str) for pulse in self.pulses]):
+            raise NotImplementedError("Only string labels are supported.")
         header = ";".join([str(pulse.label) for pulse in self.pulses])
         if inctime:
             shp = coeffs.T.shape

tests/test_device.py

@@ -186,3 +186,19 @@ def test_numerical_circuit(circuit, device_class, kwargs, schedule_mode):
     elif isinstance(device, SCQubits):
         target = _ket_expaned_dims(target, device.dims)
     assert _tol > abs(1 - qutip.metrics.fidelity(result.final_state, target))
+
+
+@pytest.mark.parametrize("processor_class", [DispersiveCavityQED, LinearSpinChain, CircularSpinChain, SCQubits])
+def test_pulse_plotting(processor_class):
+    try:
+        import matplotlib.pyplot as plt
+    except Exception:
+        return True
+    qc = QubitCircuit(3)
+    qc.add_gate("CNOT", 1, 0)
+    qc.add_gate("X", 1)
+
+    processor = processor_class(3)
+    processor.load_circuit(qc)
+    fig, ax = processor.plot_pulses()
+    plt.close(fig)

tests/test_model.py

@@ -0,0 +1,93 @@
+import pytest
+import numpy as np
+from numpy.testing import assert_array_equal, assert_allclose
+
+import qutip
+from qutip_qip.device import (
+    Model,
+    Processor,
+    DispersiveCavityQED,
+    LinearSpinChain,
+    CircularSpinChain,
+    SCQubits,
+)
+from qutip_qip.noise import RelaxationNoise
+
+
+def test_cavityqed_model():
+    model = DispersiveCavityQED(3, epsmax=[1.1, 1, 0, 0.8], w0=7.0)
+    assert model.get_all_drift() == []
+    model.get_control_labels()
+    model.get_control("g1")
+    model.get_control("sx0")
+    assert_array_equal(model.params["deltamax"], np.array([1.0] * 3))
+    assert_array_equal(model.params["w0"], 7.0)
+    assert_array_equal(model.params["epsmax"], [1.1, 1, 0, 0.8])
+    assert model.get_control(0) == model.get_control("sx0")
+    model.get_latex_str()
+
+
+@pytest.mark.parametrize(("model_class"), [LinearSpinChain, CircularSpinChain])
+def test_spinchain_model(model_class):
+    model = LinearSpinChain(3, sx=[1.1, 1, 0, 0.8], sz=7.0, t1=10.0)
+    assert model.get_all_drift() == []
+    model.get_control_labels()
+    if isinstance(model, LinearSpinChain):
+        assert len(model.get_control_labels()) == 3 * 3 - 1
+    elif isinstance(model, CircularSpinChain):
+        assert len(model.get_control_labels()) == 3 * 3
+    model.get_control("g1")
+    model.get_control("sx0")
+    assert_array_equal(model.params["sz"], 7.0)
+    assert_array_equal(model.params["sx"], [1.1, 1, 0, 0.8])
+    assert model.get_control(0) == model.get_control("sx0")
+    model.get_latex_str()
+    assert model.params["t1"] == 10.0
+
+
+def test_scqubits_model():
+    model = SCQubits(
+        3, dims=[4, 3, 4], omega_single=0.02, alpha=[-0.02, -0.015, -0.025]
+    )
+    assert model.dims == [4, 3, 4]
+    model.get_all_drift()
+    model.get_control_labels()
+    model.get_control("sx2")
+    model.get_control("zx10")
+    assert_array_equal(model.params["omega_single"], np.array([0.02] * 3))
+    assert_array_equal(model.params["alpha"], [-0.02, -0.015, -0.025])
+    assert model.get_control(0) == model.get_control("sx0")
+    model.get_latex_str()
+
+
+def test_define_model_in_processor():
+    processor = Processor(3)
+
+    # Define Hamiltonian model
+    processor.add_drift(qutip.sigmax(), 1)
+    assert processor.get_all_drift() == [(qutip.sigmax(), [1])]
+    processor.add_drift(qutip.sigmaz(), cyclic_permutation=True)
+    assert len(processor.drift) == 4
+    processor.add_control(qutip.sigmaz(), targets=2, label="sz")
+    processor.set_coeffs({"sz": np.array([[0.0, 1.0]])})
+    processor.controls[0] == qutip.tensor(
+        [qutip.qeye(2), qutip.qeye(2), qutip.sigmax()]
+    )
+    processor.add_control(qutip.sigmax(), label="sx", cyclic_permutation=True)
+    label_dict = {"sz", ("sx", (0,)), ("sx", (1,)), ("sx", (2,))}
+    for label in processor.get_control_labels():
+        assert label in label_dict
+    assert processor.get_control("sz") == (qutip.sigmaz(), [2])
+
+    # Relaxation time
+    processor.t1 = 100.0
+    processor.t2 = 60.0
+    assert processor.t1 == 100.0
+    assert processor.t2 == 60.0
+
+    # Noise
+    noise_object = RelaxationNoise(t1=20.0)
+    processor.add_noise(noise_object)
+    assert noise_object in processor.noise
+    with pytest.raises(TypeError):
+        processor.add_noise("non-noise-object")

tests/test_optpulseprocessor.py

@@ -39,17 +39,6 @@ def test_simple_hadamard(self):
         result = test.run_state(rho0)
         assert_allclose(fidelity(result.states[-1], plus), 1, rtol=1.0e-6)
 
-        # test add/remove ctrl
-        test.add_control(sigmay())
-        test.remove_pulse(0)
-        assert_(
-            len(test.pulses) == 1,
-            msg="Method of remove_pulse could be wrong.")
-        assert_allclose(test.drift.drift_hamiltonians[0].qobj, H_d)
-        assert_(
-            sigmay() in test.ctrls,
-            msg="Method of remove_pulse could be wrong.")
-
     def test_multi_qubits(self):
         """
         Test for multi-qubits system.
@@ -68,11 +57,7 @@ def test_multi_qubits(self):
         # test periodically adding ctrls
         sx = sigmax()
         iden = identity(2)
-        # print(test.ctrls)
-        # print(Qobj(tensor([sx, iden, sx])))
-        assert_(Qobj(tensor([sx, iden, sx])) in test.ctrls)
-        assert_(Qobj(tensor([iden, sx, sx])) in test.ctrls)
-        assert_(Qobj(tensor([sx, sx, iden])) in test.ctrls)
+        assert(len(test.get_control_labels()) == 3)
         test.add_control(sigmax(), cyclic_permutation=True)
         test.add_control(sigmay(), cyclic_permutation=True)
 
@@ -86,17 +71,6 @@ def test_multi_qubits(self):
             rho0, options=Options(store_states=True))
         assert_(fidelity(result.states[-1], rho1) > 1-1.0e-6)
 
-        # # test save and read coeffs
-        # test.save_coeff("qutip_test_multi_qubits.txt")
-        # test2 = OptPulseProcessor(N, H_d, H_c)
-        # test2.drift = test.drift
-        # test2.ctrls = test.ctrls
-        # test2.read_coeff("qutip_test_multi_qubits.txt")
-        # os.remove("qutip_test_multi_qubits.txt")
-        # assert_(np.max((test.coeffs-test2.coeffs)**2) < 1.0e-13)
-        # result = test2.run_state(rho0,)
-        # assert_(fidelity(result.states[-1], rho1) > 1-1.0e-6)
-
     def test_multi_gates(self):
         N = 2
         H_d = tensor([sigmaz()]*2)