def test_fit_coeffs_init_twice_largest_signal_same_sign(self):
        """Check fit coefficients have correct #, the right ones are trainable,
            and that randomisation/defaults work
        """
        largest_signal = [0.0] * bmfc.count
        for model in bmfc.signal_models:
            coeffs = bmfc.signal(model)
            for idx, coeff in enumerate(coeffs):
                if tf.math.abs(coeff).numpy() > tf.math.abs(
                        largest_signal[idx]).numpy():
                    largest_signal[idx] = coeff

        fit = bmfc.fit(bmfc.FIT_INIT_TWICE_LARGEST_SIGNAL_SAME_SIGN)

        for i in range(48):
            if i in self.fit_trainable_ids:
                # Randomization should be enabled. Check from 0 to 2x largest signal
                if largest_signal[i].numpy() < 0:
                    random_min = 2.0 * largest_signal[i].numpy()
                    random_max = 0.0
                else:
                    random_min = 0.0
                    random_max = 2.0 * largest_signal[i].numpy()
                self.assertTrue(
                    random_min < fit[i].numpy() < random_max,
                    'Coeff {} fails {} < {} < {}'.format(
                        i, random_min, fit[i].numpy(), random_max))
    def test_signal_coeffs(self):
        """Check signal coefficients have correct # of all constants"""
        signal = bmfc.signal(bmfc.SM)

        self.assertEqual(48, len(signal))
        for i in range(48):
            self.assertFalse(bmfc.is_trainable(signal[i]))
    def test_fit_write_rows(self):
        """Check writing rows to new and existing fit files works as expected"""
        tmp_file = tempfile.mktemp()
        self.maxDiff = None

        signal_coeffs = bmfc.signal(bmfc.SM)

        fit_writer = bmfw.FitWriter(tmp_file, signal_coeffs)
        fit_writer.write_coeffs(1.2, bmfc.signal(bmfc.SM), 14.8)
        fit_writer.write_coeffs(3.4, bmfc.fit(12.345), 15.3)
        fit_writer.write_coeffs(5.6, bmfc.signal(bmfc.NP), 13.9)
        self._compare('csv_writer_fit_rows_first_write.csv', tmp_file,
                      'Non-existent file gets rows written correctly')

        fit_writer = bmfw.FitWriter(tmp_file, signal_coeffs)
        fit_writer.write_coeffs(7.8, bmfc.fit(67.890), 14.0)
        self._compare('csv_writer_fit_rows_append.csv', tmp_file,
                      'Existing file gets rows appended correctly')
    def test_fit_write_headers_only(self):
        """Check opening a fit file with only headers works correctly. In normal operation this should not happen"""
        tmp_file = tempfile.mktemp()
        self.maxDiff = None

        signal_coeffs = bmfc.signal(bmfc.SM)

        shutil.copyfile(
            self._test_data_path('csv_writer_fit_headers_only.csv'), tmp_file)
        bmfw.FitWriter(tmp_file, signal_coeffs)
        self._compare('csv_writer_fit_headers_and_signal.csv', tmp_file,
                      'Existing file with only headers gets signal written')
    def test_fit_coeffs_init_current_signal(self):
        """Check fit coefficients have correct #, the right ones are trainable,
            and that randomisation/defaults work
        """
        signal = bmfc.signal(bmfc.SM)
        fit = bmfc.fit(bmfc.FIT_INIT_CURRENT_SIGNAL, bmfc.SM)

        for i in range(48):
            if i in self.fit_trainable_ids:
                nt.assert_allclose(signal[i].numpy(),
                                   fit[i].numpy(),
                                   atol=0,
                                   rtol=1e-10)
    def test_fit_write_headers_and_signal(self):
        """Check that fit file header and signal writing works correctly"""
        tmp_file = tempfile.mktemp()
        self.maxDiff = None

        signal_coeffs = bmfc.signal(bmfc.SM)

        bmfw.FitWriter(tmp_file, signal_coeffs)
        self._compare('csv_writer_fit_headers_and_signal.csv', tmp_file,
                      'Non-existent file gets headers and signal')
        bmfw.FitWriter(tmp_file, signal_coeffs)
        self._compare(
            'csv_writer_fit_headers_and_signal.csv', tmp_file,
            'Existing file does not get duplicate headers or signal')
    def test_fit_coeffs_init_twice_current_signal_any_sign(self):
        """Check fit coefficients have correct #, the right ones are trainable,
            and that randomisation/defaults work
        """
        signal = bmfc.signal(bmfc.SM)
        fit = bmfc.fit(bmfc.FIT_INIT_TWICE_CURRENT_SIGNAL_ANY_SIGN, bmfc.SM)

        for i in range(48):
            if i in self.fit_trainable_ids:
                # Randomization should be enabled. Check from 0 to 2x SM value
                random_min = -2.0 * tf.math.abs(signal[i]).numpy()
                random_max = 2.0 * tf.math.abs(signal[i]).numpy()
                self.assertTrue(
                    random_min < fit[i].numpy() < random_max,
                    'Coeff {} fails {} < {} < {}'.format(
                        i, random_min, fit[i].numpy(), random_max))
    def test_q_write_rows(self):
        """Check writing rows to new and existing Q files works as expected"""
        tmp_file = tempfile.mktemp()
        self.maxDiff = None

        fit_writer = bmfw.QWriter(tmp_file)
        fit_writer.write_q(bmfc.signal(bmfc.NP), 12.34, bmfc.signal(bmfc.SM),
                           23.45, 34.56, 45.67)
        fit_writer.write_q(bmfc.signal(bmfc.NP), 23.45, bmfc.signal(bmfc.SM),
                           34.56, 45.67, 56.78)
        self._compare('csv_writer_q_rows_first_write.csv', tmp_file,
                      'Non-existent file gets rows written correctly')

        fit_writer = bmfw.QWriter(tmp_file)
        fit_writer.write_q(bmfc.signal(bmfc.NP), 34.56, bmfc.signal(bmfc.SM),
                           45.67, 56.78, 67.89)
        self._compare('csv_writer_q_rows_append.csv', tmp_file,
                      'Existing file gets rows appended correctly')
    def test_fit_coeffs_fix_p_wave(self):
        """Check fit coefficients works properly when passing `fix_p_wave_model`"""
        signal = bmfc.signal(bmfc.SM)
        fit = bmfc.fit(12.345, fix_p_wave_model=bmfc.SM)

        for i in range(48):
            if i in self.fit_trainable_ids:
                if i < 36:
                    # P-wave coeffs should be locked to the signal model values
                    nt.assert_allclose(signal[i].numpy(),
                                       fit[i].numpy(),
                                       atol=0,
                                       rtol=1e-10)
                    self.assertFalse(
                        bmfc.is_trainable(fit[i]),
                        'Coeff {} should not be trainable'.format(i))
                else:
                    # S-wave coeffs should be set to the constant value and be trainable
                    nt.assert_allclose(tf.constant(12.345).numpy(),
                                       fit[i].numpy(),
                                       atol=0,
                                       rtol=1e-10)
                    self.assertTrue(bmfc.is_trainable(fit[i]),
                                    'Coeff {} should be trainable'.format(i))
 def test_exception_raised_if_signal_does_not_match(self):
     """Check that signal coefficients must match when opening an existing file"""
     tmp_file = tempfile.mktemp()
     bmfw.FitWriter(tmp_file, bmfc.signal(bmfc.SM))
     with self.assertRaises(RuntimeError):
         bmfw.FitWriter(tmp_file, bmfc.signal(bmfc.NP))
 def test_generate_returns_correct_shape(self):
     """Check generate() returns a tensor of shape (events_total, 4)"""
     events = bmfs.generate(bmfc.signal(bmfc.SM), 123_456)
     self.longMessage = True
     self.assertEqual(123_456, tf.shape(events)[0].numpy())
     self.assertEqual(4, tf.shape(events)[1].numpy())
class TestSignal(unittest.TestCase):

    # Fields are: name, list of coefficients, "true" decay rate (As generated by odeint_fixed() attempt across all vars)
    # To add more to this list:
    #  1. Add a new line with "true" decay rate set to something like 1.0.
    #  2. Run the test_data/signal_integrator.py file to get the "true" values
    #  3. Set the found decay rate in your new line
    test_coeffs = [
        ('signal', bmfc.signal(bmfc.NP), 543.44989,),
        ('ones', [tf.constant(1.0)] * bmfc.count, 3082.3848,),
        ('integers', [tf.constant(float(i)) for i in range(int(-bmfc.count / 2), int(bmfc.count / 2))], 515583.66,),
        ('minus_point_ones', [tf.constant(-0.1)] * bmfc.count, 30.823853,),
    ]

    def test_decay_rate_integration_methods_approx_equal(self):
        """
        Check that the _integrate_decay_rate() method that integrates a previously angle-integrated decay rate
        over q^2 returns something approximately equal to running odeint_fixed() over all variables.

        Checks to within 1% as both methods use bins and add errors.
        """
        # Check for different lists of coefficients
        for c_name, coeffs, expected_decay_rate in self.test_coeffs:
            with self.subTest(c_name=c_name):
                actual = bmfs.integrate_decay_rate(coeffs)
                # Check values are the same to within 0.1%
                nt.assert_allclose(expected_decay_rate, actual.numpy(), atol=0, rtol=0.01)

    def test_integrate_decay_rate_within_tolerance(self):
        """
        Check that the tolerances set in _integrate_decay_rate() have not been relaxed so much that
        they mess up the accuracy more than 0.1% from using odeint() on the previously angle integrated decay rate.
        """
        for c_name, coeffs, _ in self.test_coeffs:
            with self.subTest(c_name=c_name):
                true = tf_integrate.odeint(
                    lambda _, q2: bmfs.decay_rate_angle_integrated(coeffs, q2),
                    0.0,
                    tf.stack([bmfs.q2_min, bmfs.q2_max]),
                )[1]

                ours = bmfs.integrate_decay_rate(coeffs)

                nt.assert_allclose(true.numpy(), ours.numpy(), atol=0, rtol=0.001)

    def test_generate_returns_correct_shape(self):
        """Check generate() returns a tensor of shape (events_total, 4)"""
        events = bmfs.generate(bmfc.signal(bmfc.SM), 123_456)
        self.longMessage = True
        self.assertEqual(123_456, tf.shape(events)[0].numpy())
        self.assertEqual(4, tf.shape(events)[1].numpy())

    def test_frac_s_methods_approx_equal(self):
        """
        Check decay_rate_frac_s() returns approximately the same values as a method that uses the moduli of the
        amplitudes. It is almost equal for the signal coefficients, but drifts more for other coeffs. I think this is
        because:
         * Coeffs are not arbitrary but actually depend on each other so artificial ones (e.g. integers)
            are not realistic
         *  a_t_l and a_t_r have been neglected
        """
        q2 = tf.linspace(bmfs.q2_min, bmfs.q2_max, 9)

        for c_name, coeffs, _ in self.test_coeffs:
            with self.subTest(c_name=c_name):
                decay_rate_frac_s = bmfs.decay_rate_frac_s(coeffs, q2)
                modulus_frac_s = bmfs.modulus_frac_s(coeffs, q2)
                for i in range(tf.shape(q2)[0].numpy()):
                    nt.assert_allclose(
                        decay_rate_frac_s[i].numpy(),
                        modulus_frac_s[i].numpy(),
                        atol=0,
                        rtol=0.025,
                        err_msg='{} not within 2.5% for q2 {}'.format(c_name, q2[i].numpy()),
                    )