DOCSEARCHbyAlgolia
Submit Docs

torchphysics

repository·main·Indexed 19 days ago

https://github.com/qewton-labs/torchphysics

A PyTorch library for solving differential equations using mesh-free deep learning methods. It provides high-level abstractions for Physics-Informed Neural Networks (PINNs), Deep Operator Networks (DeepONets), Fourier Neural Operators (FNOs), and the Deep Ritz Method (DRM). The library includes utilities for defining problem domains, samplers, and differential operators such as gradient, laplacian, and divergence to implement PDE residuals and boundary conditions.

Tokens
33.8K
Snippets
107
Records
128
Agent score
65%

What's inside torchphysics

  1. Overview of TorchPhysics capabilities

    main

    TorchPhysics is a Python library built on PyTorch designed for mesh-free deep learning methods to solve differential equations. It provides high-level implementations for several key mathematical approaches:

    • Physics-informed neural networks (PINN)
    • The Deep Ritz method
    • DeepONets and physics-informed DeepONets
    • Fourier Neural Operators (FNO) and physics-informed FNO
    • Model order reduction networks (PCANN)

    Common use cases include solving ordinary and partial differential equations, training neural networks to approximate solutions for varying parameters, solving inverse problems, interpolating external data, and learning function operators.

  2. Understand the Domain abstraction

    main

    In torchphysics, a Domain represents a geometric region within a specific space. All domain classes inherit from the base Domain class and provide a consistent interface for geometric queries and properties.

    Core Methods and Properties

    • __contains__: Check if a point lies inside the domain.
    • volume / set_volume: Retrieve or explicitly set the volume of the domain. Note that for complex domains created via operations, the volume might not be automatically computable; use set_volume if an exact value is required.
    • bounding_box: Returns the bounding box of the domain.
    • boundary: Returns a new domain representing the boundary of the current domain. The boundary object implements the same methods as a standard domain and includes information about normal vectors, but a boundary itself has no boundary.

    All domains are located in the torchphysics.domains module.

  4. Compute derivatives with respect to multiple variables

    main

    When using operators like laplacian, you can compute the derivative with respect to multiple input variables simultaneously by passing all relevant tensors to the method. For example, to compute the Laplacian across both spatial coordinates $x$ and time $t$ (i.e., $\partial_{x_1}^2f + \partial_{x_2}^2f + \partial_t^2f$), pass both x and t to the operator.

    import torchphysics as tp

# Computes the Laplacian with respect to both x and t
laplace_xt = tp.utils.laplacian(out, x, t)

  5. What are Conditions in TorchPhysics

    main

    In TorchPhysics, Conditions are the central abstraction used to transform the mathematical conditions of a differential equation (like boundary conditions or PDE residuals) into training objectives for a neural network.

    Most conditions require the following five arguments:

    • module: The neural network model being trained.
    • sampler: The sampler providing the points where the condition is applied.
    • residual_fn: A function that computes the residual (the difference between the model's output and the expected mathematical value).
    • name: A string identifier used for logging losses during training.
    • weight: A scalar multiplier applied to the condition's loss during training.

    By default, the loss is computed as the Mean Squared Error (MSE) of the residual, but this can be customized using error_fn and reduce_fn. Since Condition inherits from torch.nn.Module, you can extend it for custom behavior.

  6. Use Sequential and Parallel model evaluation

    main

    For complex architectures, torchphysics.models provides two evaluation patterns:

    • Sequential: Evaluates multiple networks in order (left to right). This is useful for applying layers like normalization before the main model.
    • Parallel: Evaluates different networks in parallel. This is useful when a solution consists of multiple distinct functions (e.g., velocity $v$ and pressure $p$) or when applying locally different networks.

    To use Sequential, pass the layers/models in the desired order to tp.models.Sequential.

    # Example: Applying a normalization layer sequentially before a model
T = tp.domains.Triangle(X, origin=[0, 0], corner_1=[1, 0], corner_2=[2.0, 0])
normal_layer = tp.models.NormalizationLayer(T)
seq_model = tp.models.Sequential(normal_layer, model)

  7. Define training conditions using PINNCondition

    main

    Training conditions in TorchPhysics transform a PDE into residuals that are minimized during training. The tp.conditions.PINNCondition is used for Physics-Informed Neural Networks (PINNs).

    To create a PINNCondition, you need:

    1. A residual function: A Python function residual_fn(u, x) that returns the difference between the current model output and the expected physical/boundary value.
    2. A sampler: An object from tp.samplers that provides the points where the condition is evaluated.
    3. The model: The neural network being trained.
    4. (Optional) A weight: A scalar to scale the loss term.

    Note: Using .make_static() on a sampler ensures the points are sampled once and remain constant during training.

    import torch
import numpy as np
import torchphysics as tp

# 1. Define residual function
def bound_residual(u, x):
    bound_values = torch.sin(np.pi/2*x[:, :1]) * torch.cos(2*np.pi*x[:, 1:])
    return u - bound_values

# 2. Define sampler
bound_sampler = tp.samplers.GridSampler(square.boundary, n_points=5000)
bound_sampler = bound_sampler.make_static()

# 3. Wrap in PINNCondition
bound_cond = tp.conditions.PINNCondition(module=model, sampler=bound_sampler, 
                                         residual_fn=bound_residual, weight=10)

  8. Core features of TorchPhysics

    main

    TorchPhysics provides a modular framework for translating mathematical problems into code. Key features include:

    Domain Generation and Manipulation

    • Mesh-free domain generation: Supports built-in types like Point, Interval, Parallelogram, Circle, Triangle, and Sphere.
    • Complex domains: Create complex geometries using boolean operators (Union, Cut, Intersection) and higher-dimensional objects via Cartesian products.
    • External objects: Load external objects via Trimesh and Shapely (soft dependencies).
    • Interdependence: Support for interdependent or moving domains.

    Sampling and Operators

    • Point sampling: Various methods including RandomUniform, Grid, Gaussian, Latin hypercube, and Adaptive sampling.
    • Differential operators: Easy definition of differential equations using built-in operators.

    Neural Network Training

    • Model Architectures: Pre-implemented fully connected networks with easy extensibility.
    • Training Utilities: Supports sequential or parallel evaluation/training, normalization layers, and adaptive weights to accelerate training.
    • PyTorch Lightning Integration: Leverages PyTorch Lightning for powerful training, including optimizer control, learning rate management, and monitoring individual condition losses.

  9. Define spaces and variables using the Space class

    main

    In TorchPhysics, a Space defines the names and dimensionalities of variables used in domains and models. You construct a Space by taking the Cartesian product of dimension subclasses: R1, R2, or R3.

    Key properties of a Space object:

    • .dim: The total dimensionality of the space.
    • .variables: An unordered set of the variable names contained in the space.
    • Subspace checking: Use the in operator to check if a variable name (string) or a subspace (another Space object) is contained within the space.

    Example: Defining a 2D spatial space x and a 1D time space t combined into a single space G.

    import torchphysics as tp

X = tp.spaces.R2('x')
T = tp.spaces.R1('t')
G = X*T

print(G.dim)        # Output: 3
print(G.variables)  # Output: {'x', 't'}
print('x' in G)    # Output: True
print(X in G)      # Output: True

  10. Combine samplers using Sampler Operations

    main

    You can combine multiple samplers to create complex sampling patterns using specific operators:

    • Concatenation (+): Uses ConcatSampler to join the outputs of two samplers. Both samplers must exist in the same space. The total number of points is the sum of the points from both samplers. This is useful for combining different sampling strategies (e.g., interior and boundary) for the same condition.
    • Column Stacking (.append): Uses AppendSampler to create a column stack of two sampler outputs. Both samplers must produce the same number of points but can be in different spaces. This is useful when points are uncorrelated.
    • Cartesian Product (*): Uses ProductSampler to create a mesh grid of points. The total number of points is the product of the points from both samplers. This is useful for creating points on higher-dimensional surfaces (e.g., a cylinder surface from a circle boundary and an interval).
    import torchphysics as tp

# 1. Concatenation: Combine interior and boundary points
sum_sampler = random_R + random_R_bound

# 2. Cartesian Product: Create grid points on a cylinder surface
T = tp.spaces.R1('t')
I = tp.domains.Interval(T, 0, 1) # height of cylinder
C_sampler = tp.samplers.GridSampler(C.boundary, n_points=50)
I_sampler = tp.samplers.GridSampler(I, n_points=10)

# Resulting sampler has 50 * 10 = 500 points
cylinder_sampler = C_sampler * I_sampler

  11. Use PlotSampler and AnimationSampler for visualization

    main

    While any implemented sampler can be used, torchphysics provides specialized samplers for visualization tasks:

    • PlotSampler: Used for creating static plots.
    • AnimationSampler: Used for creating animations.

    How they work: These samplers take a single domain as input and generate a point grid that includes both the interior and the boundaries. Because plotting is performed over a specific domain, all other input variables for the model must be provided as constant values via a dictionary passed to the sampler.

  12. Use Callbacks to extend training logic

    main

    Callbacks allow you to monitor training or inject custom logic (e.g., plotting solutions to TensorBoard or saving the network periodically) without modifying the solver code.

    • You can use standard callbacks from pytorch_lightning.callbacks.
    • TorchPhysics provides specialized callbacks in the torchphysics.utils module.
    • To use them, pass a list of callback instances to the callbacks keyword in the pl.Trainer constructor.