Fix beachball offsetting for ndarray input

pygmt/src/meca.py

@@ -398,7 +398,10 @@ def meca(
     {transparency}
     """
     # pylint: disable=too-many-arguments,too-many-locals,too-many-branches
+    # pylint: disable=too-many-statements
     kwargs = self._preprocess(**kwargs)  # pylint: disable=protected-access
+
+    # Convert spec to pandas.DataFrame unless it's a file
     if isinstance(spec, (dict, pd.DataFrame)):  # spec is a dict or pd.DataFrame
         # determine convention from dict keys or pd.DataFrame column names
         for conv in ["aki", "gcmt", "mt", "partial", "pricipal_axis"]:
@@ -412,8 +415,42 @@ def meca(
                 msg = "Column names in pd.DataFrame 'spec' do not match known conventions."
             raise GMTError(msg)
 
-        # override the values in dict/pd.DataFrame if parameters are explicity
-        # specified
+        # convert dict to pd.DataFrame so columns can be reordered
+        if isinstance(spec, dict):
+            # convert values to ndarray so pandas doesn't complain about "all
+            # scalar values". See
+            # https://github.com/GenericMappingTools/pygmt/pull/2174
+            spec = pd.DataFrame(
+                {key: np.atleast_1d(value) for key, value in spec.items()}
+            )
+    elif isinstance(spec, np.ndarray):  # spec is a numpy array
+        if convention is None:
+            raise GMTInvalidInput("'convention' must be specified for an array input.")
+        # make sure convention is a name, not a code
+        convention = convention_name(convention)
+
+        # Convert array to pd.DataFrame and assign column names
+        spec = pd.DataFrame(np.atleast_2d(spec))
+        colnames = ["longitude", "latitude", "depth"] + convention_params(convention)
+        # check if spec has the expected number of columns
+        ncolsdiff = len(spec.columns) - len(colnames)
+        if ncolsdiff == 0:
+            pass
+        elif ncolsdiff == 1:
+            colnames += ["event_name"]
+        elif ncolsdiff == 2:
+            colnames += ["plot_longitude", "plot_latitude"]
+        elif ncolsdiff == 3:
+            colnames += ["plot_longitude", "plot_latitude", "event_name"]
+        else:
+            raise GMTInvalidInput(
+                f"Input array must have {len(colnames)} to {len(colnames) + 3} columns."
+            )
+        spec.columns = colnames
+
+    # Now spec is a pd.DataFrame or a file
+    if isinstance(spec, pd.DataFrame):
+        # override the values in pd.DataFrame if parameters are given
         if longitude is not None:
             spec["longitude"] = np.atleast_1d(longitude)
         if latitude is not None:
@@ -425,35 +462,30 @@ def meca(
         if plot_latitude is not None:
             spec["plot_latitude"] = np.atleast_1d(plot_latitude)
         if event_name is not None:
-            spec["event_name"] = np.atleast_1d(event_name).astype(str)
+            spec["event_name"] = np.atleast_1d(event_name)
 
-        # convert dict to pd.DataFrame so columns can be reordered
-        if isinstance(spec, dict):
-            # convert values to ndarray so pandas doesn't complain about "all
-            # scalar values". See
-            # https://github.com/GenericMappingTools/pygmt/pull/2174
-            spec = {key: np.atleast_1d(value) for key, value in spec.items()}
-            spec = pd.DataFrame(spec)
+        # Due to the internal implementation of the meca module, we need to
+        # convert the following columns to strings if they exist
+        if "plot_longitude" in spec.columns and "plot_latitude" in spec.columns:
+            spec["plot_longitude"] = spec["plot_longitude"].astype(str)
+            spec["plot_latitude"] = spec["plot_latitude"].astype(str)
+        if "event_name" in spec.columns:
+            spec["event_name"] = spec["event_name"].astype(str)
 
+        # Reorder columns in DataFrame to match convention if necessary
         # expected columns are:
         # longitude, latitude, depth, focal_parameters,
         #   [plot_longitude, plot_latitude] [event_name]
         newcols = ["longitude", "latitude", "depth"] + convention_params(convention)
         if "plot_longitude" in spec.columns and "plot_latitude" in spec.columns:
             newcols += ["plot_longitude", "plot_latitude"]
-            spec[["plot_longitude", "plot_latitude"]] = spec[
-                ["plot_longitude", "plot_latitude"]
-            ].astype(str)
             if kwargs.get("A") is None:
                 kwargs["A"] = True
         if "event_name" in spec.columns:
             newcols += ["event_name"]
-            spec["event_name"] = spec["event_name"].astype(str)
         # reorder columns in DataFrame
-        spec = spec.reindex(newcols, axis=1)
-    elif isinstance(spec, np.ndarray) and spec.ndim == 1:
-        # Convert 1-D array into 2-D array
-        spec = np.atleast_2d(spec)
+        if spec.columns.tolist() != newcols:
+            spec = spec.reindex(newcols, axis=1)
 
     # determine data_format from convention and component
     data_format = convention_code(convention=convention, component=component)

pygmt/tests/test_meca.py

@@ -143,12 +143,12 @@ def test_meca_spec_multiple_focalmecha(inputtype):
 
 
 @pytest.mark.mpl_image_compare(filename="test_meca_offset.png")
-@pytest.mark.parametrize("inputtype", ["offset_args", "offset_dict", "ndarray"])
+@pytest.mark.parametrize("inputtype", ["args", "dict", "ndarray"])
 def test_meca_offset(inputtype):
     """
     Test offsetting beachballs.
     """
-    if inputtype == "offset_args":
+    if inputtype == "args":
         args = {
             "spec": {"strike": 330, "dip": 30, "rake": 90, "magnitude": 3},
             "longitude": -124,
@@ -157,7 +157,7 @@ def test_meca_offset(inputtype):
             "plot_longitude": -124.5,
             "plot_latitude": 47.5,
         }
-    elif inputtype == "offset_dict":
+    elif inputtype == "dict":
         # Test https://github.com/GenericMappingTools/pygmt/issues/2016
         # offset parameters are in the dict.
         args = {