CLI¶

def parameter_check(
    args: argparse.Namespace,
) -> tuple[int, pl.DataFrame, list[str], pl.DataFrame, list[str], Callable]:
    """The function to process raw phenotype, genotype, covariates data across ancestries
        for individual-level data fine-mapping.

    Args:
        args: The command line parameter input.

    Returns:
        `Tuple[int, pl.DataFrame, List[str], Callable]`:
            A tuple of
                - an integer to indicate how many ancestries,
                - a DataFrame that contains ancestry index (can be none),
                - a list that contains subject ID that fine-mapping performs on.
                - a DataFrame that contains prior probability for each SNP to be causal.
                - a list of genotype data paths (`List[str]`),
                - genotype read-in function (`Callable`).

    """
    if args.pheno is None:
        raise ValueError("No phenotype file specified. Specify --summary if summary-level fine-mapping is wanted.")

    if args.ancestry_index is not None:
        log.logger.debug("Read in ancestry index file.")
        ancestry_index = pl.read_csv(args.ancestry_index[0], has_header=False, separator="\t")
        old_pt = ancestry_index.shape[0]
        ancestry_index = ancestry_index.unique(maintain_order=True)

        if old_pt != ancestry_index.shape[0]:
            log.logger.debug(f"Index file has {old_pt - ancestry_index.shape[0]} duplicated subjects.")

        if ancestry_index.get_column("column_1").is_duplicated().sum() != 0:
            raise ValueError(
                "The ancestry index file contains subjects with multiple ancestry index. Check the source."
            )

        ancestry_ids = np.sort(ancestry_index.get_column("column_2").unique().to_numpy())
        n_pop = len(ancestry_ids)
        index_check = np.array_equal(
            ancestry_ids,
            np.arange(1, n_pop + 1),
        )

        if not index_check:
            raise ValueError(
                "The ancestry index doesn't start from 1 continuously to the total number of ancestry."
                + f" Check {args.ancestry_index}."
            )

        if len(args.pheno) > 1:
            raise ValueError(
                "Multiple phenotype files are detected. Expectation is one when --ancestry-index is specified."
            )

        log.logger.debug(
            "Detect ancestry index file, so it expects to have single phenotype, genotype, and covariates files"
        )

    else:
        ancestry_index = _empty_frame()
        n_pop = len(args.pheno)

    name_ancestry = "ancestry" if n_pop == 1 else "ancestries"

    log.logger.info(f"Detect phenotypes for {args.trait} for {n_pop} {name_ancestry}.")

    n_geno = (
        int(args.plink is not None)
        + int(args.plink2 is not None)
        + int(args.plink2_dosage is not None)
        + int(args.vcf is not None)
        + int(args.bgen is not None)
    )

    if n_geno > 1:
        log.logger.info(
            f"Detect {n_geno} genotypes, will only use one type of genotypes in the order of "
            + "'plink, plink2, plink2-dosage, vcf, and bgen'"
        )

    # decide genotype data
    if args.plink is not None:
        if args.ancestry_index is not None:
            if len(args.plink) > 1:
                raise ValueError(
                    "Multiple plink files are detected. Expectation is one when --ancestry-index is specified."
                )
        else:
            if len(args.plink) != n_pop:
                raise ValueError(
                    "The numbers of ancestries in plink geno and pheno data does not match. Check the source."
                )

        log.logger.info(f"Detect genotype data in plink format for {n_pop} {name_ancestry}.")
        geno_path = args.plink
        geno_func = io.read_triplet
    elif args.plink2 is not None:
        if args.ancestry_index is not None:
            if len(args.plink2) > 1:
                raise ValueError(
                    "Multiple plink2 files are detected. Expectation is one when --ancestry-index is specified."
                )
        else:
            if len(args.plink2) != n_pop:
                raise ValueError(
                    "The numbers of ancestries in plink2 geno and pheno data does not match. Check the source."
                )

        log.logger.info(f"Detect genotype data in plink2 format for {n_pop} {name_ancestry}.")
        geno_path = args.plink2
        geno_func = io.read_pfile
    elif args.plink2_dosage is not None:
        if args.ancestry_index is not None:
            if len(args.plink2_dosage) > 1:
                raise ValueError(
                    "Multiple plink2 dosage files are detected. Expectation is one when --ancestry-index is specified."
                )
        else:
            if len(args.plink2_dosage) != n_pop:
                raise ValueError(
                    "The numbers of ancestries in plink2 dosage geno and pheno data does not match. Check the source."
                )

        log.logger.info(f"Detect genotype dosage data in plink2 format for {n_pop} {name_ancestry}.")
        geno_path = args.plink2_dosage
        geno_func = partial(io.read_pfile, dosage=True)
    elif args.vcf is not None:
        if args.ancestry_index is not None:
            if len(args.vcf) > 1:
                raise ValueError(
                    "Multiple vcf files are detected. Expectation is one when --ancestry-index is specified."
                )
        else:
            if len(args.vcf) != n_pop:
                raise ValueError(
                    "The numbers of ancestries in vcf geno and pheno data does not match. Check the source."
                )
        log.logger.info(f"Detect genotype data in vcf format for {n_pop} {name_ancestry}.")
        geno_path = args.vcf
        geno_func = io.read_vcf
    elif args.bgen is not None:
        if args.ancestry_index is not None:
            if len(args.bgen) > 1:
                raise ValueError(
                    "Multiple bgen files are detected. Expectation is one when --ancestry-index is specified."
                )
        else:
            if len(args.bgen) != n_pop:
                raise ValueError(
                    "The numbers of ancestries in bgen geno and pheno data does not match. Check the source."
                )

        log.logger.info(f"Detect genotype data in bgen format for {n_pop} {name_ancestry}.")
        geno_path = args.bgen
        geno_func = io.read_bgen
    else:
        raise ValueError(
            "No genotype data specified in either plink, plink2, plink2-dosage, vcf, or bgen format. Check the source."
        )

    if args.covar is not None:
        if args.ancestry_index is not None:
            if len(args.covar) > 1:
                raise ValueError(
                    "Multiple covariates files are detected. Expectation is one when --ancestry-index is specified."
                )
        else:
            if len(args.covar) != n_pop:
                raise ValueError("The number of covariates data does not match geno data.")
        log.logger.info("Detect covariates data.")
    else:
        log.logger.info("No covariates detected for this analysis.")

    keep_subject = []
    if args.keep is not None:
        log.logger.info("Detect keep subject file. The inference only performs on the subjects listed in the file.")
        df_keep = pl.read_csv(args.keep[0], has_header=False, separator="\t").select(["column_1"])
        if df_keep.shape[0] == 0:
            raise ValueError("No subjects are listed in the keep subject file. Check the source.")
        old_pt = df_keep.shape[0]
        df_keep = df_keep.unique(maintain_order=True)

        if old_pt != df_keep.shape[0]:
            log.logger.debug(f"The keep subject file has {old_pt - df_keep.shape[0]} duplicated subjects.")
        keep_subject = df_keep.get_column("column_1").to_list()

    if args.pi != "uniform":
        log.logger.info("Detect file that contains prior weights for each SNP to be causal.")
        pi = pl.read_csv(args.pi, has_header=False, separator="\t")
        if pi.shape[0] == 0:
            raise ValueError("No prior weights are listed in the prior file. Check the source.")

        if pi.shape[1] < 2:
            raise ValueError(
                "The prior file has less than 2 columns. It has to be at least two columns."
                + " The first column is the SNP ID, and the second column the prior probability."
            )

        if pi.shape[1] > 2:
            log.logger.debug("The prior file has more than 2 columns. Will only use the first two columns.")

        pi = pi.select(pi.columns[0:2]).rename({pi.columns[0]: "snp", pi.columns[1]: "pi"})
    else:
        pi = _empty_frame()

    if args.seed <= 0:
        raise ValueError(
            "The seed specified for randomization must be greater than 0. Choose a positive integer using --seed."
        )

    if args.cv:
        if args.cv_num <= 1:
            raise ValueError(
                "The number of folds in cross validation must be greater than 1." + " Update with --cv-num.",
            )

    if args.maf <= 0 or args.maf > 0.5:
        raise ValueError(
            "The minor allele frequency (MAF) has to be between 0 (exclusive) and 0.5 (inclusive)."
            + " Choose a valid frequency using --maf."
        )

    if (args.meta or args.mega) and n_pop == 1:
        log.logger.debug(
            "The number of ancestry is 1, but --meta or --mega is specified. Will skip meta or mega SuSiE."
        )

    if args.chrom is None and args.start is None and args.end is None:
        log.logger.debug("No region is specified. Will use all SNPs available in the data.")

    elif args.chrom is not None and args.start is not None and args.end is not None:
        if args.start <= 0:
            raise ValueError("The start position for the region must be greater than 0. Update with --start.")

        if args.end <= 0:
            raise ValueError("The end position for the region must be greater than 0. Update with --end.")

        if args.end <= args.start:
            raise ValueError(
                "The end position for the region must be greater than --start. Update with" + " --start or --end."
            )

        log.logger.info(
            f"Detect region (chrom{args.chrom}:{args.start}:{args.end}) to be fine-mapped."
            + " Will only use SNPs within the region."
        )
    else:
        raise ValueError(
            "The region is not specified correctly. Please provide --chrom, --start, and --end together,"
            + " or omit all of them."
        )

    log.logger.debug("Finish parameter check for individual-level fine-mapping.")

    return n_pop, ancestry_index, keep_subject, pi, geno_path, geno_func

def process_raw(
    rawData: list[io.RawData],
    keep_subject: list[str],
    pi: pl.DataFrame,
    keep_ambiguous: bool,
    maf: float,
    rint: bool,
    no_regress: bool,
    mega: bool,
    cv: bool,
    cv_num: int,
    seed: int,
    chrom: utils.IntOrNone,
    start: utils.IntOrNone,
    end: utils.IntOrNone,
) -> tuple[
    pl.DataFrame,
    io.CleanData,
    io.CleanData | None,
    list[io.CVData] | None,
]:
    """The function to process raw phenotype, genotype, covariates data across ancestries.

    Args:
        rawData: Raw data for phenotypes, genotypes, covariates across ancestries.
        keep_subject: The DataFrame that contains subject ID that fine-mapping performs on.
        pi: The DataFrame that contains prior weights for each SNP to be causal.
        keep_ambiguous: The indicator whether to keep ambiguous SNPs.
        maf: The minor allele frequency threshold to filter the genotypes.
        rint: The indicator whether to perform rank inverse normalization on each phenotype data.
        no_regress: The indicator whether to regress genotypes on covariates.
        mega: The indicator whether to prepare datasets for mega SuShiE.
        cv: The indicator whether to prepare datasets for cross-validation.
        cv_num: The number for $X$-fold cross-validation.
        seed: The random seed for row-wise shuffling the datasets for cross validation.
        chrom: The chromosome to filter SNPs.
        start: The start position to filter SNPs.
        end: The end position to filter SNPs.


    Returns:
        `Tuple[pl.DataFrame, io.CleanData, Optional[io.CleanData], Optional[List[io.CVData]]]`:
        A tuple of
            - SNP information (`pl.DataFrame`),
            - dataset for running SuShiE (`io.CleanData`),
            - dataset for mega SuShiE (`Optional[io.CleanData]`),
            - dataset for cross-validation (`Optional[List[io.CVData]]`).

    """

    n_pop = len(rawData)

    for idx in range(n_pop):
        # keep subjects that are listed in the keep subject file
        if len(keep_subject) != 0:
            rawData[idx] = _keep_file_subjects(rawData[idx], keep_subject, idx)

        # remove NA/inf value for subjects across phenotype or covariates data
        rawData[idx] = _drop_na_subjects(rawData[idx], idx)

        # impute genotype data even though we suggest users to impute the genotypes beforehand
        rawData[idx] = _impute_geno(rawData[idx], idx)

        # remove SNPs that cannot pass MAF threshold
        rawData[idx] = _filter_maf(rawData[idx], maf, idx)

        # remove duplicates SNPs based on rsid even though we suggest users to do some QC on this
        rawData[idx] = _remove_dup_geno(rawData[idx], idx)

        # reset index and add index column to all dataset for future inter-ancestry or inter-dataset processing
        rawData[idx] = _reset_idx(rawData[idx], idx)

        # find common individuals across geno, pheno, and covar within an ancestry
        rawData[idx] = _filter_common_ind(rawData[idx], idx)

    # find common snps across ancestries
    log.logger.debug("Fine common SNPs across ancestries.")

    if n_pop > 1:
        snps = rawData[0].bim.join(
            rawData[1].bim,
            how="inner",
            on=["chrom", "snp"],
            maintain_order="left",
        )

        for idx in range(n_pop - 2):
            snps = snps.join(
                rawData[idx + 2].bim,
                how="inner",
                on=["chrom", "snp"],
                maintain_order="left",
            )
        if snps.shape[0] == 0:
            raise ValueError("Ancestries have no common SNPs. Check the source.")
        # report how many snps we removed due to independent SNPs
        for idx in range(n_pop):
            snps_num_diff = rawData[idx].bim.shape[0] - snps.shape[0]
            log.logger.debug(
                f"Ancestry{idx + 1} has {snps_num_diff} independent SNPs and {snps.shape[0]}"
                + " common SNPs. Inference only performs on common SNPs.",
            )
    else:
        snps = rawData[0].bim

    # remove biallelic SNPs dont match across ancestries
    # e.g., A/T for EUR but A/C for AFR
    if n_pop > 1:
        log.logger.debug("Remove SNPs that do not have same alleles across ancestries.")
        for idx in range(1, n_pop):
            _, _, remove_idx = _allele_check(
                _col_array(snps, "a0_1"),
                _col_array(snps, "a1_1"),
                _col_array(snps, f"a0_{idx + 1}"),
                _col_array(snps, f"a1_{idx + 1}"),
            )

            if len(remove_idx) != 0:
                snps = _drop_rows(snps, remove_idx)
                log.logger.debug(
                    f"Ancestry{idx + 1} has {len(remove_idx)} alleles that"
                    + "couldn't match to ancestry 1 and couldn't be flipped. Will remove these SNPs."
                )

            if snps.shape[0] == 0:
                raise ValueError(
                    f"Ancestry {idx + 1} has none of correct or flippable SNPs matching to ancestry 1."
                    + "Check the source.",
                )

    # remove ambiguous SNPs (i.e., A/T, T/A, C/G, G/C pairs) in genotype data
    if not keep_ambiguous:
        log.logger.debug("Remove ambiguous SNPs.")

        ambiguous_snps = ["AT", "TA", "CG", "GC"]
        if_ambig = snps.select((pl.col("a0_1") + pl.col("a1_1")).is_in(ambiguous_snps).alias("ambig")).get_column(
            "ambig"
        )
        del_num = if_ambig.sum()
        snps = snps.filter(if_ambig.not_())

        if snps.shape[0] == 0:
            raise ValueError("All SNPs are ambiguous in genotype data. Check the source.")

        if del_num != 0:
            log.logger.debug(f"Drop {del_num} ambiguous SNPs in genotype data.")

    log.logger.debug("Filter SNPs based on Chrom, Start, and End using coordinates of first ancestry.")
    snps = snps.with_columns(pl.col("chrom").cast(pl.Int64))

    if chrom is not None:
        old_num = snps.shape[0]
        snps = snps.filter(pl.col("chrom") == chrom)
        del_num = old_num - snps.shape[0]

        if snps.shape[0] == 0:
            raise ValueError(f"No SNPs remain after filtering on chromosome {chrom}.")

        if del_num != 0:
            log.logger.debug(f"Drop {del_num} SNPs that are not on chromosome {chrom}.")

        old_num = snps.shape[0]
        snps = snps.filter(pl.col("pos_1") >= start)
        del_num = old_num - snps.shape[0]

        if snps.shape[0] == 0:
            raise ValueError(f"No SNPs are located after position {start} on chromosome {chrom}.")

        if del_num != 0:
            log.logger.debug(f"Drop {del_num} SNPs that are located before position {start} on chromosome {chrom}.")

        old_num = snps.shape[0]
        snps = snps.filter(pl.col("pos_1") <= end)
        del_num = old_num - snps.shape[0]

        if snps.shape[0] == 0:
            raise ValueError(f"No SNPs are located before position {end} on chromosome {chrom}.")

        if del_num != 0:
            log.logger.debug(f"Drop {del_num} SNPs that are located after position {end} on chromosome {chrom}.")

    # find flipped reference alleles across ancestries
    flip_idx = []
    if n_pop > 1:
        log.logger.debug("Flip the alleles of subsequent ancestries to match those of the first ancestry.")
        for idx in range(1, n_pop):
            _, tmp_flip_idx, _ = _allele_check(
                _col_array(snps, "a0_1"),
                _col_array(snps, "a1_1"),
                _col_array(snps, f"a0_{idx + 1}"),
                _col_array(snps, f"a1_{idx + 1}"),
            )

            if len(tmp_flip_idx) != 0:
                log.logger.debug(
                    f"Ancestry{idx + 1} has {len(tmp_flip_idx)} flipped alleles from ancestry 1. Will flip these SNPs."
                )

            # save the index for future swapping
            flip_idx.append(tmp_flip_idx)

            # drop unused columns
            snps = snps.drop([f"a0_{idx + 1}", f"a1_{idx + 1}", f"pos_{idx + 1}"])

    # rename columns for better indexing in the future
    snps = snps.with_row_index("SNPIndex").rename({"a0_1": "a0", "a1_1": "a1", "pos_1": "pos"})

    if pi.shape[0] != 0:
        # append prior weights to the snps
        log.logger.debug("Process prior weight file.")
        snps = snps.join(pi, how="left", on="snp", maintain_order="left")
        snps = snps.with_columns(pl.col("pi").cast(pl.Float64))
        nan_count = snps["pi"].null_count() + snps["pi"].is_nan().sum()
        if nan_count > 0:
            log.logger.debug(
                f"{nan_count} SNP(s) have missing prior weights. Will replace them with the mean value of the rest."
            )
        # if the column pi has nan value, replace it with the mean value of the rest of the column
        snps = snps.with_columns(pl.when(pl.col("pi").is_nan()).then(None).otherwise(pl.col("pi")).alias("pi"))
        snps = snps.with_columns(pl.col("pi").fill_null(pl.col("pi").mean()))
        pi_array = snps.select("pi").to_jax(dtype=pl.Float64).reshape(-1)
    else:
        snps = snps.with_columns(pl.lit(1.0 / snps.height).alias("pi"))
        pi_array = None

    geno = []
    pheno = []
    covar = []
    total_ind = 0
    # filter on geno, pheno, and covar
    for idx in range(n_pop):
        log.logger.debug(f"Final process the rawdata for ancestry {idx + 1}.")
        _, tmp_fam, tmp_geno, tmp_pheno, tmp_covar = rawData[idx]

        # get common individual and snp id
        common_ind_id = _col_array(tmp_fam, f"famIDX_{idx + 1}")
        common_snp_id = _col_array(snps, f"bimIDX_{idx + 1}")
        snps = snps.drop([f"bimIDX_{idx + 1}"])

        # filter on individuals who have both geno, pheno, and covar (if applicable)
        # filter on shared snps across ancestries
        tmp_geno = tmp_geno[common_ind_id, :][:, common_snp_id]

        # flip genotypes for bed files starting second ancestry
        # flip index is the positional index based on snps data frame, so we have to subset genotype
        # data based on the common snps (i.e., snps data frame).
        if idx > 0 and len(flip_idx[idx - 1]) != 0:
            tmp_geno = tmp_geno.at[:, flip_idx[idx - 1]].set(2 - tmp_geno[:, flip_idx[idx - 1]])

        # swap pheno and covar rows order to match fam/bed file, and then select the
        # values for future fine-mapping
        common_pheno_id = _col_array(tmp_fam, f"phenoIDX_{idx + 1}")
        tmp_pheno = tmp_pheno.select("pheno").to_jax().reshape(-1)[common_pheno_id]
        total_ind += tmp_pheno.shape[0]
        geno.append(tmp_geno)

        if rint:
            tmp_pheno = utils.rint(tmp_pheno)

        pheno.append(tmp_pheno)

        if tmp_covar is not None:
            # select the common individual for covar
            common_covar_id = _col_array(tmp_fam, f"covarIDX_{idx + 1}")
            n_covar = tmp_covar.shape[1]
            tmp_covar = tmp_covar.gather(_as_row_indices(common_covar_id)).select(tmp_covar.columns[2:n_covar]).to_jax()
            covar.append(tmp_covar)

    if len(covar) == 0:
        data_covar = None
    else:
        data_covar = covar

    regular_data = io.CleanData(geno=geno, pheno=pheno, covar=data_covar, pi=pi_array)

    name_ancestry = "ancestry" if n_pop == 1 else "ancestries"

    log.logger.info(
        f"Prepare {geno[0].shape[1]} SNPs for {total_ind} individuals from {n_pop} {name_ancestry} after"
        + " data cleaning. Specify --verbose for details.",
    )

    mega_data = None
    cv_data = None
    # when doing mega or cross validation, we need to regress out covariates first
    if mega or cv:
        cv_geno = copy.deepcopy(geno)
        cv_pheno = copy.deepcopy(pheno)
        if data_covar is not None:
            for idx in range(n_pop):
                cv_geno[idx], cv_pheno[idx] = utils.regress_covar(geno[idx], pheno[idx], data_covar[idx], no_regress)

        if cv:
            cv_data = _prepare_cv(cv_geno, cv_pheno, cv_num, seed)

        # prepare mega dataset
        # it's possible that different ancestries have different number of covariates,
        # so we need to regress out first
        if mega:
            mega_geno = cv_geno[0]
            mega_pheno = cv_pheno[0]
            for idx in range(1, n_pop):
                mega_geno = jnp.append(mega_geno, cv_geno[idx], axis=0)
                mega_pheno = jnp.append(mega_pheno, cv_pheno[idx], axis=0)

            # because it row-binds the phenotype data for each ancestry, we want to rint again
            mega_pheno = utils.rint(mega_pheno)
            mega_data = io.CleanData(
                geno=[mega_geno],
                pheno=[mega_pheno],
                covar=None,
                pi=pi_array,
            )

    log.logger.debug("Finish preparing data for cross-validation and mega fine-mapping.")

    return snps, regular_data, mega_data, cv_data

def sushie_wrapper(
    data: io.CleanData,
    cv_data: list[io.CVData] | None,
    args: argparse.Namespace,
    snps: pl.DataFrame,
    meta: bool = False,
    mega: bool = False,
) -> None:
    """The wrapper function to run SuShiE in regular, meta, or mega.

    Args:
        data: The clean data for SuShiE inference.
        cv_data: The cross-validation dataset.
        args: The command line parameter input.
        snps: The SNP information.
        meta: The indicator whether to prepare datasets for meta SuShiE.
        mega: The indicator whether to prepare datasets for mega SuShiE.

    """

    n_pop = len(data.geno)

    if meta:
        output = f"{args.output}.meta"
        method_type = "meta"
    elif mega:
        output = f"{args.output}.mega"
        method_type = "mega"
    else:
        output = f"{args.output}.sushie"
        method_type = "sushie"

    resid_var = None if mega else args.resid_var
    effect_var = None if mega else args.effect_var
    rho = None if mega else args.rho

    # Inference pads ancestry arrays in-place, so heritability keeps a clean copy.
    heri_data = copy.deepcopy(data)

    single_pip_all: list[Array] = []
    single_pip_cs: list[Array] = []
    meta_pips: list[Array] | None = None
    result: list[infer.SushieResult] = []
    if meta:
        for idx in range(n_pop):
            resid_var = None if args.resid_var is None else [args.resid_var[idx]]
            effect_var = None if args.effect_var is None else [args.effect_var[idx]]
            covar = None if data.covar is None else [data.covar[idx]]

            log.logger.info(
                f"Start fine-mapping using SuSiE on ancestry {idx + 1} with {args.L} effects"
                + " because --meta is specified."
            )

            tmp_result = infer.infer_sushie(
                [data.geno[idx]],
                [data.pheno[idx]],
                covar,
                L=args.L,
                no_scale=args.no_scale,
                no_regress=args.no_regress,
                no_update=args.no_update,
                pi=data.pi,
                resid_var=resid_var,
                effect_var=effect_var,
                rho=None,
                max_iter=args.max_iter,
                min_tol=args.min_tol,
                threshold=args.threshold,
                purity=args.purity,
                purity_method=args.purity_method,
                max_select=args.max_select,
                min_snps=args.min_snps,
                no_reorder=args.no_reorder,
                seed=args.seed,
            )
            single_pip_all.append(tmp_result.pip_all[:, jnp.newaxis])
            single_pip_cs.append(tmp_result.pip_cs[:, jnp.newaxis])
            result.append(tmp_result)

        meta_pips = [
            utils.make_pip(jnp.concatenate(single_pip_all, axis=1).T),
            utils.make_pip(jnp.concatenate(single_pip_cs, axis=1).T),
        ]
    else:
        # normal sushie and mega sushie can use the same wrapper function
        if mega:
            log.logger.info(f"Start fine-mapping using Mega SuSiE with {args.L} effects because --mega is specified.")
        else:
            log.logger.info(f"Start fine-mapping using SuShiE with {args.L} effects.")

        tmp_result = infer.infer_sushie(
            data.geno,
            data.pheno,
            data.covar,
            L=args.L,
            no_scale=args.no_scale,
            no_regress=args.no_regress,
            no_update=args.no_update,
            pi=data.pi,
            resid_var=resid_var,
            effect_var=effect_var,
            rho=rho,
            max_iter=args.max_iter,
            min_tol=args.min_tol,
            threshold=args.threshold,
            purity=args.purity,
            purity_method=args.purity_method,
            max_select=args.max_select,
            min_snps=args.min_snps,
            no_reorder=args.no_reorder,
            seed=args.seed,
        )
        result.append(tmp_result)

    io.output_cs(result, meta_pips, snps, output, args.trait, args.compress, method_type)
    io.output_weights(result, meta_pips, snps, output, args.trait, args.compress, method_type)

    if args.numpy:
        log.logger.info("Save all the inference results in numpy file because --numpy is specified ")
        io.output_numpy(result, snps, output)

    if args.alphas:
        log.logger.info("Save all credible set results before pruning as --alphas is specified ")

        io.output_alphas(
            result,
            snps,
            output,
            args.trait,
            args.compress,
            method_type,
            args.purity,
        )

    if not (mega or meta):
        io.output_corr(result, output, args.trait, args.compress)

        if args.her:
            log.logger.info("Save heritability analysis results as --her is specified")
            io.output_her(heri_data, output, args.trait, args.compress)

        if args.cv:
            log.logger.info(f"Start {args.cv_num}-fold cross validation as --cv is specified ")
            if cv_data is None:
                raise RuntimeError("Cross-validation data were not prepared.")
            cv_res = _run_cv(args, cv_data, data.pi)
            sample_size = np.asarray(jnp.squeeze(tmp_result.sample_size)).astype(int).tolist()
            io.output_cv(cv_res, sample_size, output, args.trait, args.compress)

    return None

def run_finemap(args: argparse.Namespace) -> int:
    """The umbrella function to run SuShiE.

    Args:
        args: The command line parameter input.

    """

    try:
        if args.jax_precision == 64:
            config.update("jax_enable_x64", True)
            config.update("jax_default_matmul_precision", "highest")

        config.update("jax_platform_name", args.platform)
        if args.summary is True:
            log.logger.info("Start fine-mapping using SuShiE on summary-level data.")

            n_pop, pi, geno_path, geno_func, ld_file = parameter_check_ss(args)

            snps, ss_data = process_raw_ss(
                geno_path,
                geno_func,
                ld_file,
                pi,
                args,
            )

            normal_data = copy.deepcopy(ss_data)
            sushie_wrapper_ss(normal_data, args, snps, meta=False)

            # if only one ancestry, no need to run mega or meta
            if n_pop != 1:
                if args.meta:
                    meta_data = copy.deepcopy(ss_data)
                    sushie_wrapper_ss(meta_data, args, snps, meta=True)

        else:
            log.logger.info("Start fine-mapping using SuShiE on individual-level data.")

            (
                n_pop,
                ancestry_index,
                keep_subject,
                pi,
                geno_path,
                geno_func,
            ) = parameter_check(args)

            rawData = io.read_data(
                n_pop,
                ancestry_index,
                args.pheno,
                args.covar,
                geno_path,
                geno_func,
            )

            snps, regular_data, mega_data, cv_data = process_raw(
                rawData,
                keep_subject,
                pi,
                args.keep_ambiguous,
                args.maf,
                args.rint,
                args.no_regress,
                args.mega,
                args.cv,
                args.cv_num,
                args.seed,
                args.chrom,
                args.start,
                args.end,
            )

            normal_data = copy.deepcopy(regular_data)
            sushie_wrapper(normal_data, cv_data, args, snps, meta=False, mega=False)

            # if only one ancestry, no need to run mega or meta
            if n_pop != 1:
                if args.meta:
                    meta_data = copy.deepcopy(regular_data)
                    sushie_wrapper(meta_data, None, args, snps, meta=True, mega=False)

                if args.mega:
                    if mega_data is None:
                        raise RuntimeError("Mega SuSiE data were not prepared.")
                    sushie_wrapper(mega_data, None, args, snps, meta=False, mega=True)

    except Exception as err:
        import traceback

        print("".join(traceback.format_exception(type(err), err, err.__traceback__)))
        log.logger.error(err)
        return 1

    finally:
        log.logger.info(
            f"Fine-mapping finishes for {args.trait}, and thanks for using our software."
            + " For bug reporting, suggestions, and comments, please go to https://github.com/mancusolab/sushie.",
        )
    return 0