jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def sample_image_prefix(
    params,
    model,
    # length,
    rng,
    dataloader,
    prefix=300,
    # bsz=32,
    imshape=(28, 28),
    n_batches=None,
    save=True,
):
    """Sample a grayscale image represented as intensities in [0, 255]"""
    import matplotlib.pyplot as plt
    import numpy as onp

    # from .data import Datasets
    # BATCH = bsz
    # start = np.zeros((BATCH, length), dtype=int)
    # start = np.zeros((BATCH, length, 1), dtype=int)
    start = np.array(next(iter(dataloader))[0].numpy())
    start = np.zeros_like(start)
    # params, prime, cache = init_recurrence(model, params, start[:, :-1], rng)
    params, prime, cache = init_recurrence(model, params, start, rng)

    BATCH = start.shape[0]
    START = prefix
    LENGTH = start.shape[1]
    assert LENGTH == onp.multiply.reduce(imshape)

    # _, dataloader, _, _, _ = Datasets["mnist"](bsz=BATCH)
    it = iter(dataloader)
    for j, im in enumerate(it):
        if n_batches is not None and j >= n_batches:
            break

        image = im[0].numpy()
        image = np.pad(
            image[:, :-1, :], [(0, 0), (1, 0), (0, 0)], constant_values=0
        )
        cur = onp.array(image)
        # cur[:, START + 1 :, 0] = 0
        # cur = np.pad(cur[:, :-1, 0], [(0, 0), (1, 0)], constant_values=256)
        cur = np.array(cur[:, :])

        # Cache the first `start` inputs.
        out, vars = model.apply(
            {"params": params, "prime": prime, "cache": cache},
            cur[:, np.arange(0, START)],
            mutable=["cache"],
        )
        cache = vars["cache"].unfreeze()
        out = sample(model, params, prime, cache, cur, START, LENGTH - 1, rng)

        # Visualization
        out = out.reshape(BATCH, *imshape)
        final = onp.zeros((BATCH, *imshape, 3))
        final2 = onp.zeros((BATCH, *imshape, 3))
        final[:, :, :, 0] = out
        f = final.reshape(BATCH, LENGTH, 3)
        i = image.reshape(BATCH, LENGTH)
        f[:, :START, 1] = i[:, :START]
        f[:, :START, 2] = i[:, :START]
        f = final2.reshape(BATCH, LENGTH, 3)
        f[:, :, 1] = i
        f[:, :START, 0] = i[:, :START]
        f[:, :START, 2] = i[:, :START]
        if save:
            for k in range(BATCH):
                fig, (ax1, ax2) = plt.subplots(ncols=2)
                ax1.set_title("Sampled")
                ax1.imshow(final[k] / 256.0)
                ax2.set_title("True")
                ax1.axis("off")
                ax2.axis("off")
                ax2.imshow(final2[k] / 256.0)
                fig.savefig("im%d.%d.png" % (j, k))
                plt.close()
                print(f"Sampled batch {j} image {k}")
    return final, final2
//  [python]
def sample_image_prefix(
    params,
    model,
    # length,
    rng,
    dataloader,
    prefix=300,
    # bsz=32,
    imshape=(28, 28),
    n_batches=None,
    save=True,
):
    """Sample a grayscale image represented as intensities in [0, 255]"""
    import matplotlib.pyplot as plt
    import numpy as onp

    # from .data import Datasets
    # BATCH = bsz
    # start = np.zeros((BATCH, length), dtype=int)
    # start = np.zeros((BATCH, length, 1), dtype=int)
    start = np.array(next(iter(dataloader))[0].numpy())
    start = np.zeros_like(start)
    # params, prime, cache = init_recurrence(model, params, start[:, :-1], rng)
    params, prime, cache = init_recurrence(model, params, start, rng)

    BATCH = start.shape[0]
    START = prefix
    LENGTH = start.shape[1]
    assert LENGTH == onp.prod(imshape)

    # _, dataloader, _, _, _ = Datasets["mnist"](bsz=BATCH)
    it = iter(dataloader)
    for j, im in enumerate(it):
        if n_batches is not None and j >= n_batches:
            break

        image = im[0].numpy()
        image = np.pad(
            image[:, :-1, :], [(0, 0), (1, 0), (0, 0)], constant_values=0
        )
        cur = onp.array(image)
        # cur[:, START + 1 :, 0] = 0
        # cur = np.pad(cur[:, :-1, 0], [(0, 0), (1, 0)], constant_values=256)
        cur = np.array(cur[:, :])

        # Cache the first `start` inputs.
        out, vars = model.apply(
            {"params": params, "prime": prime, "cache": cache},
            cur[:, np.arange(0, START)],
            mutable=["cache"],
        )
        cache = vars["cache"].unfreeze()
        out = sample(model, params, prime, cache, cur, START, LENGTH - 1, rng)

        # Visualization
        out = out.reshape(BATCH, *imshape)
        final = onp.zeros((BATCH, *imshape, 3))
        final2 = onp.zeros((BATCH, *imshape, 3))
        final[:, :, :, 0] = out
        f = final.reshape(BATCH, LENGTH, 3)
        i = image.reshape(BATCH, LENGTH)
        f[:, :START, 1] = i[:, :START]
        f[:, :START, 2] = i[:, :START]
        f = final2.reshape(BATCH, LENGTH, 3)
        f[:, :, 1] = i
        f[:, :START, 0] = i[:, :START]
        f[:, :START, 2] = i[:, :START]
        if save:
            for k in range(BATCH):
                fig, (ax1, ax2) = plt.subplots(ncols=2)
                ax1.set_title("Sampled")
                ax1.imshow(final[k] / 256.0)
                ax2.set_title("True")
                ax1.axis("off")
                ax2.axis("off")
                ax2.imshow(final2[k] / 256.0)
                fig.savefig("im%d.%d.png" % (j, k))
                plt.close()
                print(f"Sampled batch {j} image {k}")
    return final, final2
//  [python]
def averageError(self, image0, image1):
    self.assertEqual(image0.shape, image1.shape)
    image0 = image0.astype(int)  # Avoid overflow
    return np.abs(image0 - image1).sum() / np.prod(image0.shape)
//  [python]
def _mult_op(*array_list): return numpy.prod(numpy.stack(array_list), axis=0)
//  [python]
def _generateArray(self, shape):
    """Generate an array of the given shape for use in testing.
    The numbers are calculated as the cumulative sum, which
    causes the difference between neighboring numbers to vary."""

    # Flattened length of the array.
    flat_len = reduce(lambda x, y: x*y, shape)

    a = np.array(range(flat_len), dtype=int)
    a = np.cumsum(a)
    a = a.reshape(shape)

    return a
//  [python]
def _generateArray(self, shape):
    """Generate an array of the given shape for use in testing.
    The numbers are calculated as the cumulative sum, which
    causes the difference between neighboring numbers to vary."""

    # Flattened length of the array.
    flat_len = np.prod(shape)

    a = np.array(range(flat_len), dtype=int)
    a = np.cumsum(a)
    a = a.reshape(shape)

    return a
//  [python]
def testCompareBilinear(self):
    if test.is_gpu_available():
      input_shape = [1, 5, 6, 3]
      target_height = 8
      target_width = 12
      for nptype in [np.float32, np.float64]:
        for align_corners in [True, False]:
          img_np = np.arange(
              0, np.multiply.reduce(input_shape), dtype=nptype).reshape(input_shape)
          value = {}
          for use_gpu in [True, False]:
            with self.test_session(use_gpu=use_gpu):
              image = constant_op.constant(img_np, shape=input_shape)
              new_size = constant_op.constant([target_height, target_width])
              out_op = image_ops.resize_images(
                  image,
                  new_size,
                  image_ops.ResizeMethod.BILINEAR,
                  align_corners=align_corners)
              value[use_gpu] = out_op.eval()
          self.assertAllClose(value[True], value[False], rtol=1e-5, atol=1e-5)
//  [python]
def testCompareBilinear(self):
    if test.is_gpu_available():
      input_shape = [1, 5, 6, 3]
      target_height = 8
      target_width = 12
      for nptype in [np.float32, np.float64]:
        for align_corners in [True, False]:
          img_np = np.arange(
              0, np.prod(input_shape), dtype=nptype).reshape(input_shape)
          value = {}
          for use_gpu in [True, False]:
            with self.test_session(use_gpu=use_gpu):
              image = constant_op.constant(img_np, shape=input_shape)
              new_size = constant_op.constant([target_height, target_width])
              out_op = image_ops.resize_images(
                  image,
                  new_size,
                  image_ops.ResizeMethod.BILINEAR,
                  align_corners=align_corners)
              value[use_gpu] = out_op.eval()
          self.assertAllClose(value[True], value[False], rtol=1e-5, atol=1e-5)
//  [python]
def testCompareNearestNeighbor(self):
    if test.is_gpu_available():
      input_shape = [1, 5, 6, 3]
      target_height = 8
      target_width = 12
      for nptype in [np.float32, np.float64]:
        for align_corners in [True, False]:
          img_np = np.arange(
              0, np.multiply.reduce(input_shape), dtype=nptype).reshape(input_shape)
          with self.test_session(use_gpu=True):
            image = constant_op.constant(img_np, shape=input_shape)
            new_size = constant_op.constant([target_height, target_width])
            out_op = image_ops.resize_images(
                image,
                new_size,
                image_ops.ResizeMethod.NEAREST_NEIGHBOR,
                align_corners=align_corners)
            gpu_val = out_op.eval()
          with self.test_session(use_gpu=False):
            image = constant_op.constant(img_np, shape=input_shape)
            new_size = constant_op.constant([target_height, target_width])
            out_op = image_ops.resize_images(
                image,
                new_size,
                image_ops.ResizeMethod.NEAREST_NEIGHBOR,
                align_corners=align_corners)
            cpu_val = out_op.eval()
          self.assertAllClose(cpu_val, gpu_val, rtol=1e-5, atol=1e-5)
//  [python]
def testCompareNearestNeighbor(self):
    if test.is_gpu_available():
      input_shape = [1, 5, 6, 3]
      target_height = 8
      target_width = 12
      for nptype in [np.float32, np.float64]:
        for align_corners in [True, False]:
          img_np = np.arange(
              0, np.prod(input_shape), dtype=nptype).reshape(input_shape)
          with self.test_session(use_gpu=True):
            image = constant_op.constant(img_np, shape=input_shape)
            new_size = constant_op.constant([target_height, target_width])
            out_op = image_ops.resize_images(
                image,
                new_size,
                image_ops.ResizeMethod.NEAREST_NEIGHBOR,
                align_corners=align_corners)
            gpu_val = out_op.eval()
          with self.test_session(use_gpu=False):
            image = constant_op.constant(img_np, shape=input_shape)
            new_size = constant_op.constant([target_height, target_width])
            out_op = image_ops.resize_images(
                image,
                new_size,
                image_ops.ResizeMethod.NEAREST_NEIGHBOR,
                align_corners=align_corners)
            cpu_val = out_op.eval()
          self.assertAllClose(cpu_val, gpu_val, rtol=1e-5, atol=1e-5)
//  [python]
def testBasic(self, data_type):
    x_shape = [13, 9, 3]
    x_np = np.arange(0, np.prod(x_shape), dtype=data_type).reshape(x_shape)
    y_np = self._NumpyPerImageWhitening(x_np)

    with self.cached_session():
      x = constant_op.constant(x_np, dtype=data_type, shape=x_shape)
      y = image_ops.per_image_standardization(x)
      y_tf = self.evaluate(y)
      self.assertAllClose(y_tf, y_np, atol=1e-4)
//  [python]
def d_up_calculation(s_bar_path, flow_accum_path, target_d_up_path):
    """Calculate d_up = s_bar * sqrt(upslope area)."""
    cell_area_m2 = abs(numpy.multiply.reduce(pygeoprocessing.get_raster_info(
        s_bar_path)['pixel_size']))
    pygeoprocessing.raster_map(
        op=lambda s_bar, flow_accum: (
            s_bar * numpy.sqrt(flow_accum * cell_area_m2)),
        rasters=[s_bar_path, flow_accum_path],
        target_path=target_d_up_path,
        target_dtype=numpy.float32,
        target_nodata=_TARGET_NODATA)
//  [python]
def d_up_calculation(s_bar_path, flow_accum_path, target_d_up_path):
    """Calculate d_up = s_bar * sqrt(upslope area)."""
    cell_area_m2 = abs(numpy.prod(pygeoprocessing.get_raster_info(
        s_bar_path)['pixel_size']))
    pygeoprocessing.raster_map(
        op=lambda s_bar, flow_accum: (
            s_bar * numpy.sqrt(flow_accum * cell_area_m2)),
        rasters=[s_bar_path, flow_accum_path],
        target_path=target_d_up_path,
        target_dtype=numpy.float32,
        target_nodata=_TARGET_NODATA)
//  [python]
def _calculate_load(lulc_raster_path, lucode_to_load, target_load_raster):
    """Calculate load raster by mapping landcover and multiplying by area.

    Args:
        lulc_raster_path (string): path to integer landcover raster.
        lucode_to_load (dict): a mapping of landcover IDs to per-area
            nutrient load.
        target_load_raster (string): path to target raster that will have
            total load per pixel.

    Returns:
        None.

    """
    cell_area_ha = abs(n.prod(pygeoprocessing.get_raster_info(
        lulc_raster_path)['pixel_size'])) * 0.0001

    def _map_load_op(lucode_array):
        """Convert unit load to total load & handle nodata."""
        result = numpy.empty(lucode_array.shape)
        for lucode in numpy.unique(lucode_array):
            try:
                result[lucode_array == lucode] = (
                    lucode_to_load[lucode] * cell_area_ha)
            except KeyError:
                raise KeyError(
                    'lucode: %d is present in the landuse raster but '
                    'missing from the biophysical table' % lucode)
        return result

    pygeoprocessing.raster_map(
        op=_map_load_op,
        rasters=[lulc_raster_path],
        target_path=target_load_raster,
        target_dtype=numpy.float32,
        target_nodata=_TARGET_NODATA)
//  [python]
def _calculate_load(lulc_raster_path, lucode_to_load, target_load_raster):
    """Calculate load raster by mapping landcover and multiplying by area.

    Args:
        lulc_raster_path (string): path to integer landcover raster.
        lucode_to_load (dict): a mapping of landcover IDs to per-area
            nutrient load.
        target_load_raster (string): path to target raster that will have
            total load per pixel.

    Returns:
        None.

    """
    cell_area_ha = abs(numpy.prod(pygeoprocessing.get_raster_info(
        lulc_raster_path)['pixel_size'])) * 0.0001

    def _map_load_op(lucode_array):
        """Convert unit load to total load & handle nodata."""
        result = numpy.empty(lucode_array.shape)
        for lucode in numpy.unique(lucode_array):
            try:
                result[lucode_array == lucode] = (
                    lucode_to_load[lucode] * cell_area_ha)
            except KeyError:
                raise KeyError(
                    'lucode: %d is present in the landuse raster but '
                    'missing from the biophysical table' % lucode)
        return result

    pygeoprocessing.raster_map(
        op=_map_load_op,
        rasters=[lulc_raster_path],
        target_path=target_load_raster,
        target_dtype=numpy.float32,
        target_nodata=_TARGET_NODATA)
//  [python]
def testWholeImageBoundingBoxStateless(self):
    height = 40
    width = 50
    image_size = [height, width, 1]
    bounding_box = [0.0, 0.0, 1.0, 1.0]
    image = np.arange(
        0, np.prod(image_size), dtype=np.int32).reshape(image_size)
    for min_obj_covered in [0.1, constant_op.constant(0.1)]:
      self._testStatelessSampleDistortedBoundingBox(
          image,
          bounding_box,
          min_object_covered=min_obj_covered,
          aspect_ratio_range=(0.75, 1.33),
          area_range=(0.05, 1.0))
//  [python]
def test_ctsplitter_train_short_intervals(point_process_dataset):
    splitter = transform.ContinuousTimeInstanceSplitter(
        0.01,
        0.01,
        instance_sampler=MockContinuousTimeSampler(
            ret_values=[1.01, 1.5, 1.99]
        ),
        freq=to_offset("H"),
    )

    iter_de = splitter(point_process_dataset, is_train=True)

    for d in iter_de:
        assert d["future_valid_length"] == d["past_valid_length"] == 0
        assert np.prod(np.shape(d["past_target"])) == 0
        assert np.prod(np.shape(d["future_target"])) == 0
//  [python]
def testCompareBilinear(self):
    if test.is_gpu_available():
      input_shape = [1, 5, 6, 3]
      target_height = 8
      target_width = 12
      for nptype in [np.float32, np.float64]:
        img_np = np.arange(
            0, np.prod(input_shape), dtype=nptype).reshape(input_shape)
        value = {}
        for use_gpu in [True, False]:
          with self.cached_session(use_gpu=use_gpu):
            image = constant_op.constant(img_np, shape=input_shape)
            new_size = constant_op.constant([target_height, target_width])
            out_op = image_ops.resize_images(image, new_size,
                                             image_ops.ResizeMethod.BILINEAR)
            value[use_gpu] = self.evaluate(out_op)
        self.assertAllClose(value[True], value[False], rtol=1e-5, atol=1e-5)
//  [python]
def testCompareNearestNeighbor(self):
    if test.is_gpu_available():
      input_shape = [1, 5, 6, 3]
      target_height = 8
      target_width = 12
      for nptype in [np.float32, np.float64]:
        img_np = np.arange(
            0, np.multiply.reduce(input_shape), dtype=nptype).reshape(input_shape)
        with self.cached_session():
          image = constant_op.constant(img_np, shape=input_shape)
          new_size = constant_op.constant([target_height, target_width])
          out_op = image_ops.resize_images_v2(
              image, new_size, image_ops.ResizeMethod.NEAREST_NEIGHBOR)
          gpu_val = self.evaluate(out_op)
        with self.cached_session(use_gpu=False):
          image = constant_op.constant(img_np, shape=input_shape)
          new_size = constant_op.constant([target_height, target_width])
          out_op = image_ops.resize_images_v2(
              image, new_size, image_ops.ResizeMethod.NEAREST_NEIGHBOR)
          cpu_val = self.evaluate(out_op)
        self.assertAllClose(cpu_val, gpu_val, rtol=1e-5, atol=1e-5)
//  [python]
def testCompareNearestNeighbor(self):
    if test.is_gpu_available():
      input_shape = [1, 5, 6, 3]
      target_height = 8
      target_width = 12
      for nptype in [np.float32, np.float64]:
        img_np = np.arange(
            0, np.prod(input_shape), dtype=nptype).reshape(input_shape)
        with self.cached_session():
          image = constant_op.constant(img_np, shape=input_shape)
          new_size = constant_op.constant([target_height, target_width])
          out_op = image_ops.resize_images_v2(
              image, new_size, image_ops.ResizeMethod.NEAREST_NEIGHBOR)
          gpu_val = self.evaluate(out_op)
        with self.cached_session(use_gpu=False):
          image = constant_op.constant(img_np, shape=input_shape)
          new_size = constant_op.constant([target_height, target_width])
          out_op = image_ops.resize_images_v2(
              image, new_size, image_ops.ResizeMethod.NEAREST_NEIGHBOR)
          cpu_val = self.evaluate(out_op)
        self.assertAllClose(cpu_val, gpu_val, rtol=1e-5, atol=1e-5)
//  [python]
def weights(dim, degree=3):
        """
        Gauss-Hermite quadrature weights.

        Parameters
        ----------
        dim : int
            Dimension of the input random variable.

        degree : int, optional
            Degree of the integration rule.

        Returns
        -------
        : (num_points, ) ndarray
            GH quadrature weights of given degree.
        """
        # 1D sigma-points (x) and weights (w)
        x, w = hermegauss(degree)
        # hermegauss() provides weights that cause posdef errors
        w = factorial(degree) / (degree ** 2 * hermeval(x, [0] * (degree - 1) + [1]) ** 2)
        return n.prod(cartesian([w] * dim), axis=1)
//  [python]
def weights(dim, degree=3):
        """
        Gauss-Hermite quadrature weights.

        Parameters
        ----------
        dim : int
            Dimension of the input random variable.

        degree : int, optional
            Degree of the integration rule.

        Returns
        -------
        : (num_points, ) ndarray
            GH quadrature weights of given degree.
        """
        # 1D sigma-points (x) and weights (w)
        x, w = hermegauss(degree)
        # hermegauss() provides weights that cause posdef errors
        w = factorial(degree) / (degree ** 2 * hermeval(x, [0] * (degree - 1) + [1]) ** 2)
        return np.prod(cartesian([w] * dim), axis=1)
//  [python]
def _testStatelessSampleDistortedBoundingBox(self, image, bounding_box,
                                               min_object_covered,
                                               aspect_ratio_range, area_range):
    with test_util.use_gpu():
      original_area = float(np.multiply.reduce(image.shape))
      bounding_box_area = float((bounding_box[3] - bounding_box[1]) *
                               (bounding_box[2] - bounding_box[0]))

      image_size_np = np.array(image.shape, dtype=np.int32)
      bounding_box_np = (
          np.array(bounding_box, dtype=np.float32).reshape([1, 1, 4]))

      iterations = 2
      test_seeds = [(1, 2), (3, 4), (5, 6)]

      for seed in test_seeds:
        aspect_ratios = []
        area_ratios = []
        fraction_object_covered = []
        for _ in range(iterations):
          image_tf = constant_op.constant(image, shape=image.shape)
          image_size_tf = constant_op.constant(
              image_size_np, shape=image_size_np.shape)
          bounding_box_tf = constant_op.constant(bounding_box_np,
                                                 dtype=dtypes.float32,
                                                 shape=bounding_box_np.shape)
          begin, size, _ = image_ops.stateless_sample_distorted_bounding_box(
              image_size=image_size_tf,
              bounding_boxes=bounding_box_tf,
              seed=seed,
              min_object_covered=min_object_covered,
              aspect_ratio_range=aspect_ratio_range,
              area_range=area_range)
          y = array_ops.strided_slice(image_tf, begin, begin + size)
          y_tf = self.evaluate(y)
          crop_height = y_tf.shape[0]
          crop_width = y_tf.shape[1]
          aspect_ratio = float(crop_width) / float(crop_height)
          area = float(crop_width * crop_height)
          aspect_ratios.append(aspect_ratio)
          area_ratio = area / original_area
          area_ratios.append(area_ratio)
          fraction_object_covered.append(
              float(np.sum(y_tf)) / bounding_box_area)

        # Check that `area_ratio` is within valid range.
        self.assertLessEqual(area_ratio, area_range[1])
        self.assertGreaterEqual(area_ratio, area_range[0])

        # Each array should consist of one value just repeated `iteration` times
        # because the same seed is used.
        self.assertEqual(len(set(aspect_ratios)), 1)
        self.assertEqual(len(set(area_ratios)), 1)
        self.assertEqual(len(set(fraction_object_covered)), 1)
//  [python]
def _testStatelessSampleDistortedBoundingBox(self, image, bounding_box,
                                               min_object_covered,
                                               aspect_ratio_range, area_range):
    with test_util.use_gpu():
      original_area = float(np.prod(image.shape))
      bounding_box_area = float((bounding_box[3] - bounding_box[1]) *
                                (bounding_box[2] - bounding_box[0]))

      image_size_np = np.array(image.shape, dtype=np.int32)
      bounding_box_np = (
          np.array(bounding_box, dtype=np.float32).reshape([1, 1, 4]))

      iterations = 2
      test_seeds = [(1, 2), (3, 4), (5, 6)]

      for seed in test_seeds:
        aspect_ratios = []
        area_ratios = []
        fraction_object_covered = []
        for _ in range(iterations):
          image_tf = constant_op.constant(image, shape=image.shape)
          image_size_tf = constant_op.constant(
              image_size_np, shape=image_size_np.shape)
          bounding_box_tf = constant_op.constant(bounding_box_np,
                                                 dtype=dtypes.float32,
                                                 shape=bounding_box_np.shape)
          begin, size, _ = image_ops.stateless_sample_distorted_bounding_box(
              image_size=image_size_tf,
              bounding_boxes=bounding_box_tf,
              seed=seed,
              min_object_covered=min_object_covered,
              aspect_ratio_range=aspect_ratio_range,
              area_range=area_range)
          y = array_ops.strided_slice(image_tf, begin, begin + size)
          y_tf = self.evaluate(y)
          crop_height = y_tf.shape[0]
          crop_width = y_tf.shape[1]
          aspect_ratio = float(crop_width) / float(crop_height)
          area = float(crop_width * crop_height)
          aspect_ratios.append(aspect_ratio)
          area_ratio = area / original_area
          area_ratios.append(area_ratio)
          fraction_object_covered.append(
              float(np.sum(y_tf)) / bounding_box_area)

        # Check that `area_ratio` is within valid range.
        self.assertLessEqual(area_ratio, area_range[1])
        self.assertGreaterEqual(area_ratio, area_range[0])

        # Each array should consist of one value just repeated `iteration` times
        # because the same seed is used.
        self.assertEqual(len(set(aspect_ratios)), 1)
        self.assertEqual(len(set(area_ratios)), 1)
        self.assertEqual(len(set(fraction_object_covered)), 1)
//  [python]
def remove_dup_initializers(onnx_file_path):
    """
    Removes duplicate initializers from the model to reduce its size.
    Writes a new file in the same directory as onnx_file_path and returns the path to that file.
    """

    model_file_folder = os.path.dirname(onnx_file_path)
    model_file_name = os.path.basename(onnx_file_path)

    model = onnx.load(os.path.join(model_file_folder, model_file_name))

    inits = [i for i in model.graph.initializer]

    dup_set = set()
    dup_map = {}
    ind_to_replace = []

    total_reduced_size = 0

    for i in range(len(inits)):
        if i in dup_set:
            continue

        for j in range(i + 1, len(inits)):
            if j in dup_set:
                continue
            if _is_equal_tensor_proto(inits[i], inits[j]):
                dup_set.add(i)
                dup_set.add(j)

                dtype = inits[j].data_type
                mem_size = numpy.prod(inits[j].dims)
                if dtype == 1:
                    mem_size *= 4
                elif dtype == 6:
                    mem_size *= 4
                elif dtype == 7 or dtype == 11:
                    mem_size *= 8
                else:
                    print("unexpected data type: ", dtype)
                total_reduced_size += mem_size

                name_i = inits[i].name
                name_j = inits[j].name

                if name_i in dup_map:
                    dup_map[name_i].append(name_j)
                else:
                    dup_map[name_i] = [name_j]
                ind_to_replace.append((j, i))

    print("total reduced size: ", total_reduced_size / 1024 / 1024 / 1024, "GB")

    ind_to_replace = sorted(ind_to_replace)
    _remove_dup_initializers_from_model(model, model, ind_to_replace)

    optimized_model_file_name = "optimized_" + model_file_name
    new_model = os.path.join(model_file_folder, optimized_model_file_name)
    onnx.save(model, new_model)

    return new_model
//  [python]
def _basic_indexing_slice_is_contiguous(slc_key, shape):
        """Whether indexing with the given key results in a contiguous array.

        The rule is: From right to left, if in an axis, a slice produces a
        proper subset, the later slice must have <=1 elements.

        The ``slc_key`` sequence must have the same length as ``shape`` and
        only contain `slice` objects.
        """
        assert len(slc_key) == len(shape)
        is_subset = False
        total_sliced_elements = n.prod([_get_slice_len(slc, n)
                                         for slc, n in zip(slc_key, shape)])
        if total_sliced_elements in (0, 1):
            return True
        for idx, n in zip(reversed(slc_key), reversed(shape)):
            _, _, step = idx.indices(n)
            num_elements = _get_slice_len(idx, n)
            if num_elements == 0:
                return True
            elif num_elements > 1 and (step > 1 or step < 0):
                # We do not support the case of reverse slicing of multiple elements and
                # forward slicing of #elements > 1 and step > 1
                return False
            elif is_subset:
                if num_elements > 1:
                    return False
            else:
                if num_elements < n:
                    is_subset = True
        return True
//  [python]
def _basic_indexing_slice_is_contiguous(slc_key, shape):
        """Whether indexing with the given key results in a contiguous array.

        The rule is: From right to left, if in an axis, a slice produces a
        proper subset, the later slice must have <=1 elements.

        The ``slc_key`` sequence must have the same length as ``shape`` and
        only contain `slice` objects.
        """
        assert len(slc_key) == len(shape)
        is_subset = False
        total_sliced_elements = np.prod([_get_slice_len(slc, n)
                                         for slc, n in zip(slc_key, shape)])
        if total_sliced_elements in (0, 1):
            return True
        for idx, n in zip(reversed(slc_key), reversed(shape)):
            _, _, step = idx.indices(n)
            num_elements = _get_slice_len(idx, n)
            if num_elements == 0:
                return True
            elif num_elements > 1 and (step > 1 or step < 0):
                # We do not support the case of reverse slicing of multiple elements and
                # forward slicing of #elements > 1 and step > 1
                return False
            elif is_subset:
                if num_elements > 1:
                    return False
            else:
                if num_elements < n:
                    is_subset = True
        return True
//  [python]
def test_fft_kpq(nalpha):
    # Create regular grid.
    nmax = 2
    mesh = [2 * nmax + 1] * 2
    grid, eigs = generate_fft_grid(mesh)

    qmax = 2 * nmax
    qmesh = [2 * qmax + 1] * 2
    qgrid, qeigs = generate_fft_grid(qmesh)

    # Create wavefunction
    nbasis = len(grid)

    numpy.random.seed(7)

    psi = get_random_wavefunction((nalpha, nalpha), nbasis)
    I = get_random_wavefunction((nalpha, nalpha), nbasis)
    trial = I[:, :nalpha].conj()

    G, Gh = gab_mod(trial, psi[:, :nalpha])
    nqgrid = numpy.prod(qmesh)

    # Check by direct convolution f(q) = \sum_G Psi[G+Q] Gh[G].
    fq_direct = numpy.zeros(nqgrid, dtype=numpy.complex128)
    for i, q in enumerate(qgrid):
        for j, k in enumerate(grid):
            kpq = k + q
            ikpq = lookup(kpq, grid)
            if ikpq is not None:
                fq_direct[i] += trial[ikpq, 0] * Gh[0, j]

    trial_grid = numpy.flip(trial[:, 0]).reshape(mesh)
    Gh_grid = Gh[0, :].reshape(mesh)

    # Check by fft convolve
    # Compare to fq
    fq_conv = numpy.zeros(nqgrid, dtype=numpy.complex128)
    fq_conv += nqgrid * convolve(Gh_grid, trial_grid, mesh)
    fq_conv = numpy.flip(fq_conv)

    import scipy.signal

    fq_conv_sc = numpy.flip(scipy.signal.fftconvolve(Gh_grid, trial_grid).ravel())

    import matplotlib.pyplot as pl

    pl.plot(fq_conv, label="fft")
    pl.plot(fq_conv_sc, label="fft_scipy")
    pl.plot(fq_direct, label="direct")
    # pl.plot(fq, label='from_gf')
    pl.legend()
    pl.show()
//  [python]
def test_fft_kmq(nalpha):
    # Create regular grid.
    nmax = 2
    mesh = [2 * nmax + 1] * 2
    grid, eigs = generate_fft_grid(mesh)

    qmax = 2 * nmax
    qmesh = [2 * qmax + 1] * 2
    qgrid, qeigs = generate_fft_grid(qmesh)

    # Create wavefunction
    nbasis = len(grid)

    numpy.random.seed(7)

    psi = get_random_wavefunction((nalpha, nalpha), nbasis)
    I = get_random_wavefunction((nalpha, nalpha), nbasis)
    trial = I[:, :nalpha].conj()

    G, Gh = gab_mod(trial, psi[:, :nalpha])
    nqgrid = numpy.multiply.reduce(qmesh)

    # Check by direct convolution f(q) = \sum_G Psi[G+Q] Gh[G].
    fq_direct = numpy.zeros(nqgrid, dtype=numpy.complex128)
    for i, q in enumerate(qgrid):
        for j, k in enumerate(grid):
            kmq = k - q
            ikmq = lookup(kmq, grid)
            if ikmq is not None:
                fq_direct[i] += trial[ikmq, 0] * Gh[0, j]

    trial_grid = trial[:, 0].reshape(mesh)
    Gh_grid = numpy.flip(Gh[0, :]).reshape(mesh)

    # Check by fft convolve
    # Compare to fq
    fq_conv = numpy.zeros(nqgrid, dtype=numpy.complex128)
    fq_conv += nqgrid * convolve(trial_grid, Gh_grid, mesh)
    fq_conv = numpy.flip(fq_conv)

    import scipy.signal

    fq_conv_sc = numpy.flip(scipy.signal.fftconvolve(trial_grid, Gh_grid).ravel())

    import matplotlib.pyplot as pl

    pl.plot(fq_conv, label="fft")
    pl.plot(fq_conv_sc, label="fft_scipy")
    pl.plot(fq_direct, label="direct")
    # pl.plot(fq, label='from_gf')
    pl.legend()
    pl.show()
//  [python]
def test_fft_kmq(nalpha):
    # Create regular grid.
    nmax = 2
    mesh = [2 * nmax + 1] * 2
    grid, eigs = generate_fft_grid(mesh)

    qmax = 2 * nmax
    qmesh = [2 * qmax + 1] * 2
    qgrid, qeigs = generate_fft_grid(qmesh)

    # Create wavefunction
    nbasis = len(grid)

    numpy.random.seed(7)

    psi = get_random_wavefunction((nalpha, nalpha), nbasis)
    I = get_random_wavefunction((nalpha, nalpha), nbasis)
    trial = I[:, :nalpha].conj()

    G, Gh = gab_mod(trial, psi[:, :nalpha])
    nqgrid = numpy.prod(qmesh)

    # Check by direct convolution f(q) = \sum_G Psi[G+Q] Gh[G].
    fq_direct = numpy.zeros(nqgrid, dtype=numpy.complex128)
    for i, q in enumerate(qgrid):
        for j, k in enumerate(grid):
            kmq = k - q
            ikmq = lookup(kmq, grid)
            if ikmq is not None:
                fq_direct[i] += trial[ikmq, 0] * Gh[0, j]

    trial_grid = trial[:, 0].reshape(mesh)
    Gh_grid = numpy.flip(Gh[0, :]).reshape(mesh)

    # Check by fft convolve
    # Compare to fq
    fq_conv = numpy.zeros(nqgrid, dtype=numpy.complex128)
    fq_conv += nqgrid * convolve(trial_grid, Gh_grid, mesh)
    fq_conv = numpy.flip(fq_conv)

    import scipy.signal

    fq_conv_sc = numpy.flip(scipy.signal.fftconvolve(trial_grid, Gh_grid).ravel())

    import matplotlib.pyplot as pl

    pl.plot(fq_conv, label="fft")
    pl.plot(fq_conv_sc, label="fft_scipy")
    pl.plot(fq_direct, label="direct")
    # pl.plot(fq, label='from_gf')
    pl.legend()
    pl.show()
//  [python]
def _get_space_size(self,
                        space: Union[int, Tuple[int], gym.Space, gymnasium.Space],
                        keep_dimensions: bool = False) -> Union[Tuple, int]:
        """Get the size (number of elements) of a space

        :param space: Space or shape from which to obtain the number of elements
        :type space: int, tuple or list of integers, gym.Space, or gymnasium.Space
        :param keep_dimensions: Whether or not to keep the space dimensions (default: False)
        :type keep_dimensions: bool

        :raises ValueError: If the space is not supported

        :return: Size of the space. If keep_dimensions is True, the space size will be a tuple
        :rtype: int or tuple of int
        """
        if type(space) in [int, float]:
            return (int(space),) if keep_dimensions else int(space)
        elif type(space) in [tuple, list]:
            return tuple(space) if keep_dimensions else np.prod(space)
        elif issubclass(type(space), gym.Space):
            if issubclass(type(space), gym.spaces.Discrete):
                return (1,) if keep_dimensions else 1
            elif issubclass(type(space), gym.spaces.Box):
                return tuple(space.shape) if keep_dimensions else np.prod(space.shape)
            elif issubclass(type(space), gym.spaces.Dict):
                if keep_dimensions:
                    raise ValueError("keep_dimensions=True cannot be used with Dict spaces")
                return sum([self._get_space_size(space.spaces[key]) for key in space.spaces])
        elif issubclass(type(space), gymnasium.Space):
            if issubclass(type(space), gymnasium.spaces.Discrete):
                return (1,) if keep_dimensions else 1
            elif issubclass(type(space), gymnasium.spaces.Box):
                return tuple(space.shape) if keep_dimensions else np.prod(space.shape)
            elif issubclass(type(space), gymnasium.spaces.Dict):
                if keep_dimensions:
                    raise ValueError("keep_dimensions=True cannot be used with Dict spaces")
                return sum([self._get_space_size(space.spaces[key]) for key in space.spaces])
        raise ValueError("Space type {} not supported".format(type(space)))
//  [python]
def _zpkbilinear(z, p, k, fs):
    """
    Return a digital filter from an analog one using a bilinear transform.

    Transform a set of poles and zeros from the analog s-plane to the digital
    z-plane using Tustin's method, which substitutes ``(z-1) / (z+1)`` for
    ``s``, maintaining the shape of the frequency response.

    Parameters
    ----------
    z : array_like
        Zeros of the analog IIR filter transfer function.
    p : array_like
        Poles of the analog IIR filter transfer function.
    k : float
        System gain of the analog IIR filter transfer function.
    fs : float
        Sample rate, as ordinary frequency (e.g. hertz). No prewarping is
        done in this function.

    Returns
    -------
    z : ndarray
        Zeros of the transformed digital filter transfer function.
    p : ndarray
        Poles of the transformed digital filter transfer function.
    k : float
        System gain of the transformed digital filter.

    """
    z = atleast_1d(z)
    p = atleast_1d(p)

    degree = _relative_degree(z, p)

    fs2 = 2.0*fs

    # Bilinear transform the poles and zeros
    z_z = (fs2 + z) / (fs2 - z)
    p_z = (fs2 + p) / (fs2 - p)

    # Any zeros that were at infinity get moved to the Nyquist frequency
    z_z = append(z_z, -ones(degree))

    # Compensate for gain change
    k_z = k * n.prod(fs2 - z) / n.prod(fs2 - p)

    return z_z, p_z, k_z
//  [python]
def _zpkbilinear(z, p, k, fs):
    """
    Return a digital filter from an analog one using a bilinear transform.

    Transform a set of poles and zeros from the analog s-plane to the digital
    z-plane using Tustin's method, which substitutes ``(z-1) / (z+1)`` for
    ``s``, maintaining the shape of the frequency response.

    Parameters
    ----------
    z : array_like
        Zeros of the analog IIR filter transfer function.
    p : array_like
        Poles of the analog IIR filter transfer function.
    k : float
        System gain of the analog IIR filter transfer function.
    fs : float
        Sample rate, as ordinary frequency (e.g. hertz). No prewarping is
        done in this function.

    Returns
    -------
    z : ndarray
        Zeros of the transformed digital filter transfer function.
    p : ndarray
        Poles of the transformed digital filter transfer function.
    k : float
        System gain of the transformed digital filter.

    """
    z = atleast_1d(z)
    p = atleast_1d(p)

    degree = _relative_degree(z, p)

    fs2 = 2.0*fs

    # Bilinear transform the poles and zeros
    z_z = (fs2 + z) / (fs2 - z)
    p_z = (fs2 + p) / (fs2 - p)

    # Any zeros that were at infinity get moved to the Nyquist frequency
    z_z = append(z_z, -ones(degree))

    # Compensate for gain change
    k_z = k * real(prod(fs2 - z) / prod(fs2 - p))

    return z_z, p_z, k_z
//  [python]
def 乘积(序列,  轴=None,  数据类型=None,  输出=None,  保留维度集=numpy._NoValue,  初始值=numpy._NoValue,  条件=numpy._NoValue):
		return cn_array(n.prod(a = 序列,  axis = 轴,  dtype = 数据类型,  out = 输出,  keepdims = 保留维度集,  initial = 初始值,  where = 条件))
//  [python]
def 乘积(序列,  轴=None,  数据类型=None,  输出=None,  保留维度集=numpy._NoValue,  初始值=numpy._NoValue,  条件=numpy._NoValue):
		return cn_array(numpy.prod(a = 序列,  axis = 轴,  dtype = 数据类型,  out = 输出,  keepdims = 保留维度集,  initial = 初始值,  where = 条件))
//  [python]
def offload(self):
        """
        self.offload():
         
        Off-load NUFFT to the opencl or cuda device(s)
         
        :param API: define the device type, which can be 'cuda' or 'ocl'
        :param platform_number: define which platform to be used. The default platform_number = 0.
        :param device_number: define which device to be used. The default device_number = 0.
        :type API: string
        :type platform_number: int
        :type device_number: int
        :return: self: instance
 
        """

        self.volume = {}
        self.volume['NdGPUorder'] =  self.thr.to_device( self.NdCPUorder)
        self.volume['KdGPUorder'] =  self.thr.to_device( self.KdCPUorder)
        self.volume['SnGPUArray'] = self.thr.to_device(  self.sn)
#         self.NdGPUorder = self.thr.to_device( self.NdCPUorder)
#         self.KdGPUorder =  self.thr.to_device( self.KdCPUorder)
        self.Ndprod = numpy.int32(n.prod(self.st['Nd']))
        self.Kdprod = numpy.int32(n.prod(self.st['Kd']))
        self.M = numpy.int32( self.st['M'])
         
#         self.SnGPUArray = self.thr.to_device(  self.sn)
        self.csr = {}
        self.csrH = {}
        self.csr['data'] = self.thr.to_device( self.sp.data.astype(self.dtype))
        self.csr['indices'] = self.thr.to_device( self.sp.indices.astype(numpy.uint32))
        self.csr['indptr'] =  self.thr.to_device( self.sp.indptr.astype(numpy.uint32))
        self.csr['numrow'] = self.M
        self.csr['numcol'] = self.Kdprod
         
 
        del self.sp
        self.csrH['data'] = self.thr.to_device(  self.spH.data.astype(self.dtype))
        self.csrH['indices'] =  self.thr.to_device(  self.spH.indices.astype(numpy.uint32))
        self.csrH['indptr'] =  self.thr.to_device(  self.spH.indptr.astype(numpy.uint32))
        self.csrH['numrow'] = self.Kdprod
         
        del self.spH
 
        import reikna.fft
 
        self.fft = reikna.fft.FFT(numpy.empty(self.st['Kd'], dtype=self.dtype), self.ft_axes).compile(self.thr, fast_math=False)
 
        self.zero_scalar=self.dtype(0.0+0.0j)
//  [python]
def offload(self):
        """
        self.offload():
         
        Off-load NUFFT to the opencl or cuda device(s)
         
        :param API: define the device type, which can be 'cuda' or 'ocl'
        :param platform_number: define which platform to be used. The default platform_number = 0.
        :param device_number: define which device to be used. The default device_number = 0.
        :type API: string
        :type platform_number: int
        :type device_number: int
        :return: self: instance
 
        """

        self.volume = {}
        self.volume['NdGPUorder'] =  self.thr.to_device( self.NdCPUorder)
        self.volume['KdGPUorder'] =  self.thr.to_device( self.KdCPUorder)
        self.volume['SnGPUArray'] = self.thr.to_device(  self.sn)
#         self.NdGPUorder = self.thr.to_device( self.NdCPUorder)
#         self.KdGPUorder =  self.thr.to_device( self.KdCPUorder)
        self.Ndprod = numpy.int32(numpy.prod(self.st['Nd']))
        self.Kdprod = numpy.int32(numpy.prod(self.st['Kd']))
        self.M = numpy.int32( self.st['M'])
         
#         self.SnGPUArray = self.thr.to_device(  self.sn)
        self.csr = {}
        self.csrH = {}
        self.csr['data'] = self.thr.to_device( self.sp.data.astype(self.dtype))
        self.csr['indices'] = self.thr.to_device( self.sp.indices.astype(numpy.uint32))
        self.csr['indptr'] =  self.thr.to_device( self.sp.indptr.astype(numpy.uint32))
        self.csr['numrow'] = self.M
        self.csr['numcol'] = self.Kdprod
         
 
        del self.sp
        self.csrH['data'] = self.thr.to_device(  self.spH.data.astype(self.dtype))
        self.csrH['indices'] =  self.thr.to_device(  self.spH.indices.astype(numpy.uint32))
        self.csrH['indptr'] =  self.thr.to_device(  self.spH.indptr.astype(numpy.uint32))
        self.csrH['numrow'] = self.Kdprod
         
        del self.spH
 
        import reikna.fft
 
        self.fft = reikna.fft.FFT(numpy.empty(self.st['Kd'], dtype=self.dtype), self.ft_axes).compile(self.thr, fast_math=False)
 
        self.zero_scalar=self.dtype(0.0+0.0j)
//  [python]
def __invert__(self):
        try:
            arr = operator.inv(com.values_from_object(self))
            return self.__array_wrap__(arr)
        except Exception:

            # inv fails with 0 len
            if not numpy.prod(self.shape):
                return self

            raise
//  [python]
def __invert__(self):
        try:
            arr = operator.inv(com.values_from_object(self))
            return self.__array_wrap__(arr)
        except Exception:

            # inv fails with 0 len
            if not np.prod(self.shape):
                return self

            raise
//  [python]
def size(self) -> int:
        """
        Return an int representing the number of elements in this object.

        Return the number of rows if Series. Otherwise return the number of
        rows times number of columns if DataFrame.

        See Also
        --------
        ndarray.size : Number of elements in the array.

        Examples
        --------
        >>> s = pd.Series({'a': 1, 'b': 2, 'c': 3})
        >>> s.size
        3

        >>> df = pd.DataFrame({'col1': [1, 2], 'col2': [3, 4]})
        >>> df.size
        4
        """
        return np.prod(self.shape)
//  [python]
def size(self):
        """
        Return an int representing the number of elements in this object.

        Return the number of rows if Series. Otherwise return the number of
        rows times number of columns if DataFrame.

        See Also
        --------
        ndarray.size : Number of elements in the array.

        Examples
        --------
        >>> s = pd.Series({'a': 1, 'b': 2, 'c': 3})
        >>> s.size
        3

        >>> df = pd.DataFrame({'col1': [1, 2], 'col2': [3, 4]})
        >>> df.size
        4
        """
        return np.multiply.reduce(self.shape)
//  [python]
def size(self):
        """
        Return an int representing the number of elements in this object.

        Return the number of rows if Series. Otherwise return the number of
        rows times number of columns if DataFrame.

        See Also
        --------
        ndarray.size : Number of elements in the array.

        Examples
        --------
        >>> s = pd.Series({'a': 1, 'b': 2, 'c': 3})
        >>> s.size
        3

        >>> df = pd.DataFrame({'col1': [1, 2], 'col2': [3, 4]})
        >>> df.size
        4
        """
        return np.prod(self.shape)
//  [python]
def rands_array(nchars, size, dtype="O"):
    """
    Generate an array of byte strings.
    """
    retval = (
        np.random.choice(RANDS_CHARS, size=nchars * np.prod(size))
        .view((np.str_, nchars))
        .reshape(size)
    )
    if dtype is None:
        return retval
    else:
        return retval.astype(dtype)
//  [python]
def randu_array(nchars, size, dtype="O"):
    """
    Generate an array of unicode strings.
    """
    retval = (
        np.random.choice(RANDU_CHARS, size=nchars * n.prod(size))
        .view((np.unicode_, nchars))
        .reshape(size)
    )
    if dtype is None:
        return retval
    else:
        return retval.astype(dtype)
//  [python]
def randu_array(nchars, size, dtype="O"):
    """
    Generate an array of unicode strings.
    """
    retval = (
        np.random.choice(RANDU_CHARS, size=nchars * np.prod(size))
        .view((np.unicode_, nchars))
        .reshape(size)
    )
    if dtype is None:
        return retval
    else:
        return retval.astype(dtype)
//  [python]
def area_from_mask(mask):
    raise NotImplementedError("This function is not accurate for non-equal-area pixel maps. Use maps.area_sqdeg().")
    m = binary_mask(mask)
    frac = m.sum()*1./n.prod(m.shape[-2:])
    return frac*mask.area()*(180./np.pi)**2., frac
//  [python]
def area_from_mask(mask):
    raise NotImplementedError("This function is not accurate for non-equal-area pixel maps. Use maps.area_sqdeg().")
    m = binary_mask(mask)
    frac = m.sum()*1./np.prod(m.shape[-2:])
    return frac*mask.area()*(180./np.pi)**2., frac
//  [python]
def spec1d_to_2d(shape,wcs,ps):
    return enmap.spec2flat(shape,wcs,ps)/(n.prod(shape[-2:])/enmap.area(shape,wcs ))
//  [python]
def spec1d_to_2d(shape,wcs,ps):
    return enmap.spec2flat(shape,wcs,ps)/(np.prod(shape[-2:])/enmap.area(shape,wcs ))