File tree Expand file tree Collapse file tree 1 file changed +14
-8
lines changed Expand file tree Collapse file tree 1 file changed +14
-8
lines changed Original file line number Diff line number Diff line change @@ -133,18 +133,24 @@ def gelman_rubin(mtrace):
133
133
134
134
def calc_rhat (x ):
135
135
136
- m , n = x .shape
136
+ try :
137
+ m , n = x .shape
138
+
139
+ # Calculate between-chain variance
140
+ B = n * np .var (np .mean (x , axis = 1 ), ddof = 1 )
137
141
138
- # Calculate between -chain variance
139
- B = n * np .var (np .mean (x , axis = 1 ) , ddof = 1 )
142
+ # Calculate within -chain variance
143
+ W = np .mean (np .var (x , axis = 1 , ddof = 1 ) )
140
144
141
- # Calculate within-chain variance
142
- W = np . mean ( np . var ( x , axis = 1 , ddof = 1 ))
145
+ # Estimate of marginal posterior variance
146
+ Vhat = W * ( n - 1 ) / n + B / n
143
147
144
- # Estimate of marginal posterior variance
145
- Vhat = W * (n - 1 )/ n + B / n
148
+ return np .sqrt (Vhat / W )
149
+
150
+ except ValueError :
146
151
147
- return np .sqrt (Vhat / W )
152
+ rotated_indices = np .roll (np .arange (x .ndim ), 1 )
153
+ return np .squeeze ([calc_rhat (xi ) for xi in x .transpose (rotated_indices )])
148
154
149
155
Rhat = {}
150
156
for var in mtrace .varnames :
You can’t perform that action at this time.
0 commit comments