inline documentation

Author: KOVALW

src/calibrationtools/particle.py

@@ -2,5 +2,13 @@
 
 
 class Particle(UserDict):
+    """
+    Particle is a subclass of `UserDict` that represents a particle with a specific state.
+
+    Attributes:
+        data (dict): The internal dictionary storing the state of the particle. This can
+        be accessed using the standard dictionary interface provided by `UserDict`.
+    """
+
     def __repr__(self):
         return f"Particle(state={self.data})"

src/calibrationtools/particle_population.py

@@ -4,11 +4,58 @@
 
 
 class ParticlePopulation:
+    """
+    ParticlePopulation is a class that represents a collection of particles, each with an associated weight.
+    It provides methods for managing the particles, normalizing their weights, and computing properties
+    such as the effective sample size (ESS).
+
+    Args:
+        states (Sequence[dict[str, Any]] | None): Optional initial states to create as Particle objects.
+        weights (Sequence[float] | None): Optional initial weights for the initial states.
+
+    Attributes:
+        particles (list[Particle]): A list of Particle objects in the population.
+        weights (list[float]): A list of weights corresponding to each particle.
+        ess (float): The effective sample size of the particle population.
+        size (int): The number of particles in the population.
+        total_weight (float): The sum of all particle weights.
+
+    Methods:
+        __init__(states, weights):
+            Initializes the ParticlePopulation with optional states and weights.
+        add_particle(particle, weight):
+            Adds a new particle to the population with the specified weight.
+        is_empty():
+            Checks if the particle population is empty.
+        normalize_weights():
+            Normalizes the weights of the particles so that they sum to 1.
+        __repr__():
+            Returns a string representation of the ParticlePopulation instance.
+
+    Errors:
+        ValueError: If the length of weights does not match the length of particles on initialization.
+    """
+
     def __init__(
         self,
         states: Sequence[dict[str, Any]] | None = None,
         weights: Sequence[float] | None = None,
     ):
+        """
+        Initializes a particle population with optional states and weights.
+
+        Args:
+            states (Sequence[dict[str, Any]] | None, optional): A sequence of dictionaries
+                representing the states of the particles. If None, an empty particle list
+                is initialized. Defaults to None.
+            weights (Sequence[float] | None, optional): A sequence of weights corresponding
+                to the particles. If None, all specified particle states are assigned equal
+                weights of 1.0. Defaults to None. Supplied weights are normalized to 1.0 upon
+                initialization.
+
+        Raises:
+            ValueError: If the length of the weights does not match the length of the particles.
+        """
         self._particles: list[Particle] = (
             [] if states is None else [Particle(x) for x in states]
         )
@@ -41,6 +88,19 @@ def add_particle(self, particle: Particle, weight: float):
 
     @property
     def ess(self) -> float:
+        """
+        Calculate the Effective Sample Size (ESS) of the particle population.
+
+        The ESS is a measure of the diversity of the particle weights. It is
+        calculated as the square of the total weight divided by the sum of the
+        squared weights. An ESS closer to the true size indicates a more uniform
+        distribution of weights, while a lower ESS indicates that the weights
+        are concentrated on fewer particles.
+
+        Returns:
+            float: The effective sample size. Returns 0.0 if the total weight
+            is zero.
+        """
         if self.total_weight == 0:
             return 0.0
         else:
@@ -60,6 +120,13 @@ def is_empty(self) -> bool:
         return self.size == 0
 
     def normalize_weights(self):
+        """
+        Normalize the weights of the particle population.
+
+        This method adjusts the weights of all particles in the population so that
+        their total sum equals 1. If the population is empty, the method exits
+        without performing any operation.
+        """
         if self.is_empty():
             return
         else:

src/calibrationtools/particle_updater.py

@@ -6,10 +6,44 @@
 from .perturbation_kernel import PerturbationKernel
 from .prior_distribution import PriorDistribution
 from .spawn_rng import spawn_rng
-from .variance_adapter import VarianceAdapter
+from .variance_adapter import AdaptIdentityVariance, VarianceAdapter
 
 
 class _ParticleUpdater:
+    """
+    A class responsible for managing and updating a population of particles
+    in an ABC-SMC framework. It provides functionality for sampling,
+    perturbing, and calculating weights for proposed particles, as well as
+    adapting the variance of the perturbation kernel.
+
+    Attributes:
+        perturbation_kernel (PerturbationKernel): The kernel used to perturb particles.
+        priors (PriorDistribution): The prior distribution of particle states. This remains fixed regardless of population changes.
+        variance_adapter (VarianceAdapter): The adapter used to adjust the variance of the perturbation kernel according to population particle state variance.
+        seed_sequence (SeedSequence | None): An optional seed sequence for random number generation.
+        particle_population (ParticlePopulation): The current population of particles.
+
+    Methods:
+        sample_particle() -> Particle:
+            Samples a particle from the current population based on their weights.
+
+        sample_and_perturb_particle(max_attempts: int = 10_000) -> Particle:
+            Samples a particle, perturbs it using the perturbation kernel, and returns
+            the perturbed particle. Raises a RuntimeError if a valid particle cannot
+            be sampled within the maximum number of attempts.
+
+        calculate_weight(particle: Particle) -> float:
+            Calculates the weight of a given particle based on the prior distribution
+            and the transition probabilities of the perturbation kernel.
+
+        adapt_variance():
+            Adapts the variance of the perturbation kernel based on the current particle population.
+
+    Raises:
+        ValueError: If the particle population is not set when attempting to sample a particle.
+        RuntimeError: If a valid perturbed particle cannot be sampled within the maximum number of attempts.
+    """
+
     def __init__(
         self,
         perturbation_kernel: PerturbationKernel,
@@ -18,6 +52,16 @@ def __init__(
         seed_sequence: SeedSequence | None = None,
         particle_population: ParticlePopulation | None = None,
     ):
+        """
+        Initializes the ParticleUpdater class.
+
+        Args:
+            perturbation_kernel (PerturbationKernel): The initial kernel used to perturb particles during proposals.
+            priors (PriorDistribution): The prior distribution used for calculating particle weights.
+            variance_adapter (VarianceAdapter): The adapter responsible for adjusting perturbation variance.
+            seed_sequence (SeedSequence | None, optional): A sequence of seeds for replicable random number generation. Defaults to None.
+            particle_population (ParticlePopulation | None, optional): An initial population of particles. If not provided, a new ParticlePopulation instance is created. Defaults to None.
+        """
         self.perturbation_kernel = perturbation_kernel
         self.priors = priors
         self.variance_adapter = variance_adapter
@@ -34,15 +78,34 @@ def particle_population(self) -> ParticlePopulation:
 
     @particle_population.setter
     def particle_population(self, particle_population: ParticlePopulation):
+        """
+        Updates the particle population and ensures its weights are normalized.
+
+        This method sets the particle population, normalizes its weights if the
+        total weight is not equal to 1.0, and adapts the perturbation variance
+        according to the new stored particle population.
+
+        Args:
+            particle_population (ParticlePopulation): The particle population to update.
+        """
         self._particle_population = particle_population
         if self._particle_population.total_weight != 1.0:
             self._particle_population.normalize_weights()
         self.adapt_variance()
 
     def sample_particle(self) -> Particle:
-        if not hasattr(self, "particle_population"):
+        """
+        Samples a particle from the particle population based on their weights.
+
+        Returns:
+            Particle: The sampled particle from the particle population.
+
+        Raises:
+            ValueError: If the particle population is not set.
+        """
+        if self.particle_population.is_empty():
             raise ValueError(
-                "Particle population is not set. Please set the particle population before sampling."
+                "Particle population is empty. Please add entries to the particle population before sampling."
             )
         idx = spawn_rng(self.seed_sequence).choice(
             self.particle_population.size,
@@ -53,6 +116,24 @@ def sample_particle(self) -> Particle:
     def sample_and_perturb_particle(
         self, max_attempts: int = 10_000
     ) -> Particle:
+        """
+        Samples a particle from the current population and applies a perturbation to it,
+        ensuring the perturbed particle satisfies the prior probability density constraints.
+        If a perturbed particle fails to meet the prior constraints, a new particle is
+        sampled with replacement and perturbed until a valid particle is obtained or the
+        maximum number of attempts is reached.
+
+        Args:
+            max_attempts (int): The maximum number of attempts to sample and perturb
+                a particle before aborting. Defaults to 10,000.
+
+        Returns:
+            Particle: A new particle object created from the perturbed particle.
+
+        Raises:
+            RuntimeError: If the method fails to sample and perturb a particle
+                within the specified maximum number of attempts.
+        """
         for _ in range(max_attempts):
             current_particle = self.sample_particle()
             new_particle = self.perturbation_kernel.perturb(
@@ -65,6 +146,17 @@ def sample_and_perturb_particle(
         )
 
     def calculate_weight(self, particle: Particle) -> float:
+        """
+        Calculate the weight of a proposed particle based on the prior probability
+        and the weighted transition probabilities from the particles of the current population.
+
+        Args:
+            particle (Particle): The particle for which the weight is to be calculated.
+
+        Returns:
+            float: The calculated weight of the particle. Returns 0.0 if the denominator
+            (weighted sum of transition probabilities) is zero.
+        """
         numerator = self.priors.probability_density(particle)
         transition_probs = np.array(
             [
@@ -86,6 +178,24 @@ def calculate_weight(self, particle: Particle) -> float:
             return 0.0
 
     def adapt_variance(self):
+        """
+        Adjusts the variance of the particle population using the variance adapter.
+
+        This method utilizes the `variance_adapter` to adapt the variance of the
+        `perturbation_kernel` based on the current `particle_population`. The
+        perturbation kernel parameters are modified during this call but the
+        particle population remaions the same.
+
+        Raises:
+            ValueError: If `variance_adapter` is not an AdaptIdentityVariance and
+            `particle_population` is empty, adapt variance will fail.
+        """
+        if self.particle_population.is_empty() and not isinstance(
+            self.variance_adapter, AdaptIdentityVariance
+        ):
+            raise ValueError(
+                "Particle population is empty and variance adapter depends on population variance. Please add entries to the particle population or use `AdaptIdentityVariance` class."
+            )
         self.variance_adapter.adapt(
             self.particle_population, self.perturbation_kernel
         )