Let time_vary elements be IndexDistributions
@@ -8,7 +8,11 @@ problem by finding a general equilibrium dynamic rule.
"""
import sys
import os
-from HARK.distribution import Distribution, TimeVaryingDiscreteDistribution
+from HARK.distribution import (
+ Distribution,
+ TimeVaryingDiscreteDistribution,
+ IndexDistribution,
+)
from distutils.dir_util import copy_tree
from .utilities import get_arg_names, NullFunc
from copy import copy, deepcopy
@@ -426,9 +430,13 @@ class AgentType(Model):
A method to check that elements of time_vary are lists.
"""
for param in self.time_vary:
- if type(getattr(self, param)) != TimeVaryingDiscreteDistribution:
+ if isinstance(
+ getattr(self, param),
+ (TimeVaryingDiscreteDistribution, IndexDistribution),
+ ):
assert type(getattr(self, param)) == list, (
- param + " is not a list or time varying distribution,"
+ param
+ + " is not a list or time varying distribution,"
+ " but should be because it is in time_vary"
)
