Implement Dataset, Dataloader and DataModule class and fix SupervisedSolver

tests/test_dataset.py

@@ -1,44 +1,45 @@
+import math
 import torch
-import pytest
-
-from pina.data.dataset import SamplePointDataset, SamplePointLoader, DataPointDataset
+from pina.data import SamplePointDataset, SupervisedDataset, PinaDataModule, UnsupervisedDataset, unsupervised_dataset
+from pina.data import PinaDataLoader
 from pina import LabelTensor, Condition
 from pina.equation import Equation
 from pina.domain import CartesianDomain
 from pina.problem import SpatialProblem
-from pina.model import FeedForward
 from pina.operators import laplacian
 from pina.equation.equation_factory import FixedValue
 
 
 def laplace_equation(input_, output_):
-    force_term = (torch.sin(input_.extract(['x'])*torch.pi) *
-                    torch.sin(input_.extract(['y'])*torch.pi))
+    force_term = (torch.sin(input_.extract(['x']) * torch.pi) *
+                  torch.sin(input_.extract(['y']) * torch.pi))
     delta_u = laplacian(output_.extract(['u']), input_)
     return delta_u - force_term
 
+
 my_laplace = Equation(laplace_equation)
 in_ = LabelTensor(torch.tensor([[0., 1.]]), ['x', 'y'])
 out_ = LabelTensor(torch.tensor([[0.]]), ['u'])
 in2_ = LabelTensor(torch.rand(60, 2), ['x', 'y'])
 out2_ = LabelTensor(torch.rand(60, 1), ['u'])
 
+
 class Poisson(SpatialProblem):
     output_variables = ['u']
     spatial_domain = CartesianDomain({'x': [0, 1], 'y': [0, 1]})
 
     conditions = {
         'gamma1': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y':  1}),
+            domain=CartesianDomain({'x': [0, 1], 'y': 1}),
             equation=FixedValue(0.0)),
         'gamma2': Condition(
-            location=CartesianDomain({'x': [0, 1], 'y': 0}),
+            domain=CartesianDomain({'x': [0, 1], 'y': 0}),
             equation=FixedValue(0.0)),
         'gamma3': Condition(
-            location=CartesianDomain({'x':  1, 'y': [0, 1]}),
+            domain=CartesianDomain({'x': 1, 'y': [0, 1]}),
             equation=FixedValue(0.0)),
         'gamma4': Condition(
-            location=CartesianDomain({'x': 0, 'y': [0, 1]}),
+            domain=CartesianDomain({'x': 0, 'y': [0, 1]}),
             equation=FixedValue(0.0)),
         'D': Condition(
             input_points=LabelTensor(torch.rand(size=(100, 2)), ['x', 'y']),
@@ -48,75 +49,114 @@ class Poisson(SpatialProblem):
             output_points=out_),
         'data2': Condition(
             input_points=in2_,
-            output_points=out2_)
+            output_points=out2_),
+        'unsupervised': Condition(
+            input_points=LabelTensor(torch.rand(size=(45, 2)), ['x', 'y']),
+            conditional_variables=LabelTensor(torch.ones(size=(45, 1)), ['alpha']),
+        ),
+        'unsupervised2': Condition(
+            input_points=LabelTensor(torch.rand(size=(90, 2)), ['x', 'y']),
+            conditional_variables=LabelTensor(torch.ones(size=(90, 1)), ['alpha']),
+        )
     }
 
+
 boundaries = ['gamma1', 'gamma2', 'gamma3', 'gamma4']
 poisson = Poisson()
 poisson.discretise_domain(10, 'grid', locations=boundaries)
 
+
 def test_sample():
     sample_dataset = SamplePointDataset(poisson, device='cpu')
     assert len(sample_dataset) == 140
-    assert sample_dataset.pts.shape == (140, 2)
-    assert sample_dataset.pts.labels == ['x', 'y']
-    assert sample_dataset.condition_indeces.dtype == torch.int64
-    assert sample_dataset.condition_indeces.max() == torch.tensor(4)
-    assert sample_dataset.condition_indeces.min() == torch.tensor(0)
+    assert sample_dataset.input_points.shape == (140, 2)
+    assert sample_dataset.input_points.labels == ['x', 'y']
+    assert sample_dataset.condition_indices.dtype == torch.uint8
+    assert sample_dataset.condition_indices.max() == torch.tensor(4)
+    assert sample_dataset.condition_indices.min() == torch.tensor(0)
+
 
 def test_data():
-    dataset = DataPointDataset(poisson, device='cpu')
+    dataset = SupervisedDataset(poisson, device='cpu')
     assert len(dataset) == 61
-    assert dataset.input_pts.shape == (61, 2)
-    assert dataset.input_pts.labels == ['x', 'y']
-    assert dataset.output_pts.shape == (61, 1 )
-    assert dataset.output_pts.labels == ['u']
-    assert dataset.condition_indeces.dtype == torch.int64
-    assert dataset.condition_indeces.max() == torch.tensor(1)
-    assert dataset.condition_indeces.min() == torch.tensor(0)
+    assert dataset['input_points'].shape == (61, 2)
+    assert dataset.input_points.shape == (61, 2)
+    assert dataset['input_points'].labels == ['x', 'y']
+    assert dataset.input_points.labels == ['x', 'y']
+    assert dataset['input_points', 3:].shape == (58, 2)
+    assert dataset[3:][1].labels == ['u']
+    assert dataset.output_points.shape == (61, 1)
+    assert dataset.output_points.labels == ['u']
+    assert dataset.condition_indices.dtype == torch.uint8
+    assert dataset.condition_indices.max() == torch.tensor(1)
+    assert dataset.condition_indices.min() == torch.tensor(0)
+
+
+def test_unsupervised():
+    dataset = UnsupervisedDataset(poisson, device='cpu')
+    assert len(dataset) == 135
+    assert dataset.input_points.shape == (135, 2)
+    assert dataset.input_points.labels == ['x', 'y']
+    assert dataset.input_points[3:].shape == (132, 2)
+
+    assert dataset.conditional_variables.shape == (135, 1)
+    assert dataset.conditional_variables.labels == ['alpha']
+    assert dataset.condition_indices.dtype == torch.uint8
+    assert dataset.condition_indices.max() == torch.tensor(1)
+    assert dataset.condition_indices.min() == torch.tensor(0)
+
+
+def test_data_module():
+    data_module = PinaDataModule(poisson, device='cpu')
+    data_module.setup()
+    loader = data_module.train_dataloader()
+    assert isinstance(loader, PinaDataLoader)
+    assert isinstance(loader, PinaDataLoader)
+
+    data_module = PinaDataModule(poisson, device='cpu', batch_size=10, shuffle=False)
+    data_module.setup()
+    loader = data_module.train_dataloader()
+    assert len(loader) == 24
+    for i in loader:
+        assert len(i) <= 10
+    len_ref = sum([math.ceil(len(dataset) * 0.7) for dataset in data_module.datasets])
+    len_real = sum([len(dataset) for dataset in data_module.splits['train'].values()])
+    assert len_ref == len_real
+
+    supervised_dataset = SupervisedDataset(poisson, device='cpu')
+    data_module = PinaDataModule(poisson, device='cpu', batch_size=10, shuffle=False, datasets=[supervised_dataset])
+    data_module.setup()
+    loader = data_module.train_dataloader()
+    for batch in loader:
+        assert len(batch) <= 10
+
+    physics_dataset = SamplePointDataset(poisson, device='cpu')
+    data_module = PinaDataModule(poisson, device='cpu', batch_size=10, shuffle=False, datasets=[physics_dataset])
+    data_module.setup()
+    loader = data_module.train_dataloader()
+    for batch in loader:
+        assert len(batch) <= 10
+
+    unsupervised_dataset = UnsupervisedDataset(poisson, device='cpu')
+    data_module = PinaDataModule(poisson, device='cpu', batch_size=10, shuffle=False, datasets=[unsupervised_dataset])
+    data_module.setup()
+    loader = data_module.train_dataloader()
+    for batch in loader:
+        assert len(batch) <= 10
+
 
 def test_loader():
-    sample_dataset = SamplePointDataset(poisson, device='cpu')
-    data_dataset = DataPointDataset(poisson, device='cpu')
-    loader = SamplePointLoader(sample_dataset, data_dataset, batch_size=10)
-
-    for batch in loader:
-        assert len(batch) in [2, 3]
-        assert batch['pts'].shape[0] <= 10
-        assert batch['pts'].requires_grad == True
-        assert batch['pts'].labels == ['x', 'y']
-
-    loader2 = SamplePointLoader(sample_dataset, data_dataset, batch_size=None)
-    assert len(list(loader2)) == 2
-
-def test_loader2():
-    poisson2 = Poisson()
-    del poisson.conditions['data2']
-    del poisson2.conditions['data']
-    poisson2.discretise_domain(10, 'grid', locations=boundaries)
-    sample_dataset = SamplePointDataset(poisson, device='cpu')
-    data_dataset = DataPointDataset(poisson, device='cpu')
-    loader = SamplePointLoader(sample_dataset, data_dataset, batch_size=10)
-
-    for batch in loader:
-        assert len(batch) == 2 # only phys condtions
-        assert batch['pts'].shape[0] <= 10
-        assert batch['pts'].requires_grad == True
-        assert batch['pts'].labels == ['x', 'y']
-
-def test_loader3():
-    poisson2 = Poisson()
-    del poisson.conditions['gamma1']
-    del poisson.conditions['gamma2']
-    del poisson.conditions['gamma3']
-    del poisson.conditions['gamma4']
-    del poisson.conditions['D']
-    sample_dataset = SamplePointDataset(poisson, device='cpu')
-    data_dataset = DataPointDataset(poisson, device='cpu')
-    loader = SamplePointLoader(sample_dataset, data_dataset, batch_size=10)
-
-    for batch in loader:
-        assert len(batch) == 2 # only phys condtions
-        assert batch['pts'].shape[0] <= 10
-        assert batch['pts'].requires_grad == True
-        assert batch['pts'].labels == ['x', 'y']
+    data_module = PinaDataModule(poisson, device='cpu', batch_size=10)
+    data_module.setup()
+    loader = data_module.train_dataloader()
+    assert isinstance(loader, PinaDataLoader)
+    assert len(loader) == 24
+    for i in loader:
+        assert len(i) <= 10
+        assert i.supervised.input_points.labels == ['x', 'y']
+        assert i.physics.input_points.labels == ['x', 'y']
+        assert i.unsupervised.input_points.labels == ['x', 'y']
+        assert i.supervised.input_points.requires_grad == True
+        assert i.physics.input_points.requires_grad == True
+        assert i.unsupervised.input_points.requires_grad == True
+test_loader()