Skip to content

Commit

Permalink
Add pygmt.read to read a dataset/grid/image into pandas.DataFrame/xar…
Browse files Browse the repository at this point in the history
…ray.DataArray
  • Loading branch information
seisman committed Dec 4, 2024
1 parent a92d9d0 commit cac7d74
Show file tree
Hide file tree
Showing 5 changed files with 84 additions and 6 deletions.
1 change: 1 addition & 0 deletions doc/api/index.rst
Original file line number Diff line number Diff line change
Expand Up @@ -172,6 +172,7 @@ Input/output
:toctree: generated

load_dataarray
read

GMT Defaults
------------
Expand Down
1 change: 1 addition & 0 deletions pygmt/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@
makecpt,
nearneighbor,
project,
read,
select,
sph2grd,
sphdistance,
Expand Down
11 changes: 5 additions & 6 deletions pygmt/helpers/testing.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,9 @@
import string
from pathlib import Path

import xarray as xr
from pygmt.exceptions import GMTImageComparisonFailure
from pygmt.io import load_dataarray
from pygmt.src import which
from pygmt.src import read


def check_figures_equal(*, extensions=("png",), tol=0.0, result_dir="result_images"):
Expand Down Expand Up @@ -144,17 +144,16 @@ def wrapper(*args, ext="png", request=None, **kwargs):
return decorator


def load_static_earth_relief():
def load_static_earth_relief() -> xr.DataArray:
"""
Load the static_earth_relief file for internal testing.
Returns
-------
data : xarray.DataArray
data
A grid of Earth relief for internal tests.
"""
fname = which("@static_earth_relief.nc", download="c")
return load_dataarray(fname)
return read("@static_earth_relief.nc", kind="grid") # type: ignore[return-value]


def skip_if_no(package):
Expand Down
1 change: 1 addition & 0 deletions pygmt/src/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,7 @@
from pygmt.src.plot3d import plot3d
from pygmt.src.project import project
from pygmt.src.psconvert import psconvert
from pygmt.src.read import read
from pygmt.src.rose import rose
from pygmt.src.select import select
from pygmt.src.shift_origin import shift_origin
Expand Down
76 changes: 76 additions & 0 deletions pygmt/src/read.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
"""
Read data from files
"""

from typing import Literal

import pandas as pd
import xarray as xr
from pygmt.clib import Session


def read(
file, kind: Literal["dataset", "grid", "image"], **kwargs
) -> pd.DataFrame | xr.DataArray:
"""
Read a dataset, grid, or image from a file and return the appropriate object.
For datasets, it returns a :class:`pandas.DataFrame`. For grids and images,
it returns a :class:`xarray.DataArray`.
Parameters
----------
file
The file name to read.
kind
The kind of data to read. Valid values are ``"dataset"``, ``"grid"``,
and ``"image"``.
For datasets, the following keyword arguments are supported:
column_names
A list of column names.
header
Row number containing column names. ``header=None`` means not to parse the
column names from table header. Ignored if the row number is larger than the
number of headers in the table.
dtype
Data type. Can be a single type for all columns or a dictionary mapping
column names to types.
index_col
Column to set as index.
Returns
-------
Return type depends on the ``kind`` argument:
- ``"dataset"``: :class:`pandas.DataFrame`
- ``"grid"`` or ``"image"``: :class:`xarray.DataArray`
Examples
--------
Read a dataset into a :class:`pandas.DataFrame` object:
>>> from pygmt import read
>>> df = read("@hotspots.txt", kind="dataset")
>>> type(df)
<class 'pandas.core.frame.DataFrame'>
Read a grid into an :class:`xarray.DataArray` object:
>>> dataarray = read("@earth_relief_01d", kind="grid")
>>> type(dataarray)
<class 'xarray.core.dataarray.DataArray'>
"""
code = {"dataset": "d", "grid": "g", "image": "i"}[kind]

with Session() as lib:
with lib.virtualfile_out(kind=kind) as voutfile:
lib.call_module("read", args=[file, voutfile, f"-T{code}"])

match kind:
case "dataset":
return lib.virtualfile_to_dataset(vfname=voutfile, **kwargs)
case "grid" | "image":
return lib.virtualfile_to_raster(vfname=voutfile, kind=kind)

0 comments on commit cac7d74

Please sign in to comment.