Understanding Poisson Distributions with PyTorch's torch.distributions.poisson.Poisson

Purpose

• Represents the Poisson distribution, a discrete probability distribution that models the number of events occurring in a fixed interval of time or space, with the assumption that these events occur independently and at a constant rate.

Functionality

• Calculates probability density functions (PDFs) for possible outcomes.
• Generates random samples (non-negative integers) from the distribution.
• Creates a Poisson distribution object parameterized by the rate (intensity or average number of events).

Key Points

• Output
  
  • Samples
    Generates a Tensor of the same size as rate, where each element is a random integer drawn from the Poisson distribution with the corresponding rate.
  • Probabilities
    Calculates the probability of each possible number of events (0, 1, 2, ...) using the Poisson probability mass function (PMF).
• Input
  Takes a non-negative Tensor rate specifying the event rate for each element.

Code Example

import torch
from torch.distributions import Poisson

# Example rate (tensor with two elements)
rate = torch.tensor([3.5, 1.2])

# Create Poisson distribution objects
dist1 = Poisson(rate=rate[0])  # For rate 3.5
dist2 = Poisson(rate=rate[1])  # For rate 1.2

# Sample 5 values from each distribution
samples1 = dist1.sample(torch.Size([5]))
samples2 = dist2.sample(torch.Size([5]))

# Print the sampled values
print(samples1)
print(samples2)

# Calculate probability of 2 events for each distribution
prob1 = dist1.pmf(2)  # Probability of 2 events with rate 3.5
prob2 = dist2.pmf(2)  # Probability of 2 events with rate 1.2

print(prob1)
print(prob2)

1. We import torch and Poisson from torch.distributions.
2. We create a rate tensor with two elements (3.5 and 1.2).
3. We instantiate two Poisson distribution objects, dist1 and dist2, with the respective rates.
4. We use sample to generate five random integers from each distribution.
5. We calculate the probability of having exactly two events for each distribution using pmf.

• PyTorch Tensors enable efficient GPU acceleration for computations.
• torch.distributions provides various probability distributions, including continuous and discrete ones.

---

Calculating Probability of Multiple Events

This code shows how to calculate the probabilities for different numbers of events (not just a single value):

import torch
from torch.distributions import Poisson

# Example rate
rate = torch.tensor(5.0)

# Create Poisson distribution
dist = Poisson(rate=rate)

# Calculate probabilities for 0, 1, 2, and 3 events
probs = dist.pmf(torch.tensor([0, 1, 2, 3]))

print(probs)

Batch Processing with Multiple Rates

This example demonstrates how to create a Poisson distribution object for a batch of rates and then generate samples:

import torch
from torch.distributions import Poisson

# Example rates (batch of 2)
rates = torch.tensor([2.0, 4.0])

# Create Poisson distribution (batch processing)
dist = Poisson(rate=rates)

# Generate 3 samples from each rate in the batch
samples = dist.sample(torch.Size([3, 2]))

print(samples)

Calculating Log Probabilities

While pmf gives the raw probabilities, log_prob provides the natural logarithm of the probabilities:

import torch
from torch.distributions import Poisson

# Example rate
rate = torch.tensor(3.0)

# Create Poisson distribution
dist = Poisson(rate=rate)

# Integer value for which to calculate log probability
value = torch.tensor(2)

log_prob = dist.log_prob(value)

print(log_prob)

Using the mean and variance Properties

These properties provide the mean and variance of the Poisson distribution for the given rate(s):

import torch
from torch.distributions import Poisson

# Example rate
rate = torch.tensor(1.5)

# Create Poisson distribution
dist = Poisson(rate=rate)

# Get mean and variance
mean = dist.mean
variance = dist.variance

print("Mean:", mean)
print("Variance:", variance)

---

1. Manual Implementation
   

   While less efficient than torch.distributions.poisson.Poisson, you can implement the Poisson probability mass function (PMF) yourself:

   import torch
   
   def poisson_pmf(k, rate):
       """Calculates the Poisson PMF for a given number of events (k) and rate."""
       return (rate**k * torch.exp(-rate)) / torch.math.factorial(k)
   
   # Example usage
   k = torch.tensor(3)  # Number of events
   rate = torch.tensor(2.5)  # Rate
   
   prob = poisson_pmf(k, rate)
   print(prob)
   

   This approach offers more control over the calculation but is less optimized and requires manual vectorization for efficient computation.

2. torch.random.poisson (if available)
   

   Some PyTorch versions might include torch.random.poisson, which directly samples from the Poisson distribution:

   import torch
   
   # Example usage (if available)
   samples = torch.random.poisson(rate=3.0, size=(5,))
   print(samples)
   

   However, this method might not be consistently available across all PyTorch versions.

3. NumPy Integration (if using NumPy)
   

   If you're already using NumPy in your project, you can leverage its np.random.poisson function and convert the results to PyTorch tensors:

   import torch
   import numpy as np
   
   # Example usage
   rate = 1.8
   samples = torch.from_numpy(np.random.poisson(rate, size=10))
   print(samples)
   

   This approach is helpful if your workflow involves both NumPy and PyTorch, but it adds an extra conversion step.

• If you're already using NumPy, integrating with np.random.poisson can be convenient, but it adds an extra conversion step.
• Consider torch.random.poisson if available for its simplicity, but be aware of potential version-specific limitations.
• If you need a more lightweight implementation and understand the PMF calculation, the manual approach might be suitable for smaller-scale tasks.
• For most cases, torch.distributions.poisson.Poisson is the recommended approach due to its efficiency, built-in functionalities like sample and log_prob, and seamless integration with PyTorch.