diff --git a/INSTRUCTIONS.md b/INSTRUCTIONS.md index 91823e39..df1f227f 100644 --- a/INSTRUCTIONS.md +++ b/INSTRUCTIONS.md @@ -12,6 +12,7 @@ This directory contains the MCP servers and infrastructure for the AssetOpsBench - [Utilities](#utilities) - [FMSRAgent](#fmsragent) - [TSFMAgent](#tsfmagent) + - [VibrationAgent](#vibrationagent) - [Plan-Execute Runner](#plan-execute-runner) - [How it works](#how-it-works) - [CLI](#cli) @@ -158,6 +159,23 @@ uv run tsfm-mcp-server | `run_tsad` | `dataset_path`, `tsfm_output_json`, `timestamp_column`, `target_columns`, `task?`, `false_alarm?`, `ad_model_type?`, ... | Conformal anomaly detection on top of a forecasting output JSON; returns CSV with anomaly labels | | `run_integrated_tsad` | `dataset_path`, `timestamp_column`, `target_columns`, `model_checkpoint?`, `false_alarm?`, `n_calibration?`, ... | End-to-end forecasting + anomaly detection in one call; returns combined CSV | +### VibrationAgent + +**Path:** `src/servers/vibration/main.py` +**Requires:** CouchDB (same env vars as IoTAgent: `COUCHDB_URL`, `COUCHDB_DBNAME`, `COUCHDB_USERNAME`, `COUCHDB_PASSWORD`); `numpy`, `scipy` +**DSP core:** `src/servers/vibration/dsp/` — adapted from [vibration-analysis-mcp](https://github.com/LGDiMaggio/claude-stwinbox-diagnostics/tree/main/mcp-servers/vibration-analysis-mcp) (Apache-2.0) + +| Tool | Arguments | Description | +|---|---|---| +| `get_vibration_data` | `site_name`, `asset_id`, `sensor_name`, `start`, `final?` | Fetch vibration time-series from CouchDB and load into the analysis store. Returns a `data_id`. | +| `list_vibration_sensors` | `site_name`, `asset_id` | List available sensor fields for an asset. | +| `compute_fft_spectrum` | `data_id`, `window?`, `top_n?` | Compute FFT amplitude spectrum (top-N peaks + statistics). | +| `compute_envelope_spectrum` | `data_id`, `band_low_hz?`, `band_high_hz?`, `top_n?` | Compute envelope spectrum for bearing fault detection (Hilbert transform). | +| `assess_vibration_severity` | `rms_velocity_mm_s`, `machine_group?` | Classify vibration severity per ISO 10816 (Zones A–D). | +| `calculate_bearing_frequencies` | `rpm`, `n_balls`, `ball_diameter_mm`, `pitch_diameter_mm`, `contact_angle_deg?`, `bearing_name?` | Compute bearing characteristic frequencies (BPFO, BPFI, BSF, FTF). | +| `list_known_bearings` | — | List all bearings in the built-in database. | +| `diagnose_vibration` | `data_id`, `rpm?`, `bearing_designation?`, `bearing_*?`, `bpfo_hz?`, `bpfi_hz?`, `bsf_hz?`, `ftf_hz?`, `machine_group?`, `machine_description?` | Full automated diagnosis: FFT + shaft features + bearing envelope + ISO 10816 + fault classification + markdown report. | + --- ## Plan-Execute Runner @@ -225,7 +243,7 @@ uv run plan-execute --show-history --json "How many observations exist for CH-1? ### End-to-end example -All four servers (IoTAgent, Utilities, FMSRAgent, TSFMAgent) are registered by default. +All five servers (IoTAgent, Utilities, FMSRAgent, TSFMAgent, VibrationAgent) are registered by default. Run a question that exercises three of them with independent parallel steps: ```bash @@ -304,6 +322,7 @@ runner = PlanExecuteRunner( "Utilities": "utilities-mcp-server", "FMSRAgent": "fmsr-mcp-server", "TSFMAgent": "tsfm-mcp-server", + "VibrationAgent": "vibration-mcp-server", }, ) ``` @@ -334,6 +353,10 @@ Add the following to your Claude Desktop `claude_desktop_config.json`: "TSFMAgent": { "command": "/path/to/uv", "args": ["run", "--project", "/path/to/AssetOpsBench", "tsfm-mcp-server"] + }, + "VibrationAgent": { + "command": "/path/to/uv", + "args": ["run", "--project", "/path/to/AssetOpsBench", "vibration-mcp-server"] } } } @@ -396,8 +419,8 @@ uv run pytest src/ -v │ │ stdio │ │ └───────────────────┼────────────┼─────────────────────┘ │ MCP protocol (stdio) - ┌──────────┼──────────┬──────────┐ - ▼ ▼ ▼ ▼ - IoTAgent Utilities FMSRAgent TSFMAgent - (tools) (tools) (tools) (tools) + ┌──────────┼──────────┬──────────┬───────────────┐ + ▼ ▼ ▼ ▼ ▼ + IoTAgent Utilities FMSRAgent TSFMAgent VibrationAgent + (tools) (tools) (tools) (tools) (tools) ``` diff --git a/pyproject.toml b/pyproject.toml index 905ed17d..04e87eb9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,7 @@ dependencies = [ "pyyaml>=6.0", "litellm>=1.0", "python-dotenv>=1.0", + "scipy>=1.10.0", ] [project.scripts] @@ -32,6 +33,7 @@ iot-mcp-server = "servers.iot.main:main" utilities-mcp-server = "servers.utilities.main:main" fmsr-mcp-server = "servers.fmsr.main:main" tsfm-mcp-server = "servers.tsfm.main:main" +vibration-mcp-server = "servers.vibration.main:main" [dependency-groups] diff --git a/src/servers/vibration/__init__.py b/src/servers/vibration/__init__.py new file mode 100644 index 00000000..1563996a --- /dev/null +++ b/src/servers/vibration/__init__.py @@ -0,0 +1,5 @@ +# Vibration Analysis MCP Server for AssetOpsBench +# +# DSP core adapted from vibration-analysis-mcp by Luigi Di Maggio (LGDiMaggio) +# https://github.com/LGDiMaggio/claude-stwinbox-diagnostics +# SPDX-License-Identifier: Apache-2.0 diff --git a/src/servers/vibration/couchdb_client.py b/src/servers/vibration/couchdb_client.py new file mode 100644 index 00000000..897474ae --- /dev/null +++ b/src/servers/vibration/couchdb_client.py @@ -0,0 +1,139 @@ +"""CouchDB client for fetching vibration sensor data. + +Reuses the same environment variables as the IoT MCP server +(COUCHDB_URL, COUCHDB_DBNAME, COUCHDB_USERNAME, COUCHDB_PASSWORD). +""" + +from __future__ import annotations + +import logging +import os +from datetime import datetime +from typing import Optional + +import couchdb3 +import numpy as np +from dotenv import load_dotenv + +load_dotenv() +logger = logging.getLogger("vibration-mcp-server") + +COUCHDB_URL = os.environ.get("COUCHDB_URL") +COUCHDB_DBNAME = os.environ.get("COUCHDB_DBNAME") +COUCHDB_USER = os.environ.get("COUCHDB_USERNAME") +COUCHDB_PASSWORD = os.environ.get("COUCHDB_PASSWORD") + + +def _get_db() -> Optional[couchdb3.Database]: + """Lazy CouchDB connection with error handling.""" + if not COUCHDB_URL: + logger.warning("COUCHDB_URL not set — vibration data from CouchDB unavailable") + return None + try: + return couchdb3.Database( + COUCHDB_DBNAME, + url=COUCHDB_URL, + user=COUCHDB_USER, + password=COUCHDB_PASSWORD, + ) + except Exception as e: + logger.error(f"CouchDB connection failed: {e}") + return None + + +def fetch_vibration_timeseries( + asset_id: str, + sensor_name: str, + start: str, + final: Optional[str] = None, + limit: int = 10000, +) -> Optional[tuple[np.ndarray, float]]: + """ + Fetch sensor time-series from CouchDB and return as numpy array. + + Queries CouchDB for documents matching the given asset_id and time range, + extracts values from the specified sensor column, and estimates the + sample rate from the timestamp spacing. + + Args: + asset_id: Asset identifier (e.g., 'Chiller 6'). + sensor_name: Name of the sensor field in CouchDB documents. + start: ISO 8601 start timestamp. + final: Optional ISO 8601 end timestamp. + limit: Maximum number of documents to fetch. + + Returns: + (signal_array, estimated_sample_rate) or None on error. + """ + db = _get_db() + if not db: + return None + + try: + selector: dict = { + "asset_id": asset_id, + "timestamp": {"$gte": datetime.fromisoformat(start).isoformat()}, + } + if final: + selector["timestamp"]["$lt"] = datetime.fromisoformat(final).isoformat() + + res = db.find( + selector, + limit=limit, + sort=[{"asset_id": "asc"}, {"timestamp": "asc"}], + ) + except Exception as e: + logger.error(f"CouchDB query failed: {e}") + return None + + docs = res.get("docs", []) + if not docs: + logger.info(f"No documents found for {asset_id}/{sensor_name}") + return None + + # Extract single sensor column + values: list[float] = [] + timestamps: list[str] = [] + for doc in docs: + if sensor_name in doc and "timestamp" in doc: + try: + values.append(float(doc[sensor_name])) + timestamps.append(doc["timestamp"]) + except (ValueError, TypeError): + continue + + if len(values) < 2: + logger.info( + f"Insufficient data points ({len(values)}) for {asset_id}/{sensor_name}" + ) + return None + + signal = np.array(values, dtype=np.float64) + + # Estimate sample rate from timestamp differences + try: + ts = [datetime.fromisoformat(t) for t in timestamps] + diffs = [(ts[i + 1] - ts[i]).total_seconds() for i in range(len(ts) - 1)] + avg_dt = sum(diffs) / len(diffs) + sample_rate = 1.0 / avg_dt if avg_dt > 0 else 1.0 + except Exception: + sample_rate = 1.0 # fallback: 1 Hz + + return signal, sample_rate + + +def list_sensor_fields(asset_id: str) -> list[str]: + """Return the sensor field names available for an asset in CouchDB.""" + db = _get_db() + if not db: + return [] + try: + res = db.find({"asset_id": asset_id}, limit=1) + if not res["docs"]: + return [] + doc = res["docs"][0] + exclude = {"_id", "_rev", "asset_id", "timestamp"} + return sorted(k for k in doc.keys() if k not in exclude) + except Exception as e: + logger.error(f"Error listing sensors for {asset_id}: {e}") + return [] diff --git a/src/servers/vibration/data_store.py b/src/servers/vibration/data_store.py new file mode 100644 index 00000000..20e0665b --- /dev/null +++ b/src/servers/vibration/data_store.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/LGDiMaggio/claude-stwinbox-diagnostics/tree/main/mcp-servers/vibration-analysis-mcp +""" +Server-side data store for vibration signals. + +Signals are kept in memory on the MCP server so that large numpy arrays +never need to transit through the LLM conversation context. The agent +only sees compact summaries (statistics, peak lists, diagnosis reports). +""" + +from __future__ import annotations + +import hashlib +import time +from dataclasses import dataclass, field + +import numpy as np +from numpy.typing import NDArray + + +def _kurtosis(x: NDArray) -> float: + """Excess kurtosis (Fisher definition, normal = 0).""" + n = len(x) + if n < 4: + return 0.0 + m = np.mean(x) + s = np.std(x, ddof=1) + if s < 1e-15: + return 0.0 + return float(np.mean(((x - m) / s) ** 4) - 3.0) + + +@dataclass +class DataEntry: + """A stored signal with metadata.""" + + signal: NDArray[np.floating] + sample_rate: float + created_at: float = field(default_factory=time.time) + metadata: dict = field(default_factory=dict) + + @property + def n_samples(self) -> int: + return self.signal.shape[0] + + @property + def n_channels(self) -> int: + return self.signal.shape[1] if self.signal.ndim > 1 else 1 + + @property + def duration_s(self) -> float: + return self.n_samples / self.sample_rate if self.sample_rate > 0 else 0 + + def summary(self) -> dict: + """Return a compact summary (no raw data).""" + sig = self.signal + if sig.ndim == 1: + sig = sig.reshape(-1, 1) + + channel_stats = {} + labels = self.metadata.get( + "axis_labels", + ["X", "Y", "Z"][: sig.shape[1]] + if sig.shape[1] <= 3 + else [f"CH{i}" for i in range(sig.shape[1])], + ) + for i, label in enumerate(labels): + if i >= sig.shape[1]: + break + col = sig[:, i] + rms = float(np.sqrt(np.mean(col**2))) + channel_stats[label] = { + "mean": round(float(np.mean(col)), 6), + "std": round(float(np.std(col)), 6), + "rms": round(rms, 6), + "peak": round(float(np.max(np.abs(col))), 6), + "crest_factor": round(float(np.max(np.abs(col)) / rms), 2) + if rms > 0 + else 0, + "kurtosis": round(float(_kurtosis(col)), 2), + } + + return { + "n_samples": self.n_samples, + "n_channels": self.n_channels, + "sample_rate_hz": self.sample_rate, + "duration_s": round(self.duration_s, 4), + "channel_stats": channel_stats, + "metadata": { + k: v for k, v in self.metadata.items() if k != "axis_labels" + }, + } + + +class DataStore: + """Simple in-memory store for vibration signals.""" + + def __init__(self) -> None: + self._entries: dict[str, DataEntry] = {} + + def put( + self, + data_id: str, + signal: NDArray[np.floating], + sample_rate: float, + metadata: dict | None = None, + ) -> str: + """Store a signal. Returns the data_id.""" + self._entries[data_id] = DataEntry( + signal=np.asarray(signal, dtype=np.float64), + sample_rate=sample_rate, + metadata=metadata or {}, + ) + return data_id + + def put_auto( + self, + signal: NDArray[np.floating], + sample_rate: float, + metadata: dict | None = None, + ) -> str: + """Store a signal with an auto-generated ID.""" + h = hashlib.md5(signal.tobytes()[:1024]).hexdigest()[:8] + data_id = f"sig_{h}_{int(time.time()) % 100000}" + return self.put(data_id, signal, sample_rate, metadata) + + def get(self, data_id: str) -> DataEntry | None: + return self._entries.get(data_id) + + def remove(self, data_id: str) -> bool: + return self._entries.pop(data_id, None) is not None + + def list_ids(self) -> list[str]: + return list(self._entries.keys()) + + def list_entries(self) -> list[dict]: + """Return summaries of all stored entries.""" + return [{"data_id": k, **v.summary()} for k, v in self._entries.items()] + + +# Global singleton — shared across all tools in this server +store = DataStore() diff --git a/src/servers/vibration/dsp/__init__.py b/src/servers/vibration/dsp/__init__.py new file mode 100644 index 00000000..4f7dc397 --- /dev/null +++ b/src/servers/vibration/dsp/__init__.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/LGDiMaggio/claude-stwinbox-diagnostics/tree/main/mcp-servers/vibration-analysis-mcp + +from .fft_analysis import compute_fft, compute_psd, compute_spectrogram, find_peaks_in_spectrum +from .envelope import envelope_spectrum, check_bearing_peaks +from .bearing_freqs import compute_bearing_frequencies, get_bearing, list_bearings, COMMON_BEARINGS +from .fault_detection import ( + assess_iso10816, + extract_shaft_features, + classify_faults, + generate_diagnosis_summary, +) + +__all__ = [ + "compute_fft", + "compute_psd", + "compute_spectrogram", + "find_peaks_in_spectrum", + "envelope_spectrum", + "check_bearing_peaks", + "compute_bearing_frequencies", + "get_bearing", + "list_bearings", + "COMMON_BEARINGS", + "assess_iso10816", + "extract_shaft_features", + "classify_faults", + "generate_diagnosis_summary", +] diff --git a/src/servers/vibration/dsp/bearing_freqs.py b/src/servers/vibration/dsp/bearing_freqs.py new file mode 100644 index 00000000..149fb997 --- /dev/null +++ b/src/servers/vibration/dsp/bearing_freqs.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/LGDiMaggio/claude-stwinbox-diagnostics/tree/main/mcp-servers/vibration-analysis-mcp +""" +Bearing characteristic frequency calculator. + +Computes BPFI, BPFO, BSF, and FTF from bearing geometry and shaft RPM. +These are the fundamental frequencies used in envelope analysis to detect +bearing faults in rotating machinery. + +Formulas based on standard rolling element bearing kinematics: + FTF = (RPM/60) * 0.5 * (1 - Bd/Pd * cos(alpha)) + BPFO = (RPM/60) * N_balls/2 * (1 - Bd/Pd * cos(alpha)) + BPFI = (RPM/60) * N_balls/2 * (1 + Bd/Pd * cos(alpha)) + BSF = (RPM/60) * Pd/(2*Bd) * (1 - (Bd/Pd * cos(alpha))^2) +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Optional + + +@dataclass +class BearingGeometry: + """Bearing physical dimensions.""" + + name: str + n_balls: int # Number of rolling elements + ball_dia: float # Ball/roller diameter (mm) + pitch_dia: float # Pitch diameter (mm) + contact_angle: float = 0.0 # Contact angle in degrees + + +@dataclass +class BearingFrequencies: + """Calculated bearing characteristic frequencies at a given RPM.""" + + rpm: float + ftf: float # Fundamental Train Frequency (cage) + bpfo: float # Ball Pass Frequency Outer race + bpfi: float # Ball Pass Frequency Inner race + bsf: float # Ball Spin Frequency + bearing_name: str = "" + + def to_dict(self) -> dict: + return { + "bearing": self.bearing_name, + "rpm": self.rpm, + "shaft_frequency_hz": self.rpm / 60.0, + "ftf_hz": round(self.ftf, 3), + "bpfo_hz": round(self.bpfo, 3), + "bpfi_hz": round(self.bpfi, 3), + "bsf_hz": round(self.bsf, 3), + "harmonics": { + "bpfo_2x": round(2 * self.bpfo, 3), + "bpfo_3x": round(3 * self.bpfo, 3), + "bpfi_2x": round(2 * self.bpfi, 3), + "bpfi_3x": round(3 * self.bpfi, 3), + "bsf_2x": round(2 * self.bsf, 3), + }, + } + + +def compute_bearing_frequencies( + rpm: float, + n_balls: int, + ball_dia: float, + pitch_dia: float, + contact_angle: float = 0.0, + bearing_name: str = "Unknown", +) -> BearingFrequencies: + """ + Compute bearing characteristic frequencies. + + Args: + rpm: Shaft rotational speed in RPM + n_balls: Number of rolling elements (balls or rollers) + ball_dia: Ball/roller diameter in mm + pitch_dia: Pitch (cage) diameter in mm + contact_angle: Contact angle in degrees (0 for radial bearings) + bearing_name: Optional bearing designation + + Returns: + BearingFrequencies with FTF, BPFO, BPFI, BSF values in Hz + """ + f_shaft = rpm / 60.0 + alpha_rad = math.radians(contact_angle) + ratio = ball_dia / pitch_dia + + ftf = f_shaft * 0.5 * (1.0 - ratio * math.cos(alpha_rad)) + bpfo = f_shaft * (n_balls / 2.0) * (1.0 - ratio * math.cos(alpha_rad)) + bpfi = f_shaft * (n_balls / 2.0) * (1.0 + ratio * math.cos(alpha_rad)) + bsf = f_shaft * (pitch_dia / (2.0 * ball_dia)) * ( + 1.0 - (ratio * math.cos(alpha_rad)) ** 2 + ) + + return BearingFrequencies( + rpm=rpm, + ftf=ftf, + bpfo=bpfo, + bpfi=bpfi, + bsf=bsf, + bearing_name=bearing_name, + ) + + +# --------------------------------------------------------------------------- +# Common bearing database +# --------------------------------------------------------------------------- + +COMMON_BEARINGS: dict[str, BearingGeometry] = { + "6205": BearingGeometry( + name="6205 (Deep groove)", + n_balls=9, + ball_dia=7.938, + pitch_dia=38.5, + contact_angle=0, + ), + "6206": BearingGeometry( + name="6206 (Deep groove)", + n_balls=9, + ball_dia=9.525, + pitch_dia=46.0, + contact_angle=0, + ), + "6207": BearingGeometry( + name="6207 (Deep groove)", + n_balls=9, + ball_dia=11.112, + pitch_dia=53.5, + contact_angle=0, + ), + "6208": BearingGeometry( + name="6208 (Deep groove)", + n_balls=9, + ball_dia=12.7, + pitch_dia=60.0, + contact_angle=0, + ), + "6305": BearingGeometry( + name="6305 (Deep groove)", + n_balls=8, + ball_dia=10.319, + pitch_dia=39.04, + contact_angle=0, + ), + "6306": BearingGeometry( + name="6306 (Deep groove)", + n_balls=8, + ball_dia=12.303, + pitch_dia=46.36, + contact_angle=0, + ), + "NU205": BearingGeometry( + name="NU205 (Cylindrical roller)", + n_balls=13, + ball_dia=7.5, + pitch_dia=38.5, + contact_angle=0, + ), + "NU206": BearingGeometry( + name="NU206 (Cylindrical roller)", + n_balls=13, + ball_dia=9.0, + pitch_dia=46.0, + contact_angle=0, + ), + "7205": BearingGeometry( + name="7205 (Angular contact)", + n_balls=12, + ball_dia=7.144, + pitch_dia=38.0, + contact_angle=25, + ), +} + + +def get_bearing(designation: str) -> Optional[BearingGeometry]: + """Look up a bearing by its designation.""" + return COMMON_BEARINGS.get(designation.upper()) + + +def list_bearings() -> list[dict]: + """List all bearings in the database.""" + return [ + { + "designation": k, + "name": v.name, + "n_balls": v.n_balls, + "ball_dia_mm": v.ball_dia, + "pitch_dia_mm": v.pitch_dia, + "contact_angle_deg": v.contact_angle, + } + for k, v in COMMON_BEARINGS.items() + ] diff --git a/src/servers/vibration/dsp/envelope.py b/src/servers/vibration/dsp/envelope.py new file mode 100644 index 00000000..f60b23ec --- /dev/null +++ b/src/servers/vibration/dsp/envelope.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/LGDiMaggio/claude-stwinbox-diagnostics/tree/main/mcp-servers/vibration-analysis-mcp +""" +Envelope analysis module for bearing fault detection. + +Envelope analysis (demodulation) uses the Hilbert transform to extract +the amplitude envelope of a band-pass filtered vibration signal. +This reveals the repetition rate of impulsive events (e.g., a defect +on a bearing race) even when the individual impacts are buried in noise. + +Typical workflow: + 1. Band-pass filter the signal around a structural resonance + 2. Compute the analytic signal via the Hilbert transform + 3. Extract the envelope (magnitude of the analytic signal) + 4. Compute the FFT of the envelope -> the envelope spectrum + 5. Look for peaks at bearing fault frequencies (BPFO, BPFI, BSF, FTF) +""" + +from __future__ import annotations + +import numpy as np +from numpy.typing import NDArray +from scipy.signal import butter, sosfilt, hilbert + + +def bandpass_filter( + signal: NDArray[np.floating], + fs: float, + low_hz: float, + high_hz: float, + order: int = 4, +) -> NDArray[np.floating]: + """ + Apply a Butterworth band-pass filter. + + Args: + signal: Input time-domain signal. + fs: Sampling frequency in Hz. + low_hz: Lower cutoff frequency. + high_hz: Upper cutoff frequency. + order: Filter order (default 4). + + Returns: + Filtered signal. + """ + nyq = fs / 2.0 + low = low_hz / nyq + high = min(high_hz / nyq, 0.99) # stay below Nyquist + sos = butter(order, [low, high], btype="band", output="sos") + return sosfilt(sos, signal).astype(signal.dtype) + + +def compute_envelope( + signal: NDArray[np.floating], +) -> NDArray[np.floating]: + """ + Compute the amplitude envelope using the Hilbert transform. + + Args: + signal: Input time-domain signal (ideally band-pass filtered). + + Returns: + Envelope signal (same length as input). + """ + analytic = hilbert(signal) + return np.abs(analytic).astype(signal.dtype) + + +def envelope_spectrum( + signal: NDArray[np.floating], + fs: float, + band_low: float | None = None, + band_high: float | None = None, + filter_order: int = 4, + n_fft: int | None = None, +) -> dict: + """ + Full envelope analysis pipeline: band-pass -> envelope -> FFT. + + Args: + signal: Raw vibration time-domain signal. + fs: Sampling frequency in Hz. + band_low: Band-pass lower cutoff (Hz). If None, defaults to fs/20. + band_high: Band-pass upper cutoff (Hz). If None, defaults to fs/2.5. + filter_order: Butterworth filter order. + n_fft: FFT length. Default = signal length. + + Returns: + dict with keys: frequencies, envelope_spectrum, filter_band, n_samples, fs + """ + n = len(signal) + if band_low is None: + band_low = fs / 20.0 + if band_high is None: + band_high = fs / 2.5 + if n_fft is None: + n_fft = n + + # Step 1: Band-pass filter + filtered = bandpass_filter(signal, fs, band_low, band_high, order=filter_order) + + # Step 2: Hilbert envelope + env = compute_envelope(filtered) + + # Remove DC from envelope before FFT + env_zero_mean = env - np.mean(env) + + # Step 3: FFT of envelope + window = np.hanning(n) + fft_vals = np.fft.rfft(env_zero_mean * window, n=n_fft) + freqs = np.fft.rfftfreq(n_fft, d=1.0 / fs) + magnitudes = (2.0 / n) * np.abs(fft_vals) + + return { + "frequencies": freqs.tolist(), + "envelope_spectrum": magnitudes.tolist(), + "filter_band": (band_low, band_high), + "n_samples": n, + "fs": fs, + } + + +def check_bearing_peaks( + freqs: NDArray[np.floating] | list[float], + magnitudes: NDArray[np.floating] | list[float], + target_freq: float, + n_harmonics: int = 3, + tolerance_pct: float = 3.0, + noise_floor_multiplier: float = 3.0, +) -> dict: + """ + Check whether a target bearing frequency (and its harmonics) is present + in an envelope spectrum. + + Args: + freqs: Frequency axis (Hz). + magnitudes: Spectrum magnitudes. + target_freq: Expected fault frequency (e.g., BPFO) in Hz. + n_harmonics: Number of harmonics to check (1x, 2x, ... Nx). + tolerance_pct: Frequency matching tolerance in % of target. + noise_floor_multiplier: Peak must exceed median * this to be significant. + + Returns: + dict with detection results per harmonic and overall verdict. + """ + freqs = np.asarray(freqs) + mags = np.asarray(magnitudes) + + noise_floor = np.median(mags) * noise_floor_multiplier + + results = [] + detected_count = 0 + + for h in range(1, n_harmonics + 1): + f_expected = h * target_freq + tol = f_expected * tolerance_pct / 100.0 + mask = (freqs >= f_expected - tol) & (freqs <= f_expected + tol) + + if np.any(mask): + peak_mag = float(np.max(mags[mask])) + peak_freq = float(freqs[mask][np.argmax(mags[mask])]) + is_detected = bool(peak_mag > noise_floor) + if is_detected: + detected_count += 1 + results.append( + { + "harmonic": h, + "expected_hz": round(f_expected, 2), + "found_hz": round(peak_freq, 2), + "amplitude": round(peak_mag, 6), + "noise_threshold": round(float(noise_floor), 6), + "detected": is_detected, + } + ) + else: + results.append( + { + "harmonic": h, + "expected_hz": round(f_expected, 2), + "found_hz": None, + "amplitude": 0.0, + "noise_threshold": round(float(noise_floor), 6), + "detected": False, + } + ) + + return { + "target_frequency_hz": target_freq, + "harmonics_checked": n_harmonics, + "harmonics_detected": detected_count, + "confidence": ( + "high" if detected_count >= 2 else ("medium" if detected_count == 1 else "none") + ), + "details": results, + } diff --git a/src/servers/vibration/dsp/fault_detection.py b/src/servers/vibration/dsp/fault_detection.py new file mode 100644 index 00000000..d613a814 --- /dev/null +++ b/src/servers/vibration/dsp/fault_detection.py @@ -0,0 +1,388 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/LGDiMaggio/claude-stwinbox-diagnostics/tree/main/mcp-servers/vibration-analysis-mcp +""" +Automated fault detection and classification for rotating machinery. + +Uses vibration spectral features to identify common mechanical faults: + - Unbalance (1x shaft speed dominant) + - Misalignment (2x shaft speed, axial component) + - Mechanical looseness (many harmonics of shaft speed) + - Bearing faults (BPFO / BPFI / BSF / FTF in envelope spectrum) + +Also provides ISO 10816 vibration severity classification. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import numpy as np +from numpy.typing import NDArray + + +# --------------------------------------------------------------------------- +# ISO 10816 Vibration Severity (RMS velocity mm/s) +# --------------------------------------------------------------------------- + +# Thresholds per machine group (mm/s RMS, 10-1000 Hz) +ISO_10816_THRESHOLDS: dict[str, dict[str, float]] = { + # Group 1: Large machines > 300 kW on rigid foundations + "group1": {"A_good": 2.8, "B_acceptable": 7.1, "C_alarm": 18.0}, + # Group 2: Medium machines 15-300 kW on rigid foundations + "group2": {"A_good": 1.4, "B_acceptable": 2.8, "C_alarm": 7.1}, + # Group 3: Large machines on flexible foundations + "group3": {"A_good": 3.5, "B_acceptable": 9.0, "C_alarm": 22.4}, + # Group 4: Small machines < 15 kW + "group4": {"A_good": 0.71, "B_acceptable": 1.8, "C_alarm": 4.5}, +} + + +def assess_iso10816( + rms_velocity_mm_s: float, + machine_group: str = "group2", +) -> dict: + """ + Classify vibration severity per ISO 10816. + + Args: + rms_velocity_mm_s: Overall RMS velocity in mm/s (10-1000 Hz band). + machine_group: One of 'group1' .. 'group4'. + + Returns: + dict with zone (A/B/C/D), description, and thresholds used. + """ + thresholds = ISO_10816_THRESHOLDS.get( + machine_group, ISO_10816_THRESHOLDS["group2"] + ) + + if rms_velocity_mm_s <= thresholds["A_good"]: + zone, desc = "A", "Good - newly commissioned machines" + elif rms_velocity_mm_s <= thresholds["B_acceptable"]: + zone, desc = "B", "Acceptable - long-term operation permitted" + elif rms_velocity_mm_s <= thresholds["C_alarm"]: + zone, desc = "C", "Alarm - not suitable for long-term operation" + else: + zone, desc = "D", "Danger - risk of damage, stop machine" + + return { + "rms_velocity_mm_s": round(rms_velocity_mm_s, 3), + "iso_zone": zone, + "description": desc, + "machine_group": machine_group, + "thresholds": thresholds, + } + + +# --------------------------------------------------------------------------- +# Spectral fault feature extraction +# --------------------------------------------------------------------------- + + +@dataclass +class ShaftFeatures: + """Spectral amplitude at shaft-frequency harmonics.""" + + f_shaft: float # Shaft frequency (Hz) = RPM/60 + amp_1x: float # 1x amplitude + amp_2x: float # 2x amplitude + amp_3x: float # 3x amplitude + amp_half_x: float # 0.5x amplitude (sub-harmonic) + rms_overall: float # Broadband RMS + crest_factor: float # Peak / RMS + kurtosis: float # Excess kurtosis (Fisher: Gaussian = 0, impulsive > 1) + + +def extract_shaft_features( + freqs: NDArray[np.floating] | list[float], + magnitudes: NDArray[np.floating] | list[float], + shaft_freq: float, + time_signal: NDArray[np.floating] | list[float] | None = None, + tolerance_pct: float = 3.0, +) -> ShaftFeatures: + """ + Extract amplitudes at shaft-frequency harmonics from a spectrum. + + Args: + freqs: Frequency axis in Hz. + magnitudes: Spectrum magnitude array. + shaft_freq: Shaft rotational frequency in Hz (RPM / 60). + time_signal: Optional raw time signal for kurtosis / crest factor. + tolerance_pct: Frequency tolerance percentage. + + Returns: + ShaftFeatures dataclass. + """ + freqs = np.asarray(freqs) + mags = np.asarray(magnitudes) + + def _peak_at(f_target: float) -> float: + tol = f_target * tolerance_pct / 100.0 + mask = (freqs >= f_target - tol) & (freqs <= f_target + tol) + if np.any(mask): + return float(np.max(mags[mask])) + return 0.0 + + amp_1x = _peak_at(shaft_freq) + amp_2x = _peak_at(2.0 * shaft_freq) + amp_3x = _peak_at(3.0 * shaft_freq) + amp_half = _peak_at(0.5 * shaft_freq) + rms_overall = float(np.sqrt(np.mean(mags**2))) + + # Crest factor & kurtosis from time signal if available + if time_signal is not None: + ts = np.asarray(time_signal) + ts_rms = float(np.sqrt(np.mean(ts**2))) + crest = float(np.max(np.abs(ts)) / ts_rms) if ts_rms > 0 else 0 + mean_ts = np.mean(ts) + std_ts = np.std(ts, ddof=1) + if std_ts > 0 and len(ts) >= 4: + # Excess kurtosis (Fisher definition): Gaussian = 0 + kurt = float(np.mean(((ts - mean_ts) / std_ts) ** 4) - 3.0) + else: + kurt = 0.0 + else: + crest = 0.0 + kurt = 0.0 + + return ShaftFeatures( + f_shaft=shaft_freq, + amp_1x=amp_1x, + amp_2x=amp_2x, + amp_3x=amp_3x, + amp_half_x=amp_half, + rms_overall=rms_overall, + crest_factor=crest, + kurtosis=kurt, + ) + + +# --------------------------------------------------------------------------- +# Fault classification rules +# --------------------------------------------------------------------------- + + +@dataclass +class FaultDiagnosis: + """Result of automated fault classification.""" + + fault_type: str + confidence: str # "high", "medium", "low", "none" + description: str + evidence: list[str] + recommendations: list[str] + + def to_dict(self) -> dict: + return { + "fault_type": self.fault_type, + "confidence": self.confidence, + "description": self.description, + "evidence": self.evidence, + "recommendations": self.recommendations, + } + + +def classify_faults( + features: ShaftFeatures, + bearing_envelope_results: dict | None = None, +) -> list[FaultDiagnosis]: + """ + Apply rule-based classification for common rotating-machinery faults. + + Args: + features: Extracted shaft-frequency features. + bearing_envelope_results: Optional dict from envelope analysis with + keys like 'bpfo', 'bpfi', 'bsf', each containing + {harmonics_detected, confidence}. + + Returns: + List of FaultDiagnosis objects, sorted by confidence (high first). + """ + diagnoses: list[FaultDiagnosis] = [] + rms = features.rms_overall if features.rms_overall > 0 else 1e-12 + + # --- Unbalance: dominant 1x with low 2x --- + ratio_1x = features.amp_1x / rms + ratio_2x = features.amp_2x / rms + if ratio_1x > 3.0 and features.amp_1x > 2.0 * features.amp_2x: + conf = "high" if ratio_1x > 5.0 else "medium" + diagnoses.append( + FaultDiagnosis( + fault_type="unbalance", + confidence=conf, + description="Mass unbalance detected - dominant 1x shaft speed component.", + evidence=[ + f"1x amplitude = {features.amp_1x:.4f} ({ratio_1x:.1f}x RMS)", + f"2x amplitude = {features.amp_2x:.4f} " + f"(ratio 1x/2x = {features.amp_1x / max(features.amp_2x, 1e-12):.1f})", + ], + recommendations=[ + "Perform balancing of the rotor.", + "Check for material build-up or loss on rotating parts.", + "Verify coupling alignment (unbalance can be masked by misalignment).", + ], + ) + ) + + # --- Misalignment: significant 2x (and sometimes 3x) --- + if ratio_2x > 2.5 and features.amp_2x > 0.5 * features.amp_1x: + conf = "high" if features.amp_2x > features.amp_1x else "medium" + diagnoses.append( + FaultDiagnosis( + fault_type="misalignment", + confidence=conf, + description="Shaft misalignment suspected - elevated 2x component.", + evidence=[ + f"2x amplitude = {features.amp_2x:.4f} ({ratio_2x:.1f}x RMS)", + f"2x/1x ratio = {features.amp_2x / max(features.amp_1x, 1e-12):.2f}", + f"3x amplitude = {features.amp_3x:.4f}", + ], + recommendations=[ + "Check shaft alignment with laser or dial indicator.", + "Inspect coupling condition and flexible element wear.", + "Verify thermal growth compensation.", + ], + ) + ) + + # --- Mechanical looseness: many harmonics + sub-harmonics --- + n_significant = sum( + 1 + for a in [features.amp_1x, features.amp_2x, features.amp_3x] + if a / rms > 1.5 + ) + if n_significant >= 3 or (features.amp_half_x / rms > 1.5): + evidence = [f"Harmonics above threshold: {n_significant}/3"] + if features.amp_half_x / rms > 1.5: + evidence.append(f"Sub-harmonic 0.5x = {features.amp_half_x:.4f}") + diagnoses.append( + FaultDiagnosis( + fault_type="mechanical_looseness", + confidence="medium", + description="Mechanical looseness suggested - multiple shaft harmonics and/or sub-harmonics.", + evidence=evidence, + recommendations=[ + "Inspect and tighten foundation bolts.", + "Check bearing housing fit and clearance.", + "Look for structural cracks or soft foot.", + ], + ) + ) + + # --- Impulsiveness (excess kurtosis > 1.0 or crest factor > 5) --- + if features.kurtosis > 1.0 or features.crest_factor > 5.0: + diagnoses.append( + FaultDiagnosis( + fault_type="impulsive_signal", + confidence="medium", + description="Impulsive content detected - may indicate bearing defect or gear tooth damage.", + evidence=[ + f"Excess kurtosis = {features.kurtosis:.2f} (healthy ~ 0.0)", + f"Crest factor = {features.crest_factor:.2f} (healthy < 4)", + ], + recommendations=[ + "Perform envelope analysis to isolate bearing fault frequencies.", + "Inspect gears for pitting or tooth breakage if applicable.", + ], + ) + ) + + # --- Bearing faults from envelope results --- + if bearing_envelope_results: + for fault_key, label in [ + ("bpfo", "Outer race defect"), + ("bpfi", "Inner race defect"), + ("bsf", "Ball/roller defect"), + ("ftf", "Cage defect"), + ]: + result = bearing_envelope_results.get(fault_key) + if result and result.get("confidence", "none") != "none": + conf = result["confidence"] + diagnoses.append( + FaultDiagnosis( + fault_type=f"bearing_{fault_key}", + confidence=conf, + description=f"{label} detected via envelope analysis.", + evidence=[ + f"Harmonics detected: {result.get('harmonics_detected', 0)}" + f"/{result.get('harmonics_checked', 3)}", + f"Fault frequency: {result.get('target_frequency_hz', 0):.2f} Hz", + ], + recommendations=[ + "Schedule bearing replacement.", + "Monitor trend - frequency of data collection should increase.", + "Check lubrication condition.", + ], + ) + ) + + # --- No faults --- + if not diagnoses: + diagnoses.append( + FaultDiagnosis( + fault_type="healthy", + confidence="medium", + description="No significant fault patterns detected in the vibration signature.", + evidence=[ + f"1x ratio = {ratio_1x:.1f}", + f"Kurtosis = {features.kurtosis:.2f}", + f"Crest factor = {features.crest_factor:.2f}", + ], + recommendations=[ + "Continue routine monitoring.", + "Store this measurement as baseline reference.", + ], + ) + ) + + # Sort: high > medium > low > none + priority = {"high": 0, "medium": 1, "low": 2, "none": 3} + diagnoses.sort(key=lambda d: priority.get(d.confidence, 3)) + return diagnoses + + +def generate_diagnosis_summary( + diagnoses: list[FaultDiagnosis], + iso_assessment: dict | None = None, + machine_context: str = "", +) -> str: + """ + Create a human-readable diagnosis summary. + + Args: + diagnoses: List of FaultDiagnosis objects. + iso_assessment: Optional ISO 10816 result dict. + machine_context: Free text describing the machine. + + Returns: + Formatted markdown string. + """ + lines = ["## Vibration Diagnosis Summary\n"] + + if machine_context: + lines.append(f"**Machine:** {machine_context}\n") + + if iso_assessment: + zone = iso_assessment.get("iso_zone", "?") + vel = iso_assessment.get("rms_velocity_mm_s", 0) + desc = iso_assessment.get("description", "") + lines.append( + f"**Overall Severity (ISO 10816):** Zone {zone} " + f"- {vel:.2f} mm/s RMS - {desc}\n" + ) + + lines.append("### Detected Conditions\n") + for i, d in enumerate(diagnoses, 1): + lines.append( + f"#### {i}. {d.fault_type.replace('_', ' ').title()} " + f"[{d.confidence.upper()}]\n" + ) + lines.append(f"{d.description}\n") + lines.append("**Evidence:**") + for e in d.evidence: + lines.append(f"- {e}") + lines.append("\n**Recommendations:**") + for r in d.recommendations: + lines.append(f"- {r}") + lines.append("") + + return "\n".join(lines) diff --git a/src/servers/vibration/dsp/fft_analysis.py b/src/servers/vibration/dsp/fft_analysis.py new file mode 100644 index 00000000..77bffc59 --- /dev/null +++ b/src/servers/vibration/dsp/fft_analysis.py @@ -0,0 +1,195 @@ +# SPDX-License-Identifier: Apache-2.0 +# Adapted from https://github.com/LGDiMaggio/claude-stwinbox-diagnostics/tree/main/mcp-servers/vibration-analysis-mcp +""" +FFT and spectral analysis module. + +Provides functions for frequency-domain analysis of vibration signals: +- FFT magnitude spectrum +- Power Spectral Density (Welch method) +- Spectrogram (STFT) +- Peak detection +""" + +from __future__ import annotations + +import numpy as np +from scipy import signal as sig +from typing import Optional + + +def compute_fft( + data: np.ndarray, + fs: float, + window: str = "hann", + n_fft: Optional[int] = None, +) -> dict: + """ + Compute single-sided FFT magnitude spectrum. + + Args: + data: 1D time-domain signal + fs: Sampling frequency in Hz + window: Window function ('hann', 'hamming', 'blackman', 'rectangular') + n_fft: FFT length (defaults to len(data), zero-padded if > len(data)) + + Returns: + Dictionary with 'frequencies' (Hz), 'magnitude' (linear), + 'magnitude_db' (dB), and 'resolution_hz' + """ + n = len(data) + if n_fft is None: + n_fft = n + + # Apply window + if window == "rectangular": + w = np.ones(n) + else: + w = sig.get_window(window, n) + + windowed = data * w + + # Compute FFT + fft_vals = np.fft.rfft(windowed, n=n_fft) + freqs = np.fft.rfftfreq(n_fft, d=1.0 / fs) + + # Magnitude (single-sided, compensated for window energy loss) + magnitude = (2.0 / n) * np.abs(fft_vals) + magnitude[0] /= 2.0 # DC component not doubled + + # Convert to dB (reference: 1.0) + magnitude_db = 20 * np.log10(magnitude + 1e-12) + + return { + "frequencies": freqs, + "magnitude": magnitude, + "magnitude_db": magnitude_db, + "resolution_hz": fs / n_fft, + "max_frequency_hz": fs / 2, + "num_points": len(freqs), + } + + +def compute_psd( + data: np.ndarray, + fs: float, + nperseg: int = 1024, + noverlap: Optional[int] = None, + window: str = "hann", +) -> dict: + """ + Compute Power Spectral Density using Welch's method. + + Args: + data: 1D time-domain signal + fs: Sampling frequency in Hz + nperseg: Segment length for Welch method + noverlap: Overlap between segments (default: nperseg // 2) + window: Window function + + Returns: + Dictionary with 'frequencies' (Hz) and 'psd' + """ + if noverlap is None: + noverlap = nperseg // 2 + + freqs, psd = sig.welch( + data, fs=fs, nperseg=nperseg, noverlap=noverlap, window=window + ) + + return { + "frequencies": freqs, + "psd": psd, + "total_power": float(np.trapezoid(psd, freqs)), + "resolution_hz": freqs[1] - freqs[0] if len(freqs) > 1 else 0, + } + + +def compute_spectrogram( + data: np.ndarray, + fs: float, + nperseg: int = 256, + noverlap: Optional[int] = None, + window: str = "hann", +) -> dict: + """ + Compute spectrogram (Short-Time Fourier Transform). + + Args: + data: 1D time-domain signal + fs: Sampling frequency in Hz + nperseg: Segment length + noverlap: Overlap (default: nperseg // 2) + window: Window function + + Returns: + Dictionary with 'frequencies', 'times', 'spectrogram_db' + """ + if noverlap is None: + noverlap = nperseg // 2 + + freqs, times, Sxx = sig.spectrogram( + data, fs=fs, nperseg=nperseg, noverlap=noverlap, window=window + ) + + Sxx_db = 10 * np.log10(Sxx + 1e-12) + + return { + "frequencies": freqs, + "times": times, + "spectrogram_db": Sxx_db, + "num_time_frames": len(times), + "num_freq_bins": len(freqs), + } + + +def find_peaks_in_spectrum( + frequencies: np.ndarray, + magnitude: np.ndarray, + num_peaks: int = 10, + min_distance_hz: float = 5.0, + threshold_db: float = -60.0, +) -> list[dict]: + """ + Find dominant peaks in a frequency spectrum. + + Args: + frequencies: Frequency axis in Hz + magnitude: Magnitude spectrum (linear) + num_peaks: Maximum number of peaks to return + min_distance_hz: Minimum distance between peaks in Hz + threshold_db: Minimum amplitude threshold in dB + + Returns: + List of dicts with 'frequency_hz', 'magnitude', 'magnitude_db' + """ + magnitude_db = 20 * np.log10(magnitude + 1e-12) + + # Convert min_distance to samples + df = frequencies[1] - frequencies[0] if len(frequencies) > 1 else 1.0 + min_distance_samples = max(1, int(min_distance_hz / df)) + + # Find peaks + peak_indices, _properties = sig.find_peaks( + magnitude_db, + distance=min_distance_samples, + height=threshold_db, + ) + + if len(peak_indices) == 0: + return [] + + # Sort by amplitude (descending) + sorted_idx = np.argsort(magnitude_db[peak_indices])[::-1] + top_peaks = peak_indices[sorted_idx[:num_peaks]] + + peaks = [] + for idx in top_peaks: + peaks.append( + { + "frequency_hz": float(frequencies[idx]), + "magnitude": float(magnitude[idx]), + "magnitude_db": float(magnitude_db[idx]), + } + ) + + return peaks diff --git a/src/servers/vibration/main.py b/src/servers/vibration/main.py new file mode 100644 index 00000000..65bf6246 --- /dev/null +++ b/src/servers/vibration/main.py @@ -0,0 +1,552 @@ +"""VibrationAgent MCP Server. + +Vibration signal analysis and rotating machinery fault detection +for AssetOpsBench. Reads sensor data from CouchDB, performs FFT, +envelope analysis, bearing fault detection, and ISO 10816 assessment. + +DSP core adapted from https://github.com/LGDiMaggio/claude-stwinbox-diagnostics/tree/main/mcp-servers/vibration-analysis-mcp +SPDX-License-Identifier: Apache-2.0 +""" + +from __future__ import annotations + +import logging +import os +from typing import Optional, Union + +import numpy as np +from dotenv import load_dotenv +from mcp.server.fastmcp import FastMCP +from pydantic import BaseModel + +from .couchdb_client import fetch_vibration_timeseries, list_sensor_fields +from .data_store import store +from .dsp.bearing_freqs import ( + COMMON_BEARINGS, + compute_bearing_frequencies, + get_bearing, + list_bearings, +) +from .dsp.envelope import check_bearing_peaks, envelope_spectrum +from .dsp.fault_detection import ( + assess_iso10816, + classify_faults, + extract_shaft_features, + generate_diagnosis_summary, +) +from .dsp.fft_analysis import compute_fft + +load_dotenv() +_log_level = getattr( + logging, os.environ.get("LOG_LEVEL", "WARNING").upper(), logging.WARNING +) +logging.basicConfig(level=_log_level) +logger = logging.getLogger("vibration-mcp-server") + +mcp = FastMCP("VibrationAgent") + + +# --------------------------------------------------------------------------- +# Pydantic response models +# --------------------------------------------------------------------------- + + +class ErrorResult(BaseModel): + error: str + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _compact_spectrum(freqs: np.ndarray, mags: np.ndarray, top_n: int = 20) -> dict: + """Summarise a spectrum: top-N peaks + global stats. No full arrays.""" + order = np.argsort(mags)[::-1] + n = min(top_n, len(order)) + top_idx = np.sort(order[:n]) # sort back by frequency + peaks = [ + {"freq_hz": round(float(freqs[i]), 3), "amplitude": round(float(mags[i]), 6)} + for i in top_idx + ] + return { + "top_peaks": peaks, + "max_amplitude": round(float(np.max(mags)), 6), + "max_amplitude_freq_hz": round(float(freqs[np.argmax(mags)]), 3), + "rms_spectral": round(float(np.sqrt(np.mean(mags**2))), 6), + "total_bins": len(freqs), + "freq_range_hz": [round(float(freqs[0]), 3), round(float(freqs[-1]), 3)], + } + + +def _accel_g_to_velocity_rms_mms( + signal_g: np.ndarray, + sample_rate: float, + f_low: float = 10.0, + f_high: float = 1000.0, +) -> float: + """ + Convert an acceleration signal (in g) to RMS velocity (in mm/s) + within the ISO 10816 band (default 10-1000 Hz). + + Method: frequency-domain integration. + """ + N = len(signal_g) + if N < 2: + return 0.0 + + accel_ms2 = signal_g * 9.80665 + fft_vals = np.fft.rfft(accel_ms2) + freqs = np.fft.rfftfreq(N, d=1.0 / sample_rate) + + # Vectorised integration: v(f) = a(f) / (j·2π·f), skip DC bin + vel_fft = np.zeros_like(fft_vals) + mask = (freqs >= f_low) & (freqs <= f_high) + vel_fft[mask] = fft_vals[mask] / (1j * 2.0 * np.pi * freqs[mask]) + + velocity_ms = np.fft.irfft(vel_fft, n=N) + velocity_mms = velocity_ms * 1000.0 + rms = float(np.sqrt(np.mean(velocity_mms**2))) + return rms + + +def _resolve_signal(data_id: str) -> tuple[np.ndarray, float]: + """Return (1-D numpy signal, sample_rate) from a data_id.""" + entry = store.get(data_id) + if entry is None: + available = store.list_ids() + raise ValueError( + f"data_id '{data_id}' not found in store. " + f"Available: {available or '(empty - use get_vibration_data first)'}." + ) + sig = entry.signal + if sig.ndim > 1: + sig = sig[:, 0] # default to first channel + return sig, entry.sample_rate + + +# --------------------------------------------------------------------------- +# MCP Tools +# --------------------------------------------------------------------------- + + +@mcp.tool() +def get_vibration_data( + site_name: str, + asset_id: str, + sensor_name: str, + start: str, + final: Optional[str] = None, +) -> Union[dict, ErrorResult]: + """Fetch vibration sensor time-series from CouchDB and load into the analysis store. + + Returns a data_id that can be passed to analysis tools (compute_fft_spectrum, + compute_envelope_spectrum, diagnose_vibration, etc.). + + Args: + site_name: Site identifier (e.g., 'MAIN'). + asset_id: Asset identifier (e.g., 'Chiller 6'). + sensor_name: Name of the sensor field in CouchDB documents. + start: ISO 8601 start timestamp (e.g., '2020-06-01T00:00:00'). + final: Optional ISO 8601 end timestamp. + """ + result = fetch_vibration_timeseries(asset_id, sensor_name, start, final) + if result is None: + return ErrorResult( + error=f"No vibration data found for asset '{asset_id}', " + f"sensor '{sensor_name}' in time range starting {start}." + ) + signal, sample_rate = result + data_id = f"{asset_id}_{sensor_name}".replace(" ", "_") + store.put( + data_id, + signal, + sample_rate, + { + "source": "couchdb", + "asset_id": asset_id, + "sensor": sensor_name, + "start": start, + "final": final, + }, + ) + entry = store.get(data_id) + return {"data_id": data_id, **entry.summary()} + + +@mcp.tool() +def list_vibration_sensors( + site_name: str, + asset_id: str, +) -> Union[dict, ErrorResult]: + """List available sensor fields for an asset. + + Args: + site_name: Site identifier (e.g., 'MAIN'). + asset_id: Asset identifier (e.g., 'Chiller 6'). + """ + sensors = list_sensor_fields(asset_id) + if not sensors: + return ErrorResult( + error=f"No sensors found for asset '{asset_id}' at site '{site_name}'." + ) + return { + "site_name": site_name, + "asset_id": asset_id, + "total_sensors": len(sensors), + "sensors": sensors, + } + + +@mcp.tool() +def compute_fft_spectrum( + data_id: str, + window: str = "hann", + top_n: int = 20, +) -> Union[dict, ErrorResult]: + """Compute FFT amplitude spectrum of a stored vibration signal. + + Returns a compact summary (top N peaks + statistics), not the full array. + + Args: + data_id: Reference to a stored signal (from get_vibration_data). + window: Window function ('hann', 'hamming', 'blackman', 'rectangular'). + top_n: Number of highest peaks to include in the summary. + """ + try: + sig, sr = _resolve_signal(data_id) + except ValueError as e: + return ErrorResult(error=str(e)) + + result = compute_fft(sig, sr, window=window) + freqs = np.asarray(result["frequencies"]) + mags = np.asarray(result["magnitude"]) + summary = _compact_spectrum(freqs, mags, top_n=top_n) + summary.update( + { + "data_id": data_id, + "n_samples": len(sig), + "sample_rate_hz": sr, + "window": window, + "frequency_resolution_hz": round(float(freqs[1]), 4) + if len(freqs) > 1 + else 0, + } + ) + return summary + + +@mcp.tool() +def compute_envelope_spectrum( + data_id: str, + band_low_hz: Optional[float] = None, + band_high_hz: Optional[float] = None, + top_n: int = 20, +) -> Union[dict, ErrorResult]: + """Compute the envelope spectrum for bearing fault detection. + + Args: + data_id: Reference to a stored signal (from get_vibration_data). + band_low_hz: Band-pass lower cutoff. Auto if None. + band_high_hz: Band-pass upper cutoff. Auto if None. + top_n: Number of highest peaks to include. + """ + try: + sig, sr = _resolve_signal(data_id) + except ValueError as e: + return ErrorResult(error=str(e)) + + result = envelope_spectrum(sig, sr, band_low=band_low_hz, band_high=band_high_hz) + freqs = np.asarray(result["frequencies"]) + env_mags = np.asarray(result["envelope_spectrum"]) + summary = _compact_spectrum(freqs, env_mags, top_n=top_n) + summary.update( + { + "data_id": data_id, + "filter_band_hz": list(result["filter_band"]), + "n_samples": result["n_samples"], + "sample_rate_hz": sr, + } + ) + return summary + + +@mcp.tool() +def assess_vibration_severity( + rms_velocity_mm_s: float, + machine_group: str = "group2", +) -> dict: + """Classify vibration severity per ISO 10816. + + IMPORTANT: expects RMS velocity in mm/s (not acceleration in g). + The diagnose_vibration tool performs the conversion automatically. + + Machine groups: + group1 - Large machines (>300 kW) on rigid foundations + group2 - Medium machines (15-300 kW) on rigid foundations + group3 - Large machines on flexible foundations + group4 - Small machines (<15 kW) + + Args: + rms_velocity_mm_s: Overall RMS velocity in mm/s (10-1000 Hz band). + machine_group: ISO 10816 machine group identifier. + """ + return assess_iso10816(rms_velocity_mm_s, machine_group) + + +@mcp.tool() +def calculate_bearing_frequencies( + rpm: float, + n_balls: int, + ball_diameter_mm: float, + pitch_diameter_mm: float, + contact_angle_deg: float = 0.0, + bearing_name: str = "", +) -> dict: + """Compute bearing characteristic frequencies (BPFO, BPFI, BSF, FTF) from geometry. + + Args: + rpm: Shaft speed in revolutions per minute. + n_balls: Number of rolling elements. + ball_diameter_mm: Ball or roller diameter in mm. + pitch_diameter_mm: Pitch (cage) diameter in mm. + contact_angle_deg: Contact angle in degrees (0 for radial bearings). + bearing_name: Optional bearing designation for labeling. + """ + result = compute_bearing_frequencies( + rpm, + n_balls, + ball_diameter_mm, + pitch_diameter_mm, + contact_angle_deg, + bearing_name or "custom", + ) + return result.to_dict() + + +@mcp.tool() +def list_known_bearings() -> dict: + """List all bearings in the built-in database with their geometric parameters.""" + return {"bearings": list_bearings()} + + +@mcp.tool() +def diagnose_vibration( + data_id: str, + rpm: Optional[float] = None, + bearing_designation: Optional[str] = None, + bearing_n_balls: Optional[int] = None, + bearing_ball_dia_mm: Optional[float] = None, + bearing_pitch_dia_mm: Optional[float] = None, + bearing_contact_angle_deg: float = 0.0, + bpfo_hz: Optional[float] = None, + bpfi_hz: Optional[float] = None, + bsf_hz: Optional[float] = None, + ftf_hz: Optional[float] = None, + machine_group: str = "group2", + machine_description: str = "", +) -> Union[dict, ErrorResult]: + """Full automated vibration diagnosis pipeline. + + 1. Compute FFT and extract shaft-frequency features (requires rpm) + 2. Optionally perform envelope analysis for bearing faults + 3. Classify faults (unbalance, misalignment, looseness, bearing) + 4. Assess ISO 10816 severity + 5. Generate human-readable report + + rpm is optional but strongly recommended. Without it the tool cannot + perform shaft-frequency analysis and will only report basic statistics. + + Signal input: provide data_id referencing a signal from get_vibration_data. + + Bearing information can be provided in three ways: + a) Direct fault frequencies in Hz: bpfo_hz, bpfi_hz, bsf_hz, ftf_hz + b) Custom geometry: bearing_n_balls, bearing_ball_dia_mm, bearing_pitch_dia_mm + c) Database lookup: bearing_designation (e.g., '6205', 'NU206') + + Args: + data_id: Reference to a stored signal (from get_vibration_data). + rpm: Shaft speed in RPM. Omit if unknown. + bearing_designation: Bearing code from built-in database. + bearing_n_balls: Number of rolling elements (custom geometry). + bearing_ball_dia_mm: Ball diameter in mm (custom geometry). + bearing_pitch_dia_mm: Pitch diameter in mm (custom geometry). + bearing_contact_angle_deg: Contact angle in degrees (default 0). + bpfo_hz: Known BPFO in Hz. + bpfi_hz: Known BPFI in Hz. + bsf_hz: Known BSF in Hz. + ftf_hz: Known FTF in Hz. + machine_group: ISO 10816 group ('group1'..'group4'). + machine_description: Free text describing the machine for the report. + """ + try: + sig, sr = _resolve_signal(data_id) + except ValueError as e: + return ErrorResult(error=str(e)) + + # Step 1: FFT + fft_result = compute_fft(sig, sr) + freqs = np.array(fft_result["frequencies"]) + mags = np.array(fft_result["magnitude"]) + + # Basic time-domain statistics + ts_rms = float(np.sqrt(np.mean(sig**2))) + ts_peak = float(np.max(np.abs(sig))) + ts_crest = ts_peak / ts_rms if ts_rms > 0 else 0.0 + ts_mean = float(np.mean(sig)) + ts_std = float(np.std(sig, ddof=1)) + ts_kurtosis = ( + float(np.mean(((sig - ts_mean) / ts_std) ** 4) - 3.0) + if ts_std > 0 and len(sig) >= 4 + else 0.0 + ) + + # ISO 10816 — convert acceleration (g) -> velocity (mm/s) + rms_vel_mms = _accel_g_to_velocity_rms_mms(sig, sr) + + # If RPM is not provided, skip shaft-frequency analysis + if rpm is None or rpm <= 0: + iso_no_rpm = assess_iso10816(rms_vel_mms, machine_group) + return { + "warning": ( + "RPM not provided — shaft-frequency analysis (1x, 2x, etc.) " + "was skipped. Only basic signal statistics are reported. " + "To enable full diagnosis, provide the shaft RPM." + ), + "signal_statistics": { + "rms_g": round(ts_rms, 6), + "peak_g": round(ts_peak, 6), + "crest_factor": round(ts_crest, 2), + "kurtosis": round(ts_kurtosis, 2), + "duration_s": round(len(sig) / sr, 4), + "sample_rate_hz": sr, + "rms_velocity_mm_s": round(rms_vel_mms, 3), + }, + "fft_summary": _compact_spectrum(freqs, mags, top_n=20), + "iso_10816": iso_no_rpm, + "diagnoses": [], + "shaft_features": None, + "bearing_analysis": None, + "bearing_info_source": None, + "report_markdown": ( + "## Vibration Analysis (RPM unknown)\n\n" + f"| Metric | Value |\n|---|---|\n" + f"| RMS (acceleration) | {ts_rms:.4f} g |\n" + f"| Peak (acceleration) | {ts_peak:.4f} g |\n" + f"| RMS (velocity, 10-1000 Hz) | {rms_vel_mms:.3f} mm/s |\n" + f"| Crest Factor | {ts_crest:.1f} |\n" + f"| Kurtosis | {ts_kurtosis:.1f} |\n\n" + f"**ISO 10816 Severity:** Zone {iso_no_rpm['iso_zone']} — " + f"{rms_vel_mms:.3f} mm/s RMS — {iso_no_rpm['description']}\n\n" + "**Note:** RPM was not provided. Fault classification " + "(unbalance, misalignment, looseness) requires shaft frequency." + ), + } + + shaft_freq = rpm / 60.0 + + # Step 2: Shaft features + features = extract_shaft_features(freqs, mags, shaft_freq, time_signal=sig) + + # Step 3: Resolve bearing fault frequencies + fault_freqs: dict[str, float] = {} + bearing_info_source = None + + # Method A: Direct fault frequencies in Hz + direct_hz = any(v is not None and v > 0 for v in [bpfo_hz, bpfi_hz, bsf_hz, ftf_hz]) + if direct_hz: + bearing_info_source = "user-provided frequencies" + for key, hz_val in [ + ("bpfo", bpfo_hz), + ("bpfi", bpfi_hz), + ("bsf", bsf_hz), + ("ftf", ftf_hz), + ]: + if hz_val is not None and hz_val > 0: + fault_freqs[key] = hz_val + # Method B: Custom geometry + elif bearing_n_balls and bearing_ball_dia_mm and bearing_pitch_dia_mm: + bearing_info_source = "custom geometry" + bf = compute_bearing_frequencies( + rpm, + bearing_n_balls, + bearing_ball_dia_mm, + bearing_pitch_dia_mm, + bearing_contact_angle_deg, + "custom", + ) + fault_freqs = {"bpfo": bf.bpfo, "bpfi": bf.bpfi, "bsf": bf.bsf, "ftf": bf.ftf} + # Method C: Database lookup + elif bearing_designation: + bearing = get_bearing(bearing_designation) + if bearing: + bearing_info_source = f"database ({bearing.name})" + bf = compute_bearing_frequencies( + rpm, + bearing.n_balls, + bearing.ball_dia, + bearing.pitch_dia, + bearing.contact_angle, + bearing.name, + ) + fault_freqs = { + "bpfo": bf.bpfo, + "bpfi": bf.bpfi, + "bsf": bf.bsf, + "ftf": bf.ftf, + } + + # Envelope analysis if any fault frequencies are available + bearing_results = None + if fault_freqs: + env = envelope_spectrum(sig, sr) + env_freqs = np.array(env["frequencies"]) + env_mags = np.array(env["envelope_spectrum"]) + bearing_results = {} + for key, freq_val in fault_freqs.items(): + bearing_results[key] = check_bearing_peaks(env_freqs, env_mags, freq_val) + + # Step 4: Classify + diagnoses = classify_faults(features, bearing_results) + + # Step 5: ISO 10816 + iso = assess_iso10816(rms_vel_mms, machine_group) + + # Step 6: Report + report = generate_diagnosis_summary(diagnoses, iso, machine_description) + + return { + "diagnoses": [d.to_dict() for d in diagnoses], + "iso_10816": iso, + "shaft_features": { + "shaft_freq_hz": features.f_shaft, + "amp_1x": round(features.amp_1x, 6), + "amp_2x": round(features.amp_2x, 6), + "amp_3x": round(features.amp_3x, 6), + "amp_half_x": round(features.amp_half_x, 6), + "kurtosis": round(features.kurtosis, 2), + "crest_factor": round(features.crest_factor, 2), + }, + "signal_statistics": { + "rms_g": round(ts_rms, 6), + "peak_g": round(ts_peak, 6), + "crest_factor": round(ts_crest, 2), + "kurtosis": round(ts_kurtosis, 2), + "duration_s": round(len(sig) / sr, 4), + "sample_rate_hz": sr, + "rms_velocity_mm_s": round(rms_vel_mms, 3), + }, + "bearing_analysis": ( + {k: v for k, v in bearing_results.items()} if bearing_results else None + ), + "bearing_info_source": bearing_info_source, + "report_markdown": report, + } + + +def main(): + mcp.run(transport="stdio") + + +if __name__ == "__main__": + main() diff --git a/src/servers/vibration/tests/__init__.py b/src/servers/vibration/tests/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/servers/vibration/tests/conftest.py b/src/servers/vibration/tests/conftest.py new file mode 100644 index 00000000..e723281c --- /dev/null +++ b/src/servers/vibration/tests/conftest.py @@ -0,0 +1,57 @@ +"""Fixtures and helpers for VibrationAgent tests.""" + +import json +import os + +from dotenv import load_dotenv +import pytest +from unittest.mock import patch + +load_dotenv() + +# --- Custom markers --- + +requires_couchdb = pytest.mark.skipif( + os.environ.get("COUCHDB_URL") is None, + reason="CouchDB not available (set COUCHDB_URL)", +) + + +# --- Fixtures --- + + +@pytest.fixture +def mock_db(): + """Patch the module-level _get_db in couchdb_client with a mock.""" + with patch("servers.vibration.couchdb_client._get_db") as mock: + yield mock + + +@pytest.fixture +def no_db(): + """Patch _get_db to return None (simulate disconnected CouchDB).""" + with patch("servers.vibration.couchdb_client._get_db", return_value=None): + yield + + +@pytest.fixture(autouse=True) +def reset_data_store(): + """Clear the in-memory data store between tests.""" + from servers.vibration.data_store import store + + store._entries.clear() + yield + store._entries.clear() + + +async def call_tool(mcp_instance, tool_name: str, args: dict) -> dict: + """Helper: call an MCP tool and return parsed JSON response.""" + result = await mcp_instance.call_tool(tool_name, args) + # FastMCP may return (contents, is_error) tuple or just the contents list + if isinstance(result, tuple): + contents = result[0] + else: + contents = result + if isinstance(contents, dict): + return contents + return json.loads(contents[0].text) diff --git a/src/servers/vibration/tests/test_dsp.py b/src/servers/vibration/tests/test_dsp.py new file mode 100644 index 00000000..181d93d0 --- /dev/null +++ b/src/servers/vibration/tests/test_dsp.py @@ -0,0 +1,245 @@ +"""Pure-function unit tests for the DSP layer — no MCP, no CouchDB.""" + +import math + +import numpy as np +import pytest + +from servers.vibration.dsp.fft_analysis import ( + compute_fft, + compute_psd, + compute_spectrogram, + find_peaks_in_spectrum, +) +from servers.vibration.dsp.envelope import ( + bandpass_filter, + compute_envelope, + envelope_spectrum, + check_bearing_peaks, +) +from servers.vibration.dsp.bearing_freqs import ( + compute_bearing_frequencies, + get_bearing, + list_bearings, +) +from servers.vibration.dsp.fault_detection import ( + assess_iso10816, + extract_shaft_features, + classify_faults, + generate_diagnosis_summary, +) + + +# --------------------------------------------------------------------------- +# Synthetic signals +# --------------------------------------------------------------------------- + +SR = 4096.0 +DURATION = 2.0 +T = np.arange(0, DURATION, 1.0 / SR) +SINE_50 = np.sin(2 * np.pi * 50 * T) +SINE_120 = np.sin(2 * np.pi * 120 * T) +COMPOSITE = SINE_50 + 0.5 * SINE_120 + + +# =================================================================== +# fft_analysis +# =================================================================== + + +class TestComputeFFT: + def test_pure_sine(self): + result = compute_fft(SINE_50, SR) + freqs = np.array(result["frequencies"]) + mags = np.array(result["magnitude"]) + peak_idx = np.argmax(mags) + assert abs(freqs[peak_idx] - 50.0) < 2.0 + + def test_window_types(self): + for win in ("hann", "hamming", "blackman", "rectangular"): + result = compute_fft(SINE_50, SR, window=win) + assert len(result["frequencies"]) == len(result["magnitude"]) + + def test_frequency_resolution(self): + result = compute_fft(SINE_50, SR) + freqs = np.array(result["frequencies"]) + expected_res = SR / len(SINE_50) + assert abs(freqs[1] - expected_res) < 0.01 + + +class TestComputePSD: + def test_basic(self): + result = compute_psd(SINE_50, SR) + assert "frequencies" in result + assert "psd" in result + assert len(result["frequencies"]) == len(result["psd"]) + + +class TestComputeSpectrogram: + def test_basic(self): + result = compute_spectrogram(SINE_50, SR) + assert "frequencies" in result + assert "times" in result + assert "spectrogram_db" in result + + +class TestFindPeaks: + def test_find_50hz_peak(self): + fft = compute_fft(SINE_50, SR) + freqs = np.array(fft["frequencies"]) + mags = np.array(fft["magnitude"]) + peaks = find_peaks_in_spectrum(freqs, mags, num_peaks=5) + top_freq = peaks[0]["frequency_hz"] + assert abs(top_freq - 50.0) < 2.0 + + +# =================================================================== +# envelope +# =================================================================== + + +class TestBandpassFilter: + def test_shape_preserved(self): + filtered = bandpass_filter(SINE_50, SR, 30.0, 70.0) + assert len(filtered) == len(SINE_50) + + +class TestComputeEnvelope: + def test_shape_and_non_negative(self): + env = compute_envelope(SINE_50) + assert len(env) == len(SINE_50) + assert np.all(env >= 0) + + +class TestEnvelopeSpectrum: + def test_basic(self): + result = envelope_spectrum(SINE_120, SR) + assert "frequencies" in result + assert "envelope_spectrum" in result + assert "filter_band" in result + + def test_custom_band(self): + result = envelope_spectrum(SINE_120, SR, band_low=50.0, band_high=200.0) + lo, hi = result["filter_band"] + assert lo == 50.0 + assert hi == 200.0 + + +class TestCheckBearingPeaks: + def test_finds_peak(self): + # Signal with a clear 100 Hz component + sig = np.sin(2 * np.pi * 100 * T) + env = envelope_spectrum(sig, SR) + freqs = np.array(env["frequencies"]) + mags = np.array(env["envelope_spectrum"]) + result = check_bearing_peaks(freqs, mags, 100.0) + assert "details" in result + assert result["harmonics_detected"] >= 1 + + +# =================================================================== +# bearing_freqs +# =================================================================== + + +class TestComputeBearingFrequencies: + def test_6205_at_1800rpm(self): + bf = compute_bearing_frequencies(1800, 9, 7.94, 39.04, 0.0, "6205") + assert bf.bpfo > 0 + assert bf.bpfi > bf.bpfo # inner race always higher + assert bf.bsf > 0 + assert bf.ftf > 0 + assert bf.ftf < bf.bpfo # cage freq is lowest + + def test_zero_rpm(self): + bf = compute_bearing_frequencies(0, 9, 7.94, 39.04, 0.0, "test") + assert bf.bpfo == 0.0 + assert bf.bpfi == 0.0 + + def test_to_dict(self): + bf = compute_bearing_frequencies(3600, 9, 7.94, 39.04, 0.0, "test") + d = bf.to_dict() + for key in ("bpfo_hz", "bpfi_hz", "bsf_hz", "ftf_hz", "rpm", "bearing"): + assert key in d + + +class TestGetBearing: + def test_known_bearing(self): + b = get_bearing("6205") + assert b is not None + assert b.n_balls == 9 + + def test_unknown_bearing(self): + b = get_bearing("NONEXISTENT-999") + assert b is None + + +class TestListBearings: + def test_returns_list(self): + bearings = list_bearings() + assert len(bearings) >= 5 + assert all("name" in b for b in bearings) + + +# =================================================================== +# fault_detection +# =================================================================== + + +class TestAssessISO10816: + def test_zone_a(self): + result = assess_iso10816(0.5, "group2") + assert result["iso_zone"] == "A" + + def test_zone_b(self): + result = assess_iso10816(2.0, "group2") + assert result["iso_zone"] == "B" + + def test_zone_c(self): + result = assess_iso10816(6.0, "group2") + assert result["iso_zone"] == "C" + + def test_zone_d(self): + result = assess_iso10816(15.0, "group2") + assert result["iso_zone"] == "D" + + def test_all_groups(self): + for grp in ("group1", "group2", "group3", "group4"): + r = assess_iso10816(4.5, grp) + assert r["iso_zone"] in ("A", "B", "C", "D") + + def test_unknown_group_defaults(self): + r = assess_iso10816(4.5, "group_INVALID") + assert r["iso_zone"] in ("A", "B", "C", "D") + + +class TestExtractShaftFeatures: + def test_basic(self): + fft = compute_fft(COMPOSITE, SR) + freqs = np.array(fft["frequencies"]) + mags = np.array(fft["magnitude"]) + shaft_freq = 50.0 # as if rpm=3000 + features = extract_shaft_features(freqs, mags, shaft_freq, + time_signal=COMPOSITE) + assert features.f_shaft == 50.0 + assert features.amp_1x > 0 + + +class TestClassifyFaults: + def test_empty_features_returns_healthy(self): + # All amplitudes zero → healthy or very low vibration + fft = compute_fft(np.zeros(1000), SR) + freqs = np.array(fft["frequencies"]) + mags = np.array(fft["magnitude"]) + features = extract_shaft_features(freqs, mags, 30.0) + diagnoses = classify_faults(features) + # Should not crash; may return healthy or empty + assert isinstance(diagnoses, list) + + +class TestGenerateDiagnosisSummary: + def test_renders_markdown(self): + iso = assess_iso10816(4.5, "group2") + md = generate_diagnosis_summary([], iso, "Test pump") + assert "ISO 10816" in md or "10816" in md + assert isinstance(md, str) diff --git a/src/servers/vibration/tests/test_tools.py b/src/servers/vibration/tests/test_tools.py new file mode 100644 index 00000000..b53d9fc2 --- /dev/null +++ b/src/servers/vibration/tests/test_tools.py @@ -0,0 +1,263 @@ +"""Tests for VibrationAgent MCP server tools. + +Unit tests use synthetic signals; integration tests require a live CouchDB +(skipped unless COUCHDB_URL is set). +""" + +import numpy as np +import pytest + +from servers.vibration.data_store import store +from servers.vibration.main import mcp +from .conftest import call_tool, requires_couchdb + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_sine(freq_hz: float = 50.0, sr: float = 2048.0, + duration: float = 1.0, amplitude: float = 1.0) -> tuple: + """Generate a pure sine wave and store it; return (data_id, signal, sr).""" + t = np.arange(0, duration, 1.0 / sr) + sig = amplitude * np.sin(2 * np.pi * freq_hz * t) + data_id = f"test_sine_{freq_hz}Hz" + store.put(data_id, sig, sr, {"test": True}) + return data_id, sig, sr + + +def _make_composite(freqs: list[float], sr: float = 4096.0, + duration: float = 2.0) -> str: + """Composite signal with multiple sine components; returns data_id.""" + t = np.arange(0, duration, 1.0 / sr) + sig = np.zeros_like(t) + for f in freqs: + sig += np.sin(2 * np.pi * f * t) + data_id = "test_composite" + store.put(data_id, sig, sr, {"freqs": freqs}) + return data_id + + +# --------------------------------------------------------------------------- +# compute_fft_spectrum +# --------------------------------------------------------------------------- + + +class TestComputeFFTSpectrum: + @pytest.mark.anyio + async def test_basic_50hz(self): + data_id, _, _ = _make_sine(50.0) + result = await call_tool(mcp, "compute_fft_spectrum", {"data_id": data_id}) + assert "error" not in result + assert result["sample_rate_hz"] == 2048.0 + # Dominant peak should be near 50 Hz + top = result["top_peaks"][0] + assert abs(top["freq_hz"] - 50.0) < 3.0 + + @pytest.mark.anyio + async def test_missing_data_id(self): + result = await call_tool(mcp, "compute_fft_spectrum", + {"data_id": "nonexistent"}) + assert "error" in result + + @pytest.mark.anyio + async def test_window_types(self): + data_id, _, _ = _make_sine(100.0) + for win in ("hann", "hamming", "blackman", "rectangular"): + result = await call_tool(mcp, "compute_fft_spectrum", + {"data_id": data_id, "window": win}) + assert "error" not in result + assert result["window"] == win + + +# --------------------------------------------------------------------------- +# compute_envelope_spectrum +# --------------------------------------------------------------------------- + + +class TestComputeEnvelopeSpectrum: + @pytest.mark.anyio + async def test_basic_run(self): + data_id, _, _ = _make_sine(120.0, sr=4096.0) + result = await call_tool(mcp, "compute_envelope_spectrum", + {"data_id": data_id}) + assert "error" not in result + assert "filter_band_hz" in result + assert result["sample_rate_hz"] == 4096.0 + + @pytest.mark.anyio + async def test_missing_data_id(self): + result = await call_tool(mcp, "compute_envelope_spectrum", + {"data_id": "nope"}) + assert "error" in result + + +# --------------------------------------------------------------------------- +# assess_vibration_severity (ISO 10816) +# --------------------------------------------------------------------------- + + +class TestAssessVibrationSeverity: + @pytest.mark.anyio + async def test_zone_a(self): + result = await call_tool(mcp, "assess_vibration_severity", + {"rms_velocity_mm_s": 0.5}) + assert result["iso_zone"] == "A" + + @pytest.mark.anyio + async def test_zone_d(self): + result = await call_tool(mcp, "assess_vibration_severity", + {"rms_velocity_mm_s": 50.0}) + assert result["iso_zone"] == "D" + + @pytest.mark.anyio + async def test_group_param(self): + for grp in ("group1", "group2", "group3", "group4"): + result = await call_tool(mcp, "assess_vibration_severity", + {"rms_velocity_mm_s": 4.5, + "machine_group": grp}) + assert result["iso_zone"] in ("A", "B", "C", "D") + + +# --------------------------------------------------------------------------- +# list_known_bearings +# --------------------------------------------------------------------------- + + +class TestListKnownBearings: + @pytest.mark.anyio + async def test_returns_bearings(self): + result = await call_tool(mcp, "list_known_bearings", {}) + assert "bearings" in result + assert len(result["bearings"]) >= 5 + names = [b["name"] for b in result["bearings"]] + assert any("6205" in n for n in names) + + +# --------------------------------------------------------------------------- +# calculate_bearing_frequencies +# --------------------------------------------------------------------------- + + +class TestCalculateBearingFrequencies: + @pytest.mark.anyio + async def test_basic(self): + result = await call_tool(mcp, "calculate_bearing_frequencies", { + "rpm": 1800, + "n_balls": 9, + "ball_diameter_mm": 7.94, + "pitch_diameter_mm": 39.04, + "contact_angle_deg": 0.0, + }) + assert "bpfo_hz" in result + assert "bpfi_hz" in result + assert "bsf_hz" in result + assert "ftf_hz" in result + assert result["bpfo_hz"] > 0 + + @pytest.mark.anyio + async def test_with_name(self): + result = await call_tool(mcp, "calculate_bearing_frequencies", { + "rpm": 3600, + "n_balls": 8, + "ball_diameter_mm": 10.0, + "pitch_diameter_mm": 46.0, + "bearing_name": "test-bearing", + }) + assert "bearing" in result + assert result["bearing"] == "test-bearing" + + +# --------------------------------------------------------------------------- +# diagnose_vibration +# --------------------------------------------------------------------------- + + +class TestDiagnoseVibration: + @pytest.mark.anyio + async def test_no_rpm(self): + """Without RPM we expect a partial result with a warning.""" + data_id, _, _ = _make_sine(120.0, sr=4096.0, duration=2.0) + result = await call_tool(mcp, "diagnose_vibration", { + "data_id": data_id, + }) + assert "error" not in result + assert "warning" in result + assert result["shaft_features"] is None + + @pytest.mark.anyio + async def test_with_rpm(self): + data_id = _make_composite([30, 60, 90], sr=4096.0, duration=2.0) + result = await call_tool(mcp, "diagnose_vibration", { + "data_id": data_id, + "rpm": 1800.0, + }) + assert "error" not in result + assert result["shaft_features"] is not None + assert result["iso_10816"] is not None + assert "report_markdown" in result + + @pytest.mark.anyio + async def test_with_bearing_designation(self): + data_id = _make_composite([30, 60, 120], sr=4096.0, duration=2.0) + result = await call_tool(mcp, "diagnose_vibration", { + "data_id": data_id, + "rpm": 1800.0, + "bearing_designation": "6205", + }) + assert "error" not in result + assert result["bearing_info_source"] is not None + assert "database" in result["bearing_info_source"] + + @pytest.mark.anyio + async def test_with_custom_bearing_geometry(self): + data_id = _make_composite([30, 60], sr=4096.0, duration=2.0) + result = await call_tool(mcp, "diagnose_vibration", { + "data_id": data_id, + "rpm": 1800.0, + "bearing_n_balls": 9, + "bearing_ball_dia_mm": 7.94, + "bearing_pitch_dia_mm": 39.04, + }) + assert "error" not in result + assert result["bearing_info_source"] == "custom geometry" + + @pytest.mark.anyio + async def test_missing_data_id(self): + result = await call_tool(mcp, "diagnose_vibration", + {"data_id": "ghost"}) + assert "error" in result + + +# --------------------------------------------------------------------------- +# get_vibration_data (integration — requires CouchDB) +# --------------------------------------------------------------------------- + + +class TestGetVibrationData: + @requires_couchdb + @pytest.mark.anyio + async def test_fetch_integration(self): + result = await call_tool(mcp, "get_vibration_data", { + "site_name": "MAIN", + "asset_id": "Chiller 6", + "sensor_name": "Current", + "start": "2020-06-01T00:00:00", + }) + assert "error" not in result or "data_id" in result + + +# --------------------------------------------------------------------------- +# list_vibration_sensors (integration — requires CouchDB) +# --------------------------------------------------------------------------- + + +class TestListVibrationSensors: + @requires_couchdb + @pytest.mark.anyio + async def test_list_integration(self): + result = await call_tool(mcp, "list_vibration_sensors", { + "site_name": "MAIN", + "asset_id": "Chiller 6", + }) + assert "sensors" in result or "error" in result diff --git a/src/tmp/assetopsbench/scenarios/single_agent/vibration_utterance.json b/src/tmp/assetopsbench/scenarios/single_agent/vibration_utterance.json new file mode 100644 index 00000000..74a1c777 --- /dev/null +++ b/src/tmp/assetopsbench/scenarios/single_agent/vibration_utterance.json @@ -0,0 +1,142 @@ +[ + { + "id": 301, + "type": "Vibration", + "text": "What vibration analysis capabilities are available?", + "category": "Knowledge Query", + "characteristic_form": "The expected response should list: FFT spectrum analysis, envelope analysis for bearing faults, bearing characteristic frequency calculation (BPFO/BPFI/BSF/FTF), ISO 10816 vibration severity assessment, and full automated diagnosis." + }, + { + "id": 302, + "type": "Vibration", + "text": "What bearings are available in the built-in database?", + "category": "Knowledge Query", + "characteristic_form": "The expected response should include bearing designations such as 6205, 6206, 6207, 6208, 6305, 6306, NU205, NU206, 7205 with their geometric parameters (number of balls, ball diameter, pitch diameter)." + }, + { + "id": 303, + "type": "Vibration", + "text": "How does the ISO 10816 vibration severity classification work?", + "category": "Knowledge Query", + "characteristic_form": "The expected response should describe four zones: A (good), B (acceptable), C (barely tolerable), D (not permissible), and four machine groups (group1 through group4) with different velocity thresholds in mm/s." + }, + { + "id": 304, + "type": "Vibration", + "text": "Calculate the bearing characteristic frequencies for a 6205 bearing running at 1800 RPM.", + "category": "Bearing Analysis", + "characteristic_form": "The expected response should provide BPFO, BPFI, BSF, and FTF frequencies in Hz computed from the 6205 geometry (9 balls, ball_dia=7.94 mm, pitch_dia=39.04 mm) at 1800 RPM." + }, + { + "id": 305, + "type": "Vibration", + "text": "Calculate bearing fault frequencies for a bearing with 8 balls, 10.3 mm ball diameter, 46.0 mm pitch diameter, 15 degree contact angle at 3600 RPM.", + "category": "Bearing Analysis", + "characteristic_form": "The expected response should provide the computed BPFO, BPFI, BSF, and FTF values in Hz using the provided custom geometry and RPM." + }, + { + "id": 306, + "type": "Vibration", + "text": "What is the vibration severity classification for a machine with an RMS velocity of 4.5 mm/s? It is a medium-sized machine on rigid foundations.", + "category": "ISO Assessment", + "characteristic_form": "The expected response should classify 4.5 mm/s as ISO 10816 Zone B (acceptable) for a group2 machine, with thresholds context." + }, + { + "id": 307, + "type": "Vibration", + "text": "Assess vibration severity for a large machine on flexible foundations with 12.0 mm/s RMS velocity.", + "category": "ISO Assessment", + "characteristic_form": "The expected response should classify 12.0 mm/s under group3 (large machine, flexible foundation) and determine the appropriate ISO 10816 zone." + }, + { + "id": 308, + "type": "Vibration", + "text": "Fetch vibration sensor data from Chiller 6, sensor Current, starting from 2020-06-01.", + "category": "Data Retrieval", + "characteristic_form": "The expected response should load time-series data from CouchDB for asset 'Chiller 6', sensor 'Current', return a data_id, sample count, duration, and basic statistics (RMS, peak, mean)." + }, + { + "id": 309, + "type": "Vibration", + "text": "What sensors are available for Chiller 6 at site MAIN?", + "category": "Data Retrieval", + "characteristic_form": "The expected response should list available sensor fields for Chiller 6 from CouchDB." + }, + { + "id": 310, + "type": "Vibration", + "text": "Compute the FFT spectrum for the vibration data I just loaded.", + "category": "Signal Analysis", + "characteristic_form": "The expected response should provide the top frequency peaks with amplitudes, maximum amplitude and its frequency, RMS spectral value, frequency resolution, and total spectral bins." + }, + { + "id": 311, + "type": "Vibration", + "text": "Generate the FFT spectrum using a Blackman window for the loaded signal.", + "category": "Signal Analysis", + "characteristic_form": "The expected response should compute an FFT with a Blackman window and return peak frequencies, confirming the window type used." + }, + { + "id": 312, + "type": "Vibration", + "text": "Perform envelope analysis on the vibration signal to look for bearing defect frequencies.", + "category": "Signal Analysis", + "characteristic_form": "The expected response should compute the envelope spectrum using the Hilbert transform, optionally report the bandpass filter band, and show the dominant peaks in the demodulated spectrum." + }, + { + "id": 313, + "type": "Vibration", + "text": "Run a full vibration diagnosis on the loaded data at 1800 RPM with bearing 6205.", + "category": "Diagnosis", + "characteristic_form": "The expected response should include: shaft feature extraction (1x, 2x, 3x amplitudes), bearing fault frequency analysis using 6205 geometry, ISO 10816 severity assessment, fault classification (unbalance, misalignment, looseness, bearing), and a markdown diagnostic report." + }, + { + "id": 314, + "type": "Vibration", + "text": "Diagnose the vibration signal at 3600 RPM using custom bearing geometry: 8 balls, 10.3 mm ball diameter, 46 mm pitch diameter.", + "category": "Diagnosis", + "characteristic_form": "The expected response should perform full diagnosis including custom bearing frequency calculation, envelope analysis, shaft features, ISO 10816 assessment, and fault classification with a summary report." + }, + { + "id": 315, + "type": "Vibration", + "text": "Run vibration diagnosis without RPM on the current data.", + "category": "Diagnosis", + "characteristic_form": "The expected response should provide a partial diagnosis with a warning that shaft-frequency analysis was skipped, basic signal statistics (RMS, peak, crest factor, kurtosis), FFT summary, and ISO 10816 severity." + }, + { + "id": 316, + "type": "Vibration", + "text": "Diagnose the vibration and provide the bearing fault frequencies directly: BPFO=87.5 Hz, BPFI=132.5 Hz, BSF=57.3 Hz, FTF=9.7 Hz at 1800 RPM.", + "category": "Diagnosis", + "characteristic_form": "The expected response should use the provided fault frequencies directly (rather than computing them), perform envelope analysis, and check for peaks at those specific frequencies and their harmonics." + }, + { + "id": 317, + "type": "Vibration", + "text": "Is unbalance detected in the vibration signal? The machine runs at 1500 RPM.", + "category": "Fault Classification", + "characteristic_form": "The expected response should analyze the 1x shaft frequency (25 Hz) amplitude relative to 2x and 3x harmonics. Unbalance is indicated by a dominant 1x component." + }, + { + "id": 318, + "type": "Vibration", + "text": "Check if there are signs of shaft misalignment in the vibration data. RPM is 3000.", + "category": "Fault Classification", + "characteristic_form": "The expected response should check if the 2x shaft frequency component (100 Hz) is elevated relative to 1x (50 Hz), which indicates shaft misalignment." + }, + { + "id": 319, + "type": "Vibration", + "text": "Analyze the vibration signal for mechanical looseness at 1800 RPM.", + "category": "Fault Classification", + "characteristic_form": "The expected response should check for sub-harmonic (0.5x) and higher harmonics (3x+) of the shaft frequency, elevated crest factor, and presence of half-shaft-frequency component." + }, + { + "id": 320, + "type": "Vibration", + "text": "Check for outer race bearing fault (BPFO) in the vibration signal. Shaft speed 1800 RPM, bearing model 6205.", + "category": "Fault Classification", + "characteristic_form": "The expected response should compute BPFO from 6205 geometry at 1800 RPM, perform envelope analysis, and check for peaks at BPFO and its harmonics (2x, 3x BPFO)." + } + ] diff --git a/src/workflow/executor.py b/src/workflow/executor.py index e84703d0..884a8c6f 100644 --- a/src/workflow/executor.py +++ b/src/workflow/executor.py @@ -35,6 +35,7 @@ "Utilities": "utilities-mcp-server", "FMSRAgent": "fmsr-mcp-server", "TSFMAgent": "tsfm-mcp-server", + "VibrationAgent": "vibration-mcp-server", } _PLACEHOLDER_RE = re.compile(r"\{step_(\d+)\}")