NLS-MB

Model Training CommandModel Evaluation CommandModel Export CommandModel Inference Command

# soliton
python NLS-MB_optical_soliton.py
# rogue wave
python NLS-MB_optical_rogue_wave.py

# soliton
python NLS-MB_optical_soliton.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/NLS-MB/NLS-MB_soliton_pretrained.pdparams
# rogue wave
python NLS-MB_optical_rogue_wave.py mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/NLS-MB/NLS-MB_rogue_wave_pretrained.pdparams

# soliton
python NLS-MB_optical_soliton.py mode=export
# rogue wave
python NLS-MB_optical_rogue_wave.py mode=export

# soliton
python NLS-MB_optical_soliton.py mode=infer
# rogue wave
python NLS-MB_optical_rogue_wave.py mode=infer

Pretrained Model	Metrics
NLS-MB_soliton_pretrained.pdparams	Residual/loss: 0.00000 Residual/MSE.Schrodinger_1: 0.00000 Residual/MSE.Schrodinger_2: 0.00000 Residual/MSE.Maxwell_1: 0.00000 Residual/MSE.Maxwell_2: 0.00000 Residual/MSE.Bloch: 0.00000
NLS-MB_optical_rogue_wave.pdparams	Residual/loss: 0.00001 Residual/MSE.Schrodinger_1: 0.00000 Residual/MSE.Schrodinger_2: 0.00000 Residual/MSE.Maxwell_1: 0.00000 Residual/MSE.Maxwell_2: 0.00000 Residual/MSE.Bloch: 0.00000

1. Background Introduction

Nonlinear localized wave dynamics, as an important branch of nonlinear science, covers basic forms of nonlinear localized waves such as solitons, breathers and rogue waves. Laser mode-locking technology provides an experimental verification platform for these theoretically predicted nonlinear localized waves. Through this technology, people have observed rich nonlinear phenomena such as soliton molecules and rogue waves, further promoting the research of nonlinear localized waves. Currently, research in this field has penetrated into many physical fields such as fluid mechanics, nonlinear optics, Bose-Einstein condensation (BEC), and plasma physics. In the field of optical fibers, the research of nonlinear dynamics is based on the principles of optical fiber optical devices, information processing, material design and signal transmission, and has played a key role in the development of fiber lasers, amplifiers, waveguides and communication technologies. The propagation dynamics of optical pulses in optical fibers are governed by nonlinear partial differential equations (such as the nonlinear Schrödinger equation NLSE). When dispersion and nonlinear effects coexist, these equations are often difficult to solve analytically. Therefore, the split-step Fourier method and its improved versions are widely used to study nonlinear effects in optical fibers. Its advantage lies in simple implementation and high relative accuracy. However, for long-distance and highly nonlinear scenarios, in order to meet accuracy requirements, the step size of the split-step Fourier method must be significantly reduced, which undoubtedly increases computational complexity, resulting in a huge number of grid point sets in the time domain and a long calculation process. PINN shows better performance than data-driven methods with much less data, and the computational complexity (expressed in multiples) is usually two orders of magnitude lower than SFM.

2. Problem Definition

In erbium-doped fiber, the propagation properties of optical pulses can be described by the coupled NLS-MB equations, which have the form

\[ \begin{cases} \dfrac{\partial E}{\partial x} = i \alpha_1 \dfrac{\partial^2 E}{\partial t ^2} - i \alpha_2 |E|^2 E+2 p \\ \dfrac{\partial p}{\partial t} = 2 i \omega_0 p+2 E \eta \\ \dfrac{\partial \eta}{\partial t} = -(E p^* + E^* p) \end{cases} \]

Among them, x, t represent the normalized propagation distance and time respectively, the complex envelope E is the slowly varying electric field, p is the measure of polarization of the resonant medium, \(\eta\) represents the degree of population inversion, and the symbol * represents complex conjugation. \(\alpha_1\) is the group velocity dispersion parameter, \(\alpha_2\) is the Kerr nonlinearity parameter, and is the offset measuring the resonance frequency. The NLS-MB system was first proposed by Maimistov and Manykin to describe the propagation of ultrashort pulses in Kerr nonlinear media. This system also plays an important role in solving the problem that optical fiber loss limits its transmission distance. In this equation, it describes the mixed state of self-induced transparency solitons and NLS solitons, called SIT-NLS solitons. These two types of solitons can coexist, and there have been many studies on their application in optical fiber communications.

2.1 Optical soliton

In the anomalous dispersion region of optical fibers, due to the interaction of dispersion and nonlinear effects, a very compelling phenomenon can be produced - optical solitons. "Soliton" is a special wave packet that can transmit long distances without deformation. Solitons have been widely studied in many branches of physics. The solitons in optical fibers discussed in this case not only have basic theoretical research value, but also have practical applications in optical fiber communications.

\[ \begin{gathered} E(x,t) = \frac{{2\exp ( - 2it)}}{{\cosh (2t + 6x)}}, \\ p(x,t) = \frac{{\exp ( - 2it)\left\{ {\exp ( - 2t - 6x) - \exp (2t + 6x)} \right\}}}{{\cosh {{(2t + 6x)}^2}}}, \\ \eta (x,t) = \frac{{\cosh {{(2t + 6x)}^2} - 2}}{{\cosh {{(2t + 6x)}^2}}}. \end{gathered} \]

We consider the computational domain as \([−1, 1] × [−1, 1]\). We first determine the optimization strategy. There are \(200\) points on each boundary, i.e., \(N_b = 2 × 200\). In order to calculate the equation loss of NLS-MB, \(20,000\) points are randomly selected within the domain.

2.2 Optical rogue wave

Optical rogue waves are a phenomenon in optics, similar to rogue waves in the ocean, but in optical systems. They are light waves that appear suddenly and have unusually high amplitudes. Optical rogue waves have some potential applications, especially in the fields of optical communications and laser technology. Some studies suggest that they can be used to enhance the transmission and processing of optical signals, or to generate ultrashort pulse lasers. We consider the computational domain as \([−0.5, 0.5] × [−2.5, 2.5]\)

3. Problem Solving

Next, we will explain how to convert the problem into PaddleScience code step by step and solve the problem using deep learning methods. In order to quickly understand PaddleScience, only key steps such as model construction, equation construction, and computational domain construction are described below, while other details please refer to API Documentation.

3.1 Model Construction

This paper uses the classic PINN MLP model for training.

model = ppsci.arch.MLP(**cfg.MODEL)

3.2 Equation Construction

Since Optical soliton uses the NLS-MB equation, NLSMB built in PaddleScience can be used directly.

equation = {
    "NLS-MB": ppsci.equation.NLSMB(alpha_1=0.5, alpha_2=-1, omega_0=-1, time=True)
}

3.3 Computational Domain Construction

In this paper, the Optical soliton problem acts on the spatiotemporal region of space (-1.0, 1.0), time (-1.0, 1.0), so the spatiotemporal geometry time_interval built in PaddleScience can be used directly as the computational domain.

geom = {
    "time_interval": ppsci.geometry.TimeXGeometry(
        ppsci.geometry.TimeDomain(t_lower, t_upper, timestamps=timestamps),
        ppsci.geometry.Interval(x_lower, x_upper),
    )
}

3.4 Constraint Construction

Since the dataset is an analytical solution, we first construct the analytical solution function

def analytic_solution(out):
    t, x = out["t"], out["x"]
    Eu_true = 2 * np.cos(2 * t) / np.cosh(2 * t + 6 * x)

    Ev_true = -2 * np.sin(2 * t) / np.cosh(2 * t + 6 * x)

    pu_true = (
        (np.exp(-2 * t - 6 * x) - np.exp(2 * t + 6 * x))
        * np.cos(2 * t)
        / np.cosh(2 * t + 6 * x) ** 2
    )
    pv_true = (
        -(np.exp(-2 * t - 6 * x) - np.exp(2 * t + 6 * x))
        * np.sin(2 * t)
        / np.cosh(2 * t + 6 * x) ** 2
    )
    eta_true = (np.cosh(2 * t + 6 * x) ** 2 - 2) / np.cosh(2 * t + 6 * x) ** 2

    return Eu_true, Ev_true, pu_true, pv_true, eta_true

3.4.1 Interior Point Constraint

Taking InteriorConstraint acting on internal points as an example, the code is as follows:

pde_constraint = ppsci.constraint.InteriorConstraint(
    equation["NLS-MB"].equations,
    {
        "Schrodinger_1": 0,
        "Schrodinger_2": 0,
        "Maxwell_1": 0,
        "Maxwell_2": 0,
        "Bloch": 0,
    },
    geom["time_interval"],
    {
        "dataset": {"name": "IterableNamedArrayDataset"},
        "batch_size": 20000,
        "iters_per_epoch": cfg.TRAIN.iters_per_epoch,
    },
    ppsci.loss.MSELoss(),
    evenly=True,
    name="EQ",
)

The first parameter of InteriorConstraint is the equation (system) expression, used to describe how to calculate the constraint target. Here, fill in equation["NLS-MB"].equations instantiated in the 3.2 Equation Construction chapter;

The second parameter is the target value of the constraint variable. In this problem, it is hoped that each equation of NLS-MB is optimized to 0;

The third parameter is the computational domain on which the constraint equation acts. Here, fill in geom["time_interval"] instantiated in the 3.3 Computational Domain Construction chapter;

The fourth parameter is the sampling configuration on the computational domain. Here batch_size is set to 20000.

The fifth parameter is the loss function. Here the commonly used MSE function is selected, and reduction is set to "mean", that is, the mean square error of all data points involved in the calculation will be calculated;

The sixth parameter is the name of the constraint condition. Each constraint condition needs to be named for subsequent indexing. Here it is named "EQ".

3.4.2 Boundary Constraint

Since our boundary points and initial value points have analytical solutions, we use supervised constraints

sup_constraint = ppsci.constraint.SupervisedConstraint(
    train_dataloader_cfg,
    ppsci.loss.MSELoss("mean"),
    name="Sup",
)

3.5 Hyperparameter Setting

Next, the number of training epochs and learning rate need to be specified. Here, based on experimental experience, 50,000 training epochs and an initial learning rate of 0.001 are used.

# training settings
TRAIN:
  epochs: 50000
  iters_per_epoch: 1
  lbfgs:
    iters_per_epoch: ${TRAIN.iters_per_epoch}
    output_dir: ${output_dir}LBFGS
    learning_rate: 1.0
    max_iter: 1
    eval_freq: ${TRAIN.eval_freq}
    eval_during_train: ${TRAIN.eval_during_train}
  eval_during_train: true
  eval_freq: 1000

3.6 Optimizer Construction

The training process will call the optimizer to update model parameters. Here, the more commonly used Adam optimizer is selected.

optimizer = ppsci.optimizer.Adam(learning_rate=cfg.TRAIN.learning_rate)(model)

3.7 Validator Construction

Usually during the training process, the training status of the current model is evaluated using the validation set (test set) at a certain epoch interval, so ppsci.validate.GeometryValidator is used to construct the validator.

residual_validator = ppsci.validate.GeometryValidator(
    equation["NLS-MB"].equations,
    {
        "Schrodinger_1": 0,
        "Schrodinger_2": 0,
        "Maxwell_1": 0,
        "Maxwell_2": 0,
        "Bloch": 0,
    },
    geom["time_interval"],
    {
        "dataset": "IterableNamedArrayDataset",
        "total_size": 20600,
    },
    ppsci.loss.MSELoss(),
    evenly=True,
    metric={"MSE": ppsci.metric.MSE()},
    with_initial=True,
    name="Residual",
)
validator = {residual_validator.name: residual_validator}

3.8 Visualizer Construction

After the model training is completed, we can take points in the computational domain for prediction, manually calculate the amplitude, and visualize the results.

vis_points = geom["time_interval"].sample_interior(20000, evenly=True)
Eu_true, Ev_true, pu_true, pv_true, eta_true = analytic_solution(vis_points)
pred = solver.predict(vis_points, return_numpy=True)
t = vis_points["t"][:, 0]
x = vis_points["x"][:, 0]
E_ref = np.sqrt(Eu_true**2 + Ev_true**2)[:, 0]
E_pred = np.sqrt(pred["Eu"] ** 2 + pred["Ev"] ** 2)[:, 0]
p_ref = np.sqrt(pu_true**2 + pv_true**2)[:, 0]
p_pred = np.sqrt(pred["pu"] ** 2 + pred["pv"] ** 2)[:, 0]
eta_ref = eta_true[:, 0]
eta_pred = pred["eta"][:, 0]

# plot
plot(t, x, E_ref, E_pred, p_ref, p_pred, eta_ref, eta_pred, cfg.output_dir)

3.9 Model Training, Evaluation and Visualization

3.9.1 Training with Adam

After completing the above settings, you only need to pass the instantiated objects to ppsci.solver.Solver in order, and then start training, evaluation, and visualization.

solver = ppsci.solver.Solver(
    model,
    constraint,
    cfg.output_dir,
    optimizer,
    epochs=cfg.TRAIN.epochs,
    iters_per_epoch=cfg.TRAIN.iters_per_epoch,
    eval_during_train=cfg.TRAIN.eval_during_train,
    eval_freq=cfg.TRAIN.eval_freq,
    equation=equation,
    geom=geom,
    validator=validator,
)
# train model
solver.train()
# evaluate after finished training
solver.eval()

3.9.2 Fine-tuning with L-BFGS [Optional]

After training with the Adam optimizer, we can replace the optimizer with the second-order optimizer L-BFGS to continue training for a small number of epochs (here we use 10% of the Adam optimization epochs), thereby further improving model accuracy.

OUTPUT_DIR = cfg.TRAIN.lbfgs.output_dir
logger.init_logger("ppsci", osp.join(OUTPUT_DIR, f"{cfg.mode}.log"), "info")
EPOCHS = cfg.TRAIN.epochs // 10
optimizer_lbfgs = ppsci.optimizer.LBFGS(
    cfg.TRAIN.lbfgs.learning_rate, cfg.TRAIN.lbfgs.max_iter
)(model)
solver = ppsci.solver.Solver(
    model,
    constraint,
    OUTPUT_DIR,
    optimizer_lbfgs,
    None,
    EPOCHS,
    cfg.TRAIN.lbfgs.iters_per_epoch,
    eval_during_train=cfg.TRAIN.lbfgs.eval_during_train,
    eval_freq=cfg.TRAIN.lbfgs.eval_freq,
    equation=equation,
    geom=geom,
    validator=validator,
)
# train model
solver.train()
# evaluate after finished training
solver.eval()

Tip

After training with conventional optimizers, using L-BFGS to fine-tune for a small number of epochs can effectively further improve model accuracy in most scenarios.

4. Complete Code

# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from os import path as osp

import hydra
import numpy as np
from matplotlib import pyplot as plt
from omegaconf import DictConfig

import ppsci
from ppsci.utils import logger


def analytic_solution(out):
    t, x = out["t"], out["x"]
    Eu_true = 2 * np.cos(2 * t) / np.cosh(2 * t + 6 * x)

    Ev_true = -2 * np.sin(2 * t) / np.cosh(2 * t + 6 * x)

    pu_true = (
        (np.exp(-2 * t - 6 * x) - np.exp(2 * t + 6 * x))
        * np.cos(2 * t)
        / np.cosh(2 * t + 6 * x) ** 2
    )
    pv_true = (
        -(np.exp(-2 * t - 6 * x) - np.exp(2 * t + 6 * x))
        * np.sin(2 * t)
        / np.cosh(2 * t + 6 * x) ** 2
    )
    eta_true = (np.cosh(2 * t + 6 * x) ** 2 - 2) / np.cosh(2 * t + 6 * x) ** 2

    return Eu_true, Ev_true, pu_true, pv_true, eta_true


def plot(
    t: np.ndarray,
    x: np.ndarray,
    E_ref: np.ndarray,
    E_pred: np.ndarray,
    p_ref: np.ndarray,
    p_pred: np.ndarray,
    eta_ref: np.ndarray,
    eta_pred: np.ndarray,
    output_dir: str,
):
    fig = plt.figure(figsize=(10, 10))
    plt.subplot(3, 3, 1)
    plt.title("E_ref")
    plt.tricontourf(x, t, E_ref, levels=256, cmap="jet")
    plt.subplot(3, 3, 2)
    plt.title("E_pred")
    plt.tricontourf(x, t, E_pred, levels=256, cmap="jet")
    plt.subplot(3, 3, 3)
    plt.title("E_diff")
    plt.tricontourf(x, t, np.abs(E_ref - E_pred), levels=256, cmap="jet")
    plt.subplot(3, 3, 4)
    plt.title("p_ref")
    plt.tricontourf(x, t, p_ref, levels=256, cmap="jet")
    plt.subplot(3, 3, 5)
    plt.title("p_pred")
    plt.tricontourf(x, t, p_pred, levels=256, cmap="jet")
    plt.subplot(3, 3, 6)
    plt.title("p_diff")
    plt.tricontourf(x, t, np.abs(p_ref - p_pred), levels=256, cmap="jet")
    plt.subplot(3, 3, 7)
    plt.title("eta_ref")
    plt.tricontourf(x, t, eta_ref, levels=256, cmap="jet")
    plt.subplot(3, 3, 8)
    plt.title("eta_pred")
    plt.tricontourf(x, t, eta_pred, levels=256, cmap="jet")
    plt.subplot(3, 3, 9)
    plt.title("eta_diff")
    plt.tricontourf(x, t, np.abs(eta_ref - eta_pred), levels=256, cmap="jet")
    fig_path = osp.join(output_dir, "pred_optical_soliton.png")
    print(f"Saving figure to {fig_path}")
    fig.savefig(fig_path, bbox_inches="tight", dpi=400)
    plt.close()


def train(cfg: DictConfig):
    # set model
    model = ppsci.arch.MLP(**cfg.MODEL)

    # set equation
    equation = {
        "NLS-MB": ppsci.equation.NLSMB(alpha_1=0.5, alpha_2=-1, omega_0=-1, time=True)
    }

    x_lower = -1
    x_upper = 1
    t_lower = -1
    t_upper = 1
    # set timestamps(including initial t0)
    timestamps = np.linspace(t_lower, t_upper, cfg.NTIME_ALL, endpoint=True)
    # set time-geometry
    geom = {
        "time_interval": ppsci.geometry.TimeXGeometry(
            ppsci.geometry.TimeDomain(t_lower, t_upper, timestamps=timestamps),
            ppsci.geometry.Interval(x_lower, x_upper),
        )
    }

    X, T = np.meshgrid(
        np.linspace(x_lower, x_upper, 256), np.linspace(t_lower, t_upper, 256)
    )
    X_star = np.hstack((X.flatten()[:, None], T.flatten()[:, None]))

    # Boundary and Initial conditions
    ic = X_star[:, 1] == t_lower
    idx_ic = np.random.choice(np.where(ic)[0], 200, replace=False)
    lb = X_star[:, 0] == x_lower
    idx_lb = np.random.choice(np.where(lb)[0], 200, replace=False)
    ub = X_star[:, 0] == x_upper
    idx_ub = np.random.choice(np.where(ub)[0], 200, replace=False)
    icbc_idx = np.hstack((idx_lb, idx_ic, idx_ub))
    X_u_train = X_star[icbc_idx].astype("float32")
    X_u_train = {"t": X_u_train[:, 1:2], "x": X_u_train[:, 0:1]}

    Eu_train, Ev_train, pu_train, pv_train, eta_train = analytic_solution(X_u_train)

    train_dataloader_cfg = {
        "dataset": {
            "name": "NamedArrayDataset",
            "input": {"t": X_u_train["t"], "x": X_u_train["x"]},
            "label": {
                "Eu": Eu_train,
                "Ev": Ev_train,
                "pu": pu_train,
                "pv": pv_train,
                "eta": eta_train,
            },
        },
        "batch_size": 600,
        "iters_per_epoch": cfg.TRAIN.iters_per_epoch,
    }

    # set constraint
    pde_constraint = ppsci.constraint.InteriorConstraint(
        equation["NLS-MB"].equations,
        {
            "Schrodinger_1": 0,
            "Schrodinger_2": 0,
            "Maxwell_1": 0,
            "Maxwell_2": 0,
            "Bloch": 0,
        },
        geom["time_interval"],
        {
            "dataset": {"name": "IterableNamedArrayDataset"},
            "batch_size": 20000,
            "iters_per_epoch": cfg.TRAIN.iters_per_epoch,
        },
        ppsci.loss.MSELoss(),
        evenly=True,
        name="EQ",
    )

    # supervised constraint s.t ||u-u_0||
    sup_constraint = ppsci.constraint.SupervisedConstraint(
        train_dataloader_cfg,
        ppsci.loss.MSELoss("mean"),
        name="Sup",
    )

    # wrap constraints together
    constraint = {
        pde_constraint.name: pde_constraint,
        sup_constraint.name: sup_constraint,
    }

    # set optimizer
    optimizer = ppsci.optimizer.Adam(learning_rate=cfg.TRAIN.learning_rate)(model)

    # set validator
    residual_validator = ppsci.validate.GeometryValidator(
        equation["NLS-MB"].equations,
        {
            "Schrodinger_1": 0,
            "Schrodinger_2": 0,
            "Maxwell_1": 0,
            "Maxwell_2": 0,
            "Bloch": 0,
        },
        geom["time_interval"],
        {
            "dataset": "IterableNamedArrayDataset",
            "total_size": 20600,
        },
        ppsci.loss.MSELoss(),
        evenly=True,
        metric={"MSE": ppsci.metric.MSE()},
        with_initial=True,
        name="Residual",
    )
    validator = {residual_validator.name: residual_validator}

    # initialize solver
    solver = ppsci.solver.Solver(
        model,
        constraint,
        cfg.output_dir,
        optimizer,
        epochs=cfg.TRAIN.epochs,
        iters_per_epoch=cfg.TRAIN.iters_per_epoch,
        eval_during_train=cfg.TRAIN.eval_during_train,
        eval_freq=cfg.TRAIN.eval_freq,
        equation=equation,
        geom=geom,
        validator=validator,
    )
    # train model
    solver.train()
    # evaluate after finished training
    solver.eval()

    # fine-tuning pretrained model with L-BFGS
    OUTPUT_DIR = cfg.TRAIN.lbfgs.output_dir
    logger.init_logger("ppsci", osp.join(OUTPUT_DIR, f"{cfg.mode}.log"), "info")
    EPOCHS = cfg.TRAIN.epochs // 10
    optimizer_lbfgs = ppsci.optimizer.LBFGS(
        cfg.TRAIN.lbfgs.learning_rate, cfg.TRAIN.lbfgs.max_iter
    )(model)
    solver = ppsci.solver.Solver(
        model,
        constraint,
        OUTPUT_DIR,
        optimizer_lbfgs,
        None,
        EPOCHS,
        cfg.TRAIN.lbfgs.iters_per_epoch,
        eval_during_train=cfg.TRAIN.lbfgs.eval_during_train,
        eval_freq=cfg.TRAIN.lbfgs.eval_freq,
        equation=equation,
        geom=geom,
        validator=validator,
    )
    # train model
    solver.train()
    # evaluate after finished training
    solver.eval()

    # visualize prediction
    vis_points = geom["time_interval"].sample_interior(20000, evenly=True)
    Eu_true, Ev_true, pu_true, pv_true, eta_true = analytic_solution(vis_points)
    pred = solver.predict(vis_points, return_numpy=True)
    t = vis_points["t"][:, 0]
    x = vis_points["x"][:, 0]
    E_ref = np.sqrt(Eu_true**2 + Ev_true**2)[:, 0]
    E_pred = np.sqrt(pred["Eu"] ** 2 + pred["Ev"] ** 2)[:, 0]
    p_ref = np.sqrt(pu_true**2 + pv_true**2)[:, 0]
    p_pred = np.sqrt(pred["pu"] ** 2 + pred["pv"] ** 2)[:, 0]
    eta_ref = eta_true[:, 0]
    eta_pred = pred["eta"][:, 0]

    # plot
    plot(t, x, E_ref, E_pred, p_ref, p_pred, eta_ref, eta_pred, cfg.output_dir)


def evaluate(cfg: DictConfig):
    # set model
    model = ppsci.arch.MLP(**cfg.MODEL)

    # set equation
    equation = {
        "NLS-MB": ppsci.equation.NLSMB(alpha_1=0.5, alpha_2=-1, omega_0=-1, time=True)
    }

    # set geometry
    x_lower = -1
    x_upper = 1
    t_lower = -1
    t_upper = 1
    # set timestamps(including initial t0)
    timestamps = np.linspace(t_lower, t_upper, cfg.NTIME_ALL, endpoint=True)
    # set time-geometry
    geom = {
        "time_interval": ppsci.geometry.TimeXGeometry(
            ppsci.geometry.TimeDomain(t_lower, t_upper, timestamps=timestamps),
            ppsci.geometry.Interval(x_lower, x_upper),
        )
    }

    # set validator
    residual_validator = ppsci.validate.GeometryValidator(
        equation["NLS-MB"].equations,
        {
            "Schrodinger_1": 0,
            "Schrodinger_2": 0,
            "Maxwell_1": 0,
            "Maxwell_2": 0,
            "Bloch": 0,
        },
        geom["time_interval"],
        {
            "dataset": "IterableNamedArrayDataset",
            "total_size": 20600,
        },
        ppsci.loss.MSELoss(),
        evenly=True,
        metric={"MSE": ppsci.metric.MSE()},
        with_initial=True,
        name="Residual",
    )
    validator = {residual_validator.name: residual_validator}

    # initialize solver
    solver = ppsci.solver.Solver(
        model,
        output_dir=cfg.output_dir,
        eval_freq=cfg.TRAIN.eval_freq,
        equation=equation,
        geom=geom,
        validator=validator,
        pretrained_model_path=cfg.EVAL.pretrained_model_path,
    )
    solver.eval()

    # visualize prediction
    vis_points = geom["time_interval"].sample_interior(20000, evenly=True)
    Eu_true, Ev_true, pu_true, pv_true, eta_true = analytic_solution(vis_points)
    pred = solver.predict(vis_points, return_numpy=True)
    t = vis_points["t"][:, 0]
    x = vis_points["x"][:, 0]
    E_ref = np.sqrt(Eu_true**2 + Ev_true**2)[:, 0]
    E_pred = np.sqrt(pred["Eu"] ** 2 + pred["Ev"] ** 2)[:, 0]
    p_ref = np.sqrt(pu_true**2 + pv_true**2)[:, 0]
    p_pred = np.sqrt(pred["pu"] ** 2 + pred["pv"] ** 2)[:, 0]
    eta_ref = eta_true[:, 0]
    eta_pred = pred["eta"][:, 0]

    # plot
    plot(t, x, E_ref, E_pred, p_ref, p_pred, eta_ref, eta_pred, cfg.output_dir)


def export(cfg: DictConfig):
    # set model
    model = ppsci.arch.MLP(**cfg.MODEL)

    # initialize solver
    solver = ppsci.solver.Solver(
        model,
        pretrained_model_path=cfg.INFER.pretrained_model_path,
    )
    # export model
    from paddle.static import InputSpec

    input_spec = [
        {key: InputSpec([None, 1], "float32", name=key) for key in model.input_keys},
    ]
    solver.export(input_spec, cfg.INFER.export_path)


def inference(cfg: DictConfig):
    from deploy.python_infer import pinn_predictor

    predictor = pinn_predictor.PINNPredictor(cfg)

    # set geometry
    x_lower = -1
    x_upper = 1
    t_lower = -1
    t_upper = 1
    # set timestamps(including initial t0)
    timestamps = np.linspace(t_lower, t_upper, cfg.NTIME_ALL, endpoint=True)
    # set time-geometry
    geom = {
        "time_interval": ppsci.geometry.TimeXGeometry(
            ppsci.geometry.TimeDomain(t_lower, t_upper, timestamps=timestamps),
            ppsci.geometry.Interval(x_lower, x_upper),
        )
    }

    NPOINT_TOTAL = cfg.NPOINT_INTERIOR + cfg.NPOINT_BC
    input_dict = geom["time_interval"].sample_interior(NPOINT_TOTAL, evenly=True)

    output_dict = predictor.predict(
        {key: input_dict[key] for key in cfg.MODEL.input_keys}, cfg.INFER.batch_size
    )

    # mapping data to cfg.INFER.output_keys
    output_dict = {
        store_key: output_dict[infer_key]
        for store_key, infer_key in zip(cfg.MODEL.output_keys, output_dict.keys())
    }

    # visualize prediction
    Eu_true, Ev_true, pu_true, pv_true, eta_true = analytic_solution(input_dict)
    t = input_dict["t"][:, 0]
    x = input_dict["x"][:, 0]
    E_ref = np.sqrt(Eu_true**2 + Ev_true**2)[:, 0]
    E_pred = np.sqrt(output_dict["Eu"] ** 2 + output_dict["Ev"] ** 2)[:, 0]
    p_ref = np.sqrt(pu_true**2 + pv_true**2)[:, 0]
    p_pred = np.sqrt(output_dict["pu"] ** 2 + output_dict["pv"] ** 2)[:, 0]
    eta_ref = eta_true[:, 0]
    eta_pred = output_dict["eta"][:, 0]

    # plot
    plot(t, x, E_ref, E_pred, p_ref, p_pred, eta_ref, eta_pred, cfg.output_dir)


@hydra.main(version_base=None, config_path="./conf", config_name="NLS-MB_soliton.yaml")
def main(cfg: DictConfig):
    if cfg.mode == "train":
        train(cfg)
    elif cfg.mode == "eval":
        evaluate(cfg)
    elif cfg.mode == "export":
        export(cfg)
    elif cfg.mode == "infer":
        inference(cfg)
    else:
        raise ValueError(
            f"cfg.mode should in ['train', 'eval', 'export', 'infer'], but got '{cfg.mode}'"
        )


if __name__ == "__main__":
    main()

5. Result Display

5.1 optical_soliton

Comparison between analytical solution results and PINN prediction results, from top to bottom: slowly varying electric field (E), resonance bias (p) and population inversion degree (eta)

5.2 optical_rogue_wave

Comparison between analytical solution results and PINN prediction results, from top to bottom: slowly varying electric field (E), resonance bias (p) and population inversion degree (eta)

It can be seen that the PINN prediction results are basically consistent with the analytical solution results.

6. References

1. S.-Y. Xu, Q. Zhou, and W. Liu, Prediction of Soliton Evolution and Equation Parameters for NLS–MB Equation Based on the phPINN Algorithm, Nonlinear Dyn (2023).