[WIP] PINO PDE

docs/src/manual/pino_pde.md

@@ -0,0 +1 @@
+   #TODO

src/discretize.jl

@@ -369,8 +369,16 @@ which is later optimized upon to give Solution or the Solution Distribution of t
 
 For more information, see `discretize` and `PINNRepresentation`.
 """
-function SciMLBase.symbolic_discretize(pde_system::PDESystem,
-        discretization::AbstractPINN)
+#TODO?
+function SciMLBase.symbolic_discretize(pde_system::PDESystem, discretization::PhysicsInformedNO)
+
+end
+#TODO?
+function SciMLBase.symbolic_discretize(
+        pde_system::PDESystem, discretization::PhysicsInformedNN)
+end
+
+function SciMLBase.symbolic_discretize(pde_system::PDESystem, discretization::AbstractPINN)
     eqs = pde_system.eqs
     bcs = pde_system.bcs
     chain = discretization.chain
@@ -718,3 +726,11 @@ function SciMLBase.discretize(pde_system::PDESystem, discretization::PhysicsInfo
         Optimization.AutoZygote())
     Optimization.OptimizationProblem(f, pinnrep.flat_init_params)
 end
+
+#TODO?
+function SciMLBase.discretize(pde_system::PDESystem, discretization::AbstractPINN)
+    pinnrep = symbolic_discretize(pde_system, discretization)
+    f = OptimizationFunction(pinnrep.loss_functions.full_loss_function,
+        Optimization.AutoZygote())
+    Optimization.OptimizationProblem(f, pinnrep.flat_init_params)
+end

src/pinn_types.jl

@@ -307,6 +307,97 @@ struct BayesianPINN{T, P, PH, DER, PE, AL, ADA, LOG, D, K} <: AbstractPINN
     end
 end
 
+
+"""
+PhysicsInformedNO
+"""
+struct PhysicsInformedNO{T, PS, P, PH, DER, AL, ADA, LOG, K} <: AbstractPINN
+    chain::Any
+    parameters::PS
+    strategy::T
+    init_params::P
+    phi::PH
+    derivative::DER
+    additional_loss::AL
+    adaptive_loss::ADA #?
+    logger::LOG #?
+    log_options::LogOptions #?
+    iteration::Vector{Int64} #?
+    self_increment::Bool#?
+    multioutput::Bool
+    kwargs::K
+
+    @add_kwonly function PhysicsInformedNN(chain,
+            strategy,
+            parameters;
+            init_params = nothing,
+            phi = nothing,
+            derivative = nothing,
+            additional_loss = nothing,
+            adaptive_loss = nothing,
+            logger = nothing,
+            log_options = LogOptions(),
+            iteration = nothing,
+            kwargs...)
+        multioutput = chain isa AbstractArray
+        if multioutput
+            !all(i -> i isa Lux.AbstractExplicitLayer, chain) &&
+                (chain = Lux.transform.(chain))
+        else
+            !(chain isa Lux.AbstractExplicitLayer) &&
+                (chain = adapt(FromFluxAdaptor(false, false), chain))
+        end
+        if phi === nothing
+            if multioutput
+                _phi = Phi.(chain)
+            else
+                _phi = Phi(chain)
+            end
+        else
+            if multioutput
+                all([phi.f[i] isa Lux.AbstractExplicitLayer for i in eachindex(phi.f)]) ||
+                    throw(ArgumentError("Only Lux Chains are supported"))
+            else
+                (phi.f isa Lux.AbstractExplicitLayer) ||
+                    throw(ArgumentError("Only Lux Chains are supported"))
+            end
+            _phi = phi
+        end
+
+        if derivative === nothing
+            _derivative = numeric_derivative
+        else
+            _derivative = derivative
+        end
+
+        if iteration isa Vector{Int64}
+            self_increment = false
+        else
+            iteration = [1]
+            self_increment = true
+        end
+
+        new{typeof(parameters), typeof(strategy), typeof(init_params),
+            typeof(_phi), typeof(_derivative),
+            typeof(additional_loss), typeof(adaptive_loss), typeof(logger), typeof(kwargs)}(
+            chain,
+            parameters,
+            strategy,
+            init_params,
+            _phi,
+            _derivative,
+            additional_loss,
+            adaptive_loss,
+            logger,
+            log_options,
+            iteration,
+            self_increment,
+            multioutput,
+            kwargs)
+    end
+end
+
+
 """
 `PINNRepresentation``
 

src/pino_pde.jl

@@ -0,0 +1,26 @@
+
+#TODO?
+function PhysicsInformedNO(
+        neural_operator,
+        parameters,#bounds
+        strategy;
+        kwargs...
+    )
+    PhysicsInformedNN(neural_operator,
+        strategy; kwargs...
+    )
+end
+#TODO?
+function SciMLBase.discretize(pde_system::PDESystem, neural_operator::PhysicsInformedNO)
+    pinnrep = symbolic_discretize(pde_system, neural_operator)
+    f = OptimizationFunction(pinnrep.loss_functions.full_loss_function,
+        Optimization.AutoZygote())
+    Optimization.OptimizationProblem(f, pinnrep.flat_init_params)
+end
+#TODO?
+function SciMLBase.discretize(pde_system::PDESystem, neural_operator::PhysicsInformed)
+    pinnrep = symbolic_discretize(pde_system, neural_operator)
+    f = OptimizationFunction(pinnrep.loss_functions.full_loss_function,
+        Optimization.AutoZygote())
+    Optimization.OptimizationProblem(f, pinnrep.flat_init_params)
+end

test/PINO_PDE_tests.jl

@@ -0,0 +1,62 @@
+using Test
+using Optimization, OptimizationOptimJL, OptimizationOptimisers
+using Lux
+using Statistics, Random
+using ModelingToolkit
+import ModelingToolkit: Interval, infimum, supremum
+using DomainSets
+using NeuralPDE
+using LuxNeuralOperators
+
+
+##example ODE
+@parameters t
+@variables u(..)
+@parameters p #[bounds = (0.1f0, pi)]
+Dt = Differential(t)
+eq = [Dt(u(t)) ~ cos(t)]
+bc = [u(0) ~ 1.0f0]
+
+dom = [t ∈ Interval(0.0, 1.0)]
+# neural_operator = SomeNeuralOperator(some_args)
+neural_operator = Lux.Chain(
+    Lux.Dense(1, 10, Lux.tanh),
+    Lux.Dense(10, 10, Lux.tanh),
+    Lux.Dense(10, 1))
+
+# pino = PhysicsInformedNO(neural_operator, sometrainig)
+pino = NeuralPDE.PhysicsInformedNN(neural_operator, NeuralPDE.GridTraining(0.1))
+
+@named pde_system = PDESystem(eq, bc, dom, [t], [u(t)],[p])  #[p]; defaults = Dict([p => 1.0 for p in [p]]))
+
+# hasbounds(pde_system.ps[1])
+# getbounds(pde_system.ps[1])
+
+prob = discretize(pde_system, pino)
+sym_prob = symbolic_discretize(pde_system, pino)
+
+res = Optimization.solve(prob, ADAM(0.1); maxiters = 4000)
+phi = discretization.phi
+
+
+@parameters x
+@parameters p [bounds = (0.1f0, pi)]
+@variables u(..)
+Dx = Differential(x)
+
+eq = Dx(u(x)) ~ cos(p * x)
+bcs = [u(0.0) ~ 0.0]
+
+domains = [x ∈ Interval(0.0, 1.0)]
+
+chain = [Lux.Chain(Lux.Dense(1, 12, Lux.tanh), Lux.Dense(12, 12, Lux.tanh), Lux.Dense(12, 1)) ]
+
+strategy = NeuralPDE.GridTraining(0.1)
+discretization = NeuralPDE.PhysicsInformedNN(chain, strategy)
+
+@named pde_system = PDESystem(eq, bcs, domains, [x], [u(x)],[p]; defaults = Dict([p => 1.0 for p in [p]]))
+hasbounds(pde_system.ps[1])
+getbounds(pde_system.ps[1])
+
+prob = discretize(pde_system, discretization)
+sym_prob = symbolic_discretize(pde_system, discretization)