text: # Copyright 2020 The Trieste Contributors
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """ This module contains utilities for :class:`~trieste.observer.Observer` data. """
  from __future__ import annotations
  
  from dataclasses import dataclass
  
  import tensorflow as tf
  
  from trieste.type import TensorType
  
  
  @dataclass(frozen=True)
  class Dataset:
      """
      Container for the query points and corresponding observations from an
      :class:`~trieste.observer.Observer`.
      """
  
      query_points: TensorType
      """ The points at which the :class:`~trieste.observer.Observer` was queried. """
  
      observations: TensorType
      """ The observed output of the :class:`~trieste.observer.Observer` for each query point. """
  
      def __post_init__(self) -> None:
          """
          :raise ValueError (or InvalidArgumentError): If ``query_points`` or ``observations`` have \
              rank less than two, or they have unequal shape in any but their last dimension.
          """
          tf.debugging.assert_rank_at_least(self.query_points, 2)
          tf.debugging.assert_rank_at_least(self.observations, 2)
  
          if 0 in (self.query_points.shape[-1], self.observations.shape[-1]):
              raise ValueError(
                  f"query_points and observations cannot have dimension 0, got shapes"
                  f" {self.query_points.shape} and {self.observations.shape}."
              )
  
          if self.query_points.shape[:-1] != self.observations.shape[:-1]:
              raise ValueError(
                  f"Leading shapes of query_points and observations must match. Got shapes"
                  f" {self.query_points.shape}, {self.observations.shape}."
              )
  
      def __add__(self, rhs: Dataset) -> Dataset:
          r"""
          Return the :class:`Dataset` whose query points are the result of concatenating the
          `query_points` in each :class:`Dataset` along the zeroth axis, and the same for the
          `observations`. For example:
  
          >>> d1 = Dataset(
          ...     tf.constant([[0.1, 0.2], [0.3, 0.4]]),
          ...     tf.constant([[0.5, 0.6], [0.7, 0.8]])
          ... )
          >>> d2 = Dataset(tf.constant([[0.9, 1.0]]), tf.constant([[1.1, 1.2]]))
          >>> (d1 + d2).query_points
          <tf.Tensor: shape=(3, 2), dtype=float32, numpy=
          array([[0.1, 0.2],
                 [0.3, 0.4],
                 [0.9, 1. ]], dtype=float32)>
          >>> (d1 + d2).observations
          <tf.Tensor: shape=(3, 2), dtype=float32, numpy=
          array([[0.5, 0.6],
                 [0.7, 0.8],
                 [1.1, 1.2]], dtype=float32)>
  
          :param rhs: A :class:`Dataset` with the same shapes as this one, except in the zeroth
              dimension, which can have any size.
          :return: The result of concatenating the :class:`Dataset`\ s.
          :raise InvalidArgumentError: If the shapes of the `query_points` in each :class:`Dataset`
              differ in any but the zeroth dimension. The same applies for `observations`.
          """
          return Dataset(
              tf.concat([self.query_points, rhs.query_points], axis=0),
              tf.concat([self.observations, rhs.observations], axis=0),
          )
  
      def __len__(self) -> TensorType:
          """
          :return: The number of query points, or equivalently the number of observations.
          """
          return tf.shape(self.observations)[0]
  
      def __deepcopy__(self, memo: dict[int, object]) -> Dataset:
          return self
  
      def astuple(self) -> tuple[TensorType, TensorType]:
          """
          **Note:** Unlike the standard library function `dataclasses.astuple`, this method does
          **not** deepcopy the attributes.
  
          :return: A 2-tuple of the :attr:`query_points` and :attr:`observations`.
          """
          return self.query_points, self.observations
  