Upload 5 files

data_loader.py

@@ -0,0 +1,239 @@
+import os
+import re
+import numpy as np
+import pandas as pd
+from glob import glob
+from sklearn.preprocessing import StandardScaler
+
+
+class DamageCalculator:
+
+    @staticmethod
+    def compute_freeze_thaw_damage(FN, FT, a1=0.002, b1=1.0, c1=0.02):
+        return a1 * (FN ** b1) * np.exp(c1 * FT)
+
+    @staticmethod
+    def compute_chemical_damage(pH, a2=0.01, b2=1.5):
+        return a2 * np.abs(pH - 7.0) ** b2
+
+    @staticmethod
+    def compute_thermal_damage(T, T0=100.0, a3=0.0003, b3=1.2):
+        if T < T0:
+            return 0.0
+        return a3 * ((T - T0) ** b3)
+
+    @staticmethod
+    def compute_total_damage(pH, FN, FT, T):
+        D_ft = DamageCalculator.compute_freeze_thaw_damage(FN, FT)
+        D_ch = DamageCalculator.compute_chemical_damage(pH)
+        D_th = DamageCalculator.compute_thermal_damage(T)
+
+        D_total = 1.0 - (1.0 - D_ft) * (1.0 - D_ch) * (1.0 - D_th)
+        return np.clip(D_total, 0.0, 0.99)
+
+    @staticmethod
+    def compute_lambda(D0):
+        return 1.0 - D0
+
+
+class CrackDataLoader:
+
+    def __init__(self, base_path, stress_type="major"):
+        self.base_path = base_path
+        self.stress_type = stress_type
+
+        if stress_type == "major":
+            self.data_dir = os.path.join(base_path, "major_principal_stress")
+        else:
+            self.data_dir = os.path.join(base_path, "minor_principal_stress")
+
+        self.scaler_X = StandardScaler()
+        self.scaler_y = StandardScaler()
+        self.damage_calculator = DamageCalculator()
+
+    def parse_filename(self, filename):
+        pattern = r'(\d+)-(\d+)-(\d+)-(\d+)'
+        match = re.search(pattern, filename)
+
+        if match:
+            pH = int(match.group(1))
+            FN = int(match.group(2))
+            FT = int(match.group(3))
+            T = int(match.group(4))
+
+            return {
+                'pH': pH,
+                'FN': FN,
+                'FT': FT,
+                'T': T
+            }
+        else:
+            raise ValueError(f"Cannot parse filename: {filename}")
+
+    def load_single_csv(self, csv_path):
+        data = pd.read_csv(csv_path, header=None, names=['angle', 'count'])
+        angles = data['angle'].values
+        counts = data['count'].values
+        return angles, counts
+
+    def load_all_data(self, phase="both"):
+        X_list = []
+        y_list = []
+        damage_list = []
+
+        if phase == "both":
+            subdirs = ["unstable_development", "peak_stress"]
+        elif phase == "early":
+            subdirs = ["unstable_development"]
+        elif phase == "peak":
+            subdirs = ["peak_stress"]
+        else:
+            raise ValueError(f"Unknown phase: {phase}")
+
+        for subdir in subdirs:
+            subdir_path = os.path.join(self.data_dir, subdir)
+
+            if not os.path.exists(subdir_path):
+                print(f"Warning: Directory does not exist {subdir_path}")
+                continue
+
+            phase_code = 0 if "unstable" in subdir else 1
+
+            csv_files = glob(os.path.join(subdir_path, "*.csv"))
+
+            print(f"Loading {len(csv_files)} files from {subdir}...")
+
+            for csv_file in csv_files:
+                try:
+                    params = self.parse_filename(os.path.basename(csv_file))
+
+                    angles, counts = self.load_single_csv(csv_file)
+
+                    D0 = DamageCalculator.compute_total_damage(
+                        params['pH'], params['FN'], params['FT'], params['T']
+                    )
+                    lambda_coef = DamageCalculator.compute_lambda(D0)
+
+                    features = np.array([
+                        params['pH'],
+                        params['FN'],
+                        params['FT'],
+                        params['T'],
+                        phase_code
+                    ], dtype=np.float32)
+
+                    X_list.append(features)
+                    y_list.append(counts)
+                    damage_list.append({'D0': D0, 'lambda': lambda_coef})
+
+                except Exception as e:
+                    print(f"Skipping file {csv_file}: {e}")
+                    continue
+
+        if len(X_list) == 0:
+            raise ValueError("No data loaded successfully!")
+
+        X = np.array(X_list)
+
+        y_length = len(y_list[0])
+        y_padded = []
+
+        for y_sample in y_list:
+            if len(y_sample) < y_length:
+                y_sample = np.pad(y_sample, (0, y_length - len(y_sample)), 'constant')
+            elif len(y_sample) > y_length:
+                y_sample = y_sample[:y_length]
+            y_padded.append(y_sample)
+
+        y = np.array(y_padded)
+
+        angles, _ = self.load_single_csv(csv_files[0])
+        angle_bins = angles[:y_length]
+
+        print(f"\nData loading complete:")
+        print(f"  Samples: {X.shape[0]}")
+        print(f"  Input features: {X.shape[1]} (pH, FN, FT, T, phase)")
+        print(f"  Output dimension: {y.shape[1]} (angle bins)")
+        print(f"  Angle range: {angle_bins[0]:.1f} - {angle_bins[-1]:.1f}")
+        print(f"  Total cracks range: {y.sum(axis=1).min():.0f} - {y.sum(axis=1).max():.0f}")
+
+        return X, y, angle_bins, damage_list
+
+    def create_synthetic_data(self, n_samples=100, output_dim=72):
+        pH_values = [1, 3, 5, 7]
+        FN_values = [5, 10, 20, 40]
+        FT_values = [10, 20, 30, 40]
+        T_values = [25, 300, 600, 900]
+        phase_values = [0, 1]
+
+        X_list = []
+        y_list = []
+
+        for _ in range(n_samples):
+            pH = np.random.choice(pH_values)
+            FN = np.random.choice(FN_values)
+            FT = np.random.choice(FT_values)
+            T = np.random.choice(T_values)
+            phase = np.random.choice(phase_values)
+
+            D0 = DamageCalculator.compute_total_damage(pH, FN, FT, T)
+
+            if self.stress_type == "major":
+                peak_angle = 90.0 + np.random.normal(0, 10)
+                spread = 15.0 + D0 * 20.0
+            else:
+                peak_angle = 45.0 + np.random.normal(0, 15)
+                spread = 20.0 + D0 * 25.0
+
+            angles = np.linspace(0, 175, output_dim)
+            distribution = np.exp(-0.5 * ((angles - peak_angle) / spread) ** 2)
+            distribution = distribution * (100 + D0 * 200) * (1 + 0.5 * phase)
+            distribution = distribution + np.random.normal(0, 5, output_dim)
+            distribution = np.maximum(distribution, 0)
+
+            X_list.append([pH, FN, FT, T, phase])
+            y_list.append(distribution)
+
+        X = np.array(X_list, dtype=np.float32)
+        y = np.array(y_list, dtype=np.float32)
+        angle_bins = np.linspace(0, 175, output_dim)
+
+        return X, y, angle_bins
+
+    def normalize_data(self, X_train, y_train, X_test=None, y_test=None):
+        X_train_norm = self.scaler_X.fit_transform(X_train)
+        y_train_norm = self.scaler_y.fit_transform(y_train)
+
+        if X_test is not None and y_test is not None:
+            X_test_norm = self.scaler_X.transform(X_test)
+            y_test_norm = self.scaler_y.transform(y_test)
+            return X_train_norm, y_train_norm, X_test_norm, y_test_norm
+        else:
+            return X_train_norm, y_train_norm
+
+    def denormalize_output(self, y_norm):
+        return self.scaler_y.inverse_transform(y_norm)
+
+    def get_statistics(self, X, y):
+        stats = {
+            'n_samples': X.shape[0],
+            'input_dim': X.shape[1],
+            'output_dim': y.shape[1],
+            'pH_range': (X[:, 0].min(), X[:, 0].max()),
+            'FN_range': (X[:, 1].min(), X[:, 1].max()),
+            'FT_range': (X[:, 2].min(), X[:, 2].max()),
+            'T_range': (X[:, 3].min(), X[:, 3].max()),
+            'total_cracks_range': (y.sum(axis=1).min(), y.sum(axis=1).max()),
+            'total_cracks_mean': y.sum(axis=1).mean(),
+            'total_cracks_std': y.sum(axis=1).std(),
+        }
+
+        D0_values = []
+        for i in range(X.shape[0]):
+            D0 = DamageCalculator.compute_total_damage(X[i, 0], X[i, 1], X[i, 2], X[i, 3])
+            D0_values.append(D0)
+
+        stats['D0_range'] = (min(D0_values), max(D0_values))
+        stats['D0_mean'] = np.mean(D0_values)
+
+        return stats

inference.py

@@ -0,0 +1,238 @@
+import os
+import numpy as np
+import torch
+import pickle
+import matplotlib.pyplot as plt
+
+from model import CrackTransformerPINN
+from data_loader import DamageCalculator
+
+
+class CrackPredictor:
+
+    def __init__(self, model_path, scaler_path, device='cpu'):
+        self.device = device
+
+        with open(scaler_path, 'rb') as f:
+            scalers = pickle.load(f)
+
+        self.scaler_X = scalers['scaler_X']
+        self.scaler_y = scalers['scaler_y']
+        self.angle_bins = scalers['angle_bins']
+
+        checkpoint = torch.load(model_path, map_location=device)
+
+        if 'model_config' in checkpoint:
+            config = checkpoint['model_config']
+            self.model = CrackTransformerPINN(
+                input_dim=config['input_dim'],
+                output_dim=config['output_dim'],
+                hidden_dims=config['hidden_dims'],
+                dropout=config['dropout']
+            )
+            self.model.load_state_dict(checkpoint['model_state_dict'])
+        else:
+            self.model = CrackTransformerPINN(
+                input_dim=5,
+                output_dim=len(self.angle_bins),
+                hidden_dims=[128, 256, 256, 128],
+                dropout=0.2
+            )
+            self.model.load_state_dict(checkpoint)
+
+        self.model.to(device)
+        self.model.eval()
+
+    def predict(self, pH, FN, FT, T, phase):
+        X = np.array([[pH, FN, FT, T, phase]], dtype=np.float32)
+
+        X_norm = self.scaler_X.transform(X)
+
+        with torch.no_grad():
+            X_tensor = torch.FloatTensor(X_norm).to(self.device)
+            pred_dist_norm, pred_total = self.model(X_tensor, return_physics=False)
+
+            pred_dist_norm = pred_dist_norm.cpu().numpy()
+            pred_total = pred_total.cpu().numpy().flatten()
+
+        pred_dist = self.scaler_y.inverse_transform(pred_dist_norm)
+
+        D0 = DamageCalculator.compute_total_damage(pH, FN, FT, T)
+        lambda_coef = DamageCalculator.compute_lambda(D0)
+
+        return {
+            'angle_distribution': pred_dist[0],
+            'total_count': pred_total[0],
+            'D0': D0,
+            'lambda': lambda_coef
+        }
+
+    def predict_with_physics(self, pH, FN, FT, T, phase):
+        X = np.array([[pH, FN, FT, T, phase]], dtype=np.float32)
+
+        X_norm = self.scaler_X.transform(X)
+
+        with torch.no_grad():
+            X_tensor = torch.FloatTensor(X_norm).to(self.device)
+            pred_dist_norm, pred_total, physics = self.model(X_tensor, return_physics=True)
+
+            pred_dist_norm = pred_dist_norm.cpu().numpy()
+
+        pred_dist = self.scaler_y.inverse_transform(pred_dist_norm)
+
+        result = {
+            'angle_distribution': pred_dist[0],
+            'total_count': pred_total.cpu().numpy().flatten()[0],
+            'D0': physics['D0'].cpu().numpy().flatten()[0],
+            'lambda': physics['lambda'].cpu().numpy().flatten()[0],
+            'D_n': physics['D_n'].cpu().numpy().flatten()[0],
+            'tau_oct': physics['tau_oct'].cpu().numpy().flatten()[0],
+            'yield_stress': physics['yield_stress'].cpu().numpy().flatten()[0],
+            'C1': physics['C1'].cpu().numpy().flatten()[0],
+            'C2': physics['C2'].cpu().numpy().flatten()[0],
+            'D_q': physics['D_q'].cpu().numpy().flatten()[0],
+            'm': physics['m'].cpu().numpy().flatten()[0],
+            'F0': physics['F0'].cpu().numpy().flatten()[0]
+        }
+
+        return result
+
+    def predict_batch(self, X):
+        X_norm = self.scaler_X.transform(X)
+
+        with torch.no_grad():
+            X_tensor = torch.FloatTensor(X_norm).to(self.device)
+            pred_dist_norm, pred_total = self.model(X_tensor, return_physics=False)
+
+            pred_dist_norm = pred_dist_norm.cpu().numpy()
+            pred_total = pred_total.cpu().numpy().flatten()
+
+        pred_dist = self.scaler_y.inverse_transform(pred_dist_norm)
+
+        return pred_dist, pred_total
+
+    def visualize(self, result, title=None, save_path=None):
+        fig, ax = plt.subplots(figsize=(8, 8), subplot_kw=dict(projection='polar'))
+
+        theta = np.deg2rad(self.angle_bins)
+        pred_dist = result['angle_distribution']
+
+        ax.plot(theta, pred_dist, 'o-', linewidth=2, markersize=4, color='blue')
+        ax.fill(theta, pred_dist, alpha=0.3, color='blue')
+
+        if title:
+            ax.set_title(title, fontsize=14, pad=20)
+        else:
+            total = result['total_count']
+            D0 = result['D0']
+            lam = result['lambda']
+            ax.set_title(f'Predicted Crack Distribution\nTotal: {total:.0f}, D0: {D0:.3f}, lambda: {lam:.3f}',
+                        fontsize=12, pad=20)
+
+        ax.set_theta_zero_location('N')
+        ax.set_theta_direction(-1)
+        ax.grid(True, alpha=0.3)
+
+        if save_path:
+            plt.savefig(save_path, dpi=300, bbox_inches='tight')
+
+        plt.show()
+
+    def compare_stress_types(self, pH, FN, FT, T, save_path=None):
+        result_unstable = self.predict(pH, FN, FT, T, phase=0)
+        result_peak = self.predict(pH, FN, FT, T, phase=1)
+
+        fig, axes = plt.subplots(1, 2, figsize=(14, 6), subplot_kw=dict(projection='polar'))
+
+        theta = np.deg2rad(self.angle_bins)
+
+        axes[0].plot(theta, result_unstable['angle_distribution'], 'o-',
+                    linewidth=2, markersize=4, color='blue')
+        axes[0].fill(theta, result_unstable['angle_distribution'], alpha=0.3, color='blue')
+        axes[0].set_title(f'Unstable Development Phase\nTotal: {result_unstable["total_count"]:.0f}',
+                         fontsize=12, pad=20)
+        axes[0].set_theta_zero_location('N')
+        axes[0].set_theta_direction(-1)
+        axes[0].grid(True, alpha=0.3)
+
+        axes[1].plot(theta, result_peak['angle_distribution'], 'o-',
+                    linewidth=2, markersize=4, color='red')
+        axes[1].fill(theta, result_peak['angle_distribution'], alpha=0.3, color='red')
+        axes[1].set_title(f'Peak Stress Phase\nTotal: {result_peak["total_count"]:.0f}',
+                         fontsize=12, pad=20)
+        axes[1].set_theta_zero_location('N')
+        axes[1].set_theta_direction(-1)
+        axes[1].grid(True, alpha=0.3)
+
+        D0 = result_unstable['D0']
+        lam = result_unstable['lambda']
+        fig.suptitle(f'pH={pH}, FN={FN}, FT={FT}, T={T}\nD0={D0:.3f}, lambda={lam:.3f}',
+                    fontsize=14, fontweight='bold')
+
+        plt.tight_layout()
+
+        if save_path:
+            plt.savefig(save_path, dpi=300, bbox_inches='tight')
+
+        plt.show()
+
+
+def main():
+    model_path = "./output/crack_transformer_pinn.pth"
+    scaler_path = "./output/scalers.pkl"
+
+    if not os.path.exists(model_path):
+        print("Model file not found. Please train the model first using train.py")
+        return
+
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    predictor = CrackPredictor(model_path, scaler_path, device=device)
+
+    print("=" * 60)
+    print("Transformer-PINN Crack Prediction - Inference Demo")
+    print("=" * 60)
+
+    test_cases = [
+        {'pH': 3, 'FN': 30, 'FT': 40, 'T': 25, 'phase': 0},
+        {'pH': 3, 'FN': 40, 'FT': 20, 'T': 300, 'phase': 1},
+        {'pH': 7, 'FN': 10, 'FT': 40, 'T': 300, 'phase': 1},
+        {'pH': 5, 'FN': 10, 'FT': 20, 'T': 900, 'phase': 0},
+    ]
+
+    for i, params in enumerate(test_cases):
+        result = predictor.predict(**params)
+
+        print(f"\nTest Case {i+1}:")
+        print(f"  Input: pH={params['pH']}, FN={params['FN']}, FT={params['FT']}, T={params['T']}, phase={params['phase']}")
+        print(f"  D0 (Initial Damage): {result['D0']:.4f}")
+        print(f"  Lambda (Damage Coefficient): {result['lambda']:.4f}")
+        print(f"  Predicted Total Cracks: {result['total_count']:.0f}")
+        print(f"  Peak Angle: {predictor.angle_bins[result['angle_distribution'].argmax()]:.1f} degrees")
+        print(f"  Peak Count: {result['angle_distribution'].max():.0f}")
+
+    print("\n" + "=" * 60)
+    print("Testing Physics Output")
+    print("=" * 60)
+
+    physics_result = predictor.predict_with_physics(pH=3, FN=30, FT=40, T=25, phase=1)
+
+    print("\nPhysics Parameters:")
+    print(f"  Mogi-Coulomb:")
+    print(f"    tau_oct: {physics_result['tau_oct']:.4f}")
+    print(f"    yield_stress: {physics_result['yield_stress']:.4f}")
+    print(f"    C1: {physics_result['C1']:.4f}")
+    print(f"    C2: {physics_result['C2']:.4f}")
+    print(f"  Weibull Distribution:")
+    print(f"    D_q: {physics_result['D_q']:.4f}")
+    print(f"    m: {physics_result['m']:.4f}")
+    print(f"    F0: {physics_result['F0']:.4f}")
+    print(f"  Energy Damage:")
+    print(f"    D_n: {physics_result['D_n']:.4f}")
+
+    print("\n" + "=" * 60)
+    print("Inference complete!")
+    print("=" * 60)
+
+
+if __name__ == "__main__":
+    main()

model.py

@@ -0,0 +1,519 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import TensorDataset, DataLoader
+from sklearn.model_selection import train_test_split
+
+
+class MogiCoulombLayer(nn.Module):
+
+    def __init__(self, hidden_dim):
+        super(MogiCoulombLayer, self).__init__()
+        self.C1_net = nn.Linear(hidden_dim, 1)
+        self.C2_net = nn.Linear(hidden_dim, 1)
+
+    def forward(self, features, sigma1, sigma2, sigma3):
+        C1 = torch.relu(self.C1_net(features))
+        C2 = torch.sigmoid(self.C2_net(features))
+
+        tau_oct = (1.0/3.0) * torch.sqrt(
+            (sigma1 - sigma2)**2 + (sigma2 - sigma3)**2 + (sigma1 - sigma3)**2
+        )
+        sigma_m2 = (sigma1 + sigma3) / 2.0
+
+        yield_stress = C1 + C2 * sigma_m2
+
+        return tau_oct, yield_stress, C1, C2
+
+
+class WeibullStrengthLayer(nn.Module):
+
+    def __init__(self, hidden_dim):
+        super(WeibullStrengthLayer, self).__init__()
+        self.m_net = nn.Sequential(
+            nn.Linear(hidden_dim, 32),
+            nn.Tanh(),
+            nn.Linear(32, 1),
+            nn.Softplus()
+        )
+        self.F0_net = nn.Sequential(
+            nn.Linear(hidden_dim, 32),
+            nn.Tanh(),
+            nn.Linear(32, 1),
+            nn.Softplus()
+        )
+
+    def forward(self, features, F):
+        m = self.m_net(features) + 1.0
+        F0 = self.F0_net(features) + 0.1
+
+        D_q = 1.0 - torch.exp(-torch.pow(F / F0, m))
+
+        return D_q, m, F0
+
+
+class EnergyDamageLayer(nn.Module):
+
+    def __init__(self, hidden_dim):
+        super(EnergyDamageLayer, self).__init__()
+        self.a_net = nn.Sequential(
+            nn.Linear(hidden_dim, 32),
+            nn.Tanh(),
+            nn.Linear(32, 1),
+            nn.Softplus()
+        )
+        self.b_net = nn.Sequential(
+            nn.Linear(hidden_dim, 32),
+            nn.Tanh(),
+            nn.Linear(32, 1),
+            nn.Softplus()
+        )
+
+    def forward(self, features, delta_sigma, D0):
+        a = self.a_net(features) + 0.1
+        b = self.b_net(features) + 0.01
+
+        effective_stress = delta_sigma / (1.0 - D0 + 1e-8)
+
+        U_p = a * torch.exp(b * effective_stress)
+
+        D_n = (2.0 / np.pi) * torch.atan(b * U_p)
+
+        return D_n, U_p, a, b
+
+
+class CrackTransformerPINN(nn.Module):
+
+    def __init__(self, input_dim=5, output_dim=72, hidden_dims=[128, 256, 256, 128], dropout=0.2):
+        super(CrackTransformerPINN, self).__init__()
+
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+
+        self.input_embedding = nn.Sequential(
+            nn.Linear(input_dim, hidden_dims[0]),
+            nn.LayerNorm(hidden_dims[0]),
+            nn.GELU(),
+            nn.Dropout(dropout * 0.5)
+        )
+
+        self.damage_encoder = nn.Sequential(
+            nn.Linear(input_dim, hidden_dims[0]),
+            nn.Tanh(),
+            nn.Linear(hidden_dims[0], hidden_dims[0])
+        )
+
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=hidden_dims[0],
+            nhead=8,
+            dim_feedforward=hidden_dims[0] * 4,
+            dropout=dropout,
+            activation='gelu',
+            batch_first=True,
+            norm_first=True
+        )
+
+        self.transformer_encoder = nn.TransformerEncoder(
+            encoder_layer,
+            num_layers=4,
+            norm=nn.LayerNorm(hidden_dims[0])
+        )
+
+        self.mogi_coulomb = MogiCoulombLayer(hidden_dims[0])
+        self.weibull_strength = WeibullStrengthLayer(hidden_dims[0])
+        self.energy_damage = EnergyDamageLayer(hidden_dims[0])
+
+        self.angle_decoder = nn.ModuleList()
+        prev_dim = hidden_dims[0] * 2
+
+        for hidden_dim in hidden_dims[1:]:
+            self.angle_decoder.append(nn.Linear(prev_dim, hidden_dim))
+            self.angle_decoder.append(nn.LayerNorm(hidden_dim))
+            self.angle_decoder.append(nn.Tanh())
+            self.angle_decoder.append(nn.Dropout(dropout))
+            prev_dim = hidden_dim
+
+        self.angle_output = nn.Sequential(
+            nn.Linear(prev_dim, output_dim),
+            nn.ReLU()
+        )
+
+        self.total_count_head = nn.Sequential(
+            nn.Linear(hidden_dims[0] * 2, 64),
+            nn.Tanh(),
+            nn.Linear(64, 1),
+            nn.ReLU()
+        )
+
+        self.damage_factor_head = nn.Sequential(
+            nn.Linear(hidden_dims[0], 32),
+            nn.Tanh(),
+            nn.Linear(32, 1),
+            nn.Sigmoid()
+        )
+
+    def compute_initial_damage(self, pH, FN, FT, T):
+        D_ft = 0.002 * FN * torch.exp(0.02 * FT)
+        D_ch = 0.01 * torch.abs(pH - 7.0) ** 1.5
+        D_th = torch.where(
+            T > 100.0,
+            0.0003 * (T - 100.0) ** 1.2,
+            torch.zeros_like(T)
+        )
+
+        D_total = 1.0 - (1.0 - D_ft) * (1.0 - D_ch) * (1.0 - D_th)
+        D_total = torch.clamp(D_total, 0.0, 0.99)
+
+        return D_total
+
+    def forward(self, x, return_physics=False):
+        batch_size = x.shape[0]
+
+        pH = x[:, 0:1]
+        FN = x[:, 1:2]
+        FT = x[:, 2:3]
+        T = x[:, 3:4]
+        phase = x[:, 4:5]
+
+        D0 = self.compute_initial_damage(pH, FN, FT, T)
+        lambda_coef = 1.0 - D0
+
+        x_embedded = self.input_embedding(x)
+        damage_features = self.damage_encoder(x)
+
+        x_seq = x_embedded.unsqueeze(1)
+        encoded = self.transformer_encoder(x_seq)
+        encoded = encoded.squeeze(1)
+
+        combined = torch.cat([encoded, damage_features], dim=-1)
+
+        h = combined
+        for layer in self.angle_decoder:
+            h = layer(h)
+
+        angle_dist = self.angle_output(h)
+
+        total_count = self.total_count_head(combined)
+
+        predicted_D0 = self.damage_factor_head(encoded)
+
+        if return_physics:
+            sigma1 = 100.0 * torch.ones(batch_size, 1, device=x.device)
+            sigma2 = 50.0 * torch.ones(batch_size, 1, device=x.device)
+            sigma3 = 30.0 * torch.ones(batch_size, 1, device=x.device)
+            delta_sigma = 20.0 * torch.ones(batch_size, 1, device=x.device)
+            F_contact = 10.0 * torch.ones(batch_size, 1, device=x.device)
+
+            tau_oct, yield_stress, C1, C2 = self.mogi_coulomb(encoded, sigma1, sigma2, sigma3)
+            D_q, m, F0 = self.weibull_strength(encoded, F_contact)
+            D_n, U_p, a, b = self.energy_damage(encoded, delta_sigma, D0)
+
+            physics_outputs = {
+                'D0': D0,
+                'lambda': lambda_coef,
+                'predicted_D0': predicted_D0,
+                'tau_oct': tau_oct,
+                'yield_stress': yield_stress,
+                'C1': C1,
+                'C2': C2,
+                'D_q': D_q,
+                'm': m,
+                'F0': F0,
+                'D_n': D_n,
+                'U_p': U_p,
+                'a': a,
+                'b': b
+            }
+
+            return angle_dist, total_count, physics_outputs
+
+        return angle_dist, total_count
+
+
+class CrackPINNLoss(nn.Module):
+
+    def __init__(self, lambda_data=1.0, lambda_physics=0.5, lambda_smooth=0.1,
+                 lambda_damage=0.3, lambda_mogi=0.2, lambda_reg=1e-4):
+        super(CrackPINNLoss, self).__init__()
+
+        self.lambda_data = lambda_data
+        self.lambda_physics = lambda_physics
+        self.lambda_smooth = lambda_smooth
+        self.lambda_damage = lambda_damage
+        self.lambda_mogi = lambda_mogi
+        self.lambda_reg = lambda_reg
+
+        self.mse_loss = nn.MSELoss()
+
+    def data_loss(self, pred_dist, true_dist):
+        return self.mse_loss(pred_dist, true_dist)
+
+    def physics_loss(self, pred_dist, pred_total, true_dist):
+        pred_sum = pred_dist.sum(dim=1, keepdim=True)
+        true_sum = true_dist.sum(dim=1, keepdim=True)
+
+        loss_consistency = self.mse_loss(pred_sum, pred_total)
+        loss_total = self.mse_loss(pred_total, true_sum)
+
+        return loss_consistency + loss_total
+
+    def smoothness_loss(self, pred_dist):
+        diff = pred_dist[:, 1:] - pred_dist[:, :-1]
+        return torch.mean(diff ** 2)
+
+    def damage_consistency_loss(self, physics_outputs):
+        D0 = physics_outputs['D0']
+        predicted_D0 = physics_outputs['predicted_D0']
+
+        loss_D0 = self.mse_loss(predicted_D0, D0)
+
+        lambda_coef = physics_outputs['lambda']
+        loss_lambda = torch.mean(torch.relu(-lambda_coef) + torch.relu(lambda_coef - 1.0))
+
+        return loss_D0 + loss_lambda
+
+    def mogi_coulomb_loss(self, physics_outputs):
+        tau_oct = physics_outputs['tau_oct']
+        yield_stress = physics_outputs['yield_stress']
+
+        loss_yield = torch.mean(torch.relu(tau_oct - yield_stress))
+
+        C1 = physics_outputs['C1']
+        C2 = physics_outputs['C2']
+        loss_params = torch.mean(torch.relu(-C1)) + torch.mean(torch.relu(C2 - 1.0))
+
+        return loss_yield + 0.1 * loss_params
+
+    def regularization_loss(self, model):
+        l2_reg = torch.tensor(0.0, device=next(model.parameters()).device)
+        for param in model.parameters():
+            l2_reg += torch.norm(param, p=2) ** 2
+        return l2_reg
+
+    def forward(self, pred_dist, pred_total, true_dist, model, physics_outputs=None):
+        loss_data = self.data_loss(pred_dist, true_dist)
+        loss_physics = self.physics_loss(pred_dist, pred_total, true_dist)
+        loss_smooth = self.smoothness_loss(pred_dist)
+        loss_reg = self.regularization_loss(model)
+
+        loss_damage = torch.tensor(0.0, device=pred_dist.device)
+        loss_mogi = torch.tensor(0.0, device=pred_dist.device)
+
+        if physics_outputs is not None:
+            loss_damage = self.damage_consistency_loss(physics_outputs)
+            loss_mogi = self.mogi_coulomb_loss(physics_outputs)
+
+        total_loss = (self.lambda_data * loss_data +
+                     self.lambda_physics * loss_physics +
+                     self.lambda_smooth * loss_smooth +
+                     self.lambda_damage * loss_damage +
+                     self.lambda_mogi * loss_mogi +
+                     self.lambda_reg * loss_reg)
+
+        loss_dict = {
+            'total': total_loss.item(),
+            'data': loss_data.item(),
+            'physics': loss_physics.item(),
+            'smooth': loss_smooth.item(),
+            'damage': loss_damage.item(),
+            'mogi': loss_mogi.item(),
+            'reg': loss_reg.item()
+        }
+
+        return total_loss, loss_dict
+
+
+class CrackPINNTrainer:
+
+    def __init__(self, model, device='cpu', lr=1e-3, weight_decay=1e-4):
+        self.model = model.to(device)
+        self.device = device
+
+        self.optimizer = optim.AdamW(
+            model.parameters(),
+            lr=lr,
+            weight_decay=weight_decay,
+            betas=(0.9, 0.999)
+        )
+
+        self.scheduler = optim.lr_scheduler.ReduceLROnPlateau(
+            self.optimizer,
+            mode='min',
+            factor=0.5,
+            patience=10,
+            verbose=True
+        )
+
+        self.criterion = CrackPINNLoss(
+            lambda_data=1.0,
+            lambda_physics=0.5,
+            lambda_smooth=0.1,
+            lambda_damage=0.3,
+            lambda_mogi=0.2,
+            lambda_reg=1e-4
+        )
+
+        self.train_losses = []
+        self.val_losses = []
+
+    def train_epoch(self, train_loader):
+        self.model.train()
+
+        epoch_losses = []
+        loss_components = {
+            'data': [], 'physics': [], 'smooth': [],
+            'damage': [], 'mogi': [], 'reg': []
+        }
+
+        for X_batch, y_batch in train_loader:
+            X_batch = X_batch.to(self.device)
+            y_batch = y_batch.to(self.device)
+
+            pred_dist, pred_total, physics_outputs = self.model(X_batch, return_physics=True)
+
+            loss, loss_dict = self.criterion(
+                pred_dist, pred_total, y_batch, self.model, physics_outputs
+            )
+
+            self.optimizer.zero_grad()
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
+            self.optimizer.step()
+
+            epoch_losses.append(loss_dict['total'])
+            for key in loss_components.keys():
+                loss_components[key].append(loss_dict[key])
+
+        avg_loss = np.mean(epoch_losses)
+        avg_components = {key: np.mean(values) for key, values in loss_components.items()}
+
+        return avg_loss, avg_components
+
+    def validate(self, val_loader):
+        self.model.eval()
+
+        val_losses = []
+        all_preds = []
+        all_trues = []
+
+        with torch.no_grad():
+            for X_batch, y_batch in val_loader:
+                X_batch = X_batch.to(self.device)
+                y_batch = y_batch.to(self.device)
+
+                pred_dist, pred_total = self.model(X_batch, return_physics=False)
+
+                loss, _ = self.criterion(pred_dist, pred_total, y_batch, self.model)
+
+                val_losses.append(loss.item())
+                all_preds.append(pred_dist.cpu().numpy())
+                all_trues.append(y_batch.cpu().numpy())
+
+        avg_loss = np.mean(val_losses)
+
+        all_preds = np.concatenate(all_preds, axis=0)
+        all_trues = np.concatenate(all_trues, axis=0)
+
+        ss_res = np.sum((all_trues - all_preds) ** 2)
+        ss_tot = np.sum((all_trues - np.mean(all_trues)) ** 2)
+        r2 = 1 - (ss_res / (ss_tot + 1e-8))
+
+        rmse = np.sqrt(np.mean((all_trues - all_preds) ** 2))
+
+        pred_total_counts = all_preds.sum(axis=1)
+        true_total_counts = all_trues.sum(axis=1)
+        total_count_error = np.mean(np.abs(pred_total_counts - true_total_counts))
+
+        metrics = {
+            'r2': r2,
+            'rmse': rmse,
+            'total_count_mae': total_count_error
+        }
+
+        return avg_loss, metrics
+
+    def fit(self, X_train, y_train, X_val, y_val, epochs=200, batch_size=16, patience=30):
+        train_dataset = TensorDataset(
+            torch.FloatTensor(X_train),
+            torch.FloatTensor(y_train)
+        )
+        train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
+
+        val_dataset = TensorDataset(
+            torch.FloatTensor(X_val),
+            torch.FloatTensor(y_val)
+        )
+        val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
+
+        best_val_loss = float('inf')
+        patience_counter = 0
+        best_model_state = None
+
+        print("\nStarting training...")
+        print("=" * 80)
+
+        for epoch in range(epochs):
+            train_loss, train_components = self.train_epoch(train_loader)
+            val_loss, val_metrics = self.validate(val_loader)
+
+            self.train_losses.append(train_loss)
+            self.val_losses.append(val_loss)
+
+            self.scheduler.step(val_loss)
+
+            if (epoch + 1) % 10 == 0 or epoch == 0:
+                print(f"Epoch {epoch+1}/{epochs}")
+                print(f"  Train Loss: {train_loss:.4f} "
+                      f"(data: {train_components['data']:.4f}, "
+                      f"phys: {train_components['physics']:.4f}, "
+                      f"damage: {train_components['damage']:.4f})")
+                print(f"  Val Loss: {val_loss:.4f} | "
+                      f"R2: {val_metrics['r2']:.4f} | "
+                      f"RMSE: {val_metrics['rmse']:.2f}")
+
+            if val_loss < best_val_loss:
+                best_val_loss = val_loss
+                patience_counter = 0
+                best_model_state = self.model.state_dict().copy()
+            else:
+                patience_counter += 1
+
+            if patience_counter >= patience:
+                print(f"\nEarly stopping. Best val loss: {best_val_loss:.4f}")
+                break
+
+        if best_model_state is not None:
+            self.model.load_state_dict(best_model_state)
+
+        print("=" * 80)
+        print(f"Training complete. Best val loss: {best_val_loss:.4f}")
+
+    def predict(self, X):
+        self.model.eval()
+
+        with torch.no_grad():
+            X_tensor = torch.FloatTensor(X).to(self.device)
+            pred_dist, pred_total = self.model(X_tensor, return_physics=False)
+
+            pred_dist = pred_dist.cpu().numpy()
+            pred_total = pred_total.cpu().numpy().flatten()
+
+        return pred_dist, pred_total
+
+    def predict_with_physics(self, X):
+        self.model.eval()
+
+        with torch.no_grad():
+            X_tensor = torch.FloatTensor(X).to(self.device)
+            pred_dist, pred_total, physics = self.model(X_tensor, return_physics=True)
+
+            result = {
+                'angle_distribution': pred_dist.cpu().numpy(),
+                'total_count': pred_total.cpu().numpy().flatten(),
+                'D0': physics['D0'].cpu().numpy().flatten(),
+                'lambda': physics['lambda'].cpu().numpy().flatten(),
+                'D_n': physics['D_n'].cpu().numpy().flatten()
+            }
+
+        return result

requirements.txt

@@ -0,0 +1,5 @@
+torch>=2.0.0
+numpy>=1.21.0
+pandas>=1.3.0
+scikit-learn>=1.0.0
+matplotlib>=3.5.0

train.py

@@ -0,0 +1,331 @@
+import os
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+from sklearn.model_selection import train_test_split
+import pickle
+
+from data_loader import CrackDataLoader, DamageCalculator
+from model import CrackTransformerPINN, CrackPINNTrainer
+
+
+def plot_training_history(trainer, save_path=None):
+    fig, axes = plt.subplots(1, 3, figsize=(15, 4))
+
+    axes[0].plot(trainer.train_losses, label='Train Loss', linewidth=2)
+    axes[0].plot(trainer.val_losses, label='Val Loss', linewidth=2)
+    axes[0].set_xlabel('Epoch', fontsize=12)
+    axes[0].set_ylabel('Loss', fontsize=12)
+    axes[0].set_title('Training History', fontsize=14, fontweight='bold')
+    axes[0].legend()
+    axes[0].grid(True, alpha=0.3)
+
+    axes[1].semilogy(trainer.train_losses, label='Train Loss', linewidth=2)
+    axes[1].semilogy(trainer.val_losses, label='Val Loss', linewidth=2)
+    axes[1].set_xlabel('Epoch', fontsize=12)
+    axes[1].set_ylabel('Loss (log scale)', fontsize=12)
+    axes[1].set_title('Training History (Log Scale)', fontsize=14, fontweight='bold')
+    axes[1].legend()
+    axes[1].grid(True, alpha=0.3)
+
+    if len(trainer.train_losses) > 1:
+        train_improvement = np.diff(trainer.train_losses)
+        axes[2].plot(train_improvement, linewidth=2, alpha=0.7)
+        axes[2].axhline(y=0, color='r', linestyle='--', alpha=0.5)
+        axes[2].set_xlabel('Epoch', fontsize=12)
+        axes[2].set_ylabel('Loss Change', fontsize=12)
+        axes[2].set_title('Convergence Rate', fontsize=14, fontweight='bold')
+        axes[2].grid(True, alpha=0.3)
+
+    plt.tight_layout()
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        print(f"Training history saved to: {save_path}")
+
+    plt.close()
+
+
+def plot_prediction_examples(X_test, y_test, trainer, angle_bins, loader, n_examples=4, save_path=None):
+    X_test_original = loader.scaler_X.inverse_transform(X_test)
+    y_test_original = loader.scaler_y.inverse_transform(y_test)
+
+    y_pred_norm, pred_totals = trainer.predict(X_test[:n_examples])
+    y_pred = loader.scaler_y.inverse_transform(y_pred_norm)
+
+    fig, axes = plt.subplots(2, 2, figsize=(14, 10), subplot_kw=dict(projection='polar'))
+    axes = axes.flatten()
+
+    for i in range(min(n_examples, len(axes))):
+        ax = axes[i]
+
+        theta = np.deg2rad(angle_bins)
+
+        ax.plot(theta, y_test_original[i], 'o-', label='True', linewidth=2, markersize=4, alpha=0.7)
+        ax.plot(theta, y_pred[i], 's-', label='Predicted', linewidth=2, markersize=3, alpha=0.7)
+
+        pH = X_test_original[i, 0]
+        FN = X_test_original[i, 1]
+        FT = X_test_original[i, 2]
+        T = X_test_original[i, 3]
+        phase = X_test_original[i, 4]
+        phase_str = "Unstable" if phase < 0.5 else "Peak"
+
+        D0 = DamageCalculator.compute_total_damage(pH, FN, FT, T)
+        lambda_coef = DamageCalculator.compute_lambda(D0)
+
+        true_total = y_test_original[i].sum()
+        pred_total = y_pred[i].sum()
+
+        title = f"pH={pH:.0f}, FN={FN:.0f}, FT={FT:.0f}, T={T:.0f}C\n"
+        title += f"D0={D0:.3f}, lambda={lambda_coef:.3f}\n"
+        title += f"{phase_str} | True: {true_total:.0f}, Pred: {pred_total:.0f}"
+
+        ax.set_title(title, fontsize=9, pad=20)
+        ax.legend(loc='upper right', fontsize=8)
+        ax.set_theta_zero_location('N')
+        ax.set_theta_direction(-1)
+        ax.grid(True, alpha=0.3)
+
+    plt.tight_layout()
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        print(f"Prediction examples saved to: {save_path}")
+
+    plt.close()
+
+
+def plot_damage_analysis(X, y, save_path=None):
+    fig, axes = plt.subplots(2, 2, figsize=(12, 10))
+
+    D0_values = []
+    total_cracks = y.sum(axis=1)
+
+    for i in range(X.shape[0]):
+        D0 = DamageCalculator.compute_total_damage(X[i, 0], X[i, 1], X[i, 2], X[i, 3])
+        D0_values.append(D0)
+
+    D0_values = np.array(D0_values)
+
+    axes[0, 0].scatter(D0_values, total_cracks, alpha=0.6, edgecolors='black', linewidth=0.5)
+    axes[0, 0].set_xlabel('Initial Damage Factor D0', fontsize=12)
+    axes[0, 0].set_ylabel('Total Crack Count', fontsize=12)
+    axes[0, 0].set_title('D0 vs Total Cracks', fontsize=14, fontweight='bold')
+    axes[0, 0].grid(True, alpha=0.3)
+
+    axes[0, 1].scatter(X[:, 0], total_cracks, alpha=0.6, c=D0_values, cmap='viridis')
+    axes[0, 1].set_xlabel('pH Value', fontsize=12)
+    axes[0, 1].set_ylabel('Total Crack Count', fontsize=12)
+    axes[0, 1].set_title('pH vs Total Cracks', fontsize=14, fontweight='bold')
+    axes[0, 1].grid(True, alpha=0.3)
+
+    axes[1, 0].scatter(X[:, 1], total_cracks, alpha=0.6, c=D0_values, cmap='viridis')
+    axes[1, 0].set_xlabel('Freeze-thaw Cycles (FN)', fontsize=12)
+    axes[1, 0].set_ylabel('Total Crack Count', fontsize=12)
+    axes[1, 0].set_title('FN vs Total Cracks', fontsize=14, fontweight='bold')
+    axes[1, 0].grid(True, alpha=0.3)
+
+    scatter = axes[1, 1].scatter(X[:, 3], total_cracks, alpha=0.6, c=D0_values, cmap='viridis')
+    axes[1, 1].set_xlabel('Damage Temperature (T)', fontsize=12)
+    axes[1, 1].set_ylabel('Total Crack Count', fontsize=12)
+    axes[1, 1].set_title('T vs Total Cracks', fontsize=14, fontweight='bold')
+    axes[1, 1].grid(True, alpha=0.3)
+
+    plt.colorbar(scatter, ax=axes[1, 1], label='D0')
+
+    plt.tight_layout()
+
+    if save_path:
+        plt.savefig(save_path, dpi=300, bbox_inches='tight')
+        print(f"Damage analysis saved to: {save_path}")
+
+    plt.close()
+
+
+def main():
+    print("=" * 80)
+    print("Transformer-PINN Crack Prediction Model")
+    print("Based on: Mechanism of micro-damage evolution in rocks")
+    print("under multiple coupled cyclic stresses")
+    print("=" * 80)
+
+    base_path = "./data"
+    output_dir = "./output"
+
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+
+    print("\n" + "=" * 80)
+    print("Step 1: Loading/Generating Data")
+    print("=" * 80)
+
+    loader = CrackDataLoader(base_path, stress_type="major")
+
+    try:
+        X, y, angle_bins, damage_list = loader.load_all_data(phase="both")
+    except:
+        print("Real data not found. Generating synthetic data...")
+        X, y, angle_bins = loader.create_synthetic_data(n_samples=200, output_dim=72)
+
+    stats = loader.get_statistics(X, y)
+    print("\nData statistics:")
+    for key, value in stats.items():
+        print(f"  {key}: {value}")
+
+    print("\n" + "=" * 80)
+    print("Step 2: Splitting Dataset (Train:Val:Test = 64:16:20)")
+    print("=" * 80)
+
+    X_train, X_test, y_train, y_test = train_test_split(
+        X, y, test_size=0.2, random_state=42
+    )
+
+    X_train, X_val, y_train, y_val = train_test_split(
+        X_train, y_train, test_size=0.2, random_state=42
+    )
+
+    print(f"Training set: {X_train.shape[0]} samples")
+    print(f"Validation set: {X_val.shape[0]} samples")
+    print(f"Test set: {X_test.shape[0]} samples")
+
+    print("\n" + "=" * 80)
+    print("Step 3: Normalizing Data")
+    print("=" * 80)
+
+    X_train_norm, y_train_norm, X_val_norm, y_val_norm = loader.normalize_data(
+        X_train, y_train, X_val, y_val
+    )
+
+    X_test_norm = loader.scaler_X.transform(X_test)
+    y_test_norm = loader.scaler_y.transform(y_test)
+
+    print("Normalization complete")
+
+    print("\n" + "=" * 80)
+    print("Step 4: Creating Model")
+    print("=" * 80)
+
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    print(f"Using device: {device}")
+
+    model = CrackTransformerPINN(
+        input_dim=5,
+        output_dim=y.shape[1],
+        hidden_dims=[128, 256, 256, 128],
+        dropout=0.2
+    )
+
+    n_params = sum(p.numel() for p in model.parameters())
+    print(f"Model parameters: {n_params:,}")
+
+    print("\nModel components:")
+    print("  - Transformer Encoder (8 heads, 4 layers)")
+    print("  - Mogi-Coulomb Yield Criterion Layer")
+    print("  - Weibull Strength Distribution Layer")
+    print("  - Energy-based Damage Evolution Layer")
+    print("  - PINN Decoder with Physics Constraints")
+
+    print("\n" + "=" * 80)
+    print("Step 5: Training Model")
+    print("=" * 80)
+
+    trainer = CrackPINNTrainer(
+        model,
+        device=device,
+        lr=1e-3,
+        weight_decay=1e-4
+    )
+
+    trainer.fit(
+        X_train_norm, y_train_norm,
+        X_val_norm, y_val_norm,
+        epochs=300,
+        batch_size=8,
+        patience=50
+    )
+
+    print("\n" + "=" * 80)
+    print("Step 6: Testing Model")
+    print("=" * 80)
+
+    test_loss, test_metrics = trainer.validate(
+        torch.utils.data.DataLoader(
+            torch.utils.data.TensorDataset(
+                torch.FloatTensor(X_test_norm),
+                torch.FloatTensor(y_test_norm)
+            ),
+            batch_size=8,
+            shuffle=False
+        )
+    )
+
+    print(f"Test set performance:")
+    print(f"  Loss: {test_loss:.4f}")
+    print(f"  R2: {test_metrics['r2']:.4f}")
+    print(f"  RMSE: {test_metrics['rmse']:.2f}")
+    print(f"  Total Count MAE: {test_metrics['total_count_mae']:.2f}")
+
+    print("\n" + "=" * 80)
+    print("Step 7: Saving Model")
+    print("=" * 80)
+
+    model_path = os.path.join(output_dir, "crack_transformer_pinn.pth")
+    torch.save({
+        'model_state_dict': model.state_dict(),
+        'model_config': {
+            'input_dim': 5,
+            'output_dim': y.shape[1],
+            'hidden_dims': [128, 256, 256, 128],
+            'dropout': 0.2
+        },
+        'test_metrics': test_metrics
+    }, model_path)
+    print(f"Model saved to: {model_path}")
+
+    scaler_path = os.path.join(output_dir, "scalers.pkl")
+    with open(scaler_path, 'wb') as f:
+        pickle.dump({
+            'scaler_X': loader.scaler_X,
+            'scaler_y': loader.scaler_y,
+            'angle_bins': angle_bins
+        }, f)
+    print(f"Scalers saved to: {scaler_path}")
+
+    print("\n" + "=" * 80)
+    print("Step 8: Generating Visualizations")
+    print("=" * 80)
+
+    history_path = os.path.join(output_dir, "training_history.png")
+    plot_training_history(trainer, save_path=history_path)
+
+    examples_path = os.path.join(output_dir, "prediction_examples.png")
+    plot_prediction_examples(
+        X_test_norm, y_test_norm,
+        trainer, angle_bins, loader,
+        n_examples=4,
+        save_path=examples_path
+    )
+
+    damage_path = os.path.join(output_dir, "damage_analysis.png")
+    plot_damage_analysis(X, y, save_path=damage_path)
+
+    print("\n" + "=" * 80)
+    print("Training Pipeline Complete!")
+    print("=" * 80)
+    print(f"\nGenerated files:")
+    print(f"  1. Model checkpoint: {model_path}")
+    print(f"  2. Scalers: {scaler_path}")
+    print(f"  3. Training history: {history_path}")
+    print(f"  4. Prediction examples: {examples_path}")
+    print(f"  5. Damage analysis: {damage_path}")
+
+    print("\nPhysics constraints applied:")
+    print("  - Mogi-Coulomb yield criterion: tau_oct = C1 + C2 * sigma_m2")
+    print("  - Weibull strength: D_q = 1 - exp(-(F/F0)^m)")
+    print("  - Energy damage: D_n = (2/pi) * arctan(b * U_p)")
+    print("  - Total damage: D_total = 1 - (1-D_ft)(1-D_ch)(1-D_th)")
+
+
+if __name__ == "__main__":
+    main()