feature(model): model configuration

app.py

@@ -0,0 +1,357 @@
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, Iterable, List, Optional, Tuple
+
+import gradio as gr
+import numpy as np
+import pandas as pd
+import plotly.graph_objects as go
+import torch
+import yaml
+from safetensors.torch import load_file
+from sbi.neural_nets.factory import posterior_nn
+
+
+MODELS_DIR = Path(__file__).resolve().parent / "models"
+DISPERSION_CURVES_DIR = Path(__file__).resolve().parent / "disp_curves"
+DEFAULT_CURVE_NONE_LABEL = "Upload custom curve"
+
+
+@dataclass
+class LoadedModel:
+    name: str
+    sampler: "PosteriorSampler"
+
+
+class PosteriorSampler:
+    """Thin wrapper around the trained neural posterior for sampling."""
+
+    def __init__(self, weights_path: Path, config_path: Path, device: Optional[str] = None) -> None:
+        self.weights_path = weights_path
+        self.config = yaml.safe_load(config_path.read_text())
+
+        dataset_cfg = self.config.get("dataset", {})
+        model_cfg = self.config.get("model", {})
+        params_cfg = model_cfg.get("parameters", {})
+
+        self.context_dim = int(dataset_cfg["input_shape"])
+        self.theta_dim = int(dataset_cfg["output_shape"])
+
+        build_kwargs: Dict[str, int] = {}
+        for key in ("hidden_features", "num_transforms", "num_bins", "num_components"):
+            if key in params_cfg and params_cfg[key] is not None:
+                build_kwargs[key] = int(params_cfg[key])
+
+        density_estimator_builder = posterior_nn(
+            model=params_cfg.get("density_estimator", "nsf"),
+            z_score_theta=params_cfg.get("z_score_theta", "independent"),
+            z_score_x=params_cfg.get("z_score_x", "independent"),
+            **build_kwargs,
+        )
+
+        # Create a dummy network to load the trained parameters. The actual statistics
+        # (e.g. z-score buffers) are restored from the safetensors file.
+        theta_prototype = torch.zeros(2, self.theta_dim)
+        context_prototype = torch.zeros(2, self.context_dim)
+        net = density_estimator_builder(theta_prototype, context_prototype)
+
+        state_dict = load_file(str(weights_path))
+        net.load_state_dict(state_dict)
+        net.eval()
+
+        runtime_device = torch.device(device) if device else torch.device("cpu")
+        self.net = net.to(runtime_device)
+        self.device = runtime_device
+
+    def sample(self, context: np.ndarray, num_samples: int) -> np.ndarray:
+        with torch.no_grad():
+            context_tensor = torch.as_tensor(context, dtype=torch.float32, device=self.device).reshape(-1)
+            if context_tensor.numel() != self.context_dim:
+                raise ValueError(
+                    f"Expected context with {self.context_dim} elements, received {context_tensor.numel()}."
+                )
+            samples = self.net.sample((num_samples,), context=context_tensor)
+        samples_np = samples.cpu().numpy()
+        if samples_np.ndim == 3:
+            samples_np = samples_np[:, 0, :]
+        elif samples_np.ndim != 2:
+            raise ValueError(f"Unexpected sample shape {samples_np.shape}.")
+        return samples_np
+
+
+def discover_dispersion_curves(curves_dir: Path) -> Dict[str, Tuple[Path, Path]]:
+    if not curves_dir.exists():
+        return {}
+
+    discovered: Dict[str, Tuple[Path, Path]] = {}
+    for curve_path in sorted(curves_dir.glob("disp_curve_*.csv")):
+        suffix = curve_path.stem.split("disp_curve_")[-1]
+        theta_path = curves_dir / f"theta_{suffix}.csv"
+        display_name = f"Curve {suffix}"
+        discovered[display_name] = (curve_path, theta_path)
+
+    return discovered
+
+
+PREDEFINED_DISPERSION_CURVES = discover_dispersion_curves(DISPERSION_CURVES_DIR)
+
+
+def discover_models(models_dir: Path) -> List[LoadedModel]:
+    if not models_dir.exists():
+        raise FileNotFoundError(f"Expected models directory at {models_dir}")
+
+    discovered: List[LoadedModel] = []
+    for weights_path in sorted(models_dir.glob("*.safetensors")):
+        config_candidates = [
+            weights_path.with_suffix(".yaml"),
+            weights_path.with_suffix(".yml"),
+            models_dir / "config.yaml",
+        ]
+        config_path = next((path for path in config_candidates if path.exists()), None)
+        if not config_path:
+            raise FileNotFoundError(f"No configuration file found for {weights_path.name}")
+
+        sampler = PosteriorSampler(weights_path, config_path)
+        display_name = weights_path.stem.replace("_", " ").title()
+        discovered.append(LoadedModel(name=display_name, sampler=sampler))
+
+    if not discovered:
+        raise FileNotFoundError(f"No .safetensors models found in {models_dir}")
+
+    return discovered
+
+
+class ModelRegistry:
+    def __init__(self, models: Iterable[LoadedModel]):
+        self._registry: Dict[str, PosteriorSampler] = {}
+        for item in models:
+            if item.name in self._registry:
+                raise ValueError(f"Duplicate model name detected: {item.name}")
+            self._registry[item.name] = item.sampler
+
+    @property
+    def names(self) -> List[str]:
+        return list(self._registry.keys())
+
+    def get(self, name: str) -> PosteriorSampler:
+        if name not in self._registry:
+            raise KeyError(f"Unknown model '{name}'")
+        return self._registry[name]
+
+
+REGISTRY = ModelRegistry(discover_models(MODELS_DIR))
+DEFAULT_MODEL_NAME = REGISTRY.names[0]
+
+
+def load_predefined_dispersion_curve(name: str) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+    if name not in PREDEFINED_DISPERSION_CURVES:
+        raise gr.Error("Unknown dispersion curve selection.")
+
+    curve_path, theta_path = PREDEFINED_DISPERSION_CURVES[name]
+    if not curve_path.exists():
+        raise gr.Error(f"Unable to find dispersion curve file at {curve_path}.")
+    if not theta_path.exists():
+        raise gr.Error(f"Unable to find theta file at {theta_path}.")
+
+    curve_df = pd.read_csv(curve_path)
+    if curve_df.shape[1] < 2:
+        raise gr.Error(f"Dispersion curve file {curve_path.name} must contain period and vg columns.")
+
+    periods = pd.to_numeric(curve_df.iloc[:, 0], errors="coerce").to_numpy(dtype=np.float32)
+    vg_values = pd.to_numeric(curve_df.iloc[:, 1], errors="coerce").to_numpy(dtype=np.float32)
+
+    theta_df = pd.read_csv(theta_path)
+    theta_values = pd.to_numeric(theta_df.to_numpy().reshape(-1), errors="coerce").astype(np.float32)
+    theta_values = theta_values[~np.isnan(theta_values)]
+
+    if periods.size != vg_values.size:
+        raise gr.Error(
+            f"Dispersion curve file {curve_path.name} contains mismatched period and vg counts."
+        )
+
+    if np.isnan(periods).any() or np.isnan(vg_values).any():
+        raise gr.Error(f"Dispersion curve file {curve_path.name} contains non-numeric entries.")
+
+    return periods, vg_values, theta_values
+
+
+def read_dispersion_curve(upload: Optional[Any], expected_length: int) -> np.ndarray:
+    if upload is None:
+        raise gr.Error("Please upload a CSV file containing the dispersion curve.")
+
+    try:
+        df = pd.read_csv(upload.name, header=None)
+    except Exception as exc:  # pylint: disable=broad-except
+        raise gr.Error(f"Unable to read CSV file: {exc}") from exc
+
+    numeric_values = pd.to_numeric(df.to_numpy().reshape(-1), errors="coerce").astype(np.float32)
+    numeric_values = numeric_values[~np.isnan(numeric_values)]
+
+    if numeric_values.size != expected_length:
+        raise gr.Error(
+            f"Expected {expected_length} values in the dispersion curve, but found {numeric_values.size}. "
+            "Please provide a CSV with exactly one value per frequency sample."
+        )
+
+    return numeric_values
+
+
+def build_plot(samples: np.ndarray) -> go.Figure:
+    depth_axis = np.arange(1, samples.shape[1] + 1)
+    fig = go.Figure()
+    for idx, sample in enumerate(samples, start=1):
+        fig.add_trace(
+            go.Scatter(
+                x=depth_axis,
+                y=sample,
+                mode="lines",
+                name=f"Sample {idx}",
+            )
+        )
+
+    fig.update_layout(
+        xaxis_title="Layer index",
+        yaxis_title="Velocity",
+        legend_title="Generated samples",
+        template="plotly_white",
+        margin=dict(l=40, r=10, t=40, b=40),
+    )
+    return fig
+
+
+def build_dispersion_plot(periods: np.ndarray, group_velocities: np.ndarray) -> go.Figure:
+    fig = go.Figure()
+    fig.add_trace(
+        go.Scatter(
+            x=periods,
+            y=group_velocities,
+            mode="lines+markers",
+            name="Dispersion curve",
+        )
+    )
+    fig.update_layout(
+        xaxis_title="Period",
+        yaxis_title="Group velocity",
+        template="plotly_white",
+        margin=dict(l=40, r=10, t=40, b=40),
+        showlegend=False,
+    )
+    return fig
+
+
+def handle_predefined_curve_selection(selection: Optional[str]) -> Tuple[Any, Optional[np.ndarray], Optional[np.ndarray]]:
+    if not selection or selection == DEFAULT_CURVE_NONE_LABEL:
+        return gr.update(value=None), None, None
+
+    periods, vg_values, theta_values = load_predefined_dispersion_curve(selection)
+    figure = build_dispersion_plot(periods, vg_values)
+    return figure, vg_values, theta_values
+
+
+def format_samples(samples: np.ndarray) -> pd.DataFrame:
+    index = [f"Layer {i}" for i in range(1, samples.shape[1] + 1)]
+    columns = [f"Sample {idx}" for idx in range(1, samples.shape[0] + 1)]
+    return pd.DataFrame(samples.T, index=index, columns=columns)
+
+
+def generate_velocity_models(
+    upload: Optional[Any],
+    model_name: str,
+    num_samples: int,
+    predefined_curve_name: Optional[str],
+    predefined_vg: Optional[np.ndarray],
+    _preloaded_theta: Optional[np.ndarray],
+) -> Tuple[go.Figure, pd.DataFrame]:
+    sampler = REGISTRY.get(model_name)
+    dispersion_curve: Optional[np.ndarray] = None
+
+    if predefined_curve_name and predefined_curve_name != DEFAULT_CURVE_NONE_LABEL:
+        if predefined_vg is None:
+            # Reload from disk if the state is empty for any reason.
+            _, vg_values, _ = load_predefined_dispersion_curve(predefined_curve_name)
+            predefined_vg = vg_values
+        dispersion_curve = np.asarray(predefined_vg, dtype=np.float32)
+    else:
+        dispersion_curve = read_dispersion_curve(upload, sampler.context_dim)
+
+    if dispersion_curve.size != sampler.context_dim:
+        raise gr.Error(
+            f"The selected dispersion curve contains {dispersion_curve.size} samples, "
+            f"but the posterior expects {sampler.context_dim}."
+        )
+
+    samples = sampler.sample(dispersion_curve, int(num_samples))
+
+    return build_plot(samples), format_samples(samples)
+
+
+with gr.Blocks(title="Surface Wave Inversion with NPE") as demo:
+    default_curve_choices = [DEFAULT_CURVE_NONE_LABEL] + list(PREDEFINED_DISPERSION_CURVES.keys())
+    selected_vg_state = gr.State(value=None)
+    selected_theta_state = gr.State(value=None)
+
+    gr.Markdown(
+        "## Neural Posterior Estimation for Surface Wave Inversion\n"
+        "Select a built-in dispersion curve or upload your own, then choose a pretrained posterior model "
+        "to draw samples of the subsurface velocity structure."
+    )
+
+    with gr.Row():
+        with gr.Column(scale=1):
+            default_curve_choice = gr.Dropdown(
+                label="Default dispersion curve",
+                choices=default_curve_choices,
+                value=DEFAULT_CURVE_NONE_LABEL,
+                interactive=len(default_curve_choices) > 1,
+                info="Pick a built-in curve or stay on Upload custom curve to provide your own file.",
+            )
+            curve_input = gr.File(
+                label="Dispersion curve (.csv)",
+                file_types=[".csv"],
+            )
+            model_choice = gr.Dropdown(
+                label="Posterior model",
+                choices=REGISTRY.names,
+                value=DEFAULT_MODEL_NAME,
+            )
+            sample_count = gr.Slider(
+                label="Number of samples",
+                minimum=1,
+                maximum=200,
+                value=20,
+                step=1,
+            )
+            generate_btn = gr.Button("Generate velocity models", variant="primary")
+
+        with gr.Column(scale=1):
+            dispersion_plot = gr.Plot(label="Selected dispersion curve")
+            plot_output = gr.Plot(label="Sampled velocity profiles")
+            table_output = gr.Dataframe(
+                headers=[f"Sample {idx}" for idx in range(1, 6)],
+                datatype="number",
+                interactive=False,
+                label="Sample values",
+            )
+
+    default_curve_choice.change(
+        handle_predefined_curve_selection,
+        inputs=default_curve_choice,
+        outputs=[dispersion_plot, selected_vg_state, selected_theta_state],
+    )
+
+    generate_btn.click(
+        generate_velocity_models,
+        inputs=[
+            curve_input,
+            model_choice,
+            sample_count,
+            default_curve_choice,
+            selected_vg_state,
+            selected_theta_state,
+        ],
+        outputs=[plot_output, table_output],
+    )
+
+
+if __name__ == "__main__":
+    demo.launch()

disp_curves/disp_curve_01.csv

@@ -0,0 +1,109 @@
+period,vg
+1.0000,1.4450
+1.0374,1.4460
+1.0748,1.4469
+1.1121,1.4476
+1.1495,1.4482
+1.1869,1.4487
+1.2243,1.4490
+1.2617,1.4493
+1.2991,1.4494
+1.3364,1.4495
+1.3738,1.4494
+1.4112,1.4494
+1.4486,1.4497
+1.4860,1.4501
+1.5234,1.4504
+1.5607,1.4509
+1.5981,1.4516
+1.6355,1.4525
+1.6729,1.4536
+1.7103,1.4553
+1.7477,1.4571
+1.7850,1.4596
+1.8224,1.4624
+1.8598,1.4663
+1.8972,1.4709
+1.9346,1.4767
+1.9720,1.4842
+2.0093,1.4929
+2.0467,1.5059
+2.0841,1.5201
+2.1215,1.5466
+2.1589,1.5816
+2.1963,1.6574
+2.2336,1.6195
+2.2710,1.3761
+2.3084,1.2702
+2.3458,1.2804
+2.3832,1.2859
+2.4206,1.2897
+2.4579,1.2911
+2.4953,1.2922
+2.5327,1.2921
+2.5701,1.2918
+2.6075,1.2909
+2.6449,1.2899
+2.6822,1.2886
+2.7196,1.2871
+2.7570,1.2855
+2.7944,1.2837
+2.8318,1.2819
+2.8692,1.2799
+2.9065,1.2780
+2.9439,1.2759
+2.9813,1.2739
+3.0187,1.2718
+3.0561,1.2696
+3.0935,1.2675
+3.1308,1.2654
+3.1682,1.2633
+3.2056,1.2611
+3.2430,1.2590
+3.2804,1.2569
+3.3178,1.2548
+3.3551,1.2527
+3.3925,1.2506
+3.4299,1.2485
+3.4673,1.2465
+3.5047,1.2444
+3.5421,1.2424
+3.5794,1.2404
+3.6168,1.2385
+3.6542,1.2365
+3.6916,1.2346
+3.7290,1.2327
+3.7664,1.2308
+3.8037,1.2290
+3.8411,1.2272
+3.8785,1.2254
+3.9159,1.2237
+3.9533,1.2219
+3.9907,1.2202
+4.0280,1.2185
+4.0654,1.2168
+4.1028,1.2152
+4.1402,1.2135
+4.1776,1.2119
+4.2149,1.2103
+4.2523,1.2088
+4.2897,1.2073
+4.3271,1.2058
+4.3645,1.2043
+4.4019,1.2029
+4.4393,1.2014
+4.4766,1.2000
+4.5140,1.1987
+4.5514,1.1973
+4.5888,1.1960
+4.6262,1.1947
+4.6636,1.1934
+4.7009,1.1922
+4.7383,1.1909
+4.7757,1.1897
+4.8131,1.1886
+4.8505,1.1874
+4.8879,1.1862
+4.9252,1.1851
+4.9626,1.1840
+5.0000,1.1830

disp_curves/theta_01.csv

@@ -0,0 +1,11 @@
+theta
+2.0000
+1.7500
+0.9000
+1.2000
+0.3400
+0.0000
+1.5000
+1.2000
+2.5000
+3.5000

requirements.txt

@@ -0,0 +1,7 @@
+gradio==5.49.1
+pandas
+plotly
+pyyaml
+safetensors
+sbi
+torch