def test_cluster_and_fit_1_cluster(self):
# Ensure that with 1 cluster, we get the expected answer.
# Grab some data.
data = self.results[0]
# Define inputs for actual zip fitting.
zip_fit_inputs = {'s_n': self.s_n[0], 'v_n': self.v_n}
# Use the 10th element for cluster selection.
selection_data = data.iloc[10][['p', 'q']]
# Call cluster_and_fit.
fit_data = zip.cluster_and_fit(data=data,
zip_fit_inputs=zip_fit_inputs,
selection_data=selection_data,
n_clusters=1, random_state=2)
self.check_pq(expected=data, predicted=fit_data)
def test_cluster_and_fit_3_cluster(self):
# Ensure that with 3 clusters, we get the expected answer.
# Grab data for fitting.
data = pd.concat([self.results[0], self.results[1], self.results[2]])
# Define inputs for zip fitting. We'll try to match results[1].
zip_fit_inputs = {'s_n': self.s_n[1], 'v_n': self.v_n}
# Use first element for cluster selection.
selection_data = self.results[1].iloc[0][['p', 'q']]
# Call cluster_and_fit.
fit_data = zip.cluster_and_fit(data=data,
zip_fit_inputs=zip_fit_inputs,
selection_data=selection_data,
n_clusters=3, random_state=2)
self.check_pq(expected=self.results[1], predicted=fit_data)
def test_cluster_and_fit_no_cluster(self):
# Ensure that with no clustering, we get the expected answer.
# Grab some data.
data = self.results[3]
# Define inputs for actual zip fitting.
zip_fit_inputs = {'s_n': self.s_n[3], 'v_n': self.v_n}
# Use the 10th element for cluster selection.
selection_data = None
# Call cluster_and_fit.
fit_data = zip.cluster_and_fit(data=data,
zip_fit_inputs=zip_fit_inputs,
selection_data=selection_data,
n_clusters=None,
random_state=None)
self.check_pq(expected=data, predicted=fit_data)
def test_cluster_and_fit_cluster_too_small(self):
# If a cluster is too small, None should be returned.
# Grab data for fitting.
data = self.results[0]
# Define inputs for zip fitting. We'll try to match results[2].
zip_fit_inputs = {'s_n': self.s_n[0], 'v_n': self.v_n}
# Use 15th element for cluster selection.
selection_data = data.iloc[15][['p', 'q']]
# Call cluster_and_fit.
fit_data = zip.cluster_and_fit(data=data,
zip_fit_inputs=zip_fit_inputs,
selection_data=selection_data,
n_clusters=data.shape[0],
random_state=2,
min_cluster_size=10)
self.assertEqual(None, fit_data)