Skip to content

Commit

Permalink
add formatting to return array or DataFrame
Browse files Browse the repository at this point in the history
  • Loading branch information
willschlitzer committed May 26, 2021
1 parent d3ed614 commit 68fab7c
Showing 1 changed file with 20 additions and 1 deletion.
21 changes: 20 additions & 1 deletion pygmt/src/grdvolume.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,10 @@
"""
grdvolume - Calculate grid volume and area constrained by a contour.
"""
import numpy as np
import pandas as pd
from pygmt.clib import Session
from pygmt.exceptions import GMTInvalidInput
from pygmt.helpers import (
GMTTempFile,
build_arg_string,
Expand All @@ -21,7 +24,7 @@
V="verbose",
)
@kwargs_to_strings(C="sequence", R="sequence")
def grdvolume(grid, **kwargs):
def grdvolume(grid, data_format="a", **kwargs):
r"""
{aliases}
Expand All @@ -47,4 +50,20 @@ def grdvolume(grid, **kwargs):
)
lib.call_module("grdvolume", arg_str)
result = outfile.read()
if data_format == "s":
return result
data_list = []
for string_entry in result.strip().split("\n"):
float_entry = []
string_list = string_entry.strip().split()
for i in string_list:
float_entry.append(float(i))
data_list.append(float_entry)
data_array = np.array(data_list)
if data_format == "a":
result = data_array
elif data_format == "d":
result = pd.DataFrame(data_array)
else:
raise GMTInvalidInput("""Must specify format as either a, d, or s.""")
return result

0 comments on commit 68fab7c

Please sign in to comment.