test: add test for generate_data #148
Conversation
tests/data/test_generate.py
Outdated
| def test_generate_data(canonical_model: es.AmplitudeModel): | ||
| n_phsp = 1000 | ||
| n_data = 100 | ||
| model = canonical_model | ||
| kinematics = HelicityKinematics.from_model(model) | ||
| phsp_sample = generate_phsp(n_phsp, kinematics) | ||
| builder = IntensityBuilder(model.particles, kinematics, phsp_sample) | ||
| intensity = builder.create_intensity(model) | ||
| data_sample = generate_data(n_data, kinematics, intensity) | ||
| assert len(data_sample) == len(model.kinematics.final_state) | ||
| for sample in data_sample: | ||
| assert len(sample) == n_data | ||
| data_sq = data_sample ** 2 | ||
| e_sq = data_sq[:, :, 3] | ||
| p3_sq = data_sq[:, :, :3] | ||
| m_sq = np.abs(e_sq - p3_sq.sum(axis=2)) | ||
| assert pytest.approx(list(np.sqrt(m_sq.mean(axis=1))), abs=1e-4) == [ | ||
| 0, | ||
| 0.135, | ||
| 0.135, | ||
| ] | ||
| data_set = kinematics.convert(data_sample) | ||
| assert set(data_set) == { | ||
| "mSq_2", | ||
| "mSq_2_3_4", | ||
| "mSq_3", | ||
| "mSq_3_4", | ||
| "mSq_4", | ||
| "phi+3+4_vs_2", | ||
| "phi+3_4+2", | ||
| "theta+3+4_vs_2", | ||
| "theta+3_4+2", | ||
| } | ||
| assert pytest.approx(data_set["mSq_2"].mean()) == 0 | ||
| assert pytest.approx(data_set["mSq_3"].mean()) == data_set["mSq_4"].mean() | ||
| assert pytest.approx(data_set["mSq_3_4"].mean(), abs=1e-1) == 1 |
There was a problem hiding this comment.
I would test these kinematic related things like masses inside the phase space generation, since that is related to that. The generate data is related to the model, so we should perform checks that test for some expected distribution or so.
There was a problem hiding this comment.
Which checks do you suggest? (The masses are already checked in test_generate_phsp.)
There was a problem hiding this comment.
There are some distribution checks in pycompwa that we could move here https://github.com/ComPWA/pycompwa/tree/master/tests/angular-distribution-tests . But they are slighly more complicated though and took a few more seconds to run.
For now we could just check that the sample sizes match and extend the tests to distribution verifications later on
Needed for #146