test_modematching.py

import unittest

import matplotlib.pyplot as plt

from openqlab import io
from openqlab.analysis.cavity import modematching


class TestModematching(unittest.TestCase):
    def setUp(self):
        self.data = io.read(
            "test_analysis/test_cavity/opo_scan.csv", importer="KeysightCSV"
        )
        plt.figure()

    def test_default_inverted(self):
        result = modematching(self.data["opo_scan_2"], plot=True)
        # plt.show()
        self.assertAlmostEqual(result, 0.955, places=2)

    def test_default_not_inverted(self):
        self.data *= -1
        result = modematching(self.data["opo_scan_2"], plot=True)
        self.assertAlmostEqual(result, 0.955, places=2)

    def test_to_many_columns(self):
        with self.assertRaisesRegex(ValueError, "only contain one single"):
            modematching(self.data)

    # @unittest.skip
    def test_with_clipped_data(self):
        self.data.clip(lower=-1, upper=1, inplace=True)
        result = modematching(self.data["opo_scan_2"], U_max=6.668)
        self.assertAlmostEqual(result, 0.955, places=2)

    def test_fail_with_only_one_peak(self):
        with self.assertRaises(ValueError) as err:
            modematching(self.data["opo_scan_2"].loc[:0.01])
        self.assertEqual(
            err.exception.args[0],
            "The main mode must occur exactly two times for the algorithm to work, "
            "but it found 1 main modes.",
        )

    def test_with_much_offset(self):
        self.data *= -1
        self.data -= 2
        result = modematching(
            self.data["opo_scan_2"], U_max=6.668 - 2, offset=-1.98, plot=True
        )
        self.assertAlmostEqual(result, 0.955, places=2)

    def test_impossible_parameter_combinations(self):
        with self.assertRaises(ValueError):
            modematching(self.data["opo_scan_2"], without_main_peaks=True)

    def test_cropped_main_peak(self):
        data = io.read(
            "test_analysis/test_cavity/opo_scan_zoom.csv", importer="KeysightCSV"
        )
        mm = modematching(
            data["opo_scan_zoom_2"].loc[0.0015:0.0155],
            offset=-0.015,
            plot=True,
            U_max=9.27,
            rel_prominence=0.003,
            without_main_peaks=True,
        )
        self.assertAlmostEqual(mm, 0.92, places=2)