fix tests
This commit is contained in:
Nicola Demo
@@ -8,160 +8,160 @@ 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))
|
||||
# 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)]
|
||||
# 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'])
|
||||
# # 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 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)
|
||||
}
|
||||
# # define conditions
|
||||
# conditions = {
|
||||
# 'data': Condition(input_points=params, output_points=snapshots)
|
||||
# }
|
||||
|
||||
|
||||
# simple Generator Network
|
||||
class Generator(nn.Module):
|
||||
# # simple Generator Network
|
||||
# class Generator(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
input_dimension,
|
||||
parameters_dimension,
|
||||
noise_dimension,
|
||||
activation=torch.nn.SiLU):
|
||||
super().__init__()
|
||||
# def __init__(self,
|
||||
# input_dimension,
|
||||
# parameters_dimension,
|
||||
# noise_dimension,
|
||||
# activation=torch.nn.SiLU):
|
||||
# super().__init__()
|
||||
|
||||
self._noise_dimension = noise_dimension
|
||||
self._activation = activation
|
||||
# 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))
|
||||
# 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.
|
||||
# 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_)
|
||||
# # conditioning by concatenation of mapped parameters
|
||||
# input_ = torch.cat((z, self.condition(param)), dim=-1)
|
||||
# out = self.model(input_)
|
||||
|
||||
return out
|
||||
# return out
|
||||
|
||||
|
||||
# Simple Discriminator Network
|
||||
class Discriminator(nn.Module):
|
||||
# # Simple Discriminator Network
|
||||
# class Discriminator(nn.Module):
|
||||
|
||||
def __init__(self,
|
||||
input_dimension,
|
||||
parameter_dimension,
|
||||
hidden_dimension,
|
||||
activation=torch.nn.ReLU):
|
||||
super().__init__()
|
||||
# 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._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))
|
||||
# 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
|
||||
# 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()
|
||||
# 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_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))
|
||||
# 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()
|
||||
# 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_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
|
||||
# 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
|
||||
|
||||
Reference in New Issue
Block a user