def subset(self, plates_index: Dict[Plate,
List[int]]) -> "EPMeanFieldSubset":
"""Given a dictionary of Plates with a subset of indexes
returns the EPMeanFieldSubset that corresponds to the subset of indexes.
"""
factor_subset_factor = {}
factor_mean_field_subset = {}
factor_mean_field_rescale = {}
for factor, mean_field in self.factor_mean_field.items():
plate_sizes = VariableData.plate_sizes(mean_field)
factor_subset_factor[factor] = subset_factor = factor.subset(
plates_index, plate_sizes)
mean_field_subset = mean_field.subset(plates_index=plates_index)
factor_mean_field_subset[subset_factor] = mean_field_subset
mean_field_size = VariableData.prod(plate_sizes)
subset_size = VariableData.prod(
VariableData.plate_sizes(mean_field_subset))
scale_factor = subset_size / mean_field_size
factor_mean_field_rescale[subset_factor] = {
v: scale_factor * mean_field[v].size / message.size
for v, message in mean_field_subset.items()
}
subset_factor_graph = FactorGraph(factor_subset_factor.values())
return EPMeanFieldSubset(
subset_factor_graph,
factor_mean_field_subset,
factor_mean_field_rescale,
factor_subset_factor,
self,
plates_index,
)
def update_mean_field(self, mean_field, plates_index=None):
if plates_index:
plate_sizes = VariableData.plate_sizes(self)
for v, new_message in mean_field.items():
index = v.make_indexes(plates_index, plate_sizes)
self[v][index] = new_message
else:
self.update(mean_field)
return self
def merge(self, index, mean_field):
new_dist = dict(self)
if index:
plate_sizes = VariableData.plate_sizes(self)
for v, message in mean_field.items():
i = v.make_indexes(index, plate_sizes)
new_dist[v] = new_dist[v].merge(i, message)
else:
new_dist.update(mean_field)
return MeanField(new_dist)
def subset(self, variables=None, plates_index=None):
cls = type(self) if isinstance(self, MeanField) else MeanField
variables = variables or self.variables
if plates_index:
plate_sizes = VariableData.plate_sizes(self)
variable_index = (
(v, v.make_indexes(plates_index, plate_sizes)) for v in variables
)
mean_field = dict((v, self[v][index]) for v, index in variable_index)
return cls(mean_field)
return cls((v, self[v]) for v in variables if v in self.keys())