quackseq-simulator/tests/simulator.py
2024-06-06 19:45:40 +02:00

115 lines
3.4 KiB
Python

import unittest
import logging
import matplotlib.pyplot as plt
from quackseq.phase_table import PhaseTable
from quackseq.pulsesequence import QuackSequence
from quackseq.event import Event
from quackseq.functions import RectFunction
from quackseq_simulator.simulator import Simulator
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class TestQuackSequence(unittest.TestCase):
def test_event_creation(self):
seq = QuackSequence("test - event creation")
seq.add_pulse_event("tx", "10u", 100, 90.0, RectFunction())
seq.add_blank_event("blank", "3u")
seq.add_readout_event("rx", "100u", phase=90.0)
seq.add_blank_event("TR", "1m")
sim = Simulator()
sim.set_averages(100)
sim.settings.noise = 0
result = sim.run_sequence(seq)
self.assertIsNotNone(result)
self.assertTrue(hasattr(result, "tdx"))
self.assertTrue(hasattr(result, "tdy"))
self.assertGreater(len(result.tdx), 0)
self.assertGreater(len(result.tdy), 0)
# Plotting the result can be useful for visual inspection during development
plt.plot(result.tdx, result.tdy.imag, label="imaginary")
plt.plot(result.tdx, result.tdy.real, label="real")
plt.plot(result.tdx, abs(result.tdy), label="abs")
plt.legend()
plt.show()
def test_simulation_run_sequence(self):
seq = QuackSequence("test - simulation run sequence")
tx = Event("tx", "10u", seq)
seq.add_event(tx)
seq.set_tx_amplitude(tx, 100)
seq.set_tx_phase(tx, 0)
rect = RectFunction()
seq.set_tx_shape(tx, rect)
blank = Event("blank", "3u", seq)
seq.add_event(blank)
rx = Event("rx", "100u", seq)
seq.set_rx(rx, True)
seq.add_event(rx)
TR = Event("TR", "1m", seq)
seq.add_event(TR)
sim = Simulator()
sim.set_averages(100)
result = sim.run_sequence(seq)
self.assertIsNotNone(result)
self.assertTrue(hasattr(result, "tdx"))
self.assertTrue(hasattr(result, "tdy"))
self.assertGreater(len(result.tdx), 0)
self.assertGreater(len(result.tdy), 0)
# Plotting the result can be useful for visual inspection during development
plt.plot(result.tdx, abs(result.tdy))
plt.show()
def test_phase_cycling(self):
seq = QuackSequence("test - phase cycling")
tx = Event("tx", "10u", seq)
seq.add_event(tx)
seq.set_tx_amplitude(tx, 100)
seq.set_tx_phase(tx, 0)
seq.set_tx_n_phase_cycles(tx, 2)
seq.set_tx_phase_cycle_group(tx, 0)
rect = RectFunction()
seq.set_tx_shape(tx, rect)
tx2 = Event("tx2", "10u", seq)
seq.add_event(tx2)
seq.set_tx_amplitude(tx2, 100)
seq.set_tx_phase(tx2, 1)
seq.set_tx_n_phase_cycles(tx2, 4)
seq.set_tx_phase_cycle_group(tx2, 1)
tx3 = Event("tx3", "10u", seq)
seq.add_event(tx3)
seq.set_tx_amplitude(tx3, 100)
seq.set_tx_phase(tx3, 2)
seq.set_tx_n_phase_cycles(tx3, 3)
seq.set_tx_phase_cycle_group(tx3, 1)
sim = Simulator()
sim.set_averages(100)
result = sim.run_sequence(seq)
plt.plot(result.tdx, abs(result.tdy))
plt.show()
phase_table = PhaseTable(seq)
if __name__ == "__main__":
unittest.main()