# Copyright 2018 Google LLC
#
# 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
#
#     https://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.

import collections
from contextlib import contextmanager
import enum
from functools import partial
import itertools
import warnings

from absl.testing import absltest
from absl.testing import parameterized

import numpy as np

import jax
from jax import dtypes
from jax import numpy as jnp
from jax import ops
from jax import test_util as jtu
from jax._src import util

from jax.config import config
config.parse_flags_with_absl()

# We disable the whitespace continuation check in this file because otherwise it
# makes the test name formatting unwieldy.
# pylint: disable=bad-continuation

ARRAY_MSG = r"Using a non-tuple sequence for multidimensional indexing is not allowed.*arr\[array\(seq\)\]"
TUPLE_MSG = r"Using a non-tuple sequence for multidimensional indexing is not allowed.*arr\[tuple\(seq\)\]"

float_dtypes = jtu.dtypes.floating
default_dtypes = float_dtypes + jtu.dtypes.integer
all_dtypes = default_dtypes + jtu.dtypes.boolean

IndexSpec = collections.namedtuple("IndexTest", ["shape", "indexer"])

def check_grads(f, args, order, atol=None, rtol=None, eps=None):
  # TODO(mattjj,dougalm): add higher-order check
  default_tol = 1e-6 if config.x64_enabled else 1e-2
  atol = atol or default_tol
  rtol = rtol or default_tol
  eps = eps or default_tol
  jtu.check_jvp(f, partial(jax.jvp, f), args, atol, rtol, eps)
  jtu.check_vjp(f, partial(jax.vjp, f), args, atol, rtol, eps)

STATIC_INDEXING_TESTS = [
    ("OneIntIndex", [
        IndexSpec(shape=(3,), indexer=1),
        IndexSpec(shape=(3, 3), indexer=0),
        IndexSpec(shape=(3, 4, 5), indexer=2),
        IndexSpec(shape=(3,), indexer=-1),
        IndexSpec(shape=(3,), indexer=-2),
    ]),
    ("TwoIntIndices", [
        IndexSpec(shape=(3, 3), indexer=(2, 1)),
        IndexSpec(shape=(3, 4, 5), indexer=(1, 2)),
        IndexSpec(shape=(3, 4, 5), indexer=(-1, 2)),
    ]),
    ("ThreeIntIndices", [IndexSpec((3, 4, 5), indexer=(1, 2, 3))]),
    ("OneSliceIndex", [
        IndexSpec(shape=(10,), indexer=slice(1, 3)),
        IndexSpec(shape=(10,), indexer=slice(1, -1)),
        IndexSpec(shape=(10,), indexer=slice(None, -1)),
        IndexSpec(shape=(10,), indexer=slice(None, None, None)),
        IndexSpec(shape=(10, 8), indexer=slice(1, 3)),
        IndexSpec(shape=(10, 8), indexer=slice(1, None)),
        IndexSpec(shape=(10, 8), indexer=slice(None, 3)),
        IndexSpec(shape=(10, 8), indexer=slice(-3, None)),
    ]),
    ("OneSliceIndexNegativeStride", [
        IndexSpec(shape=(10,), indexer=slice(3, 1, -1)),
        IndexSpec(shape=(10,), indexer=slice(1, 8, -1)),  # empty result
        IndexSpec(shape=(10,), indexer=slice(None, 1, -2)),
        IndexSpec(shape=(10,), indexer=slice(None, None, -1)),
        IndexSpec(shape=(10, 8), indexer=slice(3, 1, -1)),
        IndexSpec(shape=(10, 8), indexer=slice(0, 8, -1)),  # empty result
        IndexSpec(shape=(10, 8), indexer=slice(None, None, -1)),
    ]),
    ("OneSliceIndexNonUnitStride", [
        IndexSpec(shape=(10,), indexer=slice(0, 8, 2)),
        IndexSpec(shape=(10,), indexer=slice(0, 8, 3)),
        IndexSpec(shape=(10,), indexer=slice(1, 3, 2)),
        IndexSpec(shape=(10,), indexer=slice(1, None, 2)),
        IndexSpec(shape=(10,), indexer=slice(None, 1, -2)),
        IndexSpec(shape=(10, 8), indexer=slice(1, 8, 3)),
        IndexSpec(shape=(10, 8), indexer=slice(None, None, 2)),
        IndexSpec(shape=(10, 8), indexer=slice(None, 1, -2)),
        IndexSpec(shape=(10, 8), indexer=slice(None, None, -2)),
    ]),
    ("TwoSliceIndices", [
        IndexSpec(shape=(10, 8), indexer=(slice(1, 3), slice(0, 2))),
        IndexSpec(shape=(10, 8), indexer=(slice(1, None), slice(None, 2))),
        IndexSpec(
            shape=(10, 8), indexer=(slice(None, None, -1), slice(None, 2))),
        IndexSpec(shape=(10, 8, 3), indexer=(slice(1, 3), slice(0, 2))),
        IndexSpec(shape=(10, 8, 3), indexer=(slice(1, 3), slice(0, None))),
        IndexSpec(shape=(10, 8, 3), indexer=(slice(1, None), slice(0, 2))),
    ]),
    ("OneColonIndex", [
        IndexSpec(shape=(3,), indexer=slice(None)),
        IndexSpec(shape=(3, 4), indexer=slice(None)),
    ]),
    ("MultipleColonIndices", [
        IndexSpec(shape=(3, 4), indexer=(slice(None), slice(None))),
        IndexSpec(shape=(3, 4, 5), indexer=(slice(None), slice(None))),
    ]),
    ("MixedSliceIndices", [
        IndexSpec(shape=(10, 4), indexer=(slice(None), slice(0, 2))),
        IndexSpec(shape=(10, 4), indexer=(1, slice(None))),
    ]),
    ("EllipsisIndex", [
        IndexSpec(shape=(3,), indexer=Ellipsis),
        IndexSpec(shape=(3, 4), indexer=Ellipsis),
        IndexSpec(shape=(3, 4, 5), indexer=(0, Ellipsis)),
        IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis, 2, 3)),
    ]),
    ("NoneIndex", [
        IndexSpec(shape=(), indexer=None),
        IndexSpec(shape=(), indexer=(None, None)),
        IndexSpec(shape=(), indexer=(Ellipsis, None)),
        IndexSpec(shape=(3,), indexer=None),
        IndexSpec(shape=(3, 4), indexer=None),
        IndexSpec(shape=(3, 4), indexer=(Ellipsis, None)),
        IndexSpec(shape=(3, 4), indexer=(0, None, Ellipsis)),
        IndexSpec(shape=(3, 4, 5), indexer=(1, None, Ellipsis)),
    ]),
    ("EmptyIndex", [
        IndexSpec(shape=(), indexer=()),
        IndexSpec(shape=(3,), indexer=()),
        IndexSpec(shape=(3, 4), indexer=()),
    ]),
    ("TupleOfIntAndSliceAndIntArray", [
        IndexSpec(shape=(3, 2, 3), indexer=(0, slice(None), np.arange(3))),
        IndexSpec(shape=(3, 2, 3), indexer=(np.int32(1), slice(None), np.arange(3))),
        IndexSpec(shape=(3, 2, 3), indexer=(np.array(2), slice(None), np.arange(3))),
    ]),
]

ADVANCED_INDEXING_TESTS = [
    ("One1DIntArrayIndex",
     [IndexSpec(shape=(3,), indexer=np.array([0, 1])),
     IndexSpec(shape=(3, 3), indexer=np.array([1, 2, 1])),
     IndexSpec(shape=(3, 4, 5), indexer=np.array([0, 2, 0, 1])),
     IndexSpec(shape=(3,), indexer=np.array([-1, 1])),
     IndexSpec(shape=(3,), indexer=np.array([-2, -1])),
     IndexSpec(shape=(0,), indexer=np.array([], dtype=np.int32)),
     ]),
    ("One2DIntArrayIndex",
     [IndexSpec(shape=(3,), indexer=np.array([[0, 0]])),
     IndexSpec(shape=(3, 3), indexer=np.array([[1, 2, 1],
                                                [0, 1, -1]])),
     IndexSpec(shape=(3, 4, 5), indexer=np.array([[0, 2, 0, 1],
                                                   [-1, -2, 1, 0]])),
     ]),
    ("Two1DIntArrayIndicesNoBroadcasting",
     [IndexSpec(shape=(3, 3), indexer=(np.array([0, 1]),
                                       np.array([1, 2]))),
     IndexSpec(shape=(3, 4, 5), indexer=(np.array([0, 2, 0, 1]),
                                         np.array([-1, 0, -1, 2]))),
     ]),
    ("Two1DIntArrayIndicesWithBroadcasting",
     [IndexSpec(shape=(3, 3), indexer=(np.array([[0, 1]]),
                                       np.array([1, 2]))),
     IndexSpec(shape=(3, 4, 5), indexer=(np.array([[0, 2, 0, 1]]),
                                         np.array([-1, 0, -1, 2]))),
     ]),
    ("ArrayOfInts",
     [IndexSpec(shape=(3,), indexer=np.array([0, 1, 0])),
     IndexSpec(shape=(3, 4, 5), indexer=np.array([0, -1])),
     ]),
    ("TupleOfListsOfPythonInts",
     [IndexSpec(shape=(3, 4, 5), indexer=([0, 1],)),
     IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]], [[2, 3, 0, 3]])),
     ]),
    ("TupleOfPythonIntsAndIntArrays",
     [IndexSpec(shape=(3, 4, 5), indexer=(0, np.array([0, 1]))),
     IndexSpec(shape=(3, 4, 5), indexer=(0, 1,
                                         np.array([[2, 3, 0, 3]]))),
     ]),
    ("TupleOfListsOfPythonIntsAndIntArrays",
     [IndexSpec(shape=(3, 4, 5), indexer=([0, 1], np.array([0]))),
     IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]],
                                         np.array([[2, 3, 0, 3]]))),
     ]),
]

ADVANCED_INDEXING_TESTS_NO_REPEATS = [
    ("One1DIntArrayIndex",
     [IndexSpec(shape=(3,), indexer=np.array([0, 1])),
      IndexSpec(shape=(3, 3), indexer=np.array([1, 2, 0])),
      IndexSpec(shape=(3, 4, 5), indexer=np.array([0, 2, 1])),
      IndexSpec(shape=(3,), indexer=np.array([-1, 1])),
      IndexSpec(shape=(3,), indexer=np.array([-2, -1])),
      IndexSpec(shape=(0,), indexer=np.array([], dtype=np.int32)),
     ]),
    ("One2DIntArrayIndex",
     [IndexSpec(shape=(3,), indexer=np.array([[0, 1]])),
      IndexSpec(shape=(6, 6), indexer=np.array([[1, 2, 0],
                                                 [3, 4, -1]])),
     ]),
    ("Two1DIntArrayIndicesNoBroadcasting",
     [IndexSpec(shape=(3, 3), indexer=(np.array([0, 1]),
                                       np.array([1, 2]))),
      IndexSpec(shape=(4, 5, 6), indexer=(np.array([0, 2, 1, 3]),
                                          np.array([-1, 0, -2, 1]))),
     ]),
    ("Two1DIntArrayIndicesWithBroadcasting",
     [IndexSpec(shape=(3, 3), indexer=(np.array([[0, 1]]),
                                       np.array([1, 2]))),
      IndexSpec(shape=(4, 5, 6), indexer=(np.array([[0, 2, -1, 1]]),
                                          np.array([-1, 0, -2, 2]))),
     ]),
    ("ArrayOfInts",
     [IndexSpec(shape=(3,), indexer=np.array([0, 2, 1])),
      IndexSpec(shape=(3, 4, 5), indexer=np.array([0, -1])),
     ]),
    ("TupleOfListsOfPythonInts",
     [IndexSpec(shape=(3, 4, 5), indexer=([0, 1],)),
      IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]], [[2, 3, 0]])),
     ]),
    ("TupleOfPythonIntsAndIntArrays",
     [IndexSpec(shape=(3, 4, 5), indexer=(0, np.array([0, 1]))),
      IndexSpec(shape=(3, 4, 5), indexer=(0, 1,
                                          np.array([[2, 3, 0]]))),
     ]),
    ("TupleOfListsOfPythonIntsAndIntArrays",
     [IndexSpec(shape=(3, 4, 5), indexer=([0, 1], np.array([0]))),
      IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]],
                                          np.array([[2, 3, 0]]))),
     ]),
]

ADVANCED_INDEXING_TESTS_NO_REPEATS_SORTED = [
    ("One1DIntArrayIndex",
     [IndexSpec(shape=(3,), indexer=np.array([0, 1])),
      IndexSpec(shape=(3, 3), indexer=np.array([0, 1, 2])),
      IndexSpec(shape=(3, 4, 5), indexer=np.array([0, 1, 2])),
      IndexSpec(shape=(3,), indexer=np.array([-1, 1])),
      IndexSpec(shape=(3,), indexer=np.array([-2, -1])),
      IndexSpec(shape=(0,), indexer=np.array([], dtype=np.int32)),
     ]),
    ("One2DIntArrayIndex",
     [IndexSpec(shape=(3,), indexer=np.array([[0, 1]])),
      IndexSpec(shape=(6, 6), indexer=np.array([[-1, 0, 1],
                                                 [ 2, 3, 4]])),
     ]),
    ("Two1DIntArrayIndicesNoBroadcasting",
     [IndexSpec(shape=(3, 3), indexer=(np.array([0, 1]),
                                       np.array([1, 2]))),
      IndexSpec(shape=(4, 5, 6), indexer=(np.array([0, 1, 2, 3]),
                                          np.array([-2, -1, 0, 1]))),
     ]),
    ("Two1DIntArrayIndicesWithBroadcasting",
     [IndexSpec(shape=(3, 3), indexer=(np.array([[0, 1]]),
                                       np.array([1, 2]))),
      IndexSpec(shape=(4, 5, 6), indexer=(np.array([[-1, 0, 1, 2]]),
                                          np.array([-2, -1, 0, 2]))),
     ]),
    ("TupleOfListsOfPythonInts",
     [IndexSpec(shape=(3, 4, 5), indexer=([0, 1],)),
      IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]], [[0, 2, 3]])),
     ]),
    ("TupleOfPythonIntsAndIntArrays",
     [IndexSpec(shape=(3, 4, 5), indexer=(0, np.array([0, 1]))),
      IndexSpec(shape=(3, 4, 5), indexer=(0, 1,
                                          np.array([[0, 2, 3]]))),
     ]),
    ("TupleOfListsOfPythonIntsAndIntArrays",
     [IndexSpec(shape=(3, 4, 5), indexer=([0, 1], np.array([0]))),
      IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]],
                                          np.array([[0, 2, 3]]))),
     ]),
]

MIXED_ADVANCED_INDEXING_TESTS_NO_REPEATS = [
    ("SlicesAndOneIntArrayIndex",
     [IndexSpec(shape=(2, 3), indexer=(np.array([0, 1]), slice(1, 2))),
     IndexSpec(shape=(2, 3), indexer=(slice(0, 2),
                                      np.array([0, 2]))),
     IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
                                         np.array([0, 2]),
                                         slice(None))),
     IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
                                         np.array([[0, 2], [1, 3]]),
                                         slice(None))),
     ]),
    ("SlicesAndTwoIntArrayIndices",
     [IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
                                          np.array([0, 2]),
                                          np.array([-1, 2]))),
     IndexSpec(shape=(3, 4, 5), indexer=(np.array([0, 2]),
                                         Ellipsis,
                                         np.array([-1, 2]))),
     IndexSpec(shape=(3, 4, 5), indexer=(np.array([0, 2]),
                                         np.array([-1, 2]),
                                         Ellipsis)),
     IndexSpec(shape=(3, 4, 5), indexer=(np.array([0, 2]),
                                         np.array([-1, 2]),
                                         slice(1, 3))),
     IndexSpec(shape=(3, 4, 5), indexer=(np.array([0, 2]),
                                         slice(1, 3),
                                         np.array([-1, 2]))),
     IndexSpec(shape=(3, 4, 5), indexer=(np.array([0, 2, -2]),
                                         slice(None, None, 2),
                                         np.array([-1, 2, 1]))),
     ]),
    ("NonesAndIntArrayIndices",
     [IndexSpec(shape=(3, 4, 5), indexer=(np.array([0, 2]),
                                          None,
                                          np.array([-1, 2]))),
     IndexSpec(shape=(3, 4, 5), indexer=(np.array([0, 2]),
                                         None,
                                         None,
                                         np.array([-1, 2]))),
     IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
                                         np.array([0, 2]),
                                         None,
                                         None,
                                         np.array([-1, 2]))),
     ]),
    ("IntArrayWithInt32Type",
     [IndexSpec(shape=(3, 4), indexer=(Ellipsis, np.array(1, dtype=np.int32)))
     ]),
]

MIXED_ADVANCED_INDEXING_TESTS = MIXED_ADVANCED_INDEXING_TESTS_NO_REPEATS + [
    ("SlicesAndOneIntArrayIndex",
     [
     IndexSpec(shape=(3, 4, 5), indexer=(Ellipsis,
                                         np.array([[0, 2], [1, 1]]),
                                         slice(None))),
     ]),
    ("SlicesAndTwoIntArrayIndices",
     [IndexSpec(shape=(3, 4, 5), indexer=(np.array([0, 2, -2]),
                                         slice(None, None, 2),
                                         np.array([-1, 2, -1]))),
      IndexSpec(shape=(3, 4, 5), indexer=(np.array([[0, 2], [2, 0]]),
                                          Ellipsis,
                                          np.array([[1, 0], [1, 0]]))),
     ]),]

@jtu.with_config(jax_numpy_rank_promotion="raise")
class IndexingTest(jtu.JaxTestCase):
  """Tests for Numpy indexing translation rules."""

  @parameterized.named_parameters(jtu.cases_from_list({
      "testcase_name": "{}_inshape={}_indexer={}".format(
          name, jtu.format_shape_dtype_string( shape, dtype), indexer),
       "shape": shape, "dtype": dtype, "indexer": indexer
  } for name, index_specs in STATIC_INDEXING_TESTS
    for shape, indexer in index_specs
    for dtype in all_dtypes))
  def testStaticIndexing(self, shape, dtype, indexer):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    np_fun = lambda x: np.asarray(x)[indexer]
    jnp_fun = lambda x: jnp.asarray(x)[indexer]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list({
      "testcase_name": "{}_inshape={}_indexer={}".format(
          name, jtu.format_shape_dtype_string( shape, dtype), indexer),
       "shape": shape, "dtype": dtype, "indexer": indexer
  } for name, index_specs in STATIC_INDEXING_TESTS
    for shape, indexer in index_specs
    for dtype in all_dtypes))
  def testStaticIndexingWithAtGet(self, shape, dtype, indexer):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype)]
    np_fun = lambda x: np.asarray(x)[indexer]
    jnp_fun = lambda x: jnp.asarray(x).at[indexer].get()
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters({
      "testcase_name":
          "{}_inshape={}_indexer={}".format(name,
                                            jtu.format_shape_dtype_string(
                                                shape, dtype), indexer),
      "shape": shape, "dtype": dtype, "indexer": indexer
  } for name, index_specs in STATIC_INDEXING_TESTS
    for shape, indexer in index_specs
    for dtype in float_dtypes)
  def testStaticIndexingGrads(self, shape, dtype, indexer):
    rng = jtu.rand_default(self.rng())
    tol = 1e-2 if jnp.finfo(dtype).bits == 32 else None
    arg = rng(shape, dtype)
    fun = lambda x: jnp.asarray(x)[indexer]**2
    check_grads(fun, (arg,), 2, tol, tol, tol)

  def _ReplaceSlicesWithTuples(self, idx):
    """Helper method to replace slices with tuples for dynamic indexing args."""
    if isinstance(idx, slice):
      triple = idx.start, idx.stop, idx.step
      isnone = [i for i, elt in enumerate(triple) if elt is None]
      zeros = itertools.repeat(0)
      nones = itertools.repeat(None)
      out = util.subvals(triple, zip(isnone, zeros))
      return out, lambda out: slice(*util.subvals(out, zip(isnone, nones)))
    elif isinstance(idx, (tuple, list)) and idx:
      t = type(idx)
      elts, packs = zip(*map(self._ReplaceSlicesWithTuples, idx))
      return elts, lambda elts: t((pack(i) for pack, i in zip(packs, elts)))
    else:
      return idx, lambda x: x

  @parameterized.named_parameters(
      {"testcase_name": "{}_inshape={}_indexer={}"
       .format(name, jtu.format_shape_dtype_string(shape, dtype), indexer),
       "shape": shape, "dtype": dtype, "indexer": indexer}
      for name, index_specs in [
          ("OneSliceIndex",
           [IndexSpec(shape=(5,), indexer=slice(1, 3)),
            IndexSpec(shape=(5, 4), indexer=slice(1, 3))]),
          ("TwoSliceIndices",
           [IndexSpec(shape=(5, 4), indexer=(slice(1, 3), slice(0, 2))),
            IndexSpec(shape=(5, 4, 3), indexer=(slice(1, 3), slice(0, 2)))]),
          ("NonUnitStrides", [
              IndexSpec(shape=(3,), indexer=slice(None, None, -1)),
              IndexSpec(shape=(3, 3), indexer=slice(0, 3, -2)),
              IndexSpec(shape=(3, 4, 5), indexer=slice(0, 4, 2))
          ]),
          ("OnlyStartOrStopDynamic", [
              IndexSpec(shape=(5, 4), indexer=(slice(None, 3), slice(0, 2))),
              IndexSpec(shape=(5, 4, 3), indexer=(slice(1, 3), slice(0, None)))
          ]),
      ]
      for shape, indexer in index_specs
      for dtype in all_dtypes)
  def testDynamicIndexingWithSlicesErrors(self, shape, dtype, indexer):
    rng = jtu.rand_default(self.rng())
    unpacked_indexer, pack_indexer = self._ReplaceSlicesWithTuples(indexer)

    @jax.jit
    def fun(x, unpacked_indexer):
      indexer = pack_indexer(unpacked_indexer)
      return x[indexer]

    args_maker = lambda: [rng(shape, dtype), unpacked_indexer]
    self.assertRaises(IndexError, lambda: fun(*args_maker()))

  @parameterized.named_parameters(
      {"testcase_name": "{}_inshape={}_indexer={}"
       .format(name, jtu.format_shape_dtype_string(shape, dtype), indexer),
       "shape": shape, "dtype": dtype, "indexer": indexer}
      for name, index_specs in [
          ("OneIntIndex",
           [IndexSpec(shape=(3,), indexer=1),
            IndexSpec(shape=(3, 3), indexer=0),
            IndexSpec(shape=(3, 4, 5), indexer=2),
            IndexSpec(shape=(3,), indexer=-1),
            IndexSpec(shape=(3,), indexer=-2)]),
          ("TwoIntIndices",
           [IndexSpec(shape=(3, 3), indexer=(2, 1)),
            IndexSpec(shape=(3, 4, 5), indexer=(1, 2)),
            IndexSpec(shape=(3, 4, 5), indexer=(-1, 2))]),
          ("ThreeIntIndices",
           [IndexSpec((3, 4, 5), indexer=(1, 2, 3))]),
      ]
      for shape, indexer in index_specs
      for dtype in all_dtypes)
  def testDynamicIndexingWithIntegers(self, shape, dtype, indexer):
    rng = jtu.rand_default(self.rng())
    unpacked_indexer, pack_indexer = self._ReplaceSlicesWithTuples(indexer)

    def np_fun(x, unpacked_indexer):
      indexer = pack_indexer(unpacked_indexer)
      return np.asarray(x)[indexer]

    def jnp_fun(x, unpacked_indexer):
      indexer = pack_indexer(unpacked_indexer)
      return jnp.array(x)[indexer]

    args_maker = lambda: [rng(shape, dtype), unpacked_indexer]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(
      {"testcase_name": "{}_inshape={}_indexer={}"
       .format(name, jtu.format_shape_dtype_string(shape, dtype), indexer),
       "shape": shape, "dtype": dtype, "indexer": indexer}
      for name, index_specs in [
          ("OneIntIndex",
           [IndexSpec(shape=(3,), indexer=1),
            IndexSpec(shape=(3, 3), indexer=0),
            IndexSpec(shape=(3, 4, 5), indexer=2),
            IndexSpec(shape=(3,), indexer=-1),
            IndexSpec(shape=(3,), indexer=-2),
            ]),
          ("TwoIntIndices",
           [IndexSpec(shape=(3, 3), indexer=(2, 1)),
            IndexSpec(shape=(3, 4, 5), indexer=(1, 2)),
            IndexSpec(shape=(3, 4, 5), indexer=(-1, 2)),
            ]),
          ("ThreeIntIndices",
           [IndexSpec((3, 4, 5), indexer=(1, 2, 3))]),
      ]
      for shape, indexer in index_specs
      for dtype in float_dtypes)
  def testDynamicIndexingWithIntegersGrads(self, shape, dtype, indexer):
    rng = jtu.rand_default(self.rng())
    tol = 1e-2 if jnp.finfo(dtype).bits == 32 else None
    unpacked_indexer, pack_indexer = self._ReplaceSlicesWithTuples(indexer)

    @jax.jit
    def fun(unpacked_indexer, x):
      indexer = pack_indexer(unpacked_indexer)
      return x[indexer]

    arr = rng(shape, dtype)
    check_grads(partial(fun, unpacked_indexer), (arr,), 2, tol, tol, tol)

  @parameterized.named_parameters(
      {"testcase_name": "{}_inshape={}_indexer={}"
       .format(name, jtu.format_shape_dtype_string(shape, dtype), indexer),
       "shape": shape, "dtype": dtype, "indexer": indexer}
      for name, index_specs in ADVANCED_INDEXING_TESTS
      for shape, indexer in index_specs
      for dtype in all_dtypes)
  def testAdvancedIntegerIndexing(self, shape, dtype, indexer):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype), indexer]
    np_fun = lambda x, idx: np.asarray(x)[idx]
    jnp_fun = lambda x, idx: jnp.asarray(x)[idx]
    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  @parameterized.named_parameters(
      {"testcase_name": "{}_inshape={}_indexer={}"
       .format(name, jtu.format_shape_dtype_string(shape, dtype), indexer),
       "shape": shape, "dtype": dtype, "indexer": indexer}
      for name, index_specs in [
          ("One1DIntArrayIndex",
           [IndexSpec(shape=(3,), indexer=np.array([0, 1])),
            IndexSpec(shape=(3, 3), indexer=np.array([1, 2, 1])),
            IndexSpec(shape=(3, 4, 5), indexer=np.array([0, 2, 0, 1])),
            IndexSpec(shape=(3,), indexer=np.array([-1, 1])),
            IndexSpec(shape=(3,), indexer=np.array([-2, -1])),
            ]),
          ("One2DIntArrayIndex",
           [IndexSpec(shape=(3,), indexer=np.array([[0, 0]])),
            IndexSpec(shape=(3, 3), indexer=np.array([[1, 2, 1],
                                                       [0, 1, -1]])),
            IndexSpec(shape=(3, 4, 5), indexer=np.array([[0, 2, 0, 1],
                                                          [-1, -2, 1, 0]])),
            ]),
          ("Two1DIntArrayIndicesNoBroadcasting",
           [IndexSpec(shape=(3, 3), indexer=(np.array([0, 1]),
                                             np.array([1, 2]))),
            IndexSpec(shape=(3, 4, 5), indexer=(np.array([0, 2, 0, 1]),
                                                np.array([-1, 0, -1, 2]))),
            ]),
          ("Two1DIntArrayIndicesWithBroadcasting",
           [IndexSpec(shape=(3, 3), indexer=(np.array([[0, 1]]),
                                             np.array([1, 2]))),
            IndexSpec(shape=(3, 4, 5), indexer=(np.array([[0, 2, 0, 1]]),
                                                np.array([-1, 0, -1, 2]))),
            ]),
          ("TupleOfPythonIntsAndIntArrays",
           [IndexSpec(shape=(3, 4, 5), indexer=(0, np.array([0, 1]))),
            IndexSpec(shape=(3, 4, 5), indexer=(0, 1,
                                                np.array([[2, 3, 0, 3]]))),
            ]),
          ("TupleOfListsOfPythonIntsAndIntArrays",
           [IndexSpec(shape=(3, 4, 5), indexer=([0, 1], np.array([0]))),
            IndexSpec(shape=(3, 4, 5), indexer=([[0], [-1]],
                                                np.array([[2, 3, 0, 3]]))),
            ]),
      ]
      for shape, indexer in index_specs
      for dtype in float_dtypes)
  def testAdvancedIntegerIndexingGrads(self, shape, dtype, indexer):
    rng = jtu.rand_default(self.rng())
    tol = 1e-2 if jnp.finfo(dtype).bits == 32 else None
    arg = rng(shape, dtype)
    fun = lambda x: jnp.asarray(x)[indexer]
    check_grads(fun, (arg,), 2, tol, tol, eps=1.)

  @parameterized.named_parameters(
      {"testcase_name": "{}_inshape={}_indexer={}"
       .format(name, jtu.format_shape_dtype_string(shape, dtype), indexer),
       "shape": shape, "dtype": dtype, "indexer": indexer}
      for name, index_specs in MIXED_ADVANCED_INDEXING_TESTS
      for shape, indexer in index_specs
      for dtype in all_dtypes)
  def testMixedAdvancedIntegerIndexing(self, shape, dtype, indexer):
    rng = jtu.rand_default(self.rng())
    indexer_with_dummies = [e if isinstance(e, np.ndarray) else ()
                            for e in indexer]
    substitutes = [(i, e) for i, e in enumerate(indexer)
                   if not isinstance(e, np.ndarray)]
    args_maker = lambda: [rng(shape, dtype), indexer_with_dummies]

    def jnp_fun(x, indexer_with_dummies):
      idx = type(indexer)(util.subvals(indexer_with_dummies, substitutes))
      return jnp.asarray(x)[idx]

    def np_fun(x, indexer_with_dummies):
      idx = type(indexer)(util.subvals(indexer_with_dummies, substitutes))
      return np.asarray(x)[idx]

    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  def testAdvancedIndexingManually(self):
    x = np.random.RandomState(0).randn(3, 4, 5)
    index_array = np.array([0, 2, -1, 0])

    op = lambda x, index_array: x[..., index_array, :]
    cop = jax.jit(op)

    a1 = op(x, index_array)
    a2 = cop(x, index_array)

    self.assertAllClose(a1, a2)

    op = lambda x, index_array: x[..., index_array, :, index_array, None]
    cop = jax.jit(op)

    a1 = op(x, index_array)
    a2 = cop(x, index_array)

    self.assertAllClose(a1, a2)

    op = lambda x, index_array: x[index_array, ..., index_array[:, None], None]
    cop = jax.jit(op)

    a1 = op(x, index_array)
    a2 = cop(x, index_array)

    self.assertAllClose(a1, a2)

  def testUnpacking(self):

    def foo(x):
      a, b, c = x
      return a + b + c

    cfoo = jax.jit(foo)

    a1 = foo(np.arange(3))
    a2 = cfoo(np.arange(3))

    self.assertAllClose(a1, a2)

  def testBooleanIndexingArray1D(self):
    idx = np.array([True, True, False])
    x = jax.device_put(np.arange(3))
    ans = x[idx]
    expected = np.arange(3)[idx]
    self.assertAllClose(ans, expected, check_dtypes=False)

  def testBooleanIndexingList1D(self):
    idx = [True, True, False]
    x = jax.device_put(np.arange(3))
    with self.assertRaisesRegex(TypeError, ARRAY_MSG):
      x[idx]

  def testBooleanIndexingArray2DBroadcast(self):
    idx = np.array([True, True, False, True])
    x = np.arange(8).reshape(4, 2)
    ans = jax.device_put(x)[idx]
    expected = x[idx]
    self.assertAllClose(ans, expected, check_dtypes=False)

  def testBooleanIndexingList2DBroadcast(self):
    idx = [True, True, False, True]
    x = np.arange(8).reshape(4, 2)
    with self.assertRaisesRegex(TypeError, ARRAY_MSG):
      jax.device_put(x)[idx]

  def testBooleanIndexingArray2D(self):
    idx = np.array([[True, False],
                     [False, True],
                     [False, False],
                     [True, True]])
    x = np.arange(8).reshape(4, 2)
    ans = jax.device_put(x)[idx]
    expected = x[idx]
    self.assertAllClose(ans, expected, check_dtypes=False)

  def testBooleanIndexingDynamicShapeError(self):
    x = np.zeros(3)
    i = np.array([True, True, False])
    self.assertRaises(IndexError, lambda: jax.jit(lambda x, i: x[i])(x, i))

  def testScalarBooleanIndexingNotImplemented(self):
    msg = "JAX arrays do not support boolean scalar indices"
    with self.assertRaisesRegex(TypeError, msg):
      jnp.arange(4)[True]
    with self.assertRaisesRegex(TypeError, msg):
      jnp.arange(4)[False]

  def testIssue187(self):
    x = jnp.ones((5, 5))
    x[[0, 2, 4], [0, 2, 4]]  # doesn't crash

    x = np.arange(25).reshape((5, 5))
    ans = jax.jit(lambda x: x[[0, 2, 4], [0, 2, 4]])(x)
    expected = x[[0, 2, 4], [0, 2, 4]]
    self.assertAllClose(ans, expected, check_dtypes=False)

  def testJVPOfGradOfIndexing(self):
    # Should return a value, even though we didn't pass a symbolic zero as the
    # index tangent.
    x = jnp.ones((3, 4), jnp.float32)
    i = jnp.ones((3,), jnp.int32)
    f = lambda x, i: jnp.sum(x[i])
    primals, tangents = jax.jvp(jax.grad(f), (x, i),
                                (x, np.zeros(i.shape, dtypes.float0)))
    expected = np.broadcast_to(
      np.array([0, 3, 0], dtype=np.float32)[:, None], (3, 4))
    self.assertAllClose(expected, primals)
    self.assertAllClose(np.zeros_like(x), tangents)

  def testTrivialGatherIsntGenerated(self):
    # https://github.com/google/jax/issues/1621
    jaxpr = jax.make_jaxpr(lambda x: x[:, None])(np.arange(4))
    self.assertEqual(len(jaxpr.jaxpr.eqns), 1)
    self.assertNotIn('gather', str(jaxpr))

  def testIndexingEmptyDimension(self):
    # Issue 2671: XLA error when indexing into dimension of size 0
    x = jnp.ones((2, 0))
    # The following work, even on axis 1 of size 0
    with jax.numpy_rank_promotion('allow'):
      _ = x[0, :] + x[0, None] + x[0, 1:] + x[0, 1:3:2]

    with self.assertRaisesRegex(IndexError,
                                "index .* is out of bounds for axis .* with size 0"):
      _ = np.ones((2, 0))[0, 0]  # The numpy error
    with self.assertRaisesRegex(IndexError,
                                "index is out of bounds for axis .* with size 0"):
      _ = x[0, 0]  # JAX indexing
    with self.assertRaisesRegex(IndexError,
                                "index is out of bounds for axis .* with size 0"):
      jax.jit(lambda i: x[0, i])(0)  # JAX indexing under jit

  def testBooleanIndexingWithEmptyResult(self):
    # based on a TensorFlow Probability test that started failing after #1622
    x = jnp.array([-1])
    mask = jnp.array([False])
    ans = x[mask]  # doesn't crash

    expected =  np.array([-1])[np.array([False])]
    self.assertAllClose(ans, expected, check_dtypes=False)

  def testBooleanIndexingShapeMismatch(self):
    # Regression test for https://github.com/google/jax/issues/7329
    x = jnp.arange(4)
    idx = jnp.array([True, False])
    with self.assertRaisesRegex(IndexError, "boolean index did not match shape.*"):
      x[idx]

  def testNontrivialBooleanIndexing(self):
    # Test nontrivial corner case in boolean indexing shape validation
    rng = jtu.rand_default(self.rng())
    index = (rng((2, 3), np.bool_), rng((6,), np.bool_))

    args_maker = lambda: [rng((2, 3, 6), np.int32)]
    np_fun = lambda x: np.asarray(x)[index]
    jnp_fun = lambda x: jnp.asarray(x)[index]

    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    self._CompileAndCheck(jnp_fun, args_maker)

  def testFloatIndexingError(self):
    BAD_INDEX_TYPE_ERROR = "Indexer must have integer or boolean type, got indexer with type"
    with self.assertRaisesRegex(TypeError, BAD_INDEX_TYPE_ERROR):
      jnp.zeros(2)[0.]
    with self.assertRaisesRegex(TypeError, BAD_INDEX_TYPE_ERROR):
      jnp.zeros((2, 2))[(0, 0.)]
    with self.assertRaisesRegex(TypeError, BAD_INDEX_TYPE_ERROR):
      jnp.zeros((2, 2))[(0, 0.)]
    with self.assertRaisesRegex(TypeError, BAD_INDEX_TYPE_ERROR):
      jax.jit(lambda idx: jnp.zeros((2, 2))[idx])((0, 0.))
    with self.assertRaisesRegex(TypeError, BAD_INDEX_TYPE_ERROR):
      ops.index_add(jnp.zeros(2), 0., 1.)
    with self.assertRaisesRegex(TypeError, BAD_INDEX_TYPE_ERROR):
      ops.index_update(jnp.zeros(2), 0., 1.)

  def testIndexOutOfBounds(self):  # https://github.com/google/jax/issues/2245
    array = jnp.ones(5)
    self.assertAllClose(array, array[:10])

def _broadcastable_shapes(shape):
  """Returns all shapes that broadcast to `shape`."""
  def f(rshape):
    yield []
    if rshape:
      for s in f(rshape[1:]):
        yield rshape[0:1] + s
      if rshape[0] != 1:
        for s in f(rshape[1:]):
          yield [1] + s
  for x in f(list(reversed(shape))):
    yield list(reversed(x))

def _update_shape(shape, indexer):
  return np.zeros(shape)[indexer].shape

class UpdateOps(enum.Enum):
  UPDATE = 0
  ADD = 1
  MUL = 2
  DIV = 3
  POW = 4
  MIN = 5
  MAX = 6

  def np_fn(op, indexer, x, y):
    x = x.copy()
    x[indexer] = {
      UpdateOps.UPDATE: lambda: y,
      UpdateOps.ADD: lambda: x[indexer] + y,
      UpdateOps.MUL: lambda: x[indexer] * y,
      UpdateOps.DIV: jtu.ignore_warning(category=RuntimeWarning)(
        lambda: x[indexer] / y.astype(x.dtype)),
      UpdateOps.POW: jtu.ignore_warning(category=RuntimeWarning)(
        lambda: x[indexer] ** y.astype(x.dtype)),
      UpdateOps.MIN: lambda: np.minimum(x[indexer], y),
      UpdateOps.MAX: lambda: np.maximum(x[indexer], y),
    }[op]()
    return x

  def jax_fn(op, indexer, x, y, indices_are_sorted=False,
             unique_indices=False):
    return {
      UpdateOps.UPDATE: ops.index_update,
      UpdateOps.ADD: ops.index_add,
      UpdateOps.MUL: ops.index_mul,
      UpdateOps.MIN: ops.index_min,
      UpdateOps.MAX: ops.index_max,
    }[op](x, indexer, y, indices_are_sorted=indices_are_sorted,
          unique_indices=unique_indices)

  def sugar_fn(op, indexer, x, y, indices_are_sorted=False,
             unique_indices=False):
    x = jnp.array(x)
    return {
      UpdateOps.UPDATE: x.at[indexer].set,
      UpdateOps.ADD: x.at[indexer].add,
      UpdateOps.MUL: x.at[indexer].multiply,
      UpdateOps.DIV: x.at[indexer].divide,
      UpdateOps.POW: x.at[indexer].power,
      UpdateOps.MIN: x.at[indexer].min,
      UpdateOps.MAX: x.at[indexer].max,
    }[op](y, indices_are_sorted=indices_are_sorted,
          unique_indices=unique_indices)

  def dtypes(op):
    if op == UpdateOps.UPDATE:
      return all_dtypes
    elif op == UpdateOps.DIV or op == UpdateOps.POW:
      return jtu.dtypes.inexact
    else:
      return default_dtypes

def _update_tol(op):
  if op == UpdateOps.POW:
    tol = {np.complex64: 1e-4 if jtu.device_under_test() == "tpu" else 1e-5,
           np.complex128: 1e-14}
  else:
    tol = {np.complex128: 1e-14}
  return tol

@jtu.with_config(jax_numpy_rank_promotion="raise")
class IndexedUpdateTest(jtu.JaxTestCase):

  @parameterized.named_parameters(jtu.named_cases_from_sampler(lambda s: ({
      "testcase_name": "{}_inshape={}_indexer={}_update={}_sugared={}_op={}".format(
          name, jtu.format_shape_dtype_string(shape, dtype), indexer,
          jtu.format_shape_dtype_string(update_shape, update_dtype), sugared, op.name),
       "shape": shape, "dtype": dtype, "indexer": indexer,
       "update_shape": update_shape, "update_dtype": update_dtype,
       "op": op, "sugared": sugared
  } for name, index_specs in s(STATIC_INDEXING_TESTS)
    for shape, indexer in s(index_specs)
    for op in s(UpdateOps)
    for dtype in s(UpdateOps.dtypes(op))
    for update_shape in s(_broadcastable_shapes(_update_shape(shape, indexer)))
    for update_dtype in s([dtype] if op == UpdateOps.ADD else all_dtypes)
    for sugared in (s([True, False]) if op not in [UpdateOps.DIV, UpdateOps.POW] else [True]))))
  def testStaticIndexing(self, shape, dtype, update_shape, update_dtype,
                         indexer, sugared, op):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype), rng(update_shape, update_dtype)]
    np_fn = lambda x, y: UpdateOps.np_fn(op, indexer, x, y)
    if sugared:
      jax_fn = lambda x, y: UpdateOps.sugar_fn(op, indexer, x, y)
    else:
      jax_fn = lambda x, y: UpdateOps.jax_fn(op, indexer, x, y)
    self._CheckAgainstNumpy(np_fn, jax_fn, args_maker, tol=_update_tol(op))
    self._CompileAndCheck(jax_fn, args_maker)

  @parameterized.named_parameters(jtu.named_cases_from_sampler(lambda s: ({
      "testcase_name": "{}_inshape={}_indexer={}_update={}_op={}".format(
          name, jtu.format_shape_dtype_string(shape, dtype), indexer,
          jtu.format_shape_dtype_string(update_shape, update_dtype), op.name),
       "shape": shape, "dtype": dtype, "indexer": indexer,
       "update_shape": update_shape, "update_dtype": update_dtype,
       "op": op
  } for name, index_specs in s(ADVANCED_INDEXING_TESTS_NO_REPEATS)
    for shape, indexer in s(index_specs)
    for op in s(UpdateOps)
    for dtype in s(UpdateOps.dtypes(op))
    for update_shape in s(_broadcastable_shapes(_update_shape(shape, indexer)))
    for update_dtype in s([dtype] if op == UpdateOps.ADD else all_dtypes))))
  def testAdvancedIndexing(self, shape, dtype, update_shape, update_dtype,
                           indexer, op):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype), rng(update_shape, update_dtype)]
    np_fn = lambda x, y: UpdateOps.np_fn(op, indexer, x, y)
    jax_fn = lambda x, y: UpdateOps.sugar_fn(op, indexer, x, y,
                                             unique_indices=True)
    self._CheckAgainstNumpy(np_fn, jax_fn, args_maker, tol=_update_tol(op))
    self._CompileAndCheck(jax_fn, args_maker)

  @parameterized.named_parameters(jtu.named_cases_from_sampler(lambda s: ({
      "testcase_name": "{}_inshape={}_indexer={}_update={}_op={}".format(
          name, jtu.format_shape_dtype_string(shape, dtype), indexer,
          jtu.format_shape_dtype_string(update_shape, update_dtype), op.name),
       "shape": shape, "dtype": dtype, "indexer": indexer,
       "update_shape": update_shape, "update_dtype": update_dtype,
       "op": op
  } for name, index_specs in s(ADVANCED_INDEXING_TESTS_NO_REPEATS_SORTED)
    for shape, indexer in s(index_specs)
    for op in s(UpdateOps)
    for dtype in s(UpdateOps.dtypes(op))
    for update_shape in s(_broadcastable_shapes(_update_shape(shape, indexer)))
    for update_dtype in s([dtype] if op == UpdateOps.ADD else all_dtypes))))
  def testAdvancedIndexingSorted(self, shape, dtype, update_shape, update_dtype,
                           indexer, op):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype), rng(update_shape, update_dtype)]
    np_fn = lambda x, y: UpdateOps.np_fn(op, indexer, x, y)
    jax_fn = lambda x, y: UpdateOps.sugar_fn(
      op, indexer, x, y, indices_are_sorted=True, unique_indices=True)
    self._CheckAgainstNumpy(np_fn, jax_fn, args_maker, check_dtypes=True,
                            tol=_update_tol(op))
    self._CompileAndCheck(jax_fn, args_maker, check_dtypes=True)

  @parameterized.named_parameters(jtu.named_cases_from_sampler(lambda s: ({
      "testcase_name": "{}_inshape={}_indexer={}_update={}_op={}".format(
          name, jtu.format_shape_dtype_string(shape, dtype), indexer,
          jtu.format_shape_dtype_string(update_shape, update_dtype), op.name),
       "shape": shape, "dtype": dtype, "indexer": indexer,
       "update_shape": update_shape, "update_dtype": update_dtype,
       "op": op
  } for name, index_specs in s(MIXED_ADVANCED_INDEXING_TESTS_NO_REPEATS)
    for shape, indexer in s(index_specs)
    for op in s(UpdateOps)
    for dtype in s(UpdateOps.dtypes(op))
    for update_shape in s(_broadcastable_shapes(_update_shape(shape, indexer)))
    for update_dtype in s([dtype] if op == UpdateOps.ADD else all_dtypes))))
  def testMixedAdvancedIndexing(self, shape, dtype, update_shape, update_dtype,
                                indexer, op):
    rng = jtu.rand_default(self.rng())
    args_maker = lambda: [rng(shape, dtype), rng(update_shape, update_dtype)]
    np_fn = lambda x, y: UpdateOps.np_fn(op, indexer, x, y)
    jax_fn = lambda x, y: UpdateOps.sugar_fn(op, indexer, x, y)
    self._CheckAgainstNumpy(np_fn, jax_fn, args_maker, tol=_update_tol(op))
    self._CompileAndCheck(jax_fn, args_maker)

  @parameterized.named_parameters(jtu.cases_from_list({
      "testcase_name": "{}_inshape={}_indexer={}_update={}_op={}".format(
          name, jtu.format_shape_dtype_string(shape, dtype), indexer,
          jtu.format_shape_dtype_string(update_shape, update_dtype), op.name),
       "shape": shape, "dtype": dtype, "indexer": indexer,
       "update_shape": update_shape, "update_dtype": update_dtype,
       "op": op
  } for name, index_specs in STATIC_INDEXING_TESTS
    for shape, indexer in index_specs
    for op in [UpdateOps.ADD, UpdateOps.MUL, UpdateOps.UPDATE]
    for dtype in float_dtypes
    for update_shape in _broadcastable_shapes(_update_shape(shape, indexer))
    for update_dtype in ([dtype] if op == UpdateOps.ADD else float_dtypes)))
  def testStaticIndexingGrads(self, shape, dtype, update_shape, update_dtype,
                              indexer, op):
    rng = jtu.rand_default(self.rng())
    jax_fn = lambda x, y: UpdateOps.sugar_fn(op, indexer, x, y)
    x = rng(shape, dtype)
    y = rng(update_shape, update_dtype)
    check_grads(jax_fn, (x, y), 2, rtol=1e-3, atol=1e-3, eps=1.)

  @parameterized.named_parameters(jtu.named_cases_from_sampler(lambda s: ({
      "testcase_name": "{}_inshape={}_indexer={}_update={}_op={}".format(
          name, jtu.format_shape_dtype_string(shape, dtype), indexer,
          jtu.format_shape_dtype_string(update_shape, update_dtype), op.name),
       "shape": shape, "dtype": dtype, "indexer": indexer,
       "update_shape": update_shape, "update_dtype": update_dtype,
       "op": op
  } for name, index_specs in s(ADVANCED_INDEXING_TESTS_NO_REPEATS)
    for shape, indexer in s(index_specs)
    for op in s([UpdateOps.ADD, UpdateOps.MUL, UpdateOps.UPDATE])
    for dtype in s(float_dtypes)
    for update_shape in s(_broadcastable_shapes(_update_shape(shape, indexer)))
    for update_dtype in s([dtype] if op == UpdateOps.ADD else float_dtypes))))
  def testAdvancedIndexingGrads(self, shape, dtype, update_shape, update_dtype,
                                indexer, op):
    rng = jtu.rand_default(self.rng())
    jax_fn = lambda x, y: UpdateOps.sugar_fn(op, indexer, x, y,
                                             unique_indices=True)
    x = rng(shape, dtype)
    y = rng(update_shape, update_dtype)
    check_grads(jax_fn, (x, y), 2, rtol=1e-3, atol=1e-3, eps=1.)

  def testSegmentSumBehavior(self):
    # testAdvancedIndexing compares against NumPy, and as a result doesn't check
    # repeated indices. This test is just a simple manual check, based on
    # https://www.tensorflow.org/api_docs/python/tf/math/segment_sum
    data = np.array([5, 1, 7, 2, 3, 4, 1, 3])
    segment_ids = np.array([0, 0, 0, 1, 2, 2, 3, 3])

    ans = ops.index_add(np.zeros(np.max(segment_ids) + 1), segment_ids, data)
    expected = np.array([13, 2, 7, 4])
    self.assertAllClose(ans, expected, check_dtypes=False)

  def testSegmentSum(self):
    data = jnp.array([5, 1, 7, 2, 3, 4, 1, 3])
    segment_ids = jnp.array([0, 0, 0, 1, 2, 2, 3, 3])

    # test with explicit num_segments
    ans = ops.segment_sum(data, segment_ids, num_segments=4)
    expected = jnp.array([13, 2, 7, 4])
    self.assertAllClose(ans, expected, check_dtypes=False)

    # test with explicit num_segments larger than the higher index.
    ans = ops.segment_sum(data, segment_ids, num_segments=5)
    expected = jnp.array([13, 2, 7, 4, 0])
    self.assertAllClose(ans, expected, check_dtypes=False)

    # test without explicit num_segments
    ans = ops.segment_sum(data, segment_ids)
    expected = jnp.array([13, 2, 7, 4])
    self.assertAllClose(ans, expected, check_dtypes=False)

    # test with negative segment ids and segment ids larger than num_segments,
    # that will be wrapped with the `mod`.
    segment_ids = jnp.array([0, 4, 8, 1, 2, -6, -1, 3])
    ans = ops.segment_sum(data, segment_ids, num_segments=4)
    expected = jnp.array([5, 2, 3, 3])
    self.assertAllClose(ans, expected, check_dtypes=False)

    # test with negative segment ids and without without explicit num_segments
    # such as num_segments is defined by the smaller index.
    segment_ids = jnp.array([3, 3, 3, 4, 5, 5, -7, -6])
    ans = ops.segment_sum(data, segment_ids)
    expected = jnp.array([0, 0, 0, 13, 2, 7])
    self.assertAllClose(ans, expected, check_dtypes=False)

  @parameterized.named_parameters(itertools.chain.from_iterable(
      jtu.cases_from_list({
        "testcase_name": "_{}_{}_num_segments={}_bucket_size={}".format(
          jtu.format_shape_dtype_string(shape, dtype),
          reducer.__name__, num_segments, bucket_size),
        "dtype": dtype, "shape": shape,
        "reducer": reducer, "op": op, "identity": identity,
        "num_segments": num_segments, "bucket_size": bucket_size}
      for dtype in default_dtypes
      for shape in [(8,), (7, 4), (6, 4, 2)]
      for bucket_size in [None, 2]
      for num_segments in [None, 1, 3])
    for reducer, op, identity in [
      (ops.segment_sum, np.add, 0),
      (ops.segment_prod, np.multiply, 1),
      (ops.segment_min, np.minimum, float('inf')),
      (ops.segment_max, np.maximum, -float('inf')),
    ]))
  def testSegmentReduce(self, shape, dtype, reducer, op, identity, num_segments, bucket_size):
    rng = jtu.rand_default(self.rng())
    idx_rng = jtu.rand_int(self.rng(), low=-2, high=3)
    args_maker = lambda: [rng(shape, dtype), idx_rng(shape[:1], jnp.int32)]

    if np.issubdtype(dtype, np.integer):
      if np.isposinf(identity):
        identity = np.iinfo(dtype).max
      elif np.isneginf(identity):
        identity = np.iinfo(dtype).min

    jnp_fun = lambda data, segment_ids: reducer(
      data, segment_ids, num_segments=num_segments, bucket_size=bucket_size)

    def np_fun(data, segment_ids):
      size = num_segments if num_segments is not None else (segment_ids.max() + 1)
      out = np.full((size,) + shape[1:], identity, dtype)
      for i, val in zip(segment_ids, data):
        if 0 <= i < size:
          out[i] = op(out[i], val).astype(dtype)
      return out

    self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
    if num_segments is not None:
      self._CompileAndCheck(jnp_fun, args_maker)

  def testIndexDtypeError(self):
    # https://github.com/google/jax/issues/2795
    jnp.array(1)  # get rid of startup warning
    with warnings.catch_warnings(record=True) as w:
      warnings.simplefilter("error")
      jnp.zeros(5).at[::2].set(1)
      self.assertLen(w, 0)

  @contextmanager
  def assertNoWarnings(self):
    with warnings.catch_warnings(record=True) as caught_warnings:
      yield
    self.assertEmpty(caught_warnings)

  @parameterized.named_parameters(jtu.cases_from_list({
      "testcase_name": "idx={}".format(idx), "idx": idx, "idx_type": idx_type}
    for idx, idx_type in [
      ([0], "array"),
      ([0, 0], "array"),
      ([[0, 0]], "tuple"),
      ([0, [0, 1]], "tuple"),
      ([0, np.arange(2)], "tuple"),
      ([0, None], "tuple"),
      ([0, slice(None)], "tuple"),
    ]))
  def testIndexSequenceDeprecation(self, idx, idx_type):
    normalize = {"array": np.array, "tuple": tuple}[idx_type]
    msg = {"array": ARRAY_MSG, "tuple": TUPLE_MSG}[idx_type]
    x = jnp.arange(6).reshape(3, 2)

    with self.assertRaisesRegex(TypeError, msg):
      x[idx]
    with self.assertNoWarnings():
      x[normalize(idx)]

    with self.assertRaisesRegex(TypeError, msg):
      x.at[idx].set(0)
    with self.assertNoWarnings():
      x.at[normalize(idx)].set(0)

if __name__ == "__main__":
  absltest.main(testLoader=jtu.JaxTestLoader())


import pandas as pd
import numpy as np

def create_dummy_df(df, cat_cols, dummy_na):
    for col in  cat_cols:
        try:
            # for each cat add dummy var, drop original column
            df = pd.concat([df.drop(col, axis=1), pd.get_dummies(df[col], prefix=col, prefix_sep='_', drop_first=True, dummy_na=dummy_na)], axis=1)
        except:
            continue
    return df

def find_missing(df):
    miss = df.isnull().sum().sort_values(ascending=False)
    percent_miss = ((df.isnull().sum() * 100) / df.isnull().count().sort_values(ascending=False))
    missing = pd.concat([miss, percent_miss], axis=1, keys=['Total Missing Instances', '% Missing'], sort=False).sort_values('Total Missing Instances', ascending=False)
    return missing


import numpy as np    # the numpy library
import matplotlib.pyplot as plt   # matplotlib's pyplot
import sys      # gives access to C-like sys library
import os    # gives access to operating system

print(sys.argv)  # command line arguments
print(os.getcwd())  # print current working directory


#!/usr/bin/env python
# -*- coding: utf-8 -*-

import numpy as np

samples = np.zeros((6,), dtype=[('sensor_code', 'S4'),('position', float), ('value', float)])
print(samples.ndim)
print(samples.shape)
print(samples.dtype.names)
samples[:] = [('ALFA', 1, 0.37), ('BETA', 1, 0.11), ('TAU', 1, 0.13),
              ('ALFA', 1.5, 0.37), ('ALFA', 3, 0.11), ('TAU', 1.2, 0.13)]
print(samples)
print(samples['sensor_code'])
print(samples['value'])
print(samples[0])
samples[0]['sensor_code'] = 'TAU'
print(samples[0])
print(samples[['position', 'value']])
print(samples[samples['sensor_code'] == b'ALFA'])


import argparse
import imutils
import cv2
import sys
import numpy as np

class ArucoTag:
    # set class variable
    path_to_save = r'output/DICT_4X4_50_1.png'
    path_to_read = r"input/image1.png"
    ARUCO_DICT = {
        "DICT_4X4_50": cv2.aruco.DICT_4X4_50,
        "DICT_4X4_100": cv2.aruco.DICT_4X4_100,
        "DICT_4X4_250": cv2.aruco.DICT_4X4_250,
        "DICT_4X4_1000": cv2.aruco.DICT_4X4_1000,
        "DICT_5X5_50": cv2.aruco.DICT_5X5_50,
        "DICT_5X5_100": cv2.aruco.DICT_5X5_100,
        "DICT_5X5_250": cv2.aruco.DICT_5X5_250,
        "DICT_5X5_1000": cv2.aruco.DICT_5X5_1000,
        "DICT_6X6_50": cv2.aruco.DICT_6X6_50,
        "DICT_6X6_100": cv2.aruco.DICT_6X6_100,
        "DICT_6X6_250": cv2.aruco.DICT_6X6_250,
        "DICT_6X6_1000": cv2.aruco.DICT_6X6_1000,
        "DICT_7X7_50": cv2.aruco.DICT_7X7_50,
        "DICT_7X7_100": cv2.aruco.DICT_7X7_100,
        "DICT_7X7_250": cv2.aruco.DICT_7X7_250,
        "DICT_7X7_1000": cv2.aruco.DICT_7X7_1000,
        "DICT_ARUCO_ORIGINAL": cv2.aruco.DICT_ARUCO_ORIGINAL,
        "DICT_APRILTAG_16h5": cv2.aruco.DICT_APRILTAG_16h5,
        "DICT_APRILTAG_25h9": cv2.aruco.DICT_APRILTAG_25h9,
        "DICT_APRILTAG_36h10": cv2.aruco.DICT_APRILTAG_36h10,
        "DICT_APRILTAG_36h11": cv2.aruco.DICT_APRILTAG_36h11
    }

    def __init__(self, arucoType="DICT_4X4_50", size=300, id=1):
        self.arucoType = self.ARUCO_DICT[arucoType]
        # will generate 300x300x1 channel image
        self.size = size
        self.id = id  # what id need to attached in the arucoImage

        # load the ArUCo dictionary
        self.arucoDict = cv2.aruco.Dictionary_get(self.arucoType)

    def generateArucoTag(self):
        tag = np.zeros((self.size, self.size, 1), dtype="uint8")
        # here 1 is padding in the image
        cv2.aruco.drawMarker(self.arucoDict, self.id, 300, tag, 1)
        cv2.imwrite(self.path_to_save, tag)
        cv2.imshow("ArUCo Tag", tag)
        print(f"Successfully saved {self.path_to_save}, ID {self.id} ")
        cv2.waitKey(0)

    def detectArucoTag(self, image):
        image = imutils.resize(image, width=600)
        arucoParams = cv2.aruco.DetectorParameters_create()
        (corners, ids, rejected) = cv2.aruco.detectMarkers(image, self.arucoDict,
                                                           parameters=arucoParams)

        # verify *at least* one ArUco marker was detected
        imls = []
        if len(corners) > 0:
            # flatten the ArUco IDs list
            ids = ids.flatten()

            # loop over the detected ArUCo corners
            for (markerCorner, markerID) in zip(corners, ids):
                # extract the marker corners (which are always returned in
                # top-left, top-right, bottom-right, and bottom-left order)
                corners = markerCorner.reshape((4, 2))
                (topLeft, topRight, bottomRight, bottomLeft) = corners

                # convert each of the (x, y)-coordinate pairs to integers
                topRight = (int(topRight[0]), int(topRight[1]))
                bottomRight = (int(bottomRight[0]), int(bottomRight[1]))
                bottomLeft = (int(bottomLeft[0]), int(bottomLeft[1]))
                topLeft = (int(topLeft[0]), int(topLeft[1]))

                # draw the bounding box of the ArUCo detection
                cv2.line(image, topLeft, topRight, (0, 255, 0), 2)
                cv2.line(image, topRight, bottomRight, (0, 255, 0), 2)
                cv2.line(image, bottomRight, bottomLeft, (0, 255, 0), 2)
                cv2.line(image, bottomLeft, topLeft, (0, 255, 0), 2)

                # compute and draw the center (x, y)-coordinates of the ArUco
                # marker
                cX = int((topLeft[0] + bottomRight[0]) / 2.0)
                cY = int((topLeft[1] + bottomRight[1]) / 2.0)
                cv2.circle(image, (cX, cY), 4, (0, 0, 255), -1)

                # draw a line from center point to the right, Horizontal Line
                cv2.line(image, (cX, cY), (cX + 40, cY), (255, 0, 0), 2)
                # draw a line from center point to the upwards, Vertical Line
                cv2.line(image, (cX, cY), (cX, cY - 40), (0, 0, 255), 2)
                # draw a line from center point to the downwards, Vertical Line
                cv2.line(image, (cX, cY), (cX, cY + 40), (0, 255, 0), 2)

                # draw the ArUco marker ID on the image
                cv2.putText(image, str(markerID),
                            (topLeft[0], topLeft[1] - 10), cv2.FONT_HERSHEY_SIMPLEX,
                            0.8, (255, 0, 0), 3)
                print("[INFO] ArUco marker ID: {}".format(markerID))
                imls.append(image.copy())
            return image, imls


import os
import unittest

from cmdstanpy.cmdstan_args import SamplerArgs, CmdStanArgs
from cmdstanpy.utils import EXTENSION
from cmdstanpy.model import Model
from cmdstanpy.stanfit import StanFit
import numpy as np

here = os.path.dirname(os.path.abspath(__file__))
datafiles_path = os.path.join(here, 'data')

code = '''data {
  int<lower=0> N;
  int<lower=0,upper=1> y[N];
}
parameters {
  real<lower=0,upper=1> theta;
}
model {
  theta ~ beta(1,1);
  for (n in 1:N)
    y[n] ~ bernoulli(theta);
}
'''

class ModelTest(unittest.TestCase):
    def test_model_good(self):
        stan = os.path.join(datafiles_path, 'bernoulli.stan')
        exe = os.path.join(datafiles_path, 'bernoulli' + EXTENSION)

        model = Model(stan_file=stan)
        self.assertEqual(stan, model.stan_file)
        self.assertEqual(None, model.exe_file)

        model = Model(stan_file=stan, exe_file=exe)
        self.assertEqual(exe, model.exe_file)

    def test_model_good_no_source(self):
        exe = os.path.join(datafiles_path, 'bernoulli' + EXTENSION)
        model = Model(exe_file=exe)
        self.assertEqual(exe, model.exe_file)
        self.assertEqual('bernoulli', model.name)

        with self.assertRaises(RuntimeError):
            model.code()
        with self.assertRaises(RuntimeError):
            model.compile()

    def test_model_none(self):
        with self.assertRaises(ValueError):
            _ = Model(exe_file=None, stan_file=None)

    def test_model_bad(self):
        with self.assertRaises(Exception):
            model = Model(stan_file='xdlfkjx', exe_file='sdfndjsds')

        stan = os.path.join(datafiles_path, 'b')
        with self.assertRaises(Exception):
            model = Model(stan_file=stan)

    def test_repr(self):
        stan = os.path.join(datafiles_path, 'bernoulli.stan')
        model = Model(stan_file=stan)
        s = repr(model)
        self.assertIn('name=bernoulli', s)

    def test_print(self):
        stan = os.path.join(datafiles_path, 'bernoulli.stan')
        model = Model(stan_file=stan)
        self.assertEqual(code, model.code())

    def test_model_compile(self):
        stan = os.path.join(datafiles_path, 'bernoulli.stan')
        exe = os.path.join(datafiles_path, 'bernoulli' + EXTENSION)
        model = Model(stan_file=stan)
        self.assertEqual(None, model.exe_file)
        model.compile()
        self.assertTrue(model.exe_file.endswith(exe.replace('\\', '/')))

        model = Model(stan_file=stan)
        if os.path.exists(exe):
            os.remove(exe)
        model.compile()
        self.assertTrue(model.exe_file.endswith(exe.replace('\\', '/')))

        stan = os.path.join(datafiles_path, 'bernoulli_include.stan')
        exe = os.path.join(datafiles_path, 'bernoulli_include' + EXTENSION)
        here = os.path.dirname(os.path.abspath(__file__))
        datafiles_abspath = os.path.join(here, 'data')
        include_paths = [datafiles_abspath]
        if os.path.exists(exe):
            os.remove(exe)
        model = Model(stan_file=stan)
        model.compile(include_paths=include_paths)
        self.assertEqual(stan, model.stan_file)
        self.assertTrue(model.exe_file.endswith(exe.replace('\\', '/')))

    # TODO: test compile with existing exe - timestamp on exe unchanged
    # TODO: test overwrite with existing exe - timestamp on exe updated

class OptimizeTest(unittest.TestCase):
    def test_optimize_works(self):
        exe = os.path.join(datafiles_path, 'bernoulli' + EXTENSION)
        stan = os.path.join(datafiles_path, 'bernoulli.stan')
        model = Model(stan_file=stan, exe_file=exe)
        jdata = os.path.join(datafiles_path, 'bernoulli.data.json')
        jinit = os.path.join(datafiles_path, 'bernoulli.init.json')
        fit = model.optimize(
            data=jdata,
            seed=1239812093,
            inits=jinit,
            algorithm='BFGS',
            init_alpha=0.001,
            iter=100,
        )

        # check if calling sample related stuff fails
        with self.assertRaises(RuntimeError):
            fit.summary()
        with self.assertRaises(RuntimeError):
            _ = fit.sample
        with self.assertRaises(RuntimeError):
            fit.diagnose()

        # test numpy output
        self.assertAlmostEqual(fit.optimized_params_np[0], -5, places=2)
        self.assertAlmostEqual(fit.optimized_params_np[1], 0.2, places=3)

        # test pandas output
        self.assertEqual(
            fit.optimized_params_np[0], fit.optimized_params_pd['lp__'][0]
        )
        self.assertEqual(
            fit.optimized_params_np[1], fit.optimized_params_pd['theta'][0]
        )

        # test dict output
        self.assertEqual(
            fit.optimized_params_np[0], fit.optimized_params_dict['lp__']
        )
        self.assertEqual(
            fit.optimized_params_np[1], fit.optimized_params_dict['theta']
        )

    def test_optimize_works_dict(self):
        import json

        exe = os.path.join(datafiles_path, 'bernoulli' + EXTENSION)
        stan = os.path.join(datafiles_path, 'bernoulli.stan')
        model = Model(stan_file=stan, exe_file=exe)
        with open(os.path.join(datafiles_path, 'bernoulli.data.json')) as d:
            data = json.load(d)
        with open(os.path.join(datafiles_path, 'bernoulli.init.json')) as d:
            init = json.load(d)
        fit = model.optimize(
            data=data,
            seed=1239812093,
            inits=init,
            algorithm='BFGS',
            init_alpha=0.001,
            iter=100,
        )

        # test numpy output
        self.assertAlmostEqual(fit.optimized_params_np[0], -5, places=2)
        self.assertAlmostEqual(fit.optimized_params_np[1], 0.2, places=3)

class SampleTest(unittest.TestCase):
    def test_bernoulli_good(self):
        stan = os.path.join(datafiles_path, 'bernoulli.stan')
        exe = os.path.join(datafiles_path, 'bernoulli' + EXTENSION)
        bern_model = Model(stan_file=stan, exe_file=exe)
        bern_model.compile()

        jdata = os.path.join(datafiles_path, 'bernoulli.data.json')
        bern_fit = bern_model.sample(
            data=jdata, chains=4, cores=2, seed=12345, sampling_iters=100
        )

        for i in range(bern_fit.chains):
            csv_file = bern_fit.csv_files[i]
            txt_file = ''.join([os.path.splitext(csv_file)[0], '.txt'])
            self.assertTrue(os.path.exists(csv_file))
            self.assertTrue(os.path.exists(txt_file))

        self.assertEqual(bern_fit.chains, 4)
        self.assertEqual(bern_fit.draws, 100)
        column_names = [
            'lp__',
            'accept_stat__',
            'stepsize__',
            'treedepth__',
            'n_leapfrog__',
            'divergent__',
            'energy__',
            'theta',
        ]
        self.assertEqual(bern_fit.column_names, tuple(column_names))

        bern_sample = bern_fit.sample
        self.assertEqual(bern_sample.shape, (100, 4, len(column_names)))

        self.assertEqual(bern_fit.metric_type, 'diag_e')
        self.assertEqual(bern_fit.stepsize.shape, (4,))
        self.assertEqual(bern_fit.metric.shape, (4, 1))

        output = os.path.join(datafiles_path, 'test1-bernoulli-output')
        bern_fit = bern_model.sample(
            data=jdata,
            chains=4,
            cores=2,
            seed=12345,
            sampling_iters=100,
            csv_basename=output,
        )
        for i in range(bern_fit.chains):
            csv_file = bern_fit.csv_files[i]
            txt_file = ''.join([os.path.splitext(csv_file)[0], '.txt'])
            self.assertTrue(os.path.exists(csv_file))
            self.assertTrue(os.path.exists(txt_file))
        bern_sample = bern_fit.sample
        self.assertEqual(bern_sample.shape, (100, 4, len(column_names)))
        for i in range(bern_fit.chains):  # cleanup datafile_path dir
            os.remove(bern_fit.csv_files[i])
            os.remove(bern_fit.console_files[i])

        rdata = os.path.join(datafiles_path, 'bernoulli.data.R')
        bern_fit = bern_model.sample(
            data=rdata, chains=4, cores=2, seed=12345, sampling_iters=100
        )
        bern_sample = bern_fit.sample
        self.assertEqual(bern_sample.shape, (100, 4, len(column_names)))

        data_dict = {'N': 10, 'y': [0, 1, 0, 0, 0, 0, 0, 0, 0, 1]}
        bern_fit = bern_model.sample(
            data=data_dict, chains=4, cores=2, seed=12345, sampling_iters=100
        )
        bern_sample = bern_fit.sample
        self.assertEqual(bern_sample.shape, (100, 4, len(column_names)))

        # check if  optimized_params_np returns first draw
        # (actually first row from csv)
        np.testing.assert_equal(
            bern_fit.get_drawset().iloc[0].values, bern_fit.optimized_params_np
        )

    def test_bernoulli_bad(self):
        stan = os.path.join(datafiles_path, 'bernoulli.stan')
        exe = os.path.join(datafiles_path, 'bernoulli' + EXTENSION)
        bern_model = Model(stan_file=stan, exe_file=exe)
        bern_model.compile()

        with self.assertRaisesRegex(Exception, 'Error during sampling'):
            bern_fit = bern_model.sample(
                chains=4, cores=2, seed=12345, sampling_iters=100
            )

class GenerateQuantitiesTest(unittest.TestCase):
    def test_gen_quantities_good(self):
        stan = os.path.join(datafiles_path, 'bernoulli_ppc.stan')
        model = Model(stan_file=stan)
        model.compile()

        jdata = os.path.join(datafiles_path, 'bernoulli.data.json')

        # synthesize stanfit object -
        # see test_stanfit.py, method 'test_validate_good_run'
        goodfiles_path = os.path.join(datafiles_path, 'runset-good')
        output = os.path.join(goodfiles_path, 'bern')
        sampler_args = SamplerArgs(
            sampling_iters=100, max_treedepth=11, adapt_delta=0.95
        )
        cmdstan_args = CmdStanArgs(
            model_name=model.name,
            model_exe=model.exe_file,
            chain_ids=[1, 2, 3, 4],
            seed=12345,
            data=jdata,
            output_basename=output,
            method_args=sampler_args,
        )
        sampler_fit = StanFit(args=cmdstan_args, chains=4)
        for i in range(4):
            sampler_fit._set_retcode(i, 0)

        bern_fit = model.run_generated_quantities(
            csv_files=sampler_fit.csv_files,
            data=jdata)

        # check results - ouput files, quantities of interest, draws
        self.assertEqual(bern_fit.chains, 4)
        for i in range(4):
            self.assertEqual(bern_fit._retcodes[i], 0)
            csv_file = bern_fit.csv_files[i]
            self.assertTrue(os.path.exists(csv_file))
        column_names = [
            'y_rep.1',
            'y_rep.2',
            'y_rep.3',
            'y_rep.4',
            'y_rep.5',
            'y_rep.6',
            'y_rep.7',
            'y_rep.8',
            'y_rep.9',
            'y_rep.10'
        ]
        self.assertEqual(bern_fit.column_names, tuple(column_names))
        self.assertEqual(bern_fit.draws, 100)

if __name__ == '__main__':
    unittest.main()


import numpy as np
import pandas as pd

np.random.seed(20180917)
dates = pd.date_range('20180917', periods=6)
df = pd.DataFrame(np.random.rand(6,4), index=dates, columns=list('ABCD'))

print(df)
print(df.head())
print(df.tail(3))

print("--"*35)
print(df.index) # 인덱스 확인
print(df.columns) # 컬럼 확인
print(df.values) # 안에 있는 데이터 확인

print("--"*35)
print(df.describe()) # 간단한 통계정보를 보여준다.

print("--"*35)
print(df.T) # index와 column을 바꾼 형태

print("--"*35)
print(df.sort_index(axis=1, ascending=False)) #  컬럼에 대해 내림차순으로 정렬
print(df.sort_values(by='B')) #  DataFrame 내부에 있는 값으로 정렬

print("--"*35)
print(df['A']) # A라는 이름을 가진 컬럼의 데이터를 가져온다
print(df.A) # 위 방법과 같음

print("--"*35)
print(df[0:3]) #맨 처음 세개의 행을 가지고 온다.
print(df['20180917':'20180920']) # 인덱스명을 사용해서 해당하는 범위의데이터를 가지고온다.
print("--"*16, "주의","--"*16)
# print(df['20180917']) # 오류
print(df['20180917':'20180917']) # 이런 형태로 입력

# 특정 인덱스 값 가지고 오기 loc, at
print("--"*35)
print(df.loc[dates[0]])
print(df.loc['20180917'])
print(df.loc['2018-09-17'])
print(df.loc[:,['A','B']])
print(df.loc[dates[0],['A','B']])
print(df.at[dates[0],'A'])

print("--"*35)
print(df.iloc[3])
print(df.iloc[2:5, 0:2])

print("--"*35)
print(df.iloc[[1,2,4],[0,2]])
print(df.iloc[1:3,:])
print(df.iloc[:,1:3])
print(df.iloc[1,1])
print(df.iat[1,1])

print("--"*35)
print(df[df.A > 0]) # 양수 값만 출력
print(df[df > 0])# 전체 출력 양수 이외의 값은 NaNㅇ로 표현

print("--"*35)
df2 = df.copy()
df2['E'] = ['one', 'one','two', 'three', 'four', 'three']
print(df2)
# 필터링을 해야하는 경우 isin를 사용
print(df2[df2['E'].isin(['two','four'])])

print("--"*35)
df.loc[:, 'D'] = np.array([5] * len(df)) # len(df) -= 6
# 우변의 값은 [5,5,5,5,5,5]
print(df)

print("--"*35)
df2 = df.copy()
df2[df2 > 0] = -df2
print(df2)


# Copyright 2019 Cloudera, Inc.
#
# 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.

# # Joining pandas DataFrames

# Import modules
import numpy as np
import pandas as pd

# Read employees data
employees = pd.read_table('data/employees/employees.txt')
employees

# Read offices data
offices = pd.read_table("data/offices/offices.txt")
offices

# The DataFrame method `merge` can be used to join two
# DataFrames together. The parameter `how` specifies
# what type of join to perform:
# - `how='inner'` for inner joins
# - `how='left'` for left outer joins
# - `how='right'` for right outer joins
# - `how='outer'` for full outer joins

# For example, use `merge` with `how='left'` to perform
# a left outer join on the `employees` and `offices`
# DataFrames
employees.merge(offices, how='left')

# pandas automatically identifies common column names in
# the two DataFrames and joins on them. To manually
# specify the join key columns, use the `on` parameter
employees.merge(offices, how='left', on='office_id')

# For more details, see the documentation for
# [`pandas.DataFrame.merge`](https://pandas.pydata.org/pandas-docs/stable/generated/pandas.DataFrame.merge.html).


"""Setting seed for reproducibility"""

import os
import random
import torch
import numpy as np

# Results can be software/hardware-dependent
# Exactly reproduciable results are expected only on the same software and hardware
def set_seed(seed=1000):
    """Sets the seed for generating random numbers to get (as) reproducible (as possible) results.

    The CuDNN options are set according to the official PyTorch guidance on reproducibility: https://pytorch.org/docs/stable/notes/randomness.html.
    Another reference is https://discuss.pytorch.org/t/difference-between-torch-manual-seed-and-torch-cuda-manual-seed/13848/6

    Args:
        seed (int, optional): The desired seed. Defaults to 1000.
    """
    # 1. Set `PYTHONHASHSEED` environment variable at a fixed value
    os.environ['PYTHONHASHSEED']=str(seed)
    # 2. Set `python` built-in pseudo-random generator at a fixed value
    random.seed(seed)
    # 3. Set `numpy` pseudo-random generator at a fixed value
    np.random.seed(seed)
    # 4. Set `pytorch` pseudo-random generator at a fixed value
    torch.manual_seed(seed)
    # if torch.cuda.is_available():
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False


# %%
"""
# Using imgaug with more Control Flow

The standard form of using `imgaug` is in a deferred way, i.e. you first "build" your augmentation sequence and then apply it many times to augment data. `imgaug` then handles the augmentation almost entirely on its own. This form of using the library is similar to e.g. tensorflow. In the case of `imgaug` it has a few advantages:
* It allows the library to randomize the order in which augmenters are applied (as is e.g. the case when calling a `Sequential(..., random_order=True)` instance). This greatly increases the space of possible augmentations. Implementing such a random order oneself will likely end up with yet another architecture that lacks control flow.
* It allows the library to easily randomize which augmenters are applied or not applied to an input, as is handled by e.g. `Sometimes(...)` or `SomeOf(...)`. This is easier to implement oneself than random order, but still leads to quite some repeated code between projects.
* It pushes random state handling to the library, decreasing the probability of bugs when re-applying augmentations to different inputs.
* It allows to have one entry point to run (tested) multicore augmentation. In the case of `imgaug` the method `Augmenter.pool()` is such an entry point. It starts a wrapper around `multiprocessing.Pool` that is optimized for augmentation (e.g. uses different random states between child processes to guarantee that each CPU core generates different augmentations).

It does however also have disadvantages, with the main one being that errors can be quite a bit harder to debug. If you don't need the above mentioned advantages, it is possible to execute `imgaug` in a non-deferred way, similar to e.g. pytorch. The implementation of this is very similar to the standard way in `imgaug`, except that you instantiate each augmenter on its own and later on also apply it on its own. The below code block shows an example:
"""

# %%
import numpy as np
import imgaug as ia
from imgaug import augmenters as iaa
%matplotlib inline

ia.seed(3)

class AugSequence:
    def __init__(self):
        # instantiate each augmenter and save it to its own variable
        self.affine = iaa.Affine(rotate=(-20, 20), translate_px={"x": (-10, 10), "y": (-5, 5)})
        self.multiply = iaa.Multiply((0.9, 1.1))
        self.contrast = iaa.LinearContrast((0.8, 1.2))
        self.gray = iaa.Grayscale((0.0, 1.0))

    def augment_images(self, x):
        # apply each augmenter on its own, one by one
        x = self.affine(images=x)
        x = self.multiply(images=x)
        x = self.contrast(images=x)
        x = self.gray(images=x)
        return x

aug = AugSequence()

image = ia.quokka_square(size=(256, 256))  # uint8 array of shape (256, 256, 3)
images_aug = aug.augment_images([image, image])

print("Before:")
ia.imshow(np.hstack([image, image]))
print("After:")
ia.imshow(np.hstack(images_aug))

# %%
"""
## Single Function
"""

# %%
"""
It is also possible to re-instantiate augmenters each time they are used for augmentation. In theory, this incurs a tiny performance penalty (as e.g. parameters have to be parsed each time). The below code block shows an example.
"""

# %%
ia.seed(3)

def augment_images(x):
    x = iaa.Affine(rotate=(-20, 20))(images=x)
    x = iaa.Multiply((0.9, 1.1))(images=x)
    x = iaa.LinearContrast((0.8, 1.2))(images=x)
    x = iaa.Grayscale((0.0, 1.0))(images=x)
    return x

images_aug = augment_images([image, image])

print("Before:")
ia.imshow(np.hstack([image, image]))
print("After:")
ia.imshow(np.hstack(images_aug))

# %%
"""
## Some Time Measurements
"""

# %%
"""
How long does it take to instantiate an augmenter? We can measure that, here with `Affine` and a large number of parameters as an example:
"""

# %%
%timeit -n 10000 iaa.Affine(translate_px=(-10, 10), scale=(0.9, 1.1), rotate=(-20, 20), shear=(-20, 20), mode=ia.ALL)

# %%
"""
As you can see, the required time is tiny, below 0.1ms (measured on rather outdated hardware). An augmentation sequence with 10 augmenters would take less than 1ms to instantiate, which is by far less than the augmentation process.

We can also measure the required time of both methods from further above, the one using a custom class and the single-function method:
"""

# %%
%timeit -n 2500 aug.augment_images([image, image])
%timeit -n 2500 augment_images([image, image])

# %%
"""
As expected, the difference per call is negligible.
"""

# %%
"""
## Single Function, Different Input Types

If you want to use the above single-function method to augment images *and* other datatypes, e.g. bounding boxes, you will have to manage random states yourself. This is however not very hard. The following example adds a `seed` argument to align the augmentations between images and bounding boxes:
"""

# %%
def augment_images(x, seed):
    x = iaa.Affine(translate_px=(-60, 60), random_state=seed)(images=x)
    x = iaa.Multiply((0.9, 1.1), random_state=seed)(images=x)
    x = iaa.LinearContrast((0.8, 1.2), random_state=seed)(images=x)
    x = iaa.Grayscale((0.0, 1.0), random_state=seed)(images=x)
    return x

# new function
def augment_bounding_boxes(x, seed):
    x = iaa.Affine(translate_px=(-60, 60), random_state=seed)(bounding_boxes=x)
    x = iaa.Multiply((0.9, 1.1), random_state=seed)(bounding_boxes=x)
    x = iaa.LinearContrast((0.8, 1.2), random_state=seed)(bounding_boxes=x)
    x = iaa.Grayscale((0.0, 1.0), random_state=seed)(bounding_boxes=x)
    return x

# bounding boxes to augment
bbsoi = ia.BoundingBoxesOnImage(
    bounding_boxes=[ia.BoundingBox(x1=40, y1=20, x2=230, y2=250)],
    shape=image.shape)

# augment images and bounding boxes
images_aug = augment_images([image, image], seed=2)
bbsois_aug = augment_bounding_boxes([bbsoi, bbsoi], seed=2)

print("Before:")
ia.imshow(
    np.hstack([
        bbsoi.draw_on_image(image, size=3),
        bbsoi.draw_on_image(image, size=3),
    ])
)

print("After:")
ia.imshow(
    np.hstack([
        bbsois_aug[0].draw_on_image(images_aug[0], size=3),
        bbsois_aug[1].draw_on_image(images_aug[1], size=3)
    ])
)

scripts/C02 - Using imgaug with more Control Flow.py


# Copyright 2021 Huawei Technologies Co., Ltd
#
# 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.
# ============================================================================

import pytest
import numpy as np
from mindspore import Tensor
from mindspore.ops import operations as P
from mindspore.common.api import ms_function
from mindspore.common.initializer import initializer
from mindspore.common.parameter import Parameter
import mindspore.nn as nn
import mindspore.context as context

context.set_context(device_target='CPU')

class Transpose(nn.Cell):
    def __init__(self):
        super(Transpose, self).__init__()
        self.transpose = P.Transpose()

        self.x_2D = Parameter(initializer(Tensor(np.arange(5 * 6).reshape(5, 6).astype(np.float32)), [5, 6]),
                              name='x_2D')
        self.perm_2D = (1, 0)

        self.x_3D = Parameter(initializer(Tensor(np.arange(2 * 2 * 4).reshape(2, 2, 4).astype(np.float32)), [2, 2, 4]),
                              name='x_3D')
        self.perm_3D = (1, 0, 2)

        self.x_4D = Parameter(
            initializer(Tensor(np.arange(2 * 3 * 4 * 5).reshape(2,
                                                                3, 4, 5).astype(np.float32)), [2, 3, 4, 5]),
            name='x_4D')
        self.perm_4D = (0, 1, 2, 3)

        self.x_5D = Parameter(
            initializer(Tensor(np.arange(1 * 2 * 3 * 4 * 5).reshape(1, 2, 3, 4, 5).astype(np.float32)),
                        [1, 2, 3, 4, 5]), name='x_5D')
        self.perm_5D = (1, 0, 3, 4, 2)

    @ms_function
    def construct(self):
        return (self.transpose(self.x_2D, self.perm_2D), self.transpose(self.x_3D, self.perm_3D),
                self.transpose(self.x_4D, self.perm_4D), self.transpose(self.x_5D, self.perm_5D))

@pytest.mark.level0
@pytest.mark.platform_x86_cpu
@pytest.mark.env_onecard
def test_transpose():
    transpose = Transpose()
    output = transpose()

    expect0 = np.array([[[0, 6, 12, 18, 24],
                         [1, 7, 13, 19, 25],
                         [2, 8, 14, 20, 26],
                         [3, 9, 15, 21, 27],
                         [4, 10, 16, 22, 28],
                         [5, 11, 17, 23, 29]]]).astype(np.float32)
    expect1 = np.array([[[[0, 1, 2, 3],
                          [8, 9, 10, 11]],
                         [[4, 5, 6, 7],
                          [12, 13, 14, 15]]]]).astype(np.float32)
    expect2 = np.array([[[[[0, 1, 2, 3, 4],
                           [5, 6, 7, 8, 9],
                           [10, 11, 12, 13, 14],
                           [15, 16, 17, 18, 19]],
                          [[20, 21, 22, 23, 24],
                           [25, 26, 27, 28, 29],
                           [30, 31, 32, 33, 34],
                           [35, 36, 37, 38, 39]],
                          [[40, 41, 42, 43, 44],
                           [45, 46, 47, 48, 49],
                           [50, 51, 52, 53, 54],
                           [55, 56, 57, 58, 59]]],

                         [[[60, 61, 62, 63, 64],
                           [65, 66, 67, 68, 69],
                           [70, 71, 72, 73, 74],
                           [75, 76, 77, 78, 79]],
                          [[80, 81, 82, 83, 84],
                           [85, 86, 87, 88, 89],
                           [90, 91, 92, 93, 94],
                           [95, 96, 97, 98, 99]],
                          [[100, 101, 102, 103, 104],
                           [105, 106, 107, 108, 109],
                           [110, 111, 112, 113, 114],
                           [115, 116, 117, 118, 119]]]]]).astype(np.float32)
    expect3 = np.array([[[[[[0, 20, 40],
                            [1, 21, 41],
                            [2, 22, 42],
                            [3, 23, 43],
                            [4, 24, 44]],
                           [[5, 25, 45],
                            [6, 26, 46],
                            [7, 27, 47],
                            [8, 28, 48],
                            [9, 29, 49]],
                           [[10, 30, 50],
                            [11, 31, 51],
                            [12, 32, 52],
                            [13, 33, 53],
                            [14, 34, 54]],
                           [[15, 35, 55],
                            [16, 36, 56],
                            [17, 37, 57],
                            [18, 38, 58],
                            [19, 39, 59]]]],

                         [[[[60, 80, 100],
                            [61, 81, 101],
                            [62, 82, 102],
                            [63, 83, 103],
                            [64, 84, 104]],
                           [[65, 85, 105],
                            [66, 86, 106],
                            [67, 87, 107],
                            [68, 88, 108],
                            [69, 89, 109]],
                           [[70, 90, 110],
                            [71, 91, 111],
                            [72, 92, 112],
                            [73, 93, 113],
                            [74, 94, 114]],
                           [[75, 95, 115],
                            [76, 96, 116],
                            [77, 97, 117],
                            [78, 98, 118],
                            [79, 99, 119]]]]]]).astype(np.float32)
    assert (output[0].asnumpy() == expect0).all()
    assert (output[1].asnumpy() == expect1).all()
    assert (output[2].asnumpy() == expect2).all()
    assert (output[3].asnumpy() == expect3).all()

test_transpose()

class Transpose_int64(nn.Cell):
    def __init__(self):
        super(Transpose_int64, self).__init__()
        self.transpose = P.Transpose()

        self.x_2D = Parameter(initializer(Tensor(np.arange(5 * 6).reshape(5, 6).astype(np.int64)), [5, 6]),
                              name='x_2D')
        self.perm_2D = (1, 0)

        self.x_3D = Parameter(initializer(Tensor(np.arange(2 * 2 * 4).reshape(2, 2, 4).astype(np.int64)), [2, 2, 4]),
                              name='x_3D')
        self.perm_3D = (1, 0, 2)

        self.x_4D = Parameter(
            initializer(Tensor(np.arange(2 * 3 * 4 * 5).reshape(2,
                                                                3, 4, 5).astype(np.int64)), [2, 3, 4, 5]),
            name='x_4D')
        self.perm_4D = (0, 1, 2, 3)

        self.x_5D = Parameter(
            initializer(Tensor(np.arange(1 * 2 * 3 * 4 * 5).reshape(1, 2, 3, 4, 5).astype(np.int64)),
                        [1, 2, 3, 4, 5]), name='x_5D')
        self.perm_5D = (1, 0, 3, 4, 2)

    @ms_function
    def construct(self):
        return (self.transpose(self.x_2D, self.perm_2D), self.transpose(self.x_3D, self.perm_3D),
                self.transpose(self.x_4D, self.perm_4D), self.transpose(self.x_5D, self.perm_5D))

@pytest.mark.level0
@pytest.mark.platform_x86_cpu
@pytest.mark.env_onecard
def test_transpose_int64():
    transpose = Transpose_int64()
    output = transpose()

    expect0 = np.array([[[0, 6, 12, 18, 24],
                         [1, 7, 13, 19, 25],
                         [2, 8, 14, 20, 26],
                         [3, 9, 15, 21, 27],
                         [4, 10, 16, 22, 28],
                         [5, 11, 17, 23, 29]]]).astype(np.int64)
    expect1 = np.array([[[[0, 1, 2, 3],
                          [8, 9, 10, 11]],
                         [[4, 5, 6, 7],
                          [12, 13, 14, 15]]]]).astype(np.int64)
    expect2 = np.array([[[[[0, 1, 2, 3, 4],
                           [5, 6, 7, 8, 9],
                           [10, 11, 12, 13, 14],
                           [15, 16, 17, 18, 19]],
                          [[20, 21, 22, 23, 24],
                           [25, 26, 27, 28, 29],
                           [30, 31, 32, 33, 34],
                           [35, 36, 37, 38, 39]],
                          [[40, 41, 42, 43, 44],
                           [45, 46, 47, 48, 49],
                           [50, 51, 52, 53, 54],
                           [55, 56, 57, 58, 59]]],

                         [[[60, 61, 62, 63, 64],
                           [65, 66, 67, 68, 69],
                           [70, 71, 72, 73, 74],
                           [75, 76, 77, 78, 79]],
                          [[80, 81, 82, 83, 84],
                           [85, 86, 87, 88, 89],
                           [90, 91, 92, 93, 94],
                           [95, 96, 97, 98, 99]],
                          [[100, 101, 102, 103, 104],
                           [105, 106, 107, 108, 109],
                           [110, 111, 112, 113, 114],
                           [115, 116, 117, 118, 119]]]]]).astype(np.int64)
    expect3 = np.array([[[[[[0, 20, 40],
                            [1, 21, 41],
                            [2, 22, 42],
                            [3, 23, 43],
                            [4, 24, 44]],
                           [[5, 25, 45],
                            [6, 26, 46],
                            [7, 27, 47],
                            [8, 28, 48],
                            [9, 29, 49]],
                           [[10, 30, 50],
                            [11, 31, 51],
                            [12, 32, 52],
                            [13, 33, 53],
                            [14, 34, 54]],
                           [[15, 35, 55],
                            [16, 36, 56],
                            [17, 37, 57],
                            [18, 38, 58],
                            [19, 39, 59]]]],

                         [[[[60, 80, 100],
                            [61, 81, 101],
                            [62, 82, 102],
                            [63, 83, 103],
                            [64, 84, 104]],
                           [[65, 85, 105],
                            [66, 86, 106],
                            [67, 87, 107],
                            [68, 88, 108],
                            [69, 89, 109]],
                           [[70, 90, 110],
                            [71, 91, 111],
                            [72, 92, 112],
                            [73, 93, 113],
                            [74, 94, 114]],
                           [[75, 95, 115],
                            [76, 96, 116],
                            [77, 97, 117],
                            [78, 98, 118],
                            [79, 99, 119]]]]]]).astype(np.int64)
    assert (output[0].asnumpy() == expect0).all()
    assert (output[1].asnumpy() == expect1).all()
    assert (output[2].asnumpy() == expect2).all()
    assert (output[3].asnumpy() == expect3).all()

test_transpose_int64()

class Transpose_uint8(nn.Cell):
    def __init__(self):
        super(Transpose_uint8, self).__init__()
        self.transpose = P.Transpose()

        self.x_2D = Parameter(initializer(Tensor(np.arange(5 * 6).reshape(5, 6).astype(np.uint8)), [5, 6]),
                              name='x_2D')
        self.perm_2D = (1, 0)

        self.x_3D = Parameter(initializer(Tensor(np.arange(2 * 2 * 4).reshape(2, 2, 4).astype(np.uint8)), [2, 2, 4]),
                              name='x_3D')
        self.perm_3D = (1, 0, 2)

        self.x_4D = Parameter(
            initializer(Tensor(np.arange(2 * 3 * 4 * 5).reshape(2,
                                                                3, 4, 5).astype(np.uint8)), [2, 3, 4, 5]),
            name='x_4D')
        self.perm_4D = (0, 1, 2, 3)

        self.x_5D = Parameter(
            initializer(Tensor(np.arange(1 * 2 * 3 * 4 * 5).reshape(1, 2, 3, 4, 5).astype(np.uint8)),
                        [1, 2, 3, 4, 5]), name='x_5D')
        self.perm_5D = (1, 0, 3, 4, 2)

    @ms_function
    def construct(self):
        return (self.transpose(self.x_2D, self.perm_2D), self.transpose(self.x_3D, self.perm_3D),
                self.transpose(self.x_4D, self.perm_4D), self.transpose(self.x_5D, self.perm_5D))

@pytest.mark.level0
@pytest.mark.platform_x86_cpu
@pytest.mark.env_onecard
def test_transpose_uint8():
    transpose = Transpose_uint8()
    output = transpose()

    expect0 = np.array([[[0, 6, 12, 18, 24],
                         [1, 7, 13, 19, 25],
                         [2, 8, 14, 20, 26],
                         [3, 9, 15, 21, 27],
                         [4, 10, 16, 22, 28],
                         [5, 11, 17, 23, 29]]]).astype(np.uint8)
    expect1 = np.array([[[[0, 1, 2, 3],
                          [8, 9, 10, 11]],
                         [[4, 5, 6, 7],
                          [12, 13, 14, 15]]]]).astype(np.uint8)
    expect2 = np.array([[[[[0, 1, 2, 3, 4],
                           [5, 6, 7, 8, 9],
                           [10, 11, 12, 13, 14],
                           [15, 16, 17, 18, 19]],
                          [[20, 21, 22, 23, 24],
                           [25, 26, 27, 28, 29],
                           [30, 31, 32, 33, 34],
                           [35, 36, 37, 38, 39]],
                          [[40, 41, 42, 43, 44],
                           [45, 46, 47, 48, 49],
                           [50, 51, 52, 53, 54],
                           [55, 56, 57, 58, 59]]],

                         [[[60, 61, 62, 63, 64],
                           [65, 66, 67, 68, 69],
                           [70, 71, 72, 73, 74],
                           [75, 76, 77, 78, 79]],
                          [[80, 81, 82, 83, 84],
                           [85, 86, 87, 88, 89],
                           [90, 91, 92, 93, 94],
                           [95, 96, 97, 98, 99]],
                          [[100, 101, 102, 103, 104],
                           [105, 106, 107, 108, 109],
                           [110, 111, 112, 113, 114],
                           [115, 116, 117, 118, 119]]]]]).astype(np.uint8)
    expect3 = np.array([[[[[[0, 20, 40],
                            [1, 21, 41],
                            [2, 22, 42],
                            [3, 23, 43],
                            [4, 24, 44]],
                           [[5, 25, 45],
                            [6, 26, 46],
                            [7, 27, 47],
                            [8, 28, 48],
                            [9, 29, 49]],
                           [[10, 30, 50],
                            [11, 31, 51],
                            [12, 32, 52],
                            [13, 33, 53],
                            [14, 34, 54]],
                           [[15, 35, 55],
                            [16, 36, 56],
                            [17, 37, 57],
                            [18, 38, 58],
                            [19, 39, 59]]]],

                         [[[[60, 80, 100],
                            [61, 81, 101],
                            [62, 82, 102],
                            [63, 83, 103],
                            [64, 84, 104]],
                           [[65, 85, 105],
                            [66, 86, 106],
                            [67, 87, 107],
                            [68, 88, 108],
                            [69, 89, 109]],
                           [[70, 90, 110],
                            [71, 91, 111],
                            [72, 92, 112],
                            [73, 93, 113],
                            [74, 94, 114]],
                           [[75, 95, 115],
                            [76, 96, 116],
                            [77, 97, 117],
                            [78, 98, 118],
                            [79, 99, 119]]]]]]).astype(np.uint8)
    assert (output[0].asnumpy() == expect0).all()
    assert (output[1].asnumpy() == expect1).all()
    assert (output[2].asnumpy() == expect2).all()
    assert (output[3].asnumpy() == expect3).all()

test_transpose_uint8()


import numpy as np
from typing import Callable, Iterable, Union

class Simulator:

    """
    A callback function to display the current value of the arguments and the
    value of the objective function at each iteration. Based on the class
    developed in https://stackoverflow.com/a/59330005/4983192

    Parameters
    ----------
    func: Callable
        The objective function to be optimized.

    Attributes
    ----------
    callback_count: int
        Number of times the callback is called.
    sol_eval: list
        A list that contains all the intermediate solutions used in the problem.
    func_eval: list
        A list that contains all the evaluations of the objective function of the problem.
    """

    def __init__(self, function: Callable):
        self.func = function
        self.callback_count = 0
        self.sol_eval = []
        self.func_eval = []

    def simulate(self, x_k: Iterable[Union[int, float]], *args) -> Union[int, float]:

        """
        Executes the actual simulation and returns the result, while updating
        the attributes `sol_eval` and `func_eval`. This must be passed to the
        optimizer without arguments or parentheses.

        Inputs
        ------
        x_k: Iterable[Union[int, float]]
            The actual value for the solution.

        Returns
        -------
        (Union[int, float]) The objective function evaluated at x_k.
        """

        result = self.func(x_k, *args)
        self.sol_eval.append(x_k)
        self.func_eval.append(result)
        return result

    def callback(self, x_k: Iterable[Union[int, float]], *_) -> None:
        """
        Callback function that can be used by optimizers of scipy.optimize.
        The third argument "*_" makes sure that it still works when the
        optimizer calls the callback function with more than one argument. Pass
        to optimizer without arguments or parentheses.

        Inputs
        ------
        x_k: list
            The actual value for the solution.
        """

        # Locate the position in sol_eval that coincides with x_k. Once this
        # position is found, break the loop to save the position.
        for i, x in reversed(list(enumerate(self.sol_eval))):
            if np.allclose(x, x_k):
                break

        sp_info = ""
        # For each component in the actual solution, generate an string containing its
        # value and finally the value of the objective function stored in func_eval
        for comp in x_k:
            sp_info += f"{comp:10.5e}\t"
        sp_info += f"{self.func_eval[i]:10.5e}"

        # Set the title of the callback, with the smoothing parameter names and the
        # objective function label
        if not self.callback_count:
            title_list = [f"sp{j+1}" for j, _ in enumerate(x_k)] + ["Objective"]
            print("Starting the optimization algorithm")
            print(*title_list, sep="\t\t")
        # Print the actual solution and the actual objective function value
        print(sp_info)
        self.callback_count += 1
        return None


import unittest

import numpy as np

import functions

class TestSquare(unittest.TestCase):

    def check_square(self, x, y_desired):
        y_actual = functions.square(x)
        np.testing.assert_equal(y_desired, y_actual)

    def test_scalar(self):
        x = 1
        y_desired = 1
        self.check_square(x, y_desired)

    def test_ndarray(self):
        x = np.array([0, 1, 2, 3], dtype=np.float32)
        y_desired = np.array([0, 1, 4, 9], dtype=np.float32)
        self.check_square(x, y_desired)

    def test_invalid_input(self):
        with self.assertRaises(TypeError):
            functions.square('a')

if __name__ == '__main__':
    unittest.main()