NeutroWater API

class DiffusingNeutrons:
    """
    Class that simulates multiple neutrons from diffusing in a medium.

    Args:
        p (Parameters): Parameters for the simulation.
    """

    def __init__(
        self,
        p: Parameters,
    ):
        self.kT = (
            1.380649e-23 * p.temperature * 6.24150907 * 10**18
        )  # [m^2 kg / (s^2 K) * K * (eV/J) = J*(eV/J) = eV]
        self.nCollisions = 0
        self.mol_struc = p.molecule_structure
        self.tank = Tank(p.radius_tank, p.height_tank, p.position_tank)

        # For interpolating the total cross section data and computing the
        # mean-free-path.
        self.total_processor = TotalProcessor(p.total_data)

        # For interpolating the scattering and absorption cross section and
        # computing the aborption ratio.
        self.absorption_processor = AbsorptionProcessor(
            p.scattering_data, p.absorption_data
        )

        # Sample from the interpolated energy spectrum
        spectrum_processor = SpectrumProcessor(p.spectrum_data)
        initial_energies = spectrum_processor.sample(num_samples=p.nNeutrons)

        # Convert the energies MeV -> eV
        initial_energies = [energy * 10**6 for energy in initial_energies]

        # All neutrons start at the origin
        initial_positions = [np.array([0, 0, 0]) for _ in range(p.nNeutrons)]
        self.neutrons = Neutrons(initial_energies, initial_positions)

        # For handling the collisions with the atomic nuclei, hydrogen and oxygen.
        self.nuclei_masses = p.nuclei_masses

        # Maxwell-Boltzmann distribution for the thermal energy
        self.mw = MaxwellBoltzmann(T=p.temperature)

        # For handling the angular distribution of the scattering.
        self.angular_processor = AngularProcessor(p.angular_data, p.nuclei_masses)

    def diffuse(self, nCollisions: int):
        """
        Let the neutrons diffuse in the medium.

        Args:
            nCollisions (int): number of times each neutron collides with an
                atomic nucleus.
        """
        self.nCollisions += nCollisions
        # Using not all cores but 2 less than the total number of cores, tends
        # to be faster
        num_processes = multiprocessing.cpu_count() - 2

        # Split the neutrons into chunks
        chunk_size = (len(self.neutrons) + num_processes - 1) // num_processes
        chunks = [
            self.neutrons[i : i + chunk_size]
            for i in range(0, len(self.neutrons), chunk_size)
        ]

        # Diffuse the chuncks of neutrons in parallel
        with multiprocessing.Pool(processes=num_processes) as pool:
            results = list(
                process_map(self._diffuse_chunk, chunks, [nCollisions] * len(chunks)),
            )

        # Update the neutrons list with the results from the parallel processess
        self.neutrons = [neutron for chunk in results for neutron in chunk]

    def _diffuse_chunk(self, chunk: Neutrons, nCollisions: int):
        """
        Diffuse a chunk of all neutrons in the medium.

        Args:
            chunk (Sequence[Neutron]): chunk of neutrons to diffuse
            nCollisions (int): number of times each neutron collides with an atomic
            nucleus.

        Returns: the neutrons after diffusing in the medium.
        """
        np.random.seed()  # Set a new seed for each process
        return [self._diffuse_neutron(neutron, nCollisions) for neutron in chunk]

    def _diffuse_neutron(self, neutron: Neutron, nCollisions: int) -> Neutron:
        """
        Diffuse a single neutron in the medium.

        Args:
            neutron (Neutron): Neutron to diffuse
            nCollisions (int): number of times the neutron collides with an atomic
            nucleus.

        Returns: the neutron after diffusing in the medium.
        """
        neutron.direction = random_direction()

        for _ in range(nCollisions):

            if not self.tank.inside(neutron.position):
                break

            # Sample the distance the neutron travels from an exponential distribution
            # with mean free path as the scale parameter.
            l = -self.total_processor.get_mfp(neutron.energy) * np.log(random())
            neutron.travel(l)

            # Determine the nucleus the neutron collides with
            mass = (
                self.nuclei_masses[0]
                if random() < self.total_processor.get_ratio(neutron.energy)
                else self.nuclei_masses[1]
            )

            collision = Collision(
                initial_E=neutron.energy,
                initial_direction=neutron.direction,
                mass=mass,
                scattering_cosine=self.angular_processor.get_CM_cosines(
                    mass, neutron.energy, 1
                )[0],
                absorption=self._absorbed(neutron, self.nuclei_masses.index(mass)),
                thermal=neutron.energy < 10 * self.kT,
            )

            # Update the neutron's energy and direction
            neutron.collide(collision.energy_loss_frac, collision.scattering_direction)

            # If the neutron is absorbed, break the loop
            if collision.energy_loss_frac == 1:
                break

        return neutron

    def _absorbed(self, neutron: Neutron, index: int) -> bool:
        """
        Check if the neutron is absorbed by the nucleus.

        Args:
            neutron (Neutron): Neutron to check.
            index (int): index of the nucleus in the molecule.

        Returns: True if the neutron is absorbed, False otherwise.
        """
        if (
            random()
            < self.absorption_processor.get_absorption_rates(neutron.energy)[index]
        ):
            return True
        else:
            return False

@dataclass
class Parameters:
    """
    Class to store the parameters for the simulation.

    Args:
        nNeutrons (int): Number of neutrons to simulate
        molecule_structure (Sequence): number of atoms in the molecule.
        nuclei_masses (Sequence): Masses of the nuclei in the molecule, should be
            in the same order as the molecule_structure.
        radius_tank (float): Radius of the tank in meters.
        height_tank (float): Height of the tank in meters.
        position_tank (np.ndarray): Position of the tank.
        temperature (float): Temperature [K] of the medium.
    """

    nNeutrons: int
    molecule_structure: Sequence = (2, 1)
    nuclei_masses: Sequence = (1, 16)
    radius_tank: float = 1
    height_tank: float = 1
    position_tank: tuple = (0.0, 0.0, 0.0)
    temperature: float = 293

    def __post_init__(self):
        files = resources.files(data)

        self.total_data: Sequence[pd.DataFrame] = [
            pd.read_csv(files / "h_cross_t.txt", sep=r"\s+"),
            pd.read_csv(files / "o_cross_t.txt", sep=r"\s+"),
        ]
        self.scattering_data: Sequence[pd.DataFrame] = [
            pd.read_csv(files / "h_cross_s.txt", sep=r"\s+"),
            pd.read_csv(files / "o_cross_s.txt", sep=r"\s+"),
        ]
        self.absorption_data: Sequence[pd.DataFrame] = [
            pd.read_csv(files / "h_cross_a.txt", sep=r"\s+"),
            pd.read_csv(files / "o_cross_a.txt", sep=r"\s+"),
        ]
        self.angular_data: Sequence[pd.DataFrame] = [
            pd.read_csv(files / "H_angular.txt", sep=r"\s+"),
            pd.read_csv(files / "O_angular.txt", sep=r","),
        ]
        self.spectrum_data: pd.DataFrame = pd.read_csv(
            files / "neutron_spectrum_normalized.txt", sep=","
        )

class Measurer:
    """Class to measure the properties of the neutrons after the simulation.

    Methods:
        positions: Get the positions of the neutrons.
        energies: Get the energies of the neutrons.
        number_total: Get the total number of neutrons.
        number_escaped: Get the number of neutrons that escaped the tank.
        number_thermal: Get the number of neutrons that are thermalized.
        thermalize_positions: Get the positions of the neutrons where their energy is thermal for the first time.
        thermalize_distances: Get the distance of the neutrons where their energy is thermal for the first time.
        number_absorbed: Get the number of neutrons that were absorbed.
        absorbed_positions: Get the positions of the neutrons that were absorbed.
        absorbed_distances: Get the distances of the neutrons that were absorbed.
        flux: Get the flux of the neutrons at a given radius r.
        energy_spectrum: Get the energy spectrum of the neutrons at a given radius r.
        energy_spectrum_escaped: Get the energy spectrum of the neutrons that escaped the tank.
        number_above_energy: Get the number of neutrons with energy above a given energy.

    """

    def __init__(self, sim: DiffusingNeutrons):
        self.neutrons = sim.neutrons
        self.sim = sim

    def positions(self) -> list[list[Vector]]:
        """
        Get the positions of the neutrons.

        Returns:
            a list where each element is a list of length nNeutrons
                [[np.ndarray, np.ndarray, ...], [np.ndarray, np.ndarray, ...], ...]
        """
        return [neutron.positions for neutron in self.neutrons]

    def energies(self) -> list[list[float]]:
        """
        Get the energies of the neutrons.

        Returns:
            a list where each element is a list of length nNeutrons
                [[float, float, ...], [float, float, ...], ...]
        """
        return [neutron.energies for neutron in self.neutrons]

    def number_total(self) -> int:
        """
        Get the total number of neutrons.

        Returns:
            an int.
        """
        return len(self.neutrons)

    def number_escaped(self) -> int:
        """
        Get the number of neutrons that escaped the tank.

        Returns:
            an int.
        """
        return sum(
            [
                1
                for neutron in self.neutrons
                if not self.sim.tank.inside(neutron.position)
            ],
            start=0,
        )

    def number_thermal(self) -> int:
        """
        Get the number of neutrons that are thermalized.

        Returns:
            an int.
        """
        n = 0
        for energies in self.energies():
            if energies[-1] < (10 * self.sim.kT):
                n += 1
        return n

    def thermalize_positions(self) -> list:
        """
        Get the positions of the neutrons where their energy is thermal for
        the first time.

        Returns:
            a list of np.ndarray.
        """
        thermalize_positions = []
        for energies, positions in zip(self.energies(), self.positions()):
            # Get the first element of the list of energies that is less than 10kT
            res = list(filter((lambda val: val < (10 * self.sim.kT)), energies))
            if len(res) > 0:
                index = energies.index(res[0])
                thermalize_positions.append(positions[index])

        return thermalize_positions

    def thermalize_distances(self) -> list:
        """
        Get the distance of the neutrons where their energy is thermal for
        the first time.

        Returns:
            a list of floats.
        """
        return [np.linalg.norm(pos) for pos in self.thermalize_positions()]

    def number_absorbed(self) -> int:
        """
        Get the number of neutrons that were absorbed.

        Returns:
            an int.
        """
        n = 0
        for neutron in self.neutrons:
            if len(neutron.positions) < self.sim.nCollisions and self.sim.tank.inside(
                neutron.position
            ):
                n += 1
        return n

    def absorbed_positions(self) -> list[Vector]:
        """
        Get the positions of the neutrons that were absorbed.

        Returns:
            a list of np.ndarray.
        """
        return [
            neutron.positions[-1]
            for neutron in self.neutrons
            if len(neutron.positions) < self.sim.nCollisions
            and self.sim.tank.inside(neutron.position)
        ]

    def absorbed_distances(self) -> list[np.float64]:
        """
        Get the distances of the neutrons that were absorbed.

        Returns:
            a list of floats.
        """
        return [np.linalg.norm(pos) for pos in self.absorbed_positions()]

    def flux(self, r: float) -> float:
        """
        Get the flux of the neutrons at a given radius r.

        Args:
            r (float): radius at which to compute the flux.

        Returns:
            a float value of the flux.
        """
        count = 0
        for positions in self.positions():
            positions = np.array(positions)
            norms = np.linalg.norm(positions, axis=1)
            # Count the number of neutrons that have been in the shell
            count += np.sum((norms[:-1] < r) & (r < norms[1:])) + np.sum(
                (norms[1:] < r) & (r < norms[:-1])
            )
        return count / (4 * np.pi * r**2)

    def energy_spectrum(self, r: float) -> list[float]:
        """
        Get the energy spectrum of the neutrons at a given radius r.

        Args:
            r (float): radius at which to compute the energy spectrum.

        Returns:
            a list of floats.
        """
        result_energies = []
        for positions, energies in zip(self.positions(), self.energies()):
            positions = np.array(positions)
            energies = np.array(energies)
            norms = np.linalg.norm(positions, axis=1)
            # Create boolean arrays for the conditions
            cond1 = np.concatenate([(norms[:-1] < r) & (r < norms[1:]), [False]])
            cond2 = np.concatenate([[False], (norms[1:] < r) & (r < norms[:-1])])
            # Get the energies of the neutrons that have been in the shell
            result_energies += list(energies[cond1]) + list(energies[cond2])
        return result_energies

    def energy_spectrum_escaped(self) -> list[float]:
        """
        Get the energy spectrum of the neutrons that escaped the tank.

        Returns:
            a list of floats.
        """
        return [
            neutron.energy
            for neutron in self.neutrons
            if not self.sim.tank.inside(neutron.position)
        ]

    def number_above_energy(self, E: float) -> int:
        """
        Get the number of neutrons with energy above a given energy.

        Args:
            E (float): energy threshold.

        Returns:
            an int.
        """
        return sum([1 for neutron in self.neutrons if neutron.energies[0] > E])