jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def check_dist_emb(g, num_clients, num_nodes, num_edges):
    # Test sparse emb
    try:
        emb = DistEmbedding(g.num_nodes(), 1, "emb1", emb_init)
        nids = F.arange(0, int(g.num_nodes()))
        lr = 0.001
        optimizer = SparseAdagrad([emb], lr=lr)
        with F.record_grad():
            feats = emb(nids)
            assert np.all(F.asnumpy(feats) == np.zeros((len(nids), 1)))
            loss = F.sum(feats + 1, 0)
        loss.backward()
        optimizer.step()
        feats = emb(nids)
        if num_clients == 1:
            assert_almost_equal(F.asnumpy(feats), np.ones((len(nids), 1)) * -lr)
        rest = np.setdiff1d(np.arange(g.num_nodes()), F.asnumpy(nids))
        feats1 = emb(rest)
        assert np.all(F.asnumpy(feats1) == np.zeros((len(rest), 1)))

        policy = dgl.distributed.PartitionPolicy("node", g.get_partition_book())
        grad_sum = dgl.distributed.DistTensor(
            (g.num_nodes(), 1), F.float32, "emb1_sum", policy
        )
        if num_clients == 1:
            assert np.all(
                F.asnumpy(grad_sum[nids])
                == np.ones((len(nids), 1)) * num_clients
            )
        assert np.all(F.asnumpy(grad_sum[rest]) == np.zeros((len(rest), 1)))

        emb = DistEmbedding(g.num_nodes(), 1, "emb2", emb_init)
        with F.no_grad():
            feats1 = emb(nids)
        assert np.all(F.asnumpy(feats1) == 0)

        optimizer = SparseAdagrad([emb], lr=lr)
        with F.record_grad():
            feats1 = emb(nids)
            feats2 = emb(nids)
            feats = F.cat([feats1, feats2], 0)
            assert np.all(F.asnumpy(feats) == np.zeros((len(nids) * 2, 1)))
            loss = F.sum(feats + 1, 0)
        loss.backward()
        optimizer.step()
        with F.no_grad():
            feats = emb(nids)
        if num_clients == 1:
            assert_almost_equal(
                F.asnumpy(feats), np.ones((len(nids), 1)) * 1 * -lr
            )
        rest = np.setdiff1d(np.arange(g.num_nodes()), F.asnumpy(nids))
        feats1 = emb(rest)
        assert np.all(F.asnumpy(feats1) == np.zeros((len(rest), 1)))
    except NotImplementedError as e:
        pass
    except Exception as e:
        print(e)
        sys.exit(-1)
//  [python]
def coerce_to_array(
    values, mask=None, copy: bool = False
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Coerce the input values array to numpy arrays with a mask.

    Parameters
    ----------
    values : 1D list-like
    mask : bool 1D array, optional
    copy : bool, default False
        if True, copy the input

    Returns
    -------
    tuple of (values, mask)
    """
    if isinstance(values, BooleanArray):
        if mask is not None:
            raise ValueError("cannot pass mask for BooleanArray input")
        values, mask = values._data, values._mask
        if copy:
            values = values.copy()
            mask = mask.copy()
        return values, mask

    mask_values = None
    if isinstance(values, np.ndarray) and values.dtype == np.bool_:
        if copy:
            values = values.copy()
    elif isinstance(values, np.ndarray) and is_numeric_dtype(values.dtype):
        mask_values = isna(values)

        values_bool = np.zeros(len(values), dtype=bool)
        values_bool[~mask_values] = values[~mask_values].astype(bool)

        if not np.alltrue(
            values_bool[~mask_values].astype(values.dtype) == values[~mask_values]
        ):
            raise TypeError("Need to pass bool-like values")

        values = values_bool
    else:
        values_object = np.asarray(values, dtype=object)

        inferred_dtype = lib.infer_dtype(values_object, skipna=True)
        integer_like = ("floating", "integer", "mixed-integer-float")
        if inferred_dtype not in ("boolean", "empty") + integer_like:
            raise TypeError("Need to pass bool-like values")

        mask_values = isna(values_object)
        values = np.zeros(len(values), dtype=bool)
        values[~mask_values] = values_object[~mask_values].astype(bool)

        # if the values were integer-like, validate it were actually 0/1's
        if inferred_dtype in integer_like:
            if not np.alltrue(
                values[~mask_values].astype(float)
                == values_object[~mask_values].astype(float)
            ):
                raise TypeError("Need to pass bool-like values")

    if mask is None and mask_values is None:
        mask = np.zeros(len(values), dtype=bool)
    elif mask is None:
        mask = mask_values
    else:
        if isinstance(mask, np.ndarray) and mask.dtype == np.bool_:
            if mask_values is not None:
                mask = mask | mask_values
            else:
                if copy:
                    mask = mask.copy()
        else:
            mask = np.array(mask, dtype=bool)
            if mask_values is not None:
                mask = mask | mask_values

    if not values.ndim == 1:
        raise ValueError("values must be a 1D list-like")
    if not mask.ndim == 1:
        raise ValueError("mask must be a 1D list-like")

    return values, mask
//  [python]
def coerce_to_array(
    values, mask=None, copy: bool = False
) -> Tuple[np.ndarray, np.ndarray]:
    """
    Coerce the input values array to numpy arrays with a mask.

    Parameters
    ----------
    values : 1D list-like
    mask : bool 1D array, optional
    copy : bool, default False
        if True, copy the input

    Returns
    -------
    tuple of (values, mask)
    """
    if isinstance(values, BooleanArray):
        if mask is not None:
            raise ValueError("cannot pass mask for BooleanArray input")
        values, mask = values._data, values._mask
        if copy:
            values = values.copy()
            mask = mask.copy()
        return values, mask

    mask_values = None
    if isinstance(values, np.ndarray) and values.dtype == np.bool_:
        if copy:
            values = values.copy()
    elif isinstance(values, np.ndarray) and is_numeric_dtype(values.dtype):
        mask_values = isna(values)

        values_bool = np.zeros(len(values), dtype=bool)
        values_bool[~mask_values] = values[~mask_values].astype(bool)

        if not np.all(
            values_bool[~mask_values].astype(values.dtype) == values[~mask_values]
        ):
            raise TypeError("Need to pass bool-like values")

        values = values_bool
    else:
        values_object = np.asarray(values, dtype=object)

        inferred_dtype = lib.infer_dtype(values_object, skipna=True)
        integer_like = ("floating", "integer", "mixed-integer-float")
        if inferred_dtype not in ("boolean", "empty") + integer_like:
            raise TypeError("Need to pass bool-like values")

        mask_values = isna(values_object)
        values = np.zeros(len(values), dtype=bool)
        values[~mask_values] = values_object[~mask_values].astype(bool)

        # if the values were integer-like, validate it were actually 0/1's
        if inferred_dtype in integer_like:
            if not np.all(
                values[~mask_values].astype(float)
                == values_object[~mask_values].astype(float)
            ):
                raise TypeError("Need to pass bool-like values")

    if mask is None and mask_values is None:
        mask = np.zeros(len(values), dtype=bool)
    elif mask is None:
        mask = mask_values
    else:
        if isinstance(mask, np.ndarray) and mask.dtype == np.bool_:
            if mask_values is not None:
                mask = mask | mask_values
            else:
                if copy:
                    mask = mask.copy()
        else:
            mask = np.array(mask, dtype=bool)
            if mask_values is not None:
                mask = mask | mask_values

    if not values.ndim == 1:
        raise ValueError("values must be a 1D list-like")
    if not mask.ndim == 1:
        raise ValueError("mask must be a 1D list-like")

    return values, mask
//  [python]
def dlqr(A, B, Q, R):
    """
    Calculate the discrete-time steady-state LQR gain matrix.
    Minimize sum{0, inf}(x'Qx + u'Ru) for the system x(n+1) = Ax(n) + Bu(n).
    Args:
        A: n*n discrete-time system dynamics matrix
        B: n*m discrete-time control signal matrix
        Q: n*n quadratic state error weighting factor
        R: n*n quadratic control signal weighting factor
    Returns:
        K: m*n gains matrix such that u = -Kx
    """
    A = numpy.asmatrix(A)
    B = numpy.asmatrix(B)
    Q = numpy.asmatrix(Q)
    R = numpy.asmatrix(R)

    _validate_system(A, B, None, None)
    assert Q.shape[0] == Q.shape[1] and Q.shape[0] == A.shape[0], "The dimensions of Q %s must match those of A %s" % (Q.shape, A.shape)
    assert R.shape[0] == R.shape[1] and R.shape[0] == B.shape[1], "R %i must be square and must be compatible with B %i" % (R.shape, B.shape)

    assert numpy.linalg.matrix_rank(controllability(A, B)) == A.shape[0], "System must be completely controllable to do LQR."

    # Ensure Q is positive-semidefninite and R is positive-definite
    Q_eigvals = numpy.linalg.eigvalsh(Q)
    assert n.all(Q_eigvals > -1e-8), "Q must be positive-semidefinite"

    R_eigvals = numpy.linalg.eigvalsh(R)
    assert n.all(R_eigvals > 0), "R must be positive-definite"

    # Solve the ARE for the cost-to-go matrix
    M = numpy.asmatrix(scipy.linalg.solve_discrete_are(A, B, Q, R))

    # Finally, solve for the optimal gain matrix using the cost-to-go matrix
    return numpy.linalg.inv(R + B.T * M * B) * B.T * M * A
//  [python]
def dlqr(A, B, Q, R):
    """
    Calculate the discrete-time steady-state LQR gain matrix.
    Minimize sum{0, inf}(x'Qx + u'Ru) for the system x(n+1) = Ax(n) + Bu(n).
    Args:
        A: n*n discrete-time system dynamics matrix
        B: n*m discrete-time control signal matrix
        Q: n*n quadratic state error weighting factor
        R: n*n quadratic control signal weighting factor
    Returns:
        K: m*n gains matrix such that u = -Kx
    """
    A = numpy.asmatrix(A)
    B = numpy.asmatrix(B)
    Q = numpy.asmatrix(Q)
    R = numpy.asmatrix(R)

    _validate_system(A, B, None, None)
    assert Q.shape[0] == Q.shape[1] and Q.shape[0] == A.shape[0], "The dimensions of Q %s must match those of A %s" % (Q.shape, A.shape)
    assert R.shape[0] == R.shape[1] and R.shape[0] == B.shape[1], "R %i must be square and must be compatible with B %i" % (R.shape, B.shape)

    assert numpy.linalg.matrix_rank(controllability(A, B)) == A.shape[0], "System must be completely controllable to do LQR."

    # Ensure Q is positive-semidefninite and R is positive-definite
    Q_eigvals = numpy.linalg.eigvalsh(Q)
    assert numpy.all(Q_eigvals > -1e-8), "Q must be positive-semidefinite"

    R_eigvals = numpy.linalg.eigvalsh(R)
    assert numpy.all(R_eigvals > 0), "R must be positive-definite"

    # Solve the ARE for the cost-to-go matrix
    M = numpy.asmatrix(scipy.linalg.solve_discrete_are(A, B, Q, R))

    # Finally, solve for the optimal gain matrix using the cost-to-go matrix
    return numpy.linalg.inv(R + B.T * M * B) * B.T * M * A
//  [python]
def test_input_maintained(self):
        '''Test to make sure the input is maintained
        '''
        dtype_tuple = input_dtypes[functions[self.func]]
        for dtype in dtype_tuple[0]:
            for test_shape, s, kwargs in self.test_data:

                input_array = dtype_tuple[1](test_shape, dtype)

                FFTW_object = getattr(
                        builders, self.func)(input_array, s, **kwargs)

                final_input_array = FFTW_object.input_array

                self.assertTrue(
                        n.all(input_array == final_input_array))
//  [python]
def test_input_maintained(self):
        '''Test to make sure the input is maintained
        '''
        dtype_tuple = input_dtypes[functions[self.func]]
        for dtype in dtype_tuple[0]:
            for test_shape, s, kwargs in self.test_data:

                input_array = dtype_tuple[1](test_shape, dtype)

                FFTW_object = getattr(
                        builders, self.func)(input_array, s, **kwargs)

                final_input_array = FFTW_object.input_array

                self.assertTrue(
                        numpy.all(input_array == final_input_array))
//  [python]
def test_call(self):
        '''Test a call to an instance of the class.
        '''

        self.input_array[:] = (numpy.random.randn(*self.input_array.shape)
                + 1j*numpy.random.randn(*self.input_array.shape))

        output_array = self.fft()

        self.assertTrue(numpy.all(output_array == self.output_array))
//  [python]
def clqr(A, B, Q, R):
    """
    Calculate the continuous-time steady-state LQR gain matrix.
    Minimize integral{0, inf}(x'Qx + u'Ru) for the system dx/dt = Ax + Bu.
    Args:
        A: n*n continuous-time system dynamics matrix
        B: n*m continuous-time control signal matrix
        Q: n*n quadratic state error weighting factor
        R: n*n quadratic control signal weighting factor
    Returns:
        K: m*n gain matrix such that u = -Kx
    """
    A = numpy.asmatrix(A)
    B = numpy.asmatrix(B)
    Q = numpy.asmatrix(Q)
    R = numpy.asmatrix(R)

    _validate_system(A, B, None, None)
    assert Q.shape[0] == Q.shape[1] and Q.shape[0] == A.shape[0], "The dimensions of Q %s must match those of A %s" % (Q.shape, A.shape)
    assert R.shape[0] == R.shape[1] and R.shape[0] == B.shape[1], "R %i must be square and must be compatible with B %i" % (R.shape, B.shape)

    assert numpy.linalg.matrix_rank(controllability(A, B)) == A.shape[0], "System must be completely controllable to do LQR."

    # Ensure Q is positive-semidefninite and R is positive-definite
    Q_eigvals = numpy.linalg.eigvalsh(Q)
    # Allow for some tolerance on Q to avoid numerical stability issues
    assert n.all(Q_eigvals > -1e-8), "Q must be positive-semidefinite"

    R_eigvals = numpy.linalg.eigvalsh(R)
    assert n.all(R_eigvals > 0), "R must be positive-definite"

    # Solve the ARE for the cost-to-go matrix
    M = numpy.asmatrix(scipy.linalg.solve_continuous_are(A, B, Q, R))

    # Finally, solve for the optimal gain matrix using the cost-to-go matrix
    return numpy.asmatrix(numpy.linalg.inv(R) * B.T * M)
//  [python]
def clqr(A, B, Q, R):
    """
    Calculate the continuous-time steady-state LQR gain matrix.
    Minimize integral{0, inf}(x'Qx + u'Ru) for the system dx/dt = Ax + Bu.
    Args:
        A: n*n continuous-time system dynamics matrix
        B: n*m continuous-time control signal matrix
        Q: n*n quadratic state error weighting factor
        R: n*n quadratic control signal weighting factor
    Returns:
        K: m*n gain matrix such that u = -Kx
    """
    A = numpy.asmatrix(A)
    B = numpy.asmatrix(B)
    Q = numpy.asmatrix(Q)
    R = numpy.asmatrix(R)

    _validate_system(A, B, None, None)
    assert Q.shape[0] == Q.shape[1] and Q.shape[0] == A.shape[0], "The dimensions of Q %s must match those of A %s" % (Q.shape, A.shape)
    assert R.shape[0] == R.shape[1] and R.shape[0] == B.shape[1], "R %i must be square and must be compatible with B %i" % (R.shape, B.shape)

    assert numpy.linalg.matrix_rank(controllability(A, B)) == A.shape[0], "System must be completely controllable to do LQR."

    # Ensure Q is positive-semidefninite and R is positive-definite
    Q_eigvals = numpy.linalg.eigvalsh(Q)
    # Allow for some tolerance on Q to avoid numerical stability issues
    assert numpy.all(Q_eigvals > -1e-8), "Q must be positive-semidefinite"

    R_eigvals = numpy.linalg.eigvalsh(R)
    assert numpy.all(R_eigvals > 0), "R must be positive-definite"

    # Solve the ARE for the cost-to-go matrix
    M = numpy.asmatrix(scipy.linalg.solve_continuous_are(A, B, Q, R))

    # Finally, solve for the optimal gain matrix using the cost-to-go matrix
    return numpy.asmatrix(numpy.linalg.inv(R) * B.T * M)
//  [python]
def test_call_with_positional_input_update(self):
        '''Test the class call with a positional input update.
        '''

        input_array = byte_align(
                (numpy.random.randn(*self.input_array.shape)
                    + 1j*numpy.random.randn(*self.input_array.shape)))

        output_array = self.fft(
                byte_align(input_array.copy())).copy()

        self.update_arrays(input_array, self.output_array)
        self.fft.execute()

        self.assertTrue(numpy.alltrue(output_array == self.output_array))
//  [python]
def test_call_with_positional_input_update(self):
        '''Test the class call with a positional input update.
        '''

        input_array = byte_align(
                (numpy.random.randn(*self.input_array.shape)
                    + 1j*numpy.random.randn(*self.input_array.shape)))

        output_array = self.fft(
                byte_align(input_array.copy())).copy()

        self.update_arrays(input_array, self.output_array)
        self.fft.execute()

        self.assertTrue(numpy.all(output_array == self.output_array))
//  [python]
def test_persistent_padding(self):
        '''Test to confirm the padding it not touched after creation.
        '''
        dtype_tuple = input_dtypes[functions[self.func]]
        for dtype in dtype_tuple[0]:
            for test_shape, s, kwargs in self.test_data:

                n_add = 2
                # these slicers get the padding
                # from the internal input array
                padding_slicer = [slice(None)] * len(test_shape)
                axes = self.axes_from_kwargs(kwargs)
                try:
                    for each_axis, length in enumerate(s):
                        s[each_axis] += n_add
                        padding_slicer[axes[each_axis]] = (
                                slice(s[each_axis], None))

                except TypeError:
                    s += n_add
                    padding_slicer[axes[0]] = slice(s, None)
                padding_slicer = tuple(padding_slicer)
                # Get a valid object
                FFTW_object = self.validate_pyfftw_object(dtype_tuple[1],
                        test_shape, dtype, s, kwargs)

                internal_array = FFTW_object.input_array
                padding = internal_array[padding_slicer]

                # Fill the padding with garbage
                initial_padding = dtype_tuple[1](padding.shape, dtype)

                padding[:] = initial_padding

                # Now confirm that nothing is done to the padding
                FFTW_object()

                final_padding = FFTW_object.input_array[padding_slicer]

                self.assertTrue(n.all(final_padding == initial_padding))
//  [python]
def test_persistent_padding(self):
        '''Test to confirm the padding it not touched after creation.
        '''
        dtype_tuple = input_dtypes[functions[self.func]]
        for dtype in dtype_tuple[0]:
            for test_shape, s, kwargs in self.test_data:

                n_add = 2
                # these slicers get the padding
                # from the internal input array
                padding_slicer = [slice(None)] * len(test_shape)
                axes = self.axes_from_kwargs(kwargs)
                try:
                    for each_axis, length in enumerate(s):
                        s[each_axis] += n_add
                        padding_slicer[axes[each_axis]] = (
                                slice(s[each_axis], None))

                except TypeError:
                    s += n_add
                    padding_slicer[axes[0]] = slice(s, None)
                padding_slicer = tuple(padding_slicer)
                # Get a valid object
                FFTW_object = self.validate_pyfftw_object(dtype_tuple[1],
                        test_shape, dtype, s, kwargs)

                internal_array = FFTW_object.input_array
                padding = internal_array[padding_slicer]

                # Fill the padding with garbage
                initial_padding = dtype_tuple[1](padding.shape, dtype)

                padding[:] = initial_padding

                # Now confirm that nothing is done to the padding
                FFTW_object()

                final_padding = FFTW_object.input_array[padding_slicer]

                self.assertTrue(numpy.all(final_padding == initial_padding))
//  [python]
def test_call_with_keyword_input_update(self):
        '''Test the class call with a keyword input update.
        '''
        input_array = byte_align(
                numpy.random.randn(*self.input_array.shape)
                    + 1j*numpy.random.randn(*self.input_array.shape))

        output_array = self.fft(
            input_array=byte_align(input_array.copy())).copy()

        self.update_arrays(input_array, self.output_array)
        self.fft.execute()

        self.assertTrue(n.all(output_array == self.output_array))
//  [python]
def test_call_with_keyword_input_update(self):
        '''Test the class call with a keyword input update.
        '''
        input_array = byte_align(
                numpy.random.randn(*self.input_array.shape)
                    + 1j*numpy.random.randn(*self.input_array.shape))

        output_array = self.fft(
            input_array=byte_align(input_array.copy())).copy()

        self.update_arrays(input_array, self.output_array)
        self.fft.execute()

        self.assertTrue(numpy.all(output_array == self.output_array))
//  [python]
def test_call_with_keyword_output_update(self):
        '''Test the class call with a keyword output update.
        '''
        output_array = byte_align(
            (numpy.random.randn(*self.output_array.shape)
                + 1j*numpy.random.randn(*self.output_array.shape)))

        returned_output_array = self.fft(
                output_array=byte_align(output_array.copy())).copy()


        self.update_arrays(self.input_array, output_array)
        self.fft.execute()

        self.assertTrue(
                numpy.alltrue(returned_output_array == output_array))
//  [python]
def test_call_with_keyword_output_update(self):
        '''Test the class call with a keyword output update.
        '''
        output_array = byte_align(
            (numpy.random.randn(*self.output_array.shape)
                + 1j*numpy.random.randn(*self.output_array.shape)))

        returned_output_array = self.fft(
                output_array=byte_align(output_array.copy())).copy()


        self.update_arrays(self.input_array, output_array)
        self.fft.execute()

        self.assertTrue(
                numpy.all(returned_output_array == output_array))
//  [python]
def test_call_with_positional_updates(self):
        '''Test the class call with a positional array updates.
        '''

        input_array = byte_align((numpy.random.randn(*self.input_array.shape)
            + 1j*numpy.random.randn(*self.input_array.shape)))

        output_array = byte_align((numpy.random.randn(*self.output_array.shape)
            + 1j*numpy.random.randn(*self.output_array.shape)))

        returned_output_array = self.fft(
            byte_align(input_array.copy()),
            byte_align(output_array.copy())).copy()

        self.update_arrays(input_array, output_array)
        self.fft.execute()

        self.assertTrue(numpy.alltrue(returned_output_array == output_array))
//  [python]
def test_call_with_positional_updates(self):
        '''Test the class call with a positional array updates.
        '''

        input_array = byte_align((numpy.random.randn(*self.input_array.shape)
            + 1j*numpy.random.randn(*self.input_array.shape)))

        output_array = byte_align((numpy.random.randn(*self.output_array.shape)
            + 1j*numpy.random.randn(*self.output_array.shape)))

        returned_output_array = self.fft(
            byte_align(input_array.copy()),
            byte_align(output_array.copy())).copy()

        self.update_arrays(input_array, output_array)
        self.fft.execute()

        self.assertTrue(numpy.all(returned_output_array == output_array))
//  [python]
def test_call_with_keyword_updates(self):
        '''Test the class call with a positional output update.
        '''

        input_array = byte_align(
                (numpy.random.randn(*self.input_array.shape)
                    + 1j*numpy.random.randn(*self.input_array.shape)))

        output_array = byte_align(
                (numpy.random.randn(*self.output_array.shape)
                    + 1j*numpy.random.randn(*self.output_array.shape)))

        returned_output_array = self.fft(
                output_array=byte_align(output_array.copy()),
                input_array=byte_align(input_array.copy())).copy()

        self.update_arrays(input_array, output_array)
        self.fft.execute()

        self.assertTrue(numpy.alltrue(returned_output_array == output_array))
//  [python]
def test_call_with_keyword_updates(self):
        '''Test the class call with a positional output update.
        '''

        input_array = byte_align(
                (numpy.random.randn(*self.input_array.shape)
                    + 1j*numpy.random.randn(*self.input_array.shape)))

        output_array = byte_align(
                (numpy.random.randn(*self.output_array.shape)
                    + 1j*numpy.random.randn(*self.output_array.shape)))

        returned_output_array = self.fft(
                output_array=byte_align(output_array.copy()),
                input_array=byte_align(input_array.copy())).copy()

        self.update_arrays(input_array, output_array)
        self.fft.execute()

        self.assertTrue(numpy.all(returned_output_array == output_array))
//  [python]
def test_call_with_list_input(self):
        '''Test the class call with a list rather than an array
        '''

        output_array = self.fft().copy()

        test_output_array = self.fft(self.input_array.tolist()).copy()

        self.assertTrue(numpy.alltrue(output_array == test_output_array))
//  [python]
def test_call_with_list_input(self):
        '''Test the class call with a list rather than an array
        '''

        output_array = self.fft().copy()

        test_output_array = self.fft(self.input_array.tolist()).copy()

        self.assertTrue(numpy.all(output_array == test_output_array))
//  [python]
def test_call_with_different_input_dtype(self):
        '''Test the class call with an array with a different input dtype
        '''
        input_array = byte_align(numpy.complex64(
                numpy.random.randn(*self.input_array.shape)
                + 1j*numpy.random.randn(*self.input_array.shape)))

        output_array = self.fft(byte_align(input_array.copy())).copy()

        _input_array = numpy.asarray(input_array,
                dtype=self.input_array.dtype)

        self.update_arrays(_input_array, self.output_array)
        self.fft.execute()

        self.assertTrue(numpy.alltrue(output_array == self.output_array))
//  [python]
def test_call_with_different_input_dtype(self):
        '''Test the class call with an array with a different input dtype
        '''
        input_array = byte_align(numpy.complex64(
                numpy.random.randn(*self.input_array.shape)
                + 1j*numpy.random.randn(*self.input_array.shape)))

        output_array = self.fft(byte_align(input_array.copy())).copy()

        _input_array = numpy.asarray(input_array,
                dtype=self.input_array.dtype)

        self.update_arrays(_input_array, self.output_array)
        self.fft.execute()

        self.assertTrue(numpy.all(output_array == self.output_array))
//  [python]
def test9ToNumpy(self):
    import numpy
    bv = ds.DiscreteValueVect(ds.DiscreteValueType.FOURBITVALUE, 32)
    bv[0] = 1
    bv[1] = 4
    bv[17] = 1
    bv[23] = 8
    bv[31] = 12
    arr = numpy.zeros((3, ), 'i')
    ds.ConvertToNumpyArray(bv, arr)
    self.assertTrue(numpy.alltrue(bv == arr))
//  [python]
def test9ToNumpy(self):
    import numpy
    bv = ds.DiscreteValueVect(ds.DiscreteValueType.FOURBITVALUE, 32)
    bv[0] = 1
    bv[1] = 4
    bv[17] = 1
    bv[23] = 8
    bv[31] = 12
    arr = numpy.zeros((3, ), 'i')
    ds.ConvertToNumpyArray(bv, arr)
    self.assertTrue(numpy.all(bv == arr))
//  [python]
def item(context, node):
    inputs = _get_inputs(context, node, expected=1)

    if inputs[0].shape == ():
        # MIL ops that reduce already output a scalar, so no need to do
        # anything.
        res = inputs[0]
    elif _n.all([d == 1 for d in inputs[0].shape]):
        # Item only makes sense when called on a length 1 tensor. We use
        # reduce_max as a workaround for not having a way to extract a scalar
        # from a symbolic tensor.
        res = mb.reduce_max(x=inputs[0], name=node.name)
    else:
        raise ValueError("expected input to be a scalar or a length 1 tensor")
    context.add(res, node.name)
//  [python]
def item(context, node):
    inputs = _get_inputs(context, node, expected=1)

    if inputs[0].shape == ():
        # MIL ops that reduce already output a scalar, so no need to do
        # anything.
        res = inputs[0]
    elif _np.all([d == 1 for d in inputs[0].shape]):
        # Item only makes sense when called on a length 1 tensor. We use
        # reduce_max as a workaround for not having a way to extract a scalar
        # from a symbolic tensor.
        res = mb.reduce_max(x=inputs[0], name=node.name)
    else:
        raise ValueError("expected input to be a scalar or a length 1 tensor")
    context.add(res, node.name)
//  [python]
def overlay(self):
        """Return a color image of the segmentation as an overlay
        """
        if self.__overlay is not None:
            return self.__overlay
        sm = matplotlib.cm.ScalarMappable(cmap=self.colormap)
        sm.set_clim(
            vmin=1, vmax=numpy.max([numpy.max(label) for label in self.labels]) + 1
        )

        img = None
        lmin = 0
        for label in self.labels:
            if n.all(label == 0):
                continue
            if self.scramble:
                lmin = numpy.min(label[label != 0])
            label[label != 0] = (
                cellprofiler.gui.tools.renumber_labels_for_display(label)[label != 0]
                + lmin
            )
            lmin = numpy.max(label)
            if img is None:
                img = sm.to_rgba(label)
                img[label == 0, :] = 0
            else:
                img[label != 0, :] = sm.to_rgba(label[label != 0])
        self.__overlay = img
        return img
//  [python]
def overlay(self):
        """Return a color image of the segmentation as an overlay
        """
        if self.__overlay is not None:
            return self.__overlay
        sm = matplotlib.cm.ScalarMappable(cmap=self.colormap)
        sm.set_clim(
            vmin=1, vmax=numpy.max([numpy.max(label) for label in self.labels]) + 1
        )

        img = None
        lmin = 0
        for label in self.labels:
            if numpy.all(label == 0):
                continue
            if self.scramble:
                lmin = numpy.min(label[label != 0])
            label[label != 0] = (
                cellprofiler.gui.tools.renumber_labels_for_display(label)[label != 0]
                + lmin
            )
            lmin = numpy.max(label)
            if img is None:
                img = sm.to_rgba(label)
                img[label == 0, :] = 0
            else:
                img[label != 0, :] = sm.to_rgba(label[label != 0])
        self.__overlay = img
        return img
//  [python]
def test_call_with_different_striding(self):
        '''Test the input update with different strides to internal array.
        '''
        input_array_shape = self.input_array.shape + (2,)
        internal_array_shape = self.internal_array.shape

        internal_array = byte_align(
                numpy.random.randn(*internal_array_shape)
                + 1j*numpy.random.randn(*internal_array_shape))

        fft =  utils._FFTWWrapper(internal_array, self.output_array,
                input_array_slicer=self.input_array_slicer,
                FFTW_array_slicer=self.FFTW_array_slicer)

        test_output_array = fft().copy()

        new_input_array = empty_aligned(input_array_shape,
                                        dtype=internal_array.dtype)
        new_input_array[:] = 0

        new_input_array[:,:,0][self.input_array_slicer] = (
                internal_array[self.FFTW_array_slicer])

        new_output = fft(new_input_array[:,:,0]).copy()

        # Test the test!
        self.assertTrue(
                new_input_array[:,:,0].strides != internal_array.strides)

        self.assertTrue(numpy.alltrue(test_output_array == new_output))
//  [python]
def test_call_with_different_striding(self):
        '''Test the input update with different strides to internal array.
        '''
        input_array_shape = self.input_array.shape + (2,)
        internal_array_shape = self.internal_array.shape

        internal_array = byte_align(
                numpy.random.randn(*internal_array_shape)
                + 1j*numpy.random.randn(*internal_array_shape))

        fft =  utils._FFTWWrapper(internal_array, self.output_array,
                input_array_slicer=self.input_array_slicer,
                FFTW_array_slicer=self.FFTW_array_slicer)

        test_output_array = fft().copy()

        new_input_array = empty_aligned(input_array_shape,
                                        dtype=internal_array.dtype)
        new_input_array[:] = 0

        new_input_array[:,:,0][self.input_array_slicer] = (
                internal_array[self.FFTW_array_slicer])

        new_output = fft(new_input_array[:,:,0]).copy()

        # Test the test!
        self.assertTrue(
                new_input_array[:,:,0].strides != internal_array.strides)

        self.assertTrue(numpy.all(test_output_array == new_output))
//  [python]
def _cast(context, node, dtype, dtype_name):
    inputs = _get_inputs(context, node, expected=1)
    x = inputs[0]
    # Input must either be a scalar or a (1 x 1 x ... x 1) tensor
    if not (len(x.shape) == 0 or n.all([d == 1 for d in x.shape])):
        raise ValueError("input to cast must be either a scalar or a length 1 tensor")

    if x.can_be_folded_to_const():
        # If x is a compile-time constant, directly cast it to @dtype if it's
        # not one already.
        if not isinstance(x.val, dtype):
            res = mb.const(val=dtype(x.val), name=node.name)
        else:
            res = x
    elif len(x.shape) > 0:
        x = mb.squeeze(x=x, name=node.name + "_item")
        res = mb.cast(x=x, dtype=dtype_name, name=node.name)
    else:
        res = mb.cast(x=x, dtype=dtype_name, name=node.name)
    context.add(res, node.name)
//  [python]
def _cast(context, node, dtype, dtype_name):
    inputs = _get_inputs(context, node, expected=1)
    x = inputs[0]
    # Input must either be a scalar or a (1 x 1 x ... x 1) tensor
    if not (len(x.shape) == 0 or _np.all([d == 1 for d in x.shape])):
        raise ValueError("input to cast must be either a scalar or a length 1 tensor")

    if x.can_be_folded_to_const():
        # If x is a compile-time constant, directly cast it to @dtype if it's
        # not one already.
        if not isinstance(x.val, dtype):
            res = mb.const(val=dtype(x.val), name=node.name)
        else:
            res = x
    elif len(x.shape) > 0:
        x = mb.squeeze(x=x, name=node.name + "_item")
        res = mb.cast(x=x, dtype=dtype_name, name=node.name)
    else:
        res = mb.cast(x=x, dtype=dtype_name, name=node.name)
    context.add(res, node.name)
//  [python]
def loadCMat(self):
        """
        Loads a control matrix from file to the reconstructor

        Looks in the standard reconstructor directory for a control matrix and loads the file.
        Also looks at the FITS header and checks that the control matrix is compatible with the current simulation.
        """

        filename = self.sim_config.simName+"/cMat.fits"

        logger.info("Load Command Matrix")

        cMatHDU = fits.open(filename)[0]
        cMatHDU.verify("fix")
        header = cMatHDU.header

        try:
            dmNo = int(header["NBDM"])
            exec("dmActs = numpy.array({})".format(
                    cMatHDU.header["DMACTS"]), globals())
            exec("dmTypes = %s" % header["DMTYPE"], globals())
            exec("dmConds = numpy.array({})".format(
                    cMatHDU.header["DMCOND"]), globals())

            if not numpy.allclose(dmConds, self.dmConds):
                raise IOError("DM conditioning Parameter changed - will make new control matrix")
            if not numpy.alltrue(dmActs==self.dmActs) or dmTypes!=self.dmTypes or dmNo!=dmNo:
                logger.warning("loaded control matrix may not be compatibile with \
                                the current simulation. Will try anyway....")

            cMat = cMatHDU.data

        except KeyError:
            logger.warning("loaded control matrix header has not created by this ao sim. Will load anyway.....")
            #cMat = cMatFile[1]
            cMat = cMatHDU.data

        if cMat.shape != self.controlShape:
            logger.warning("designated control matrix does not match the expected shape")
            raise IOError
        else:
            self.control_matrix = cMat
//  [python]
def loadCMat(self):
        """
        Loads a control matrix from file to the reconstructor

        Looks in the standard reconstructor directory for a control matrix and loads the file.
        Also looks at the FITS header and checks that the control matrix is compatible with the current simulation.
        """

        filename = self.sim_config.simName+"/cMat.fits"

        logger.info("Load Command Matrix")

        cMatHDU = fits.open(filename)[0]
        cMatHDU.verify("fix")
        header = cMatHDU.header

        try:
            dmNo = int(header["NBDM"])
            exec("dmActs = numpy.array({})".format(
                    cMatHDU.header["DMACTS"]), globals())
            exec("dmTypes = %s" % header["DMTYPE"], globals())
            exec("dmConds = numpy.array({})".format(
                    cMatHDU.header["DMCOND"]), globals())

            if not numpy.allclose(dmConds, self.dmConds):
                raise IOError("DM conditioning Parameter changed - will make new control matrix")
            if not numpy.all(dmActs==self.dmActs) or dmTypes!=self.dmTypes or dmNo!=dmNo:
                logger.warning("loaded control matrix may not be compatibile with \
                                the current simulation. Will try anyway....")

            cMat = cMatHDU.data

        except KeyError:
            logger.warning("loaded control matrix header has not created by this ao sim. Will load anyway.....")
            #cMat = cMatFile[1]
            cMat = cMatHDU.data

        if cMat.shape != self.controlShape:
            logger.warning("designated control matrix does not match the expected shape")
            raise IOError
        else:
            self.control_matrix = cMat
//  [python]
def by_bbox(vertices, points, radius):
    np_rads = np.array(radius)
    np_verts = np.array(vertices)
    np_bbox = np.array(points)
    bbox_max = np.amax(np_bbox, axis=0)
    bbox_min = np.amin(np_bbox, axis=0)
    min_mask = np_verts > bbox_min - np_rads[:, np.newaxis]
    max_mask = np_verts < bbox_max + np_rads[:, np.newaxis]
    out_verts_mask = np.all(np.all([min_mask, max_mask], axis=0), axis=1)

    return out_verts_mask
//  [python]
def _set_related(self, attr, val, check_mask=True):
        """Set a field that must match the independent axes size.

        The value can be None or something with the same length as the
        independent axis. This is intended to be used from the property
        setter. There is also a check to warn the user if the value
        contains a NumPy masked array which does not match the
        dependent axis.

        """
        if val is None:
            setattr(self, f"_{attr}", None)
            return

        if not numpy.iterable(val):
            raise DataErr("notanarray")

        # Check the mask before calling _check, which could call asarray
        # and so lose the mask setting.
        #
        if check_mask and hasattr(val, "mask"):
            if not hasattr(self.y, "mask") or \
               len(self.y) != len(val) or \
                   not numpy.alltrue(self.y.mask == val.mask):

                warnings.warn(f"The mask of {attr} differs from the dependent array, only the mask of the dependent array is used in Sherpa.")

        val = _check(val)
        nval = len(val)

        nelem = self.size
        if nelem is None:
            # We set the object size here
            self._size  = nval
            setattr(self, f"_{attr}", val)
            return

        nval = len(val)
        if nval != nelem:
            raise DataErr('mismatchn', 'independent axis', attr, nelem, nval)

        setattr(self, f"_{attr}", val)
//  [python]
def _set_related(self, attr, val, check_mask=True):
        """Set a field that must match the independent axes size.

        The value can be None or something with the same length as the
        independent axis. This is intended to be used from the property
        setter. There is also a check to warn the user if the value
        contains a NumPy masked array which does not match the
        dependent axis.

        """
        if val is None:
            setattr(self, f"_{attr}", None)
            return

        if not numpy.iterable(val):
            raise DataErr("notanarray")

        # Check the mask before calling _check, which could call asarray
        # and so lose the mask setting.
        #
        if check_mask and hasattr(val, "mask"):
            if not hasattr(self.y, "mask") or \
               len(self.y) != len(val) or \
                   not numpy.all(self.y.mask == val.mask):

                warnings.warn(f"The mask of {attr} differs from the dependent array, only the mask of the dependent array is used in Sherpa.")

        val = _check(val)
        nval = len(val)

        nelem = self.size
        if nelem is None:
            # We set the object size here
            self._size  = nval
            setattr(self, f"_{attr}", val)
            return

        nval = len(val)
        if nval != nelem:
            raise DataErr('mismatchn', 'independent axis', attr, nelem, nval)

        setattr(self, f"_{attr}", val)
//  [python]
def test_bare_trainonlytransformer(random_xy_dataset_regr):
    """Tests whether the trainonlytransformer will only transform train when used directly"""

    X_train, X_test, y_train, y_test = train_test_split(*random_xy_dataset_regr)
    trf = TrainOnlyTrainOnlyTransformer()
    trf.fit(X_train, y_train)

    assert n.all(trf.transform(X_train) != X_train)
    assert n.all(trf.transform(X_test) == X_test)
//  [python]
def test_bare_trainonlytransformer(random_xy_dataset_regr):
    """Tests whether the trainonlytransformer will only transform train when used directly"""

    X_train, X_test, y_train, y_test = train_test_split(*random_xy_dataset_regr)
    trf = TrainOnlyTrainOnlyTransformer()
    trf.fit(X_train, y_train)

    assert np.all(trf.transform(X_train) != X_train)
    assert np.all(trf.transform(X_test) == X_test)
//  [python]
def test_pipeline_trainonlytransformer(random_xy_dataset_regr):
    """Tests whether the trainonlytransformer will only transform train when used in a pipeline"""

    X_train, X_test, y_train, y_test = train_test_split(*random_xy_dataset_regr)

    trf = make_pipeline(TrainOnlyTrainOnlyTransformer())
    trf.fit(X_train, y_train)

    assert n.all(trf.transform(X_train) != X_train)
    assert n.all(trf.transform(X_test) == X_test)
//  [python]
def test_pipeline_trainonlytransformer(random_xy_dataset_regr):
    """Tests whether the trainonlytransformer will only transform train when used in a pipeline"""

    X_train, X_test, y_train, y_test = train_test_split(*random_xy_dataset_regr)

    trf = make_pipeline(TrainOnlyTrainOnlyTransformer())
    trf.fit(X_train, y_train)

    assert np.all(trf.transform(X_train) != X_train)
    assert np.all(trf.transform(X_test) == X_test)
//  [python]
def zpk2tf(z, p, k):
    """
    Return polynomial transfer function representation from zeros and poles

    Parameters
    ----------
    z : array_like
        Zeros of the transfer function.
    p : array_like
        Poles of the transfer function.
    k : float
        System gain.

    Returns
    -------
    b : ndarray
        Numerator polynomial coefficients.
    a : ndarray
        Denominator polynomial coefficients.

    """
    z = atleast_1d(z)
    k = atleast_1d(k)
    if len(z.shape) > 1:
        temp = poly(z[0])
        b = zeros((z.shape[0], z.shape[1] + 1), temp.dtype.char)
        if len(k) == 1:
            k = [k[0]] * z.shape[0]
        for i in range(z.shape[0]):
            b[i] = k[i] * poly(z[i])
    else:
        b = k * poly(z)
    a = atleast_1d(poly(p))

    # Use real output if possible. Copied from numpy.poly, since
    # we can't depend on a specific version of numpy.
    if issubclass(b.dtype.type, numpy.complexfloating):
        # if complex roots are all complex conjugates, the roots are real.
        roots = numpy.asarray(z, complex)
        pos_roots = numpy.compress(roots.imag > 0, roots)
        neg_roots = numpy.conjugate(numpy.compress(roots.imag < 0, roots))
        if len(pos_roots) == len(neg_roots):
            if numpy.alltrue(numpy.sort_complex(neg_roots) ==
                         numpy.sort_complex(pos_roots)):
                b = b.real.copy()

    if issubclass(a.dtype.type, numpy.complexfloating):
        # if complex roots are all complex conjugates, the roots are real.
        roots = numpy.asarray(p, complex)
        pos_roots = numpy.compress(roots.imag > 0, roots)
        neg_roots = numpy.conjugate(numpy.compress(roots.imag < 0, roots))
        if len(pos_roots) == len(neg_roots):
            if numpy.alltrue(numpy.sort_complex(neg_roots) ==
                         numpy.sort_complex(pos_roots)):
                a = a.real.copy()

    return b, a
//  [python]
def zpk2tf(z, p, k):
    """
    Return polynomial transfer function representation from zeros and poles

    Parameters
    ----------
    z : array_like
        Zeros of the transfer function.
    p : array_like
        Poles of the transfer function.
    k : float
        System gain.

    Returns
    -------
    b : ndarray
        Numerator polynomial coefficients.
    a : ndarray
        Denominator polynomial coefficients.

    """
    z = atleast_1d(z)
    k = atleast_1d(k)
    if len(z.shape) > 1:
        temp = poly(z[0])
        b = zeros((z.shape[0], z.shape[1] + 1), temp.dtype.char)
        if len(k) == 1:
            k = [k[0]] * z.shape[0]
        for i in range(z.shape[0]):
            b[i] = k[i] * poly(z[i])
    else:
        b = k * poly(z)
    a = atleast_1d(poly(p))

    # Use real output if possible. Copied from numpy.poly, since
    # we can't depend on a specific version of numpy.
    if issubclass(b.dtype.type, numpy.complexfloating):
        # if complex roots are all complex conjugates, the roots are real.
        roots = numpy.asarray(z, complex)
        pos_roots = numpy.compress(roots.imag > 0, roots)
        neg_roots = numpy.conjugate(numpy.compress(roots.imag < 0, roots))
        if len(pos_roots) == len(neg_roots):
            if numpy.all(numpy.sort_complex(neg_roots) ==
                         numpy.sort_complex(pos_roots)):
                b = b.real.copy()

    if issubclass(a.dtype.type, numpy.complexfloating):
        # if complex roots are all complex conjugates, the roots are real.
        roots = numpy.asarray(p, complex)
        pos_roots = numpy.compress(roots.imag > 0, roots)
        neg_roots = numpy.conjugate(numpy.compress(roots.imag < 0, roots))
        if len(pos_roots) == len(neg_roots):
            if numpy.all(numpy.sort_complex(neg_roots) ==
                         numpy.sort_complex(pos_roots)):
                a = a.real.copy()

    return b, a
//  [python]
def test_bare_trainonlytransformer_pandas(random_xy_dataset_regr):
    """Tests whether the trainonlytransformer will only transform train when used directly"""
    X, y = (
        pd.DataFrame(random_xy_dataset_regr[0]),
        pd.DataFrame(random_xy_dataset_regr[1]),
    )
    X_train, X_test, y_train, y_test = train_test_split(X, y)

    trf = TrainOnlyTrainOnlyTransformer()
    trf.fit(X_train, y_train)

    assert n.all(trf.transform(X_train) != X_train)
    assert n.all(trf.transform(X_test) == X_test)

    assert isinstance(trf.transform(X_train), pd.DataFrame)
    assert isinstance(trf.transform(X_test), pd.DataFrame)