Source code for pygmt.src.blockmedian

"""
blockmedian - Block average (x,y,z) data tables by median estimation.
"""
import pandas as pd
from pygmt.clib import Session
from pygmt.exceptions import GMTInvalidInput
from pygmt.helpers import (
    GMTTempFile,
    build_arg_string,
    data_kind,
    dummy_context,
    fmt_docstring,
    kwargs_to_strings,
    use_alias,
)


@fmt_docstring
@use_alias(I="spacing", R="region", V="verbose")
@kwargs_to_strings(R="sequence")
def blockmedian(table, outfile=None, **kwargs):
    """
    Block average (x,y,z) data tables by median estimation.

    Reads arbitrarily located (x,y,z) triples [or optionally weighted
    quadruples (x,y,z,w)] from a table and writes to the output a median
    position and value for every non-empty block in a grid region defined by
    the region and spacing arguments.

    Full option list at :gmt-docs:`blockmedian.html`

    {aliases}

    Parameters
    ----------
    table : pandas.DataFrame or str
        Either a pandas dataframe with (x, y, z) or (longitude, latitude,
        elevation) values in the first three columns, or a file name to an
        ASCII data table.

    spacing : str
        ``'xinc[unit][+e|n][/yinc[unit][+e|n]]'``.
        x_inc [and optionally y_inc] is the grid spacing.

    region : str or list
        ``'xmin/xmax/ymin/ymax[+r][+uunit]'``.
        Specify the region of interest.

    outfile : str
        Required if 'table' is a file. The file name for the output ASCII file.

    {V}

    Returns
    -------
    output : pandas.DataFrame or None
        Return type depends on whether the outfile parameter is set:

        - pandas.DataFrame table with (x, y, z) columns if outfile is not set
        - None if outfile is set (filtered output will be stored in outfile)
    """
    kind = data_kind(table)
    with GMTTempFile(suffix=".csv") as tmpfile:
        with Session() as lib:
            if kind == "matrix":
                if not hasattr(table, "values"):
                    raise GMTInvalidInput(f"Unrecognized data type: {type(table)}")
                file_context = lib.virtualfile_from_matrix(table.values)
            elif kind == "file":
                if outfile is None:
                    raise GMTInvalidInput("Please pass in a str to 'outfile'")
                file_context = dummy_context(table)
            else:
                raise GMTInvalidInput(f"Unrecognized data type: {type(table)}")

            with file_context as infile:
                if outfile is None:
                    outfile = tmpfile.name
                arg_str = " ".join([infile, build_arg_string(kwargs), "->" + outfile])
                lib.call_module(module="blockmedian", args=arg_str)

        # Read temporary csv output to a pandas table
        if outfile == tmpfile.name:  # if user did not set outfile, return pd.DataFrame
            result = pd.read_csv(tmpfile.name, sep="\t", names=table.columns)
        elif outfile != tmpfile.name:  # return None if outfile set, output in outfile
            result = None

    return result