VibroFlow AI

# Test imports
from src.data.loader import HydraulicDataLoader, CWRUBearingDataLoader
from src.models.deep_learning import CNN1D
from src.flow.estimator import FlowEstimator

print(" VibroFlow AI installed successfully!")

from src.data.loader import HydraulicDataLoader, CWRUBearingDataLoader

# ═══════════════════════════════════════════════════════
#  Hydraulic System Dataset
# ═══════════════════════════════════════════════════════
hydraulic_loader = HydraulicDataLoader("dataset0")

# Load sensor data
pressure_data = hydraulic_loader.load_sensor("PS1")     # Shape: (2205, 6000)
vibration_data = hydraulic_loader.load_sensor("VS1")    # Shape: (2205, 60)
flow_data = hydraulic_loader.load_sensor("FS1")         # Shape: (2205, 600)

# Load target conditions
targets = hydraulic_loader.load_targets()
# Returns: cooler, valve, pump_leakage, accumulator, stable_flag

# ═══════════════════════════════════════════════════════
#  CWRU Bearing Dataset
# ═══════════════════════════════════════════════════════
cwru_loader = CWRUBearingDataLoader("dataset1")

# Load preprocessed features
features_df = cwru_loader.load_features_csv()

# Load raw signals for deep learning
X, y = cwru_loader.load_cnn_data()  # Shape: (n_samples, 2048)

from src.data.features import TimeFeatureExtractor, FrequencyFeatureExtractor

# ═══════════════════════════════════════════════════════
# ⏱ Time-Domain Features
# ═══════════════════════════════════════════════════════
time_extractor = TimeFeatureExtractor()

# Single signal
features = time_extractor.extract_all(signal)
# Returns: mean, std, rms, max, min, peak_to_peak,
#          skewness, kurtosis, crest_factor, shape_factor,
#          impulse_factor, margin_factor, energy, entropy

# Batch processing
feature_matrix = time_extractor.extract_batch(signals)  # (n_samples, 14 features)

# ═══════════════════════════════════════════════════════
#  Frequency-Domain Features
# ═══════════════════════════════════════════════════════
freq_extractor = FrequencyFeatureExtractor(sampling_freq=48000)

spectral_features = freq_extractor.extract_all(signal)
# Returns: spectral_centroid, spectral_bandwidth, spectral_rolloff,
#          spectral_flatness, spectral_entropy, spectral_peak_freq,
#          band_power_low, band_power_mid, band_power_high, total_power

from src.models.baseline import BaselineClassifier, ModelComparison
from sklearn.model_selection import train_test_split

# Prepare data
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

# ═══════════════════════════════════════════════════════
#  Train Random Forest Classifier
# ═══════════════════════════════════════════════════════
rf_classifier = BaselineClassifier(model_type='rf')
rf_classifier.fit(X_train, y_train)

# Evaluate
metrics = rf_classifier.evaluate(X_test, y_test)
print(f"Accuracy: {metrics['accuracy']:.4f}")
print(f"F1-Score: {metrics['f1_score']:.4f}")

# ═══════════════════════════════════════════════════════
#  Compare Multiple Models
# ═══════════════════════════════════════════════════════
comparison = ModelComparison()
results = comparison.compare_all(X_train, y_train, X_test, y_test)

# Available models: 'svm', 'rf', 'gb', 'xgb', 'knn'

from src.models.deep_learning import CNN1D, LSTM, HybridCNNLSTM, DeepLearningTrainer
import torch

# ═══════════════════════════════════════════════════════
#  CNN 1D Model
# ═══════════════════════════════════════════════════════
cnn_model = CNN1D(
    input_size=2048,          # Signal length
    num_classes=10,           # Number of fault types
    channels=[32, 64, 128]    # Conv layer channels
)

trainer = DeepLearningTrainer(
    model=cnn_model,
    learning_rate=0.001,
    device='cuda' if torch.cuda.is_available() else 'cpu'
)

# Train
history = trainer.train(
    X_train, y_train,
    X_val, y_val,
    epochs=50,
    batch_size=32
)

# ═══════════════════════════════════════════════════════
#  LSTM Model
# ═══════════════════════════════════════════════════════
lstm_model = LSTM(
    input_size=2048,
    num_classes=10,
    hidden_size=128,
    num_layers=2,
    bidirectional=True
)

# ═══════════════════════════════════════════════════════
#  Hybrid CNN-LSTM Model
# ═══════════════════════════════════════════════════════
hybrid_model = HybridCNNLSTM(
    input_size=2048,
    num_classes=10
)

from src.flow.estimator import FlowEstimator

# ═══════════════════════════════════════════════════════
#  Train Flow Estimator
# ═══════════════════════════════════════════════════════
estimator = FlowEstimator(model_type='rf')  # Options: 'linear', 'ridge', 'rf', 'gb', 'mlp'

# Extract features from vibration data
vibration_features = estimator.extract_vibration_features(vibration_data)

# Prepare flow targets
flow_targets = estimator.prepare_flow_targets(flow_sensor_data)

# Train model
estimator.fit(vibration_features, flow_targets)

# Predict flow from new vibration data
predicted_flow = estimator.predict(new_vibration_features)

# Evaluate
metrics = estimator.evaluate(X_test, y_test)
print(f"R² Score: {metrics['r2']:.4f}")
print(f"MAE: {metrics['mae']:.4f}")

# Save model
estimator.save("models/flow_estimator.joblib")