feat: integrate XJTU-SY dataset and validate cross-domain reliability

Author: TasfinMahmud

cnsd/datasets/xjtusy.py

@@ -0,0 +1,132 @@
+import glob
+import os
+
+import numpy as np
+import pandas as pd
+
+from cnsd.datasets.contract import Dataset
+from cnsd.physics.configs import XJTUSY_PHYSICS
+
+# Known fault mappings per XJTU-SY paper
+# We map Bearing ID to CNSD fault class: 1=Outer, 2=Inner, 3=Cage
+XJTUSY_FAULTS = {
+    '35Hz12kN': {
+        'Bearing1_1': 1,  # Outer
+        'Bearing1_2': 1,  # Outer
+        'Bearing1_3': 1,  # Outer
+    },
+    '37.5Hz11kN': {
+        'Bearing2_1': 2,  # Inner
+        'Bearing2_2': 1,  # Outer
+        'Bearing2_4': 1,  # Outer
+        'Bearing2_5': 1,  # Outer
+    },
+    '40Hz10kN': {
+        'Bearing3_1': 1,  # Outer
+        'Bearing3_2': 2,  # Inner
+        'Bearing3_3': 2,  # Inner
+        'Bearing3_4': 2,  # Inner
+        'Bearing3_5': 1,  # Outer
+    },
+}
+
+
+def load_xjtusy_domain_split(
+    data_dir=r'E:\301\CNSD\data\XJTU-SY\XJTU-SY_Bearing_Datasets',
+    window_size=32768,
+    train_cond='35Hz12kN',
+    test_cond='37.5Hz11kN',
+):
+    """
+    Loads authentic XJTU-SY dataset and strictly splits by condition (Domain Shift).
+    Uses the run-to-failure nature to grab the first 15% as Healthy (0) and the
+    last 15% as the Fault label.
+    """
+
+    def _load_condition(cond_folder, rpm_val):
+        X, y, cond = [], [], []
+        cond_path = os.path.join(data_dir, cond_folder)
+
+        if not os.path.exists(cond_path):
+            return [], [], []
+
+        for bearing_folder in os.listdir(cond_path):
+            bearing_path = os.path.join(cond_path, bearing_folder)
+            if not os.path.isdir(bearing_path):
+                continue
+
+            fault_label = XJTUSY_FAULTS.get(cond_folder, {}).get(bearing_folder, None)
+            if (
+                fault_label is None or bearing_folder == 'Bearing1_5'
+            ):  # skip 1_5 to avoid mixed labels
+                continue
+
+            csv_files = sorted(
+                glob.glob(os.path.join(bearing_path, '*.csv')),
+                key=lambda x: int(os.path.splitext(os.path.basename(x))[0]),
+            )
+
+            total_files = len(csv_files)
+            if total_files < 10:
+                continue  # ignore wildly corrupted directories
+
+            healthy_count = max(1, int(total_files * 0.20))
+            fault_count = max(1, int(total_files * 0.20))
+
+            # Sliding window parameters
+            step_size = 1024
+
+            # Extract Healthy
+            for fpath in csv_files[:healthy_count]:
+                df = pd.read_csv(fpath)
+                sig = df.iloc[:, 0].values  # Horizontal acceleration
+                sig = (sig - np.mean(sig)) / (np.std(sig) + 1e-8)
+
+                # Slicing the 32768 array into overlapping 4096 windows
+                for start_idx in range(0, len(sig) - window_size + 1, step_size):
+                    X.append(sig[start_idx : start_idx + window_size])
+                    y.append(0)
+                    cond.append(rpm_val)
+
+            # Extract Fault
+            for fpath in csv_files[-fault_count:]:
+                df = pd.read_csv(fpath)
+                sig = df.iloc[:, 0].values
+                sig = (sig - np.mean(sig)) / (np.std(sig) + 1e-8)
+
+                for start_idx in range(0, len(sig) - window_size + 1, step_size):
+                    X.append(sig[start_idx : start_idx + window_size])
+                    y.append(fault_label)
+                    cond.append(rpm_val)
+
+        return X, y, cond
+
+    X_train, y_train, c_train = _load_condition(train_cond, 2100.0)
+    X_test, y_test, c_test = _load_condition(test_cond, 2250.0)
+
+    if not X_train or not X_test:
+        raise FileNotFoundError(
+            f'Could not load data. Ensure {data_dir} contains extracted {train_cond} and {test_cond} folders.'
+        )
+
+    ds_train = Dataset.from_arrays(
+        X=np.array(X_train, dtype=np.float32),
+        y=np.array(y_train, dtype=np.int32),
+        cond=np.array(c_train, dtype=np.float32),
+        fs=25600,
+        physics=XJTUSY_PHYSICS,
+        taxonomy={0: ('Normal', 'None'), 1: ('Outer Race', 'Medium'), 2: ('Inner Race', 'High')},
+        name='XJTUSY_Train',
+    )
+
+    ds_test = Dataset.from_arrays(
+        X=np.array(X_test, dtype=np.float32),
+        y=np.array(y_test, dtype=np.int32),
+        cond=np.array(c_test, dtype=np.float32),
+        fs=25600,
+        physics=XJTUSY_PHYSICS,
+        taxonomy={0: ('Normal', 'None'), 1: ('Outer Race', 'Medium'), 2: ('Inner Race', 'High')},
+        name='XJTUSY_Test',
+    )
+
+    return ds_train, ds_test

cnsd/physics/configs.py

@@ -31,3 +31,10 @@ class PhysicsConfig:
     fs=20000,
     name='SEU-Gearbox',
 )
+
+XJTUSY_PHYSICS = PhysicsConfig(
+    bearing={'n_balls': 8, 'd_ball': 7.94, 'd_pitch': 34.55, 'contact_angle': 0.0},
+    cond_to_rpm={2100.0: 2100.0, 2250.0: 2250.0, 2400.0: 2400.0},
+    fs=25600,
+    name='XJTU-SY-LDK-UER204',
+)

evaluate_baselines.py

@@ -1,3 +1,4 @@
+import argparse
 import os
 import sys
 import traceback
@@ -7,39 +8,81 @@
 
 try:
     os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
-
     import tensorflow as tf
 
     from cnsd import Dataset
     from cnsd.diagnosis.system import CNSD
     from cnsd.perception.cnn import _train_cnn
-    from cnsd.physics import PhysicsConfig
-    from validate_pu import load_pu_domain_split
 
-    print('Loading Authentic PU dataset (Cross-Domain RPM Split)...')
-    (X_train_full, y_train_full, cond_train_full), (X_target, y_target, cond_target) = (
-        load_pu_domain_split()
-    )
+    # Argparse
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--dataset', type=str, required=True, choices=['cwru', 'pu', 'xjtusy'])
+    args = parser.parse_args()
 
-    unique_rpm = set(cond_train_full).union(set(cond_target))
-    rpm_map = {float(r): float(r) for r in unique_rpm}
-    pu_physics = PhysicsConfig(
-        bearing={'n_balls': 8, 'd_ball': 6.75, 'd_pitch': 28.5, 'contact_angle': 0.0},
-        cond_to_rpm=rpm_map,
-        fs=64000,
-        name='PU-6203',
-    )
-    pu_taxonomy = {
-        0: ('Normal', 'None'),
-        1: ('Outer Race', 'Medium'),
-        2: ('Inner Race', 'High'),
-    }
+    print(f'Loading {args.dataset.upper()} dataset...')
+
+    if args.dataset == 'pu':
+        from cnsd.physics import PhysicsConfig
+        from validate_pu import load_pu_domain_split
+
+        (X_train_full, y_train_full, cond_train_full), (X_target, y_target, cond_target) = (
+            load_pu_domain_split()
+        )
+        unique_rpm = set(cond_train_full).union(set(cond_target))
+        rpm_map = {float(r): float(r) for r in unique_rpm}
+        physics = PhysicsConfig(
+            bearing={'n_balls': 8, 'd_ball': 6.75, 'd_pitch': 28.5, 'contact_angle': 0.0},
+            cond_to_rpm=rpm_map,
+            fs=64000,
+            name='PU-6203',
+        )
+        taxonomy = {0: ('Normal', 'None'), 1: ('Outer Race', 'Medium'), 2: ('Inner Race', 'High')}
+        fs = 64000
+
+    elif args.dataset == 'cwru':
+        from validate_run import CWRU, TAXONOMY, load_cwru
+
+        X, y, cond = load_cwru()
+        X = np.asarray(X, np.float32)
+        y = np.asarray(y)
+        cond = np.asarray(cond)
+
+        train_mask = cond < 3
+        target_mask = cond == 3
+
+        X_train_full = X[train_mask]
+        y_train_full = y[train_mask]
+        cond_train_full = cond[train_mask]
+
+        X_target = X[target_mask]
+        y_target = y[target_mask]
+        cond_target = cond[target_mask]
+
+        physics = CWRU
+        taxonomy = TAXONOMY
+        fs = 12000
+
+    elif args.dataset == 'xjtusy':
+        from cnsd.datasets.xjtusy import load_xjtusy_domain_split
+
+        train_ds, target_ds = load_xjtusy_domain_split(window_size=4096)
+
+        X_train_full = train_ds.X
+        y_train_full = train_ds.y
+        cond_train_full = train_ds.cond
+
+        X_target = target_ds.X
+        y_target = target_ds.y
+        cond_target = target_ds.cond
+
+        physics = target_ds.physics
+        taxonomy = target_ds.taxonomy
+        fs = target_ds.fs
 
     def get_matched_coverage_gap(score, correct, target_n):
         if target_n == 0:
             return float('nan')
         # Score is higher for MORE confident
-        # Sort descending by score
         sorted_indices = np.argsort(score)[::-1]
         hi_indices = sorted_indices[:target_n]
 
@@ -76,10 +119,10 @@ def clone_for_mc(layer):
         X_target,
         y_target,
         cond_target,
-        fs=64000,
-        physics=pu_physics,
-        taxonomy=pu_taxonomy,
-        name='PU_Test',
+        fs=fs,
+        physics=physics,
+        taxonomy=taxonomy,
+        name=f'{args.dataset.upper()}_Test',
     )
     sig_te = np.stack([test_ds.X[i].reshape(-1) for i in range(len(test_ds.X))]).astype(np.float32)
     yte = test_ds.y
@@ -112,10 +155,22 @@ def clone_for_mc(layer):
         )
 
         train_ds = Dataset.from_arrays(
-            X_tr, y_tr, cond_tr, fs=64000, physics=pu_physics, taxonomy=pu_taxonomy, name='PU_Train'
+            X_tr,
+            y_tr,
+            cond_tr,
+            fs=fs,
+            physics=physics,
+            taxonomy=taxonomy,
+            name=f'{args.dataset.upper()}_Train',
         )
         calib_ds = Dataset.from_arrays(
-            X_ca, y_ca, cond_ca, fs=64000, physics=pu_physics, taxonomy=pu_taxonomy, name='PU_Calib'
+            X_ca,
+            y_ca,
+            cond_ca,
+            fs=fs,
+            physics=physics,
+            taxonomy=taxonomy,
+            name=f'{args.dataset.upper()}_Calib',
         )
 
         # 2. Train Primary Model (Bypass SCM to prevent multiprocess deadlocks)
@@ -209,13 +264,10 @@ def clone_for_mc(layer):
         results['mc_gap'].append(mc_gap)
 
         # Ensemble at matched coverage
-        ens_preds_probs = np.stack(
-            [m.predict(Xin_te, batch_size=128, verbose=0) for m in ens]
-        )  # (3, n, c)
-        ens_mean = ens_preds_probs.mean(0)  # (n, c)
+        ens_preds_probs = np.stack([m.predict(Xin_te, batch_size=128, verbose=0) for m in ens])
+        ens_mean = ens_preds_probs.mean(0)
         ens_pred_class = ens_mean.argmax(1)
         ens_correct = ens_pred_class == yte
-        # Score = negative entropy of ensemble mean
         ens_score = (ens_mean * np.log(ens_mean + eps)).sum(1)
         ens_gap = get_matched_coverage_gap(ens_score, ens_correct, target_n)
         results['ens_gap'].append(ens_gap)
@@ -235,7 +287,6 @@ def clone_for_mc(layer):
                 sig_n = sig_te + rng.randn(*sig_te.shape).astype(np.float32) * np.sqrt(npow)
 
             Xin_n = sig_n[..., None]
-            # Use ensemble mode vote to define 'unanimous' exactly like Abhi's template
             ep_class = np.stack(
                 [m.predict(Xin_n, batch_size=128, verbose=0).argmax(1) for m in ens]
             )