🎨 Format Python code with psf/black

pina/dataset.py

@@ -17,38 +17,41 @@ class SamplePointDataset(Dataset):
         self.condition_names = []
 
         for name, condition in problem.conditions.items():
-            if not hasattr(condition, 'output_points'):
+            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.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 = [] 
+        output_list = []
         self.condition_names = []
 
         for name, condition in problem.conditions.items():
-            if hasattr(condition, 'output_points'):
+            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)
@@ -57,11 +60,14 @@ class DataPointDataset(Dataset):
         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.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([])
@@ -83,7 +89,9 @@ class SamplePointLoader:
     :vartype condition_names: list[str]
     """
 
-    def __init__(self, sample_dataset, data_dataset, batch_size=None, shuffle=True) -> None:
+    def __init__(
+        self, sample_dataset, data_dataset, batch_size=None, shuffle=True
+    ) -> None:
         """
         Constructor.
 
@@ -94,9 +102,13 @@ class SamplePointLoader:
             Default is ``True``.
         """
         if not isinstance(sample_dataset, SamplePointDataset):
-            raise TypeError(f'Expected SamplePointDataset, got {type(sample_dataset)}')
+            raise TypeError(
+                f"Expected SamplePointDataset, got {type(sample_dataset)}"
+            )
         if not isinstance(data_dataset, DataPointDataset):
-            raise TypeError(f'Expected DataPointDataset, got {type(data_dataset)}')
+            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)
@@ -106,25 +118,21 @@ class SamplePointLoader:
         self._prepare_data_dataset(data_dataset, batch_size, shuffle)
 
         self.condition_names = (
-            sample_dataset.condition_names + data_dataset.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)
-            )
+            self.batch_list.append(("sample", i))
 
         for i in range(len(self.batch_input_pts)):
-            self.batch_list.append(
-                ('data', i)
-            )
+            self.batch_list.append(("data", i))
 
         if shuffle:
-            self.random_idx = torch.randperm(len(self.batch_list))   
+            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.
@@ -157,17 +165,18 @@ class SamplePointLoader:
             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_input_pts = torch.tensor_split(dataset.input_pts, batch_num)
         self.batch_output_pts = torch.tensor_split(
-            dataset.output_pts, batch_num)
+            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)
+            self.tensor_conditions, batch_num
+        )
 
     def _prepare_sample_dataset(self, dataset, batch_size, shuffle):
         """
@@ -190,7 +199,7 @@ class SamplePointLoader:
         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
 
@@ -198,13 +207,14 @@ class SamplePointLoader:
         #     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)
+            self.tensor_conditions, batch_num
+        )
 
     def __iter__(self):
         """
@@ -222,20 +232,20 @@ class SamplePointLoader:
         :return: An iterator over the points.
         :rtype: iter
         """
-        #for i in self.random_idx:
+        # for i in self.random_idx:
         for i in range(len(self.batch_list)):
             type_, idx_ = self.batch_list[i]
 
-            if type_ == 'sample':
+            if type_ == "sample":
                 d = {
-                    'pts': self.batch_sample_pts[idx_].requires_grad_(True),
-                    'condition': self.batch_sample_conditions[idx_],
+                    "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_],
+                    "pts": self.batch_input_pts[idx_].requires_grad_(True),
+                    "output": self.batch_output_pts[idx_],
+                    "condition": self.batch_data_conditions[idx_],
                 }
             yield d
 
@@ -246,4 +256,4 @@ class SamplePointLoader:
         :return: The number of batches.
         :rtype: int
         """
-        return len(self.batch_list)
+        return len(self.batch_list)