diff --git a/mesa/experimental/measure.py b/mesa/experimental/measure.py new file mode 100644 index 00000000000..95f74984ae4 --- /dev/null +++ b/mesa/experimental/measure.py @@ -0,0 +1,84 @@ +from collections import defaultdict + + +class Group: + def __init__(self, model, fn): + self.model = model + self.fn = fn + + @property + def value(self): + return self.fn(self.model) + + +class Measure: + def __init__(self, group, measurer): + self.group = group + self.measurer = measurer + + def _measure_group(self, group, measurer): + # get an attribute + if isinstance(measurer, str): + return getattr(group, measurer) + # apply + return measurer(group) + + @property + def value(self): + group_object = self.group + if isinstance(self.group, Group): + group_object = self.group.value + return self._measure_group(group_object, self.measurer) + + +class DataCollector: + """ + Example: a model consisting of a hybrid of Boltzmann wealth model and + Epstein civil violence. + + class EpsteinBoltzmannModel: + def __init__(self): + # Groups + self.quiescents = Group( + lambda model: model.agents.select( + agent_type=Citizen, filter_func=lambda a: a.condition == "Quiescent" + ) + ) + self.citizens = Group(lambda model: model.get_agents_of_type(Citizen)) + + # Measurements + self.num_quiescents = Measure(self.quiescents, len) + self.gini = Measure( + self.agents, lambda agents: calculate_gini(agents.get("wealth")) + ) + self.gini_quiescents = Measure( + self.quiescents, lambda agents: calculate_gini(agents.get("wealth")) + ) + self.condition = Measure(self.citizens, "condition") + self.wealth = Measure(self.agents, "wealth") + + + def run(): + model = EpsteinBoltzmannModel() + datacollector = DataCollector( + model, ["num_quiescents", "gini_quiescents", "wealth"] + ) + + for _ in range(10): + model.step() + datacollector.collect() + """ + + def __init__(self, model, attributes): + self.model = model + self.attributes = attributes + self.data_collection = defaultdict(lambda: defaultdict(list)) + + def collect(self): + for name in self.attributes: + attribute = getattr(self.model, name) + group = "model" + if isinstance(attribute, Measure): + group = attribute.group + attribute = attribute.value + self.data_collection[group][name].append(attribute)