pulsarfitpy Technical Information

Implementing 2D Physics Informed Neural Networks (PINNs) with pulsarfitpy

pulsarfitpy offers a comprehensive framework for solving 2D partial differential equations using Physics Informed Neural Networks (PINNs). The 2D implementation extends builds on the 1D framework to handle complex physical systems. The 2D PINN framework is particularly useful for studying phenomena that require spatial variation across two dimensions, enabling scientists to test theoretical models of pulsar physics against observational data.

Understanding the PulsarPINN2D Framework

The 2D PINN framework in pulsarfitpy is built on a modular architecture consisting of three main components: the primary PulsarPINN2D solver, the PyTorchBackend for automatic differentiation, and the DeepXDEBackend for high-level PDE solving. This separation of concerns allows flexibility in choosing computational strategies while maintaining a unified user interface.

PulsarPINN2D Class Overview

The PulsarPINN2D class serves as the main interface for solving 2D partial differential equations using physics-informed neural networks. It provides a unified interface that abstracts away backend-specific details while offering comprehensive visualization and analysis capabilities.

The class requires the following parameters:

Backend Architecture

The PulsarPINN2D framework supports two distinct computational backends, each offering different advantages for solving 2D PDEs.

PyTorch Backend

The PyTorch backend (pytorch_solver_2D.py) provides a low-level, flexible implementation using automatic differentiation. This backend is ideal for custom PDE formulations and when fine-grained control over the computation is needed.

Key Features:

The PyTorchBackend class provides the following core functionality:

.build_network(input_dim, hidden_layers, output_dim): Constructs the neural network architecture with Tanh activation functions for smooth derivatives needed in PDE computations.

.set_training_data(collocation_points, boundary_points, boundary_values): Loads training data consisting of interior collocation points for PDE residuals and boundary points with known values.

.compile_pde_loss(pde_expr, variables): Converts a SymPy PDE expression into a PyTorch computational graph using automatic differentiation.

.train(epochs, learning_rate, callback_interval): Executes the training loop, optimizing network weights and computing loss metrics at regular intervals.

.predict(x, y): Evaluates the trained network at specified 2D coordinates.

DeepXDE Backend

The DeepXDE backend (deepxde_solver_2D.py) provides a high-level API built on the DeepXDE library, which is specifically designed for physics-informed neural networks. This backend is excellent for rapid prototyping and standard PDE problems.

Key Features:

The DeepXDEBackend class provides the following core functionality:

.build_network(input_dim, hidden_layers, output_dim): Constructs the neural network using DeepXDE’s FNN class with Tanh activation and Glorot normal initialization.

.set_geometry(x_range, y_range): Defines the rectangular computational domain with specified x and y bounds.

.add_boundary_condition(bc_func, on_boundary): Adds Dirichlet boundary conditions to the domain, with automatic boundary detection.

.compile_pde_loss(pde_expr, variables): Creates a PDE residual function compatible with DeepXDE’s automatic differentiation system.

.create_training_data(num_domain, num_boundary, train_distribution): Generates training data by sampling collocation points and boundary points from the domain.

.compile_model(): Assembles the network, data, and PDE into a complete model ready for training.

.train(epochs, learning_rate, callback_interval): Executes training with the specified optimizer and convergence criteria.

.predict(x, y): Evaluates the trained network at specified 2D coordinates.

PulsarPINN2D Methods

The primary methods of the PulsarPINN2D class for 2D PDE solving and analysis are as follows:

1. .set_training_data(collocation_points, boundary_points, boundary_values):

Sets the training data for the PINN model. This method prepares interior points where the PDE residual will be enforced and boundary points with known values.

Inputs:

Outputs:

2. .train(epochs, learning_rate=1e-3, callback_interval=500):

Trains the physics-informed neural network by minimizing combined physics and boundary condition losses.

Inputs:

Outputs:

3. .predict(x, y):

Evaluates the trained network at specified 2D coordinates to obtain solution predictions.

Inputs:

Outputs:

4. .set_visualization_config(colormap=None, dpi=None, figsize=None, style=None):

Configures visualization settings for all plotting functions.

Inputs:

Outputs:

5. .plot_loss_history(log_scale=True, separate_losses=True, savefig=None):

Visualizes training loss evolution over epochs, showing total loss, PDE residual loss, and boundary condition loss.

Inputs:

Outputs:

6. .plot_loss_heatmap(window_size=100, savefig=None):

Creates a heatmap visualization showing how different loss components evolve during training with rolling average smoothing.

Inputs:

Outputs:

7. .plot_convergence_rate(savefig=None):

Analyzes and visualizes the convergence rate of training by fitting exponential decay and computing loss reduction rates.

Inputs:

Outputs:

8. .plot_solution_2d(resolution=100, x_range=(0, 1), y_range=(0, 1), colormap=None, show_colorbar=True, contour_levels=15, savefig=None):

Visualizes the 2D PINN solution as a filled contour plot with optional contour lines.

Inputs:

Outputs:

9. .plot_solution_3d(resolution=50, x_range=(0, 1), y_range=(0, 1), colormap=None, elevation=30, azimuth=45, savefig=None):

Visualizes the 2D PINN solution as a 3D surface plot with adjustable viewpoint.

Inputs:

Outputs:

10. .plot_residual_distribution(resolution=100, x_range=(0, 1), y_range=(0, 1), savefig=None):

Visualizes the spatial distribution of PDE residuals across the domain to assess where the model satisfies the physics constraints.

Inputs:

Outputs:

11. .plot_comparison_with_analytical(analytical_solution, resolution=100, x_range=(0, 1), y_range=(0, 1), savefig=None):

Compares PINN solution with an analytical reference solution, computing error metrics and visualizing differences.

Inputs:

Outputs:

12. .create_comprehensive_report(resolution=100, x_range=(0, 1), y_range=(0, 1), savefig=None):

Generates a comprehensive analysis report combining multiple visualizations including loss history, solution plots, residuals, and convergence metrics.

Inputs:

Outputs:

13. .plot_training_animation(resolution=50, x_range=(0, 1), y_range=(0, 1), interval=200, save_path=None):

Creates an animated visualization of how the PINN solution evolves during training.

Inputs:

Outputs:

14. .get_metrics_summary():

Retrieves a summary dictionary of key metrics from the most recent training session.

Inputs:

Outputs:

15. .save_model(filepath):

Saves the trained PINN model to disk for later inference or resumption of training.

Inputs:

Outputs:

16. .load_model(filepath):

Loads a previously trained PINN model from disk.

Inputs:

Outputs:

PulsarPINN2D Key Attributes

Here we describe the main attributes of the PulsarPINN2D class:

Model Configuration

Backend and Computation

Training History

Visualization Configuration

PyTorch Backend Key Attributes

The PyTorchBackend class maintains the following important attributes:

Network Architecture

Training Data

Loss Functions

DeepXDE Backend Key Attributes

The DeepXDEBackend class maintains the following important attributes:

Network and Geometry

Training Configuration

PDE Definition

Training Metrics Dataclass

Both backends use a TrainingMetrics dataclass to record training progress. Each metric includes:

Exporting Trained Models

The pulsarfitpy framework provides comprehensive export functionality through the PINNSolutionExporter class. This allows you to save training results and persist trained models for future use.

Key export methods:

1. .save_predictions_to_csv(filepath):

Exports model predictions on training, validation, and test sets to a CSV file with input-output pairs and model evaluations.

Inputs:

2. .save_learned_constants_to_csv(filepath):

Saves the learned physical constants from the trained model to a CSV file with constant names and their optimized values.

Inputs:

3. .save_metrics_to_csv(filepath):

Exports comprehensive evaluation metrics (R^2, RMSE, MAE, chi^2) for all data splits (train/val/test) to a CSV file.

Inputs:

4. .save_loss_history_to_csv(filepath):

Saves the complete training loss history including total loss, physics loss, and data loss across all epochs to a CSV file.

Inputs:

5. .save_model_checkpoint(filepath, include_metadata=True):

Saves the trained model state to a PyTorch checkpoint file (.pt) for later inference or continued training. This preserves the neural network weights, learned constants, and optionally includes training metadata.

Inputs:

Outputs:

Usage Notes

The following practical guidance applies when using the 2D PINN framework:

Physics Loss Computation: The physics loss measures how well the network solution satisfies the governing partial differential equation through residual computation at interior collocation points.

Boundary Condition Enforcement: Boundary condition loss ensures the network output matches known values at domain boundaries, which is essential for well-posed PDE problems.

Grid Resolution: Higher resolution in collocation points (set via set_training_data) and visualization (set via plot_ methods) improves accuracy but increases computational cost. Balance is needed based on problem complexity.

Visualization Interpretation: The contour plots show solution topology directly, while 3D surface plots provide intuitive understanding of solution magnitude variation. Loss heatmaps reveal which training phase had performance issues.

Backend Selection: Choose PyTorchBackend for maximum flexibility with custom PDEs, or DeepXDEBackend for rapid prototyping of standard PDE types with built-in optimizations.

Training Convergence: Monitor both physics loss and boundary loss separately using plot_loss_history. If physics loss stagnates while boundary loss decreases, increase physics_weight in training parameters.

Solution Validation: When analytical solutions are available, use plot_comparison_with_analytical to quantify solution accuracy and identify remaining errors in the physics approximation.

Residual Analysis: Use plot_residual_distribution to locate regions where the learned solution violates physics constraints most significantly. High residuals indicate need for domain refinement or model architecture adjustment.

Memory Management: For large grids or high-resolution visualizations, use resolution parameter carefully. Memory usage scales with resolution squared for 2D problems.

Checkpoint Management: Use save_model and load_model for long training sessions to preserve progress and enable resumption of training without restarting from scratch.

← Back to Technical Information Home