TimeDependentProblem, SpatialProblem, ParametricProblem

[karthikncsuiisc]

TimeDependentProblem, SpatialProblem, ParametricProblem

[ndem0]

Yes, it should be sufficient to create a new problem by inheriting by those classes. Like

class MyProblem(TimeDependentProblem, SpatialProblem, ParametricProblem):
    ...

Just remember to specify also the variable domains!

[karthikncsuiisc]

""" Burgers' problem. """

# ===================================================== #
#                                                       #
#  This script implements the one dimensional Burger    #
#  problem. The Burgers1D class is defined inheriting   #
#  from TimeDependentProblem, SpatialProblem and we     #
#  denote:                                              #
#           u --> field variable                        #
#           x --> spatial variable                      #
#           t --> temporal variable                     #
#                                                       #
# ===================================================== #


import torch
from pina.geometry import CartesianDomain
from pina import Condition
from pina.problem import TimeDependentProblem, SpatialProblem, ParametricProblem
from pina.operators import grad
from pina.equation import FixedValue, Equation


class Burgers1D(TimeDependentProblem, SpatialProblem, ParametricProblem):

    # define the burger equation
    def burger_equation(input_, output_):
        du = grad(output_, input_)
        ddu = grad(du, input_, components=['dudx'])
        mu1 = input_.extract(['mu1'])
        return (
                du.extract(['dudt']) +
                output_.extract(['u']) * du.extract(['dudx']) -
                (mu1 / torch.pi) * ddu.extract(['ddudxdx'])
        )

    # define initial condition
    def initial_condition(input_, output_):
        u_expected = -torch.sin(torch.pi * input_.extract(['x']))
        return output_.extract(['u']) - u_expected

    # assign output/ spatial and temporal variables
    output_variables = ['u']
    spatial_domain = CartesianDomain({'x': [-1, 1]})
    temporal_domain = CartesianDomain({'t': [0, 1]})
    mu1 = CartesianDomain({'mu1': [0.01, 0.02]})

    # problem condition statement
    conditions = {
        'gamma1': Condition(location=CartesianDomain({'x': -1, 't': [0, 1],'mu1':[0.01,0.02]}), equation=FixedValue(0.)),
        'gamma2': Condition(location=CartesianDomain({'x': 1, 't': [0, 1],'mu1':[0.01,0.02]}), equation=FixedValue(0.)),
        't0': Condition(location=CartesianDomain({'x': [-1, 1], 't': 0,'mu1':[0.01,0.02]}), equation=Equation(initial_condition)),
        'D': Condition(location=CartesianDomain({'x': [-1, 1], 't': [0, 1],'mu1':[0.01,0.02]}), equation=Equation(burger_equation)),
    }

""" Run PINA on Burgers equation. """

import argparse
import torch
from torch.nn import Softplus

from pina import LabelTensor
from pina.model import FeedForward
from pina.solvers import PINN
from pina.plotter import Plotter
from pina.trainer import Trainer


class myFeature(torch.nn.Module):
    """
    Feature: sin(pi*x)
    """

    def __init__(self):
        super(myFeature, self).__init__()

    def forward(self, x):
        return LabelTensor(torch.sin(torch.pi * x.extract(['x'])), ['sin(x)'])


if __name__ == "__main__":

    parser = argparse.ArgumentParser(description="Run PINA")
    parser.add_argument("--load", help="directory to save or load file", type=str)
    parser.add_argument("--features", help="extra features", type=int)
    parser.add_argument("--epochs", help="extra features", type=int, default=1000)
    args = parser.parse_args()

    if args.features is None:
        args.features = 0

    # extra features
    feat = [myFeature()] if args.features else []

    # create problem and discretise domain
    burgers_problem = Burgers1D()
    burgers_problem.discretise_domain(n=200, mode='grid', variables = 't', locations=['D'])
    burgers_problem.discretise_domain(n=20, mode='grid', variables = 'x', locations=['D'])
    burgers_problem.discretise_domain(n=20, mode='grid', variables = 'mu1', locations=['D'])
    burgers_problem.discretise_domain(n=150, mode='random', locations=['gamma1', 'gamma2', 't0'])

    # create model
    model = FeedForward(
        layers=[30, 20, 10, 5],
        output_dimensions=len(burgers_problem.output_variables),
        input_dimensions=len(burgers_problem.input_variables) + len(feat),
        func=Softplus
    )

    # create solver
    pinn = PINN(
        problem=burgers_problem,
        model=model,
        extra_features=feat,
        optimizer_kwargs={'lr' : 0.006}
    )

    # create trainer
    directory = 'pina.burger_extrafeats_{}'.format(bool(args.features))
    trainer = Trainer(solver=pinn, accelerator='cpu', max_epochs=args.epochs, default_root_dir=directory)


    if args.load:
        pinn = PINN.load_from_checkpoint(checkpoint_path=args.load, problem=burgers_problem, model=model)
        plotter = Plotter()
        plotter.plot(pinn)
    else:
        trainer.train()