ruff formatting

Author: ColtAllen

pymc_extras/distributions/discrete.py

@@ -15,8 +15,8 @@
 import numpy as np
 import pymc as pm
 
-from pymc.distributions.distribution import Discrete
 from pymc.distributions.dist_math import betaln, check_parameters, factln, logpow
+from pymc.distributions.distribution import Discrete
 from pymc.distributions.shape_utils import rv_size_is_none
 from pytensor import tensor as pt
 from pytensor.tensor.random.op import RandomVariable
@@ -432,10 +432,10 @@ def sim_data(lam):
                 lam,  # For small lambda, p ≈ lambda
                 1 - np.exp(-lam)  # Standard formula for larger lambda
             )
-            
+
             # Ensure p is in valid range for geometric distribution
             p = np.clip(p, 1e-5, 1.)
-            
+
             t = rng.geometric(p)
             return np.array([t])
 
@@ -529,15 +529,15 @@ def logcdf(value, r, alpha):
 
     def support_point(rv, size, r, alpha):
         """Calculate a reasonable starting point for sampling.
-        
+
         For the GrassiaIIGeometric distribution, we use a point estimate based on
         the expected value of the mixing distribution. Since the mixing distribution
         is Gamma(r, 1/alpha), its mean is r/alpha. We then transform this through
         the geometric link function and round to ensure an integer value.
         """
         mean = pt.ceil(pt.exp(alpha/r))
-        
+
         if not rv_size_is_none(size):
             mean = pt.full(size, mean)
-        
-        return mean
+
+        return mean

tests/distributions/test_discrete.py

@@ -33,8 +33,8 @@
 from pymc_extras.distributions import (
     BetaNegativeBinomial,
     GeneralizedPoisson,
-    Skellam,
     GrassiaIIGeometric,
+    Skellam,
 )
 
 
@@ -233,7 +233,7 @@ def test_random_basic_properties(self):
                     1 - np.exp(-np.random.gamma(r, 1/alpha, size=size)), size=size
                 ),
             )
-            
+
             # Test small parameter values that could generate small lambda values
             discrete_random_tester(
                 dist=self.pymc_dist,
@@ -278,14 +278,14 @@ def test_logp_basic(self):
         test_value = np.array([1, 2, 3, 4, 5])
         test_r = 1.0
         test_alpha = 1.0
-        
+
         logp_vals = logp_fn(test_value, test_r, test_alpha)
         assert not np.any(np.isnan(logp_vals))
         assert np.all(np.isfinite(logp_vals))
 
         # Test invalid values
         assert logp_fn(np.array([0]), test_r, test_alpha) == np.inf  # Value must be > 0
-        
+
         with pytest.raises(TypeError):
             logp_fn(np.array([1.5]), test_r, test_alpha) == -np.inf  # Value must be integer
 
@@ -328,16 +328,16 @@ def test_support_point(self, r, alpha, size, expected_shape):
         """Test that support_point returns reasonable values with correct shapes"""
         with pm.Model() as model:
             GrassiaIIGeometric("x", r=r, alpha=alpha, size=size)
-        
+
         init_point = model.initial_point()["x"]
-        
+
         # Check shape
         assert init_point.shape == expected_shape
-        
+
         # Check values are positive integers
         assert np.all(init_point > 0)
         assert np.all(init_point.astype(int) == init_point)
-        
+
         # Check values are finite and reasonable
         assert np.all(np.isfinite(init_point))
         assert np.all(init_point < 1e6)  # Should not be extremely large