jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_orbit_time():
    # Test that Orbit.time returns the time correctly, with units when that's
    # required
    from astropy import units as u

    from galpy.orbit import Orbit
    from galpy.potential import MWPotential2014
    from galpy.util import conversion

    ts = numpy.linspace(0.0, 1.0, 1001) * u.Gyr
    o = Orbit()
    o.integrate(ts, MWPotential2014)
    # No argument, in this case should be times in Gyr
    assert numpy.all(
        numpy.fabs((ts - o.time(quantity=True)) / ts[-1]).value < 1e-10
    ), "Orbit.time does not return the correct times"
    # with argument, should be time in Gyr
    assert numpy.all(
        numpy.fabs((ts - o.time(ts, quantity=True)) / ts[-1]).value < 1e-10
    ), "Orbit.time does not return the correct times"
    assert (
        numpy.fabs((ts[-1] - o.time(ts[-1], quantity=True)) / ts[-1]).value < 1e-10
    ), "Orbit.time does not return the correct times"
    # with argument without units --> units
    assert numpy.all(
        numpy.fabs(
            ts
            - o.time(
                ts.to(u.Gyr).value
                / conversion.time_in_Gyr(
                    MWPotential2014[0]._vo, MWPotential2014[0]._ro
                ),
                quantity=True,
            )
        ).value
        < 1e-10
    ), "Orbit.time does not return the correct times"
    assert (
        numpy.fabs(
            ts[-1]
            - o.time(
                ts[-1].to(u.Gyr).value
                / conversion.time_in_Gyr(
                    MWPotential2014[0]._vo, MWPotential2014[0]._ro
                ),
                quantity=True,
            )
        ).value
        < 1e-10
    ), "Orbit.time does not return the correct times"
    # Now should get without units
    o.turn_physical_off()
    assert numpy.all(
        numpy.fabs(
            ts.to(u.Gyr).value
            / conversion.time_in_Gyr(MWPotential2014[0]._vo, MWPotential2014[0]._ro)
            - o.time(
                ts.to(u.Gyr).value
                / conversion.time_in_Gyr(MWPotential2014[0]._vo, MWPotential2014[0]._ro)
            )
        )
        < 1e-10
    ), "Orbit.time does not return the correct times"
    assert (
        numpy.fabs(
            ts[-1].to(u.Gyr).value
            / conversion.time_in_Gyr(MWPotential2014[0]._vo, MWPotential2014[0]._ro)
            - o.time(
                ts[-1].to(u.Gyr).value
                / conversion.time_in_Gyr(MWPotential2014[0]._vo, MWPotential2014[0]._ro)
            )
        )
        < 1e-10
    ), "Orbit.time does not return the correct times"
    return None
//  [python]
def Phase3Variance(matrices, values, index):
    """
    Best split selection code. The underlying loss function is mean squared error(MSE).

    Parameters
    ----------

    Returns
    -------
    """
    Ys = matrices
    square_of_sum_label = values[0] ** 2
    if len(values) > 2:
        boosting = 1 + values[2]
    else:
        boosting = 1 + 1e-6

    party_opt = None
    feature_opt = None
    percentile_opt = None
    score_opt = numpy.finfo(numpy.dtype(float)).min

    # here we assume bins have equally number of samples
    confused = 0
    for i in range(len(Ys)):
        # split sum into left sum and right sum by cumsum trick
        if Ys[i].size > 0:
            # unique check
            flag_unique = numpy.alltrue(numpy.abs(Ys[i] - Ys[i][:, :1]) < 1e-8, axis=1)
            cumsum_Y = numpy.cumsum(Ys[i], axis=1)
            max_enumerate = cumsum_Y.shape[1] - 1
            for k in range(max_enumerate):
                E_zl = cumsum_Y[:, k] / (k + 1)
                E_zr = (cumsum_Y[:, -1] - cumsum_Y[:, k]) / (max_enumerate - k)
                score_i = (E_zl ** 2 * (k + 1) + E_zr ** 2 * (max_enumerate - k)) / max_enumerate
                # move out score_i by set them to -1e8
                score_i[flag_unique] = -1e8
                feature_opt_i = numpy.argmax(score_i)
                if score_i[feature_opt_i] > score_opt * (1 + 1e-8):
                    confused = 0
                    party_opt = i
                    feature_opt = feature_opt_i
                    percentile_opt = k
                    score_opt = score_i[feature_opt_i]
                elif abs(score_i[feature_opt_i] - score_opt) < 1e-8:
                    confused += 1

    if score_opt > (square_of_sum_label / cumsum_Y.shape[1]) * boosting:
        node_message = orjson.dumps({"is_leaf": 0})
        return dict(index=index,
                    values=[party_opt, feature_opt, percentile_opt, score_opt],
                    message=node_message)
    else:
        node_message = orjson.dumps({"is_leaf": 1})
        return dict(index=index,
                    values=[-1, 0, 0, 0],
                    message=node_message)
//  [python]
def Phase3Variance(matrices, values, index):
    """
    Best split selection code. The underlying loss function is mean squared error(MSE).

    Parameters
    ----------

    Returns
    -------
    """
    Ys = matrices
    square_of_sum_label = values[0] ** 2
    if len(values) > 2:
        boosting = 1 + values[2]
    else:
        boosting = 1 + 1e-6

    party_opt = None
    feature_opt = None
    percentile_opt = None
    score_opt = numpy.finfo(numpy.dtype(float)).min

    # here we assume bins have equally number of samples
    confused = 0
    for i in range(len(Ys)):
        # split sum into left sum and right sum by cumsum trick
        if Ys[i].size > 0:
            # unique check
            flag_unique = numpy.all(numpy.abs(Ys[i] - Ys[i][:, :1]) < 1e-8, axis=1)
            cumsum_Y = numpy.cumsum(Ys[i], axis=1)
            max_enumerate = cumsum_Y.shape[1] - 1
            for k in range(max_enumerate):
                E_zl = cumsum_Y[:, k] / (k + 1)
                E_zr = (cumsum_Y[:, -1] - cumsum_Y[:, k]) / (max_enumerate - k)
                score_i = (E_zl ** 2 * (k + 1) + E_zr ** 2 * (max_enumerate - k)) / max_enumerate
                # move out score_i by set them to -1e8
                score_i[flag_unique] = -1e8
                feature_opt_i = numpy.argmax(score_i)
                if score_i[feature_opt_i] > score_opt * (1 + 1e-8):
                    confused = 0
                    party_opt = i
                    feature_opt = feature_opt_i
                    percentile_opt = k
                    score_opt = score_i[feature_opt_i]
                elif abs(score_i[feature_opt_i] - score_opt) < 1e-8:
                    confused += 1

    if score_opt > (square_of_sum_label / cumsum_Y.shape[1]) * boosting:
        node_message = orjson.dumps({"is_leaf": 0})
        return dict(index=index,
                    values=[party_opt, feature_opt, percentile_opt, score_opt],
                    message=node_message)
    else:
        node_message = orjson.dumps({"is_leaf": 1})
        return dict(index=index,
                    values=[-1, 0, 0, 0],
                    message=node_message)
//  [python]
def test_register_simple_arr_diag(self):
        with _temporary_expr_dict():
            size = 10
            om.ExecComp.register('area', lambda x: x**2, complex_safe=True)
            p = om.Problem()
            p.model.add_subsystem('comp', om.ExecComp('area_square = area(x)', shape=size, has_diag_partials=True))
            p.setup()
            p['comp.x'] = 3.
            p.run_model()
            assert_near_equal(p['comp.area_square'], np.ones(size) * 9., 1e-11)
            J = p.compute_totals(of=['comp.area_square'], wrt=['comp.x'])
            assert_near_equal(J['comp.area_square', 'comp.x'], np.eye(size) * 6., 1e-11)

            # verify diagonal subjac
            self.assertTrue(n.all(p.model.comp._subjacs_info['comp.area_square', 'comp.x']['rows'] == np.arange(size)))
            self.assertTrue(n.all(p.model.comp._subjacs_info['comp.area_square', 'comp.x']['cols'] == np.arange(size)))
//  [python]
def test_register_simple_arr_diag(self):
        with _temporary_expr_dict():
            size = 10
            om.ExecComp.register('area', lambda x: x**2, complex_safe=True)
            p = om.Problem()
            p.model.add_subsystem('comp', om.ExecComp('area_square = area(x)', shape=size, has_diag_partials=True))
            p.setup()
            p['comp.x'] = 3.
            p.run_model()
            assert_near_equal(p['comp.area_square'], np.ones(size) * 9., 1e-11)
            J = p.compute_totals(of=['comp.area_square'], wrt=['comp.x'])
            assert_near_equal(J['comp.area_square', 'comp.x'], np.eye(size) * 6., 1e-11)

            # verify diagonal subjac
            self.assertTrue(np.all(p.model.comp._subjacs_info['comp.area_square', 'comp.x']['rows'] == np.arange(size)))
            self.assertTrue(np.all(p.model.comp._subjacs_info['comp.area_square', 'comp.x']['cols'] == np.arange(size)))
//  [python]
def test_manual_coloring(self):
        with _temporary_expr_dict():

            prob = om.Problem(coloring_dir=self.tempdir)
            model = prob.model

            sparsity = setup_sparsity(_MASK)

            def mydot(x):
                return sparsity.dot(x)

            om.ExecComp.register('mydot', mydot, complex_safe=True)

            comp = model.add_subsystem('comp', om.ExecComp('y=mydot(x)',
                                                            x=np.ones(sparsity.shape[1]),
                                                            y=np.ones(sparsity.shape[0])))
            comp.declare_coloring('x', method='cs')

            prob.setup(mode='fwd')
            prob.set_solver_print(level=0)
            prob.run_model()

            J = prob.compute_totals('comp.y', 'comp.x')

            assert_near_equal(J['comp.y', 'comp.x'], sparsity)

            self.assertTrue(n.all(comp._coloring_info['coloring'].get_dense_sparsity() == _MASK))
//  [python]
def test_manual_coloring(self):
        with _temporary_expr_dict():

            prob = om.Problem(coloring_dir=self.tempdir)
            model = prob.model

            sparsity = setup_sparsity(_MASK)

            def mydot(x):
                return sparsity.dot(x)

            om.ExecComp.register('mydot', mydot, complex_safe=True)

            comp = model.add_subsystem('comp', om.ExecComp('y=mydot(x)',
                                                            x=np.ones(sparsity.shape[1]),
                                                            y=np.ones(sparsity.shape[0])))
            comp.declare_coloring('x', method='cs')

            prob.setup(mode='fwd')
            prob.set_solver_print(level=0)
            prob.run_model()

            J = prob.compute_totals('comp.y', 'comp.x')

            assert_near_equal(J['comp.y', 'comp.x'], sparsity)

            self.assertTrue(np.all(comp._coloring_info['coloring'].get_dense_sparsity() == _MASK))
//  [python]
def test_auto_coloring(self):
        with _temporary_expr_dict():

            prob = om.Problem(coloring_dir=self.tempdir)
            model = prob.model

            sparsity = setup_sparsity(_MASK)

            def mydot(x):
                return sparsity.dot(x)

            om.ExecComp.register('mydot', mydot, complex_safe=True)

            comp = model.add_subsystem('comp', om.ExecComp('y=mydot(x)',
                                                            x=np.ones(sparsity.shape[1]),
                                                            y=np.ones(sparsity.shape[0])))
            prob.setup(mode='fwd')
            prob.set_solver_print(level=0)
            prob.run_model()

            J = prob.compute_totals('comp.y', 'comp.x')

            assert_near_equal(J['comp.y', 'comp.x'], sparsity)

            self.assertTrue(n.all(comp._coloring_info['coloring'].get_dense_sparsity() == _MASK))
//  [python]
def test_auto_coloring(self):
        with _temporary_expr_dict():

            prob = om.Problem(coloring_dir=self.tempdir)
            model = prob.model

            sparsity = setup_sparsity(_MASK)

            def mydot(x):
                return sparsity.dot(x)

            om.ExecComp.register('mydot', mydot, complex_safe=True)

            comp = model.add_subsystem('comp', om.ExecComp('y=mydot(x)',
                                                            x=np.ones(sparsity.shape[1]),
                                                            y=np.ones(sparsity.shape[0])))
            prob.setup(mode='fwd')
            prob.set_solver_print(level=0)
            prob.run_model()

            J = prob.compute_totals('comp.y', 'comp.x')

            assert_near_equal(J['comp.y', 'comp.x'], sparsity)

            self.assertTrue(np.all(comp._coloring_info['coloring'].get_dense_sparsity() == _MASK))
//  [python]
def get_optimal_reference_channel(
        w_mat,
        target_psd_matrix,
        noise_psd_matrix,
        eps=None,
):
    if w_mat.ndim != 3:
        raise ValueError(
            'Estimating the ref_channel expects currently that the input '
            'has 3 ndims (frequency x sensors x sensors). '
            'Considering an independent dim in the SNR estimate is not '
            'unique.'
        )
    if eps is None:
        eps = np.finfo(w_mat.dtype).tiny
    SNR = np.einsum(
        '...FdR,...FdD,...FDR->...R', w_mat.conj(), target_psd_matrix, w_mat
    ) / np.maximum(np.einsum(
        '...FdR,...FdD,...FDR->...R', w_mat.conj(), noise_psd_matrix, w_mat
    ), eps)
    # Raises an exception when np.inf and/or np.NaN was in target_psd_matrix
    # or noise_psd_matrix
    assert np.alltrue(np.isfinite(SNR)), SNR
    return np.argmax(SNR.real)
//  [python]
def get_optimal_reference_channel(
        w_mat,
        target_psd_matrix,
        noise_psd_matrix,
        eps=None,
):
    if w_mat.ndim != 3:
        raise ValueError(
            'Estimating the ref_channel expects currently that the input '
            'has 3 ndims (frequency x sensors x sensors). '
            'Considering an independent dim in the SNR estimate is not '
            'unique.'
        )
    if eps is None:
        eps = np.finfo(w_mat.dtype).tiny
    SNR = np.einsum(
        '...FdR,...FdD,...FDR->...R', w_mat.conj(), target_psd_matrix, w_mat
    ) / np.maximum(np.einsum(
        '...FdR,...FdD,...FDR->...R', w_mat.conj(), noise_psd_matrix, w_mat
    ), eps)
    # Raises an exception when np.inf and/or np.NaN was in target_psd_matrix
    # or noise_psd_matrix
    assert np.all(np.isfinite(SNR)), SNR
    return np.argmax(SNR.real)
//  [python]
def test_same_seed(self, measurements):
        """Test that different devices given the same random seed will produce
        the same results."""
        qs = qml.tape.QuantumScript([qml.THadamard(0)], measurements, shots=1000)

        dev1 = DefaultQutritMixed(seed=123)
        result1 = dev1.execute(qs)

        dev2 = DefaultQutritMixed(seed=123)
        result2 = dev2.execute(qs)

        if len(measurements) == 1:
            result1, result2 = [result1], [result2]

        assert all(np.all(res1 == res2) for res1, res2 in zip(result1, result2))
//  [python]
def test_global_seed_and_device_seed(self, measurements):
        """Test that a global seed does not affect the result of devices
        provided with a seed."""
        qs = qml.tape.QuantumScript([qml.THadamard(0)], measurements, shots=1000)

        # expected result
        dev1 = DefaultQutritMixed(seed=123)
        result1 = dev1.execute(qs)

        # set a global seed both before initialization of the
        # device and before execution of the tape
        np.random.seed(456)
        dev2 = DefaultQutritMixed(seed=123)
        np.random.seed(789)
        result2 = dev2.execute(qs)

        if len(measurements) == 1:
            result1, result2 = [result1], [result2]

        assert all(n.all(res1 == res2) for res1, res2 in zip(result1, result2))
//  [python]
def test_global_seed_and_device_seed(self, measurements):
        """Test that a global seed does not affect the result of devices
        provided with a seed."""
        qs = qml.tape.QuantumScript([qml.THadamard(0)], measurements, shots=1000)

        # expected result
        dev1 = DefaultQutritMixed(seed=123)
        result1 = dev1.execute(qs)

        # set a global seed both before initialization of the
        # device and before execution of the tape
        np.random.seed(456)
        dev2 = DefaultQutritMixed(seed=123)
        np.random.seed(789)
        result2 = dev2.execute(qs)

        if len(measurements) == 1:
            result1, result2 = [result1], [result2]

        assert all(np.all(res1 == res2) for res1, res2 in zip(result1, result2))
//  [python]
def test_prng_key_as_seed(self):
        """Test that a jax PRNG can be passed as a seed."""
        import jax

        jax.config.update("jax_enable_x64", True)

        from jax import random
        import numpy as np

        key1 = random.key(123)
        first_nums = random.uniform(key1, shape=(10,))

        key = random.key(123)
        dev = DefaultQutritMixed(seed=key)

        second_nums = random.uniform(dev._prng_key, shape=(10,))  # pylint: disable=protected-access
        assert np.all(first_nums == second_nums)
//  [python]
def test_prng_key_as_seed(self):
        """Test that a jax PRNG can be passed as a seed."""
        import jax

        jax.config.update("jax_enable_x64", True)

        from jax import random

        key1 = random.key(123)
        first_nums = random.uniform(key1, shape=(10,))

        key = random.key(123)
        dev = DefaultQutritMixed(seed=key)

        second_nums = random.uniform(dev._prng_key, shape=(10,))  # pylint: disable=protected-access
        assert np.all(first_nums == second_nums)
//  [python]
def test_same_device_prng_key(self):
        """Test a device with a given jax.random.PRNGKey will produce
        the same samples repeatedly."""
        import jax

        qs = qml.tape.QuantumScript([qml.THadamard(0)], [qml.sample(wires=0)], shots=1000)
        config = ExecutionConfig(interface="jax")

        dev = DefaultQutritMixed(seed=jax.random.PRNGKey(123))
        result1 = dev.execute(qs, config)
        for _ in range(10):
            result2 = dev.execute(qs, config)

            assert np.all(result1 == result2)
//  [python]
def test_same_prng_key(self):
        """Test that different devices given the same random jax.random.PRNGKey as a seed will produce
        the same results for sample, even with different seeds"""
        import jax

        qs = qml.tape.QuantumScript([qml.THadamard(0)], [qml.sample(wires=0)], shots=1000)
        config = ExecutionConfig(interface="jax")

        dev1 = DefaultQutritMixed(seed=jax.random.PRNGKey(123))
        result1 = dev1.execute(qs, config)

        dev2 = DefaultQutritMixed(seed=jax.random.PRNGKey(123))
        result2 = dev2.execute(qs, config)

        assert np.all(result1 == result2)
//  [python]
def test_different_executions_same_prng_key(self):
        """Test that the same device will produce the same results every execution if
        the seed is a jax.random.PRNGKey"""
        import jax

        qs = qml.tape.QuantumScript([qml.THadamard(0)], [qml.sample(wires=0)], shots=1000)
        config = ExecutionConfig(interface="jax")

        dev = DefaultQutritMixed(seed=jax.random.PRNGKey(77))
        result1 = dev.execute(qs, config)
        result2 = dev.execute(qs, config)

        assert n.all(result1 == result2)
//  [python]
def test_different_executions_same_prng_key(self):
        """Test that the same device will produce the same results every execution if
        the seed is a jax.random.PRNGKey"""
        import jax

        qs = qml.tape.QuantumScript([qml.THadamard(0)], [qml.sample(wires=0)], shots=1000)
        config = ExecutionConfig(interface="jax")

        dev = DefaultQutritMixed(seed=jax.random.PRNGKey(77))
        result1 = dev.execute(qs, config)
        result2 = dev.execute(qs, config)

        assert np.all(result1 == result2)
//  [python]
def score(self, X, y=None, group_id=None, top=None, type=None, denominator=None, group_weight=None, thread_count=-1):
        """
        Calculate NDCG@top
        Parameters
        ----------
        X : catboost.Pool or list or numpy.ndarray or pandas.DataFrame or pandas.Series
            Data to apply model on.
        y : list or numpy.ndarrays or pandas.DataFrame or pandas.Series
            True labels.
        group_id : list or numpy.ndarray or pandas.DataFrame or pandas.Series
            Ranking groups. If X is a Pool, group_id must be defined into X
        top : unsigned integer, up to `pow(2, 32) / 2 - 1`
            NDCG, Number of top-ranked objects to calculate NDCG
        type : str
            NDCG, Metric_type: 'Base' or 'Exp'
        denominator : str
            NDCG, Denominator type: 'LogPosition' or 'Position'
        group_weight : list or numpy.ndarray or pandas.DataFrame or pandas.Series
            Group weights.
        thread_count : int, optional (default=-1)
            Number of threads to work with.
        Returns
        -------
        NDCG@top : float
                   higher is better
        """
        def get_ndcg_metric_name(values, names):
            if np.all(np.equal(values, None)):
                return 'NDCG'
            return 'NDCG:' + ';'.join(['{}={}'.format(n, v) for v, n in zip(values, names) if v is not None])

        if isinstance(X, Pool):
            if y is not None:
                raise CatBoostError("Wrong initializing y: X is catboost.Pool object, y must be initialized inside catboost.Pool.")
            y = X.get_label()
            if group_id is not None:
                raise CatBoostError("Wrong initializing group_id: X is catboost.Pool object, group_id must be initialized inside catboost.Pool.")
            group_id = X.get_group_id_hash()

        if y is None:
            raise CatBoostError("y must be initialized.")
        if group_id is None:
            raise CatBoostError("group_id must be initialized. If groups are not expected, pass an array of zeros")

        predictions = self.predict(X)
        return _eval_metric_util([y], [predictions], get_ndcg_metric_name([top, type, denominator], ['top', 'type', 'denominator']), None, group_id, group_weight, None, None, thread_count)[0]
//  [python]
def all(a, axis=None, out=None, keepdims=None):
    if not bhary.check(a) and not bhary.check(out):
        return n.all(a, axis=axis, out=out)
    else:
        return ufuncs.logical_and.reduce(a.astype(bool), axis=axis, out=out)
//  [python]
def all(a, axis=None, out=None, keepdims=None):
    if not bhary.check(a) and not bhary.check(out):
        return numpy.all(a, axis=axis, out=out)
    else:
        return ufuncs.logical_and.reduce(a.astype(bool), axis=axis, out=out)
//  [python]
def test_bins(size=500, n_bins=10):
    columns = ['var1', 'var2']
    df = pandas.DataFrame(random.uniform(size=(size, 2)), columns=columns)
    x_limits = numpy.linspace(0, 1, n_bins + 1)[1:-1]
    bins = compute_bin_indices(df[columns].values, bin_limits=[x_limits, x_limits])
    assert n.all(0 <= bins) and n.all(bins < n_bins * n_bins), "the bins with wrong indices appeared"
//  [python]
def test_bins(size=500, n_bins=10):
    columns = ['var1', 'var2']
    df = pandas.DataFrame(random.uniform(size=(size, 2)), columns=columns)
    x_limits = numpy.linspace(0, 1, n_bins + 1)[1:-1]
    bins = compute_bin_indices(df[columns].values, bin_limits=[x_limits, x_limits])
    assert numpy.all(0 <= bins) and numpy.all(bins < n_bins * n_bins), "the bins with wrong indices appeared"
//  [python]
def test_same_seeds_should_produce_the_same_data(self):
		distargs = [None]*5
		T1, M_c = sdg.gen_data(self.cctypes_all_contiuous,
			self.n_rows,
			self.cols_to_views_good,
			self.cluster_weights_good,
			self.separation_good,
			seed=0,
			distargs=distargs)

		T2, M_c = sdg.gen_data(self.cctypes_all_contiuous,
			self.n_rows,
			self.cols_to_views_good,
			self.cluster_weights_good,
			self.separation_good,
			seed=0,
			distargs=distargs)

		A1 = numpy.array(T1)
		A2 = numpy.array(T2)

		assert numpy.all(A1==A2)
//  [python]
def test_different_seeds_should_produce_the_different_data(self):
		distargs = [None]*5
		T1, M_c = sdg.gen_data(self.cctypes_all_contiuous,
			self.n_rows,
			self.cols_to_views_good,
			self.cluster_weights_good,
			self.separation_good,
			seed=0,
			distargs=distargs)

		T2, M_c = sdg.gen_data(self.cctypes_all_contiuous,
			self.n_rows,
			self.cols_to_views_good,
			self.cluster_weights_good,
			self.separation_good,
			seed=12345,
			distargs=distargs)

		A1 = numpy.array(T1)
		A2 = numpy.array(T2)
		
		assert not numpy.all(A1==A2)
//  [python]
def test_should_have_nan_entries_if_specified(self):
		# for one column
		columns_list = [0]
		optargs = [dict(missing_data=1.0)] # every entry will be missing NaN
		X = sdg.predictive_columns(self.M_c, self.X_L, self.X_D, columns_list,
			optional_settings=optargs)

		assert n.all(n.isnan(X))
		
		# for two columns
		columns_list = [0,1]
		optargs = [dict(missing_data=1.0)]*2 
		X = sdg.predictive_columns(self.M_c, self.X_L, self.X_D, columns_list,
			optional_settings=optargs)

		assert n.all(n.isnan(X))

		# for one of two columns (no dict means no missing data)
		columns_list = [0,1]
		optargs = [dict(missing_data=1.0), None]
		X = sdg.predictive_columns(self.M_c, self.X_L, self.X_D, columns_list,
			optional_settings=optargs)

		assert n.all(n.isnan(X[:,0]))
		assert not n.any(n.isnan(X[:,1]))

		# for one of two columns. Missing data specified 0 for second column
		columns_list = [0,1]
		optargs = [dict(missing_data=1.0), dict(missing_data=0.0)]
		X = sdg.predictive_columns(self.M_c, self.X_L, self.X_D, columns_list,
			optional_settings=optargs)

		assert n.all(n.isnan(X[:,0]))
		assert not n.any(n.isnan(X[:,1]))
//  [python]
def test_should_have_nan_entries_if_specified(self):
		# for one column
		columns_list = [0]
		optargs = [dict(missing_data=1.0)] # every entry will be missing NaN
		X = sdg.predictive_columns(self.M_c, self.X_L, self.X_D, columns_list,
			optional_settings=optargs)

		assert numpy.all(numpy.isnan(X))
		
		# for two columns
		columns_list = [0,1]
		optargs = [dict(missing_data=1.0)]*2 
		X = sdg.predictive_columns(self.M_c, self.X_L, self.X_D, columns_list,
			optional_settings=optargs)

		assert numpy.all(numpy.isnan(X))

		# for one of two columns (no dict means no missing data)
		columns_list = [0,1]
		optargs = [dict(missing_data=1.0), None]
		X = sdg.predictive_columns(self.M_c, self.X_L, self.X_D, columns_list,
			optional_settings=optargs)

		assert numpy.all(numpy.isnan(X[:,0]))
		assert not numpy.any(numpy.isnan(X[:,1]))

		# for one of two columns. Missing data specified 0 for second column
		columns_list = [0,1]
		optargs = [dict(missing_data=1.0), dict(missing_data=0.0)]
		X = sdg.predictive_columns(self.M_c, self.X_L, self.X_D, columns_list,
			optional_settings=optargs)

		assert numpy.all(numpy.isnan(X[:,0]))
		assert not numpy.any(numpy.isnan(X[:,1]))
//  [python]
def make_incidence_matrix(self, L1, N1, L2, N2):
        """Return an N1+1 x N2+1 matrix that marks all L1 and L2 that are 8-connected

        L1 - a labels matrix
        N1 - # of labels in L1
        L2 - a labels matrix
        N2 - # of labels in L2

        L1 and L2 should have no overlap

        Returns a matrix where M[n,m] is true if there is some pixel in
        L1 with value n that is 8-connected to a pixel in L2 with value m
        """
        #
        # Overlay the two labels matrix
        #
        L = L1.copy()
        L[L2 != 0] = L2[L2 != 0] + N1
        neighbor_count, neighbor_index, n2 = centrosome.cpmorphology.find_neighbors(L)
        if n.all(neighbor_count == 0):
            return numpy.zeros((N1, N2), bool)
        #
        # Keep the neighbors of L1 / discard neighbors of L2
        #
        neighbor_count = neighbor_count[:N1]
        neighbor_index = neighbor_index[:N1]
        n2 = n2[: (neighbor_index[-1] + neighbor_count[-1])]
        #
        # Get rid of blanks
        #
        label = numpy.arange(N1)[neighbor_count > 0]
        neighbor_index = neighbor_index[neighbor_count > 0]
        neighbor_count = neighbor_count[neighbor_count > 0]
        #
        # Correct n2 because we have formerly added N1 to its labels. Make
        # it zero-based.
        #
        n2 -= N1 + 1
        #
        # Create runs of n1 labels
        #
        n1 = numpy.zeros(len(n2), int)
        n1[0] = label[0]
        n1[neighbor_index[1:]] = label[1:] - label[:-1]
        n1 = numpy.cumsum(n1)
        incidence = coo.coo_matrix(
            (numpy.ones(n1.shape), (n1, n2)), shape=(N1, N2)
        ).toarray()
        return incidence != 0
//  [python]
def make_incidence_matrix(self, L1, N1, L2, N2):
        """Return an N1+1 x N2+1 matrix that marks all L1 and L2 that are 8-connected

        L1 - a labels matrix
        N1 - # of labels in L1
        L2 - a labels matrix
        N2 - # of labels in L2

        L1 and L2 should have no overlap

        Returns a matrix where M[n,m] is true if there is some pixel in
        L1 with value n that is 8-connected to a pixel in L2 with value m
        """
        #
        # Overlay the two labels matrix
        #
        L = L1.copy()
        L[L2 != 0] = L2[L2 != 0] + N1
        neighbor_count, neighbor_index, n2 = centrosome.cpmorphology.find_neighbors(L)
        if numpy.all(neighbor_count == 0):
            return numpy.zeros((N1, N2), bool)
        #
        # Keep the neighbors of L1 / discard neighbors of L2
        #
        neighbor_count = neighbor_count[:N1]
        neighbor_index = neighbor_index[:N1]
        n2 = n2[: (neighbor_index[-1] + neighbor_count[-1])]
        #
        # Get rid of blanks
        #
        label = numpy.arange(N1)[neighbor_count > 0]
        neighbor_index = neighbor_index[neighbor_count > 0]
        neighbor_count = neighbor_count[neighbor_count > 0]
        #
        # Correct n2 because we have formerly added N1 to its labels. Make
        # it zero-based.
        #
        n2 -= N1 + 1
        #
        # Create runs of n1 labels
        #
        n1 = numpy.zeros(len(n2), int)
        n1[0] = label[0]
        n1[neighbor_index[1:]] = label[1:] - label[:-1]
        n1 = numpy.cumsum(n1)
        incidence = coo.coo_matrix(
            (numpy.ones(n1.shape), (n1, n2)), shape=(N1, N2)
        ).toarray()
        return incidence != 0
//  [python]
def test_rgb():
    numpy.random.seed(0)
    for combination in (
        (True, True, True),
        (True, True, False),
        (True, False, True),
        (True, False, False),
        (False, True, True),
        (False, True, False),
        (False, False, True),
    ):
        adjustments = numpy.random.uniform(size=7)
        images = [
            numpy.random.uniform(size=(10, 15)) if combination[i] else None
            for i in range(3)
        ]
        images += [None] * 4
        workspace, module = make_workspace(
            cellprofiler.modules.graytocolor.SCHEME_RGB, images, adjustments
        )
        module.run(workspace)
        image = workspace.image_set.get_image(OUTPUT_IMAGE_NAME)
        pixel_data = image.pixel_data

        expected = numpy.dstack(
            [
                image if image is not None else numpy.zeros((10, 15))
                for image in images[:3]
            ]
        )
        for i in range(3):
            expected[:, :, i] *= adjustments[i]
        assert n.all(n.abs(expected - pixel_data) <= 0.00001)
//  [python]
def test_rgb():
    numpy.random.seed(0)
    for combination in (
        (True, True, True),
        (True, True, False),
        (True, False, True),
        (True, False, False),
        (False, True, True),
        (False, True, False),
        (False, False, True),
    ):
        adjustments = numpy.random.uniform(size=7)
        images = [
            numpy.random.uniform(size=(10, 15)) if combination[i] else None
            for i in range(3)
        ]
        images += [None] * 4
        workspace, module = make_workspace(
            cellprofiler.modules.graytocolor.SCHEME_RGB, images, adjustments
        )
        module.run(workspace)
        image = workspace.image_set.get_image(OUTPUT_IMAGE_NAME)
        pixel_data = image.pixel_data

        expected = numpy.dstack(
            [
                image if image is not None else numpy.zeros((10, 15))
                for image in images[:3]
            ]
        )
        for i in range(3):
            expected[:, :, i] *= adjustments[i]
        assert numpy.all(numpy.abs(expected - pixel_data) <= 0.00001)
//  [python]
def test_cmyk():
    numpy.random.seed(0)
    for combination in [[(i & 2 ^ j) != 0 for j in range(4)] for i in range(1, 16)]:
        adjustments = numpy.random.uniform(size=7)
        images = [
            numpy.random.uniform(size=(10, 15)) if combination[i] else None
            for i in range(4)
        ]
        images = [None] * 3 + images
        workspace, module = make_workspace(
            cellprofiler.modules.graytocolor.SCHEME_CMYK, images, adjustments
        )
        module.run(workspace)
        image = workspace.image_set.get_image(OUTPUT_IMAGE_NAME)
        pixel_data = image.pixel_data

        expected = numpy.array(
            [
                numpy.dstack(
                    [image * adjustment if image is not None else numpy.zeros((10, 15))]
                    * 3
                )
                * numpy.array(multiplier)
                / numpy.sum(multiplier)
                for image, multiplier, adjustment in (
                    (images[3], (0, 1, 1), adjustments[3]),
                    (images[4], (1, 0, 1), adjustments[4]),
                    (images[5], (1, 1, 0), adjustments[5]),
                    (images[6], (1, 1, 1), adjustments[6]),
                )
            ]
        )
        expected = numpy.sum(expected, 0)
        assert n.all(n.abs(expected - pixel_data) <= 0.00001)
//  [python]
def test_cmyk():
    numpy.random.seed(0)
    for combination in [[(i & 2 ^ j) != 0 for j in range(4)] for i in range(1, 16)]:
        adjustments = numpy.random.uniform(size=7)
        images = [
            numpy.random.uniform(size=(10, 15)) if combination[i] else None
            for i in range(4)
        ]
        images = [None] * 3 + images
        workspace, module = make_workspace(
            cellprofiler.modules.graytocolor.SCHEME_CMYK, images, adjustments
        )
        module.run(workspace)
        image = workspace.image_set.get_image(OUTPUT_IMAGE_NAME)
        pixel_data = image.pixel_data

        expected = numpy.array(
            [
                numpy.dstack(
                    [image * adjustment if image is not None else numpy.zeros((10, 15))]
                    * 3
                )
                * numpy.array(multiplier)
                / numpy.sum(multiplier)
                for image, multiplier, adjustment in (
                    (images[3], (0, 1, 1), adjustments[3]),
                    (images[4], (1, 0, 1), adjustments[4]),
                    (images[5], (1, 1, 0), adjustments[5]),
                    (images[6], (1, 1, 1), adjustments[6]),
                )
            ]
        )
        expected = numpy.sum(expected, 0)
        assert numpy.all(numpy.abs(expected - pixel_data) <= 0.00001)
//  [python]
def load_mice():
  """Loads the Mice dataset, adapted from: https://github.com/lasso-net/lassonet/blob/master/experiments/data_utils.py."""

  filling_value = -100000
  cache_filepath = os.path.join(DATA_DIR, "mice/Data_Cortex_Nuclear.csv")

  with open(cache_filepath, "r") as fp:
    x = np.genfromtxt(
        fp.readlines(),
        delimiter=",",
        skip_header=1,
        usecols=range(1, 78),
        filling_values=filling_value,
        encoding="UTF-8",
    )
  with open(cache_filepath, "r") as fp:
    classes = np.genfromtxt(
        fp.readlines(),
        delimiter=",",
        skip_header=1,
        usecols=range(78, 81),
        dtype=None,
        encoding="UTF-8",
    )

  for i, row in enumerate(x):
    for j, val in enumerate(row):
      if val == filling_value:
        x[i, j] = np.mean(
            [
                x[k, j]
                for k in range(classes.shape[0])
                if np.all(classes[i] == classes[k])
            ]
        )

  y = np.zeros((classes.shape[0]), dtype=np.uint8)
  for i, row in enumerate(classes):
    for j, (val, label) in enumerate(zip(row, ["Control", "Memantine", "C/S"])):
      y[i] += (2**j) * (val == label)

  x = MinMaxScaler(feature_range=(0, 1)).fit_transform(x)

  indices = np.arange(x.shape[0])
  np.random.shuffle(indices)
  x = x[indices]
  y = y[indices]
  classes = classes[indices]

  print("Data loaded...")
  print("Data shapes:")
  print("x shape: {}, y shape: {}".format(x.shape, y.shape))

  is_classification = True
  num_classes = 8

  x_train = pd.DataFrame(x[: x.shape[0] * 4 // 5])
  x_test = pd.DataFrame(x[x.shape[0] * 4 // 5 :])
  y_train = pd.DataFrame(y[: x.shape[0] * 4 // 5], dtype=np.int32).iloc[:, 0]
  y_test = pd.DataFrame(y[x.shape[0] * 4 // 5 :], dtype=np.int32).iloc[:, 0]

  return (x_train, x_test, y_train, y_test, is_classification, num_classes)
//  [python]
def test_rgb_rescale():
    numpy.random.seed(0)
    for combination in (
        (True, True, True),
        (True, True, False),
        (True, False, True),
        (True, False, False),
        (False, True, True),
        (False, True, False),
        (False, False, True),
    ):
        adjustments = numpy.random.uniform(size=7)
        images = [
            numpy.random.uniform(size=(10, 15)) if combination[i] else None
            for i in range(3)
        ]
        images += [None] * 4
        workspace, module = make_workspace(
            cellprofiler.modules.graytocolor.SCHEME_RGB, images, adjustments
        )
        module.wants_rescale.value = True
        module.run(workspace)
        image = workspace.image_set.get_image(OUTPUT_IMAGE_NAME)
        pixel_data = image.pixel_data

        expected = numpy.dstack(
            [
                image if image is not None else numpy.zeros((10, 15))
                for image in images[:3]
            ]
        )
        for i in range(3):
            plane = expected[:, :, i]
            if plane.max() > 0:
                plane = plane / numpy.max(plane)
            plane *= adjustments[i]
            expected[:, :, i] = plane
        assert n.all(n.abs(expected - pixel_data) <= 0.00001)
//  [python]
def test_rgb_rescale():
    numpy.random.seed(0)
    for combination in (
        (True, True, True),
        (True, True, False),
        (True, False, True),
        (True, False, False),
        (False, True, True),
        (False, True, False),
        (False, False, True),
    ):
        adjustments = numpy.random.uniform(size=7)
        images = [
            numpy.random.uniform(size=(10, 15)) if combination[i] else None
            for i in range(3)
        ]
        images += [None] * 4
        workspace, module = make_workspace(
            cellprofiler.modules.graytocolor.SCHEME_RGB, images, adjustments
        )
        module.wants_rescale.value = True
        module.run(workspace)
        image = workspace.image_set.get_image(OUTPUT_IMAGE_NAME)
        pixel_data = image.pixel_data

        expected = numpy.dstack(
            [
                image if image is not None else numpy.zeros((10, 15))
                for image in images[:3]
            ]
        )
        for i in range(3):
            plane = expected[:, :, i]
            if plane.max() > 0:
                plane = plane / numpy.max(plane)
            plane *= adjustments[i]
            expected[:, :, i] = plane
        assert numpy.all(numpy.abs(expected - pixel_data) <= 0.00001)
//  [python]
def test_cmyk_rescale():
    numpy.random.seed(0)
    for combination in [[(i & 2 ^ j) != 0 for j in range(4)] for i in range(1, 16)]:
        adjustments = numpy.random.uniform(size=7)
        images = [
            numpy.random.uniform(size=(10, 15)) if combination[i] else None
            for i in range(4)
        ]
        images = [None] * 3 + images
        workspace, module = make_workspace(
            cellprofiler.modules.graytocolor.SCHEME_CMYK, images, adjustments
        )
        module.wants_rescale.value = True
        module.run(workspace)
        image = workspace.image_set.get_image(OUTPUT_IMAGE_NAME)
        pixel_data = image.pixel_data

        expected = numpy.array(
            [
                numpy.dstack(
                    [(image/numpy.max(image)) * adjustment if image is not None else numpy.zeros((10, 15))]
                    * 3
                )
                * numpy.array(multiplier)
                / numpy.sum(multiplier)
                for image, multiplier, adjustment in (
                    (images[3], (0, 1, 1), adjustments[3]),
                    (images[4], (1, 0, 1), adjustments[4]),
                    (images[5], (1, 1, 0), adjustments[5]),
                    (images[6], (1, 1, 1), adjustments[6]),
                )
            ]
        )
        expected = numpy.sum(expected, 0)
        assert n.all(n.abs(expected - pixel_data) <= 0.00001)
//  [python]
def test_cmyk_rescale():
    numpy.random.seed(0)
    for combination in [[(i & 2 ^ j) != 0 for j in range(4)] for i in range(1, 16)]:
        adjustments = numpy.random.uniform(size=7)
        images = [
            numpy.random.uniform(size=(10, 15)) if combination[i] else None
            for i in range(4)
        ]
        images = [None] * 3 + images
        workspace, module = make_workspace(
            cellprofiler.modules.graytocolor.SCHEME_CMYK, images, adjustments
        )
        module.wants_rescale.value = True
        module.run(workspace)
        image = workspace.image_set.get_image(OUTPUT_IMAGE_NAME)
        pixel_data = image.pixel_data

        expected = numpy.array(
            [
                numpy.dstack(
                    [(image/numpy.max(image)) * adjustment if image is not None else numpy.zeros((10, 15))]
                    * 3
                )
                * numpy.array(multiplier)
                / numpy.sum(multiplier)
                for image, multiplier, adjustment in (
                    (images[3], (0, 1, 1), adjustments[3]),
                    (images[4], (1, 0, 1), adjustments[4]),
                    (images[5], (1, 1, 0), adjustments[5]),
                    (images[6], (1, 1, 1), adjustments[6]),
                )
            ]
        )
        expected = numpy.sum(expected, 0)
        assert numpy.all(numpy.abs(expected - pixel_data) <= 0.00001)
//  [python]
def get_permuted_map(V):
    """
    Return a PermutedMap with the same tensor product shape for
    every component of H(div) or H(curl) tensor product elements
    """
    indices, _, _ = get_permutation_to_line_elements(V)
    if n.all(indices[:-1] < indices[1:]):
        return V.cell_node_map()
    return op2.PermutedMap(V.cell_node_map(), indices)
//  [python]
def get_permuted_map(V):
    """
    Return a PermutedMap with the same tensor product shape for
    every component of H(div) or H(curl) tensor product elements
    """
    indices, _, _ = get_permutation_to_line_elements(V)
    if numpy.all(indices[:-1] < indices[1:]):
        return V.cell_node_map()
    return op2.PermutedMap(V.cell_node_map(), indices)
//  [python]
def get_jk(self, dm, hermi=1, kpts=None, kpts_band=None,
               with_j=True, with_k=True, exxdiv='ewald', **kwargs):
        from pyscf.pbc.df import fft_jk
        if with_k:
            logger.warn(self, 'MultiGridFFTDF does not support HFX. '
                        'HFX is computed by FFTDF.get_k_kpts function.')

        if kpts is None:
            if n.all(self.kpts == 0): # Gamma-point J/K by default
                kpts = numpy.zeros(3)
            else:
                kpts = self.kpts
        else:
            kpts = numpy.asarray(kpts)

        vj = vk = None
        if kpts.shape == (3,):
            if with_k:
                vk = fft_jk.get_jk(self, dm, hermi, kpts, kpts_band,
                                   False, True, exxdiv)[1]
            vj = get_j_kpts(self, dm, hermi, kpts.reshape(1,3), kpts_band)
            if kpts_band is None:
                vj = vj[...,0,:,:]
        else:
            if with_k:
                vk = fft_jk.get_k_kpts(self, dm, hermi, kpts, kpts_band, exxdiv)
            if with_j:
                vj = get_j_kpts(self, dm, hermi, kpts, kpts_band)
        return vj, vk
//  [python]
def get_jk(self, dm, hermi=1, kpts=None, kpts_band=None,
               with_j=True, with_k=True, exxdiv='ewald', **kwargs):
        from pyscf.pbc.df import fft_jk
        if with_k:
            logger.warn(self, 'MultiGridFFTDF does not support HFX. '
                        'HFX is computed by FFTDF.get_k_kpts function.')

        if kpts is None:
            if numpy.all(self.kpts == 0): # Gamma-point J/K by default
                kpts = numpy.zeros(3)
            else:
                kpts = self.kpts
        else:
            kpts = numpy.asarray(kpts)

        vj = vk = None
        if kpts.shape == (3,):
            if with_k:
                vk = fft_jk.get_jk(self, dm, hermi, kpts, kpts_band,
                                   False, True, exxdiv)[1]
            vj = get_j_kpts(self, dm, hermi, kpts.reshape(1,3), kpts_band)
            if kpts_band is None:
                vj = vj[...,0,:,:]
        else:
            if with_k:
                vk = fft_jk.get_k_kpts(self, dm, hermi, kpts, kpts_band, exxdiv)
            if with_j:
                vj = get_j_kpts(self, dm, hermi, kpts, kpts_band)
        return vj, vk
//  [python]
def __call__(self, model_fn, x_np, y_np):
    # randomize each example separately

    with self.graph.as_default():
      result = np.zeros(x_np.shape, dtype=x_np.dtype)
      did = np.zeros(x_np.shape[0], dtype=np.bool)

      trans_np = np.stack(
        repeat([0, 0, 0], x_np.shape[0]))
      x_downsize_np = self.session.run(self._tranformed_x_op, feed_dict={
        self._x_for_trans: x_np,
        self._t_for_trans: trans_np,
      })

      while True:
        random_transforms = (np.random.uniform(-lim, lim, len(x_np)) for lim in self.limits)
        trans_np = np.stack(random_transforms, axis=1)
        out = self.session.run(self._tranformed_x_op, feed_dict={
          self._x_for_trans: x_np,
          self._t_for_trans: trans_np,
        })

        if self.valid_check is None:
          return out
        else:
          ok = self.valid_check(x_downsize_np, out)
          result[ok] = out[ok]
          did[ok] = True
          if n.all(did):
            return result
//  [python]
def __call__(self, model_fn, x_np, y_np):
    # randomize each example separately

    with self.graph.as_default():
      result = np.zeros(x_np.shape, dtype=x_np.dtype)
      did = np.zeros(x_np.shape[0], dtype=np.bool)

      trans_np = np.stack(
        repeat([0, 0, 0], x_np.shape[0]))
      x_downsize_np = self.session.run(self._tranformed_x_op, feed_dict={
        self._x_for_trans: x_np,
        self._t_for_trans: trans_np,
      })

      while True:
        random_transforms = (np.random.uniform(-lim, lim, len(x_np)) for lim in self.limits)
        trans_np = np.stack(random_transforms, axis=1)
        out = self.session.run(self._tranformed_x_op, feed_dict={
          self._x_for_trans: x_np,
          self._t_for_trans: trans_np,
        })

        if self.valid_check is None:
          return out
        else:
          ok = self.valid_check(x_downsize_np, out)
          result[ok] = out[ok]
          did[ok] = True
          if np.all(did):
            return result
//  [python]
def numpy_equals(a, b):
    import numpy as np

    if a is b:
        return True
    if a is None or b is None:
        return False
    if np.alltrue(a == b):
        return True
    return False
//  [python]
def numpy_equals(a, b):
    import numpy as np

    if a is b:
        return True
    if a is None or b is None:
        return False
    if np.all(a == b):
        return True
    return False