Skip to content

Commit

Permalink
pygmt.filter1d: Improve performance by getting rid of temporary files
Browse files Browse the repository at this point in the history
  • Loading branch information
seisman committed Mar 8, 2024
1 parent 9640c26 commit 8857ce6
Showing 1 changed file with 30 additions and 37 deletions.
67 changes: 30 additions & 37 deletions pygmt/src/filter1d.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,11 @@
filter1d - Time domain filtering of 1-D data tables
"""

import pandas as pd
from typing import Literal

from pygmt.clib import Session
from pygmt.exceptions import GMTInvalidInput
from pygmt.helpers import (
GMTTempFile,
build_arg_string,
fmt_docstring,
use_alias,
Expand All @@ -20,7 +20,12 @@
F="filter_type",
N="time_col",
)
def filter1d(data, output_type="pandas", outfile=None, **kwargs):
def filter1d(
data,
output_type: Literal["pandas", "numpy", "file"] = "pandas",
outfile: str | None = None,
**kwargs,
):
r"""
Time domain filtering of 1-D data tables.
Expand All @@ -38,6 +43,15 @@ def filter1d(data, output_type="pandas", outfile=None, **kwargs):
Parameters
----------
output_type
Desired output type of the result data.
- ``pandas`` will return a :class:`pandas.DataFrame` object.
- ``numpy`` will return a :class:`numpy.ndarray` object.
- ``file`` will save the result to the file given by the ``outfile`` parameter.
outfile
The file name for saving the result. If specified, ``output_type`` will be
forced to be ``"file"``.
filter_type : str
**type**\ *width*\ [**+h**].
Set the filter **type**. Choose among convolution and non-convolution
Expand Down Expand Up @@ -91,48 +105,27 @@ def filter1d(data, output_type="pandas", outfile=None, **kwargs):
left-most column is 0, while the right-most is (*n_cols* - 1)
[Default is ``0``].
output_type : str
Determine the format the xyz data will be returned in [Default is
``pandas``]:
- ``numpy`` - :class:`numpy.ndarray`
- ``pandas``- :class:`pandas.DataFrame`
- ``file`` - ASCII file (requires ``outfile``)
outfile : str
The file name for the output ASCII file.
Returns
-------
ret : pandas.DataFrame or numpy.ndarray or None
Return type depends on ``outfile`` and ``output_type``:
- None if ``outfile`` is set (output will be stored in file set by
``outfile``)
- :class:`pandas.DataFrame` or :class:`numpy.ndarray` if ``outfile`` is
not set (depends on ``output_type`` [Default is
:class:`pandas.DataFrame`])
- None if ``outfile`` is set (output will be stored in file set by ``outfile``)
- :class:`pandas.DataFrame` or :class:`numpy.ndarray` if ``outfile`` is not set
(depends on ``output_type``)
"""
if kwargs.get("F") is None:
raise GMTInvalidInput("Pass a required argument to 'filter_type'.")

output_type = validate_output_table_type(output_type, outfile=outfile)

with GMTTempFile() as tmpfile:
with Session() as lib:
with lib.virtualfile_in(check_kind="vector", data=data) as vintbl:
if outfile is None:
outfile = tmpfile.name
lib.call_module(
module="filter1d",
args=build_arg_string(kwargs, infile=vintbl, outfile=outfile),
)

# 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", header=None, comment=">")
elif outfile != tmpfile.name: # return None if outfile set, output in outfile
result = None

if output_type == "numpy":
result = result.to_numpy()
return result
with Session() as lib:
with (
lib.virtualfile_in(check_kind="vector", data=data) as vintbl,
lib.virtualfile_out(kind="dataset", fname=outfile) as vouttbl,
):
lib.call_module(
module="filter1d",
args=build_arg_string(kwargs, infile=vintbl, outfile=vouttbl),
)
return lib.return_dataset(output_type=output_type, vfile=vouttbl)

0 comments on commit 8857ce6

Please sign in to comment.