Let x2sys_file handle multiple input tracks and external crossovers

Author: weiji14

pygmt/tests/test_x2sys_cross.py

@@ -5,6 +5,7 @@
 import os
 from tempfile import TemporaryDirectory
 
+import numpy as np
 import pandas as pd
 import pytest
 
@@ -25,14 +26,15 @@ def fixture_mock_x2sys_home(monkeypatch):
 
 def test_x2sys_cross_input_file_output_file(mock_x2sys_home):
     """
-    Run x2sys_cross by passing in a filename and output to an ASCII txt file
+    Run x2sys_cross by passing in a filename, and output internal crossovers to
+    an ASCII txt file
     """
     with TemporaryDirectory(prefix="X2SYS", dir=os.getcwd()) as tmpdir:
         tag = os.path.basename(tmpdir)
         x2sys_init(tag=tag, fmtfile="xyz", force=True)
         outfile = os.path.join(tmpdir, "tmp_coe.txt")
         output = x2sys_cross(
-            tracks=["@tut_ship.xyz"], tag=tag, coe="i", outfile=outfile, verbose="d"
+            tracks=["@tut_ship.xyz"], tag=tag, coe="i", outfile=outfile, verbose="i"
         )
 
         assert output is None  # check that output is None since outfile is set
@@ -44,12 +46,13 @@ def test_x2sys_cross_input_file_output_file(mock_x2sys_home):
 
 def test_x2sys_cross_input_file_output_dataframe(mock_x2sys_home):
     """
-    Run x2sys_cross by passing in a filename and output to a pandas.DataFrame
+    Run x2sys_cross by passing in a filename, and output internal crossovers to
+    a pandas.DataFrame
     """
     with TemporaryDirectory(prefix="X2SYS", dir=os.getcwd()) as tmpdir:
         tag = os.path.basename(tmpdir)
         x2sys_init(tag=tag, fmtfile="xyz", force=True)
-        output = x2sys_cross(tracks=["@tut_ship.xyz"], tag=tag, coe="i", verbose="d")
+        output = x2sys_cross(tracks=["@tut_ship.xyz"], tag=tag, coe="i", verbose="i")
 
         assert isinstance(output, pd.DataFrame)
         assert output.shape == (14294, 12)
@@ -58,3 +61,31 @@ def test_x2sys_cross_input_file_output_dataframe(mock_x2sys_home):
         assert columns[6:] == ["head_1", "head_2", "vel_1", "vel_2", "z_X", "z_M"]
 
     return output
+
+
+def test_x2sys_cross_input_two_filenames(mock_x2sys_home):
+    """
+    Run x2sys_cross by passing in two filenames, and output external crossovers
+    to a pandas.DataFrame
+    """
+    with TemporaryDirectory(prefix="X2SYS", dir=os.getcwd()) as tmpdir:
+        tag = os.path.basename(tmpdir)
+        x2sys_init(tag=tag, fmtfile="xyz", force=True)
+
+        # Create temporary xyz files
+        for i in range(2):
+            np.random.seed(seed=i)
+            with open(os.path.join(os.getcwd(), f"track_{i}.xyz"), mode="w") as fname:
+                np.savetxt(fname=fname, X=np.random.rand(10, 3))
+
+        output = x2sys_cross(
+            tracks=["track_0.xyz", "track_1.xyz"], tag=tag, coe="e", verbose="i"
+        )
+
+        assert isinstance(output, pd.DataFrame)
+        assert output.shape == (24, 12)
+        columns = list(output.columns)
+        assert columns[:6] == ["x", "y", "i_1", "i_2", "dist_1", "dist_2"]
+        assert columns[6:] == ["head_1", "head_2", "vel_1", "vel_2", "z_X", "z_M"]
+
+    return output

pygmt/x2sys.py

@@ -1,20 +1,22 @@
 """
 GMT supplementary X2SYS module for crossover analysis.
 """
+import contextlib
+
 import pandas as pd
 
 from .clib import Session
+
+# from .exceptions import GMTInvalidInput
 from .helpers import (
-    build_arg_string,
-    fmt_docstring,
     GMTTempFile,
-    use_alias,
+    build_arg_string,
     data_kind,
     dummy_context,
+    fmt_docstring,
+    use_alias,
 )
 
-# from .exceptions import GMTInvalidInput
-
 
 @fmt_docstring
 @use_alias(D="fmtfile", F="force", V="verbose")
@@ -86,10 +88,19 @@ def x2sys_cross(tracks=None, outfile=None, **kwargs):
 
     Parameters
     ----------
-    tracks : pandas.DataFrame or str
-        Either a table with (x, y) or (lon, lat) values in the first two
-        columns, or a filename (e.g. csv, txt format). More columns may be
-        present.
+    tracks : pandas.DataFrame or str or list
+        Either a single pandas.DataFrame table with (x, y) or (lon, lat) values
+        in the first two columns, or a single filename (supported formats are
+        ASCII .txt, native binary, or COARDS netCDF 1-D .nc), or a combination
+        of tables and/or filenames in a list. More columns may also be present.
+
+        If the filenames are missing their file extension, we will append the
+        suffix specified for this TAG. Track files will be searched for first
+        in the current directory and second in all directories listed in
+        $X2SYS_HOME/TAG/TAG_paths.txt (if it exists). [If $X2SYS_HOME is not
+        set it will default to $GMT_SHAREDIR/x2sys]. (Note: MGD77 files will
+        also be looked for via $MGD77_HOME/mgd77_paths.txt and *.gmt files will
+        be searched for via $GMT_SHAREDIR/mgg/gmtfile_paths).
 
     outfile : str
         Optional. The file name for the output ASCII txt file to store the
@@ -115,22 +126,21 @@ def x2sys_cross(tracks=None, outfile=None, **kwargs):
         - None if outfile is set (track output will be stored in outfile)
 
     """
-    # kinds = [data_kind(data=track) for track in tracks]
-    track = tracks[0]
-    kind = data_kind(track)
-
     with Session() as lib:
-        # for kind, track in zip(kinds, tracks):
-        if kind == "file":
-            file_context = dummy_context(track)
-        elif kind == "matrix":
-            file_context = lib.virtualfile_from_matrix(track.values)
+        file_contexts = []
+        for track in tracks:
+            kind = data_kind(track)
+            if kind == "file":
+                file_contexts.append(dummy_context(track))
+            elif kind == "matrix":
+                file_contexts.append(lib.virtualfile_from_matrix(track.values))
 
         with GMTTempFile(suffix=".txt") as tmpfile:
-            with file_context as fname:
+            with contextlib.ExitStack() as stack:
+                fnames = [stack.enter_context(c) for c in file_contexts]
                 if outfile is None:
                     outfile = tmpfile.name
-                arg_str = " ".join([fname, build_arg_string(kwargs), "->" + outfile])
+                arg_str = " ".join([*fnames, build_arg_string(kwargs), "->" + outfile])
                 lib.call_module(module="x2sys_cross", args=arg_str)
 
             # Read temporary csv output to a pandas table