def test_read_write(self):
household_model, _ = self._mock_household_collection()
def _check_network(household_model_new):
self.assertSequenceEqual(household_model.fields,
household_model_new.fields)
self._check_household_generate(household_model_new)
builtin_module_name = 'builtins' if sys.version_info.major == 3 else '__builtin__'
with patch('{}.open'.format(builtin_module_name),
new_callable=mock_open()) as open_mock:
household_model.write('file')
open_mock.assert_called_once_with('file', 'w')
# Check that correct json is written
json_net = open_mock.return_value.__enter__.return_value.write.call_args[
0][0]
household_model_new = BayesianNetworkModel.from_json(
json_net, self._household_segmenter())
_check_network(household_model_new)
# Check that json is correctly read
open_mock.return_value.__enter__.return_value.read.return_value = json_net
household_model_new = BayesianNetworkModel.from_file(
'file', self._household_segmenter())
_check_network(household_model_new)
def create_bayes_net(state_id, puma_id, output_dir, households_data, persons_data, configuration,
person_segmenter, household_segmenter):
'''Create a bayes net from pums dataframes and a configuration.
Args:
state_id: 2-digit state fips code
puma_id: 5-digit puma code
output_dir: dir to write out the generated bayesian nets to
households_data: pums households data frame
persons_data: pums persons data frame
configuration: specifies the structure of the bayes net
person_segmenter: function of inputs data to segment on a person variable
household_segmenter: function of inputs data to segment on a household variable
Returns:
household and person bayesian models
'''
# Write the persons bayes net to disk
person_training_data = SegmentedData.from_data(
cleaned_data=persons_data,
fields=list(configuration.person_fields),
weight_field=inputs.PERSON_WEIGHT.name,
segmenter=person_segmenter
)
person_model = BayesianNetworkModel.train(
input_data=person_training_data,
structure=configuration.person_structure,
fields=configuration.person_fields
)
person_model_filename = os.path.join(
output_dir, FILE_PATTERN.format(state_id, puma_id, 'person_model.json')
)
person_model.write(person_model_filename)
# Write the households bayes net to disk
household_training_data = SegmentedData.from_data(
cleaned_data=households_data,
fields=list(configuration.household_fields),
weight_field=inputs.HOUSEHOLD_WEIGHT.name,
segmenter=household_segmenter,
)
household_model = BayesianNetworkModel.train(
input_data=household_training_data,
structure=configuration.household_structure,
fields=configuration.household_fields
)
household_model_filename = os.path.join(
output_dir, FILE_PATTERN.format(state_id, puma_id, 'household_model.json')
)
household_model.write(household_model_filename)
return household_model, person_model
def test_to_from_json(self):
household_model, _ = self._mock_household_collection()
household_string = household_model.to_json()
household_model_new = BayesianNetworkModel.from_json(household_string)
self.assertSequenceEqual(household_model.fields,
household_model_new.fields)
self._check_household_generate(household_model_new)
def test_generate_with_prior(self):
network = BayesianNetworkModel.train(
bayesnets.SegmentedData({'one_bucket': [('35-64', 'F', '40k+')]}),
self._person_structure(),
self._person_fields(),
prior_data={('35-64', 'F', '40k+')})
person = network.generate('one_bucket', ())[0]
self.assertEquals(person, ('35-64', 'F', '40k+'))
def test_read_write(self):
household_model, _ = self._mock_household_collection()
def _check_network(household_model_new):
self.assertSequenceEqual(household_model.fields,
household_model_new.fields)
self._check_household_generate(household_model_new)
with patch('__builtin__.open') as open_mock:
household_model.write('file')
open_mock.assert_called_once_with('file', 'w')
# Check that correct json is written
json_net = open_mock.return_value.__enter__.return_value.write.call_args[
0][0]
household_model_new = BayesianNetworkModel.from_json(
json_net, self._household_segmenter())
_check_network(household_model_new)
# Check that json is correctly read
open_mock.return_value.__enter__.return_value.read.return_value = json_net
household_model_new = BayesianNetworkModel.from_file(
'file', self._household_segmenter())
_check_network(household_model_new)