update app.py

app.py

@@ -0,0 +1,162 @@
+import os
+from json import dumps
+from typing import Dict, List, Union
+
+import numpy as np
+import pandas as pd
+from fastapi import FastAPI
+from kafka import KafkaProducer
+from pydantic import BaseModel
+from pyproj import Proj
+
+from gamma.utils import association
+
+app = FastAPI()
+
+
+@app.get("/")
+def greet_json():
+    return {"message": "Hello, World!"}
+
+
+@app.post("/predict/")
+def predict(picks: dict, stations: dict, config: dict):
+    picks = picks["data"]
+    stations = stations["data"]
+    picks = pd.DataFrame(picks)
+    picks["phase_time"] = pd.to_datetime(picks["phase_time"])
+    stations = pd.DataFrame(stations)
+    print(stations)
+    events_, picks_ = run_gamma(picks, stations, config)
+    picks_ = picks_.to_dict(orient="records")
+    events_ = events_.to_dict(orient="records")
+
+    return {"picks": picks_, "events": events_}
+
+
+def set_config(region="ridgecrest"):
+
+    config = {
+        "min_picks": 8,
+        "min_picks_ratio": 0.2,
+        "max_residual_time": 1.0,
+        "max_residual_amplitude": 1.0,
+        "min_score": 0.6,
+        "min_s_picks": 2,
+        "min_p_picks": 2,
+        "use_amplitude": False,
+    }
+
+    # ## Domain
+    if region.lower() == "ridgecrest":
+        config.update(
+            {
+                "region": "ridgecrest",
+                "minlongitude": -118.004,
+                "maxlongitude": -117.004,
+                "minlatitude": 35.205,
+                "maxlatitude": 36.205,
+                "mindepth_km": 0.0,
+                "maxdepth_km": 30.0,
+            }
+        )
+
+    lon0 = (config["minlongitude"] + config["maxlongitude"]) / 2
+    lat0 = (config["minlatitude"] + config["maxlatitude"]) / 2
+    proj = Proj(f"+proj=sterea +lon_0={lon0} +lat_0={lat0}  +units=km")
+    xmin, ymin = proj(config["minlongitude"], config["minlatitude"])
+    xmax, ymax = proj(config["maxlongitude"], config["maxlatitude"])
+    zmin, zmax = config["mindepth_km"], config["maxdepth_km"]
+    xlim_km = (xmin, xmax)
+    ylim_km = (ymin, ymax)
+    zlim_km = (zmin, zmax)
+
+    config.update(
+        {
+            "xlim_km": xlim_km,
+            "ylim_km": ylim_km,
+            "zlim_km": zlim_km,
+            "proj": proj,
+        }
+    )
+
+    config.update(
+        {
+            "min_picks_per_eq": 5,
+            "min_p_picks_per_eq": 0,
+            "min_s_picks_per_eq": 0,
+            "max_sigma11": 3.0,
+            "max_sigma22": 1.0,
+            "max_sigma12": 1.0,
+        }
+    )
+
+    config["use_dbscan"] = False
+    config["use_amplitude"] = True
+    config["oversample_factor"] = 8.0
+    config["dims"] = ["x(km)", "y(km)", "z(km)"]
+    config["method"] = "BGMM"
+    config["ncpu"] = 1
+    vel = {"p": 6.0, "s": 6.0 / 1.75}
+    config["vel"] = vel
+
+    config["bfgs_bounds"] = (
+        (xlim_km[0] - 1, xlim_km[1] + 1),  # x
+        (ylim_km[0] - 1, ylim_km[1] + 1),  # y
+        (0, zlim_km[1] + 1),  # z
+        (None, None),  # t
+    )
+
+    config["event_index"] = 0
+
+    return config
+
+
+config = set_config()
+
+
+def run_gamma(picks, stations, config_):
+
+    # %%
+    config.update(config_)
+
+    proj = config["proj"]
+
+    picks = picks.rename(
+        columns={
+            "station_id": "id",
+            "phase_time": "timestamp",
+            "phase_type": "type",
+            "phase_score": "prob",
+            "phase_amplitude": "amp",
+        }
+    )
+    stations[["x(km)", "y(km)"]] = stations.apply(
+        lambda x: pd.Series(proj(longitude=x.longitude, latitude=x.latitude)), axis=1
+    )
+    stations["z(km)"] = stations["elevation_m"].apply(lambda x: -x / 1e3)
+    stations = stations.rename(columns={"station_id": "id"})
+
+    events, assignments = association(picks, stations, config, 0, config["method"])
+
+    print(events)
+    events = pd.DataFrame(events)
+    events[["longitude", "latitude"]] = events.apply(
+        lambda x: pd.Series(proj(longitude=x["x(km)"], latitude=x["y(km)"], inverse=True)), axis=1
+    )
+    events["depth_km"] = events["z(km)"]
+    events.drop(columns=["x(km)", "y(km)", "z(km)"], inplace=True, errors="ignore")
+    picks = picks.rename(
+        columns={
+            "id": "station_id",
+            "timestamp": "phase_time",
+            "type": "phase_type",
+            "prob": "phase_score",
+            "amp": "phase_amplitude",
+        }
+    )
+
+    assignments = pd.DataFrame(assignments, columns=["pick_index", "event_index", "gamma_score"])
+    picks = picks.join(assignments.set_index("pick_index")).fillna(-1).astype({"event_index": int})
+
+    return events, picks