diff --git a/deepxde/icbc/boundary_conditions.py b/deepxde/icbc/boundary_conditions.py index 85732f4a4..27ddea6e5 100644 --- a/deepxde/icbc/boundary_conditions.py +++ b/deepxde/icbc/boundary_conditions.py @@ -248,19 +248,41 @@ class PointSetOperatorBC: and outputs a tensor of size `N x 1`, where `N` is the length of `inputs`. `inputs` and `outputs` are the network input and output tensors, respectively; `X` are the NumPy array of the `inputs`. + batch_size: The number of points per minibatch, or `None` to return all points. + This is only supported for the backend PyTorch and PaddlePaddle. + Note, If you want to use batch size here, you should also set callback + 'dde.callbacks.PDEPointResampler(bc_points=True)' in training. + shuffle: Randomize the order on each pass through the data when batching. """ - def __init__(self, points, values, func): + def __init__(self, points, values, func, batch_size=None, shuffle=True): self.points = np.array(points, dtype=config.real(np)) if not isinstance(values, numbers.Number) and values.shape[1] != 1: raise RuntimeError("PointSetOperatorBC should output 1D values") self.values = bkd.as_tensor(values, dtype=config.real(bkd.lib)) self.func = func + self.batch_size = batch_size + + if batch_size is not None: # batch iterator and state + if backend_name not in ["pytorch", "paddle"]: + raise RuntimeError( + "batch_size only implemented for pytorch and paddle backend" + ) + self.batch_sampler = data.sampler.BatchSampler(len(self), shuffle=shuffle) + self.batch_indices = None + + def __len__(self): + return self.points.shape[0] def collocation_points(self, X): + if self.batch_size is not None: + self.batch_indices = self.batch_sampler.get_next(self.batch_size) + return self.points[self.batch_indices] return self.points def error(self, X, inputs, outputs, beg, end, aux_var=None): + if self.batch_size is not None: + return self.func(inputs, outputs, X)[beg:end] - self.values[self.batch_indices] return self.func(inputs, outputs, X)[beg:end] - self.values