API Reference - hNMF
Author: Eric Stasney
Description: Rank-2 NMF for Hierarchical Clustering

Skip to content

# API Reference

## hNMF.model

### `hnmf.model.HierarchicalNMF(k, unbalanced=0.1, init=None, solver='cd', beta_loss=0, alpha_W=0.0, alpha_H='same', random_state=42, trial_allowance=100, tol=1e-06, maxiter=10000, dtype=np.float64)`

Bases: `BaseEstimator`

Methods:

Name | Description
--- | ---
`_init_fit` |
`_stack_H_buffer` |
`_stack_clusters` |
`cluster_assignments` | Returns a mapping of features and their assigned cluster(s)
`cluster_features` | Returns the features assigned as a cluster to nodes
`fit` | Fit `HierarchicalNMF` to data
`top_discriminative_samples_in_node` | Computes most discriminative samples (node vs rest)
`top_features_in_node` | For a given node, return the top n features and values
`top_nodes_in_feature` | Returns the top nodes for a specified feature
`top_nodes_in_samples` | Returns the top nodes for each sample.
`top_samples_in_nodes` | Returns the top samples for each node

Attributes:

Name | Type | Description
--- | --- | ---
`H_buffer_` | `NDArray | None` |
`Hs_` | `NDArray | None` |
`W_buffer_` | `NDArray | None` |
`Ws_` | `NDArray | None` |
`alpha_H` | `Literal['same'] | float` |
`alpha_W` | `float` |
`beta_loss` | `Literal['FRO', 0, 'KL', 1, 'IS', 2]` |
`clusters_` | `NDArray | None` |
`dtype` | `DTypeLike` |
`feature2id_` | `dict[str, int] | None` |
`id2feature_` | `dict[int, str] | None` |
`id2sample_` | `dict[int, str] | None` |
`init` | `Literal[None, 'random', 'nndsvd', 'nndsvda', 'nndsvdar']` |
`is_leaf_` | `NDArray | None` |
`k` | `int` |
`maxiter` | `int` |
`n_features_` | `int | None` |
`n_leaves_` | `int` |
`n_nodes_` | `int` |
`n_samples_` | `int | None` |
`priorities_` | `NDArray | None` |
`random_state` | `RandomState` |
`solver` | `Literal['cd', 'mu']` |
`splits_` | `NDArray | None` |
`tol` | `float` |
`tree_` | `NDArray | None` |
`trial_allowance` | `int` |
`unbalanced` | `float` |

Source code in `src/hnmf/model.py`

     59
     60
     61
     62
     63
     64
     65
     66
     67
     68
     69
     70
     71
     72
     73
     74
     75
     76
     77
     78
     79
     80
     81
     82
     83
     84
     85
     86
     87
     88
     89
     90
     91
     92
     93
     94
     95
     96
     97
     98
     99
    100
    101
    102
    def __init__(
        self,
        k: int,
        unbalanced: float = 0.1,
        init: Literal[None, "random", "nndsvd", "nndsvda", "nndsvdar"] = None,
        solver: Literal["cd", "mu"] = "cd",
        beta_loss: Literal["FRO", 0, "KL", 1, "IS", 2] = 0,
        alpha_W: float = 0.0,
        alpha_H: Literal["same"] | float = "same",
        random_state: int = 42,
        trial_allowance: int = 100,
        tol: float = 1e-6,
        maxiter: int = 10000,
        dtype: npt.DTypeLike = np.float64,
    ):
        self.k = k
        self.unbalanced = unbalanced
        self.init = init
        self.solver = solver
        self.beta_loss = beta_loss
        self.alpha_W = alpha_W
        self.alpha_H = alpha_H
        self.random_state = np.random.RandomState(seed=random_state)
        self.trial_allowance = trial_allowance
        self.tol = tol
        self.maxiter = maxiter
        self.dtype = dtype

        self.n_samples_ = None
        self.n_features_ = None
        self.n_nodes_ = 0
        self.n_leaves_ = 0
        self.tree_ = None
        self.splits_ = None
        self.is_leaf_ = None
        self.clusters_ = None
        self.Ws_ = None
        self.Hs_ = None
        self.W_buffer_ = None
        self.H_buffer_ = None
        self.priorities_ = None
        self.id2sample_ = None
        self.id2feature_ = None
        self.feature2id_ = None

#### `H_buffer_ = None` `instance-attribute`

#### `Hs_ = None` `instance-attribute`

#### `W_buffer_ = None` `instance-attribute`

#### `Ws_ = None` `instance-attribute`

#### `alpha_H = alpha_H` `instance-attribute`

#### `alpha_W = alpha_W` `instance-attribute`

#### `beta_loss = beta_loss` `instance-attribute`

#### `clusters_ = None` `instance-attribute`

#### `dtype = dtype` `instance-attribute`

#### `feature2id_ = None` `instance-attribute`

#### `id2feature_ = None` `instance-attribute`

#### `id2sample_ = None` `instance-attribute`

#### `init = init` `instance-attribute`

#### `is_leaf_ = None` `instance-attribute`

#### `k = k` `instance-attribute`

#### `maxiter = maxiter` `instance-attribute`

#### `n_features_ = None` `instance-attribute`

#### `n_leaves_ = 0` `instance-attribute`

#### `n_nodes_ = 0` `instance-attribute`

#### `n_samples_ = None` `instance-attribute`

#### `priorities_ = None` `instance-attribute`

#### `random_state = np.random.RandomState(seed=random_state)` `instance-attribute`

#### `solver = solver` `instance-attribute`

#### `splits_ = None` `instance-attribute`

#### `tol = tol` `instance-attribute`

#### `tree_ = None` `instance-attribute`

#### `trial_allowance = trial_allowance` `instance-attribute`

#### `unbalanced = unbalanced` `instance-attribute`

#### `_init_fit(X, term_subset)`

Source code in `src/hnmf/model.py`

    183
    184
    185
    186
    187
    188
    189
    190
    191
    192
    193
    194
    195
    196
    197
    198
    199
    200
    201
    202
    203
    204
    205
    206
    207
    def _init_fit(
        self, X: npt.NDArray, term_subset: npt.NDArray
    ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
        if not self.n_samples_:
            raise ValueError("n_samples_ not set before _init_fit called")

        nmf = NMF(
            n_components=2,
            random_state=self.random_state,
            tol=self.tol,
            max_iter=self.maxiter,
            init=self.init,
        )

        if len(term_subset) == self.n_samples_:
            W = nmf.fit_transform(X)
            H = nmf.components_
            return W, H

        W_tmp = nmf.fit_transform(X[term_subset, :])
        H = nmf.components_
        W = np.zeros((self.n_samples_, 2), dtype=self.dtype)
        W[term_subset, :] = W_tmp

        return W, H

#### `_stack_H_buffer(buffer)`

Source code in `src/hnmf/model.py`

    402
    403
    404
    405
    406
    407
    408
    409
    410
    411
    412
    413
    def _stack_H_buffer(self, buffer: list) -> npt.NDArray:
        if self.n_features_ is None:
            raise ValueError("n_features_ not set before _stack_H_buffer called")
        if self.clusters_ is None:
            raise ValueError("clusters_ not set before _stack_H_buffer called")

        result = np.zeros((len(buffer), 2, self.n_features_), dtype=self.dtype)
        for i, buff in enumerate(buffer):
            cluster_nz_idx = np.argwhere(self.clusters_[i]).flatten()
            result[i, 0, cluster_nz_idx] = buff[0, :]
            result[i, 1, cluster_nz_idx] = buff[1, :]
        return result

#### `_stack_clusters(clusters)`

Source code in `src/hnmf/model.py`

    394
    395
    396
    397
    398
    399
    400
    def _stack_clusters(self, clusters: list[npt.NDArray | None]) -> npt.NDArray:
        if not self.n_features_:
            raise ValueError("n_features_ not set before _stack_clusters called")
        result = np.zeros((len(clusters), self.n_features_), dtype=np.int64)
        for i, cluster in enumerate(clusters):
            result[i, cluster] = 1
        return result

#### `cluster_assignments(leaves_only=True, include_outliers=True)`

Returns a mapping of features and their assigned cluster(s)

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`leaves_only` | `bool` | Whether to return only leaf nodes | `True`
`include_outliers` | `bool` | If True, include feature_idx keys that are not assigned a cluster. | `True`

Source code in `src/hnmf/model.py`

    634
    635
    636
    637
    638
    639
    640
    641
    642
    643
    644
    645
    646
    647
    648
    649
    650
    651
    652
    653
    654
    655
    656
    657
    658
    659
    660
    661
    662
    663
    664
    665
    666
    667
    668
    669
    670
    671
    672
    def cluster_assignments(
        self,
        leaves_only: bool = True,
        include_outliers: bool = True,
    ) -> dict[int, set[int]]:
        """
        Returns a mapping of features and their assigned cluster(s)

        Parameters
        ----------
        leaves_only
            Whether to return only leaf nodes
        include_outliers
            If True, include feature_idx keys that are not assigned a cluster.

        """

        if self.clusters_ is None:
            raise ValueError("Model not fitted, clusters_ is None")

        node_leaf_idx = np.where(self.is_leaf_ == 1)[0]

        clusters = self.clusters_
        output = defaultdict(set)
        assignments = np.argwhere(clusters)
        if leaves_only:
            assignments = assignments[
                np.where(np.isin(assignments[:, 0], node_leaf_idx))[0]
            ]

        for cluster_idx, feature_idx in assignments:
            output[cluster_idx].add(feature_idx)

        if include_outliers:
            outliers = np.where(clusters.sum(axis=0) == 0)[0]
            for outlier in outliers:
                output[outlier] = set()

        return dict(output)

#### `cluster_features(leaves_only=True, include_outliers=True)`

Returns the features assigned as a cluster to nodes

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`leaves_only` | `bool` | Whether to return only leaf nodes | `True`
`include_outliers` | `bool` | If True, features without a node assignment are returned under the key -1 | `True`

Source code in `src/hnmf/model.py`

    597
    598
    599
    600
    601
    602
    603
    604
    605
    606
    607
    608
    609
    610
    611
    612
    613
    614
    615
    616
    617
    618
    619
    620
    621
    622
    623
    624
    625
    626
    627
    628
    629
    630
    631
    632
    def cluster_features(
        self,
        leaves_only: bool = True,
        include_outliers: bool = True,
    ) -> dict[int, list[int]]:
        """
        Returns the features assigned as a cluster to nodes

        Parameters
        ----------
        leaves_only
            Whether to return only leaf nodes
        include_outliers
            If True, features without a node assignment are returned under the key -1

        """

        if self.clusters_ is None:
            raise ValueError("Model not fitted, clusters_ is None")

        output = defaultdict(list)

        node_leaf_idx = np.where(self.is_leaf_ == 1)[0]

        clusters = self.clusters_[node_leaf_idx] if leaves_only else self.clusters_

        assignments = np.argwhere(clusters)

        for cluster_idx, feature_idx in assignments:
            output[cluster_idx].append(feature_idx)

        if include_outliers:
            outliers = np.where(clusters.sum(axis=0) == 0)[0]
            output[-1] = outliers

        return dict(output)

#### `fit(X)`

Fit `HierarchicalNMF` to data

Source code in `src/hnmf/model.py`

    209
    210
    211
    212
    213
    214
    215
    216
    217
    218
    219
    220
    221
    222
    223
    224
    225
    226
    227
    228
    229
    230
    231
    232
    233
    234
    235
    236
    237
    238
    239
    240
    241
    242
    243
    244
    245
    246
    247
    248
    249
    250
    251
    252
    253
    254
    255
    256
    257
    258
    259
    260
    261
    262
    263
    264
    265
    266
    267
    268
    269
    270
    271
    272
    273
    274
    275
    276
    277
    278
    279
    280
    281
    282
    283
    284
    285
    286
    287
    288
    289
    290
    291
    292
    293
    294
    295
    296
    297
    298
    299
    300
    301
    302
    303
    304
    305
    306
    307
    308
    309
    310
    311
    312
    313
    314
    315
    316
    317
    318
    319
    320
    321
    322
    323
    324
    325
    326
    327
    328
    329
    330
    331
    332
    333
    334
    335
    336
    337
    338
    339
    340
    341
    342
    343
    344
    345
    346
    347
    348
    349
    350
    351
    352
    353
    354
    355
    356
    357
    358
    359
    360
    361
    362
    363
    364
    365
    366
    367
    368
    369
    370
    371
    372
    373
    374
    375
    376
    377
    378
    379
    380
    381
    382
    383
    384
    385
    386
    387
    388
    389
    390
    391
    392
    def fit(self, X: npt.NDArray) -> Self:
        """
        Fit `HierarchicalNMF` to data
        """
        shape: tuple[int, int] = X.shape
        n_samples, n_features = shape
        self.n_samples_ = n_samples
        self.n_features_ = n_features

        # TODO Expect different sized ranks
        clusters: list[npt.NDArray[np.int64] | None] = [None] * (2 * (self.k - 1))
        Ws = [None] * (2 * (self.k - 1))
        Hs = [None] * (2 * (self.k - 1))
        W_buffer = [None] * (2 * (self.k - 1))
        H_buffer = [None] * (2 * (self.k - 1))
        priorities = np.zeros(2 * (self.k - 1), dtype=self.dtype)
        is_leaf = np.zeros(2 * (self.k - 1), dtype=np.bool)  # No leaves at start
        tree = np.zeros((2, 2 * (self.k - 1)), dtype=np.int64)
        splits = -np.ones(self.k - 1, dtype=np.int64)

        # Where X has at least one non-zero
        term_subset = np.flatnonzero(np.sum(X, axis=1))

        W, H = self._init_fit(X, term_subset)

        result_used = 0

        with ProgressTree() as pt:
            for i in range(self.k - 1):
                if i == 0:
                    split_node = 0
                    new_nodes = [0, 1]
                    min_priority = 1e40
                    split_subset = np.arange(n_features)
                else:
                    leaves = np.where(is_leaf == 1)[0]
                    temp_priority = priorities[leaves]

                    if len(np.where(temp_priority > 0)[0]) > 0:
                        min_priority = np.min(temp_priority[temp_priority > 0])
                        split_node = np.argmax(temp_priority)
                    else:  # There are no more candidates stop early
                        min_priority = -1
                        split_node = 0

                    if temp_priority[split_node] < 0 or min_priority == -1:
                        logger.warning(
                            f"Cannot generate all {self.k} leaf clusters, stopping at {i} leaf clusters"
                        )

                        Ws = [i for i in Ws if i is not None]
                        W_buffer = [i for i in W_buffer if i is not None]

                        Hs = [i for i in Hs if i is not None]
                        H_buffer = [i for i in H_buffer if i is not None]

                        # Resize attributes
                        tree = tree[:, :result_used]
                        splits = splits[:result_used]
                        is_leaf = is_leaf[:result_used]
                        clusters = clusters[:result_used]
                        priorities = priorities[:result_used]

                        self.tree_ = tree.T
                        self.splits_ = splits
                        self.is_leaf_ = is_leaf
                        self.n_nodes_ = self.is_leaf_.shape[0]
                        self.n_leaves_ = int(np.count_nonzero(self.is_leaf_))
                        self.clusters_ = self._stack_clusters(clusters)
                        self.Ws_ = np.array(Ws)
                        self.Hs_ = np.array(Hs)
                        self.W_buffer_ = np.array(W_buffer)
                        self.H_buffer_ = self._stack_H_buffer(H_buffer)
                        self.priorities_ = priorities
                        return self

                    split_node = leaves[split_node]  # Attempt to split this node
                    is_leaf[split_node] = 0
                    W = W_buffer[split_node]
                    H = H_buffer[split_node]

                    # Find which features are clustered on this node
                    split_subset = clusters[split_node]
                    new_nodes = [result_used, result_used + 1]
                    tree[:, split_node] = new_nodes

                result_used += 2
                # For each row find where it is more greatly represented
                cluster_subset = np.argmax(H, axis=0)

                subset_0 = np.flatnonzero(cluster_subset == 0)
                subset_1 = np.flatnonzero(cluster_subset == 1)
                ls0 = len(subset_0)
                ls1 = len(subset_1)

                if i == 0:
                    pt.add_branch("Root", new_nodes[0], ls0)
                    pt.add_branch("Root", new_nodes[1], ls1)
                else:
                    pt.add_branch(split_node, new_nodes[0], ls0)
                    pt.add_branch(split_node, new_nodes[1], ls1)

                clusters[new_nodes[0]] = split_subset[subset_0]
                clusters[new_nodes[1]] = split_subset[subset_1]
                Ws[new_nodes[0]] = W[:, 0]
                Ws[new_nodes[1]] = W[:, 1]

                # These will not have shape of (2, n_features) because they are fitting a subset
                # Create zero filled array of shape (2, n_features)
                h_temp = np.zeros(shape=(2, self.n_features_), dtype=self.dtype)
                # Which features are present in H

                h_temp[0, split_subset] = H[0]
                h_temp[1, split_subset] = H[1]

                Hs[new_nodes[0]] = h_temp[0]
                Hs[new_nodes[1]] = h_temp[1]

                splits[i] = split_node
                is_leaf[new_nodes] = 1

                subset = clusters[new_nodes[0]]
                (
                    subset,
                    W_buffer_one,
                    H_buffer_one,
                    priority_one,
                ) = trial_split_sklearn(
                    min_priority=min_priority,
                    X=X,
                    subset=subset,
                    W_parent=W[:, 0],
                    random_state=self.random_state,
                    trial_allowance=self.trial_allowance,
                    unbalanced=self.unbalanced,
                    dtype=self.dtype,
                    tol=self.tol,
                    maxiter=self.maxiter,
                    init=self.init,
                    alpha_W=self.alpha_W,
                    alpha_H=self.alpha_H,
                )
                clusters[new_nodes[0]] = subset
                W_buffer[new_nodes[0]] = W_buffer_one
                H_buffer[new_nodes[0]] = H_buffer_one
                priorities[new_nodes[0]] = priority_one

                subset = clusters[new_nodes[1]]
                (
                    subset,
                    W_buffer_one,
                    H_buffer_one,
                    priority_one,
                ) = trial_split_sklearn(
                    min_priority=min_priority,
                    X=X,
                    subset=subset,
                    W_parent=W[:, 1],
                    random_state=self.random_state,
                    trial_allowance=self.trial_allowance,
                    unbalanced=self.unbalanced,
                    dtype=self.dtype,
                    tol=self.tol,
                    maxiter=self.maxiter,
                    init=self.init,
                    alpha_W=self.alpha_W,
                    alpha_H=self.alpha_H,
                )
                clusters[new_nodes[1]] = subset
                W_buffer[new_nodes[1]] = W_buffer_one
                H_buffer[new_nodes[1]] = H_buffer_one
                priorities[new_nodes[1]] = priority_one
        self.tree_ = tree.T
        self.splits_ = splits
        self.is_leaf_ = is_leaf
        self.clusters_ = self._stack_clusters(clusters)
        self.Ws_ = np.array(Ws)
        self.Hs_ = np.array(Hs)
        self.W_buffer_ = np.array(W_buffer)
        self.H_buffer_ = self._stack_H_buffer(H_buffer)
        self.priorities_ = priorities
        self.n_nodes_ = self.is_leaf_.shape[0]
        self.n_leaves_ = int(np.count_nonzero(self.is_leaf_))
        return self

#### `top_discriminative_samples_in_node(node, n=10, sign='abs')`

Computes most discriminative samples (node vs rest)

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`node` | `int` |  | required
`n` | `int` | The number of features to return | `10`
`sign` | `Literal['positive', 'negative', 'abs']` | One of `['positive', 'negative', 'abs']`. | `'abs'`

Returns:

Type | Description
--- | ---
`list of dict with form::` | sample: Any node: int node_value: float others_value: float

Source code in `src/hnmf/model.py`

    535
    536
    537
    538
    539
    540
    541
    542
    543
    544
    545
    546
    547
    548
    549
    550
    551
    552
    553
    554
    555
    556
    557
    558
    559
    560
    561
    562
    563
    564
    565
    566
    567
    568
    569
    570
    571
    572
    573
    574
    575
    576
    577
    578
    579
    580
    581
    582
    583
    584
    585
    586
    587
    588
    589
    590
    591
    592
    593
    594
    595
    def top_discriminative_samples_in_node(
        self,
        node: int,
        n: int = 10,
        sign: Literal["positive", "negative", "abs"] = "abs",
    ) -> "list[DiscriminatedSample]":
        """
        Computes most discriminative samples (node vs rest)

        Parameters
        ----------
        node
        n
            The number of features to return
        sign
            One of `['positive', 'negative', 'abs']`.

        Returns
        --------
        list of dict with form::

            sample: Any
            node: int
            node_value: float
            others_value: float

        """

        if self.Ws_ is None:
            raise ValueError("Model not fitted, Ws_ is None")
        if sign not in ("positive", "negative", "abs"):
            raise ValueError("Sign must be one of 'positive', 'negative' or 'abs'")

        # Masks
        member_mask = np.array(node, dtype=np.int64)
        non_member_mask = np.array(
            [x for x in np.arange(0, self.n_nodes_) if x != node]
        )

        member_values = self.Ws_[member_mask].ravel()
        other_means = self.Ws_[non_member_mask].mean(axis=0)

        diffs = (
            np.abs(member_values - other_means)
            if sign == "positive"
            else member_values - other_means
            if sign == "positive"
            else other_means - member_values
        )

        diff_tops = diffs.argsort()[::-1][:n]

        return [
            DiscriminatedSample(
                sample=diff,
                node=node,
                node_value=member_values[diff],
                others_value=other_means[diff],
            )
            for diff in diff_tops
        ]

#### `top_features_in_node(node, n=10)`

For a given node, return the top n features and values

Source code in `src/hnmf/model.py`

    415
    416
    417
    418
    419
    420
    421
    422
    423
    424
    425
    def top_features_in_node(self, node: int, n: int = 10) -> list[tuple]:
        """
        For a given node, return the top n features and values
        """

        if self.Hs_ is None:
            raise ValueError("Model not fitted, Hs_ is None")

        node_i = self.Hs_[node]
        ranks = node_i.argsort()[::-1][:n]
        return [(i, node_i[i]) for i in ranks if node_i[i] > 0]

#### `top_nodes_in_feature(feature_idx, n=10, leaves_only=True)`

Returns the top nodes for a specified feature

Source code in `src/hnmf/model.py`

    427
    428
    429
    430
    431
    432
    433
    434
    435
    436
    437
    438
    439
    440
    441
    442
    443
    444
    445
    446
    447
    def top_nodes_in_feature(
        self,
        feature_idx: int | str,
        n: int = 10,
        leaves_only: bool = True,
    ) -> list[tuple]:
        """
        Returns the top nodes for a specified feature
        """
        if self.Hs_ is None:
            raise ValueError("Model not fitted, Hs_ is None")

        node_leaf_idx = np.where(self.is_leaf_ == 1)[0]
        node_weights = self.Hs_.T[feature_idx]
        ranks = node_weights.argsort()[::-1]
        if leaves_only:
            ranks = ranks[np.isin(ranks, node_leaf_idx)]

        ranks = ranks[:n]

        return [(i, node_weights[i]) for i in ranks if node_weights[i] > 0]

#### `top_nodes_in_samples(n=10, leaves_only=True)`

Returns the top nodes for each sample.

Source code in `src/hnmf/model.py`

    449
    450
    451
    452
    453
    454
    455
    456
    457
    458
    459
    460
    461
    462
    463
    464
    465
    466
    467
    468
    469
    470
    471
    472
    473
    474
    475
    476
    477
    478
    479
    480
    481
    482
    483
    484
    485
    486
    487
    488
    489
    490
    491
    def top_nodes_in_samples(self, n: int = 10, leaves_only: bool = True):
        """
        Returns the top nodes for each sample.
        """

        if self.Ws_ is None or self.n_nodes_ is None:
            raise ValueError("Model not fitted, Ws_ is None")

        # Idx of leaves
        node_leaf_idx = np.where(self.is_leaf_ == 1)[0]
        # Keep map of enumerated -> actual cluster
        if leaves_only:
            node_map = dict(enumerate(node_leaf_idx))
        else:
            node_map = dict(enumerate(range(self.n_nodes_)))

        # A dictionary of {sample : [top_nodes]}

        output = {}

        # Ws_ is shape n_nodes, n_samples
        # Transpose weights so it has samples as rows, nodes as columns

        weights = self.Ws_.T[node_leaf_idx].T if leaves_only else self.Ws_.T

        # The ellipsis indicates that the selection is done row wise
        sample_tops = weights.argsort()[:, ::-1][:, :n]

        # Create an array with samples as rows, top n weights as columns
        sample_top_weights = np.take_along_axis(weights, sample_tops, axis=1)

        for sample_idx, (node_ids, node_weights) in enumerate(
            zip(sample_tops, sample_top_weights, strict=True)
        ):
            tops = [
                (node_map[node_id], weight)
                for node_id, weight in zip(node_ids, node_weights, strict=True)
                if weight > 0
            ]
            tops.sort(key=itemgetter(1), reverse=True)
            output[sample_idx] = tops

        return output

#### `top_samples_in_nodes(n=10, leaves_only=True)`

Returns the top samples for each node

Source code in `src/hnmf/model.py`

    493
    494
    495
    496
    497
    498
    499
    500
    501
    502
    503
    504
    505
    506
    507
    508
    509
    510
    511
    512
    513
    514
    515
    516
    517
    518
    519
    520
    521
    522
    523
    524
    525
    526
    527
    528
    529
    530
    531
    532
    533
    def top_samples_in_nodes(self, n: int = 10, leaves_only: bool = True):
        """
        Returns the top samples for each node
        """

        if self.Ws_ is None:
            raise ValueError("Model not fitted, Ws_ is None")

        # Idx of leaves
        node_leaf_idx = np.where(self.is_leaf_ == 1)[0]

        # A dictionary of {nodes : [sample]}

        output = {}

        # Ws_ is shape n_nodes, n_samples

        weights = self.Ws_

        # The ellipsis indicates that the selection is done row wise
        node_tops = weights.argsort()[:, ::-1][:, :n]

        # Create an array with samples as rows, top n weights as columns
        node_top_weights = np.take_along_axis(weights, node_tops, axis=1)

        for node_idx, (sample_ids, sample_weights) in enumerate(
            zip(node_tops, node_top_weights, strict=True)
        ):
            if leaves_only and node_idx not in node_leaf_idx:
                continue
            tops = [
                (sample_id, weight)
                for sample_id, weight in zip(sample_ids, sample_weights, strict=True)
                if weight > 0
            ]
            tops.sort(key=itemgetter(1), reverse=True)
            # Decode samples if available

            output[node_idx] = tops

        return output

## hNMF.helpers

### `hnmf.helpers`

#### `nmfsh_comb_rank2(A, Winit, Hinit, anls_alg, vec_norm, normW, tol, maxiter, dtype)`

Source code in `src/hnmf/helpers.py`

    407
    408
    409
    410
    411
    412
    413
    414
    415
    416
    417
    418
    419
    420
    421
    422
    423
    424
    425
    426
    427
    428
    429
    430
    431
    432
    433
    434
    435
    436
    437
    438
    439
    440
    441
    442
    443
    444
    445
    446
    447
    448
    449
    450
    451
    452
    453
    454
    455
    456
    457
    458
    459
    460
    461
    462
    463
    464
    465
    466
    467
    468
    469
    470
    471
    472
    473
    474
    475
    476
    477
    478
    479
    480
    481
    482
    483
    484
    485
    486
    487
    488
    489
    490
    491
    492
    493
    494
    495
    496
    497
    498
    499
    500
    501
    502
    503
    504
    505
    506
    507
    def nmfsh_comb_rank2(
        A: npt.NDArray,
        Winit: npt.NDArray,
        Hinit: npt.NDArray,
        anls_alg: AnlsAlgorithm,
        vec_norm: float,
        normW: bool,
        tol: float,
        maxiter: int,
        dtype: npt.DTypeLike,
    ) -> tuple[npt.NDArray, npt.NDArray]:
        """"""
        eps = 1e-6
        shape: tuple[int, int] = A.shape
        m, n = shape
        W, H = Winit, Hinit

        if W.shape[1] != 2:
            warnings.warn(
                f"Error: Wrong size of W! Expected shape of (n, 2) but received W of shape ({W.shape[0]}, {W.shape[1]})",
                stacklevel=2,
            )

        if H.shape[0] != 2:
            warnings.warn(
                f"Error: Wrong size of H! Expected shape of (2, n) but received H of shape ({H.shape[0]}, {H.shape[1]})",
                stacklevel=2,
            )

        left = H.dot(H.T)
        right = A.dot(H.T)
        for iter_ in range(maxiter):
            if matrix_rank(left) < 2:
                W = np.zeros((m, 2), dtype=dtype)
                H = np.zeros((2, n), dtype=dtype)
                if sp.issparse(A):
                    U, _S, V = svd(A.toarray(), full_matrices=False)  # type: ignore[attr-defined]  # A can be sparse
                else:
                    U, _S, V = svd(A, full_matrices=False)
                U, V = U[:, 0], V[0, :]
                if sum(U) < 0:
                    U, V = -U, -V

                W[:, 0] = U
                H[0, :] = V

                return W, H

            W = anls_alg(left, right, W, dtype)
            norms_W = norm(W, axis=0)
            if np.min(norms_W) < eps:
                logger.warning("Error: Some column of W is essentially zero")

            W *= 1.0 / norms_W
            left = W.T.dot(W)
            right = A.T.dot(W)
            if matrix_rank(left) < 2:
                W = np.zeros((m, 2), dtype=dtype)
                H = np.zeros((2, n), dtype=dtype)
                if sp.issparse(A):
                    U, _S, V = svd(A.toarray(), full_matrices=False)  # type: ignore[attr-defined]  # A can be sparse
                else:
                    U, _S, V = svd(A, full_matrices=False)
                U, V = U[:, 0], V[0, :]
                if sum(U) < 0:
                    U, V = -U, -V

                W[:, 0] = U
                H[0, :] = V

                return W, H

            H = anls_alg(left, right, H.T, dtype).T
            gradH = left.dot(H) - right.T
            left = H.dot(H.T)
            right = A.dot(H.T)
            gradW = W.dot(left) - right
            initgrad = 1
            if iter_ == 0:
                gradW_square = np.sum(np.power(gradW[np.logical_or(gradW <= 0, W > 0)], 2))
                gradH_square = np.sum(np.power(gradH[np.logical_or(gradH <= 0, H > 0)], 2))
                initgrad = np.sqrt(gradW_square + gradH_square)
                continue
            gradW_square = np.sum(np.power(gradW[np.logical_or(gradW <= 0, W > 0)], 2))
            gradH_square = np.sum(np.power(gradH[np.logical_or(gradH <= 0, H > 0)], 2))
            projnorm = np.sqrt(gradW_square + gradH_square)

            if projnorm < tol * initgrad:
                break

        if vec_norm != 0:
            if normW:
                norms = np.power(np.sum(np.power(W, vec_norm), axis=0), 1 / vec_norm)
                W /= norms
                H *= norms[:, None]
            else:
                norms = np.power(np.sum(np.power(H, vec_norm), axis=1), 1 / vec_norm)
                W *= norms[None, :]
                H /= norms

        return W, H