def test_ChainCollection_Hmatrix_calculation(self):
antibody_collection_1 = ChainCollection(
path='./tests/Data/chain_collection_fasta_test.fasta')
antibody_collection_1.load(show_progressbar=False, verbose=False)
# if this fails it means that abysis has been updated
self.assertEqual(antibody_collection_1.hydrophobicity_matrix().shape,
(1, 158))
def test_ChainCollection_append_2(self):
antibody_collection_1 = ChainCollection(
path='./tests/Data/chain_collection_1_heavy.json')
antibody_collection_1.load(show_progressbar=False, verbose=False)
antibody_collection_2 = ChainCollection(
path='./tests/Data/chain_collection_2_heavy.json')
antibody_collection_2.load(show_progressbar=False, verbose=False)
antibody_collection_1.append(antibody_collection_2)
self.assertEqual(antibody_collection_1.hydrophobicity_matrix().shape,
(2, 158))
class Cluster:
def __init__(self,
antibodies,
metric='hydrophobicity',
clustering_method='kmeans',
decomposition_method='PCA'):
if isinstance(antibodies, ChainCollection):
self.antibodies = antibodies
elif isinstance(antibodies, str):
self.antibodies = ChainCollection(path=antibodies)
if self.antibodies.n_ab == 0:
self.antibodies.load()
self.metric = metric
self.clustering_method = clustering_method
self.decomposition_method = decomposition_method
self.cluster_assignment = np.zeros(self.antibodies.n_ab, dtype=int)
self.cluster_assignment_dict = dict()
self._data = None
def _collect_data(self):
if self.metric == 'hydrophobicity':
return self.antibodies.hydrophobicity_matrix()
def cluster(self, n_components=0.95, n_clusters=3):
if self.decomposition_method == 'PCA':
decomposition_obj = decomposition.PCA(n_components)
self._data = decomposition_obj.fit_transform(self._collect_data())
if self.clustering_method == 'kmeans':
clustering_obj = cluster.KMeans(n_clusters=n_clusters)
self.cluster_assignment = clustering_obj.fit_predict(self._data)
for i, antibody_obj in enumerate(self.antibodies.antibody_objects):
assignment = 'Cluster_{}'.format(self.cluster_assignment[i])
if assignment not in self.cluster_assignment_dict:
self.cluster_assignment_dict[assignment] = list()
self.cluster_assignment_dict[assignment].append(antibody_obj)
def plot_cluster(self):
if len(self.cluster_assignment_dict) == 0:
self.cluster()
color = iter(plt.get_cmap('Vega20').colors)
plt.figure(figsize=(8, 8))
for assignment in np.unique(self.cluster_assignment):
c = next(color)
plt.scatter(self._data[self.cluster_assignment == assignment, 0],
self._data[self.cluster_assignment == assignment, 1],
c=c,
label='Cluster {}'.format(assignment))
plt.legend(loc='best', prop={"size": 14})
