Source code for immuneML.encodings.kmer_frequency.KmerFreqRepertoireEncoder

from collections import Counter
from multiprocessing.pool import Pool

from immuneML.caching.CacheHandler import CacheHandler
from immuneML.caching.CacheObjectType import CacheObjectType
from immuneML.data_model.dataset.RepertoireDataset import RepertoireDataset
from immuneML.encodings.EncoderParams import EncoderParams
from immuneML.encodings.kmer_frequency.KmerFrequencyEncoder import KmerFrequencyEncoder
from immuneML.util.Logger import log


[docs]class KmerFreqRepertoireEncoder(KmerFrequencyEncoder): @log def _encode_new_dataset(self, dataset, params: EncoderParams): encoded_data = self._encode_data(dataset, params) encoded_dataset = RepertoireDataset(repertoires=dataset.repertoires, encoded_data=encoded_data, labels=dataset.labels, metadata_file=dataset.metadata_file) return encoded_dataset @log def _encode_examples(self, dataset, params: EncoderParams): arguments = [(repertoire, params) for repertoire in dataset.repertoires] with Pool(params.pool_size) as pool: repertoires = pool.starmap(self.get_encoded_repertoire, arguments) encoded_repertoire_list, repertoire_names, labels, feature_annotation_names = zip(*repertoires) encoded_labels = {k: [dic[k] for dic in labels] for k in labels[0]} if params.encode_labels else None feature_annotation_names = feature_annotation_names[0] return list(encoded_repertoire_list), list(repertoire_names), encoded_labels, feature_annotation_names
[docs] def get_encoded_repertoire(self, repertoire, params: EncoderParams): params.model = vars(self) return CacheHandler.memo_by_params((("encoding_model", params.model), ("type", "kmer_encoding"), ("labels", params.label_config.get_labels_by_name()), ("repertoire_id", repertoire.identifier)), lambda: self.encode_repertoire(repertoire, params), CacheObjectType.ENCODING_STEP)
[docs] def encode_repertoire(self, repertoire, params: EncoderParams): counts = Counter() sequence_encoder = self._prepare_sequence_encoder() feature_names = sequence_encoder.get_feature_names(params) for sequence in repertoire.sequences: counts = self._encode_sequence(sequence, params, sequence_encoder, counts) label_config = params.label_config labels = dict() if params.encode_labels else None if labels is not None: for label_name in label_config.get_labels_by_name(): label = repertoire.metadata[label_name] labels[label_name] = label # TODO: refactor this not to return 4 values but e.g. a dict or split into different functions? return counts, repertoire.identifier, labels, feature_names