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)