jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def test_pdf_logpdf(self):
        """ Checks that the probability density function (PDF) is within some
        tolerance of the natural logarithm of the PDF. This implementation is
        for one dimensional distributions.
        """

        # assign tolerance for element-wise ratio of logarithm of PDF
        tolerance = 0.01

        # step size to take in PDF evaluation
        step = 0.1

        # loop over distributions
        for dist in self.dists:
            if dist.name in EXCLUDE_DIST_NAMES:
                continue
            for param in dist.params:

                # get min and max
                hist_min = dist.bounds[param][0]
                hist_max = dist.bounds[param][1]

                # get the PDF and logarithm of the PDF from the distriubtion
                x = numpy.arange(hist_min, hist_max, step)
                pdf = numpy.array([dist.pdf(**{param : xx}) for xx in x])
                logpdf = numpy.array([dist.logpdf(**{param : xx}) for xx in x])

                # find the logarithm of the PDF
                pdf_log = numpy.log(pdf)

                # see if each element in ratio of these two logarithm of PDF
                # values are within the specified tolerance
                if not numpy.all(abs(1.0 - logpdf / pdf_log) < tolerance):
                    raise ValueError("The PDF and logarithm of the PDF "
                                     "functions for distribution {} "
                                     "do not agree".format(dist.name))
//  [python]
def is_null(obj):
    """Check if an obj is null (nan, None or array of nan)"""
    array_nan = False
    all_none = False
    try:
        tmp = to_numpy(obj)
        array_nan = np.all(np.isnan(tmp))
    except Exception:
        pass
    try:
        all_none = all(x is None for x in obj)
    except Exception:
        pass
    return obj is None or all_none or array_nan
//  [python]
def test_do_not_shrink_identical():
    """Test a case where the primary and secondary objects are identical"""
    primary_labels = numpy.zeros((20, 20), int)
    secondary_labels = numpy.zeros((20, 20), int)
    expected_labels = numpy.zeros((20, 20), int)

    # first and third objects have different sizes
    primary_labels[3:6, 4:7] = 1
    secondary_labels[2:7, 3:8] = 1
    expected_labels[2:7, 3:8] = 1
    expected_labels[3:6, 4:7] = 0

    primary_labels[13:16, 4:7] = 3
    secondary_labels[12:17, 3:8] = 3
    expected_labels[12:17, 3:8] = 3
    expected_labels[13:16, 4:7] = 0

    # second object and fourth have same size

    primary_labels[3:6, 14:17] = 2
    secondary_labels[3:6, 14:17] = 2
    primary_labels[13:16, 14:17] = 4
    secondary_labels[13:16, 14:17] = 4
    workspace = make_workspace(primary_labels, secondary_labels)

    module = workspace.module
    module.shrink_primary.value = False
    module.run(workspace)
    output_objects = workspace.object_set.get_objects(TERTIARY)
    assert n.all(output_objects.segmented == expected_labels)

    measurements = workspace.measurements
    count_feature = "Count_%s" % TERTIARY
    value = measurements.get_current_measurement("Image", count_feature)
    assert value == 3

    child_count_feature = "Children_%s_Count" % TERTIARY
    for parent_name in PRIMARY, SECONDARY:
        parent_of_feature = "Parent_%s" % parent_name
        parent_of = measurements.get_current_measurement(TERTIARY, parent_of_feature)
        child_count = measurements.get_current_measurement(
            parent_name, child_count_feature
        )
        for parent, expected_child_count in ((1, 1), (2, 0), (3, 1), (4, 0)):
            assert child_count[parent - 1] == expected_child_count
        for child in (1, 3):
            assert parent_of[child - 1] == child

    for location_feature in (
        M_LOCATION_CENTER_X,
        M_LOCATION_CENTER_Y,
    ):
        values = measurements.get_current_measurement(TERTIARY, location_feature)
        assert n.all(numpy.isnan(values) == [False, True, False])
//  [python]
def test_do_not_shrink_identical():
    """Test a case where the primary and secondary objects are identical"""
    primary_labels = numpy.zeros((20, 20), int)
    secondary_labels = numpy.zeros((20, 20), int)
    expected_labels = numpy.zeros((20, 20), int)

    # first and third objects have different sizes
    primary_labels[3:6, 4:7] = 1
    secondary_labels[2:7, 3:8] = 1
    expected_labels[2:7, 3:8] = 1
    expected_labels[3:6, 4:7] = 0

    primary_labels[13:16, 4:7] = 3
    secondary_labels[12:17, 3:8] = 3
    expected_labels[12:17, 3:8] = 3
    expected_labels[13:16, 4:7] = 0

    # second object and fourth have same size

    primary_labels[3:6, 14:17] = 2
    secondary_labels[3:6, 14:17] = 2
    primary_labels[13:16, 14:17] = 4
    secondary_labels[13:16, 14:17] = 4
    workspace = make_workspace(primary_labels, secondary_labels)

    module = workspace.module
    module.shrink_primary.value = False
    module.run(workspace)
    output_objects = workspace.object_set.get_objects(TERTIARY)
    assert numpy.all(output_objects.segmented == expected_labels)

    measurements = workspace.measurements
    count_feature = "Count_%s" % TERTIARY
    value = measurements.get_current_measurement("Image", count_feature)
    assert value == 3

    child_count_feature = "Children_%s_Count" % TERTIARY
    for parent_name in PRIMARY, SECONDARY:
        parent_of_feature = "Parent_%s" % parent_name
        parent_of = measurements.get_current_measurement(TERTIARY, parent_of_feature)
        child_count = measurements.get_current_measurement(
            parent_name, child_count_feature
        )
        for parent, expected_child_count in ((1, 1), (2, 0), (3, 1), (4, 0)):
            assert child_count[parent - 1] == expected_child_count
        for child in (1, 3):
            assert parent_of[child - 1] == child

    for location_feature in (
        M_LOCATION_CENTER_X,
        M_LOCATION_CENTER_Y,
    ):
        values = measurements.get_current_measurement(TERTIARY, location_feature)
        assert numpy.all(numpy.isnan(values) == [False, True, False])
//  [python]
def __eq__(self, other):
        return (
            isinstance(other, Trace)
            and self.network == other.network
            and self.station == other.station
            and self.location == other.location
            and self.channel == other.channel
            and (abs(self.deltat - other.deltat)
                 < (self.deltat + other.deltat)*1e-6)
            and abs(self.tmin-other.tmin) < self.deltat*0.01
            and abs(self.tmax-other.tmax) < self.deltat*0.01
            and n.all(self.ydata == other.ydata))
//  [python]
def __eq__(self, other):
        return (
            isinstance(other, Trace)
            and self.network == other.network
            and self.station == other.station
            and self.location == other.location
            and self.channel == other.channel
            and (abs(self.deltat - other.deltat)
                 < (self.deltat + other.deltat)*1e-6)
            and abs(self.tmin-other.tmin) < self.deltat*0.01
            and abs(self.tmax-other.tmax) < self.deltat*0.01
            and num.all(self.ydata == other.ydata))
//  [python]
def test_do_not_shrink_missing():
    # Regression test of 705

    for missing in range(1, 3):
        for missing_primary in False, True:
            primary_labels = numpy.zeros((20, 20), int)
            secondary_labels = numpy.zeros((20, 20), int)
            expected_labels = numpy.zeros((20, 20), int)
            centers = ((5, 5), (15, 5), (5, 15))
            pidx = 1
            sidx = 1
            for idx, (i, j) in enumerate(centers):
                if (idx + 1 != missing) or not missing_primary:
                    primary_labels[(i - 1) : (i + 2), (j - 1) : (j + 2)] = pidx
                    pidx += 1
                if (idx + 1 != missing) or missing_primary:
                    secondary_labels[(i - 2) : (i + 3), (j - 2) : (j + 3)] = sidx
                    sidx += 1
            expected_labels = secondary_labels * (primary_labels == 0)
            workspace = make_workspace(primary_labels, secondary_labels)

            module = workspace.module
            module.shrink_primary.value = False
            module.run(workspace)
            output_objects = workspace.object_set.get_objects(TERTIARY)
            assert n.all(output_objects.segmented == expected_labels)

            m = workspace.measurements

            child_name = module.subregion_objects_name.value
            primary_name = module.primary_objects_name.value
            ftr = FF_PARENT % primary_name
            pparents = m[child_name, ftr]
            assert len(pparents) == (3 if missing_primary else 2)
            if missing_primary:
                assert pparents[missing - 1] == 0

            secondary_name = module.secondary_objects_name.value
            ftr = FF_PARENT % secondary_name
            pparents = m[child_name, ftr]
            assert len(pparents) == (3 if missing_primary else 2)
            if not missing_primary:
                assert all([x in pparents for x in range(1, 3)])

            ftr = FF_CHILDREN_COUNT % child_name
            children = m[primary_name, ftr]
            assert len(children) == (2 if missing_primary else 3)
            if not missing_primary:
                assert children[missing - 1] == 0
                assert n.all(n.delete(children, missing - 1) == 1)
            else:
                assert n.all(children == 1)

            children = m[secondary_name, ftr]
            assert len(children) == (3 if missing_primary else 2)
            assert n.all(children == 1)
//  [python]
def test_do_not_shrink_missing():
    # Regression test of 705

    for missing in range(1, 3):
        for missing_primary in False, True:
            primary_labels = numpy.zeros((20, 20), int)
            secondary_labels = numpy.zeros((20, 20), int)
            expected_labels = numpy.zeros((20, 20), int)
            centers = ((5, 5), (15, 5), (5, 15))
            pidx = 1
            sidx = 1
            for idx, (i, j) in enumerate(centers):
                if (idx + 1 != missing) or not missing_primary:
                    primary_labels[(i - 1) : (i + 2), (j - 1) : (j + 2)] = pidx
                    pidx += 1
                if (idx + 1 != missing) or missing_primary:
                    secondary_labels[(i - 2) : (i + 3), (j - 2) : (j + 3)] = sidx
                    sidx += 1
            expected_labels = secondary_labels * (primary_labels == 0)
            workspace = make_workspace(primary_labels, secondary_labels)

            module = workspace.module
            module.shrink_primary.value = False
            module.run(workspace)
            output_objects = workspace.object_set.get_objects(TERTIARY)
            assert numpy.all(output_objects.segmented == expected_labels)

            m = workspace.measurements

            child_name = module.subregion_objects_name.value
            primary_name = module.primary_objects_name.value
            ftr = FF_PARENT % primary_name
            pparents = m[child_name, ftr]
            assert len(pparents) == (3 if missing_primary else 2)
            if missing_primary:
                assert pparents[missing - 1] == 0

            secondary_name = module.secondary_objects_name.value
            ftr = FF_PARENT % secondary_name
            pparents = m[child_name, ftr]
            assert len(pparents) == (3 if missing_primary else 2)
            if not missing_primary:
                assert all([x in pparents for x in range(1, 3)])

            ftr = FF_CHILDREN_COUNT % child_name
            children = m[primary_name, ftr]
            assert len(children) == (2 if missing_primary else 3)
            if not missing_primary:
                assert children[missing - 1] == 0
                assert numpy.all(numpy.delete(children, missing - 1) == 1)
            else:
                assert numpy.all(children == 1)

            children = m[secondary_name, ftr]
            assert len(children) == (3 if missing_primary else 2)
            assert numpy.all(children == 1)
//  [python]
def numpy_has_correlate_flip_bug():
    '''
    Check if NumPy's correlate function reveals old behaviour.
    '''

    if _globals._numpy_has_correlate_flip_bug is None:
        a = num.array([0, 0, 1, 0, 0, 0, 0])
        b = num.array([0, 0, 0, 0, 1, 0, 0, 0])
        ab = num.correlate(a, b, mode='same')
        ba = num.correlate(b, a, mode='same')
        _globals._numpy_has_correlate_flip_bug = n.all(ab == ba)

    return _globals._numpy_has_correlate_flip_bug
//  [python]
def numpy_has_correlate_flip_bug():
    '''
    Check if NumPy's correlate function reveals old behaviour.
    '''

    if _globals._numpy_has_correlate_flip_bug is None:
        a = num.array([0, 0, 1, 0, 0, 0, 0])
        b = num.array([0, 0, 0, 0, 1, 0, 0, 0])
        ab = num.correlate(a, b, mode='same')
        ba = num.correlate(b, a, mode='same')
        _globals._numpy_has_correlate_flip_bug = num.all(ab == ba)

    return _globals._numpy_has_correlate_flip_bug
//  [python]
def find_nearest_sample_points_by_angle(thetas, phi, return_NaN=False):
    """Returns two successive indices of the angles in the array argument
    thetas, for which theta values straddle angle phi.

    This function is useful if phi represents the angle from the vertical of
    the normal to a curve's tangent line, and thetas are angles from the
    vertical of another curve in the plane. This function can therefore help
    find closest or furthest perpendicular distances between curves.

    Assumes thetas are in increasing order, and that phi is properly contained
      between min(thetas) and max(thetas) (otherwise raises a ValueError).

    For the corner case where one difference between angles is exactly 0 (not
      generic), one of the thetas values given by the returned indices is
      exactly phi."""
    global plotter
    pos_diffs = (thetas - phi) > 0
    if n.all(pos_diffs) or not n.any(pos_diffs):
        if return_NaN:
            return np.NaN
        else:
            raise ValueError("Outside domain")
    else:
        if pos_diffs[0]:
            # Trues then Falses in pos_diffs
            # first_neg_ix guaranteed not to be 0
            first_neg_ix = np.argmin(np.asarray(pos_diffs, int))
            straddling_ixs = (first_neg_ix-1, first_neg_ix)
        else:
            # Falses then Trues in pos_diffs
            first_pos_ix = np.argmax(np.asarray(pos_diffs, int))
            straddling_ixs = (first_pos_ix-1, first_pos_ix)
        return straddling_ixs
//  [python]
def find_nearest_sample_points_by_angle(thetas, phi, return_NaN=False):
    """Returns two successive indices of the angles in the array argument
    thetas, for which theta values straddle angle phi.

    This function is useful if phi represents the angle from the vertical of
    the normal to a curve's tangent line, and thetas are angles from the
    vertical of another curve in the plane. This function can therefore help
    find closest or furthest perpendicular distances between curves.

    Assumes thetas are in increasing order, and that phi is properly contained
      between min(thetas) and max(thetas) (otherwise raises a ValueError).

    For the corner case where one difference between angles is exactly 0 (not
      generic), one of the thetas values given by the returned indices is
      exactly phi."""
    global plotter
    pos_diffs = (thetas - phi) > 0
    if np.all(pos_diffs) or not np.any(pos_diffs):
        if return_NaN:
            return np.NaN
        else:
            raise ValueError("Outside domain")
    else:
        if pos_diffs[0]:
            # Trues then Falses in pos_diffs
            # first_neg_ix guaranteed not to be 0
            first_neg_ix = np.argmin(np.asarray(pos_diffs, int))
            straddling_ixs = (first_neg_ix-1, first_neg_ix)
        else:
            # Falses then Trues in pos_diffs
            first_pos_ix = np.argmax(np.asarray(pos_diffs, int))
            straddling_ixs = (first_pos_ix-1, first_pos_ix)
        return straddling_ixs
//  [python]
def MakeSmh(S):
    ew,ev = np.linalg.eigh(S)
    assert (n.all(ew > 1e-10))
    v = ev * (ew**-0.25)[np.newaxis,:]
    return np.dot(v, v.T)
//  [python]
def MakeSmh(S):
    ew,ev = np.linalg.eigh(S)
    assert (np.all(ew > 1e-10))
    v = ev * (ew**-0.25)[np.newaxis,:]
    return np.dot(v, v.T)
//  [python]
def zero_pad(self, a_xg):
        """Pad array with zeros as first element along non-periodic directions.

        Should only be invoked on global arrays.
        """
        assert np.all(a_xg.shape[-3:] == (self.N_c + self.pbc_c - 1))
        if n.all(self.pbc_c):
            return a_xg

        npbx, npby, npbz = 1 - self.pbc_c
        b_xg = np.zeros(a_xg.shape[:-3] + tuple(self.N_c), dtype=a_xg.dtype)
        b_xg[..., npbx:, npby:, npbz:] = a_xg
        return b_xg
//  [python]
def zero_pad(self, a_xg):
        """Pad array with zeros as first element along non-periodic directions.

        Should only be invoked on global arrays.
        """
        assert np.all(a_xg.shape[-3:] == (self.N_c + self.pbc_c - 1))
        if self.pbc_c.all():
            return a_xg

        npbx, npby, npbz = 1 - self.pbc_c
        b_xg = np.zeros(a_xg.shape[:-3] + tuple(self.N_c), dtype=a_xg.dtype)
        b_xg[..., npbx:, npby:, npbz:] = a_xg
        return b_xg
//  [python]
def crop(self, xdom, ydom, resample_frac=0.1):
        """Returns a new Nullcline object cropped to the (x, y) domain
        given by pairs xdom and ydom.

        If resample_frac > 0 (default 0.1), points at the domain
        ends and additional sample points will be included in the new nullcline
        object regardless of whether they were in the original sample point set
        for the cropped nullcline. New points will be selected at the rate
        indicated, measured as a fraction of the domain width. Points that are
        as close as 1/10th of this rate to existing sample points will be
        excluded, to avoid ill-conditioning the spline shape in the resampling.

        Given the assumption that nullclines are montonic functions of x,
        this guarantees a unique, contiguous piece of the nullcline is returned.
        """
        xinterval = Interval('xdom', float, xdom)
        yinterval = Interval('ydom', float, ydom)
        sample_vals = list(crop_2D(self.array, xinterval, yinterval))
        if len(sample_vals) == 0:
            # force resampling if not already used or if too sparse
            if resample_frac == 0 or resample_frac > 0.2:
                resample_frac = 0.1
            #raise ValueError("No nullcline sample points in given region")
        if resample_frac > 0:
            width = xdom[1] - xdom[0]
            new_xs = xinterval.sample(dt=width*resample_frac)
            tol = width*resample_frac*0.1
            for x in new_xs:
                # exclude any that are within 5% of domain extent of
                # existing sample points and within original point extents
                if n.all(abs(x - self.array[:,0])>tol) and \
                   x >= self.array[0,0] and x <= self.array[-1,0]:
                    y = self(x)
                    if yinterval.contains(y) is not notcontained:
                        sample_vals.append((x,y))
            sample_vals = np.array(sample_vals)
            ixs = argsort(sample_vals[:,0])  # sort by x
            sample_vals = sample_vals[ixs]
        try:
            return nullcline(self.xname, self.yname, np.array(sample_vals),
                             x_relative_scale=self.x_relative_scale_fac)
        except:
            print("Error cropping nullcline at sample points", sample_vals)
            print("MAYBE TOO FEW VALUES SAMPLED: number was", len(sample_vals))
            raise
//  [python]
def crop(self, xdom, ydom, resample_frac=0.1):
        """Returns a new Nullcline object cropped to the (x, y) domain
        given by pairs xdom and ydom.

        If resample_frac > 0 (default 0.1), points at the domain
        ends and additional sample points will be included in the new nullcline
        object regardless of whether they were in the original sample point set
        for the cropped nullcline. New points will be selected at the rate
        indicated, measured as a fraction of the domain width. Points that are
        as close as 1/10th of this rate to existing sample points will be
        excluded, to avoid ill-conditioning the spline shape in the resampling.

        Given the assumption that nullclines are montonic functions of x,
        this guarantees a unique, contiguous piece of the nullcline is returned.
        """
        xinterval = Interval('xdom', float, xdom)
        yinterval = Interval('ydom', float, ydom)
        sample_vals = list(crop_2D(self.array, xinterval, yinterval))
        if len(sample_vals) == 0:
            # force resampling if not already used or if too sparse
            if resample_frac == 0 or resample_frac > 0.2:
                resample_frac = 0.1
            #raise ValueError("No nullcline sample points in given region")
        if resample_frac > 0:
            width = xdom[1] - xdom[0]
            new_xs = xinterval.sample(dt=width*resample_frac)
            tol = width*resample_frac*0.1
            for x in new_xs:
                # exclude any that are within 5% of domain extent of
                # existing sample points and within original point extents
                if np.all(abs(x - self.array[:,0])>tol) and \
                   x >= self.array[0,0] and x <= self.array[-1,0]:
                    y = self(x)
                    if yinterval.contains(y) is not notcontained:
                        sample_vals.append((x,y))
            sample_vals = np.array(sample_vals)
            ixs = argsort(sample_vals[:,0])  # sort by x
            sample_vals = sample_vals[ixs]
        try:
            return nullcline(self.xname, self.yname, np.array(sample_vals),
                             x_relative_scale=self.x_relative_scale_fac)
        except:
            print("Error cropping nullcline at sample points", sample_vals)
            print("MAYBE TOO FEW VALUES SAMPLED: number was", len(sample_vals))
            raise
//  [python]
def popbuffers_finalize(
    out, array, validbits, awkwardarrow_type, generate_bitmasks, fix_offsets=True
):
    # Every buffer from Arrow must be offsets-corrected.
    if fix_offsets and (array.offset != 0 or len(array) != len(out)):
        out = out[array.offset : array.offset + len(array)]

    # Everything must leave popbuffers as option-type; the mask_node will be
    # removed by the next level up in popbuffers recursion if appropriate.

    if isinstance(awkwardarrow_type, AwkwardArrowType):
        if awkwardarrow_type.mask_type == "UnmaskedArray":
            assert validbits is None or n.all(
                n.frombuffer(validbits, dtype=n.uint8)[: len(out) // 8] == 0xFF
            )
            return revertable(
                ak.contents.UnmaskedArray.simplified(
                    out, parameters=awkwardarrow_type.mask_parameters
                ),
                out,
            )

        else:
            if validbits is None:
                if generate_bitmasks:
                    return revertable(
                        ak.contents.BitMaskedArray.simplified(
                            # ceildiv(len(out), 8) = -(len(out) // -8)
                            ak.index.IndexU8(
                                n.full(
                                    -(len(out) // -8), n.uint8(0xFF), dtype=n.uint8
                                )
                            ),
                            out,
                            valid_when=True,
                            length=len(out),
                            lsb_order=True,
                            parameters=awkwardarrow_type.mask_parameters,
                        ),
                        out,
                    )
                else:
                    return revertable(
                        ak.contents.UnmaskedArray.simplified(
                            out, parameters=awkwardarrow_type.mask_parameters
                        ),
                        out,
                    )
            else:
                return revertable(
                    ak.contents.BitMaskedArray.simplified(
                        ak.index.IndexU8(n.frombuffer(validbits, dtype=n.uint8)),
                        out,
                        valid_when=True,
                        length=len(out),
                        lsb_order=True,
                        parameters=awkwardarrow_type.mask_parameters,
                    ),
                    out,
                )

    else:
        if validbits is None and generate_bitmasks:
            # ceildiv(len(out), 8) = -(len(out) // -8)
            validbits = n.full(-(len(out) // -8), n.uint8(0xFF), dtype=n.uint8)

        if validbits is None:
            return revertable(ak.contents.UnmaskedArray.simplified(out), out)
        else:
            return revertable(
                ak.contents.BitMaskedArray.simplified(
                    ak.index.IndexU8(n.frombuffer(validbits, dtype=n.uint8)),
                    out,
                    valid_when=True,
                    length=len(out),
                    lsb_order=True,
                ),
                out,
            )
//  [python]
def test_array_broadcasting(self):
        """
        Regression test for https://github.com/spacetelescope/PyFITS/pull/48
        """

        with fits.open(self.data("table.fits")) as hdu:
            data = hdu[1].data
            data["V_mag"] = 0
            assert data["V_mag"].all() == 0

            data["V_mag"] = 1
            assert data["V_mag"].all() == 1

            for container in (list, tuple, np.array):
                data["V_mag"] = container([1, 2, 3])
                assert np.array_equal(data["V_mag"], np.array([1, 2, 3]))
//  [python]
def test_array_broadcasting(self):
        """
        Regression test for https://github.com/spacetelescope/PyFITS/pull/48
        """

        with fits.open(self.data("table.fits")) as hdu:
            data = hdu[1].data
            data["V_mag"] = 0
            assert np.all(data["V_mag"] == 0)

            data["V_mag"] = 1
            assert np.all(data["V_mag"] == 1)

            for container in (list, tuple, np.array):
                data["V_mag"] = container([1, 2, 3])
                assert np.array_equal(data["V_mag"], np.array([1, 2, 3]))
//  [python]
def test_dim_column_byte_order_mismatch(self):
        """
        When creating a table column with non-trivial TDIMn, and
        big-endian array data read from an existing FITS file, the data
        should not be unnecessarily byteswapped.

        Regression test for https://github.com/astropy/astropy/issues/3561
        """

        data = fits.getdata(self.data("random_groups.fits"))["DATA"]
        col = fits.Column(name="TEST", array=data, dim="(3,1,128,1,1)", format="1152E")
        thdu = fits.BinTableHDU.from_columns([col])
        thdu.writeto(self.temp("test.fits"))

        with fits.open(self.temp("test.fits")) as hdul:
            assert np.all(hdul[1].data["TEST"] == data)
//  [python]
def test_bin_table_init_from_string_array_column(self):
        """
        Tests two ways of creating a new `BinTableHDU` from a column of
        string arrays.

        This tests for a couple different regressions, and ensures that
        both BinTableHDU(data=arr) and BinTableHDU.from_columns(arr) work
        equivalently.

        Some of this is redundant with the following test, but checks some
        subtly different cases.
        """

        data = [[b"abcd", b"efgh"], [b"ijkl", b"mnop"], [b"qrst", b"uvwx"]]

        arr = np.array(
            [(data,), (data,), (data,), (data,), (data,)], dtype=[("S", "(3, 2)S4")]
        )

        tbhdu1 = fits.BinTableHDU(data=arr)

        def test_dims_and_roundtrip(tbhdu):
            assert tbhdu.data["S"].shape == (5, 3, 2)
            assert tbhdu.data["S"].dtype.str.endswith("U4")

            tbhdu.writeto(self.temp("test.fits"), overwrite=True)

            with fits.open(self.temp("test.fits")) as hdul:
                tbhdu2 = hdul[1]
                assert tbhdu2.header["TDIM1"] == "(4,2,3)"
                assert tbhdu2.data["S"].shape == (5, 3, 2)
                assert tbhdu.data["S"].dtype.str.endswith("U4")
                assert np.alltrue(tbhdu2.data["S"] == tbhdu.data["S"])

        test_dims_and_roundtrip(tbhdu1)

        tbhdu2 = fits.BinTableHDU.from_columns(arr)
        test_dims_and_roundtrip(tbhdu2)
//  [python]
def test_bin_table_init_from_string_array_column(self):
        """
        Tests two ways of creating a new `BinTableHDU` from a column of
        string arrays.

        This tests for a couple different regressions, and ensures that
        both BinTableHDU(data=arr) and BinTableHDU.from_columns(arr) work
        equivalently.

        Some of this is redundant with the following test, but checks some
        subtly different cases.
        """

        data = [[b"abcd", b"efgh"], [b"ijkl", b"mnop"], [b"qrst", b"uvwx"]]

        arr = np.array(
            [(data,), (data,), (data,), (data,), (data,)], dtype=[("S", "(3, 2)S4")]
        )

        tbhdu1 = fits.BinTableHDU(data=arr)

        def test_dims_and_roundtrip(tbhdu):
            assert tbhdu.data["S"].shape == (5, 3, 2)
            assert tbhdu.data["S"].dtype.str.endswith("U4")

            tbhdu.writeto(self.temp("test.fits"), overwrite=True)

            with fits.open(self.temp("test.fits")) as hdul:
                tbhdu2 = hdul[1]
                assert tbhdu2.header["TDIM1"] == "(4,2,3)"
                assert tbhdu2.data["S"].shape == (5, 3, 2)
                assert tbhdu.data["S"].dtype.str.endswith("U4")
                assert np.all(tbhdu2.data["S"] == tbhdu.data["S"])

        test_dims_and_roundtrip(tbhdu1)

        tbhdu2 = fits.BinTableHDU.from_columns(arr)
        test_dims_and_roundtrip(tbhdu2)
//  [python]
def test_bintable_to_asciitable(self):
        """Tests initializing a TableHDU with the data from a BinTableHDU."""

        with fits.open(self.data("tb.fits")) as hdul:
            tbdata = hdul[1].data
            tbhdu = fits.TableHDU(data=tbdata)
            tbhdu.writeto(self.temp("test.fits"), overwrite=True)
            with fits.open(self.temp("test.fits")) as hdul2:
                tbdata2 = hdul2[1].data
                assert n.all(tbdata["c1"] == tbdata2["c1"])
                assert n.all(tbdata["c2"] == tbdata2["c2"])
                # c3 gets converted from float32 to float64 when writing
                # test.fits, so cast to float32 before testing that the correct
                # value is retrieved
                assert n.all(
                    tbdata["c3"].astype(np.float32) == tbdata2["c3"].astype(np.float32)
                )
                # c4 is a boolean column in the original table; we want ASCII
                # columns to convert these to columns of 'T'/'F' strings
                assert n.all(np.where(tbdata["c4"], "T", "F") == tbdata2["c4"])
//  [python]
def test_bintable_to_asciitable(self):
        """Tests initializing a TableHDU with the data from a BinTableHDU."""

        with fits.open(self.data("tb.fits")) as hdul:
            tbdata = hdul[1].data
            tbhdu = fits.TableHDU(data=tbdata)
            tbhdu.writeto(self.temp("test.fits"), overwrite=True)
            with fits.open(self.temp("test.fits")) as hdul2:
                tbdata2 = hdul2[1].data
                assert np.all(tbdata["c1"] == tbdata2["c1"])
                assert np.all(tbdata["c2"] == tbdata2["c2"])
                # c3 gets converted from float32 to float64 when writing
                # test.fits, so cast to float32 before testing that the correct
                # value is retrieved
                assert np.all(
                    tbdata["c3"].astype(np.float32) == tbdata2["c3"].astype(np.float32)
                )
                # c4 is a boolean column in the original table; we want ASCII
                # columns to convert these to columns of 'T'/'F' strings
                assert np.all(np.where(tbdata["c4"], "T", "F") == tbdata2["c4"])
//  [python]
def test_pseudo_unsigned_ints(self):
        """
        Tests updating a table column containing pseudo-unsigned ints.
        """

        data = np.array([1, 2, 3], dtype=np.uint32)
        col = fits.Column(name="A", format="1J", bzero=2**31, array=data)
        thdu = fits.BinTableHDU.from_columns([col])
        thdu.writeto(self.temp("test.fits"))

        # Test that the file wrote out correctly
        with fits.open(self.temp("test.fits"), uint=True) as hdul:
            hdu = hdul[1]
            assert "TZERO1" in hdu.header
            assert hdu.header["TZERO1"] == 2**31
            assert hdu.data["A"].dtype == np.dtype("uint32")
            assert n.all(hdu.data["A"] == data)

            # Test updating the unsigned int data
            hdu.data["A"][0] = 99
            hdu.writeto(self.temp("test2.fits"))

        with fits.open(self.temp("test2.fits"), uint=True) as hdul:
            hdu = hdul[1]
            assert "TZERO1" in hdu.header
            assert hdu.header["TZERO1"] == 2**31
            assert hdu.data["A"].dtype == np.dtype("uint32")
            assert n.all(hdu.data["A"] == [99, 2, 3])
//  [python]
def test_pseudo_unsigned_ints(self):
        """
        Tests updating a table column containing pseudo-unsigned ints.
        """

        data = np.array([1, 2, 3], dtype=np.uint32)
        col = fits.Column(name="A", format="1J", bzero=2**31, array=data)
        thdu = fits.BinTableHDU.from_columns([col])
        thdu.writeto(self.temp("test.fits"))

        # Test that the file wrote out correctly
        with fits.open(self.temp("test.fits"), uint=True) as hdul:
            hdu = hdul[1]
            assert "TZERO1" in hdu.header
            assert hdu.header["TZERO1"] == 2**31
            assert hdu.data["A"].dtype == np.dtype("uint32")
            assert np.all(hdu.data["A"] == data)

            # Test updating the unsigned int data
            hdu.data["A"][0] = 99
            hdu.writeto(self.temp("test2.fits"))

        with fits.open(self.temp("test2.fits"), uint=True) as hdul:
            hdu = hdul[1]
            assert "TZERO1" in hdu.header
            assert hdu.header["TZERO1"] == 2**31
            assert hdu.data["A"].dtype == np.dtype("uint32")
            assert np.all(hdu.data["A"] == [99, 2, 3])
//  [python]
def test_x_column_deepcopy(self):
        """
        Regression test for https://github.com/astropy/astropy/pull/4514

        Tests that columns with the X (bit array) format can be deep-copied.
        """

        c = fits.Column("xcol", format="5X", array=[1, 0, 0, 1, 0])
        c2 = copy.deepcopy(c)
        assert c2.name == c.name
        assert c2.format == c.format
        assert np.all(c2.array == c.array)
//  [python]
def test_zero_precision_float_column(self):
        """
        Regression test for https://github.com/astropy/astropy/issues/3422
        """

        c = fits.Column("TEST", "F5.0", array=[1.1, 2.2, 3.3])
        # The decimal places will be clipped
        t = fits.TableHDU.from_columns([c])
        t.writeto(self.temp("test.fits"))

        with fits.open(self.temp("test.fits")) as hdul:
            assert hdul[1].header["TFORM1"] == "F5.0"
            assert hdul[1].data["TEST"].dtype == np.dtype("float64")
            assert np.all(hdul[1].data["TEST"] == [1.0, 2.0, 3.0])

            # Check how the raw data looks
            raw = np.rec.recarray.field(hdul[1].data, "TEST")
            assert raw.tobytes() == b"   1.   2.   3."
//  [python]
def test_multidimension_table_from_numpy_rec_columns(self):
        """Regression test for https://github.com/astropy/astropy/issues/5280
        and https://github.com/astropy/astropy/issues/5287

        multidimentional tables can now be written with the correct TDIM.
        Author: Stephen Bailey.
        """

        dtype = [
            ("x", (str, 5)),  # 1D column of 5-character strings
            ("y", (str, 3), (4,)),  # 2D column; each row is four 3-char strings
        ]
        data = np.zeros(2, dtype=dtype)
        data["x"] = ["abcde", "xyz"]
        data["y"][0] = ["A", "BC", "DEF", "123"]
        data["y"][1] = ["X", "YZ", "PQR", "999"]
        table = Table(data)

        # Test convenience functions io.fits.writeto / getdata
        fits.writeto(self.temp("test.fits"), data)
        dx = fits.getdata(self.temp("test.fits"))
        assert data["x"].dtype == dx["x"].dtype
        assert data["y"].dtype == dx["y"].dtype
        assert np.alltrue(data["x"] == dx["x"]), f"x: {data['x']} != {dx['x']}"
        assert np.alltrue(data["y"] == dx["y"]), f"y: {data['y']} != {dx['y']}"

        # Test fits.BinTableHDU(data) and avoid convenience functions
        hdu0 = fits.PrimaryHDU()
        hdu1 = fits.BinTableHDU(data)
        hx = fits.HDUList([hdu0, hdu1])
        hx.writeto(self.temp("test2.fits"))
        fx = fits.open(self.temp("test2.fits"))
        dx = fx[1].data
        fx.close()
        assert data["x"].dtype == dx["x"].dtype
        assert data["y"].dtype == dx["y"].dtype
        assert np.alltrue(data["x"] == dx["x"]), f"x: {data['x']} != {dx['x']}"
        assert np.alltrue(data["y"] == dx["y"]), f"y: {data['y']} != {dx['y']}"

        # Test Table write and read
        table.write(self.temp("test3.fits"))
        tx = Table.read(self.temp("test3.fits"), character_as_bytes=False)
        assert table["x"].dtype == tx["x"].dtype
        assert table["y"].dtype == tx["y"].dtype
        assert np.alltrue(table["x"] == tx["x"]), f"x: {table['x']} != {tx['x']}"
        assert np.alltrue(table["y"] == tx["y"]), f"y: {table['y']} != {tx['y']}"
//  [python]
def test_multidimension_table_from_numpy_rec_columns(self):
        """Regression test for https://github.com/astropy/astropy/issues/5280
        and https://github.com/astropy/astropy/issues/5287

        multidimentional tables can now be written with the correct TDIM.
        Author: Stephen Bailey.
        """

        dtype = [
            ("x", (str, 5)),  # 1D column of 5-character strings
            ("y", (str, 3), (4,)),  # 2D column; each row is four 3-char strings
        ]
        data = np.zeros(2, dtype=dtype)
        data["x"] = ["abcde", "xyz"]
        data["y"][0] = ["A", "BC", "DEF", "123"]
        data["y"][1] = ["X", "YZ", "PQR", "999"]
        table = Table(data)

        # Test convenience functions io.fits.writeto / getdata
        fits.writeto(self.temp("test.fits"), data)
        dx = fits.getdata(self.temp("test.fits"))
        assert data["x"].dtype == dx["x"].dtype
        assert data["y"].dtype == dx["y"].dtype
        assert np.all(data["x"] == dx["x"]), f"x: {data['x']} != {dx['x']}"
        assert np.all(data["y"] == dx["y"]), f"y: {data['y']} != {dx['y']}"

        # Test fits.BinTableHDU(data) and avoid convenience functions
        hdu0 = fits.PrimaryHDU()
        hdu1 = fits.BinTableHDU(data)
        hx = fits.HDUList([hdu0, hdu1])
        hx.writeto(self.temp("test2.fits"))
        fx = fits.open(self.temp("test2.fits"))
        dx = fx[1].data
        fx.close()
        assert data["x"].dtype == dx["x"].dtype
        assert data["y"].dtype == dx["y"].dtype
        assert np.all(data["x"] == dx["x"]), f"x: {data['x']} != {dx['x']}"
        assert np.all(data["y"] == dx["y"]), f"y: {data['y']} != {dx['y']}"

        # Test Table write and read
        table.write(self.temp("test3.fits"))
        tx = Table.read(self.temp("test3.fits"), character_as_bytes=False)
        assert table["x"].dtype == tx["x"].dtype
        assert table["y"].dtype == tx["y"].dtype
        assert np.all(table["x"] == tx["x"]), f"x: {table['x']} != {tx['x']}"
        assert np.all(table["y"] == tx["y"]), f"y: {table['y']} != {tx['y']}"
//  [python]
def test_p_column_deepcopy(self):
        """
        Regression test for https://github.com/astropy/astropy/pull/4514

        Tests that columns with the P/Q formats (variable length arrays) can be
        deep-copied.
        """

        c = fits.Column("pcol", format="PJ", array=[[1, 2], [3, 4, 5]])
        c2 = copy.deepcopy(c)
        assert c2.name == c.name
        assert c2.format == c.format
        assert n.all(c2.array[0] == c.array[0])
        assert n.all(c2.array[1] == c.array[1])

        c3 = fits.Column("qcol", format="QJ", array=[[1, 2], [3, 4, 5]])
        c4 = copy.deepcopy(c3)
        assert c4.name == c3.name
        assert c4.format == c3.format
        assert n.all(c4.array[0] == c3.array[0])
        assert n.all(c4.array[1] == c3.array[1])
//  [python]
def test_p_column_deepcopy(self):
        """
        Regression test for https://github.com/astropy/astropy/pull/4514

        Tests that columns with the P/Q formats (variable length arrays) can be
        deep-copied.
        """

        c = fits.Column("pcol", format="PJ", array=[[1, 2], [3, 4, 5]])
        c2 = copy.deepcopy(c)
        assert c2.name == c.name
        assert c2.format == c.format
        assert np.all(c2.array[0] == c.array[0])
        assert np.all(c2.array[1] == c.array[1])

        c3 = fits.Column("qcol", format="QJ", array=[[1, 2], [3, 4, 5]])
        c4 = copy.deepcopy(c3)
        assert c4.name == c3.name
        assert c4.format == c3.format
        assert np.all(c4.array[0] == c3.array[0])
        assert np.all(c4.array[1] == c3.array[1])
//  [python]
def test_mbget_multidim(self):
        sz = 17 * 31
        arr = np.random.rand(sz)
        arr = arr.reshape((17, 31))
        idx = np.random.randint(0, 17, 100)

        result_np = arr[idx]
        result_rt = mbget(arr, idx)

        self.assertTrue(bool(np.alltrue(result_np == result_rt)))
        self.assertTrue(isinstance(result_rt, FastArray))
//  [python]
def test_mbget_multidim(self):
        sz = 17 * 31
        arr = np.random.rand(sz)
        arr = arr.reshape((17, 31))
        idx = np.random.randint(0, 17, 100)

        result_np = arr[idx]
        result_rt = mbget(arr, idx)

        self.assertTrue(bool(np.all(result_np == result_rt)))
        self.assertTrue(isinstance(result_rt, FastArray))
//  [python]
def test_isnan_sentinel(self):
        dts = [
            np.int8,
            np.uint8,
            np.int16,
            np.uint16,
            np.int32,
            np.uint32,
            np.int64,
            np.uint64,
            np.float32,
            np.float64,
        ]
        arr = FastArray([-128, 2, 3, 4, 5], dtype=np.int8)
        correct = [True, False, False, False, False]

        for dt in dts:
            arr = arr.astype(dt)
            if not isnan(INVALID_DICT[arr.dtype.num]):
                self.assertEqual(INVALID_DICT[arr.dtype.num], arr[0])
            result = arr.isnan()
            self.assertTrue(
                bool(n.all(result == correct)),
                msg=f"isnan failed for dtype {np.dtype(dt)}",
            )
//  [python]
def test_isnan_sentinel(self):
        dts = [
            np.int8,
            np.uint8,
            np.int16,
            np.uint16,
            np.int32,
            np.uint32,
            np.int64,
            np.uint64,
            np.float32,
            np.float64,
        ]
        arr = FastArray([-128, 2, 3, 4, 5], dtype=np.int8)
        correct = [True, False, False, False, False]

        for dt in dts:
            arr = arr.astype(dt)
            if not isnan(INVALID_DICT[arr.dtype.num]):
                self.assertEqual(INVALID_DICT[arr.dtype.num], arr[0])
            result = arr.isnan()
            self.assertTrue(
                bool(np.all(result == correct)),
                msg=f"isnan failed for dtype {np.dtype(dt)}",
            )
//  [python]
def test_overlappedarrays(self):
        a = arange(21).reshape((7, 3), order="F")
        b = arange(21).reshape((7, 3), order="F")._np
        a[0][1:] -= a[0][:-1]
        b[0][1:] -= b[0][:-1]
        self.assertTrue(bool(n.all(a == b)))
//  [python]
def test_overlappedarrays(self):
        a = arange(21).reshape((7, 3), order="F")
        b = arange(21).reshape((7, 3), order="F")._np
        a[0][1:] -= a[0][:-1]
        b[0][1:] -= b[0][:-1]
        self.assertTrue(bool(np.all(a == b)))
//  [python]
def test_sign(self):
        x = [1, -2, 3, np.nan, np.nan, np.nan]
        result = FastArray(x).sign().astype(np.int32)
        goodresult = np.sign(x).astype(np.int32)
        self.assertTrue(np.all(result == goodresult))
//  [python]
def test_rollingsum(self):
        x = FastArray([1, 2, 3, np.nan, np.nan, np.nan])
        result = x.rolling_nansum()
        goodresult = FastArray([1.0, 3.0, 6.0, 5.0, 3.0, 0.0])
        self.assertTrue(n.all(result == goodresult))

        # int based test of invalids
        result = x.astype(np.int64).rolling_nansum()
        goodresult = FastArray([1.0, 3.0, 6.0, 5.0, 3.0, 0.0]).astype(np.int64)
        self.assertTrue(n.all(result == goodresult))
//  [python]
def test_rollingsum(self):
        x = FastArray([1, 2, 3, np.nan, np.nan, np.nan])
        result = x.rolling_nansum()
        goodresult = FastArray([1.0, 3.0, 6.0, 5.0, 3.0, 0.0])
        self.assertTrue(np.all(result == goodresult))

        # int based test of invalids
        result = x.astype(np.int64).rolling_nansum()
        goodresult = FastArray([1.0, 3.0, 6.0, 5.0, 3.0, 0.0]).astype(np.int64)
        self.assertTrue(np.all(result == goodresult))
//  [python]
def test_nan_sentinels(self):
        all_dtypes = [np.dtype(c) for c in rt.rt_enum.NumpyCharTypes.All]
        nan_dtypes = [
            dt
            for dt in all_dtypes
            if (
                dt.num in rt.rt_enum.INVALID_DICT
                and dt.char not in rt.rt_enum.NumpyCharTypes.Unsupported
                and dt.char != "?"  # bool nan is broken (neither supported nor unsupported)
            )
        ]
        ops = {
            True: ["isnan", "isnanorzero", "isnotfinite"],
            False: ["isnotnan", "isfinite"],
        }

        for dt in nan_dtypes:
            arr = rt.FA([rt.rt_enum.INVALID_DICT[dt.num], 1], dtype=dt)
            try:
                is_nan = isnan(INVALID_DICT[arr.dtype.num])
            except:
                is_nan = False
            if not is_nan:
                self.assertEqual(INVALID_DICT[arr.dtype.num], arr[0], msg=f"failed for {dt!r}")

            for expected, names in ops.items():
                correct = [expected, not expected]

                for name in names:
                    # test the fa.xxx() variant
                    fn = getattr(arr, name)
                    result = fn()
                    self.assertTrue(
                        bool(n.all(result == correct)),
                        msg=f"fa.{name} failed for dtype {dt!r}",
                    )

                    # test the rt.xxx(fa) variant
                    fn = getattr(rt, name)
                    result = fn(arr)
                    self.assertTrue(
                        bool(n.all(result == correct)),
                        msg=f"rt.{name} failed for dtype {dt!r}",
                    )