|
| 1 | +from dataclasses import asdict |
| 2 | + |
1 | 3 | from langchain.tools import tool |
2 | 4 | from utils.segy_handler import SegyHandler |
3 | 5 | from utils.miniseed_handler import MiniSEEDHandler |
4 | 6 | from utils.hdf5_handler import HDF5Handler |
| 7 | +from utils.phase_picker import pick_phases, summarize_pick_results |
5 | 8 | import json |
6 | 9 | from typing import Union |
7 | 10 | import numpy as np |
@@ -71,6 +74,121 @@ def _coerce_int(value, *, allow_none=False, default=None, field_name="value"): |
71 | 74 | raise ValueError(f"{field_name} must be an integer, got {value!r}") |
72 | 75 |
|
73 | 76 |
|
| 77 | +def _coerce_float(value, *, allow_none=False, default=None, field_name="value"): |
| 78 | + if value is None: |
| 79 | + if allow_none: |
| 80 | + return default |
| 81 | + if default is not None: |
| 82 | + return float(default) |
| 83 | + raise ValueError(f"{field_name} must be provided") |
| 84 | + if isinstance(value, (int, float, np.integer, np.floating)): |
| 85 | + return float(value) |
| 86 | + if isinstance(value, str): |
| 87 | + lowered = value.strip() |
| 88 | + if allow_none and lowered.lower() in {"none", "null", ""}: |
| 89 | + return default |
| 90 | + return float(lowered) |
| 91 | + raise ValueError(f"{field_name} must be a float, got {value!r}") |
| 92 | + |
| 93 | + |
| 94 | +def _normalize_method_list(raw_value): |
| 95 | + if raw_value is None: |
| 96 | + return None |
| 97 | + if isinstance(raw_value, str): |
| 98 | + candidate = raw_value.strip() |
| 99 | + if not candidate: |
| 100 | + return None |
| 101 | + if candidate.startswith("["): |
| 102 | + try: |
| 103 | + data = json.loads(candidate) |
| 104 | + if isinstance(data, list): |
| 105 | + return [str(item).strip() for item in data if str(item).strip()] |
| 106 | + except json.JSONDecodeError: |
| 107 | + pass |
| 108 | + return [part.strip() for part in candidate.split(",") if part.strip()] |
| 109 | + if isinstance(raw_value, (list, tuple, set)): |
| 110 | + return [str(item).strip() for item in raw_value if str(item).strip()] |
| 111 | + return None |
| 112 | + |
| 113 | + |
| 114 | +def _parse_method_params(raw_value): |
| 115 | + if raw_value is None: |
| 116 | + return None |
| 117 | + if isinstance(raw_value, dict): |
| 118 | + return raw_value |
| 119 | + if isinstance(raw_value, str): |
| 120 | + candidate = raw_value.strip() |
| 121 | + if not candidate: |
| 122 | + return None |
| 123 | + try: |
| 124 | + data = json.loads(candidate) |
| 125 | + if isinstance(data, dict): |
| 126 | + return data |
| 127 | + except json.JSONDecodeError: |
| 128 | + return None |
| 129 | + return None |
| 130 | + |
| 131 | + |
| 132 | +def _normalize_source_type(value: str | None): |
| 133 | + if not value: |
| 134 | + return None |
| 135 | + normalized = value.lower().strip() |
| 136 | + alias = { |
| 137 | + "miniseed": "mseed", |
| 138 | + "mseed": "mseed", |
| 139 | + "segy": "segy", |
| 140 | + "sgy": "segy", |
| 141 | + "hdf5": "hdf5", |
| 142 | + "h5": "hdf5", |
| 143 | + "npy": "npy", |
| 144 | + "npz": "npz", |
| 145 | + "sac": "sac", |
| 146 | + } |
| 147 | + return alias.get(normalized, normalized) |
| 148 | + |
| 149 | + |
| 150 | +def _infer_file_type_from_path(path: str | None): |
| 151 | + if not path: |
| 152 | + return None |
| 153 | + ext = os.path.splitext(path)[1].lower() |
| 154 | + mapping = { |
| 155 | + ".segy": "segy", |
| 156 | + ".sgy": "segy", |
| 157 | + ".mseed": "mseed", |
| 158 | + ".miniseed": "mseed", |
| 159 | + ".h5": "hdf5", |
| 160 | + ".hdf5": "hdf5", |
| 161 | + ".npy": "npy", |
| 162 | + ".npz": "npz", |
| 163 | + ".sac": "sac", |
| 164 | + } |
| 165 | + return mapping.get(ext) |
| 166 | + |
| 167 | + |
| 168 | +def _resolve_source_path(path: str | None, source_type: str | None): |
| 169 | + normalized_type = _normalize_source_type(source_type) |
| 170 | + if path: |
| 171 | + inferred = normalized_type or _infer_file_type_from_path(path) |
| 172 | + return path, inferred |
| 173 | + |
| 174 | + candidates = [ |
| 175 | + ("segy", CURRENT_SEGY_PATH), |
| 176 | + ("mseed", CURRENT_MINISEED_PATH), |
| 177 | + ("hdf5", CURRENT_HDF5_PATH), |
| 178 | + ] |
| 179 | + |
| 180 | + if normalized_type: |
| 181 | + for ctype, cpath in candidates: |
| 182 | + if ctype == normalized_type and cpath: |
| 183 | + return cpath, ctype |
| 184 | + return None, normalized_type |
| 185 | + |
| 186 | + for ctype, cpath in candidates: |
| 187 | + if cpath: |
| 188 | + return cpath, ctype |
| 189 | + return None, None |
| 190 | + |
| 191 | + |
74 | 192 | def _resolve_output_path(output_path: str | None, *, default_filename: str) -> str: |
75 | 193 | """Resolve output file path. |
76 | 194 |
|
@@ -603,3 +721,64 @@ def convert_miniseed_to_sac(params: Union[str, dict, None] = None): |
603 | 721 | if "error" in result: |
604 | 722 | return json.dumps(result, indent=2) |
605 | 723 | return json.dumps(result, indent=2) |
| 724 | + |
| 725 | +# 运行传统拾取方法 |
| 726 | +@tool |
| 727 | +def run_phase_picking(params: Union[str, dict, None] = None): |
| 728 | + """ |
| 729 | + Run classical phase picking on the loaded file (SEGY/MiniSEED/HDF5/NumPy) or a specified path. |
| 730 | + Args: path (optional), file_type/source_type, dataset (HDF5), sampling_rate (for NumPy), methods, method_params. |
| 731 | + """ |
| 732 | + """在已加载或指定的地震数据文件上运行传统初至拾取。参数：path（可选）、file_type/source_type、dataset、sampling_rate、methods、method_params。""" |
| 733 | + parsed = _parse_param_dict(params) |
| 734 | + |
| 735 | + requested_path = parsed.get("path") |
| 736 | + source_type = parsed.get("file_type") or parsed.get("source_type") |
| 737 | + dataset = parsed.get("dataset") |
| 738 | + |
| 739 | + try: |
| 740 | + sampling_rate = _coerce_float( |
| 741 | + parsed.get("sampling_rate"), |
| 742 | + allow_none=True, |
| 743 | + default=None, |
| 744 | + field_name="sampling_rate", |
| 745 | + ) |
| 746 | + except ValueError as exc: |
| 747 | + return str(exc) |
| 748 | + |
| 749 | + methods = _normalize_method_list(parsed.get("methods")) |
| 750 | + method_params = _parse_method_params(parsed.get("method_params")) |
| 751 | + |
| 752 | + resolved_path, inferred_type = _resolve_source_path(requested_path, source_type) |
| 753 | + if not resolved_path: |
| 754 | + return "No suitable data file is currently loaded. Please upload or specify a file path." |
| 755 | + |
| 756 | + file_type = _normalize_source_type(source_type) or inferred_type |
| 757 | + |
| 758 | + try: |
| 759 | + picks = pick_phases( |
| 760 | + resolved_path, |
| 761 | + file_type=file_type, |
| 762 | + dataset=dataset, |
| 763 | + sampling_rate=sampling_rate, |
| 764 | + methods=methods, |
| 765 | + method_params=method_params, |
| 766 | + ) |
| 767 | + except Exception as exc: |
| 768 | + return json.dumps({"error": str(exc)}, indent=2) |
| 769 | + |
| 770 | + if not picks: |
| 771 | + return "No phase arrivals were detected." |
| 772 | + |
| 773 | + serialized = [asdict(item) for item in picks] |
| 774 | + summary = summarize_pick_results(picks) |
| 775 | + return json.dumps( |
| 776 | + { |
| 777 | + "file": resolved_path, |
| 778 | + "file_type": file_type, |
| 779 | + "count": len(serialized), |
| 780 | + "results": serialized, |
| 781 | + "summary": summary, |
| 782 | + }, |
| 783 | + indent=2, |
| 784 | + ) |
0 commit comments