<filename>tests/test_exact_distribution_calculator.py
__author__ = "<NAME>"

import unittest

from numpy import array, complex128, int64, allclose

from theboss.boson_sampling_utilities.permanent_calculators.bs_permanent_calculator_factory import (
    BSPermanentCalculatorFactory,
    PermanentCalculatorType,
)
from theboss.distribution_calculators.bs_exact_distribution_with_uniform_losses import (
    BSDistributionCalculatorWithUniformLosses,
    BSDistributionCalculatorWithFixedLosses,
    BosonSamplingExperimentConfiguration,
)


class TestExactLossyBosonSamplingDistributionCalculator(unittest.TestCase):
    def setUp(self) -> None:
        # Define some additional variables for more clear experiment configuration assignment.
        self.number_of_particles_lost = 2

        # Generate permutation matrix and define initial state.
        self.permutation_matrix = array(
            [
                [0, 0, 1, 0, 0],
                [1, 0, 0, 0, 0],
                [0, 0, 0, 1, 0],
                [0, 0, 0, 0, 1],
                [0, 1, 0, 0, 0],
            ],
            dtype=complex128,
        )

        self.initial_state = array([1, 1, 1, 0, 0], dtype=int64)

        # Create configuration object.
        self.experiment_configuration = BosonSamplingExperimentConfiguration(
            interferometer_matrix=self.permutation_matrix,
            initial_state=self.initial_state,
            number_of_modes=len(self.initial_state),
            initial_number_of_particles=sum(self.initial_state),
            number_of_particles_lost=self.number_of_particles_lost,
            number_of_particles_left=sum(self.initial_state)
            - self.number_of_particles_lost,
            uniform_transmissivity=0.8,
        )

        self._calculator_type = PermanentCalculatorType.PARALLEL_CHIN_HUH
        self._permanent_calculator_factory = BSPermanentCalculatorFactory(
            matrix=None,
            input_state=None,
            output_state=None,
            calculator_type=self._calculator_type,
        )
        self._permanent_calculator = (
            self._permanent_calculator_factory.generate_calculator()
        )

    def test_probabilities_sum_in_fixed_losses_scenario(self) -> None:
        exact_distribution_calculator = BSDistributionCalculatorWithFixedLosses(
            self.experiment_configuration, self._permanent_calculator
        )
        exact_distribution = exact_distribution_calculator.calculate_distribution()
        self.assertAlmostEqual(sum(exact_distribution), 1.0, delta=1e-4)

    def test_probabilities_sum_in_uniform_losses_scenario(self) -> None:
        exact_distribution_calculator = BSDistributionCalculatorWithUniformLosses(
            self.experiment_configuration, self._permanent_calculator
        )
        exact_distribution = exact_distribution_calculator.calculate_distribution()
        self.assertAlmostEqual(sum(exact_distribution), 1.0, delta=1e-4)

    def test_probabilities(self) -> None:
        exact_distribution_calculator = BSDistributionCalculatorWithUniformLosses(
            self.experiment_configuration, self._permanent_calculator
        )
        exact_distribution = exact_distribution_calculator.calculate_distribution()
        self.assertTrue(
            allclose(
                exact_distribution,
                [
                    0.008,
                    0.032,
                    0.032,
                    0.0,
                    0.0,
                    0.032,
                    0.0,
                    0.128,
                    0.0,
                    0.0,
                    0.128,
                    0.0,
                    0.0,
                    0.0,
                    0.128,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.512,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                    0.0,
                ],
            )
        )
