artist.raytracing.sampling
==========================

.. py:module:: artist.raytracing.sampling


Classes
-------

.. autoapisummary::

   artist.raytracing.sampling.DistortionsDataset
   artist.raytracing.sampling.RestrictedDistributedSampler


Module Contents
---------------

.. py:class:: DistortionsDataset(light_source: artist.scene.light_source.LightSource, number_of_points_per_heliostat: int, number_of_active_heliostats: int, random_seed: int = 7)

   Bases: :py:obj:`torch.utils.data.Dataset`


   Initialize the dataset.

   This class implements a custom dataset according to the ``torch`` interface. The content of this
   dataset are the distortions. The distortions are used in our version of "heliostat"-tracing to
   indicate how each incoming ray must be multiplied and scattered on the heliostat. According to
   ``torch``, this dataset must implement a function to return the length of the dataset and one function
   to retrieve an element through an index.

   Parameters
   ----------
   light_source : LightSource
       The light source used to model the distortions.
   number_of_points_per_heliostat : int
       The number of points on the heliostats for which distortions are created.
   number_of_active_heliostats : int
       The number of active heliostats in the scenario.
   random_seed : int
       The random seed used for generating the distortions (default is 7).


   .. py:method:: __len__() -> int

      Calculate the length of the dataset, i.e., the number of items contained.

      Returns
      -------
      int
          The length of the dataset.



   .. py:method:: __getitem__(idx: int) -> tuple[torch.Tensor, torch.Tensor]

      Select an item from the dataset.

      Parameters
      ----------
      idx : int
          The index of the item to select.

      Returns
      -------
      torch.Tensor
          The distortions in the up direction for the given index.
      torch.Tensor
          The distortions in the east direction for the given index.



.. py:class:: RestrictedDistributedSampler(number_of_samples: int, number_of_active_heliostats: int, world_size: int = 1, rank: int = 0)

   Bases: :py:obj:`torch.utils.data.Sampler`


   Set up a custom distributed sampler to assign data to each rank or leave them idle.

   Parameters
   ----------
   number_of_samples : int
       Length of the dataset or total number of samples.
   number_of_active_heliostats : int
       Number of active heliostats.
   world_size : int
       World size or total number of processes (default is 1).
   rank : int
       Rank of the current process (default is 0).


   .. py:attribute:: rank_indices
      :value: []



   .. py:method:: __iter__() -> Iterator[int]

      Generate a sequence of indices for the current rank's portion of the dataset.

      Returns
      -------
      Iterator[int]
          An iterator over indices for the current rank.