260 lines
8.9 KiB
Python
260 lines
8.9 KiB
Python
from torch.utils.data import Dataset
|
|
import torch
|
|
from ..label_tensor import LabelTensor
|
|
|
|
|
|
class SamplePointDataset(Dataset):
|
|
"""
|
|
This class is used to create a dataset of sample points.
|
|
"""
|
|
|
|
def __init__(self, problem, device) -> None:
|
|
"""
|
|
:param dict input_pts: The input points.
|
|
"""
|
|
super().__init__()
|
|
pts_list = []
|
|
self.condition_names = []
|
|
|
|
for name, condition in problem.conditions.items():
|
|
if not hasattr(condition, "output_points"):
|
|
pts_list.append(problem.input_pts[name])
|
|
self.condition_names.append(name)
|
|
|
|
self.pts = LabelTensor.vstack(pts_list)
|
|
|
|
if self.pts != []:
|
|
self.condition_indeces = torch.cat(
|
|
[
|
|
torch.tensor([i] * len(pts_list[i]))
|
|
for i in range(len(self.condition_names))
|
|
],
|
|
dim=0,
|
|
)
|
|
else: # if there are no sample points
|
|
self.condition_indeces = torch.tensor([])
|
|
self.pts = torch.tensor([])
|
|
|
|
self.pts = self.pts.to(device)
|
|
self.condition_indeces = self.condition_indeces.to(device)
|
|
|
|
def __len__(self):
|
|
return self.pts.shape[0]
|
|
|
|
|
|
class DataPointDataset(Dataset):
|
|
|
|
def __init__(self, problem, device) -> None:
|
|
super().__init__()
|
|
input_list = []
|
|
output_list = []
|
|
self.condition_names = []
|
|
|
|
for name, condition in problem.conditions.items():
|
|
if hasattr(condition, "output_points"):
|
|
input_list.append(problem.conditions[name].input_points)
|
|
output_list.append(problem.conditions[name].output_points)
|
|
self.condition_names.append(name)
|
|
|
|
self.input_pts = LabelTensor.vstack(input_list)
|
|
self.output_pts = LabelTensor.vstack(output_list)
|
|
|
|
if self.input_pts != []:
|
|
self.condition_indeces = torch.cat(
|
|
[
|
|
torch.tensor([i] * len(input_list[i]))
|
|
for i in range(len(self.condition_names))
|
|
],
|
|
dim=0,
|
|
)
|
|
else: # if there are no data points
|
|
self.condition_indeces = torch.tensor([])
|
|
self.input_pts = torch.tensor([])
|
|
self.output_pts = torch.tensor([])
|
|
|
|
self.input_pts = self.input_pts.to(device)
|
|
self.output_pts = self.output_pts.to(device)
|
|
self.condition_indeces = self.condition_indeces.to(device)
|
|
|
|
def __len__(self):
|
|
return self.input_pts.shape[0]
|
|
|
|
|
|
class SamplePointLoader:
|
|
"""
|
|
This class is used to create a dataloader to use during the training.
|
|
|
|
:var condition_names: The names of the conditions. The order is consistent
|
|
with the condition indeces in the batches.
|
|
:vartype condition_names: list[str]
|
|
"""
|
|
|
|
def __init__(
|
|
self, sample_dataset, data_dataset, batch_size=None, shuffle=True
|
|
) -> None:
|
|
"""
|
|
Constructor.
|
|
|
|
:param SamplePointDataset sample_pts: The sample points dataset.
|
|
:param int batch_size: The batch size. If ``None``, the batch size is
|
|
set to the number of sample points. Default is ``None``.
|
|
:param bool shuffle: If ``True``, the sample points are shuffled.
|
|
Default is ``True``.
|
|
"""
|
|
if not isinstance(sample_dataset, SamplePointDataset):
|
|
raise TypeError(
|
|
f"Expected SamplePointDataset, got {type(sample_dataset)}"
|
|
)
|
|
if not isinstance(data_dataset, DataPointDataset):
|
|
raise TypeError(
|
|
f"Expected DataPointDataset, got {type(data_dataset)}"
|
|
)
|
|
|
|
self.n_data_conditions = len(data_dataset.condition_names)
|
|
self.n_phys_conditions = len(sample_dataset.condition_names)
|
|
data_dataset.condition_indeces += self.n_phys_conditions
|
|
|
|
self._prepare_sample_dataset(sample_dataset, batch_size, shuffle)
|
|
self._prepare_data_dataset(data_dataset, batch_size, shuffle)
|
|
|
|
self.condition_names = (
|
|
sample_dataset.condition_names + data_dataset.condition_names
|
|
)
|
|
|
|
self.batch_list = []
|
|
for i in range(len(self.batch_sample_pts)):
|
|
self.batch_list.append(("sample", i))
|
|
|
|
for i in range(len(self.batch_input_pts)):
|
|
self.batch_list.append(("data", i))
|
|
|
|
if shuffle:
|
|
self.random_idx = torch.randperm(len(self.batch_list))
|
|
else:
|
|
self.random_idx = torch.arange(len(self.batch_list))
|
|
|
|
def _prepare_data_dataset(self, dataset, batch_size, shuffle):
|
|
"""
|
|
Prepare the dataset for data points.
|
|
|
|
:param SamplePointDataset dataset: The dataset.
|
|
:param int batch_size: The batch size.
|
|
:param bool shuffle: If ``True``, the sample points are shuffled.
|
|
"""
|
|
self.sample_dataset = dataset
|
|
|
|
if len(dataset) == 0:
|
|
self.batch_data_conditions = []
|
|
self.batch_input_pts = []
|
|
self.batch_output_pts = []
|
|
return
|
|
|
|
if batch_size is None:
|
|
batch_size = len(dataset)
|
|
batch_num = len(dataset) // batch_size
|
|
if len(dataset) % batch_size != 0:
|
|
batch_num += 1
|
|
|
|
output_labels = dataset.output_pts.labels
|
|
input_labels = dataset.input_pts.labels
|
|
self.tensor_conditions = dataset.condition_indeces
|
|
|
|
if shuffle:
|
|
idx = torch.randperm(dataset.input_pts.shape[0])
|
|
self.input_pts = dataset.input_pts[idx]
|
|
self.output_pts = dataset.output_pts[idx]
|
|
self.tensor_conditions = dataset.condition_indeces[idx]
|
|
|
|
self.batch_input_pts = torch.tensor_split(dataset.input_pts, batch_num)
|
|
self.batch_output_pts = torch.tensor_split(
|
|
dataset.output_pts, batch_num
|
|
)
|
|
|
|
for i in range(len(self.batch_input_pts)):
|
|
self.batch_input_pts[i].labels = input_labels
|
|
self.batch_output_pts[i].labels = output_labels
|
|
|
|
self.batch_data_conditions = torch.tensor_split(
|
|
self.tensor_conditions, batch_num
|
|
)
|
|
|
|
def _prepare_sample_dataset(self, dataset, batch_size, shuffle):
|
|
"""
|
|
Prepare the dataset for sample points.
|
|
|
|
:param DataPointDataset dataset: The dataset.
|
|
:param int batch_size: The batch size.
|
|
:param bool shuffle: If ``True``, the sample points are shuffled.
|
|
"""
|
|
|
|
self.sample_dataset = dataset
|
|
if len(dataset) == 0:
|
|
self.batch_sample_conditions = []
|
|
self.batch_sample_pts = []
|
|
return
|
|
|
|
if batch_size is None:
|
|
batch_size = len(dataset)
|
|
|
|
batch_num = len(dataset) // batch_size
|
|
if len(dataset) % batch_size != 0:
|
|
batch_num += 1
|
|
|
|
self.tensor_pts = dataset.pts
|
|
self.tensor_conditions = dataset.condition_indeces
|
|
|
|
# if shuffle:
|
|
# idx = torch.randperm(self.tensor_pts.shape[0])
|
|
# self.tensor_pts = self.tensor_pts[idx]
|
|
# self.tensor_conditions = self.tensor_conditions[idx]
|
|
|
|
self.batch_sample_pts = torch.tensor_split(self.tensor_pts, batch_num)
|
|
for i in range(len(self.batch_sample_pts)):
|
|
self.batch_sample_pts[i].labels = dataset.pts.labels
|
|
|
|
self.batch_sample_conditions = torch.tensor_split(
|
|
self.tensor_conditions, batch_num
|
|
)
|
|
|
|
def __iter__(self):
|
|
"""
|
|
Return an iterator over the points. Any element of the iterator is a
|
|
dictionary with the following keys:
|
|
- ``pts``: The input sample points. It is a LabelTensor with the
|
|
shape ``(batch_size, input_dimension)``.
|
|
- ``output``: The output sample points. This key is present only
|
|
if data conditions are present. It is a LabelTensor with the
|
|
shape ``(batch_size, output_dimension)``.
|
|
- ``condition``: The integer condition indeces. It is a tensor
|
|
with the shape ``(batch_size, )`` of type ``torch.int64`` and
|
|
indicates for any ``pts`` the corresponding problem condition.
|
|
|
|
:return: An iterator over the points.
|
|
:rtype: iter
|
|
"""
|
|
# for i in self.random_idx:
|
|
for i in range(len(self.batch_list)):
|
|
type_, idx_ = self.batch_list[i]
|
|
|
|
if type_ == "sample":
|
|
d = {
|
|
"pts": self.batch_sample_pts[idx_].requires_grad_(True),
|
|
"condition": self.batch_sample_conditions[idx_],
|
|
}
|
|
else:
|
|
d = {
|
|
"pts": self.batch_input_pts[idx_].requires_grad_(True),
|
|
"output": self.batch_output_pts[idx_],
|
|
"condition": self.batch_data_conditions[idx_],
|
|
}
|
|
yield d
|
|
|
|
def __len__(self):
|
|
"""
|
|
Return the number of batches.
|
|
|
|
:return: The number of batches.
|
|
:rtype: int
|
|
"""
|
|
return len(self.batch_list)
|