diff --git a/pygmt/src/_common.py b/pygmt/src/_common.py index 96ae9791b71..408204d050a 100644 --- a/pygmt/src/_common.py +++ b/pygmt/src/_common.py @@ -2,9 +2,12 @@ Common functions used in multiple PyGMT functions/methods. """ +from collections.abc import Sequence +from enum import StrEnum from pathlib import Path -from typing import Any +from typing import Any, ClassVar, Literal +from pygmt.exceptions import GMTInvalidInput from pygmt.src.which import which @@ -39,3 +42,212 @@ def _data_geometry_is_point(data: Any, kind: str) -> bool: except FileNotFoundError: pass return False + + +class _FocalMechanismConventionCode(StrEnum): + """ + Enum to handle focal mechanism convention codes. + + The enum names are in the format of ``CONVENTION_COMPONENT``, where ``CONVENTION`` + is the focal mechanism convention and ``COMPONENT`` is the component of the seismic + moment tensor to plot. The enum values are the single-letter codes that can be used + in meca/coupe's ``-S`` option. + + For some conventions, ``COMPONENT`` is not applicable, but we still define the enums + to simplify the code logic. + """ + + AKI_DC = "a" + AKI_DEVIATORIC = "a" + AKI_FULL = "a" + GCMT_DC = "c" + GCMT_DEVIATORIC = "c" + GCMT_FULL = "c" + PARTIAL_DC = "p" + PARTIAL_DEVIATORIC = "p" + PARTIAL_FULL = "p" + MT_DC = "d" + MT_DEVIATORIC = "z" + MT_FULL = "m" + PRINCIPAL_AXIS_DC = "y" + PRINCIPAL_AXIS_DEVIATORIC = "t" + PRINCIPAL_AXIS_FULL = "x" + + +class _FocalMechanismConvention: + """ + Class to handle focal mechanism convention, code, and associated parameters. + + Examples + -------- + >>> from pygmt.src._common import _FocalMechanismConvention + + >>> conv = _FocalMechanismConvention("aki") + >>> conv.code + <_FocalMechanismConventionCode.AKI_DC: 'a'> + >>> conv.params + ['strike', 'dip', 'rake', 'magnitude'] + + >>> conv = _FocalMechanismConvention("mt") + >>> conv.code + <_FocalMechanismConventionCode.MT_FULL: 'm'> + >>> conv.params + ['mrr', 'mtt', 'mff', 'mrt', 'mrf', 'mtf', 'exponent'] + + >>> conv = _FocalMechanismConvention("mt", component="dc") + >>> conv.code + <_FocalMechanismConventionCode.MT_DC: 'd'> + >>> conv.params + ['mrr', 'mtt', 'mff', 'mrt', 'mrf', 'mtf', 'exponent'] + + >>> conv = _FocalMechanismConvention("a") + >>> conv.code + <_FocalMechanismConventionCode.AKI_DC: 'a'> + >>> conv.params + ['strike', 'dip', 'rake', 'magnitude'] + + >>> conv = _FocalMechanismConvention.from_params( + ... ["strike", "dip", "rake", "magnitude"] + ... ) + >>> conv.code + <_FocalMechanismConventionCode.AKI_DC: 'a'> + + >>> conv = _FocalMechanismConvention(convention="invalid") + Traceback (most recent call last): + ... + pygmt.exceptions.GMTInvalidInput: Invalid focal mechanism ...'. + + >>> conv = _FocalMechanismConvention("mt", component="invalid") + Traceback (most recent call last): + ... + pygmt.exceptions.GMTInvalidInput: Invalid focal mechanism ...'. + + >>> _FocalMechanismConvention.from_params(["strike", "dip", "rake"]) + Traceback (most recent call last): + ... + pygmt.exceptions.GMTInvalidInput: Fail to determine focal mechanism convention... + """ + + # Mapping of focal mechanism conventions to their parameters. + _params: ClassVar = { + "aki": ["strike", "dip", "rake", "magnitude"], + "gcmt": [ + "strike1", + "dip1", + "rake1", + "strike2", + "dip2", + "rake2", + "mantissa", + "exponent", + ], + "partial": ["strike1", "dip1", "strike2", "fault_type", "magnitude"], + "mt": ["mrr", "mtt", "mff", "mrt", "mrf", "mtf", "exponent"], + "principal_axis": [ + "t_value", + "t_azimuth", + "t_plunge", + "n_value", + "n_azimuth", + "n_plunge", + "p_value", + "p_azimuth", + "p_plunge", + "exponent", + ], + } + + def __init__( + self, + convention: Literal["aki", "gcmt", "partial", "mt", "principal_axis"], + component: Literal["full", "deviatoric", "dc"] = "full", + ): + """ + Initialize the ``_FocalMechanismConvention`` object from ``convention`` and + ``component``. + + If the convention is specified via a single-letter code, ``convention`` and + ``component`` are determined from the code. + + Parameters + ---------- + convention + The focal mechanism convention. Valid values are: + + - ``"aki"``: Aki and Richards convention. + - ``"gcmt"``: Global CMT (Centroid Moment Tensor) convention. + - ``"partial"``: Partial focal mechanism convention. + - ``"mt"``: Moment Tensor convention. + - ``"principal_axis"``: Principal axis convention. + component + The component of the seismic moment tensor to plot. Valid values are: + + - ``"full"``: the full seismic moment tensor + - ``"dc"``: the closest double couple defined from the moment tensor (zero + trace and zero determinant) + - ``"deviatoric"``: deviatoric part of the moment tensor (zero trace) + + Doesn't apply to the conventions ``"aki"``, ``"gcmt"``, and ``"partial"``. + """ + # TODO(Python>=3.12): Simplify to "convention in _FocalMechanismConventionCode". + if convention in _FocalMechanismConventionCode.__members__.values(): + # Convention is specified via the actual single-letter convention code. + self.code = _FocalMechanismConventionCode(convention) + # Parse the convention from the convention code name. + self._convention = "_".join(self.code.name.split("_")[:-1]).lower() + else: # Convention is specified via "convention" and "component". + name = f"{convention.upper()}_{component.upper()}" # e.g., "AKI_DC" + if name not in _FocalMechanismConventionCode.__members__: + msg = ( + "Invalid focal mechanism convention with " + f"convention='{convention}' and component='{component}'." + ) + raise GMTInvalidInput(msg) + self.code = _FocalMechanismConventionCode[name] + self._convention = convention + + @property + def params(self): + """ + The parameters associated with the focal mechanism convention. + """ + return self._params[self._convention] + + @classmethod + def from_params( + cls, + params: Sequence[str], + component: Literal["full", "deviatoric", "dc"] = "full", + ) -> "_FocalMechanismConvention": + """ + Create a _FocalMechanismConvention object from a sequence of parameters. + + The method checks if the given parameters are a superset of a supported focal + mechanism convention to determine the convention. If the parameters are not + sufficient to determine the convention, an exception is raised. + + Parameters + ---------- + params + Sequence of parameters to determine the focal mechanism convention. The + order of the parameters does not matter. + + Returns + ------- + _FocalMechanismConvention + The _FocalMechanismConvention object. + + Raises + ------ + GMTInvalidInput + If the focal mechanism convention cannot be determined from the given + parameters. + """ + for convention, param_list in cls._params.items(): + if set(param_list).issubset(set(params)): + return cls(convention, component=component) + msg = ( + "Fail to determine focal mechanism convention from the given parameters: " + f"{', '.join(params)}." + ) + raise GMTInvalidInput(msg) diff --git a/pygmt/src/meca.py b/pygmt/src/meca.py index e17cdceb4fd..2fb9638e8e7 100644 --- a/pygmt/src/meca.py +++ b/pygmt/src/meca.py @@ -5,183 +5,9 @@ import numpy as np import pandas as pd from pygmt.clib import Session -from pygmt.exceptions import GMTError, GMTInvalidInput +from pygmt.exceptions import GMTInvalidInput from pygmt.helpers import build_arg_list, fmt_docstring, kwargs_to_strings, use_alias - - -def convention_code(convention, component="full"): - """ - Determine the convention code for focal mechanisms. - - The convention code can be used in meca's -S option. - - Parameters - ---------- - convention : str - The focal mechanism convention. Can be one of the following: - - - ``"aki"``: Aki and Richards - - ``"gcmt"``: Global Centroid Moment Tensor - - ``"partial"``: Partial focal mechanism - - ``"mt"``: Moment tensor - - ``"principal_axis"``: Principal axis - - Single letter convention codes like ``"a"`` and ``"c"`` are also - supported but undocumented. - - component : str - The component of the focal mechanism. Only used when ``convention`` is - ``"mt"`` or ``"principal_axis"``. Can be one of the following: - - - ``"full"``: Full moment tensor - - ``"deviatoric"``: Deviatoric moment tensor - - ``"dc"``: Double couple - - Returns - ------- - str - The single-letter convention code used in meca's -S option. - - Examples - -------- - >>> convention_code("aki") - 'a' - >>> convention_code("gcmt") - 'c' - >>> convention_code("partial") - 'p' - - >>> convention_code("mt", component="full") - 'm' - >>> convention_code("mt", component="deviatoric") - 'z' - >>> convention_code("mt", component="dc") - 'd' - >>> convention_code("principal_axis", component="full") - 'x' - >>> convention_code("principal_axis", component="deviatoric") - 't' - >>> convention_code("principal_axis", component="dc") - 'y' - - >>> for code in ["a", "c", "m", "d", "z", "p", "x", "y", "t"]: - ... assert convention_code(code) == code - - >>> convention_code("invalid") - Traceback (most recent call last): - ... - pygmt.exceptions.GMTInvalidInput: Invalid convention 'invalid'. - - >>> convention_code("mt", "invalid") # doctest: +NORMALIZE_WHITESPACE - Traceback (most recent call last): - ... - pygmt.exceptions.GMTInvalidInput: - Invalid component 'invalid' for convention 'mt'. - """ - # Codes for focal mechanism formats determined by "convention" - codes1 = {"aki": "a", "gcmt": "c", "partial": "p"} - # Codes for focal mechanism formats determined by both "convention" and - # "component" - codes2 = { - "mt": {"deviatoric": "z", "dc": "d", "full": "m"}, - "principal_axis": {"deviatoric": "t", "dc": "y", "full": "x"}, - } - - if convention in codes1: - return codes1[convention] - if convention in codes2: - if component not in codes2[convention]: - msg = f"Invalid component '{component}' for convention '{convention}'." - raise GMTInvalidInput(msg) - return codes2[convention][component] - if convention in {"a", "c", "m", "d", "z", "p", "x", "y", "t"}: - return convention - msg = f"Invalid convention '{convention}'." - raise GMTInvalidInput(msg) - - -def convention_name(code): - """ - Determine the name of a focal mechanism convention from its code. - - Parameters - ---------- - code : str - The single-letter convention code. - - Returns - ------- - str - The name of the focal mechanism convention. - - Examples - -------- - >>> convention_name("a") - 'aki' - >>> convention_name("aki") - 'aki' - """ - name = { - "a": "aki", - "c": "gcmt", - "p": "partial", - "z": "mt", - "d": "mt", - "m": "mt", - "x": "principal_axis", - "y": "principal_axis", - "t": "principal_axis", - }.get(code) - return name if name is not None else code - - -def convention_params(convention): - """ - Return the list of focal mechanism parameters for a given convention. - - Parameters - ---------- - convention : str - The focal mechanism convention. Can be one of the following: - - - ``"aki"``: Aki and Richards - - ``"gcmt"``: Global Centroid Moment Tensor - - ``"partial"``: Partial focal mechanism - - ``"mt"``: Moment tensor - - ``"principal_axis"``: Principal axis - - Returns - ------- - list - The list of focal mechanism parameters. - """ - return { - "aki": ["strike", "dip", "rake", "magnitude"], - "gcmt": [ - "strike1", - "dip1", - "rake1", - "strike2", - "dip2", - "rake2", - "mantissa", - "exponent", - ], - "mt": ["mrr", "mtt", "mff", "mrt", "mrf", "mtf", "exponent"], - "partial": ["strike1", "dip1", "strike2", "fault_type", "magnitude"], - "principal_axis": [ - "t_value", - "t_azimuth", - "t_plunge", - "n_value", - "n_azimuth", - "n_plunge", - "p_value", - "p_azimuth", - "p_plunge", - "exponent", - ], - }[convention] +from pygmt.src._common import _FocalMechanismConvention @fmt_docstring @@ -204,7 +30,7 @@ def convention_params(convention): t="transparency", ) @kwargs_to_strings(R="sequence", c="sequence_comma", p="sequence") -def meca( # noqa: PLR0912, PLR0913, PLR0915 +def meca( # noqa: PLR0912, PLR0913 self, spec, scale, @@ -410,17 +236,10 @@ def meca( # noqa: PLR0912, PLR0913, PLR0915 # 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", "principal_axis"]: - if set(convention_params(conv)).issubset(set(spec.keys())): - convention = conv - break - else: - if isinstance(spec, dict): - msg = "Keys in dict 'spec' do not match known conventions." - else: - msg = "Column names in pd.DataFrame 'spec' do not match known conventions." - raise GMTError(msg) + # Determine convention from dict keys or pd.DataFrame column names + _convention = _FocalMechanismConvention.from_params( + spec.keys(), component=component + ) # convert dict to pd.DataFrame so columns can be reordered if isinstance(spec, dict): @@ -434,12 +253,14 @@ def meca( # noqa: PLR0912, PLR0913, PLR0915 if convention is None: msg = "'convention' must be specified for an array input." raise GMTInvalidInput(msg) - # make sure convention is a name, not a code - convention = convention_name(convention) + + _convention = _FocalMechanismConvention( + convention=convention, component=component + ) # Convert array to pd.DataFrame and assign column names spec = pd.DataFrame(np.atleast_2d(spec)) - colnames = ["longitude", "latitude", "depth", *convention_params(convention)] + colnames = ["longitude", "latitude", "depth", *_convention.params] # check if spec has the expected number of columns ncolsdiff = len(spec.columns) - len(colnames) if ncolsdiff == 0: @@ -456,6 +277,10 @@ def meca( # noqa: PLR0912, PLR0913, PLR0915 ) raise GMTInvalidInput(msg) spec.columns = colnames + else: + _convention = _FocalMechanismConvention( + convention=convention, component=component + ) # Now spec is a pd.DataFrame or a file if isinstance(spec, pd.DataFrame): @@ -483,7 +308,7 @@ def meca( # noqa: PLR0912, PLR0913, PLR0915 # expected columns are: # longitude, latitude, depth, focal_parameters, # [plot_longitude, plot_latitude] [event_name] - newcols = ["longitude", "latitude", "depth", *convention_params(convention)] + newcols = ["longitude", "latitude", "depth", *_convention.params] if "plot_longitude" in spec.columns and "plot_latitude" in spec.columns: newcols += ["plot_longitude", "plot_latitude"] if kwargs.get("A") is None: @@ -494,11 +319,7 @@ def meca( # noqa: PLR0912, PLR0913, PLR0915 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) - - # Assemble -S flag - kwargs["S"] = f"{data_format}{scale}" + kwargs["S"] = f"{_convention.code}{scale}" with Session() as lib: with lib.virtualfile_in(check_kind="vector", data=spec) as vintbl: lib.call_module(module="meca", args=build_arg_list(kwargs, infile=vintbl))