Skip to content

Commit 92dc36a

Browse files
committed
pygmt.filter1d: Improve performance by getting rid of temporary files
1 parent 68bfe57 commit 92dc36a

File tree

2 files changed

+30
-94
lines changed

2 files changed

+30
-94
lines changed

pygmt/src/filter1d.py

+30-39
Original file line numberDiff line numberDiff line change
@@ -2,11 +2,11 @@
22
filter1d - Time domain filtering of 1-D data tables
33
"""
44

5-
import pandas as pd
5+
from typing import Literal
6+
67
from pygmt.clib import Session
78
from pygmt.exceptions import GMTInvalidInput
89
from pygmt.helpers import (
9-
GMTTempFile,
1010
build_arg_string,
1111
fmt_docstring,
1212
use_alias,
@@ -20,7 +20,12 @@
2020
F="filter_type",
2121
N="time_col",
2222
)
23-
def filter1d(data, output_type="pandas", outfile=None, **kwargs):
23+
def filter1d(
24+
data,
25+
output_type: Literal["pandas", "numpy", "file"] = "pandas",
26+
outfile: str | None = None,
27+
**kwargs,
28+
):
2429
r"""
2530
Time domain filtering of 1-D data tables.
2631
@@ -38,6 +43,14 @@ def filter1d(data, output_type="pandas", outfile=None, **kwargs):
3843
3944
Parameters
4045
----------
46+
output_type
47+
Desired output type of the result data.
48+
- ``pandas`` will return a :class:`pandas.DataFrame` object.
49+
- ``numpy`` will return a :class:`numpy.ndarray` object.
50+
- ``file`` will save the result to the file given by the ``outfile`` parameter.
51+
outfile
52+
File name for saving the result data. Required if ``output_type`` is ``"file"``.
53+
If specified, ``output_type`` will be forced to be ``"file"``.
4154
filter_type : str
4255
**type**\ *width*\ [**+h**].
4356
Set the filter **type**. Choose among convolution and non-convolution
@@ -91,48 +104,26 @@ def filter1d(data, output_type="pandas", outfile=None, **kwargs):
91104
left-most column is 0, while the right-most is (*n_cols* - 1)
92105
[Default is ``0``].
93106
94-
output_type : str
95-
Determine the format the xyz data will be returned in [Default is
96-
``pandas``]:
97-
98-
- ``numpy`` - :class:`numpy.ndarray`
99-
- ``pandas``- :class:`pandas.DataFrame`
100-
- ``file`` - ASCII file (requires ``outfile``)
101-
outfile : str
102-
The file name for the output ASCII file.
103-
104107
Returns
105108
-------
106-
ret : pandas.DataFrame or numpy.ndarray or None
109+
ret
107110
Return type depends on ``outfile`` and ``output_type``:
108-
109-
- None if ``outfile`` is set (output will be stored in file set by
110-
``outfile``)
111-
- :class:`pandas.DataFrame` or :class:`numpy.ndarray` if ``outfile`` is
112-
not set (depends on ``output_type`` [Default is
113-
:class:`pandas.DataFrame`])
111+
- None if ``outfile`` is set (output will be stored in file set by ``outfile``)
112+
- :class:`pandas.DataFrame` or :class:`numpy.ndarray` if ``outfile`` is not set
113+
(depends on ``output_type``)
114114
"""
115115
if kwargs.get("F") is None:
116116
raise GMTInvalidInput("Pass a required argument to 'filter_type'.")
117117

118118
output_type = validate_output_table_type(output_type, outfile=outfile)
119119

120-
with GMTTempFile() as tmpfile:
121-
with Session() as lib:
122-
with lib.virtualfile_in(check_kind="vector", data=data) as vintbl:
123-
if outfile is None:
124-
outfile = tmpfile.name
125-
lib.call_module(
126-
module="filter1d",
127-
args=build_arg_string(kwargs, infile=vintbl, outfile=outfile),
128-
)
129-
130-
# Read temporary csv output to a pandas table
131-
if outfile == tmpfile.name: # if user did not set outfile, return pd.DataFrame
132-
result = pd.read_csv(tmpfile.name, sep="\t", header=None, comment=">")
133-
elif outfile != tmpfile.name: # return None if outfile set, output in outfile
134-
result = None
135-
136-
if output_type == "numpy":
137-
result = result.to_numpy()
138-
return result
120+
with Session() as lib:
121+
with (
122+
lib.virtualfile_in(check_kind="vector", data=data) as vintbl,
123+
lib.virtualfile_out(kind="dataset", fname=outfile) as vouttbl,
124+
):
125+
lib.call_module(
126+
module="filter1d",
127+
args=build_arg_string(kwargs, infile=vintbl, outfile=vouttbl),
128+
)
129+
return lib.virtualfile_to_dataset(output_type=output_type, vfname=vouttbl)

pygmt/tests/test_filter1d.py

-55
Original file line numberDiff line numberDiff line change
@@ -2,15 +2,11 @@
22
Test pygmt.filter1d.
33
"""
44

5-
from pathlib import Path
6-
75
import numpy as np
86
import pandas as pd
97
import pytest
108
from pygmt import filter1d
119
from pygmt.datasets import load_sample_data
12-
from pygmt.exceptions import GMTInvalidInput
13-
from pygmt.helpers import GMTTempFile
1410

1511

1612
@pytest.fixture(scope="module", name="data")
@@ -29,57 +25,6 @@ def test_filter1d_no_outfile(data):
2925
assert result.shape == (671, 2)
3026

3127

32-
def test_filter1d_file_output(data):
33-
"""
34-
Test that filter1d returns a file output when it is specified.
35-
"""
36-
with GMTTempFile(suffix=".txt") as tmpfile:
37-
result = filter1d(
38-
data=data, filter_type="g5", outfile=tmpfile.name, output_type="file"
39-
)
40-
assert result is None # return value is None
41-
assert Path(tmpfile.name).stat().st_size > 0 # check that outfile exists
42-
43-
44-
def test_filter1d_invalid_format(data):
45-
"""
46-
Test that filter1d fails with an incorrect format for output_type.
47-
"""
48-
with pytest.raises(GMTInvalidInput):
49-
filter1d(data=data, filter_type="g5", output_type="a")
50-
51-
52-
def test_filter1d_no_filter(data):
53-
"""
54-
Test that filter1d fails with an argument is missing for filter.
55-
"""
56-
with pytest.raises(GMTInvalidInput):
57-
filter1d(data=data)
58-
59-
60-
def test_filter1d_no_outfile_specified(data):
61-
"""
62-
Test that filter1d fails when outpput_type is set to 'file' but no output file name
63-
is specified.
64-
"""
65-
with pytest.raises(GMTInvalidInput):
66-
filter1d(data=data, filter_type="g5", output_type="file")
67-
68-
69-
def test_filter1d_outfile_incorrect_output_type(data):
70-
"""
71-
Test that filter1d raises a warning when an outfile filename is set but the
72-
output_type is not set to 'file'.
73-
"""
74-
with pytest.warns(RuntimeWarning):
75-
with GMTTempFile(suffix=".txt") as tmpfile:
76-
result = filter1d(
77-
data=data, filter_type="g5", outfile=tmpfile.name, output_type="numpy"
78-
)
79-
assert result is None # return value is None
80-
assert Path(tmpfile.name).stat().st_size > 0 # check that outfile exists
81-
82-
8328
@pytest.mark.benchmark
8429
def test_filter1d_format(data):
8530
"""

0 commit comments

Comments
 (0)