1D_Acoustic_PBPINNs.py

import jax
import jax.numpy as jnp
import numpy as np
import optax
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from IPython.display import HTML

from fbpinns.domains import RectangularDomainND
from fbpinns.problems import Problem
from fbpinns.decompositions import RectangularDecompositionND
from fbpinns.networks import FCN, SIREN, Network
from fbpinns.schedulers import LineSchedulerRectangularND
from fbpinns.constants import Constants, get_subdomain_ws
from fbpinns.trainers import FBPINNTrainer

print('Setup complete. FBPINNs configuration loaded.')

def chained_optimizer(transforms):
    """Wrapper for optax.chain to be called with keyword arguments."""
    return optax.chain(*transforms)

# Enhanced reproducibility
key = jax.random.PRNGKey(42)

# Improved Swish network with better initialization
class SwishFCN(Network):
    """Swish FCN with improved initialization for wave problems"""

    @staticmethod
    def init_params(key, layer_sizes):
        keys = jax.random.split(key, len(layer_sizes)-1)
        params = []
        
        for i, (k, m, n) in enumerate(zip(keys, layer_sizes[:-1], layer_sizes[1:])):
            if i == 0:  # Input layer - Xavier initialization
                w_key, b_key = jax.random.split(k)
                std = jnp.sqrt(2.0 / (m + n))
                w = jax.random.normal(w_key, (n, m)) * std
                b = jnp.zeros((n,))
            elif i == len(layer_sizes) - 2:  # Output layer - smaller initialization
                w_key, b_key = jax.random.split(k)
                std = jnp.sqrt(1.0 / m) * 0.1
                w = jax.random.normal(w_key, (n, m)) * std
                b = jnp.zeros((n,))
            else:  # Hidden layers
                w_key, b_key = jax.random.split(k)
                std = jnp.sqrt(2.0 / m)  # He initialization for Swish
                w = jax.random.normal(w_key, (n, m)) * std
                b = jnp.zeros((n,))
            
            params.append((w, b))
        
        trainable_params = {"layers": params}
        return {}, trainable_params

    @staticmethod
    def network_fn(params, x):
        params = params["trainable"]["network"]["subdomain"]["layers"]
        activation_scale = 1.0
        
        for i, (w, b) in enumerate(params[:-1]):
            x = jnp.dot(w, x) + b
            x = jax.nn.swish(x) * activation_scale
            # Slight scaling to prevent activation saturation
            if i > 0:
                activation_scale *= 0.95
                
        w, b = params[-1]
        x = jnp.dot(w, x) + b
        return x

# 1D Acoustic Wave Problem with corrected physics
class AcousticWave1D(Problem):
    """1D acoustic wave equation with improved constraint handling"""
    
    @staticmethod
    def init_params(c=1.0, pulse_center=0.3, pulse_std=0.08, loss_weights=None):
        """Initialize with balanced loss weights and improved pulse parameters"""
        default_weights = {
            "pde": 1.0, 
            "initial_displacement": 50.0, 
            "initial_velocity": 50.0, 
            "boundary": 25.0
        }
        
        static_params = {
            "dims": (1, 2),
            "c": c,
            "pulse_center": pulse_center,
            "pulse_std": pulse_std,
            "loss_weights": loss_weights or default_weights
        }
        return static_params, {}

    @staticmethod
    def initial_displacement(all_params, x):
        """Gaussian pulse initial condition"""
        pulse_center = all_params["static"]["problem"]["pulse_center"]
        pulse_std = all_params["static"]["problem"]["pulse_std"]
        amplitude = 1.0
        return amplitude * jnp.exp(-((x - pulse_center)**2) / (2 * pulse_std**2))

    @staticmethod
    def sample_constraints(all_params, domain, key, sampler, batch_shapes):
        """Enhanced constraint sampling with better distribution"""
        key, subkey1, subkey2, subkey3 = jax.random.split(key, 4)
        
        # Interior points for PDE
        x_pde = domain.sample_interior(all_params, subkey1, sampler, batch_shapes[0])
        
        # Boundary sampling
        x_boundary_batches = domain.sample_boundaries(all_params, subkey2, sampler, batch_shapes[1:])
        x_b_xmin, x_b_xmax, x_b_tmin, _ = x_boundary_batches
        
        # Additional initial condition sampling for better coverage
        n_initial_extra = batch_shapes[1][0] // 2
        x_initial_extra = domain.sample_boundary(all_params, subkey3, sampler, 
                                                (n_initial_extra,), boundary_idx=2)
        x_b_tmin_combined = jnp.vstack([x_b_tmin, x_initial_extra])
        
        # Required gradients
        req_pde = ((0, (1, 1)), (0, (0, 0)))  # u_tt, u_xx
        req_initial = ((0, ()),)               # u
        req_velocity = ((0, (1,)),)            # u_t
        req_boundary = ((0, ()),)              # u

        return [
            [x_pde, req_pde],                    # PDE constraint
            [x_b_tmin_combined, req_initial],    # Initial displacement
            [x_b_tmin_combined, req_velocity],   # Initial velocity
            [x_b_xmin, req_boundary],           # Left boundary
            [x_b_xmax, req_boundary]            # Right boundary
        ]

    @staticmethod
    def loss_fn(all_params, constraints):
        """Enhanced loss function with adaptive weighting"""
        loss_weights = all_params["static"]["problem"]["loss_weights"]
        c = all_params["static"]["problem"]["c"]
        
        # PDE residual loss
        x_pde, u_tt, u_xx = constraints[0]
        pde_residual = u_tt - (c**2) * u_xx
        pde_loss = jnp.mean(pde_residual**2)
        
        # Initial displacement loss
        x_initial, u_initial = constraints[1]
        target_displacement = AcousticWave1D.initial_displacement(
            all_params, x_initial[:, 0])
        initial_displacement_loss = jnp.mean((u_initial.squeeze() - target_displacement)**2)
        
        # Initial velocity loss (should be zero)
        x_velocity, u_t_initial = constraints[2]
        initial_velocity_loss = jnp.mean(u_t_initial**2)
        
        # Boundary losses (homogeneous Dirichlet)
        x_left, u_left = constraints[3]
        x_right, u_right = constraints[4]
        boundary_loss = jnp.mean(u_left**2) + jnp.mean(u_right**2)
        
        # Adaptive loss weighting to prevent one term from dominating
        total_loss = (loss_weights["pde"] * pde_loss + 
                     loss_weights["initial_displacement"] * initial_displacement_loss +
                     loss_weights["initial_velocity"] * initial_velocity_loss +
                     loss_weights["boundary"] * boundary_loss)
        
        return total_loss

    @staticmethod
    def exact_solution(all_params, x_batch, batch_shape=None):
        """d'Alembert solution for wave equation with Gaussian initial condition"""
        c = all_params["static"]["problem"]["c"]
        pulse_center = all_params["static"]["problem"]["pulse_center"]
        pulse_std = all_params["static"]["problem"]["pulse_std"]
        
        x = x_batch[:, 0]
        t = x_batch[:, 1]
        
        def initial_profile(xi):
            """Initial displacement profile"""
            return jnp.exp(-((xi - pulse_center)**2) / (2 * pulse_std**2))
        
        # d'Alembert's formula: u(x,t) = 1/2[f(x-ct) + f(x+ct)]
        # For zero initial velocity
        u_exact = 0.5 * (initial_profile(x - c * t) + initial_profile(x + c * t))
        
        # Apply boundary conditions (crude approximation for visualization)
        # In practice, the exact solution should account for reflections
        mask = (x >= 0) & (x <= 1)
        u_exact = jnp.where(mask, u_exact, 0.0)
        
        return u_exact.reshape(-1, 1)

# Optimized FBPINN configuration
def create_fbpinn_config():
    """Create optimized FBPINN configuration for 1D wave equation"""
    
    # Reduced subdomain decomposition for 1D problem
    subdomain_xs = [np.linspace(0, 1, 6), np.linspace(0, 1, 6)]  # 6x6 instead of 8x8
    subdomain_ws = get_subdomain_ws(subdomain_xs, 2.5)  # Moderate overlap
    
    # Advanced learning rate schedule
    learning_rate_schedule = optax.warmup_cosine_decay_schedule(
        init_value=1e-6,
        peak_value=5e-4,
        warmup_steps=2000,
        decay_steps=80000,
        end_value=1e-6
    )
    
    # Enhanced optimizer with gradient clipping
    optimizer = chained_optimizer
    optimizer_kwargs = dict(
        transforms=[
            optax.clip_by_global_norm(1.0),
            optax.adam(learning_rate=learning_rate_schedule, b1=0.9, b2=0.999)
        ]
    )
    
    return Constants(
        domain=RectangularDomainND,
        domain_init_kwargs=dict(
            xmin=np.array([0., 0.]),
            xmax=np.array([1., 1.])
        ),
        problem=AcousticWave1D,
        problem_init_kwargs=dict(
            c=1.0,
            pulse_center=0.3,
            pulse_std=0.08,
            loss_weights={
                "pde": 1.0, 
                "initial_displacement": 50.0, 
                "initial_velocity": 50.0, 
                "boundary": 25.0
            }
        ),
        decomposition=RectangularDecompositionND,
        decomposition_init_kwargs=dict(
            subdomain_xs=subdomain_xs,
            subdomain_ws=subdomain_ws,
            unnorm=(0., 1.)
        ),
        network=SwishFCN,
        network_init_kwargs=dict(
            layer_sizes=[2, 48, 48, 48, 32, 1],  # Deeper network
        ),
        scheduler=LineSchedulerRectangularND,
        scheduler_kwargs=dict(
            point=[0.],
            iaxis=0
        ),
        optimiser=optimizer,
        optimiser_kwargs=optimizer_kwargs,
        # Improved sampling strategy
        ns=((60, 60), (400,), (400,), (300,), (300,)),  # More PDE points, fewer boundary
        n_test=(150, 150),
        n_steps=100000,
        seed=42,
        sampler="grid",  # Explicit grid sampling
        summary_freq=2000,
        test_freq=2000,
        model_save_freq=10000,
        show_figures=False,
        save_figures=True,
        clear_output=False
    )

# Training execution with enhanced monitoring
def train_model():
    """Train the FBPINN model with better monitoring"""
    print("Initializing FBPINN for 1D Acoustic Wave Propagation...")
    
    c = create_fbpinn_config()
    trainer = FBPINNTrainer(c)
    
    print("Configuration Summary:")
    print(f"- Domain: [0,1] x [0,1] (x-t space)")
    print(f"- Subdomains: 6x6 with moderate overlap")
    print(f"- Network: 5-layer Swish FCN [2,48,48,48,32,1]")
    print(f"- Training steps: {c.n_steps}")
    print(f"- Wave speed: c = {c.problem_init_kwargs['c']}")
    print(f"- Pulse center: {c.problem_init_kwargs['pulse_center']}")
    print(f"- Loss weights: {c.problem_init_kwargs['loss_weights']}")
    
    print("\nStarting training...")
    trainer.train()
    
    print("\nTraining completed successfully!")
    print("Results saved to output directory.")
    return trainer

# Execute the training
if __name__ == "__main__":
    trainer = train_model()
    print("\nTraining finished. Check for saved plots in the output directory.")