gsMap 1.71.2__py3-none-any.whl → 1.73.0__py3-none-any.whl

gsMap/utils/regression_read.py

@@ -1,55 +1,55 @@
+import os
+
 import numpy as np
 import pandas as pd
-import os
 
 
 # Fun for reading gwas data
 def _read_sumstats(fh, alleles=False, dropna=False):
-    '''
+    """
     Parse gwas summary statistics.
-    '''
-    print('Reading summary statistics from {S} ...'.format(S=fh))
+    """
+    print(f"Reading summary statistics from {fh} ...")
     sumstats = ps_sumstats(fh, alleles=alleles, dropna=dropna)
-    print('Read summary statistics for {N} SNPs.'.format(N=len(sumstats)))
+    print(f"Read summary statistics for {len(sumstats)} SNPs.")
 
     m = len(sumstats)
-    sumstats = sumstats.drop_duplicates(subset='SNP')
+    sumstats = sumstats.drop_duplicates(subset="SNP")
     if m > len(sumstats):
-        print('Dropped {M} SNPs with duplicated rs numbers.'.format(M=m - len(sumstats)))
+        print(f"Dropped {m - len(sumstats)} SNPs with duplicated rs numbers.")
 
     return sumstats
 
 
 def ps_sumstats(fh, alleles=False, dropna=True):
-    '''
+    """
     Parses .sumstats files. See docs/file_formats_sumstats.txt.
-    '''
-
-    dtype_dict = {'SNP': str, 'Z': float, 'N': float, 'A1': str, 'A2': str}
+    """
+    dtype_dict = {"SNP": str, "Z": float, "N": float, "A1": str, "A2": str}
     compression = get_compression(fh)
-    usecols = ['SNP', 'Z', 'N']
+    usecols = ["SNP", "Z", "N"]
     if alleles:
-        usecols += ['A1', 'A2']
+        usecols += ["A1", "A2"]
 
     try:
         x = read_csv(fh, usecols=usecols, dtype=dtype_dict, compression=compression)
     except (AttributeError, ValueError) as e:
-        raise ValueError('Improperly formatted sumstats file: ' + str(e.args))
+        raise ValueError("Improperly formatted sumstats file: " + str(e.args)) from e
 
     if dropna:
-        x = x.dropna(how='any')
+        x = x.dropna(how="any")
 
     return x
 
 
 def get_compression(fh):
-    '''
+    """
     Determin the format of compression used with read_csv?
-    '''
-    if fh.endswith('gz'):
-        compression = 'gzip'
-    elif fh.endswith('bz2'):
-        compression = 'bz2'
+    """
+    if fh.endswith("gz"):
+        compression = "gzip"
+    elif fh.endswith("bz2"):
+        compression = "bz2"
     else:
         compression = None
     # -
@@ -57,120 +57,96 @@ def get_compression(fh):
 
 
 def read_csv(fh, **kwargs):
-    '''
+    """
     Read the csv data
-    '''
-    return pd.read_csv(fh, sep='\s+', na_values='.', **kwargs)
+    """
+    return pd.read_csv(fh, sep=r"\s+", na_values=".", **kwargs)
 
 
-# Fun for reading loading LD scores 
+# Fun for reading loading LD scores
 def which_compression(fh):
-    '''
+    """
     Given a file prefix, figure out what sort of compression to use.
-    '''
-    if os.access(fh + '.bz2', 4):
-        suffix = '.bz2'
-        compression = 'bz2'
-    elif os.access(fh + '.gz', 4):
-        suffix = '.gz'
-        compression = 'gzip'
-    elif os.access(fh + '.parquet', 4):
-        suffix = '.parquet'
-        compression = 'parquet'
-    elif os.access(fh + '.feather', 4):
-        suffix = '.feather'
-        compression = 'feather'
+    """
+    if os.access(fh + ".bz2", 4):
+        suffix = ".bz2"
+        compression = "bz2"
+    elif os.access(fh + ".gz", 4):
+        suffix = ".gz"
+        compression = "gzip"
+    elif os.access(fh + ".parquet", 4):
+        suffix = ".parquet"
+        compression = "parquet"
+    elif os.access(fh + ".feather", 4):
+        suffix = ".feather"
+        compression = "feather"
     elif os.access(fh, 4):
-        suffix = ''
+        suffix = ""
         compression = None
     else:
-        raise IOError('Could not open {F}[./gz/bz2/parquet/feather]'.format(F=fh))
+        raise OSError(f"Could not open {fh}[./gz/bz2/parquet/feather]")
     # -
     return suffix, compression
 
 
-def _read_ref_ld(ld_file):
-    suffix = '.l2.ldscore'
-    file = ld_file
-    first_fh = f'{file}1{suffix}'
-    s, compression = which_compression(first_fh)
-    #
-    ldscore_array = []
-    print(f'Reading ld score annotations from {file}[1-22]{suffix}.{compression}')
-
-    for chr in range(1, 23):
-        file_chr = f'{file}{chr}{suffix}{s}'
-        #
-        if compression == 'parquet':
-            x = pd.read_parquet(file_chr)
-        elif compression == 'feather':
-            x = pd.read_feather(file_chr)
-        else:
-            x = pd.read_csv(file_chr, compression=compression, sep='\t')
-
-        x = x.sort_values(by=['CHR', 'BP'])  # SEs will be wrong unless sorted
-
-        columns_to_drop = ['MAF', 'CM', 'Gene', 'TSS', 'CHR', 'BP']
-        columns_to_drop = [col for col in columns_to_drop if col in x.columns]
-        x = x.drop(columns_to_drop, axis=1)
-
-        ldscore_array.append(x)
-    #
-    ref_ld = pd.concat(ldscore_array, axis=0)
-    return ref_ld
-
-
 def _read_ref_ld_v2(ld_file):
-    suffix = '.l2.ldscore'
+    suffix = ".l2.ldscore"
     file = ld_file
-    first_fh = f'{file}1{suffix}'
+    first_fh = f"{file}1{suffix}"
     s, compression = which_compression(first_fh)
-    print(f'Reading ld score annotations from {file}[1-22]{suffix}.{compression}')
+    print(f"Reading ld score annotations from {file}[1-22]{suffix}.{compression}")
     ref_ld = pd.concat(
-        [pd.read_feather(f'{file}{chr}{suffix}{s}') for chr in range(1, 23)], axis=0
+        [pd.read_feather(f"{file}{chr}{suffix}{s}") for chr in range(1, 23)], axis=0
     )
     # set first column as index
-    ref_ld.rename(columns={'index': 'SNP'}, inplace=True)
-    ref_ld.set_index('SNP', inplace=True)
+    ref_ld.rename(columns={"index": "SNP"}, inplace=True)
+    ref_ld.set_index("SNP", inplace=True)
     return ref_ld
 
+
 def _read_M_v2(ld_file, n_annot, not_M_5_50):
-    suffix = '.l2.M'
+    suffix = ".l2.M"
     if not not_M_5_50:
-        suffix += '_5_50'
-    M_annot= np.array(
+        suffix += "_5_50"
+    M_annot = np.array(
         [
-            np.loadtxt(f'{ld_file}{chr}{suffix}', )
-         for chr in range(1, 23)]
-
+            np.loadtxt(
+                f"{ld_file}{chr}{suffix}",
+            )
+            for chr in range(1, 23)
+        ]
     )
     assert M_annot.shape == (22, n_annot)
     return M_annot.sum(axis=0).reshape((1, n_annot))
-# Fun for reading M annotations 
+
+
+# Fun for reading M annotations
 def _read_M(ld_file, n_annot, not_M_5_50):
-    '''
+    """
     Read M (--M, --M-file, etc).
-    '''
+    """
     M_annot = M(ld_file, common=(not not_M_5_50))
 
     try:
         M_annot = np.array(M_annot).reshape((1, n_annot))
     except ValueError as e:
-        raise ValueError('# terms in --M must match # of LD Scores in --ref-ld.\n' + str(e.args))
+        raise ValueError(
+            "# terms in --M must match # of LD Scores in --ref-ld.\n" + str(e.args)
+        ) from e
     return M_annot
 
 
 def M(fh, common=False):
-    '''
+    """
     Parses .l{N}.M files, split across num chromosomes.
-    '''
-    suffix = '.l2.M'
+    """
+    suffix = ".l2.M"
     if common:
-        suffix += '_5_50'
+        suffix += "_5_50"
     # -
     M_array = []
     for i in range(1, 23):
-        M_current = pd.read_csv(f'{fh}{i}' + suffix, header=None)
+        M_current = pd.read_csv(f"{fh}{i}" + suffix, header=None)
         M_array.append(M_current)
 
     M_array = pd.concat(M_array, axis=1).sum(axis=1)
@@ -178,117 +154,48 @@ def M(fh, common=False):
     return np.array(M_array).reshape((1, len(M_array)))
 
 
-def _check_variance(M_annot, ref_ld):
-    '''
-    Remove zero-variance LD Scores.
-    '''
-    ii = ref_ld.iloc[:, 1:].var() == 0  # NB there is a SNP column here
-    if ii.all():
-        raise ValueError('All LD Scores have zero variance.')
-    else:
-        print('Removing partitioned LD Scores with zero variance.')
-        ii_snp = np.array([True] + list(~ii))
-        ii_m = np.array(~ii)
-        ref_ld = ref_ld.iloc[:, ii_snp]
-        M_annot = M_annot[:, ii_m]
-    # -
-    return M_annot, ref_ld, ii
 def _check_variance_v2(M_annot, ref_ld):
     ii = ref_ld.var() == 0
     if ii.all():
-        raise ValueError('All LD Scores have zero variance.')
+        raise ValueError("All LD Scores have zero variance.")
     elif not ii.any():
-        print('No partitioned LD Scores have zero variance.')
+        print("No partitioned LD Scores have zero variance.")
     else:
-        ii_snp= ii_m = np.array(~ii)
-        print(f'Removing {sum(ii)} partitioned LD Scores with zero variance.')
+        ii_snp = ii_m = np.array(~ii)
+        print(f"Removing {sum(ii)} partitioned LD Scores with zero variance.")
         ref_ld = ref_ld.iloc[:, ii_snp]
         M_annot = M_annot[:, ii_m]
     return M_annot, ref_ld
 
 
-# Fun for reading regression weights
-def which_compression(fh):
-    '''
-    Given a file prefix, figure out what sort of compression to use.
-    '''
-    if os.access(fh + '.bz2', 4):
-        suffix = '.bz2'
-        compression = 'bz2'
-    elif os.access(fh + '.gz', 4):
-        suffix = '.gz'
-        compression = 'gzip'
-    elif os.access(fh + '.parquet', 4):
-        suffix = '.parquet'
-        compression = 'parquet'
-    elif os.access(fh + '.feather', 4):
-        suffix = '.feather'
-        compression = 'feather'
-    elif os.access(fh, 4):
-        suffix = ''
-        compression = None
-    else:
-        raise IOError('Could not open {F}[./gz/bz2/parquet/feather]'.format(F=fh))
-    # -
-    return suffix, compression
-
-
 def _read_w_ld(w_file):
-    suffix = '.l2.ldscore'
+    suffix = ".l2.ldscore"
     file = w_file
-    first_fh = f'{file}1{suffix}'
+    first_fh = f"{file}1{suffix}"
     s, compression = which_compression(first_fh)
     #
     w_array = []
-    print(f'Reading ld score annotations from {file}[1-22]{suffix}.{compression}')
+    print(f"Reading ld score annotations from {file}[1-22]{suffix}.{compression}")
 
     for chr in range(1, 23):
-        file_chr = f'{file}{chr}{suffix}{s}'
+        file_chr = f"{file}{chr}{suffix}{s}"
         #
-        if compression == 'parquet':
+        if compression == "parquet":
             x = pd.read_parquet(file_chr)
-        elif compression == 'feather':
+        elif compression == "feather":
             x = pd.read_feather(file_chr)
         else:
-            x = pd.read_csv(file_chr, compression=compression, sep='\t')
+            x = pd.read_csv(file_chr, compression=compression, sep="\t")
 
-        x = x.sort_values(by=['CHR', 'BP'])
+        x = x.sort_values(by=["CHR", "BP"])
 
-        columns_to_drop = ['MAF', 'CM', 'Gene', 'TSS', 'CHR', 'BP']
+        columns_to_drop = ["MAF", "CM", "Gene", "TSS", "CHR", "BP"]
         columns_to_drop = [col for col in columns_to_drop if col in x.columns]
         x = x.drop(columns_to_drop, axis=1)
 
         w_array.append(x)
     #
     w_ld = pd.concat(w_array, axis=0)
-    w_ld.columns = ['SNP', 'LD_weights']
+    w_ld.columns = ["SNP", "LD_weights"]
 
     return w_ld
-
-
-# Fun for merging
-def _merge_and_log(ld, sumstats, noun):
-    '''
-    Wrap smart merge with log messages about # of SNPs.
-    '''
-    sumstats = smart_merge(ld, sumstats)
-    msg = 'After merging with {F}, {N} SNPs remain.'
-    if len(sumstats) == 0:
-        raise ValueError(msg.format(N=len(sumstats), F=noun))
-    else:
-        print(msg.format(N=len(sumstats), F=noun))
-    # -
-    return sumstats
-
-
-def smart_merge(x, y):
-    '''
-    Check if SNP columns are equal. If so, save time by using concat instead of merge.
-    '''
-    if len(x) == len(y) and (x.index == y.index).all() and (x.SNP == y.SNP).all():
-        x = x.reset_index(drop=True)
-        y = y.reset_index(drop=True).drop('SNP', 1)
-        out = pd.concat([x, y], axis=1)
-    else:
-        out = pd.merge(x, y, how='inner', on='SNP')
-    return out

gsMap/visualize.py

@@ -11,35 +11,41 @@ from gsMap.config import VisualizeConfig
 
 
 def load_ldsc(ldsc_input_file):
-    ldsc = pd.read_csv(ldsc_input_file, compression='gzip')
-    ldsc.spot = ldsc.spot.astype(str).replace('\.0', '', regex=True)
-    ldsc.index = ldsc.spot
-    ldsc['logp'] = -np.log10(ldsc.p)
+    ldsc = pd.read_csv(
+        ldsc_input_file,
+        compression="gzip",
+        dtype={"spot": str, "p": float},
+        index_col="spot",
+        usecols=["spot", "p"],
+    )
+    ldsc["logp"] = -np.log10(ldsc.p)
     return ldsc
 
 
 # %%
 def load_st_coord(adata, feature_series: pd.Series, annotation):
     spot_name = adata.obs_names.to_list()
-    assert 'spatial' in adata.obsm.keys(), 'spatial coordinates are not found in adata.obsm'
+    assert "spatial" in adata.obsm.keys(), "spatial coordinates are not found in adata.obsm"
 
     # to DataFrame
-    space_coord = adata.obsm['spatial']
+    space_coord = adata.obsm["spatial"]
     if isinstance(space_coord, np.ndarray):
-        space_coord = pd.DataFrame(space_coord, columns=['sx', 'sy'], index=spot_name)
+        space_coord = pd.DataFrame(space_coord, columns=["sx", "sy"], index=spot_name)
     else:
-        space_coord = pd.DataFrame(space_coord.values, columns=['sx', 'sy'], index=spot_name)
+        space_coord = pd.DataFrame(space_coord.values, columns=["sx", "sy"], index=spot_name)
 
     space_coord = space_coord[space_coord.index.isin(feature_series.index)]
     space_coord_concat = pd.concat([space_coord.loc[feature_series.index], feature_series], axis=1)
     space_coord_concat.head()
     if annotation is not None:
-        annotation = pd.Series(adata.obs[annotation].values, index=adata.obs_names, name='annotation')
+        annotation = pd.Series(
+            adata.obs[annotation].values, index=adata.obs_names, name="annotation"
+        )
         space_coord_concat = pd.concat([space_coord_concat, annotation], axis=1)
     return space_coord_concat
 
 
-def estimate_point_size_for_plot(coordinates, DEFAULT_PIXEL_WIDTH = 1000):
+def estimate_point_size_for_plot(coordinates, DEFAULT_PIXEL_WIDTH=1000):
     tree = KDTree(coordinates)
     distances, _ = tree.query(coordinates, k=2)
     avg_min_distance = np.mean(distances[:, 1])
@@ -55,34 +61,42 @@ def estimate_point_size_for_plot(coordinates, DEFAULT_PIXEL_WIDTH = 1000):
     return (pixel_width, pixel_height), point_size
 
 
-def draw_scatter(space_coord_concat, title=None, fig_style: Literal['dark', 'light'] = 'light',
-                 point_size: int = None, width=800, height=600, annotation=None, color_by='logp'):
+def draw_scatter(
+    space_coord_concat,
+    title=None,
+    fig_style: Literal["dark", "light"] = "light",
+    point_size: int = None,
+    width=800,
+    height=600,
+    annotation=None,
+    color_by="logp",
+):
     # Set theme based on fig_style
-    if fig_style == 'dark':
+    if fig_style == "dark":
         px.defaults.template = "plotly_dark"
     else:
         px.defaults.template = "plotly_white"
 
     custom_color_scale = [
-        (1, '#d73027'),  # Red
-        (7 / 8, '#f46d43'),  # Red-Orange
-        (6 / 8, '#fdae61'),  # Orange
-        (5 / 8, '#fee090'),  # Light Orange
-        (4 / 8, '#e0f3f8'),  # Light Blue
-        (3 / 8, '#abd9e9'),  # Sky Blue
-        (2 / 8, '#74add1'),  # Medium Blue
-        (1 / 8, '#4575b4'),  # Dark Blue
-        (0, '#313695')  # Deep Blue
+        (1, "#d73027"),  # Red
+        (7 / 8, "#f46d43"),  # Red-Orange
+        (6 / 8, "#fdae61"),  # Orange
+        (5 / 8, "#fee090"),  # Light Orange
+        (4 / 8, "#e0f3f8"),  # Light Blue
+        (3 / 8, "#abd9e9"),  # Sky Blue
+        (2 / 8, "#74add1"),  # Medium Blue
+        (1 / 8, "#4575b4"),  # Dark Blue
+        (0, "#313695"),  # Deep Blue
     ]
     custom_color_scale.reverse()
 
     # Create the scatter plot
     fig = px.scatter(
         space_coord_concat,
-        x='sx',
-        y='sy',
+        x="sx",
+        y="sy",
         color=color_by,
-        symbol='annotation' if annotation is not None else None,
+        symbol="annotation" if annotation is not None else None,
         title=title,
         color_continuous_scale=custom_color_scale,
         range_color=[0, max(space_coord_concat[color_by])],
@@ -90,7 +104,7 @@ def draw_scatter(space_coord_concat, title=None, fig_style: Literal['dark', 'lig
 
     # Update marker size if specified
     if point_size is not None:
-        fig.update_traces(marker=dict(size=point_size, symbol='circle'))
+        fig.update_traces(marker=dict(size=point_size, symbol="circle"))
 
     # Update layout for figure size
     fig.update_layout(
@@ -108,45 +122,45 @@ def draw_scatter(space_coord_concat, title=None, fig_style: Literal['dark', 'lig
             x=1.0,
             font=dict(
                 size=10,
-            )
+            ),
         )
     )
 
     # Update colorbar to be at the bottom and horizontal
     fig.update_layout(
         coloraxis_colorbar=dict(
-            orientation='h',  # Make the colorbar horizontal
+            orientation="h",  # Make the colorbar horizontal
             x=0.5,  # Center the colorbar horizontally
             y=-0.0,  # Position below the plot
-            xanchor='center',  # Anchor the colorbar at the center
-            yanchor='top',  # Anchor the colorbar at the top to keep it just below the plot
+            xanchor="center",  # Anchor the colorbar at the center
+            yanchor="top",  # Anchor the colorbar at the top to keep it just below the plot
             len=0.75,  # Length of the colorbar relative to the plot width
             title=dict(
-                text='-log10(p)' if color_by == 'logp' else color_by,  # Colorbar title
-                side='top'  # Place the title at the top of the colorbar
-            )
+                text="-log10(p)" if color_by == "logp" else color_by,  # Colorbar title
+                side="top",  # Place the title at the top of the colorbar
+            ),
         )
     )
     # Remove gridlines, axis labels, and ticks
     fig.update_xaxes(
-        showgrid=False,   # Hide x-axis gridlines
-        zeroline=False,   # Hide x-axis zero line
+        showgrid=False,  # Hide x-axis gridlines
+        zeroline=False,  # Hide x-axis zero line
         showticklabels=False,  # Hide x-axis tick labels
-        title=None,       # Remove x-axis title
-        scaleanchor='y',  # Link the x-axis scale to the y-axis scale
+        title=None,  # Remove x-axis title
+        scaleanchor="y",  # Link the x-axis scale to the y-axis scale
     )
 
     fig.update_yaxes(
-        showgrid=False,   # Hide y-axis gridlines
-        zeroline=False,   # Hide y-axis zero line
+        showgrid=False,  # Hide y-axis gridlines
+        zeroline=False,  # Hide y-axis zero line
         showticklabels=False,  # Hide y-axis tick labels
-        title=None        # Remove y-axis title
+        title=None,  # Remove y-axis title
     )
 
     # Adjust margins to ensure no clipping and equal axis ratio
     fig.update_layout(
         margin=dict(l=0, r=0, t=20, b=10),  # Adjust margins to prevent clipping
-        height=width  # Ensure the figure height matches the width for equal axis ratio
+        height=width,  # Ensure the figure height matches the width for equal axis ratio
     )
 
     # Adjust the title location and font size
@@ -154,46 +168,50 @@ def draw_scatter(space_coord_concat, title=None, fig_style: Literal['dark', 'lig
         title=dict(
             y=0.98,
             x=0.5,  # Center the title horizontally
-            xanchor='center',  # Anchor the title at the center
-            yanchor='top',  # Anchor the title at the top
+            xanchor="center",  # Anchor the title at the center
+            yanchor="top",  # Anchor the title at the top
             font=dict(
                 size=20  # Increase the title font size
-            )
-        ))
+            ),
+        )
+    )
 
     return fig
 
 
-
 def run_Visualize(config: VisualizeConfig):
-    print(f'------Loading LDSC results of {config.ldsc_save_dir}...')
-    ldsc = load_ldsc(ldsc_input_file=Path(config.ldsc_save_dir) / f'{config.sample_name}_{config.trait_name}.csv.gz')
+    print(f"------Loading LDSC results of {config.ldsc_save_dir}...")
+    ldsc = load_ldsc(
+        ldsc_input_file=Path(config.ldsc_save_dir)
+        / f"{config.sample_name}_{config.trait_name}.csv.gz"
+    )
 
-    print(f'------Loading ST data of {config.sample_name}...')
-    adata = sc.read_h5ad(f'{config.hdf5_with_latent_path}')
+    print(f"------Loading ST data of {config.sample_name}...")
+    adata = sc.read_h5ad(f"{config.hdf5_with_latent_path}")
 
     space_coord_concat = load_st_coord(adata, ldsc, annotation=config.annotation)
-    fig = draw_scatter(space_coord_concat,
-                       title=config.fig_title,
-                       fig_style=config.fig_style,
-                       point_size=config.point_size,
-                       width=config.fig_width,
-                       height=config.fig_height,
-                       annotation=config.annotation
-                       )
-
+    fig = draw_scatter(
+        space_coord_concat,
+        title=config.fig_title,
+        fig_style=config.fig_style,
+        point_size=config.point_size,
+        width=config.fig_width,
+        height=config.fig_height,
+        annotation=config.annotation,
+    )
 
     # Visualization
     output_dir = Path(config.output_dir)
     output_dir.mkdir(parents=True, exist_ok=True, mode=0o755)
-    output_file_html = output_dir / f'{config.sample_name}_{config.trait_name}.html'
-    output_file_pdf = output_dir / f'{config.sample_name}_{config.trait_name}.pdf'
-    output_file_csv = output_dir / f'{config.sample_name}_{config.trait_name}.csv'
+    output_file_html = output_dir / f"{config.sample_name}_{config.trait_name}.html"
+    output_file_pdf = output_dir / f"{config.sample_name}_{config.trait_name}.pdf"
+    output_file_csv = output_dir / f"{config.sample_name}_{config.trait_name}.csv"
 
     fig.write_html(str(output_file_html))
     fig.write_image(str(output_file_pdf))
     space_coord_concat.to_csv(str(output_file_csv))
 
     print(
-        f'------The visualization result is saved in a html file: {output_file_html} which can interactively viewed in a web browser and a pdf file: {output_file_pdf}.')
-    print(f'------The visualization data is saved in a csv file: {output_file_csv}.')
+        f"------The visualization result is saved in a html file: {output_file_html} which can interactively viewed in a web browser and a pdf file: {output_file_pdf}."
+    )
+    print(f"------The visualization data is saved in a csv file: {output_file_csv}.")