update app.py

app.py

@@ -5,165 +5,158 @@ from typing import Dict, List, Union
 import numpy as np
 import pandas as pd
 from fastapi import FastAPI
-from gamma.utils import association
+from kafka import KafkaProducer
 from pydantic import BaseModel
 from pyproj import Proj
 
-app = FastAPI()
+from gamma.utils import association
 
-PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), ".."))
-# STATION_CSV = os.path.join(PROJECT_ROOT, "tests/stations_hawaii.csv")
-# STATION_CSV = os.path.join(PROJECT_ROOT, "tests/stations.csv")  ## ridgecrest
-
-
-def default_config(config):
-    if "degree2km" not in config:
-        config["degree2km"] = 111.195
-    if "use_amplitude" not in config:
-        config["use_amplitude"] = True
-    if "use_dbscan" not in config:
-        config["use_dbscan"] = True
-    if "dbscan_eps" not in config:
-        config["dbscan_eps"] = 30.0
-    if "dbscan_min_samples" not in config:
-        config["dbscan_min_samples"] = 3
-    if "method" not in config:
-        config["method"] = "BGMM"
-    if "oversample_factor" not in config:
-        config["oversample_factor"] = 5
-    if "min_picks_per_eq" not in config:
-        config["min_picks_per_eq"] = 10
-    if "max_sigma11" not in config:
-        config["max_sigma11"] = 2.0
-    if "max_sigma22" not in config:
-        config["max_sigma22"] = 1.0
-    if "max_sigma12" not in config:
-        config["max_sigma12"] = 1.0
-    if "dims" not in config:
-        config["dims"] = ["x(km)", "y(km)", "z(km)"]
-    return config
+app = FastAPI()
 
 
-## set config
-config = {"xlim_degree": [-156.32, -154.32], "ylim_degree": [18.39, 20.39], "z(km)": [0, 41]}  ## hawaii
-# config = {'xlim_degree': [-118.004, -117.004], 'ylim_degree': [35.205, 36.205], "z(km)": [0, 41]}  ## ridgecrest
+@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,
+            }
+        )
 
-config = default_config(config)
-config["center"] = [np.mean(config["xlim_degree"]), np.mean(config["ylim_degree"])]
-config["x(km)"] = (np.array(config["xlim_degree"]) - config["center"][0]) * config["degree2km"]
-config["y(km)"] = (np.array(config["ylim_degree"]) - config["center"][1]) * config["degree2km"]
-config["bfgs_bounds"] = [list(config[x]) for x in config["dims"]] + [[None, None]]
+    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,
+        }
+    )
 
-for k, v in config.items():
-    print(f"{k}: {v}")
+    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,
+        }
+    )
 
-# ## read stations
-# stations = pd.read_csv(STATION_CSV, delimiter="\t")
-# stations = stations.rename(columns={"station": "id"})
-# stations["x(km)"] = stations["longitude"].apply(lambda x: (x - config["center"][0]) * config["degree2km"])
-# stations["y(km)"] = stations["latitude"].apply(lambda x: (x - config["center"][1]) * config["degree2km"])
-# stations["z(km)"] = stations["elevation(m)"].apply(lambda x: -x / 1e3)
+    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
+    )
 
-# print(stations)
+    config["event_index"] = 0
 
+    return config
 
-class Data(BaseModel):
-    picks: List[Dict[str, Union[float, str]]]
-    stations: List[Dict[str, Union[float, str]]]
-    config: Dict[str, Union[List[float], List[int], List[str], float, int, str]]
 
+config = set_config()
 
-class Pick(BaseModel):
-    picks: List[Dict[str, Union[float, str]]]
 
+def run_gamma(picks, stations, config_):
 
-def run_gamma(picks, config, stations):
+    # %%
+    config.update(config_)
 
-    proj = Proj(f"+proj=sterea +lon_0={config['center'][0]} +lat_0={config['center'][1]} +units=km")
+    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)
-
-    print(f"{len(picks)} picks, {len(stations)} stations")
-    catalogs, assignments = association(picks, stations, config, 0, config["method"])
+    stations["z(km)"] = stations["elevation_m"].apply(lambda x: -x / 1e3)
+    stations = stations.rename(columns={"station_id": "id"})
 
-    catalogs = pd.DataFrame(
-        catalogs,
-        columns=["time"]
-        + config["dims"]
-        + ["magnitude", "sigma_time", "sigma_amp", "cov_time_amp", "event_index", "gamma_score"],
-    )
-    if len(catalogs) == 0:
-        print("No events associated")
-        return pd.DataFrame(), pd.DataFrame()
+    events, assignments = association(picks, stations, config, 0, config["method"])
 
-    catalogs[["longitude", "latitude"]] = catalogs.apply(
+    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
     )
-    catalogs["depth(m)"] = catalogs["z(km)"].apply(lambda x: x * 1e3)
+    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_gamma = picks.join(assignments.set_index("pick_index")).fillna(-1).astype({"event_index": int})
-
-    return catalogs, picks_gamma
-
-
-# @app.post("/predict_stream")
-# def predict(data: Pick):
-
-#     picks = pd.DataFrame(data.picks)
-#     if len(picks) == 0:
-#         return {"catalog": [], "picks": []}
-
-#     catalogs, picks_gamma = run_gamma(data, config, stations)
-
-#     return {"catalog": catalogs.to_dict(orient="records"), "picks": picks_gamma.to_dict(orient="records")}
-
-
-@app.post("/predict")
-def predict(data: Data):
-
-    picks = pd.DataFrame(data.picks)
-    if len(picks) == 0:
-        return {"catalog": [], "picks": []}
-
-    stations = pd.DataFrame(data.stations)
-    if len(stations) == 0:
-        return {"catalog": [], "picks": []}
-
-    assert "latitude" in stations
-    assert "longitude" in stations
-    assert "elevation(m)" in stations
-
-    config = data.config
-    config = default_config(config)
-
-    if "xlim_degree" not in config:
-        config["xlim_degree"] = (stations["longitude"].min(), stations["longitude"].max())
-    if "ylim_degree" not in config:
-        config["ylim_degree"] = (stations["latitude"].min(), stations["latitude"].max())
-    if "center" not in config:
-        config["center"] = [np.mean(config["xlim_degree"]), np.mean(config["ylim_degree"])]
-    if "x(km)" not in config:
-        config["x(km)"] = (
-            (np.array(config["xlim_degree"]) - config["center"][0])
-            * config["degree2km"]
-            * np.cos(np.deg2rad(config["center"][1]))
-        )
-    if "y(km)" not in config:
-        config["y(km)"] = (np.array(config["ylim_degree"]) - config["center"][1]) * config["degree2km"]
-    if "z(km)" not in config:
-        config["z(km)"] = (0, 41)
-    if "bfgs_bounds" not in config:
-        config["bfgs_bounds"] = [list(config[x]) for x in config["dims"]] + [[None, None]]
-
-    catalogs, picks_gamma = run_gamma(picks, config, stations)
-
-    return {"catalog": catalogs.to_dict(orient="records"), "picks": picks_gamma.to_dict(orient="records")}
-
+    picks = picks.join(assignments.set_index("pick_index")).fillna(-1).astype({"event_index": int})
 
-@app.get("/healthz")
-def healthz():
-    return {"status": "ok"}
+    return events, picks