make tests pass - formatting of ZeroSumNormal.py

Author: drbenvincent

examples/generalized_linear_models/ZeroSumNormal.py

@@ -10,12 +10,13 @@
 from scipy import stats
 from pymc3.distributions.distribution import generate_samples, draw_values
 
+
 def extend_axis_aet(array, axis):
     n = array.shape[axis] + 1
     sum_vals = array.sum(axis, keepdims=True)
     norm = sum_vals / (np.sqrt(n) + n)
     fill_val = norm - sum_vals / np.sqrt(n)
-    
+
     out = aet.concatenate([array, fill_val.astype(str(array.dtype))], axis=axis)
     return out - norm.astype(str(array.dtype))
 
@@ -27,7 +28,7 @@ def extend_axis_rev_aet(array: np.ndarray, axis: int):
 
     n = array.shape[axis]
     last = aet.take(array, [-1], axis=axis)
-    
+
     sum_vals = -last * np.sqrt(n)
     norm = sum_vals / (np.sqrt(n) + n)
     slice_before = (slice(None, None),) * axis
@@ -39,15 +40,15 @@ def extend_axis(array, axis):
     sum_vals = array.sum(axis, keepdims=True)
     norm = sum_vals / (np.sqrt(n) + n)
     fill_val = norm - sum_vals / np.sqrt(n)
-    
+
     out = np.concatenate([array, fill_val.astype(str(array.dtype))], axis=axis)
     return out - norm.astype(str(array.dtype))
 
 
 def extend_axis_rev(array, axis):
     n = array.shape[axis]
     last = np.take(array, [-1], axis=axis)
-    
+
     sum_vals = -last * np.sqrt(n)
     norm = sum_vals / (np.sqrt(n) + n)
     slice_before = (slice(None, None),) * len(array.shape[:axis])
@@ -56,60 +57,60 @@ def extend_axis_rev(array, axis):
 
 class ZeroSumTransform(pm.distributions.transforms.Transform):
     name = "zerosum"
-    
+
     _active_dims: List[int]
-    
+
     def __init__(self, active_dims):
         self._active_dims = active_dims
-    
+
     def forward(self, x):
         for axis in self._active_dims:
             x = extend_axis_rev_aet(x, axis=axis)
         return x
-    
+
     def forward_val(self, x, point=None):
         for axis in self._active_dims:
             x = extend_axis_rev(x, axis=axis)
         return x
-    
+
     def backward(self, z):
         z = aet.as_tensor_variable(z)
         for axis in self._active_dims:
             z = extend_axis_aet(z, axis=axis)
         return z
-    
+
     def jacobian_det(self, x):
-        return aet.constant(0.)
-    
-    
+        return aet.constant(0.0)
+
+
 class ZeroSumNormal(pm.Continuous):
     def __init__(self, sigma=1, *, active_dims=None, active_axes=None, **kwargs):
         shape = kwargs.get("shape", ())
         dims = kwargs.get("dims", None)
         if isinstance(shape, int):
             shape = (shape,)
-        
+
         if isinstance(dims, str):
             dims = (dims,)
 
         self.mu = self.median = self.mode = aet.zeros(shape)
         self.sigma = aet.as_tensor_variable(sigma)
-        
+
         if active_dims is None and active_axes is None:
             if shape:
                 active_axes = (-1,)
             else:
                 active_axes = ()
-        
+
         if isinstance(active_axes, int):
             active_axes = (active_axes,)
-        
+
         if isinstance(active_dims, str):
             active_dims = (active_dims,)
-        
+
         if active_axes is not None and active_dims is not None:
             raise ValueError("Only one of active_axes and active_dims can be specified.")
-        
+
         if active_dims is not None:
             model = pm.modelcontext(None)
             print(model.RV_dims)
@@ -118,19 +119,19 @@ def __init__(self, sigma=1, *, active_dims=None, active_axes=None, **kwargs):
             active_axes = []
             for dim in active_dims:
                 active_axes.append(dims.index(dim))
-        
+
         super().__init__(**kwargs, transform=ZeroSumTransform(active_axes))
 
     def logp(self, x):
         return pm.Normal.dist(sigma=self.sigma).logp(x)
-    
+
     @staticmethod
     def _random(scale, size):
         samples = stats.norm.rvs(loc=0, scale=scale, size=size)
         return samples - np.mean(samples, axis=-1, keepdims=True)
-    
+
     def random(self, point=None, size=None):
-        sigma, = draw_values([self.sigma], point=point, size=size)
+        (sigma,) = draw_values([self.sigma], point=point, size=size)
         return generate_samples(self._random, scale=sigma, dist_shape=self.shape, size=size)
 
     def _distr_parameters_for_repr(self):