jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def project_simplex(v, s=1):
    """
    Taken from https://gist.github.com/daien/1272551/edd95a6154106f8e28209a1c7964623ef8397246.

    Compute the Euclidean projection on a positive simplex
    Solves the optimisation problem (using the algorithm from [1]):
        min_w 0.5 * || w - v ||_2^2 , s.t. \sum_i w_i = s, w_i >= 0
    Parameters
    ----------
    v: (n,) numpy array,
       n-dimensional vector to project
    s: int, optional, default: 1,
       radius of the simplex
    Returns
    -------
    w: (n,) numpy array,
       Euclidean projection of v on the simplex
    Notes
    -----
    The complexity of this algorithm is in O(n log(n)) as it involves sorting v.
    Better alternatives exist for high-dimensional sparse vectors (cf. [1])
    However, this implementation still easily scales to millions of dimensions.
    References
    ----------
    [1] Efficient Projections onto the .1-Ball for Learning in High Dimensions
        John Duchi, Shai Shalev-Shwartz, Yoram Singer, and Tushar Chandra.
        International Conference on Machine Learning (ICML 2008)
        http://www.cs.berkeley.edu/~jduchi/projects/DuchiSiShCh08.pdf
    """

    assert s > 0, "Radius s must be strictly positive (%d <= 0)" % s

    n, = v.shape  # will raise ValueError if v is not 1-D
    # check if we are already on the simplex
    if v.sum() == s and numpy.alltrue(v >= 0):
        # best projection: itself!
        return v
    # get the array of cumulative sums of a sorted (decreasing) copy of v
    u = numpy.sort(v)[::-1]
    cssv = numpy.cumsum(u)
    # get the number of > 0 components of the optimal solution
    rho = numpy.nonzero(u * numpy.arange(1, n+1) > (cssv - s))[0][-1]
    # compute the Lagrange multiplier associated to the simplex constraint
    theta = float(cssv[rho] - s) / rho
    # compute the projection by thresholding v using theta
    w = (v - theta).clip(min=0)
    return w
//  [python]
def project_simplex(v, s=1):
    """
    Taken from https://gist.github.com/daien/1272551/edd95a6154106f8e28209a1c7964623ef8397246.

    Compute the Euclidean projection on a positive simplex
    Solves the optimisation problem (using the algorithm from [1]):
        min_w 0.5 * || w - v ||_2^2 , s.t. \sum_i w_i = s, w_i >= 0
    Parameters
    ----------
    v: (n,) numpy array,
       n-dimensional vector to project
    s: int, optional, default: 1,
       radius of the simplex
    Returns
    -------
    w: (n,) numpy array,
       Euclidean projection of v on the simplex
    Notes
    -----
    The complexity of this algorithm is in O(n log(n)) as it involves sorting v.
    Better alternatives exist for high-dimensional sparse vectors (cf. [1])
    However, this implementation still easily scales to millions of dimensions.
    References
    ----------
    [1] Efficient Projections onto the .1-Ball for Learning in High Dimensions
        John Duchi, Shai Shalev-Shwartz, Yoram Singer, and Tushar Chandra.
        International Conference on Machine Learning (ICML 2008)
        http://www.cs.berkeley.edu/~jduchi/projects/DuchiSiShCh08.pdf
    """

    assert s > 0, "Radius s must be strictly positive (%d <= 0)" % s

    n, = v.shape  # will raise ValueError if v is not 1-D
    # check if we are already on the simplex
    if v.sum() == s and numpy.all(v >= 0):
        # best projection: itself!
        return v
    # get the array of cumulative sums of a sorted (decreasing) copy of v
    u = numpy.sort(v)[::-1]
    cssv = numpy.cumsum(u)
    # get the number of > 0 components of the optimal solution
    rho = numpy.nonzero(u * numpy.arange(1, n+1) > (cssv - s))[0][-1]
    # compute the Lagrange multiplier associated to the simplex constraint
    theta = float(cssv[rho] - s) / rho
    # compute the projection by thresholding v using theta
    w = (v - theta).clip(min=0)
    return w
//  [python]
def test02b_write(self):
        """Test for point-selections (write, integer keys)."""

        nparr = self.nparr
        tbarr = self.tbarr
        for value1, value2 in self.limits:
            key = numpy.where((nparr >= value1) & (nparr < value2))
            if common.verbose:
                print("Selection to test:", key)
            s = nparr[key]
            nparr[key] = s * 2
            tbarr[key] = s * 2
            a = nparr[:]
            b = tbarr[:]
            self.assertTrue(
                numpy.alltrue(a == b),
                "NumPy array and PyTables modifications does not match.")
//  [python]
def test02b_write(self):
        """Test for point-selections (write, integer keys)."""

        nparr = self.nparr
        tbarr = self.tbarr
        for value1, value2 in self.limits:
            key = numpy.where((nparr >= value1) & (nparr < value2))
            if common.verbose:
                print("Selection to test:", key)
            s = nparr[key]
            nparr[key] = s * 2
            tbarr[key] = s * 2
            a = nparr[:]
            b = tbarr[:]
            self.assertTrue(
                numpy.all(a == b),
                "NumPy array and PyTables modifications does not match.")
//  [python]
def test01a_read(self):
        """Test for fancy-selections (working selections, read)."""

        nparr = self.nparr
        tbarr = self.tbarr
        for key in self.working_keyset:
            if common.verbose:
                print("Selection to test:", key)
            a = nparr[key]
            b = tbarr[key]
            self.assertTrue(
                numpy.alltrue(a == b),
                "NumPy array and PyTables selections does not match.")
//  [python]
def test01a_read(self):
        """Test for fancy-selections (working selections, read)."""

        nparr = self.nparr
        tbarr = self.tbarr
        for key in self.working_keyset:
            if common.verbose:
                print("Selection to test:", key)
            a = nparr[key]
            b = tbarr[key]
            self.assertTrue(
                numpy.all(a == b),
                "NumPy array and PyTables selections does not match.")
//  [python]
def test02c_write(self):
        """Test for point-selections (write, integer values, broadcast)."""

        nparr = self.nparr
        tbarr = self.tbarr
        for value1, value2 in self.limits:
            key = numpy.where((nparr >= value1) & (nparr < value2))
            if common.verbose:
                print("Selection to test:", key)
            # s = nparr[key]
            nparr[key] = 2   # force a broadcast
            tbarr[key] = 2   # force a broadcast
            a = nparr[:]
            b = tbarr[:]
            self.assertTrue(
                numpy.alltrue(a == b),
                "NumPy array and PyTables modifications does not match.")
//  [python]
def test02c_write(self):
        """Test for point-selections (write, integer values, broadcast)."""

        nparr = self.nparr
        tbarr = self.tbarr
        for value1, value2 in self.limits:
            key = numpy.where((nparr >= value1) & (nparr < value2))
            if common.verbose:
                print("Selection to test:", key)
            # s = nparr[key]
            nparr[key] = 2   # force a broadcast
            tbarr[key] = 2   # force a broadcast
            a = nparr[:]
            b = tbarr[:]
            self.assertTrue(
                numpy.all(a == b),
                "NumPy array and PyTables modifications does not match.")
//  [python]
def test02a_write(self):
        """Test for fancy-selections (working selections, write)."""

        nparr = self.nparr
        tbarr = self.tbarr
        for key in self.working_keyset:
            if common.verbose:
                print("Selection to test:", key)
            s = nparr[key]
            nparr[key] = s * 2
            tbarr[key] = s * 2
            a = nparr[:]
            b = tbarr[:]
            self.assertTrue(
                numpy.alltrue(a == b),
                "NumPy array and PyTables modifications does not match.")
//  [python]
def test02a_write(self):
        """Test for fancy-selections (working selections, write)."""

        nparr = self.nparr
        tbarr = self.tbarr
        for key in self.working_keyset:
            if common.verbose:
                print("Selection to test:", key)
            s = nparr[key]
            nparr[key] = s * 2
            tbarr[key] = s * 2
            a = nparr[:]
            b = tbarr[:]
            self.assertTrue(
                numpy.all(a == b),
                "NumPy array and PyTables modifications does not match.")
//  [python]
def test02b_write(self):
        """Test for fancy-selections (working selections, write, broadcast)."""

        nparr = self.nparr
        tbarr = self.tbarr
        for key in self.working_keyset:
            if common.verbose:
                print("Selection to test:", key)
            # s = nparr[key]
            nparr[key] = 2   # broadcast value
            tbarr[key] = 2   # broadcast value
            a = nparr[:]
            b = tbarr[:]
#             if common.verbose:
#                 print("NumPy modified array:", a)
#                 print("PyTables modifyied array:", b)
            self.assertTrue(
                numpy.alltrue(a == b),
                "NumPy array and PyTables modifications does not match.")
//  [python]
def test02b_write(self):
        """Test for fancy-selections (working selections, write, broadcast)."""

        nparr = self.nparr
        tbarr = self.tbarr
        for key in self.working_keyset:
            if common.verbose:
                print("Selection to test:", key)
            # s = nparr[key]
            nparr[key] = 2   # broadcast value
            tbarr[key] = 2   # broadcast value
            a = nparr[:]
            b = tbarr[:]
#             if common.verbose:
#                 print("NumPy modified array:", a)
#                 print("PyTables modifyied array:", b)
            self.assertTrue(
                numpy.all(a == b),
                "NumPy array and PyTables modifications does not match.")
//  [python]
def test_SobolGeneratorAllTunable(self) -> None:
        generator = SobolGenerator(seed=0)
        bounds = self._create_bounds(n_tunable=3, n_fixed=0)
        generated_points, weights = generator.gen(
            n=3, bounds=bounds, rounding_func=lambda x: x
        )
        self.assertEqual(np.shape(generated_points), (3, 3))
        np_bounds = np.array(bounds)
        self.assertTrue(np.alltrue(generated_points >= np_bounds[:, 0]))
        self.assertTrue(np.alltrue(generated_points <= np_bounds[:, 1]))
        self.assertTrue(np.all(weights == 1.0))
        self.assertEqual(generator._get_state().get("init_position"), 3)
//  [python]
def test_SobolGeneratorAllTunable(self) -> None:
        generator = SobolGenerator(seed=0)
        bounds = self._create_bounds(n_tunable=3, n_fixed=0)
        generated_points, weights = generator.gen(
            n=3, bounds=bounds, rounding_func=lambda x: x
        )
        self.assertEqual(np.shape(generated_points), (3, 3))
        np_bounds = np.array(bounds)
        self.assertTrue(np.all(generated_points >= np_bounds[:, 0]))
        self.assertTrue(np.all(generated_points <= np_bounds[:, 1]))
        self.assertTrue(np.all(weights == 1.0))
        self.assertEqual(generator._get_state().get("init_position"), 3)
//  [python]
def test_SobolGeneratorNoScramble(self) -> None:
        generator = SobolGenerator(scramble=False)
        n_tunable = fixed_param_index = 3
        bounds = self._create_bounds(n_tunable=n_tunable, n_fixed=1)
        generated_points, weights = generator.gen(
            n=3,
            bounds=bounds,
            fixed_features={fixed_param_index: 1},
            rounding_func=lambda x: x,
        )
        self.assertEqual(np.shape(generated_points), (3, 4))
        np_bounds = np.array(bounds)
        self.assertTrue(np.alltrue(generated_points >= np_bounds[:, 0]))
        self.assertTrue(np.alltrue(generated_points <= np_bounds[:, 1]))
//  [python]
def test_SobolGeneratorNoScramble(self) -> None:
        generator = SobolGenerator(scramble=False)
        n_tunable = fixed_param_index = 3
        bounds = self._create_bounds(n_tunable=n_tunable, n_fixed=1)
        generated_points, weights = generator.gen(
            n=3,
            bounds=bounds,
            fixed_features={fixed_param_index: 1},
            rounding_func=lambda x: x,
        )
        self.assertEqual(np.shape(generated_points), (3, 4))
        np_bounds = np.array(bounds)
        self.assertTrue(np.all(generated_points >= np_bounds[:, 0]))
        self.assertTrue(np.all(generated_points <= np_bounds[:, 1]))
//  [python]
def test_times(self):
        """
        Test if the correct times array is returned for normal traces and
        traces with gaps.
        """
        from matplotlib import __version__
        tr = Trace(data=np.ones(100))
        tr.stats.sampling_rate = 20
        delta = tr.stats.delta
        start = UTC(2000, 1, 1, 0, 0, 0, 0)
        tr.stats.starttime = start
        tm = tr.times()
        diff = tr.stats.endtime - tr.stats.starttime
        assert np.isclose(tm[-1], diff)
        tr.data = np.ma.ones(100)
        tr.data[30:40] = np.ma.masked
        tm = tr.times()
        assert np.alltrue(tr.data.mask == tm.mask)
        # test relative with reftime
        tr.data = np.ones(100)
        shift = 9.5
        reftime = start - shift
        got = tr.times(reftime=reftime)
        assert len(got) == tr.stats.npts
        expected = np.arange(shift, shift + 4.5 * delta, delta)
        np.testing.assert_allclose(got[:5], expected, rtol=1e-8)
        # test other options
        got = tr.times("UTCDateTime")
        expected = np.array([
            UTC(2000, 1, 1, 0, 0),
            UTC(2000, 1, 1, 0, 0, 0, 50000),
            UTC(2000, 1, 1, 0, 0, 0, 100000),
            UTC(2000, 1, 1, 0, 0, 0, 150000),
            UTC(2000, 1, 1, 0, 0, 0, 200000)], dtype=UTC)
        assert isinstance(got[0], UTC)
        np.testing.assert_allclose(
            [t_.timestamp for t_ in got[:5]],
            [t_.timestamp for t_ in expected], rtol=1e-17)
        got = tr.times("timestamp")
        expected = np.arange(0, 4.5 * delta, delta) + 946684800.0
        np.testing.assert_allclose(got[:5], expected, rtol=1e-17)
        got = tr.times("matplotlib")
        expected = np.array([
                10957.000000000000, 10957.000000578704, 10957.000001157407,
                10957.000001736111, 10957.000002314815])
        if parse_version(__version__) < parse_version('3.3'):
            expected = np.array([
                730120.00000000000000000000, 730120.00000057870056480169,
                730120.00000115740112960339, 730120.00000173610169440508,
                730120.00000231480225920677])
        np.testing.assert_allclose(got[:5], expected, rtol=1e-17)
//  [python]
def test_times(self):
        """
        Test if the correct times array is returned for normal traces and
        traces with gaps.
        """
        from matplotlib import __version__
        tr = Trace(data=np.ones(100))
        tr.stats.sampling_rate = 20
        delta = tr.stats.delta
        start = UTC(2000, 1, 1, 0, 0, 0, 0)
        tr.stats.starttime = start
        tm = tr.times()
        diff = tr.stats.endtime - tr.stats.starttime
        assert np.isclose(tm[-1], diff)
        tr.data = np.ma.ones(100)
        tr.data[30:40] = np.ma.masked
        tm = tr.times()
        assert np.all(tr.data.mask == tm.mask)
        # test relative with reftime
        tr.data = np.ones(100)
        shift = 9.5
        reftime = start - shift
        got = tr.times(reftime=reftime)
        assert len(got) == tr.stats.npts
        expected = np.arange(shift, shift + 4.5 * delta, delta)
        np.testing.assert_allclose(got[:5], expected, rtol=1e-8)
        # test other options
        got = tr.times("UTCDateTime")
        expected = np.array([
            UTC(2000, 1, 1, 0, 0),
            UTC(2000, 1, 1, 0, 0, 0, 50000),
            UTC(2000, 1, 1, 0, 0, 0, 100000),
            UTC(2000, 1, 1, 0, 0, 0, 150000),
            UTC(2000, 1, 1, 0, 0, 0, 200000)], dtype=UTC)
        assert isinstance(got[0], UTC)
        np.testing.assert_allclose(
            [t_.timestamp for t_ in got[:5]],
            [t_.timestamp for t_ in expected], rtol=1e-17)
        got = tr.times("timestamp")
        expected = np.arange(0, 4.5 * delta, delta) + 946684800.0
        np.testing.assert_allclose(got[:5], expected, rtol=1e-17)
        got = tr.times("matplotlib")
        expected = np.array([
                10957.000000000000, 10957.000000578704, 10957.000001157407,
                10957.000001736111, 10957.000002314815])
        if parse_version(__version__) < parse_version('3.3'):
            expected = np.array([
                730120.00000000000000000000, 730120.00000057870056480169,
                730120.00000115740112960339, 730120.00000173610169440508,
                730120.00000231480225920677])
        np.testing.assert_allclose(got[:5], expected, rtol=1e-17)
//  [python]
def test_SobolGeneratorFixedSpace(self) -> None:
        generator = SobolGenerator(seed=0, deduplicate=False)
        bounds = self._create_bounds(n_tunable=0, n_fixed=2)
        generated_points, _ = generator.gen(
            n=3,
            bounds=bounds,
            fixed_features={0: 1, 1: 2},
            rounding_func=lambda x: x,
        )
        self.assertEqual(np.shape(generated_points), (3, 2))
        np_bounds = np.array(bounds)
        self.assertTrue(np.alltrue(generated_points >= np_bounds[:, 0]))
        self.assertTrue(np.alltrue(generated_points <= np_bounds[:, 1]))
        # Should error out if deduplicating since there's only one feasible point.
        generator = SobolGenerator(seed=0, deduplicate=True)
        with self.assertRaisesRegex(SearchSpaceExhausted, "Rejection sampling"):
            generated_points, _ = generator.gen(
                n=3,
                bounds=bounds,
                fixed_features={0: 1, 1: 2},
                rounding_func=lambda x: x,
            )
        # But we can generate 1 point.
        generated_points, _ = generator.gen(
            n=1,
            bounds=bounds,
            fixed_features={0: 1, 1: 2},
            rounding_func=lambda x: x,
        )
        self.assertEqual(np.shape(generated_points), (1, 2))
//  [python]
def test_SobolGeneratorFixedSpace(self) -> None:
        generator = SobolGenerator(seed=0, deduplicate=False)
        bounds = self._create_bounds(n_tunable=0, n_fixed=2)
        generated_points, _ = generator.gen(
            n=3,
            bounds=bounds,
            fixed_features={0: 1, 1: 2},
            rounding_func=lambda x: x,
        )
        self.assertEqual(np.shape(generated_points), (3, 2))
        np_bounds = np.array(bounds)
        self.assertTrue(np.all(generated_points >= np_bounds[:, 0]))
        self.assertTrue(np.all(generated_points <= np_bounds[:, 1]))
        # Should error out if deduplicating since there's only one feasible point.
        generator = SobolGenerator(seed=0, deduplicate=True)
        with self.assertRaisesRegex(SearchSpaceExhausted, "Rejection sampling"):
            generated_points, _ = generator.gen(
                n=3,
                bounds=bounds,
                fixed_features={0: 1, 1: 2},
                rounding_func=lambda x: x,
            )
        # But we can generate 1 point.
        generated_points, _ = generator.gen(
            n=1,
            bounds=bounds,
            fixed_features={0: 1, 1: 2},
            rounding_func=lambda x: x,
        )
        self.assertEqual(np.shape(generated_points), (1, 2))
//  [python]
def test_SobolGeneratorOnline(self) -> None:
        # Verify that the generator will return the expected arms if called
        # one at a time.
        bulk_generator = SobolGenerator(seed=0)
        generator = SobolGenerator(seed=0)
        n_tunable = fixed_param_index = 3
        bounds = self._create_bounds(n_tunable=n_tunable, n_fixed=1)
        bulk_generated_points, bulk_weights = bulk_generator.gen(
            n=3,
            bounds=bounds,
            fixed_features={fixed_param_index: 1},
            rounding_func=lambda x: x,
        )
        np_bounds = np.array(bounds)
        for expected_points in bulk_generated_points:
            generated_points, weights = generator.gen(
                n=1,
                bounds=bounds,
                fixed_features={fixed_param_index: 1},
                rounding_func=lambda x: x,
            )
            self.assertEqual(weights, [1])
            self.assertTrue(np.alltrue(generated_points >= np_bounds[:, 0]))
            self.assertTrue(np.alltrue(generated_points <= np_bounds[:, 1]))
            self.assertTrue(generated_points[..., -1] == 1)
            self.assertTrue(np.array_equal(expected_points, generated_points.flatten()))
//  [python]
def test_SobolGeneratorOnline(self) -> None:
        # Verify that the generator will return the expected arms if called
        # one at a time.
        bulk_generator = SobolGenerator(seed=0)
        generator = SobolGenerator(seed=0)
        n_tunable = fixed_param_index = 3
        bounds = self._create_bounds(n_tunable=n_tunable, n_fixed=1)
        bulk_generated_points, bulk_weights = bulk_generator.gen(
            n=3,
            bounds=bounds,
            fixed_features={fixed_param_index: 1},
            rounding_func=lambda x: x,
        )
        np_bounds = np.array(bounds)
        for expected_points in bulk_generated_points:
            generated_points, weights = generator.gen(
                n=1,
                bounds=bounds,
                fixed_features={fixed_param_index: 1},
                rounding_func=lambda x: x,
            )
            self.assertEqual(weights, [1])
            self.assertTrue(np.all(generated_points >= np_bounds[:, 0]))
            self.assertTrue(np.all(generated_points <= np_bounds[:, 1]))
            self.assertTrue(generated_points[..., -1] == 1)
            self.assertTrue(np.array_equal(expected_points, generated_points.flatten()))
//  [python]
def test_SobolGeneratorWithOrderConstraints(self) -> None:
        # Enforce dim_0 <= dim_1 <= dim_2 <= dim_3.
        # Enforce both fixed and tunable constraints.
        generator = SobolGenerator(seed=0)
        n_tunable = fixed_param_index = 3
        bounds = self._create_bounds(n_tunable=n_tunable, n_fixed=1)
        A = np.array([[1, -1, 0, 0], [0, 1, -1, 0], [0, 0, 1, -1]])
        b = np.array([0, 0, 0])
        generated_points, weights = generator.gen(
            n=3,
            bounds=bounds,
            linear_constraints=(A, b),
            fixed_features={fixed_param_index: 0.5},
            rounding_func=lambda x: x,
        )
        self.assertEqual(np.shape(generated_points), (3, 4))
        self.assertTrue(np.alltrue(generated_points[..., -1] == 0.5))
        self.assertTrue(
            np.array_equal(
                np.sort(generated_points[..., :-1], axis=-1),
                generated_points[..., :-1],
            )
        )
//  [python]
def test_SobolGeneratorWithOrderConstraints(self) -> None:
        # Enforce dim_0 <= dim_1 <= dim_2 <= dim_3.
        # Enforce both fixed and tunable constraints.
        generator = SobolGenerator(seed=0)
        n_tunable = fixed_param_index = 3
        bounds = self._create_bounds(n_tunable=n_tunable, n_fixed=1)
        A = np.array([[1, -1, 0, 0], [0, 1, -1, 0], [0, 0, 1, -1]])
        b = np.array([0, 0, 0])
        generated_points, weights = generator.gen(
            n=3,
            bounds=bounds,
            linear_constraints=(A, b),
            fixed_features={fixed_param_index: 0.5},
            rounding_func=lambda x: x,
        )
        self.assertEqual(np.shape(generated_points), (3, 4))
        self.assertTrue(np.all(generated_points[..., -1] == 0.5))
        self.assertTrue(
            np.array_equal(
                np.sort(generated_points[..., :-1], axis=-1),
                generated_points[..., :-1],
            )
        )
//  [python]
def test_SobolGeneratorWithLinearConstraints(self) -> None:
        # Enforce dim_0 <= dim_1 <= dim_2 <= dim_3.
        # Enforce both fixed and tunable constraints.
        generator = SobolGenerator(seed=0)
        n_tunable = fixed_param_index = 3
        bounds = self._create_bounds(n_tunable=n_tunable, n_fixed=1)
        A = np.array([[1, 1, 0, 0], [0, 1, 1, 0]])
        b = np.array([1, 1])

        generated_points, weights = generator.gen(
            n=3,
            bounds=bounds,
            linear_constraints=(
                A,
                b,
            ),
            fixed_features={fixed_param_index: 1},
            rounding_func=lambda x: x,
        )
        self.assertTrue(np.shape(generated_points) == (3, 4))
        self.assertTrue(np.alltrue(generated_points[..., -1] == 1))
        self.assertTrue(np.alltrue(generated_points @ A.transpose() <= b))
//  [python]
def test_SobolGeneratorWithLinearConstraints(self) -> None:
        # Enforce dim_0 <= dim_1 <= dim_2 <= dim_3.
        # Enforce both fixed and tunable constraints.
        generator = SobolGenerator(seed=0)
        n_tunable = fixed_param_index = 3
        bounds = self._create_bounds(n_tunable=n_tunable, n_fixed=1)
        A = np.array([[1, 1, 0, 0], [0, 1, 1, 0]])
        b = np.array([1, 1])

        generated_points, weights = generator.gen(
            n=3,
            bounds=bounds,
            linear_constraints=(
                A,
                b,
            ),
            fixed_features={fixed_param_index: 1},
            rounding_func=lambda x: x,
        )
        self.assertTrue(np.shape(generated_points) == (3, 4))
        self.assertTrue(np.all(generated_points[..., -1] == 1))
        self.assertTrue(np.all(generated_points @ A.transpose() <= b))
//  [python]
def test_monitorGetGamma():
    #create our monitor object
    gammaVal = [2.2, 2.2, 2.2]
    mon = monitors.Monitor('test')
    mon.setGamma(gammaVal)
    #create window using that monitor
    win = visual.Window([100,100], monitor=mon, autoLog=False)
    assert numpy.alltrue(win.gamma==gammaVal)
    win.close()
//  [python]
def test_monitorGetGamma():
    #create our monitor object
    gammaVal = [2.2, 2.2, 2.2]
    mon = monitors.Monitor('test')
    mon.setGamma(gammaVal)
    #create window using that monitor
    win = visual.Window([100,100], monitor=mon, autoLog=False)
    assert numpy.all(win.gamma==gammaVal)
    win.close()
//  [python]
def is_np_all(self):
        ret = self._flatten(self._recursive_do_on_memory(self.memory, is_np))
        return np.alltrue([v for k, v in ret.items()]) if isinstance(ret, dict) else ret
//  [python]
def is_np_all(self):
        ret = self._flatten(self._recursive_do_on_memory(self.memory, is_np))
        return np.all([v for k, v in ret.items()]) if isinstance(ret, dict) else ret
//  [python]
def test_SobolGeneratorFallbackToPolytopeSamplerWithFixedParam(self) -> None:
        # Ten parameters with sum less than 1. In this example, the rejection
        # sampler gives a search space exhausted error.  Testing fallback to
        # polytope sampler when encountering this error.
        generator = SobolGenerator(seed=0, fallback_to_sample_polytope=True)
        bounds = self._create_bounds(n_tunable=10, n_fixed=1)
        A = np.insert(np.ones((1, 10)), 10, 0, axis=1)
        b = np.array([1]).reshape((1, 1))
        generated_points, weights = generator.gen(
            n=3,
            bounds=bounds,
            linear_constraints=(
                A,
                b,
            ),
            fixed_features={10: 1},
            rounding_func=lambda x: x,
        )
        self.assertTrue(np.shape(generated_points) == (3, 11))
        self.assertTrue(np.alltrue(generated_points[..., -1] == 1))
        self.assertTrue(np.alltrue(generated_points @ A.transpose() <= b))
//  [python]
def test_SobolGeneratorFallbackToPolytopeSamplerWithFixedParam(self) -> None:
        # Ten parameters with sum less than 1. In this example, the rejection
        # sampler gives a search space exhausted error.  Testing fallback to
        # polytope sampler when encountering this error.
        generator = SobolGenerator(seed=0, fallback_to_sample_polytope=True)
        bounds = self._create_bounds(n_tunable=10, n_fixed=1)
        A = np.insert(np.ones((1, 10)), 10, 0, axis=1)
        b = np.array([1]).reshape((1, 1))
        generated_points, weights = generator.gen(
            n=3,
            bounds=bounds,
            linear_constraints=(
                A,
                b,
            ),
            fixed_features={10: 1},
            rounding_func=lambda x: x,
        )
        self.assertTrue(np.shape(generated_points) == (3, 11))
        self.assertTrue(np.all(generated_points[..., -1] == 1))
        self.assertTrue(np.all(generated_points @ A.transpose() <= b))
//  [python]
def test_monitorGetGammaGrid():
    #create (outdated) gamma grid (new one is [4,6])
    newGrid = numpy.array([[0,150,2.0],#lum
                           [0,30,2.0],#r
                           [0,110,2.0],#g
                           [0,10,2.0]],#b
                           )
    mon = monitors.Monitor('test')
    mon.setGammaGrid(newGrid)
    win = visual.Window([100,100], monitor=mon, autoLog=False)
    assert numpy.alltrue(win.gamma==numpy.array([2.0, 2.0, 2.0]))
    win.close()
//  [python]
def test_monitorGetGammaGrid():
    #create (outdated) gamma grid (new one is [4,6])
    newGrid = numpy.array([[0,150,2.0],#lum
                           [0,30,2.0],#r
                           [0,110,2.0],#g
                           [0,10,2.0]],#b
                           )
    mon = monitors.Monitor('test')
    mon.setGammaGrid(newGrid)
    win = visual.Window([100,100], monitor=mon, autoLog=False)
    assert numpy.all(win.gamma==numpy.array([2.0, 2.0, 2.0]))
    win.close()
//  [python]
def test_monitorGetGammaAndGrid():
    """test what happens if gamma (old) and gammaGrid (new) are both present"""
    #create (outdated) gamma grid (new one is [4,6])
    newGrid = numpy.array([[0,150,2.0],#lum
                           [0,30,2.0],#r
                           [0,110,2.0],#g
                           [0,10,2.0]],#b
                           )
    mon = monitors.Monitor('test')
    mon.setGammaGrid(newGrid)
    mon.setGamma([3,3,3])
    #create window using that monitor
    win = visual.Window([100,100], monitor=mon, autoLog=False)
    assert numpy.alltrue(win.gamma==numpy.array([2.0, 2.0, 2.0]))
    win.close()
//  [python]
def test_monitorGetGammaAndGrid():
    """test what happens if gamma (old) and gammaGrid (new) are both present"""
    #create (outdated) gamma grid (new one is [4,6])
    newGrid = numpy.array([[0,150,2.0],#lum
                           [0,30,2.0],#r
                           [0,110,2.0],#g
                           [0,10,2.0]],#b
                           )
    mon = monitors.Monitor('test')
    mon.setGammaGrid(newGrid)
    mon.setGamma([3,3,3])
    #create window using that monitor
    win = visual.Window([100,100], monitor=mon, autoLog=False)
    assert numpy.all(win.gamma==numpy.array([2.0, 2.0, 2.0]))
    win.close()
//  [python]
def test_SobolGeneratorFallbackToPolytopeSampler(self) -> None:
        # Ten parameters with sum less than 1. In this example, the rejection
        # sampler gives a search space exhausted error.  Testing fallback to
        # polytope sampler when encountering this error.
        generator = SobolGenerator(seed=0, fallback_to_sample_polytope=True)
        bounds = self._create_bounds(n_tunable=10, n_fixed=0)
        A = np.ones((1, 10))
        b = np.array([1]).reshape((1, 1))
        with mock.patch("ax.models.random.base.logger.info") as mock_logger, mock.patch(
            "botorch.utils.sampling.sample_polytope",
            wraps=sample_polytope,
        ) as wrapped_sampler:
            generated_points, weights = generator.gen(
                n=3,
                bounds=bounds,
                linear_constraints=(
                    A,
                    b,
                ),
                rounding_func=lambda x: x,
            )
        # First call uses the original seed since no candidates are generated.
        self.assertEqual(wrapped_sampler.call_args[-1]["seed"], 0)
        self.assertTrue(
            "exceeded specified maximum draws" in mock_logger.call_args[0][0]
        )
        self.assertTrue(np.shape(generated_points) == (3, 10))
        self.assertTrue(np.alltrue(generated_points @ A.transpose() <= b))

        with mock.patch(
            "botorch.utils.sampling.sample_polytope",
            wraps=sample_polytope,
        ) as wrapped_sampler:
            generator.gen(
                n=3,
                bounds=bounds,
                linear_constraints=(
                    A,
                    b,
                ),
                rounding_func=lambda x: x,
            )
        # Second call uses seed 3 since 3 candidates are already generated.
        self.assertEqual(wrapped_sampler.call_args[-1]["seed"], 3)
//  [python]
def test_SobolGeneratorFallbackToPolytopeSampler(self) -> None:
        # Ten parameters with sum less than 1. In this example, the rejection
        # sampler gives a search space exhausted error.  Testing fallback to
        # polytope sampler when encountering this error.
        generator = SobolGenerator(seed=0, fallback_to_sample_polytope=True)
        bounds = self._create_bounds(n_tunable=10, n_fixed=0)
        A = np.ones((1, 10))
        b = np.array([1]).reshape((1, 1))
        with mock.patch("ax.models.random.base.logger.info") as mock_logger, mock.patch(
            "botorch.utils.sampling.sample_polytope",
            wraps=sample_polytope,
        ) as wrapped_sampler:
            generated_points, weights = generator.gen(
                n=3,
                bounds=bounds,
                linear_constraints=(
                    A,
                    b,
                ),
                rounding_func=lambda x: x,
            )
        # First call uses the original seed since no candidates are generated.
        self.assertEqual(wrapped_sampler.call_args[-1]["seed"], 0)
        self.assertTrue(
            "exceeded specified maximum draws" in mock_logger.call_args[0][0]
        )
        self.assertTrue(np.shape(generated_points) == (3, 10))
        self.assertTrue(np.all(generated_points @ A.transpose() <= b))

        with mock.patch(
            "botorch.utils.sampling.sample_polytope",
            wraps=sample_polytope,
        ) as wrapped_sampler:
            generator.gen(
                n=3,
                bounds=bounds,
                linear_constraints=(
                    A,
                    b,
                ),
                rounding_func=lambda x: x,
            )
        # Second call uses seed 3 since 3 candidates are already generated.
        self.assertEqual(wrapped_sampler.call_args[-1]["seed"], 3)
//  [python]
def testFull(self, num_best=None):
        # Override testFull.

        if not PYEMD_EXT:
            return

        index = self.cls(texts, self.w2v_model)
        index.num_best = num_best
        query = texts[0]
        sims = index[query]

        if num_best is not None:
            # Sparse array.
            for i, sim in sims:
                # Note that similarities are bigger than zero, as they are the 1/ 1 + distances.
                self.assertTrue(numpy.alltrue(sim > 0.0))
        else:
            self.assertTrue(sims[0] == 1.0)  # Similarity of a document with itself is 0.0.
            self.assertTrue(numpy.alltrue(sims[1:] > 0.0))
            self.assertTrue(numpy.alltrue(sims[1:] < 1.0))
//  [python]
def testFull(self, num_best=None):
        # Override testFull.

        if not PYEMD_EXT:
            return

        index = self.cls(texts, self.w2v_model)
        index.num_best = num_best
        query = texts[0]
        sims = index[query]

        if num_best is not None:
            # Sparse array.
            for i, sim in sims:
                # Note that similarities are bigger than zero, as they are the 1/ 1 + distances.
                self.assertTrue(numpy.all(sim > 0.0))
        else:
            self.assertTrue(sims[0] == 1.0)  # Similarity of a document with itself is 0.0.
            self.assertTrue(numpy.all(sims[1:] > 0.0))
            self.assertTrue(numpy.all(sims[1:] < 1.0))
//  [python]
def testIter(self):
        # Override testIter.

        index = self.cls(self.corpus, self.similarity_matrix)
        for sims in index:
            self.assertTrue(numpy.alltrue(sims >= 0.0))
            self.assertTrue(numpy.alltrue(sims <= 1.0))
//  [python]
def testIter(self):
        # Override testIter.

        index = self.cls(self.corpus, self.similarity_matrix)
        for sims in index:
            self.assertTrue(numpy.all(sims >= 0.0))
            self.assertTrue(numpy.all(sims <= 1.0))
//  [python]
def testChunking(self):
        # Override testChunking.

        index = self.cls(self.corpus, self.similarity_matrix)
        query = [self.dictionary.doc2bow(document) for document in texts[:3]]
        sims = index[query]

        for i in range(3):
            self.assertTrue(numpy.alltrue(sims[i, i] == 1.0))  # Similarity of a document with itself is 1.0.

        # test the same thing but with num_best
        index.num_best = 5
        sims = index[query]
        for i, chunk in enumerate(sims):
            expected = i
            self.assertAlmostEquals(expected, chunk[0][0], places=2)
            expected = 1.0
            self.assertAlmostEquals(expected, chunk[0][1], places=2)
//  [python]
def testChunking(self):
        # Override testChunking.

        index = self.cls(self.corpus, self.similarity_matrix)
        query = [self.dictionary.doc2bow(document) for document in texts[:3]]
        sims = index[query]

        for i in range(3):
            self.assertTrue(numpy.all(sims[i, i] == 1.0))  # Similarity of a document with itself is 1.0.

        # test the same thing but with num_best
        index.num_best = 5
        sims = index[query]
        for i, chunk in enumerate(sims):
            expected = i
            self.assertAlmostEquals(expected, chunk[0][0], places=2)
            expected = 1.0
            self.assertAlmostEquals(expected, chunk[0][1], places=2)
//  [python]
def testFull(self, num_best=None):
        # Override testFull.

        index = self.cls(self.corpus, self.similarity_matrix, num_best=num_best)
        query = self.dictionary.doc2bow(texts[0])
        sims = index[query]

        if num_best is not None:
            # Sparse array.
            for i, sim in sims:
                self.assertTrue(numpy.alltrue(sim <= 1.0))
                self.assertTrue(numpy.alltrue(sim >= 0.0))
        else:
            self.assertTrue(sims[0] == 1.0)  # Similarity of a document with itself is 1.0.
            self.assertTrue(numpy.alltrue(sims[1:] >= 0.0))
            self.assertTrue(numpy.alltrue(sims[1:] < 1.0))
            expected = 2.1889350195476758
            self.assertAlmostEqual(expected, numpy.sum(sims))
//  [python]
def testFull(self, num_best=None):
        # Override testFull.

        index = self.cls(self.corpus, self.similarity_matrix, num_best=num_best)
        query = self.dictionary.doc2bow(texts[0])
        sims = index[query]

        if num_best is not None:
            # Sparse array.
            for i, sim in sims:
                self.assertTrue(numpy.all(sim <= 1.0))
                self.assertTrue(numpy.all(sim >= 0.0))
        else:
            self.assertTrue(sims[0] == 1.0)  # Similarity of a document with itself is 1.0.
            self.assertTrue(numpy.all(sims[1:] >= 0.0))
            self.assertTrue(numpy.all(sims[1:] < 1.0))
            expected = 2.1889350195476758
            self.assertAlmostEqual(expected, numpy.sum(sims))
//  [python]
def test_predict_gpu(self):
        clf = Regressor(self.predictor, device=0)
        actual_t = clf.predict(self.x)
        assert numpy.alltrue(actual_t == self.t)
//  [python]
def test_predict_gpu(self):
        clf = Regressor(self.predictor, device=0)
        actual_t = clf.predict(self.x)
        assert numpy.all(actual_t == self.t)
//  [python]
def fail_if_array_equal(x, y, err_msg='', verbose=True):
    """
    Raises an assertion error if two masked arrays are not equal elementwise.

    """
    def compare(x, y):
        return (not np.alltrue(approx(x, y)))
    assert_array_compare(compare, x, y, err_msg=err_msg, verbose=verbose,
                         header='Arrays are not equal')
//  [python]
def fail_if_array_equal(x, y, err_msg='', verbose=True):
    """
    Raises an assertion error if two masked arrays are not equal elementwise.

    """
    def compare(x, y):
        return (not np.all(approx(x, y)))
    assert_array_compare(compare, x, y, err_msg=err_msg, verbose=verbose,
                         header='Arrays are not equal')
//  [python]
def test_predict_cpu(self):
        clf = Regressor(self.predictor)
        actual_t = clf.predict(self.x)
        assert actual_t.shape == (3, 2)
        assert actual_t.dtype == numpy.float32
        assert numpy.alltrue(actual_t == self.t)
//  [python]
def test_predict_cpu(self):
        clf = Regressor(self.predictor)
        actual_t = clf.predict(self.x)
        assert actual_t.shape == (3, 2)
        assert actual_t.dtype == numpy.float32
        assert numpy.all(actual_t == self.t)
//  [python]
def find_K_index(self, K0):
        '''
        Find the index of K0.
        '''

        for ii in range(self.wfc._nkpts):
            if np.alltrue(
                np.abs(self.wfc._kvecs[ii] - K0) < 1E-5
            ):
                return ii + 1
        # the for-else
        else:
            raise ValueError(
                'Cannot find the corresponding K-points in WAVECAR!'
            )
//  [python]
def find_K_index(self, K0):
        '''
        Find the index of K0.
        '''

        for ii in range(self.wfc._nkpts):
            if np.all(
                np.abs(self.wfc._kvecs[ii] - K0) < 1E-5
            ):
                return ii + 1
        # the for-else
        else:
            raise ValueError(
                'Cannot find the corresponding K-points in WAVECAR!'
            )
//  [python]
def test_cfworld(med_vars, cfworld, world, world0, world1):
    assert np.alltrue(cfworld.Observables()[med_vars.A] == world1.Observables()[med_vars.A] )
    assert np.alltrue(cfworld.Observables()[med_vars.M] == world0.Observables()[med_vars.M] )
    assert np.alltrue(cfworld.data[med_vars.A.Noise()] == world.data[med_vars.A.Noise()] )
    assert np.alltrue(cfworld.data[med_vars.M.Noise()] == world.data[med_vars.M.Noise()] )
    assert np.alltrue(cfworld.data[med_vars.Y.Noise()] == world.data[med_vars.Y.Noise()] )
//  [python]
def test_cfworld(med_vars, cfworld, world, world0, world1):
    assert np.all(cfworld.Observables()[med_vars.A] == world1.Observables()[med_vars.A] )
    assert np.all(cfworld.Observables()[med_vars.M] == world0.Observables()[med_vars.M] )
    assert np.all(cfworld.data[med_vars.A.Noise()] == world.data[med_vars.A.Noise()] )
    assert np.all(cfworld.data[med_vars.M.Noise()] == world.data[med_vars.M.Noise()] )
    assert np.all(cfworld.data[med_vars.Y.Noise()] == world.data[med_vars.Y.Noise()] )
//  [python]
def get_ovlap_G(self, ikpt=1, epsilon=1E-5):
        '''
        Get subset of the reciprocal space vectors of the supercell,
        specifically the ones that match the reciprocal space vectors of the
        primitive cell.
        '''

        assert 1 <= ikpt <= self.wfc._nkpts, 'Invalid K-point index!'

        # Reciprocal space vectors of the supercell in fractional unit
        Gvecs = self.wfc.gvectors(ikpt=ikpt)
        # Gvecs = self.allGvecs[ikpt - 1]

        if self._lgam:
            nplw = Gvecs.shape[0]
            tmp = np.zeros((nplw * 2 - 1, 3), dtype=int)
            # the gvectors of Gamma version only contains half the number of a
            # normal version.
            tmp[:nplw, ...] = Gvecs
            tmp[nplw:, ...] = -Gvecs[1:, ...]            # G' = -G

            Gvecs = tmp

        # Shape of Gvecs: (nplws, 3)
        # iGvecs = np.arange(Gvecs.shape[0], dtype=int)

        # Reciprocal space vectors of the primitive cell
        gvecs = np.dot(Gvecs, np.linalg.inv(self.M).T)
        # Deviation from the perfect sites
        gd = gvecs - np.round(gvecs)
        # match = np.linalg.norm(gd, axis=1) < epsilon
        match = np.alltrue(
            np.abs(gd) < epsilon, axis=1
        )

        # return Gvecs[match], iGvecs[match]
        return Gvecs[match], Gvecs
//  [python]
def get_ovlap_G(self, ikpt=1, epsilon=1E-5):
        '''
        Get subset of the reciprocal space vectors of the supercell,
        specifically the ones that match the reciprocal space vectors of the
        primitive cell.
        '''

        assert 1 <= ikpt <= self.wfc._nkpts, 'Invalid K-point index!'

        # Reciprocal space vectors of the supercell in fractional unit
        Gvecs = self.wfc.gvectors(ikpt=ikpt)
        # Gvecs = self.allGvecs[ikpt - 1]

        if self._lgam:
            nplw = Gvecs.shape[0]
            tmp = np.zeros((nplw * 2 - 1, 3), dtype=int)
            # the gvectors of Gamma version only contains half the number of a
            # normal version.
            tmp[:nplw, ...] = Gvecs
            tmp[nplw:, ...] = -Gvecs[1:, ...]            # G' = -G

            Gvecs = tmp

        # Shape of Gvecs: (nplws, 3)
        # iGvecs = np.arange(Gvecs.shape[0], dtype=int)

        # Reciprocal space vectors of the primitive cell
        gvecs = np.dot(Gvecs, np.linalg.inv(self.M).T)
        # Deviation from the perfect sites
        gd = gvecs - np.round(gvecs)
        # match = np.linalg.norm(gd, axis=1) < epsilon
        match = np.all(
            np.abs(gd) < epsilon, axis=1
        )

        # return Gvecs[match], iGvecs[match]
        return Gvecs[match], Gvecs
//  [python]
def test_boxcar_nan(self):
        
        """ test boxcar performance with a nan in data"""
                
        data=[1.0]*200+[2.0]*100+[1.0]*100
        data=numpy.array(data)
        data[201]=numpy.nan
        st=datetime.datetime(year=1990,month=2,day=3,hour=11, minute=15)
        delta=time_interval(hours=1)
        test_ts=rts(data,st,delta,{})
        nt=boxcar(test_ts,11,12)
        
        ## besides the leading and trailing nan, there will be
        ## 24 nan in the middle 
        self.assert_(numpy.alltrue(numpy.isnan(nt.data[190:213])))
//  [python]
def test_boxcar_nan(self):
        
        """ test boxcar performance with a nan in data"""
                
        data=[1.0]*200+[2.0]*100+[1.0]*100
        data=numpy.array(data)
        data[201]=numpy.nan
        st=datetime.datetime(year=1990,month=2,day=3,hour=11, minute=15)
        delta=time_interval(hours=1)
        test_ts=rts(data,st,delta,{})
        nt=boxcar(test_ts,11,12)
        
        ## besides the leading and trailing nan, there will be
        ## 24 nan in the middle 
        self.assert_(numpy.all(numpy.isnan(nt.data[190:213])))
//  [python]
def test_godin_15min(self):        
        """ test godin filtering on a 15min constant values 
            data series with a nan.
        """        
        data=[1.0]*800+[2.0]*400+[1.0]*400
        data=numpy.array(data)
        data[336]=numpy.nan
        st=datetime.datetime(year=1990,month=2,day=3,hour=11, minute=15)
        delta=time_interval(minutes=15)
        test_ts=rts(data,st,delta,{})
        nt3=godin(test_ts)
        self.assert_(numpy.alltrue(numpy.isnan(nt3.data[0:144])))
        assert_array_almost_equal(nt3.data[144:192],[1]*48,12)
        self.assert_(numpy.alltrue(numpy.isnan(nt3.data[192:481])))
        assert_array_almost_equal(nt3.data[481:656],[1]*175,12)
        self.assert_(numpy.alltrue(numpy.greater(nt3.data[656:944],1)))
        self.assertAlmostEqual(nt3.data[868],1.916618441)
        assert_array_almost_equal(nt3.data[944:1056],[2]*112,12)
        self.assert_(numpy.alltrue(numpy.greater(nt3.data[1056:1344],1)))
        self.assertAlmostEqual(nt3.data[1284],1.041451845)
        assert_array_almost_equal(nt3.data[1344:1456],[1]*112,12)
        self.assert_(numpy.alltrue(numpy.isnan(nt3.data[1456:1600])))