43f69242ab
* add Averaging Neural Operator with tests * add backward test * minor changes * doc addition --------- Co-authored-by: Dario Coscia <dariocoscia@Dario-Coscia.local>
63 lines
1.7 KiB
Python
63 lines
1.7 KiB
Python
import torch
|
|
from pina.model import AveragingNeuralOperator
|
|
from pina import LabelTensor
|
|
|
|
output_numb_fields = 5
|
|
batch_size = 15
|
|
|
|
|
|
def test_constructor():
|
|
input_numb_fields = 1
|
|
output_numb_fields = 1
|
|
#minimuum constructor
|
|
AveragingNeuralOperator(input_numb_fields,
|
|
output_numb_fields,
|
|
coordinates_indices=['p'],
|
|
field_indices=['v'])
|
|
|
|
#all constructor
|
|
AveragingNeuralOperator(input_numb_fields,
|
|
output_numb_fields,
|
|
inner_size=5,
|
|
n_layers=5,
|
|
func=torch.nn.ReLU,
|
|
coordinates_indices=['p'],
|
|
field_indices=['v'])
|
|
|
|
|
|
def test_forward():
|
|
input_numb_fields = 1
|
|
output_numb_fields = 1
|
|
dimension = 1
|
|
input_ = LabelTensor(
|
|
torch.rand(batch_size, 1000, input_numb_fields + dimension), ['p', 'v'])
|
|
ano = AveragingNeuralOperator(input_numb_fields,
|
|
output_numb_fields,
|
|
dimension=dimension,
|
|
coordinates_indices=['p'],
|
|
field_indices=['v'])
|
|
out = ano(input_)
|
|
assert out.shape == torch.Size(
|
|
[batch_size, input_.shape[1], output_numb_fields])
|
|
|
|
|
|
def test_backward():
|
|
input_numb_fields = 1
|
|
dimension = 1
|
|
output_numb_fields = 1
|
|
input_ = LabelTensor(
|
|
torch.rand(batch_size, 1000, dimension + input_numb_fields),
|
|
['p', 'v'])
|
|
input_ = input_.requires_grad_()
|
|
avno = AveragingNeuralOperator(input_numb_fields,
|
|
output_numb_fields,
|
|
dimension=dimension,
|
|
coordinates_indices=['p'],
|
|
field_indices=['v'])
|
|
out = avno(input_)
|
|
tmp = torch.linalg.norm(out)
|
|
tmp.backward()
|
|
grad = input_.grad
|
|
assert grad.shape == torch.Size(
|
|
[batch_size, input_.shape[1], dimension + input_numb_fields])
|