AI/ML for Scientific Discovery Using PyTorch and JAX

Aim

Workshop Objectives

Workshop Structure

Day 1: Scientific ML Foundations & PyTorch-Based Property Prediction

• AI/ML for scientific discovery, simulation acceleration, and inverse design
• Scientific datasets: molecular, materials, physics, experimental, and simulation data
• Scientific data challenges: small datasets, noise, sparsity, uncertainty, and extrapolation
• Tensor representation of scientific systems using PyTorch
• Automatic differentiation for scientific modeling
• Designing neural networks for scientific property prediction
• Model validation, residual analysis, and scientific interpretability

Hands-on

PyTorch-Based Scientific Property Prediction Model
Participants will build a neural network for predicting a physics, chemistry, or materials-related property using structured scientific data.

Day 2: Physics-Informed Neural Networks & Surrogate Modeling

• Physics-informed neural networks
• Scientific loss functions: data loss, physics loss, constraint loss
• Incorporating differential equations, boundary conditions, and conservation laws
• Surrogate models for expensive simulations
• Parameter estimation and inverse modeling
• Failure modes: non-physical predictions, overfitting, and unstable training
• Applications in diffusion, heat transfer, reaction kinetics, and material response modeling

Hands-on

Physics-Informed Neural Network Prototype
Participants will implement a model that combines data-driven learning with physics-based constraints.

Day 3: JAX for Differentiable Scientific Computing & Optimization

• JAX for high-performance scientific ML
• Key transformations: grad, jit, and vmap
• Differentiable scientific computing and parameter optimization
• Vectorized and accelerated scientific model training
• Optax-based optimization workflows
• Comparing PyTorch and JAX for research applications
• Advanced trends: graph neural networks, neural operators, equivariant models, foundation models, and AI-driven materials discovery

Hands-on

JAX-Based Differentiable Optimization Workflow
Participants will build a JAX model for fitting a physical response curve, optimizing a scientific parameter, or modeling an energy/function landscape.

Who Should Enrol?

Important Dates

Workshop Outcomes

Meet Your Mentor(s)

Abedeera Gunarathna Patabandige Gayathri Maduwanthi

more