Skip to content

cd2d

Loss Function Implementation for 2D Convection-Diffusion Problem.

This module implements the loss function for solving 2D convection-diffusion equations using neural networks. It focuses on computing residuals in the weak form of the PDE with known coefficients.

Key functions
  • pde_loss_cd2d: Computes domain-based PDE loss
Note

The implementation is based on the FastVPINNs methodology [1] for efficient computation of Variational residuals of PDEs.

References

[1] FastVPINNs: Tensor-Driven Acceleration of VPINNs for Complex Geometries DOI: https://arxiv.org/abs/2404.12063

pde_loss_cd2d(test_shape_val_mat, test_grad_x_mat, test_grad_y_mat, pred_nn, pred_grad_x_nn, pred_grad_y_nn, forcing_function, bilinear_params)

Calculates residual for 2D convection-diffusion problem.

Implements the FastVPINNs methodology for computing variational residuals in 2D convection-diffusion equations with known coefficients using efficient tensor operations.

Parameters:

Name Type Description Default
test_shape_val_mat Tensor

Test function values at quadrature points Shape: (n_elements, n_test_functions, n_quad_points)

 required
test_grad_x_mat Tensor

Test function x-derivatives at quadrature points Shape: (n_elements, n_test_functions, n_quad_points)

 required
test_grad_y_mat Tensor

Test function y-derivatives at quadrature points Shape: (n_elements, n_test_functions, n_quad_points)

 required
pred_nn Tensor

Neural network solution at quadrature points Shape: (n_elements, n_quad_points)

 required
pred_grad_x_nn Tensor

x-derivative of NN solution at quadrature points Shape: (n_elements, n_quad_points)

 required
pred_grad_y_nn Tensor

y-derivative of NN solution at quadrature points Shape: (n_elements, n_quad_points)

 required
forcing_function callable

Right-hand side forcing term

 required
bilinear_params dict

Dictionary containing: eps: Diffusion coefficient b_x: x-direction convection coefficient b_y: y-direction convection coefficient c: reaction coefficient

 required

Returns:

Type Description
Tensor

Cell-wise residuals averaged over test functions Shape: (n_cells,)
Note

The weak form includes: - Diffusion term: ∫ε∇u·∇v dΩ - Convection term: ∫(b·∇u)v dΩ - Reaction term: ∫cuv dΩ where ε, b, and c are known coefficients.

Source code in scirex/core/sciml/fastvpinns/physics/cd2d.py 
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
def pde_loss_cd2d(
    test_shape_val_mat: tf.Tensor,
    test_grad_x_mat: tf.Tensor,
    test_grad_y_mat: tf.Tensor,
    pred_nn: tf.Tensor,
    pred_grad_x_nn: tf.Tensor,
    pred_grad_y_nn: tf.Tensor,
    forcing_function: callable,
    bilinear_params: dict,
) -> tf.Tensor:
    """Calculates residual for 2D convection-diffusion problem.

    Implements the FastVPINNs methodology for computing variational residuals
    in 2D convection-diffusion equations with known coefficients using
    efficient tensor operations.

    Args:
        test_shape_val_mat: Test function values at quadrature points
            Shape: (n_elements, n_test_functions, n_quad_points)
        test_grad_x_mat: Test function x-derivatives at quadrature points
            Shape: (n_elements, n_test_functions, n_quad_points)
        test_grad_y_mat: Test function y-derivatives at quadrature points
            Shape: (n_elements, n_test_functions, n_quad_points)
        pred_nn: Neural network solution at quadrature points
            Shape: (n_elements, n_quad_points)
        pred_grad_x_nn: x-derivative of NN solution at quadrature points
            Shape: (n_elements, n_quad_points)
        pred_grad_y_nn: y-derivative of NN solution at quadrature points
            Shape: (n_elements, n_quad_points)
        forcing_function: Right-hand side forcing term
        bilinear_params: Dictionary containing:
            eps: Diffusion coefficient
            b_x: x-direction convection coefficient
            b_y: y-direction convection coefficient
            c: reaction coefficient

    Returns:
        Cell-wise residuals averaged over test functions
            Shape: (n_cells,)

    Note:
        The weak form includes:
        - Diffusion term: ∫ε∇u·∇v dΩ
        - Convection term: ∫(b·∇u)v dΩ
        - Reaction term: ∫cuv dΩ
        where ε, b, and c are known coefficients.
    """

    # Loss Function : ∫du/dx. dv/dx  +  ∫du/dy. dv/dy - ∫f.v

    # ∫du/dx. dv/dx dΩ
    pde_diffusion_x = tf.transpose(tf.linalg.matvec(test_grad_x_mat, pred_grad_x_nn))

    # ∫du/dy. dv/dy dΩ
    pde_diffusion_y = tf.transpose(tf.linalg.matvec(test_grad_y_mat, pred_grad_y_nn))

    # eps * ∫ (du/dx. dv/dx + du/dy. dv/dy) dΩ
    pde_diffusion = bilinear_params["eps"] * (pde_diffusion_x + pde_diffusion_y)

    # ∫du/dx. v dΩ
    conv_x = tf.transpose(tf.linalg.matvec(test_shape_val_mat, pred_grad_x_nn))

    # # ∫du/dy. v dΩ
    conv_y = tf.transpose(tf.linalg.matvec(test_shape_val_mat, pred_grad_y_nn))

    # # b(x) * ∫du/dx. v dΩ + b(y) * ∫du/dy. v dΩ
    conv = bilinear_params["b_x"] * conv_x + bilinear_params["b_y"] * conv_y

    # reaction term
    # ∫c.u.v dΩ
    reaction = bilinear_params["c"] * tf.transpose(
        tf.linalg.matvec(test_shape_val_mat, pred_nn)
    )

    residual_matrix = (pde_diffusion + conv + reaction) - forcing_function

    # Perform Reduce mean along the axis 0
    residual_cells = tf.reduce_mean(tf.square(residual_matrix), axis=0)

    return residual_cells