欢迎使用 SeisPolarity

SeisPolarity 是一个全面的地震初动极性拾取框架。

SeisPolarity 为将深度学习模型应用于地震波形极性分类提供了统一的 API，支持多种数据集和模型。

特性

统一数据接口：支持 SCSN、Txed、DiTing、Instance、PNW 数据集，并支持自动下载
多种模型：Ross、Eqpolarity、APP、DiTingMotion、CFM、RPNet、PolarCAP，均提供预训练权重
灵活的数据加载：支持 RAM / 磁盘流式加载任意规模的数据集
高级推理：Predictor 类支持从 Hugging Face / ModelScope 自动下载预训练模型
数据增强：平衡采样及多种数据增强技术
统一训练：支持检查点保存和早停机制
跨平台支持：Linux、macOS 和 Windows

快速开始

安装

pip install seispolarity

git clone https://github.com/Chuan1937/SeisPolarity.git
cd SeisPolarity
pip install -e .

推理

from seispolarity.inference import Predictor
import numpy as np

predictor = Predictor("ROSS_SCSN")
waveforms = np.random.randn(10, 400)  # (Batch, Length)
predictions = predictor.predict(waveforms)

数据加载

from seispolarity.data import WaveformDataset

# 大型数据集使用磁盘流式加载
dataset = WaveformDataset(path="data.hdf5", name="SCSN", preload=False)
loader = dataset.get_dataloader(batch_size=1024, num_workers=4)

# 小型数据集使用 RAM 预加载
dataset_ram = WaveformDataset(path="data.hdf5", name="SCSN", preload=True)

自动下载数据集

from seispolarity.data import PNW
from pathlib import Path

output_dir = Path('datasets/PNW')
processor = PNW(
    csv_path=str(output_dir / 'PNW.csv'),
    hdf5_path=str(output_dir / 'PNW.hdf5'),
    output_polarity=str(output_dir),
    auto_download=True,  # 自动下载缺失的数据
    use_hf=False,        # 使用 ModelScope 而不是 Hugging Face
    component='Z',       # 提取垂直分量
    sampling_rate=100    # 目标采样率 (Hz)
)
processor.process()

训练

from seispolarity.models.scsn import SCSN
from seispolarity.training import Trainer, TrainingConfig

model = SCSN(num_fm_classes=3)
config = TrainingConfig(
    batch_size=256,
    epochs=50,
    learning_rate=1e-4,
    device='cuda'  # 或 'cpu'
)
trainer = Trainer(model=model, dataset=dataset, config=config)
trainer.train()

数据增强

from seispolarity.generate import GenericGenerator
from seispolarity.generate.augmentation import (
    Demean, Normalize, RandomTimeShift,
    BandpassFilter, PolarityInversion
)

generator = GenericGenerator(dataset)
generator.add_augmentations([
    Demean(),
    Normalize(amp_norm_type="peak"),
    RandomTimeShift(max_shift=10),
    BandpassFilter(freqmin=1, freqmax=20),
    PolarityInversion(prob=0.5)
])
loader = generator.get_dataloader(batch_size=256)

支持的数据集

数据集	说明	样本数	自动下载
SCSN	南加州地震台网	100k+	是
Txed	得克萨斯地震数据集	50k+	是
DiTing	中国地震台网中心	80k+	否（需申请）
Instance	基于实例的数据集	30k+	是
PNW	太平洋西北地区	20k+	是

支持的模型

模型	输入长度	类别
Ross (SCSN)	400	3 (U/D/N)
Eqpolarity	600	2 (U/D)
DiTingMotion	128	3 (U/D/N)
CFM	160	2 (U/D)
RPNet	400	2 (U/D)
PolarCAP	64	2 (U/D)
APP	400	3 (U/D/N)

目录

文档内容

示例

examples/ 目录中包含完整的示例：

数据集 API 用法 - 数据集加载与使用示例
预测 API 用法 - 模型推理示例
模型 API 用法 - 模型结构与初始化示例
训练 API 用法 - 训练流程示例

引用

论文将后续发表，敬请期待。

链接

联系方式

如有问题或需要支持，请在 GitHub 上提交 issue 或联系： [chuanjun1978@gmail.com]