41 lines
1.4 KiB
Python
41 lines
1.4 KiB
Python
from torch.utils.data import Dataset
|
|
import torch
|
|
from ..label_tensor import LabelTensor
|
|
|
|
|
|
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.stack(input_list)
|
|
self.output_pts = LabelTensor.stack(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] |