add a more thorough test for model

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_processor.py

@@ -1,4 +1,5 @@
 import os
+import pytest
 
 from numpy.testing import (
     assert_, run_module_suite, assert_allclose, assert_equal)
@@ -15,25 +16,27 @@
 
 
 class TestCircuitProcessor:
-    def test_modify_ctrls(self):
-        """
-        Test for modifying Hamiltonian, add_control, remove_pulse
-        """
-        N = 2
-        proc = Processor(N=N)
-        proc.ctrls
-        proc.add_control(sigmaz())
-        assert_(tensor([sigmaz(), identity(2)]), proc.get_control(0))
-        proc.add_control(sigmax(), cyclic_permutation=True, label="sx")
-        # assert_allclose(len(proc.get_control_labels()), 3)
-        # assert_allclose(tensor([sigmax(), identity(2)]).full(), proc.get_control("sx[0]")[0].full())
-        # assert_allclose(tensor([identity(2), sigmax()]).full(), proc.get_control("sx[1]").full())
-        # proc.add_control(sigmay(), targets=1)
-        # assert_allclose(tensor([identity(2), sigmay()]), proc.ctrls[3])
-        # proc.remove_pulse([0, 1, 2])
-        # assert_allclose(tensor([identity(2), sigmay()]), proc.ctrls[0])
-        # proc.remove_pulse(0)
-        # assert_allclose(len(proc.ctrls), 0)
+    def test_control_and_coeffs(self):
+        processor = Processor(2)
+        processor.add_control(sigmax())
+        processor.add_control(sigmaz())
+
+        # Set coeffs and tlist without a label
+        coeffs = np.array([[1., 2., 3.], [3., 2., 1.]])
+        processor.set_coeffs(coeffs)
+        assert_allclose(coeffs, processor.coeffs)
+        tlist = np.array([0., 1., 2., 3.])
+        processor.set_tlist(tlist)
+        assert_allclose(tlist, processor.get_full_tlist())
+        processor.set_tlist({0: tlist, 1: tlist})
+        assert_allclose(tlist, processor.get_full_tlist())
+
+        # Pulses
+        assert(len(processor.pulses) == 2)
+        assert(processor.find_pulse(0) == processor.pulses[0])
+        assert(processor.find_pulse(1) == processor.pulses[1])
+        with pytest.raises(KeyError):
+            processor.find_pulse("non_exist_pulse")
 
     def test_save_read(self):
         """
@@ -331,15 +334,23 @@ def test_no_saving_intermidiate_state(self):
         result = processor.run_state(basis(2,0), tlist=[0,1])
         assert(len(result.states) == 2)
 
-    def test_control_label(self):
-        processor = Processor(1)
-        processor.add_control(sigmax(), 0, label="test")
-        assert("test" in processor.get_control_labels())
-
     def test_repeated_use_of_processor(self):
         processor = Processor(1, t1=1.)
         processor.add_pulse(
             Pulse(sigmax(), targets=0, coeff=True))
         result1 = processor.run_state(basis(2, 0), tlist=np.linspace(0, 1, 10))
         result2 = processor.run_state(basis(2, 0), tlist=np.linspace(0, 1, 10))
         assert_allclose(result1.states[-1].full(), result2.states[-1].full())
+
+    def test_pulse_mode(self):
+        processor = Processor(2)
+        processor.add_control(sigmax(), targets=0, label="sx")
+        processor.set_coeffs({"sx": np.array([1., 2., 3.])})
+        processor.set_tlist({"sx": np.array([0., 1., 2., 3.])})
+
+        processor.pulse_mode = "continuous"
+        assert(processor.pulse_mode == "continuous")
+        assert(processor.pulses[0].spline_kind == "cubic")
+        processor.pulse_mode = "discrete"
+        assert(processor.pulse_mode == "discrete")
+        assert(processor.pulses[0].spline_kind == "step_func")