Create and store auxiliary distributions instead of re-creating them every time
@@ -79,13 +79,22 @@ class IndexDistribution(Distribution):
self.conditional = conditional
self.engine = engine
- def __getitem__(self, y):
- # test one item to determine case handling
+
+ self.dstns = []
+
+ # Test one item to determine case handling
item0 = list(self.conditional.values())[0]
if type(item0) is list:
+ # Create and store all the conditional distributions
+ for y in range(len(item0)):
cond = {key: val[y] for (key, val) in self.conditional.items()}
- return self.engine(seed=self.RNG.randint(0, 2 ** 31 - 1), **cond)
+ self.dstns.append(self.engine(seed=self.RNG.randint(0, 2 ** 31 - 1), **cond))
+
+ elif type(item0) is float:
+
+ self.dstns = [self.engine(seed=self.RNG.randint(0, 2 ** 31 - 1), **conditional)]
+
else:
raise (
Exception(
@@ -93,6 +102,10 @@ class IndexDistribution(Distribution):
)
)
+ def __getitem__(self, y):
+
+ return self.dstns[y]
+
def approx(self, N, **kwds):
"""
Approximation of the distribution.
@@ -124,9 +137,7 @@ class IndexDistribution(Distribution):
if type(item0) is float:
# degenerate case. Treat the parameterization as constant.
- return self.engine(
- seed=self.RNG.randint(0, 2 ** 31 - 1), **self.conditional
- ).approx(N, **kwds)
+ return self.dstns[0].approx(N, **kwds)
if type(item0) is list:
return TimeVaryingDiscreteDistribution(
