14class DSMDataConverter:
 15    """
 16    Detailed explanation of this class can be found in the above [guide](#data-conversion).
 17
 18    It can be used for conversion for both AE and HAN. Their conversion are similar so there is a common
 19    method. They are different in that HAN needs a hierarchical structure that has one more dimension than that for AE.
 20
 21    This class uses an already trained word2vec model to convert brain regions into something numeric, but you can
 22    train one yourself and provide it in initialization. By default, it extracts three features from the original
 23    dataset: region, node_type and xyz.
 24
 25    In most cases, 'region' is enough, so it's the default option. The 'node_type' is useful for finding words level
 26    features in data conversion for HAN. If you need more features, consider making a new class like this one.
 27    """
 28
 29    def __init__(self, dataset: Union[dict, OrderedDict], w2v_model: Union[str, 'os.PathLike[str]', Word2Vec] = None):
 30        """Load the word2vec model and preprocess the dataset(common steps for AE and HAN data preparation).
 31        Input dataset should be OrderedDict, consistent to the output of the `preprocessing` module.
 32        If it's just a dict, it will be converted as ordered.
 33
 34        :param dataset: An OrderedDict for
 35        :param w2v_model: Another word2vec model rather than the default, can be the model or path.
 36        """
 37
 38        self._origin = OrderedDict(dataset)
 39        self._region_vec = OrderedDict()
 40        self._scaled_loc = OrderedDict()
 41        self._type_dummy = OrderedDict()
 42
 43        # load word2vec model
 44        use_default = False
 45        if isinstance(w2v_model, Word2Vec):
 46            model = w2v_model
 47        elif isinstance(w2v_model, str) or isinstance(w2v_model, os.PathLike):
 48            model = Word2Vec.load(str(w2v_model))
 49        else:
 50            with as_file(files('dsmtools.modeling') / 'w2v_model_6dim_from_1200_seu.model') as path:
 51                model = Word2Vec.load(str(path))
 52                use_default = True
 53        self._region2vec = dict(zip(model.wv.index_to_key, StandardScaler().fit_transform(model.wv.vectors)))
 54        if use_default and 'unknow' in self._region2vec:  # sorry, we named it wrong while training
 55            self._region2vec['unknown'] = self._region2vec.pop('unknow')
 56
 57        self.preprocess()
 58
 59    def preprocess(self):
 60        """Common steps for AE and HAN data preparations: turn the brain regions to vector, scale the coordinates,
 61        and make the topological type field as dummy. These three features are currently the ones you can choose
 62        from to make up the deep learning data.
 63        """
 64
 65        for key, df in self._origin.items():
 66            self._region_vec[key] = pd.DataFrame(
 67                map(lambda w: self._region2vec[w] if w in self._region2vec else self._region2vec['unknown'],
 68                    df['region'].str.replace('fiber tracts', 'fiber_tracts')), index=df.index)
 69            self._scaled_loc[key] = pd.DataFrame(StandardScaler().fit_transform(df[['x', 'y', 'z']].values),
 70                                                 index=df.index)
 71            self._type_dummy[key] = pd.get_dummies(list(df['node_type']))
 72
 73    def convert_for_ae(self, seq_len=2000, features=('region',)):
 74        """Data conversion for autoencoder.
 75
 76        Specify the sequence length, so that all neuron sequences are padded or truncated to the same length.
 77
 78        The default parameters are consistent with our trained model. You don't have to change if you are going
 79        to use.
 80
 81        :param seq_len: The final sequence length.
 82        :param features: The features to use, among the three in the preprocessing.
 83        :return: An array of shape (#neuron, seq_len, #feature)
 84        """
 85
 86        final = {
 87            'region': pad_sequences([*self._region_vec.values()], maxlen=seq_len, dtype='float32'),
 88            'coord': pad_sequences([*self._scaled_loc.values()], maxlen=seq_len, dtype='float32'),
 89            'topology': pad_sequences([*self._type_dummy.values()], maxlen=seq_len, dtype='float32')
 90        }
 91        return np.concatenate([final[f] for f in features], axis=-1)
 92
 93    def convert_for_han(self, seq_len_word=20, seq_len_sentence=300, features=('region',),
 94                        labels: dict[str, str] = None):
 95        """Data conversion for autoencoder.
 96
 97        The sequences will be further transformed so that they grouped in a two level structure.
 98        word and sentence, each of which will be given a sequence length, and they will be padded
 99        or truncated to the same length.
100
101        The two level structure will use the topological information--terminals to use to break up words, i.e.
102        a sequence between two terminals is seen as a word and its length will be unified. A neuron sequence will
103        be turned into multiple words to make up a sentence, and the number of words will be seen as the length
104        of the sentence and will also be unified.
105
106        If labels are provided as a dict mapping the key from the input dataset to their types, this conversion will
107        output the encoded labels as vectors and the label encoder, suitable for training.
108
109        The default parameters are consistent with our trained model. You don't have to change if you are going
110        to use.
111
112        :param seq_len_word: The sequence length of word level
113        :param seq_len_sentence: The sequence length of sentence level.
114        :param features: The features to use, among the three in the preprocessing.
115        :param labels: A dict mapping from the key of each sequence to their cell type.
116        :return: An array of shape (#neuron, seq_len_sentence, seq_len_word, #feature)
117        """
118
119        region_hierarch = []
120        loc_hierarch = []
121        type_hierarch = []
122        for df, region_vec, scaled_loc, type_dummy in zip(self._origin.values(), self._region_vec.values(),
123                                                          self._scaled_loc.values(), self._type_dummy.values()):
124            higher_region = []
125            higher_loc = []
126            higher_type = []
127            last_node = df.index[0]
128            for terminal in df.index[df['node_type'] == 'T']:  # you can see terminals as punctuations
129                higher_region.append(region_vec.loc[last_node:terminal].values.astype(float))
130                higher_loc.append(scaled_loc.loc[last_node:terminal].values.astype(float))
131                higher_type.append(type_dummy.loc[last_node:terminal].values.astype(float))
132                last_node = terminal
133            region_hierarch.append(pad_sequences(higher_region, maxlen=seq_len_word, dtype='float32'))
134            loc_hierarch.append(pad_sequences(higher_loc, maxlen=seq_len_word, dtype='float32'))
135            type_hierarch.append(pad_sequences(higher_type, maxlen=seq_len_word, dtype='float32'))
136
137        final = {
138            'region': pad_sequences(region_hierarch, maxlen=seq_len_sentence, dtype='float32'),
139            'coord': pad_sequences(region_hierarch, maxlen=seq_len_sentence, dtype='float32'),
140            'topology': pad_sequences(type_hierarch, maxlen=seq_len_sentence, dtype='float32')
141        }
142
143        x = np.concatenate([final[f] for f in features], axis=-1)
144
145        if labels is None:
146            return x
147
148        label_encoder = LabelEncoder()
149        y = label_encoder.fit_transform([labels[k] for k in self._origin.keys()])
150        y = to_categorical(y, num_classes=len(label_encoder.classes_))
151        return x, y, label_encoder

 15class DSMAutoencoder(ModelAbstract):
 16    """
 17    The class to compile an autoencoder, train/load model parameters and use it to encode your neuron sequences,
 18    based on keras API.
 19    Detailed explanations can be found in [the above guide](#deep-sequential-models).
 20
 21    To use this class, you can start by invoking `DSMAutoencoder.load_1282_seu` to get our trained model and encode
 22    your neuron sequences, provided that they are converted to the shape that this model accepts (same as the default
 23    parameters of the `DSMDataConverter.convert_for_ae`). And then use `DSMAutoencoder.predict` to get the latent
 24    variables.
 25
 26    To use this class from the beginning, after initialization, use `DSMAutoencoder.compile`
 27    to set up a new model, whose dimension should be consistent with your input.
 28    If you'd like to train your own model, use `DSMAutoencoder.fit`, and `DSMAutoencoder.load` to load a saved one.
 29    Use `DSMAutoencoder.evaluate` for further evaluation.
 30
 31    The training history and the model are available as `DSMAutoencoder.history` and `DSMAutoencoder.model`. If you
 32    are familiar with keras, you know what to do with them, but we offer some basic functions to plot learning stats and
 33    save the model parameters. Their paths can be specified in each method.
 34    """
 35
 36    def compile(self, seq_len=2000, feature_dim=6, latent_dim=32,
 37                optimizer=Adam(learning_rate=0.001, decay=0.001, clipnorm=1.0)):
 38        """Compile the AE model with specified parameters, before you can train, load or predict.
 39
 40        The dimensions of sequence length and feature should be consistent with those in the data conversion,
 41        used by `DSMDataConverter.convert_for_ae`. The dimension of the latent space is dimension of the encoder output.
 42        They should also be the same with a trained model, if used, and the default parameters are used by our
 43        prepared model.
 44
 45        :param seq_len: The length of the sequence, the second dimension of the input array.
 46        :param feature_dim: The number of features, the last dimension of the input array.
 47        :param latent_dim: The dimension of the encoded vector.
 48        :param optimizer: A keras optimizer with learning parameters.
 49        """
 50
 51        input = Input((seq_len, feature_dim,))
 52        x = Masking(mask_value=0.0)(input)
 53        x = GRU(128, kernel_regularizer=regularizers.l2(0.001), dropout=0.5, return_sequences=True)(x)
 54        x = GRU(64, kernel_regularizer=regularizers.l2(0.001), dropout=0.5, return_sequences=True)(x)
 55
 56        encoder = GRU(latent_dim, kernel_regularizer=regularizers.l2(0.001),
 57                      dropout=0.5, return_sequences=False, name='encoder')(x)   # return sequence is false here
 58
 59        x = RepeatVector(seq_len)(encoder)
 60        x = GRU(64, kernel_regularizer=regularizers.l2(0.001), dropout=0.5, return_sequences=True)(x)
 61        x = GRU(128, kernel_regularizer=regularizers.l2(0.001), dropout=0.5, return_sequences=True)(x)
 62
 63        decoder = GRU(feature_dim, kernel_regularizer=regularizers.l2(0.001), dropout=0.5, return_sequences=True)(x)
 64
 65        self._model = Model(inputs=input, outputs=decoder)
 66        self._model.summary()
 67        self._model.compile(optimizer=optimizer, loss=mean_squared_error)
 68
 69    def load(self, path: Union[str, 'PathLike[str]']):
 70        super(DSMAutoencoder, self).load(path)
 71
 72    def fit(self, data_train, data_test, batch_size=64, epochs=300,
 73            model_save_path: Union[str, 'PathLike[str]'] = None, shuffle=False):
 74        """
 75        Model training. Since AE is unsupervised, you only need to provide the input data, for train and test.
 76
 77        If a model save path is specified, the model will be saved when improved by loss evaluation. This save path
 78        can also be formatted as a keras model path. Check keras documentation to write the format string.
 79
 80        :param data_train: An input array for training, with shape of (#neuron, #seq_len, #feature).
 81        :param data_test: An input array for testing, with shape of (#neuron, #seq_len, #feature).
 82        :param batch_size: The training batch size.
 83        :param epochs: The number of training epochs.
 84        :param model_save_path: The path to save the model by checkpoint.
 85        :param shuffle: Whether to shuffle the input data.
 86        """
 87
 88        super(DSMAutoencoder, self)._fit(data_train, data_train, data_test, data_test,
 89                                         batch_size, epochs, model_save_path, 'val_loss', shuffle)
 90
 91    def plot_learning_curve(self, path: Union[str, 'PathLike[str]'] = None, fig_size=(20, 20), dpi=200):
 92        """Plot the learning curve using the history data after training.
 93
 94        It will plot a new figure and return the matplotlib handles.
 95        If a path is specified, the figure will be saved.
 96
 97        :param path: The save path of the figure.
 98        :param fig_size: The size of the figure.
 99        :param dpi: resolution setting.
100        :return: The plotted figure's handles, figure and axes.
101        """
102
103        return super(DSMAutoencoder, self)._plot_learning_curve(path, fig_size, dpi, ['loss', 'val_loss'])
104
105    def predict(self, x: np.array):
106        """Give the encoder's output with given data.
107
108        This will temporarily build a new model using the parameters from the whole model. For efficiency, predicting
109        everything in a single shot is preferred.
110
111        :param x: Data to encode, with shape of (#neuron, #seq_len, #feature).
112        :return: The predicted latent variables, with shape of (#neuron, #latent_dim).
113        """
114        encoder = Model(inputs=self._model.input, outputs=self._model.get_layer('encoder').output)
115        return encoder.predict(x)
116
117    def evaluate(self, x):
118        """Test the model with some input data. As this model is unsupervised, it only needs the input data.
119        Your input array should share its dimensions other than the number of neurons with the model.
120
121        :param x: An array for evaluation, with shape of (#neuron, #seq_len, #feature).
122        :return: The loss of evaluation.
123        """
124        return super(DSMAutoencoder, self)._evaluate(x, x)
125
126    @staticmethod
127    def load_1282_seu():
128        """A static method to build and load an already trained model on 1282 neurons from Southeast University.
129
130        :return: A `DSMAutoencoder` instance with the trained model.
131        """
132        model = DSMAutoencoder()
133        model.compile()
134        with as_file(files('dsmtools.modeling') / 'AE_1282_SEU.h5') as path:
135            model.load(path)
136        return model

 45class DSMHierarchicalAttentionNetwork(ModelAbstract):
 46    """
 47    The class to compile a hierarchical attention network,
 48    train/load model parameters and use it to encode your neuron sequences, based on keras API.
 49    Detailed explanations can be found in [the above guide](#deep-sequential-models).
 50
 51    To use this class, you can start by invoking `DSMHierarchicalAttentionNetwork.load_1282_seu` to get our trained
 52    model and encode your neuron sequences, provided that they are converted to the shape that this model accepts
 53    (same as the default parameters of the `DSMDataConverter.convert_for_han`).
 54    And then use `DSMHierarchicalAttentionNetwork.predict` to get the classification. The result is a one-hot
 55    matrix. To get their labels, your can find the largest position of each neuron, and invoke
 56    `DSMHierarchicalAttentionNetwork.label_encoder_1282_seu` to get the label for that position. There are 12 types
 57    in our trained classification model.
 58
 59    To use this class from the beginning, after initialization, use `DSMHierarchicalAttentionNetwork.compile` to set up
 60    a new model, whose dimension should be consistent with your input.
 61    If you'd like to train your own model, use `DSMAutoencoder.fit`, and `DSMAutoencoder.load` to load a saved one.
 62    Use `DSMAutoencoder.evaluate` for further evaluation. The label encoder is obtained by
 63    `DSMDataConverter.convert_for_han`, and it's yours to maintain.
 64
 65    The training history and the model are available as `DSMHierarchicalAttentionNetwork.history` and
 66    `DSMHierarchicalAttentionNetwork.model`. If you
 67    are familiar with keras, you know what to do with them, but we offer some basic functions to plot learning stats and
 68    save the model parameters. Their paths can be specified in each method.
 69    """
 70
 71    def compile(self, seq_len_sentence=300, seq_len_word=20, feature_dim=6, class_num=12,
 72                last_activation='softmax', index_test=False,
 73                optimizer=RMSprop(learning_rate=0.001, clipvalue=0.5, decay=0.002)):
 74        """Compile the HAN model with specified parameters, before you can train, load or predict.
 75
 76        The dimensions of word and sentence length and feature should be consistent with those in the data conversion,
 77        used by `DSMDataConverter.convert_for_han`. The class number is dimension of the classification result,
 78        a one-hot matrix.
 79        They should also be the same with a trained model, if used, and the default parameters are used by our
 80        prepared model.
 81
 82        You can also specify the way how the last activation is calculated, and how word and sentence layers
 83        are built.
 84
 85        :param seq_len_sentence: The length of the sentence level, the second dimension of the input array.
 86        :param seq_len_word: The length of the word level, the third dimension of the input array.
 87        :param feature_dim: The number of features, the last dimension of the input array.
 88        :param class_num: The number of classes, the second dimension of the output.
 89        :param last_activation: The activation function for the last layer.
 90        :param index_test: Whether to use time distributed rather than GRU to build word and sentence layers.
 91        :param optimizer: A keras optimizer with learning parameters.
 92        """
 93
 94        input_word = Input(shape=(seq_len_word, feature_dim,))  # [batchsize, word_length, dim(8)]
 95        x_word = Masking(mask_value=0.0)(input_word)
 96        em1 = 128
 97
 98        if not index_test:
 99            x_word = GRU(em1, return_sequences=False, kernel_regularizer=regularizers.l2(0.001), dropout=0.5)(
100                x_word)  # LSTM or GRU
101        elif index_test:
102            x_word = TimeDistributed(Dense(em1, kernel_regularizer=regularizers.l2(0.001)), name='segment_encoder')(
103                x_word)
104
105        model_word = Model(input_word, x_word)
106
107        # Sentence part
108        input = Input(
109            shape=(seq_len_sentence, seq_len_word, feature_dim,))  # [batchsize, sentence_length, word_length, dim(8)]
110        x_sentence = TimeDistributed(model_word, name='segment_encoder')(input)
111        em2 = 128
112
113        if not index_test:
114            x_sentence = GRU(em2, return_sequences=False, kernel_regularizer=regularizers.l2(0.001), dropout=0.5,
115                             name='sentence_out')(x_sentence)  # LSTM or GRU
116        elif index_test:
117            x_sentence = TimeDistributed(Dense(em2, kernel_regularizer=regularizers.l2(0.001)))(x_sentence)
118
119        x_sentence = Dense(64, kernel_regularizer=regularizers.l2(0.01), name='dense_out')(x_sentence)
120        x_sentence = Dropout(rate=0.5)(x_sentence)
121        x_sentence = ReLU()(x_sentence)
122
123        output = Dense(class_num, activation=last_activation, kernel_regularizer=regularizers.l2(0.01))(x_sentence)
124        self._model = Model(inputs=input, outputs=output)
125        self._model.summary()
126        self._model.compile(optimizer=optimizer, loss='categorical_crossentropy',
127                            metrics=[CategoricalAccuracy(name='accuracy')])
128
129    def load(self, path: Union[str, 'PathLike[str]']):
130        super(DSMHierarchicalAttentionNetwork, self).load(path)
131
132    def fit(self, x_train, y_train, x_test, y_test, batch_size=64, epochs=400, model_save_path=None, shuffle=False):
133        """Model training.
134
135        If a model save path is specified, the model will be saved when improved by loss evaluation. This save path
136        can also be formatted as a keras model path. Check keras documentation to write the format string.
137
138        :param x_train: An input array for training, with shape of (#neuron, #seq_len, #feature).
139        :param y_train: An output array for training, with shape of (#neuron, #classes).
140        :param x_test: An input array for test, with shape of (#neuron, #seq_len, #feature).
141        :param y_test: An output array for testing, with shape of (#neuron, #classes).
142        :param batch_size: The training batch size.
143        :param epochs: The number of training epochs.
144        :param model_save_path: The path to save the model by checkpoint.
145        :param shuffle: Whether to shuffle the input data.
146        :return:
147        """
148
149        super(DSMHierarchicalAttentionNetwork, self)._fit(x_train, y_train, x_test, y_test,
150                                                          batch_size, epochs, model_save_path, 'val_accuracy', shuffle)
151
152    def plot_learning_curve(self, path: Union[str, 'PathLike[str]'] = None, fig_size=(20, 20), dpi=200):
153        """Plot the learning curve using the history data after training.
154
155        It will plot a new figure and return the matplotlib handles.
156        If a path is specified, the figure will be saved.
157
158        :param path: The save path of the figure.
159        :param fig_size: The size of the figure.
160        :param dpi: resolution setting.
161        :return: The plotted figure's handles, figure and axes.
162        """
163
164        return super(DSMHierarchicalAttentionNetwork, self).\
165            _plot_learning_curve(path, fig_size, dpi, ['accuracy', 'val_accuracy', 'loss', 'val_loss'])
166
167    def predict(self, x: np.array):
168        return super(DSMHierarchicalAttentionNetwork, self).predict(x)
169
170    def evaluate(self, x, y):
171        """Test the model with a given pair of input data and output class one-hot matrix.
172        Your input and output array should share their dimensions other than the number of neurons with the model.
173
174        :param x: An array for evaluation, with shape of (#neuron, #sentence_len, #word_len, #feature).
175        :param y: A one-hot matrix of class results, with shape of (#neuron, #class).
176        :return: The loss and accuracy of evaluation.
177        """
178
179        return super(DSMHierarchicalAttentionNetwork, self)._evaluate(x, y)
180
181    @staticmethod
182    def load_1282_seu():
183        """A static method to build and load an already trained model on 1282 neurons from Southeast University.
184
185        :return: A `DSMHierarchicalAttentionNetwork` instance with the trained model.
186        """
187        model = DSMHierarchicalAttentionNetwork()
188        model.compile()
189        with as_file(files('dsmtools.modeling') / 'HAN_1282_SEU.h5') as path:
190            model.load(path)
191        return model
192
193    @staticmethod
194    def label_encoder_1282_seu():
195        """A static method to get the label encoder for the trained model on 1282 neurons from Southeast University.
196        You can use it to transform the number places inferred from the one-hot matrix back to their labels.
197
198        :return: A fitted `LabelEncoder` instance.
199        """
200        return LabelEncoder().fit(['CP_GPe', 'CP_SNr', 'ET_MO', 'ET_SS', 'IT_MO', 'IT_SS', 'IT_VIS',
201                                   'LGd', 'MG', 'RT', 'VPL', 'VPM'])