168 lines
6.1 KiB
Python
168 lines
6.1 KiB
Python
import torch
|
|
|
|
from pina.problem import AbstractProblem
|
|
from pina import Condition, LabelTensor
|
|
from pina.solvers import GAROM
|
|
from pina.trainer import Trainer
|
|
import torch.nn as nn
|
|
import matplotlib.tri as tri
|
|
|
|
|
|
# def func(x, mu1, mu2):
|
|
# import torch
|
|
# x_m1 = (x[:, 0] - mu1).pow(2)
|
|
# x_m2 = (x[:, 1] - mu2).pow(2)
|
|
# norm = x[:, 0]**2 + x[:, 1]**2
|
|
# return torch.exp(-(x_m1 + x_m2))
|
|
|
|
|
|
# class ParametricGaussian(AbstractProblem):
|
|
# output_variables = [f'u_{i}' for i in range(900)]
|
|
|
|
# # params
|
|
# xx = torch.linspace(-1, 1, 20)
|
|
# yy = xx
|
|
# params = LabelTensor(torch.cartesian_prod(xx, yy), labels=['mu1', 'mu2'])
|
|
|
|
# # define domain
|
|
# x = torch.linspace(-1, 1, 30)
|
|
# domain = torch.cartesian_prod(x, x)
|
|
# triang = tri.Triangulation(domain[:, 0], domain[:, 1])
|
|
# sol = []
|
|
# for p in params:
|
|
# sol.append(func(domain, p[0], p[1]))
|
|
# snapshots = LabelTensor(torch.stack(sol), labels=output_variables)
|
|
|
|
# # define conditions
|
|
# conditions = {
|
|
# 'data': Condition(input_points=params, output_points=snapshots)
|
|
# }
|
|
|
|
|
|
# # simple Generator Network
|
|
# class Generator(nn.Module):
|
|
|
|
# def __init__(self,
|
|
# input_dimension,
|
|
# parameters_dimension,
|
|
# noise_dimension,
|
|
# activation=torch.nn.SiLU):
|
|
# super().__init__()
|
|
|
|
# self._noise_dimension = noise_dimension
|
|
# self._activation = activation
|
|
|
|
# self.model = torch.nn.Sequential(
|
|
# torch.nn.Linear(6 * self._noise_dimension, input_dimension // 6),
|
|
# self._activation(),
|
|
# torch.nn.Linear(input_dimension // 6, input_dimension // 3),
|
|
# self._activation(),
|
|
# torch.nn.Linear(input_dimension // 3, input_dimension))
|
|
# self.condition = torch.nn.Sequential(
|
|
# torch.nn.Linear(parameters_dimension, 2 * self._noise_dimension),
|
|
# self._activation(),
|
|
# torch.nn.Linear(2 * self._noise_dimension,
|
|
# 5 * self._noise_dimension))
|
|
|
|
# def forward(self, param):
|
|
# # uniform sampling in [-1, 1]
|
|
# z = torch.rand(size=(param.shape[0], self._noise_dimension),
|
|
# device=param.device,
|
|
# dtype=param.dtype,
|
|
# requires_grad=True)
|
|
# z = 2. * z - 1.
|
|
|
|
# # conditioning by concatenation of mapped parameters
|
|
# input_ = torch.cat((z, self.condition(param)), dim=-1)
|
|
# out = self.model(input_)
|
|
|
|
# return out
|
|
|
|
|
|
# # Simple Discriminator Network
|
|
# class Discriminator(nn.Module):
|
|
|
|
# def __init__(self,
|
|
# input_dimension,
|
|
# parameter_dimension,
|
|
# hidden_dimension,
|
|
# activation=torch.nn.ReLU):
|
|
# super().__init__()
|
|
|
|
# self._activation = activation
|
|
# self.encoding = torch.nn.Sequential(
|
|
# torch.nn.Linear(input_dimension, input_dimension // 3),
|
|
# self._activation(),
|
|
# torch.nn.Linear(input_dimension // 3, input_dimension // 6),
|
|
# self._activation(),
|
|
# torch.nn.Linear(input_dimension // 6, hidden_dimension))
|
|
# self.decoding = torch.nn.Sequential(
|
|
# torch.nn.Linear(2 * hidden_dimension, input_dimension // 6),
|
|
# self._activation(),
|
|
# torch.nn.Linear(input_dimension // 6, input_dimension // 3),
|
|
# self._activation(),
|
|
# torch.nn.Linear(input_dimension // 3, input_dimension),
|
|
# )
|
|
|
|
# self.condition = torch.nn.Sequential(
|
|
# torch.nn.Linear(parameter_dimension, hidden_dimension // 2),
|
|
# self._activation(),
|
|
# torch.nn.Linear(hidden_dimension // 2, hidden_dimension))
|
|
|
|
# def forward(self, data):
|
|
# x, condition = data
|
|
# encoding = self.encoding(x)
|
|
# conditioning = torch.cat((encoding, self.condition(condition)), dim=-1)
|
|
# decoding = self.decoding(conditioning)
|
|
# return decoding
|
|
|
|
|
|
# problem = ParametricGaussian()
|
|
|
|
|
|
# def test_constructor():
|
|
# GAROM(problem=problem,
|
|
# generator=Generator(input_dimension=900,
|
|
# parameters_dimension=2,
|
|
# noise_dimension=12),
|
|
# discriminator=Discriminator(input_dimension=900,
|
|
# parameter_dimension=2,
|
|
# hidden_dimension=64))
|
|
|
|
|
|
# def test_train_cpu():
|
|
# solver = GAROM(problem=problem,
|
|
# generator=Generator(input_dimension=900,
|
|
# parameters_dimension=2,
|
|
# noise_dimension=12),
|
|
# discriminator=Discriminator(input_dimension=900,
|
|
# parameter_dimension=2,
|
|
# hidden_dimension=64))
|
|
|
|
# trainer = Trainer(solver=solver, max_epochs=4, accelerator='cpu', batch_size=20)
|
|
# trainer.train()
|
|
|
|
|
|
# def test_sample():
|
|
# solver = GAROM(problem=problem,
|
|
# generator=Generator(input_dimension=900,
|
|
# parameters_dimension=2,
|
|
# noise_dimension=12),
|
|
# discriminator=Discriminator(input_dimension=900,
|
|
# parameter_dimension=2,
|
|
# hidden_dimension=64))
|
|
# solver.sample(problem.params)
|
|
# assert solver.sample(problem.params).shape == problem.snapshots.shape
|
|
|
|
|
|
# def test_forward():
|
|
# solver = GAROM(problem=problem,
|
|
# generator=Generator(input_dimension=900,
|
|
# parameters_dimension=2,
|
|
# noise_dimension=12),
|
|
# discriminator=Discriminator(input_dimension=900,
|
|
# parameter_dimension=2,
|
|
# hidden_dimension=64))
|
|
# solver(problem.params, mc_steps=100, variance=True)
|
|
# assert solver(problem.params).shape == problem.snapshots.shape
|