Integration Test for Model Training and Energy Estimation

[davidfitzek]

Description

As a part of improving test coverage and ensuring the correctness of the training process, adding an integration test that trains a small model on a dataset and verifies the energy estimation using the get_rydberg_energy function would be useful.

Acceptance Criteria

1. Data Preparation:

   • Use the dataset located at src/rydberggpt/tests/dataset_test/.

2. Model Training:

   • Train a small model on the specified dataset for a predefined number of steps (e.g., 100 steps).
   • Ensure the training process completes without errors.

3. Energy Estimation:

   • Use the trained model to estimate the energy of the system using the RydbergDecoderWavefunction.get_rydberg_energy function. We can estimate the energy for a configuration that is present in the dataset.
   • Validate the estimated energy against the true energy values based on the data.

Additional Context

This test will serve as a crucial validation check for future changes, ensuring that the training process and energy estimation are functioning as expected.

Implementation Suggestions

• It may be beneficial to create a simplified model architecture for this test to reduce the execution time.
• The test could be structured as a new pytest function or module within the existing testing framework.
• It may be worth considering setting a random seed to ensure reproducibility between test runs.

Questions

1. What should be the acceptable margin of error between the estimated and true energy values?
2. Are there any specific logging or reporting requirements for this test?