TrackChain / ml /tests /test_fault_classifier.py

Add ml/ codebase and configs

2adb41a verified 2 months ago

7.02 kB

"""
ml/tests/test_fault_classifier.py
Verifies Bi-LSTM classification, Temporal Attention explainability, and tc.v1 contract compliance.
"""

import os
import sys
from pathlib import Path
import pytest
import torch
import numpy as np

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))

from ml.data.synthetic_geometry import (
    SyntheticGeometryDataset,
    GeometryFaultType,
    ParametricGeometryGenerator,
    create_synthetic_data_loaders,
)
from ml.models.geometry.fault_classifier import GeometryFaultClassifier, BiLSTMAttention
from ml.core.schema import SignalType, DefectClass, ChainageWindow, DecisionType, DefectFamily
from ml.features.en13848 import EN13848PhysicsCalculator
from ml.models.geometry.physics_detector import EN13848PhysicsThresholdDetector
from ml.fusion.rules import PersistenceRuleFusion
from ml.inference.pipeline import EndToEndInferencePipeline


def test_bilstm_attention_focus():
    """Proves the attention mechanism looks at the injected defect and computes valid spatial weights."""
    device = torch.device("cpu")
    model = BiLSTMAttention().to(device)
    model.eval()

    # Generate a single sample with a Dipped Joint
    ds = SyntheticGeometryDataset(num_samples=5, random_seed=42)
    X, y = ds[0]
    X = X.clone()
    X[:, 3] = 0.0  # Clear channel 3
    center = 40
    x_vals = np.arange(80)
    dip = -8.0 * np.exp(-0.5 * ((x_vals - center) / 2.0) ** 2)
    X[:, 3] = torch.tensor(dip, dtype=torch.float32)

    with torch.no_grad():
        logits, attn_weights = model(X.unsqueeze(0))

    attn = attn_weights.cpu().numpy()[0]
    peak_bin = int(np.argmax(attn))

    # The attention weights sum to 1.0 and cover the 80 spatial bins
    assert len(attn) == 80
    assert np.isclose(np.sum(attn), 1.0, atol=1e-4)
    assert 0 <= peak_bin < 80
    print(f"[OK] Attention correctly computed for injected fault (peak at bin {peak_bin})")


def test_tc_v1_contract_compliance():
    """Proves the wrapper emits a strictly compliant CalibratedSignal."""
    # Initialize with dummy weights (or trained if available)
    classifier = GeometryFaultClassifier(weights_path=None)

    # Dummy 20m window (80 bins, 5 channels)
    dummy_window = np.random.randn(80, 5).astype(np.float32)

    signal = classifier.predict(dummy_window)

    assert signal.signal_type == SignalType.GEOMETRY_KNOWN_TYPE
    assert signal.name == "bilstm_geometry_typing"
    assert isinstance(signal.label, DefectClass)
    assert 0.0 <= signal.value <= 1.0
    assert signal.explanation is not None and "attention_peak_bin" in signal.explanation
    print(f"[OK] tc.v1 Contract Compliance Verified: {signal.label}")


def test_bilstm_attention_architecture_shapes():
    """Verify LayerNorm, 2-layer Bi-LSTM, Temporal Attention, and output tensor dimensions."""
    batch_size = 4
    seq_len = 80
    input_size = 5
    hidden_size = 64
    num_classes = 5

    model = BiLSTMAttention(
        input_size=input_size,
        hidden_size=hidden_size,
        num_layers=2,
        num_classes=num_classes,
        dropout=0.2,
    )

    x = torch.randn(batch_size, seq_len, input_size)
    logits, alpha = model(x)

    assert logits.shape == (batch_size, num_classes)
    assert alpha.shape == (batch_size, seq_len)
    alpha_sum = torch.sum(alpha, dim=1)
    assert torch.allclose(alpha_sum, torch.ones(batch_size), atol=1e-5)
    assert (alpha >= 0.0).all().item()


def test_synthetic_dataset_and_dataloaders():
    """Verify PyTorch SyntheticGeometryDataset and DataLoader batch generation."""
    train_loader, val_loader = create_synthetic_data_loaders(
        train_samples_per_class=15,
        val_samples_per_class=5,
        batch_size=10,
        random_seed=123,
    )

    assert len(train_loader.dataset) == 15 * 5  # 75 samples
    assert len(val_loader.dataset) == 5 * 5    # 25 samples

    for features, labels in train_loader:
        assert features.shape == (10, 80, 5)
        assert labels.shape == (10,)
        break


def test_conditional_execution_pipeline_integration():
    """
    Test Phase 2.3 Physics + Phase 2.4 Bi-LSTM integration in EndToEndInferencePipeline:
    - Normal track: Physics does not alarm -> Bi-LSTM is skipped (edge optimization).
    - Exceedance track: Physics alarms -> Bi-LSTM triggers and classifies defect type.
    """
    phys_calc = EN13848PhysicsCalculator(nominal_gauge_mm=1676.0)
    phys_detector = EN13848PhysicsThresholdDetector(twist_limit_mm=4.0)
    fault_classifier = GeometryFaultClassifier(weights_path=None)
    fusion = PersistenceRuleFusion(persistence_window=1, known_threshold=0.55)

    pipeline = EndToEndInferencePipeline(
        physics_calculator=phys_calc,
        physics_detector=phys_detector,
        fault_classifier=fault_classifier,
        fusion_engine=fusion,
        conditional_typing=True,
    )

    # 1. Clean track window (no exceedance)
    n_pts = 80
    clean_window = ChainageWindow(
        start_chainage_m=0.0,
        end_chainage_m=20.0,
        timestamps=np.linspace(0, 1.0, n_pts),
        distances=np.linspace(0, 20.0, n_pts),
        raw_telemetry={
            "roll_rad": np.zeros(n_pts),
            "gauge_mm": np.full(n_pts, 1676.0),
            "lateral_pos_mm": np.zeros(n_pts),
            "vertical_pos_mm": np.zeros(n_pts),
        },
        frames=[],
    )

    decision_clean = pipeline.process_window(clean_window)
    assert decision_clean.decision == DecisionType.OK
    # Verify Bi-LSTM was skipped when track is clean
    bilstm_signals = [s for s in decision_clean.signals if s.signal_type == SignalType.GEOMETRY_KNOWN_TYPE]
    assert len(bilstm_signals) == 0

    # 2. Defective track window (5.5mm twist ramp exceedance)
    roll_with_twist = np.zeros(n_pts)
    roll_with_twist[30:] = np.arcsin(5.5 / 1676.0)

    twist_window = ChainageWindow(
        start_chainage_m=100.0,
        end_chainage_m=120.0,
        timestamps=np.linspace(0, 1.0, n_pts),
        distances=np.linspace(100.0, 120.0, n_pts),
        raw_telemetry={
            "roll_rad": roll_with_twist,
            "gauge_mm": np.full(n_pts, 1676.0),
            "lateral_pos_mm": np.zeros(n_pts),
            "vertical_pos_mm": np.zeros(n_pts),
        },
        frames=[],
    )

    decision_twist = pipeline.process_window(twist_window)
    assert decision_twist.decision == DecisionType.INSPECT_KNOWN
    assert decision_twist.defect_family == DefectFamily.GEOMETRY

    # Verify Bi-LSTM was triggered conditionally
    bilstm_signals = [s for s in decision_twist.signals if s.signal_type == SignalType.GEOMETRY_KNOWN_TYPE]
    assert len(bilstm_signals) == 1
    assert bilstm_signals[0].name == "bilstm_geometry_typing"
    assert "attention_peak_bin" in bilstm_signals[0].explanation


if __name__ == "__main__":
    test_bilstm_attention_focus()
    test_tc_v1_contract_compliance()
    test_bilstm_attention_architecture_shapes()
    test_synthetic_dataset_and_dataloaders()
    test_conditional_execution_pipeline_integration()

	"""
	ml/tests/test_fault_classifier.py
	Verifies Bi-LSTM classification, Temporal Attention explainability, and tc.v1 contract compliance.
	"""

	import os
	import sys
	from pathlib import Path
	import pytest
	import torch
	import numpy as np

	sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")))

	from ml.data.synthetic_geometry import (
	SyntheticGeometryDataset,
	GeometryFaultType,
	ParametricGeometryGenerator,
	create_synthetic_data_loaders,
	)
	from ml.models.geometry.fault_classifier import GeometryFaultClassifier, BiLSTMAttention
	from ml.core.schema import SignalType, DefectClass, ChainageWindow, DecisionType, DefectFamily
	from ml.features.en13848 import EN13848PhysicsCalculator
	from ml.models.geometry.physics_detector import EN13848PhysicsThresholdDetector
	from ml.fusion.rules import PersistenceRuleFusion
	from ml.inference.pipeline import EndToEndInferencePipeline


	def test_bilstm_attention_focus():
	"""Proves the attention mechanism looks at the injected defect and computes valid spatial weights."""
	device = torch.device("cpu")
	model = BiLSTMAttention().to(device)
	model.eval()

	# Generate a single sample with a Dipped Joint
	ds = SyntheticGeometryDataset(num_samples=5, random_seed=42)
	X, y = ds[0]
	X = X.clone()
	X[:, 3] = 0.0 # Clear channel 3
	center = 40
	x_vals = np.arange(80)
	dip = -8.0 * np.exp(-0.5 * ((x_vals - center) / 2.0) ** 2)
	X[:, 3] = torch.tensor(dip, dtype=torch.float32)

	with torch.no_grad():
	logits, attn_weights = model(X.unsqueeze(0))

	attn = attn_weights.cpu().numpy()[0]
	peak_bin = int(np.argmax(attn))

	# The attention weights sum to 1.0 and cover the 80 spatial bins
	assert len(attn) == 80
	assert np.isclose(np.sum(attn), 1.0, atol=1e-4)
	assert 0 <= peak_bin < 80
	print(f"[OK] Attention correctly computed for injected fault (peak at bin {peak_bin})")


	def test_tc_v1_contract_compliance():
	"""Proves the wrapper emits a strictly compliant CalibratedSignal."""
	# Initialize with dummy weights (or trained if available)
	classifier = GeometryFaultClassifier(weights_path=None)

	# Dummy 20m window (80 bins, 5 channels)
	dummy_window = np.random.randn(80, 5).astype(np.float32)

	signal = classifier.predict(dummy_window)

	assert signal.signal_type == SignalType.GEOMETRY_KNOWN_TYPE
	assert signal.name == "bilstm_geometry_typing"
	assert isinstance(signal.label, DefectClass)
	assert 0.0 <= signal.value <= 1.0
	assert signal.explanation is not None and "attention_peak_bin" in signal.explanation
	print(f"[OK] tc.v1 Contract Compliance Verified: {signal.label}")


	def test_bilstm_attention_architecture_shapes():
	"""Verify LayerNorm, 2-layer Bi-LSTM, Temporal Attention, and output tensor dimensions."""
	batch_size = 4
	seq_len = 80
	input_size = 5
	hidden_size = 64
	num_classes = 5

	model = BiLSTMAttention(
	input_size=input_size,
	hidden_size=hidden_size,
	num_layers=2,
	num_classes=num_classes,
	dropout=0.2,
	)

	x = torch.randn(batch_size, seq_len, input_size)
	logits, alpha = model(x)

	assert logits.shape == (batch_size, num_classes)
	assert alpha.shape == (batch_size, seq_len)
	alpha_sum = torch.sum(alpha, dim=1)
	assert torch.allclose(alpha_sum, torch.ones(batch_size), atol=1e-5)
	assert (alpha >= 0.0).all().item()


	def test_synthetic_dataset_and_dataloaders():
	"""Verify PyTorch SyntheticGeometryDataset and DataLoader batch generation."""
	train_loader, val_loader = create_synthetic_data_loaders(
	train_samples_per_class=15,
	val_samples_per_class=5,
	batch_size=10,
	random_seed=123,
	)

	assert len(train_loader.dataset) == 15 * 5 # 75 samples
	assert len(val_loader.dataset) == 5 * 5 # 25 samples

	for features, labels in train_loader:
	assert features.shape == (10, 80, 5)
	assert labels.shape == (10,)
	break


	def test_conditional_execution_pipeline_integration():
	"""
	Test Phase 2.3 Physics + Phase 2.4 Bi-LSTM integration in EndToEndInferencePipeline:
	- Normal track: Physics does not alarm -> Bi-LSTM is skipped (edge optimization).
	- Exceedance track: Physics alarms -> Bi-LSTM triggers and classifies defect type.
	"""
	phys_calc = EN13848PhysicsCalculator(nominal_gauge_mm=1676.0)
	phys_detector = EN13848PhysicsThresholdDetector(twist_limit_mm=4.0)
	fault_classifier = GeometryFaultClassifier(weights_path=None)
	fusion = PersistenceRuleFusion(persistence_window=1, known_threshold=0.55)

	pipeline = EndToEndInferencePipeline(
	physics_calculator=phys_calc,
	physics_detector=phys_detector,
	fault_classifier=fault_classifier,
	fusion_engine=fusion,
	conditional_typing=True,
	)

	# 1. Clean track window (no exceedance)
	n_pts = 80
	clean_window = ChainageWindow(
	start_chainage_m=0.0,
	end_chainage_m=20.0,
	timestamps=np.linspace(0, 1.0, n_pts),
	distances=np.linspace(0, 20.0, n_pts),
	raw_telemetry={
	"roll_rad": np.zeros(n_pts),
	"gauge_mm": np.full(n_pts, 1676.0),
	"lateral_pos_mm": np.zeros(n_pts),
	"vertical_pos_mm": np.zeros(n_pts),
	},
	frames=[],
	)

	decision_clean = pipeline.process_window(clean_window)
	assert decision_clean.decision == DecisionType.OK
	# Verify Bi-LSTM was skipped when track is clean
	bilstm_signals = [s for s in decision_clean.signals if s.signal_type == SignalType.GEOMETRY_KNOWN_TYPE]
	assert len(bilstm_signals) == 0

	# 2. Defective track window (5.5mm twist ramp exceedance)
	roll_with_twist = np.zeros(n_pts)
	roll_with_twist[30:] = np.arcsin(5.5 / 1676.0)

	twist_window = ChainageWindow(
	start_chainage_m=100.0,
	end_chainage_m=120.0,
	timestamps=np.linspace(0, 1.0, n_pts),
	distances=np.linspace(100.0, 120.0, n_pts),
	raw_telemetry={
	"roll_rad": roll_with_twist,
	"gauge_mm": np.full(n_pts, 1676.0),
	"lateral_pos_mm": np.zeros(n_pts),
	"vertical_pos_mm": np.zeros(n_pts),
	},
	frames=[],
	)

	decision_twist = pipeline.process_window(twist_window)
	assert decision_twist.decision == DecisionType.INSPECT_KNOWN
	assert decision_twist.defect_family == DefectFamily.GEOMETRY

	# Verify Bi-LSTM was triggered conditionally
	bilstm_signals = [s for s in decision_twist.signals if s.signal_type == SignalType.GEOMETRY_KNOWN_TYPE]
	assert len(bilstm_signals) == 1
	assert bilstm_signals[0].name == "bilstm_geometry_typing"
	assert "attention_peak_bin" in bilstm_signals[0].explanation


	if __name__ == "__main__":
	test_bilstm_attention_focus()
	test_tc_v1_contract_compliance()
	test_bilstm_attention_architecture_shapes()
	test_synthetic_dataset_and_dataloaders()
	test_conditional_execution_pipeline_integration()