diff --git a/pygmt/src/grdvolume.py b/pygmt/src/grdvolume.py index 02c5567f1b9..d7f232512e2 100644 --- a/pygmt/src/grdvolume.py +++ b/pygmt/src/grdvolume.py @@ -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, @@ -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} @@ -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