exprmat 0.1.41__tar.gz → 0.1.43__tar.gz

{exprmat-0.1.41 → exprmat-0.1.43}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: exprmat
-Version: 0.1.41
+Version: 0.1.43
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: Implementation :: CPython
 Classifier: Programming Language :: Python :: Implementation :: PyPy
@@ -75,6 +75,8 @@ Requires-Dist: pyfastx
 Requires-Dist: genomicranges
 Requires-Dist: rustworkx
 Requires-Dist: xgboost
+Requires-Dist: simpleppt
+Requires-Dist: elpigraph-python
 Requires-Dist: dask-expr>=2.0
 Requires-Dist: sphinx ; extra == 'dev'
 Requires-Dist: twine ; extra == 'dev'

{exprmat-0.1.41 → exprmat-0.1.43}/exprmat/__init__.py

@@ -108,10 +108,11 @@ for finder in default_finders:
 basepath = config['data']
 
 # core method exports
-from exprmat.utils import setup_styles
+from exprmat.utils import setup_styles, savefig
 from exprmat.ansi import error, warning, info, pprog, pproga
 from exprmat.reader.experiment import experiment, load_experiment
 from exprmat.reader.metadata import metadata, load_metadata
+import exprmat.reader.static as st
 
 
 setup_styles()
@@ -127,7 +128,8 @@ def version_db():
         with open(version_file, 'r') as vf:
             db_ver = vf.read().strip()
     
-    return db_ver
+    if db_ver is None: return None
+    return tuple([int(x) for x in db_ver.split('.')])
 
 
 # perform database integrity check.
@@ -136,9 +138,9 @@ if PERFORM_DATABASE_CHECK:
     if db_ver is None: 
         warning('the database is not installed.')
         print(DATABASE_SETUP_MESSAGE)
-    elif db_ver != importlib.metadata.version("exprmat"):
+    elif ".".join([str(x) for x in db_ver]) != importlib.metadata.version("exprmat"):
         warning('the database version do not match the package version')
-        warning(f'db version: {db_ver}  package version: {importlib.metadata.version("exprmat")}')
+        warning(f'db version: {".".join([str(x) for x in db_ver])}  package version: {importlib.metadata.version("exprmat")}')
 
 
 def cuda():
@@ -168,11 +170,11 @@ def version():
     from exprmat.ansi import info, error, format_file_size
 
     db_ver = version_db()
-    info(f'exprmat {MAJOR}.{MINOR}.{REVISION} / exprmat-db {db_ver if db_ver is not None else "(Not installed)"}')
+    info(f'exprmat {MAJOR}.{MINOR}.{REVISION} / exprmat-db {".".join([str(x) for x in db_ver]) if db_ver is not None else "(Not installed)"}')
     info(f'os: {os.name} ({sys.platform})  platform version: {platform.release()}')
     info(f'loaded configuration from {finder}')
     info(f'current working directory: {os.getcwd()}')
-    info(f'current database directory: {config["data"]} ({db_ver if db_ver is not None else "(Not found)"})')
+    info(f'current database directory: {config["data"]} ({".".join([str(x) for x in db_ver]) if db_ver is not None else "(Not found)"})')
     memory()
     return (MAJOR, MINOR, REVISION)
 
@@ -206,5 +208,8 @@ __all__ = [
     'warning',
     'info',
     'pprog',
-    'pproga'
+    'pproga',
+    'savefig',
+
+    'st',
 ]

{exprmat-0.1.41 → exprmat-0.1.43}/exprmat/ansi.py

@@ -1,6 +1,7 @@
 
 import re
 import os
+import sys
 
 from functools import partial
 from tqdm import tqdm
@@ -91,12 +92,12 @@ def common_length(string: str, limit: int) -> str:
         return x_str
     
     elif len(string) + cn_length > limit:
-        x_str = string[-(limit - 4):]
+        x_str = string[-(limit - 3):]
         cn = len(re.findall(cn_re, x_str))
-        while cn + len(x_str) > limit - 4:
+        while cn + len(x_str) > limit - 3:
             x_str = x_str[1:]
             cn = len(re.findall(cn_re, x_str))
-        return '... ' + x_str
+        return '.. ' + x_str
         
     else: return ('{:<' + str(limit - cn_length) + '}').format(string)
 
@@ -156,13 +157,22 @@ def yellow(x): return '\033[33m' + x + '\033[0m'
 def blue(x): return '\033[34m' + x + '\033[0m'
 def purple(x): return '\033[35m' + x + '\033[0m'
 def cyan(x): return '\033[36m' + x + '\033[0m'
+def annot(x): return '\033[90;3m' + x + '\033[0m'
 
 
 progress_styles = {
     'ncols': 80,
     'ascii': '-━',
-    'bar_format': '   {bar} {desc:20} {n:5d} / {total:<5d} ({elapsed} < {remaining})'
+    'bar_format': '   {bar} {desc:20} {n:5d} / {total:<5d} ({elapsed} < {remaining})',
+    'file': sys.stderr
 }
 
-pprog = partial(tqdm, **progress_styles)
-pproga = partial(tqdma, **progress_styles)
+
+class pprog(tqdm):
+    def __init__(self, iterable = None, **kwargs):
+        super().__init__(iterable, **progress_styles, **kwargs)
+
+
+class pproga(tqdma):
+    def __init__(self, iterable = None, **kwargs):
+        super().__init__(iterable, **progress_styles, **kwargs)

{exprmat-0.1.41 → exprmat-0.1.43}/exprmat/cnv/infercnv.py

@@ -12,7 +12,7 @@ from anndata import AnnData
 from scanpy import logging
 from tqdm.contrib.concurrent import process_map
 
-from exprmat.ansi import error, warning, info, pproga
+from exprmat.ansi import error, warning, info, pproga, tqdma
 from exprmat.utils import ensure_array
 
 
@@ -143,9 +143,9 @@ def infercnv(
             itertools.repeat(dynamic_threshold),
             itertools.repeat(calculate_gene_values),
             tqdm_class = pproga,
-            max_workers=cpu_count() if n_jobs is None else n_jobs,
+            max_workers = cpu_count() if n_jobs is None else n_jobs,
         ),
-        strict=False,
+        strict = False,
     )
 
     res = scipy.sparse.vstack(chunks)

{exprmat-0.1.41 → exprmat-0.1.43}/exprmat/descriptive/tcr.py

@@ -101,8 +101,6 @@ def transition(adata, clonotype = 'clone.id', cluster = 'leiden', key = 'tcr.clo
         adata.obs.loc[adata.obs[clonotype] == cl, key] = shannon_e
 
 
-
-
 def migration(adata, clonotype = 'clone.id', cluster = 'tissue', key = 'tcr.clone.migr'):
     # completely the same.
     transition(adata, clonotype = clonotype, cluster = cluster, key = key)

{exprmat-0.1.41 → exprmat-0.1.43}/exprmat/plotting/cnmf.py

@@ -129,6 +129,10 @@ def matrix_plot(
     if annotations is not None and isinstance(annotations, list):
         ax_annot = ax.inset_axes([1.02, 0, 0.05, 1], sharey = ax)
         annot_u = list(set(annotations))
+        # sort it
+        if all([str(x).isdigit() for x in annot_u]):
+            annot_u = sorted(annot_u, key = int)
+        else: annot_u = sorted(annot_u, key = str)
         annot_n = [annot_u.index(x) for x in annotations]
 
         from exprmat.plotting import palettes

{exprmat-0.1.41 → exprmat-0.1.43}/exprmat/plotting/cnv.py

@@ -60,7 +60,7 @@ def chromosome_heatmap(
 
     labels = list(chr_pos_dict.keys())
     if draw_autosome_only:
-        labels = [x if len(x) <= 2 else '' for x in labels]
+        labels = [x if len(x) <= 5 else '' for x in labels]
 
     return_ax_dic = sc.pl.heatmap(
         tmp_adata,
@@ -70,7 +70,7 @@ def chromosome_heatmap(
         cmap = mpl(cmap),
         show_gene_labels = False,
         var_group_positions = var_group_positions,
-        var_group_labels = list(chr_pos_dict.keys()),
+        var_group_labels = labels,
         show = False,
         ** kwargs,
     )

exprmat-0.1.43/exprmat/plotting/milestone.py

@@ -0,0 +1,118 @@
+
+from typing_extensions import Literal
+import igraph
+import numpy as np
+import pandas as pd
+from typing import Union, Optional, Iterable
+import matplotlib
+import matplotlib.cm as cm
+import matplotlib.colors as mcolors
+import matplotlib.pyplot as plt
+import scanpy as sc
+
+
+def milestones(
+    adata,
+    basis: Union[None, str] = 'umap',
+    trajectory_key = 'ppt',
+    annotate: bool = False,
+    title = None,
+    subset: Optional[Iterable] = None,
+    ax = None,
+    sort = False,
+    cmap = 'turbo',
+    **kwargs,
+):
+    
+    emb = adata.obsm[basis]
+    kwargs['sort_order'] = sort
+
+    if not sort: order = adata.obs[f'{trajectory_key}.pseudotime'].sort_values().index
+    else: order = adata.obs[f'{trajectory_key}.pseudotime'].sort_values().index
+    order = order[order.isin(subset)] if subset is not None else order
+
+    if "color" in kwargs:
+        kwargs.pop("color")
+        
+    if ax is None:
+        fig = plt.figure()
+        ax = fig.add_subplot(111)
+
+    if "edgecolor" not in kwargs: kwargs["edgecolor"] = "none"
+    if "s" not in kwargs: kwargs["s"] = 120000 / adata.shape[0]
+
+    ax.scatter(
+        adata[order].obsm[basis][:, 0],
+        adata[order].obsm[basis][:, 1],
+        c = gen_milestones_gradients(adata, key = trajectory_key, palette = cmap)[order].values,
+        marker = ".",
+        rasterized = True,
+        plotnonfinite = True,
+        ** kwargs,
+    )
+
+    ax.set_yticks([])
+    ax.set_xticks([])
+    ax.autoscale_view()
+
+    if annotate:
+        
+        R = adata.obsm[trajectory_key]
+        proj = (np.dot(emb.T, R) / R.sum(axis=0)).T
+
+        X = proj[list(adata.uns[f"{trajectory_key}.graph"]["milestones"].values()), :]
+        adata_m = sc.AnnData(
+            X, dtype = X.dtype,
+            obs = dict(mil = list(adata.uns[f"{trajectory_key}.graph"]["milestones"].keys())),
+            obsm = {basis: X},
+        )
+
+        adata_m.obs["mil"] = adata_m.obs["mil"].astype("category")
+
+        sc.pl.embedding(
+            adata_m,
+            basis,
+            color = "mil",
+            title = title,
+            legend_loc = "on data",
+            ax = ax,
+            alpha = 0,
+            legend_fontoutline = True,
+            show = False,
+        )
+
+
+def gen_milestones_gradients(adata, palette = 'turbo', key = 'ppt', seg_order = None):
+
+    seg_order = adata.obs[f'{key}.segs'].unique() \
+        if seg_order is None else seg_order
+
+    if "milestones.colors" not in adata.uns or len(adata.uns["milestones.colors"]) == 1:
+        from exprmat.plotting.palettes import get_palette
+        adata.uns['milestones.colors'] = get_palette(palette, len(seg_order))
+
+    def milestones_prog(s):
+        
+        cfrom = adata.obs.t[adata.obs.seg == s].idxmin()
+        cto = adata.obs.t[adata.obs.seg == s].idxmax()
+        mfrom = adata.obs.milestones[cfrom]
+        mto = adata.obs.milestones[cto]
+        mfrom_c = adata.uns["milestones_colors"][
+            np.argwhere(adata.obs.milestones.cat.categories == mfrom)[0][0]
+        ]
+        mto_c = adata.uns["milestones_colors"][
+            np.argwhere(adata.obs.milestones.cat.categories == mto)[0][0]
+        ]
+
+        from matplotlib.colors import LinearSegmentedColormap, to_hex
+        cm = LinearSegmentedColormap.from_list("test", [mfrom_c, mto_c], N=1000)
+        pst = (
+            adata.obs.t[adata.obs.seg == s] - adata.obs.t[adata.obs.seg == s].min()
+        ) / (
+            adata.obs.t[adata.obs.seg == s].max()
+            - adata.obs.t[adata.obs.seg == s].min()
+        )
+        return pd.Series(list(map(to_hex, cm(pst))), index = pst.index)
+
+    return pd.concat(list(map(milestones_prog, seg_order)))
+

exprmat-0.1.43/exprmat/plotting/trajectory.py

@@ -0,0 +1,345 @@
+
+import numpy as np
+import pandas as pd
+import igraph
+from anndata import AnnData
+import matplotlib.collections
+from typing import Union, Optional, Sequence, Tuple, List
+import scanpy as sc
+from cycler import Cycler
+
+from pandas.api.types import is_categorical_dtype
+from scanpy.plotting._utils import savefig_or_show
+
+import matplotlib.patheffects as path_effects
+import matplotlib.text as mtext
+from matplotlib.cm import ScalarMappable
+from matplotlib.colors import Normalize, hex2color, rgb2hex
+from numba import njit
+import math
+
+from exprmat.plotting.milestone import milestones as milestones_plot
+from exprmat import error, warning, info
+
+
+def graph(
+    adata: AnnData,
+    background_color = None,
+    basis: Union[None, str] = None,
+    size_nodes: float = None,
+    alpha_nodes: float = 1,
+    linewidth: float = 2,
+    alpha_seg: float = 1,
+    color_cells: Union[str, None] = None,
+    tips: bool = True,
+    forks: bool = True,
+    nodes: Optional[List] = [],
+    rasterized: bool = True,
+    ax = None,
+    trajectory_key = 'ppt',
+    embedding_dims = [0, 1],
+    contour_plot = False,
+    figsize = (3, 3),
+    ** kwargs,
+):
+
+    if f"{trajectory_key}.graph" not in adata.uns:
+        error(f'do not find trajectory key `{trajectory_key}`')
+
+    graph = adata.uns[f"{trajectory_key}.graph"]
+    emb = adata.obsm[basis]
+    emb = emb[:, embedding_dims]
+    R = adata.obsm[trajectory_key]
+    proj = (np.dot(emb.T, R) / R.sum(axis = 0)).T
+    B = graph["adjacencies"]
+
+    from exprmat.reduction.plot import embedding
+    
+    if ax: embedding(
+        adata, basis = basis, ax = ax, embedding_dims = embedding_dims, 
+        contour_plot = contour_plot, color = background_color, **kwargs
+    )
+        
+    else: ax = embedding(
+        adata, basis = basis, embedding_dims = embedding_dims, figsize = figsize,
+        contour_plot = contour_plot, color = background_color, **kwargs
+    ).axes[0]
+
+    al = np.array(igraph.Graph.Adjacency((B > 0).tolist(), mode = "undirected").get_edgelist())
+    segs = al.tolist()
+    vertices = proj.tolist()
+    lines = [[tuple(vertices[j]) for j in i] for i in segs]
+    lc = matplotlib.collections.LineCollection(
+        lines, colors = "k", linewidths = linewidth, 
+        alpha = alpha_seg, rasterized = rasterized
+    )
+
+    ax.add_collection(lc)
+    ax.scatter(
+        proj[:, 0], proj[:, 1], 
+        s = size_nodes, c = "k", alpha = alpha_nodes, 
+        rasterized = rasterized
+    )
+
+    bbox = dict(facecolor = "white", alpha = 0.6, edgecolor = "white", pad = 0.1)
+
+    if tips:
+        for tip in graph["tips"]:
+            ax.annotate(
+                tip,
+                (proj[tip, 0], proj[tip, 1]),
+                ha = "center",
+                va = "center",
+                xytext = (0, 0),
+                textcoords = "offset points",
+                bbox = bbox,
+            )
+
+    if forks:
+        for fork in graph["forks"]:
+            ax.annotate(
+                fork,
+                (proj[fork, 0], proj[fork, 1]),
+                ha = "center",
+                va = "center",
+                xytext = (0, 0),
+                textcoords = "offset points",
+                bbox = bbox,
+            )
+
+    if nodes:
+        for node in nodes:
+            ax.annotate(
+                node,
+                (proj[node, 0], proj[node, 1]),
+                ha = "center",
+                va = "center",
+                xytext = (0, 0),
+                textcoords = "offset points",
+                bbox = bbox,
+            )
+    
+    return ax.figure
+            
+
+def trajectory(
+    adata: AnnData,
+    basis: Union[None, str] = None,
+    color_seg = "segments",
+    background_color = None,
+    cmap_seg: str = "viridis",
+    layer_seg: Union[str, None] = "fitted",
+    perc_seg: Union[List, None] = None,
+    color_cells: Union[str, None] = None,
+    scale_path: float = 1,
+    rasterized: bool = True,
+    ax = None,
+    trajectory_key = 'ppt',
+    embedding_dims = [0, 1],
+    contour_plot = False,
+    figsize = (3, 3),
+    annotate = False,
+    annotate_fontsize = 9,
+    **kwargs,
+):
+    
+    from exprmat.reduction.plot import embedding
+    if ax: embedding(
+        adata, basis = basis, ax = ax, embedding_dims = embedding_dims, annotate = False,
+        contour_plot = contour_plot, color = background_color, **kwargs
+    )
+        
+    else: ax = embedding(
+        adata, basis = basis, embedding_dims = embedding_dims, figsize = figsize, annotate = False,
+        contour_plot = contour_plot, color = background_color, **kwargs
+    ).axes[0]
+
+    if f"{trajectory_key}.graph" not in adata.uns:
+        raise ValueError("You need to run `principle_tree_pseudotime` first before plotting.")
+
+    graph = adata.uns[f"{trajectory_key}.graph"]
+    emb = adata.obsm[f"{basis}"]
+    emb = emb[:, embedding_dims]
+    R = adata.obsm[trajectory_key]
+
+    nodes = graph["pp_info"].index
+    proj = pd.DataFrame((np.dot(emb.T, R) / R.sum(axis=0)).T, index = nodes)
+
+    B = graph["adjacencies"]
+    g = igraph.Graph.Adjacency((B > 0).tolist(), mode = "undirected")
+    g.vs[:]["name"] = [v.index for v in g.vs]
+
+    miles_ids = np.concatenate([graph["tips"], graph["forks"]])
+
+    al = np.array(g.get_edgelist())
+    edges = [g.vs[e.tolist()]["name"] for e in al]
+    lines = [[tuple(proj.loc[j]) for j in i] for i in edges]
+    miles_ids = miles_ids[np.isin(miles_ids, proj.index)]
+
+    if color_seg == "milestones":
+        from matplotlib.colors import LinearSegmentedColormap
+
+        rev_dict = dict(zip(graph["milestones"].values(), graph["milestones"].keys()))
+        miles_cat = adata.obs[f"{trajectory_key}.milestones"].cat.categories
+        mil_col = np.array(adata.uns[f"{trajectory_key}.milestones.colors"])
+
+        def get_milestones_gradients(i):
+            start = graph["pp_seg"].iloc[i, 1]
+            end = graph["pp_seg"].iloc[i, 2]
+            mil_path = graph["pp_info"].time[g.get_all_shortest_paths(start, end)[0]]
+            mil_path = (mil_path - mil_path.min()) / (mil_path.max() - mil_path.min())
+            start_col = mil_col[miles_cat == rev_dict[start]][0] if start in rev_dict.keys() else '#e0e0e0'
+            end_col = mil_col[miles_cat == rev_dict[end]][0] if end in rev_dict.keys() else '#e0e0e0'
+
+            edges_mil = pd.Series([
+                mil_path[[first, second]].mean()
+                for first, second in zip(mil_path.index, mil_path.index[1:])],
+                index = [(first, second) for first, second in zip(mil_path.index, mil_path.index[1:])]
+            )
+
+            cmap = LinearSegmentedColormap.from_list("mil", [start_col, end_col])
+            return pd.Series(
+                [rgb2hex(c) for c in cmap(edges_mil)], index = edges_mil.index
+            )
+
+        edge_colors = pd.concat(
+            [get_milestones_gradients(i) for i in range(graph["pp_seg"].shape[0])]
+        )
+
+        edges_tuples = [tuple(e) for e in edges]
+        edge_colors.index = [
+            e if any((np.array(edges_tuples) == e).sum(axis=1) == 2) else e[::-1]
+            for e in edge_colors.index
+        ]
+        edge_colors = edge_colors[edges_tuples]
+
+        color_segs = [hex2color(c) for c in edge_colors]
+        color_mils = []
+        for m in miles_ids:
+            if m in rev_dict.keys(): color_mils.append(mil_col[miles_cat.tolist().index(rev_dict[m])])
+            else: color_mils.append('#e0e0e0')
+
+
+    elif color_seg == "segments":
+        
+        seg_edges = (
+            graph["pp_info"]
+            .loc[np.array(edges).ravel(), "seg"]
+            .values.reshape(-1, 2)[:, 0]
+        )
+
+        seg_col = pd.Series(
+            adata.uns[f"{trajectory_key}.seg.colors"], 
+            index = graph["pp_seg"].n
+        )
+
+        color_segs = [hex2color(seg_col.loc[s]) for s in seg_edges]
+        
+        color_mils = []
+        for m in miles_ids:
+            if m in graph["milestones"].values(): 
+                color_mils.append(hex2color(seg_col.loc[graph["pp_info"].loc[m].seg]))
+            else: color_mils.append('#e0e0e0')
+    
+
+    lc = matplotlib.collections.LineCollection(
+        lines,
+        colors = "k",
+        linewidths = 7.5 * scale_path,
+        zorder = 100,
+        path_effects = [path_effects.Stroke(capstyle = "round")],
+        rasterized = rasterized,
+    )
+
+    ax.add_collection(lc)
+    g = igraph.Graph.Adjacency((B > 0).tolist(), mode = "undirected")
+    seg = graph["pp_seg"].loc[:, ["from", "to"]].values.tolist()
+    bbox = dict(facecolor = "white", alpha = 0.6, edgecolor = "white", pad = 0.1)
+
+    lc = matplotlib.collections.LineCollection(
+        lines,
+        colors = color_segs,
+        linewidths = 5 * scale_path,
+        zorder = 104,
+        path_effects = [path_effects.Stroke(capstyle = "round")],
+        rasterized = rasterized,
+    )
+
+    # milestones
+
+    ax.scatter(
+        proj.loc[miles_ids, 0],
+        proj.loc[miles_ids, 1],
+        zorder = 103,
+        c = "k",
+        s = 200 * scale_path,
+        rasterized = False,
+    )
+
+
+    ax.add_collection(lc)
+    ax.scatter(
+        proj.loc[miles_ids, 0],
+        proj.loc[miles_ids, 1],
+        zorder = 105,
+        c = color_mils,
+        s = 140 * scale_path,
+        rasterized = rasterized,
+    )
+
+    if annotate:
+
+        if color_seg == 'milestones':
+
+            for x in miles_ids.tolist():
+
+                text = mtext.Text(
+                    x = proj.loc[x, 0], y = proj.loc[x, 1], # fontproperties = 'bold',
+                    text = str(x), 
+                    color = 'black', zorder = 106,
+                    ha = 'center', va = 'center', size = annotate_fontsize
+                )
+
+                text.set_path_effects([
+                    matplotlib.patheffects.Stroke(linewidth = 3, foreground = 'white'),
+                    matplotlib.patheffects.Normal()
+                ])
+
+                ax.add_artist(text)
+                pass
+            
+        elif color_seg == 'segments':
+
+            unique_edges = list(set(seg_edges))
+            for uniseg in unique_edges:
+                line_segs = []
+                for l, seg in zip(lines, seg_edges):
+                    if seg == uniseg: line_segs.append(l)
+
+                points = len(line_segs)
+                (fromx, fromy), (tox, toy) = line_segs[points // 2]
+                text = mtext.Text(
+                    x = (fromx + tox) / 2, y = (fromy + toy) / 2, # fontproperties = 'bold',
+                    text = str(uniseg), 
+                    color = 'black', zorder = 106,
+                    ha = 'center', va = 'center', size = annotate_fontsize
+                )
+
+                text.set_path_effects([
+                    matplotlib.patheffects.Stroke(linewidth = 3, foreground = 'white'),
+                    matplotlib.patheffects.Normal()
+                ])
+
+                ax.add_artist(text)
+                pass
+    
+    return ax.figure
+
+
+@njit()
+def cdist_numba(coords, out):
+    for i in range(0, coords.shape[0] - 1):
+        out[i] = math.sqrt(
+            (coords[i, 0] - coords[i + 1, 0]) ** 2
+            + (coords[i, 1] - coords[i + 1, 1]) ** 2
+        )

{exprmat-0.1.41 → exprmat-0.1.43}/exprmat/preprocessing/linear.py

@@ -68,7 +68,15 @@ def normalize_total(
     if target_total is None:
         target_total = np.mean(tots_use)
 
-    b = np.array(float(target_total) / tots_use.T)[0]
+    b = np.array(float(target_total) / tots_use.T)
+    e_norm = mul_row(E, b)
+
+    return e_norm
+
+
+def mul_row(E, b):
+
+    if len(b.shape) == 2: b = b[0]
     if not sp.issparse(E):
         # this will densify the matrix, causing incredibly large memory comsumption
         # when treating with an sparse matrix.
@@ -83,7 +91,7 @@ def normalize_total(
     elif sp.isspmatrix_csc(E):
         e_norm = E.copy()
         e_norm.data *= b[E.indices]
-
+    
     return e_norm