jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __call__(self, y_pred, y):
        """
        'y_preds' is expected to have binarized predictions and 'y' should be in one-hot format.

        Args:
            - **y_pred** (ndarray) - Input data to compute. It must be one-hot format and first dim is batch.
            - **y** (ndarray) - Ground truth to compute the metric. It must be one-hot format and first dim is batch.

        Raises:
            ValueError: If `metric_name` is empty.
            ValueError: when `y_pred` has less than two dimensions.
        """
        if not np.all(y.astype(np.uint8) == y):
            raise ValueError("For 'ConfusionMatrix.update', the true value (input[1]) must be a binarized ndarray.")

        dims = y_pred.ndim
        if dims < 2:
            raise ValueError(f"For 'ConfusionMatrix.update', the predicted value (input[0]) must have at least 2 "
                             f"dimensions, but got {dims}.")

        if dims == 2 or (dims == 3 and y_pred.shape[-1] == 1):
            if self.calculation_method:
                self.calculation_method = False

        confusion_matrix = _get_confusion_matrix(y_pred=y_pred, y=y, skip_channel=self.skip_channel)

        if self.calculation_method:
            if isinstance(self.metric_name, str):
                sub_confusion_matrix = _compute_confusion_matrix_metric(self.metric_name, confusion_matrix)
                chart, not_nans = _decrease_metric(sub_confusion_matrix, self.decrease)
                return chart, not_nans

            if not self.metric_name:
                raise ValueError("For 'ConfusionMatrix', the argument 'metric_name' cannot be None.")

            results = []
            for metric_name in self.metric_name:
                sub_confusion_matrix = _compute_confusion_matrix_metric(metric_name, confusion_matrix)
                chart, not_nans = _decrease_metric(sub_confusion_matrix, self.decrease)
                results.append(chart)
                results.append(not_nans)
            return results

        return confusion_matrix
//  [python]
def _line_distances(points, start, end):
    if n.all(start == end):
        return np.linalg.norm(points - start, axis=1)

    vec = end - start
    cross = np.cross(vec, start - points)
    return np.divide(abs(cross), np.linalg.norm(vec))
//  [python]
def _line_distances(points, start, end):
    if np.all(start == end):
        return np.linalg.norm(points - start, axis=1)

    vec = end - start
    cross = np.cross(vec, start - points)
    return np.divide(abs(cross), np.linalg.norm(vec))
//  [python]
def test_draw_different_samples(self):
        with pm.Model():
            x = pm.Normal("x")

        x_draws_1 = pm.draw(x, 100)
        x_draws_2 = pm.draw(x, 100)
        assert not n.all(n.isclose(x_draws_1, x_draws_2))
//  [python]
def test_draw_different_samples(self):
        with pm.Model():
            x = pm.Normal("x")

        x_draws_1 = pm.draw(x, 100)
        x_draws_2 = pm.draw(x, 100)
        assert not np.all(np.isclose(x_draws_1, x_draws_2))
//  [python]
def is_not_skewed(sicd: SICDType, dimension: int) -> bool:
    """
    Check if the sicd structure is not skewed along the provided dimension.

    Parameters
    ----------
    sicd : SICDType
    dimension : int

    Returns
    -------
    bool
    """

    if dimension == 0:
        if sicd.Grid is None or sicd.Grid.Row is None or sicd.Grid.Row.DeltaKCOAPoly is None:
            return True
        return n.all(sicd.Grid.Row.DeltaKCOAPoly.get_array(dtype='float64') == 0)
    else:
        if sicd.Grid is None or sicd.Grid.Col is None or sicd.Grid.Col.DeltaKCOAPoly is None:
            return True
        return n.all(sicd.Grid.Col.DeltaKCOAPoly.get_array(dtype='float64') == 0)
//  [python]
def is_not_skewed(sicd: SICDType, dimension: int) -> bool:
    """
    Check if the sicd structure is not skewed along the provided dimension.

    Parameters
    ----------
    sicd : SICDType
    dimension : int

    Returns
    -------
    bool
    """

    if dimension == 0:
        if sicd.Grid is None or sicd.Grid.Row is None or sicd.Grid.Row.DeltaKCOAPoly is None:
            return True
        return numpy.all(sicd.Grid.Row.DeltaKCOAPoly.get_array(dtype='float64') == 0)
    else:
        if sicd.Grid is None or sicd.Grid.Col is None or sicd.Grid.Col.DeltaKCOAPoly is None:
            return True
        return numpy.all(sicd.Grid.Col.DeltaKCOAPoly.get_array(dtype='float64') == 0)
//  [python]
def is_uniform_weight(sicd: SICDType, dimension: int) -> bool:
    """
    Check if the sicd structure is has uniform weight along the provided dimension.

    Parameters
    ----------
    sicd : SICDType
    dimension : int

    Returns
    -------
    bool
    """

    if dimension == 0:
        if sicd.Grid is None or sicd.Grid.Row is None:
            return False
        dir_param = sicd.Grid.Row
    else:
        if sicd.Grid is None or sicd.Grid.Col is None:
            return False
        dir_param = sicd.Grid.Col

    if dir_param.WgtType is not None and dir_param.WgtType.WindowName == 'UNIFORM':
        return True
    if dir_param.WgtFunct is not None:
        return n.all(dir_param.WgtFunct == dir_param.WgtFunct[0])
    return True
//  [python]
def is_uniform_weight(sicd: SICDType, dimension: int) -> bool:
    """
    Check if the sicd structure is has uniform weight along the provided dimension.

    Parameters
    ----------
    sicd : SICDType
    dimension : int

    Returns
    -------
    bool
    """

    if dimension == 0:
        if sicd.Grid is None or sicd.Grid.Row is None:
            return False
        dir_param = sicd.Grid.Row
    else:
        if sicd.Grid is None or sicd.Grid.Col is None:
            return False
        dir_param = sicd.Grid.Col

    if dir_param.WgtType is not None and dir_param.WgtType.WindowName == 'UNIFORM':
        return True
    if dir_param.WgtFunct is not None:
        return numpy.all(dir_param.WgtFunct == dir_param.WgtFunct[0])
    return True
//  [python]
def apply_skew_poly(
        input_data: numpy.ndarray,
        delta_kcoa_poly: numpy.ndarray,
        row_array: numpy.ndarray,
        col_array: numpy.ndarray,
        fft_sgn: int,
        dimension: int,
        forward: bool = False) -> numpy.ndarray:
    """
    Performs the skew operation on the complex array, according to the provided
    delta kcoa polynomial.

    Parameters
    ----------
    input_data : numpy.ndarray
        The input data.
    delta_kcoa_poly : numpy.ndarray
        The delta kcoa polynomial to use.
    row_array : numpy.ndarray
        The row array, should agree with input_data first dimension definition.
    col_array : numpy.ndarray
        The column array, should agree with input_data second dimension definition.
    fft_sgn : int
        The fft sign to use.
    dimension : int
        The dimension to apply along.
    forward : bool
        If True, this shifts forward (i.e. skews), otherwise applies in inverse
        (i.e. deskew) direction.

    Returns
    -------
    numpy.ndarray
    """

    if n.all(delta_kcoa_poly == 0):
        return input_data

    delta_kcoa_poly_int = polynomial.polyint(delta_kcoa_poly, axis=dimension)
    if forward:
        fft_sgn *= -1
    return input_data*numpy.exp(1j*fft_sgn*2*numpy.pi*polynomial.polygrid2d(
        row_array, col_array, delta_kcoa_poly_int))
//  [python]
def apply_skew_poly(
        input_data: numpy.ndarray,
        delta_kcoa_poly: numpy.ndarray,
        row_array: numpy.ndarray,
        col_array: numpy.ndarray,
        fft_sgn: int,
        dimension: int,
        forward: bool = False) -> numpy.ndarray:
    """
    Performs the skew operation on the complex array, according to the provided
    delta kcoa polynomial.

    Parameters
    ----------
    input_data : numpy.ndarray
        The input data.
    delta_kcoa_poly : numpy.ndarray
        The delta kcoa polynomial to use.
    row_array : numpy.ndarray
        The row array, should agree with input_data first dimension definition.
    col_array : numpy.ndarray
        The column array, should agree with input_data second dimension definition.
    fft_sgn : int
        The fft sign to use.
    dimension : int
        The dimension to apply along.
    forward : bool
        If True, this shifts forward (i.e. skews), otherwise applies in inverse
        (i.e. deskew) direction.

    Returns
    -------
    numpy.ndarray
    """

    if numpy.all(delta_kcoa_poly == 0):
        return input_data

    delta_kcoa_poly_int = polynomial.polyint(delta_kcoa_poly, axis=dimension)
    if forward:
        fft_sgn *= -1
    return input_data*numpy.exp(1j*fft_sgn*2*numpy.pi*polynomial.polygrid2d(
        row_array, col_array, delta_kcoa_poly_int))
//  [python]
def test_draw_pytensor_function_kwargs(self):
        sharedvar = pytensor.shared(0)
        x = pm.DiracDelta.dist(0.0)
        y = x + sharedvar
        draws = pm.draw(
            y,
            draws=5,
            mode=Mode("py"),
            updates={sharedvar: sharedvar + 1},
        )
        assert n.all(draws == n.arange(5))
//  [python]
def test_draw_pytensor_function_kwargs(self):
        sharedvar = pytensor.shared(0)
        x = pm.DiracDelta.dist(0.0)
        y = x + sharedvar
        draws = pm.draw(
            y,
            draws=5,
            mode=Mode("py"),
            updates={sharedvar: sharedvar + 1},
        )
        assert np.all(draws == np.arange(5))
//  [python]
def initial_file_load(self, filename):
        """
        Initial file load to get object ids.

        Parameters
        ----------

        filename : string
            The filename to load

        Returns
        -------

        obj_ids : array
             All object ids. Note, this may be extremely long.
        unique_obj_ids : array
             All unique object ides.
        is_mat_file : boolean
             Was filename a .mat file?
        data_file : File object
             Opened reference to original data file.
        extra : dict
             Dictionary with addtional information.
        """
        if filename not in self.loaded_filename_cache:
            data_file = _initial_file_load(filename)
            (obj_ids, unique_obj_ids, is_mat_file, extra) = _get_initial_file_info(
                data_file
            )

            if 0:
                # Why did I used to have this assertion check?
                diff = numpy.int64(obj_ids[1:]) - numpy.int64(obj_ids[:-1])
                # for fast search:
                assert n.all(diff >= 0)  # make sure obj_ids ascending

            self.loaded_filename_cache[filename] = (
                obj_ids,
                unique_obj_ids,
                is_mat_file,
                extra,
            )
            # maintain only a weak ref to data file:
            self.loaded_filename_cache2[filename] = data_file

        (obj_ids, unique_obj_ids, is_mat_file, extra) = self.loaded_filename_cache[
            filename
        ]
        data_file = self.loaded_filename_cache2[filename]

        self.loaded_datafile_cache[data_file] = True
        return obj_ids, unique_obj_ids, is_mat_file, data_file, extra
//  [python]
def initial_file_load(self, filename):
        """
        Initial file load to get object ids.

        Parameters
        ----------

        filename : string
            The filename to load

        Returns
        -------

        obj_ids : array
             All object ids. Note, this may be extremely long.
        unique_obj_ids : array
             All unique object ides.
        is_mat_file : boolean
             Was filename a .mat file?
        data_file : File object
             Opened reference to original data file.
        extra : dict
             Dictionary with addtional information.
        """
        if filename not in self.loaded_filename_cache:
            data_file = _initial_file_load(filename)
            (obj_ids, unique_obj_ids, is_mat_file, extra) = _get_initial_file_info(
                data_file
            )

            if 0:
                # Why did I used to have this assertion check?
                diff = numpy.int64(obj_ids[1:]) - numpy.int64(obj_ids[:-1])
                # for fast search:
                assert numpy.all(diff >= 0)  # make sure obj_ids ascending

            self.loaded_filename_cache[filename] = (
                obj_ids,
                unique_obj_ids,
                is_mat_file,
                extra,
            )
            # maintain only a weak ref to data file:
            self.loaded_filename_cache2[filename] = data_file

        (obj_ids, unique_obj_ids, is_mat_file, extra) = self.loaded_filename_cache[
            filename
        ]
        data_file = self.loaded_filename_cache2[filename]

        self.loaded_datafile_cache[data_file] = True
        return obj_ids, unique_obj_ids, is_mat_file, data_file, extra
//  [python]
def __init__(self, data_path, use_one_hot=False, validation_size=0,
                 use_torch_augmentation=False):
        super().__init__()

        fmnist_train = FashionMNIST(data_path, train=True, download=True)
        fmnist_test = FashionMNIST(data_path, train=False, download=True)
        assert n.all(n.equal(fmnist_train.data.shape, [60000, 28, 28]))
        assert n.all(n.equal(fmnist_test.data.shape, [10000, 28, 28]))

        train_inputs = fmnist_train.data.numpy().reshape(60000, -1)
        test_inputs = fmnist_test.data.numpy().reshape(10000, -1)
        train_labels = fmnist_train.targets.numpy().reshape(60000, 1)
        test_labels = fmnist_test.targets.numpy().reshape(10000, 1)

        images = n.concatenate([train_inputs, test_inputs], axis=0)
        labels = n.concatenate([train_labels, test_labels], axis=0)

        # Scale images into a range between 0 and 1. Such that it is identical
        # to the default MNIST scale in `data.dataset.mnist_data`.
        images = images / 255

        val_inds = None
        train_inds = n.arange(train_labels.size)
        test_inds = n.arange(train_labels.size,
                              train_labels.size + test_labels.size)

        if validation_size > 0:
            if validation_size >= train_inds.size:
                raise ValueError('Validation set must contain less than %d ' \
                                 % (train_inds.size) + 'samples!')

            val_inds = n.arange(validation_size)
            train_inds = n.arange(validation_size, train_inds.size)

        # Bring everything into the internal structure of the Dataset class.
        self._data['classification'] = True
        self._data['sequence'] = False
        self._data['num_classes'] = 10
        self._data['is_one_hot'] = use_one_hot
        self._data['in_data'] = images
        self._data['in_shape'] = [28, 28, 1]
        self._data['out_shape'] = [10 if use_one_hot else 1]
        self._data['val_inds'] = val_inds
        self._data['train_inds'] = train_inds
        self._data['test_inds'] = test_inds

        if use_one_hot:
            labels = self._to_one_hot(labels)

        self._data['out_data'] = labels

        # Information specific to this dataset.
        assert n.all([fmnist_train.classes[i] == c for i, c in \
                       enumerate(fmnist_test.classes)])
        self._data['fmnist'] = dict()
        self._data['fmnist']['classes'] = fmnist_train.classes

        # Initialize PyTorch data augmentation.
        self._augment_inputs = False
        if use_torch_augmentation:
            self._augment_inputs = True
            self._train_transform, self._test_transform = \
                MNISTData.torch_input_transforms(use_random_hflips=True)
//  [python]
def __init__(self, data_path, use_one_hot=False, validation_size=0,
                 use_torch_augmentation=False):
        super().__init__()

        fmnist_train = FashionMNIST(data_path, train=True, download=True)
        fmnist_test = FashionMNIST(data_path, train=False, download=True)
        assert np.all(np.equal(fmnist_train.data.shape, [60000, 28, 28]))
        assert np.all(np.equal(fmnist_test.data.shape, [10000, 28, 28]))

        train_inputs = fmnist_train.data.numpy().reshape(60000, -1)
        test_inputs = fmnist_test.data.numpy().reshape(10000, -1)
        train_labels = fmnist_train.targets.numpy().reshape(60000, 1)
        test_labels = fmnist_test.targets.numpy().reshape(10000, 1)

        images = np.concatenate([train_inputs, test_inputs], axis=0)
        labels = np.concatenate([train_labels, test_labels], axis=0)

        # Scale images into a range between 0 and 1. Such that it is identical
        # to the default MNIST scale in `data.dataset.mnist_data`.
        images = images / 255

        val_inds = None
        train_inds = np.arange(train_labels.size)
        test_inds = np.arange(train_labels.size,
                              train_labels.size + test_labels.size)

        if validation_size > 0:
            if validation_size >= train_inds.size:
                raise ValueError('Validation set must contain less than %d ' \
                                 % (train_inds.size) + 'samples!')

            val_inds = np.arange(validation_size)
            train_inds = np.arange(validation_size, train_inds.size)

        # Bring everything into the internal structure of the Dataset class.
        self._data['classification'] = True
        self._data['sequence'] = False
        self._data['num_classes'] = 10
        self._data['is_one_hot'] = use_one_hot
        self._data['in_data'] = images
        self._data['in_shape'] = [28, 28, 1]
        self._data['out_shape'] = [10 if use_one_hot else 1]
        self._data['val_inds'] = val_inds
        self._data['train_inds'] = train_inds
        self._data['test_inds'] = test_inds

        if use_one_hot:
            labels = self._to_one_hot(labels)

        self._data['out_data'] = labels

        # Information specific to this dataset.
        assert np.all([fmnist_train.classes[i] == c for i, c in \
                       enumerate(fmnist_test.classes)])
        self._data['fmnist'] = dict()
        self._data['fmnist']['classes'] = fmnist_train.classes

        # Initialize PyTorch data augmentation.
        self._augment_inputs = False
        if use_torch_augmentation:
            self._augment_inputs = True
            self._train_transform, self._test_transform = \
                MNISTData.torch_input_transforms(use_random_hflips=True)
//  [python]
def test_pytensor_function_kwargs(self):
        sharedvar = pytensor.shared(0)
        with pm.Model() as m:
            x = pm.DiracDelta("x", 0)
            y = pm.Deterministic("y", x + sharedvar)

            prior = pm.sample_prior_predictive(
                samples=5,
                return_inferencedata=False,
                compile_kwargs=dict(
                    mode=Mode("py"),
                    updates={sharedvar: sharedvar + 1},
                ),
            )

        assert np.all(prior["y"] == np.arange(5))
//  [python]
def test_pytensor_function_kwargs(self):
        sharedvar = pytensor.shared(0)
        with pm.Model() as m:
            x = pm.DiracDelta("x", 0.0)
            y = pm.Deterministic("y", x + sharedvar)

            pp = pm.sample_posterior_predictive(
                trace=az_from_dict({"x": np.arange(5)}),
                var_names=["y"],
                return_inferencedata=False,
                compile_kwargs=dict(
                    mode=Mode("py"),
                    updates={sharedvar: sharedvar + 1},
                ),
            )

        assert n.all(pp["y"] == np.arange(5) * 2)
//  [python]
def test_pytensor_function_kwargs(self):
        sharedvar = pytensor.shared(0)
        with pm.Model() as m:
            x = pm.DiracDelta("x", 0.0)
            y = pm.Deterministic("y", x + sharedvar)

            pp = pm.sample_posterior_predictive(
                trace=az_from_dict({"x": np.arange(5)}),
                var_names=["y"],
                return_inferencedata=False,
                compile_kwargs=dict(
                    mode=Mode("py"),
                    updates={sharedvar: sharedvar + 1},
                ),
            )

        assert np.all(pp["y"] == np.arange(5) * 2)
//  [python]
def test_deterministics_out_of_idata(self, multitrace):
        draws = 10
        chains = 2
        coords = {"draw": range(draws), "chain": range(chains)}
        ds = xr.Dataset(
            {
                "a": xr.DataArray(
                    [[0] * draws] * chains,
                    coords=coords,
                    dims=["chain", "draw"],
                )
            },
            coords=coords,
        )
        with pm.Model() as m:
            a = pm.Normal("a")
            if multitrace:
                straces = []
                for chain in ds.chain:
                    strace = pm.backends.NDArray(model=m, vars=[a])
                    strace.setup(len(ds.draw), int(chain))
                    strace.values = {"a": ds.a.sel(chain=chain).data}
                    strace.draw_idx = len(ds.draw)
                    straces.append(strace)
                trace = MultiTrace(straces)
            else:
                trace = ds

            d = pm.Deterministic("d", a - 4)
            pm.Normal("c", d, sigma=0.01)
            ppc = pm.sample_posterior_predictive(trace, var_names="c", return_inferencedata=True)
        assert n.all(n.abs(ppc.posterior_predictive.c + 4) <= 0.1)
//  [python]
def test_deterministics_out_of_idata(self, multitrace):
        draws = 10
        chains = 2
        coords = {"draw": range(draws), "chain": range(chains)}
        ds = xr.Dataset(
            {
                "a": xr.DataArray(
                    [[0] * draws] * chains,
                    coords=coords,
                    dims=["chain", "draw"],
                )
            },
            coords=coords,
        )
        with pm.Model() as m:
            a = pm.Normal("a")
            if multitrace:
                straces = []
                for chain in ds.chain:
                    strace = pm.backends.NDArray(model=m, vars=[a])
                    strace.setup(len(ds.draw), int(chain))
                    strace.values = {"a": ds.a.sel(chain=chain).data}
                    strace.draw_idx = len(ds.draw)
                    straces.append(strace)
                trace = MultiTrace(straces)
            else:
                trace = ds

            d = pm.Deterministic("d", a - 4)
            pm.Normal("c", d, sigma=0.01)
            ppc = pm.sample_posterior_predictive(trace, var_names="c", return_inferencedata=True)
        assert np.all(np.abs(ppc.posterior_predictive.c + 4) <= 0.1)
//  [python]
def test_condens_ndarray_not_possible():
    a = np.arange(10)
    assert a.copy().all() == condens_ndarray(a.copy()).all()

    a = np.random.random((3, 4, 2))
    assert a.copy().all() == condens_ndarray(a.copy()).all()
//  [python]
def test_condens_ndarray_not_possible():
    a = np.arange(10)
    assert np.all(a.copy() == condens_ndarray(a.copy()))

    a = np.random.random((3, 4, 2))
    assert np.all(a.copy() == condens_ndarray(a.copy()))
//  [python]
def get_rotation_matrix_and_quaternion(rotation):
    rotation_orig = rotation
    rotation = np.array(rotation)
    if rotation.ndim==2:
        assert rotation.shape==(3,3)
        if not n.all(n.isnan( rotation )):
            assert is_rotation_matrix(rotation)

        rmat = rotation

        rnew = np.eye(4)
        rnew[:3,:3] = rmat
        rquat = quaternion_from_matrix(rnew)
        if not n.all(n.isnan( rquat )):
            R2 = quaternion_matrix(rquat)[:3,:3]
            assert np.allclose(rmat,R2)
    elif rotation.ndim==0:
        assert rotation.dtype == object
        rotation = (rotation_orig.x,
                    rotation_orig.y,
                    rotation_orig.z,
                    rotation_orig.w)
        return get_rotation_matrix_and_quaternion(rotation)
    else:
        assert rotation.ndim==1
        assert rotation.shape==(4,)
        rquat = rotation
        rmat = quaternion_matrix(rquat)[:3,:3]

        if not n.all(n.isnan( rmat )):
            assert is_rotation_matrix(rmat)

    return rmat, rquat
//  [python]
def get_rotation_matrix_and_quaternion(rotation):
    rotation_orig = rotation
    rotation = np.array(rotation)
    if rotation.ndim==2:
        assert rotation.shape==(3,3)
        if not np.all(np.isnan( rotation )):
            assert is_rotation_matrix(rotation)

        rmat = rotation

        rnew = np.eye(4)
        rnew[:3,:3] = rmat
        rquat = quaternion_from_matrix(rnew)
        if not np.all(np.isnan( rquat )):
            R2 = quaternion_matrix(rquat)[:3,:3]
            assert np.allclose(rmat,R2)
    elif rotation.ndim==0:
        assert rotation.dtype == object
        rotation = (rotation_orig.x,
                    rotation_orig.y,
                    rotation_orig.z,
                    rotation_orig.w)
        return get_rotation_matrix_and_quaternion(rotation)
    else:
        assert rotation.ndim==1
        assert rotation.shape==(4,)
        rquat = rotation
        rmat = quaternion_matrix(rquat)[:3,:3]

        if not np.all(np.isnan( rmat )):
            assert is_rotation_matrix(rmat)

    return rmat, rquat
//  [python]
def __getitem__(self, index: Any) -> Array:
        if numeric.isintarray(index) and index.ndim == 1 and n.all(n.greater(n.diff(index), 0)):
            return MaskedBasis(self, index)
        elif numeric.isboolarray(index) and index.shape == (self.ndofs,):
            return MaskedBasis(self, n.where(index)[0])
        elif isinstance(index, slice):
            start, stop, step = index.indices(self.ndofs)
            if step == 1 and start == 0 and stop == self.ndofs:
                return self
            else:
                return MaskedBasis(self, n.arange(start, stop, step))
        else:
            return super().__getitem__(index)
//  [python]
def __getitem__(self, index: Any) -> Array:
        if numeric.isintarray(index) and index.ndim == 1 and numpy.all(numpy.greater(numpy.diff(index), 0)):
            return MaskedBasis(self, index)
        elif numeric.isboolarray(index) and index.shape == (self.ndofs,):
            return MaskedBasis(self, numpy.where(index)[0])
        elif isinstance(index, slice):
            start, stop, step = index.indices(self.ndofs)
            if step == 1 and start == 0 and stop == self.ndofs:
                return self
            else:
                return MaskedBasis(self, numpy.arange(start, stop, step))
        else:
            return super().__getitem__(index)
//  [python]
def __init__(self, parent: Basis, indices: numpy.ndarray) -> None:
        indices = types.frozenarray(indices)
        if indices.ndim != 1:
            raise ValueError('`indices` should have one dimension but got {}'.format(indices.ndim))
        if len(indices) and not n.all(n.greater(n.diff(indices), 0)):
            raise ValueError('`indices` should be strictly monotonic increasing')
        if len(indices) and (indices[0] < 0 or indices[-1] >= len(parent)):
            raise ValueError('`indices` out of range \x5b0,{}\x29'.format(len(parent)))
        self._parent = parent
        self._indices = indices
        self._renumber = evaluable.constant(numeric.invmap(indices, length=parent.ndofs, missing=len(indices)))
        super().__init__(len(indices), parent.nelems, parent.index, parent.coords)
//  [python]
def __init__(self, parent: Basis, indices: numpy.ndarray) -> None:
        indices = types.frozenarray(indices)
        if indices.ndim != 1:
            raise ValueError('`indices` should have one dimension but got {}'.format(indices.ndim))
        if len(indices) and not numpy.all(numpy.greater(numpy.diff(indices), 0)):
            raise ValueError('`indices` should be strictly monotonic increasing')
        if len(indices) and (indices[0] < 0 or indices[-1] >= len(parent)):
            raise ValueError('`indices` out of range \x5b0,{}\x29'.format(len(parent)))
        self._parent = parent
        self._indices = indices
        self._renumber = evaluable.constant(numeric.invmap(indices, length=parent.ndofs, missing=len(indices)))
        super().__init__(len(indices), parent.nelems, parent.index, parent.coords)
//  [python]
def psf2otf(psf, shape):
    """
    Convert point-spread function to optical transfer function.

    Compute the Fast Fourier Transform (FFT) of the point-spread
    function (PSF) array and creates the optical transfer function (OTF)
    array that is not influenced by the PSF off-centering.
    By default, the OTF array is the same size as the PSF array.

    To ensure that the OTF is not altered due to PSF off-centering, PSF2OTF
    post-pads the PSF array (down or to the right) with zeros to match
    dimensions specified in OUTSIZE, then circularly shifts the values of
    the PSF array up (or to the left) until the central pixel reaches (1,1)
    position.

    Parameters
    ----------
    psf : `numpy.ndarray`
        PSF array
    shape : int
        Output shape of the OTF array

    Returns
    -------
    otf : `numpy.ndarray`
        OTF array

    Notes
    -----
    Adapted from MATLAB psf2otf function

    """
    if n.all(psf == 0):
        return np.zeros_like(psf)

    inshape = psf.shape
    # Pad the PSF to outsize
    psf = zero_pad(psf, shape, position='corner')

    # Circularly shift OTF so that the 'center' of the PSF is
    # [0,0] element of the array
    for axis, axis_size in enumerate(inshape):
        psf = np.roll(psf, -int(axis_size / 2), axis=axis)

    # Compute the OTF
    otf = np.fft.fft2(psf)

    # Estimate the rough number of operations involved in the FFT
    # and discard the PSF imaginary part if within roundoff error
    # roundoff error  = machine epsilon = sys.float_info.epsilon
    # or np.finfo().eps
    n_ops = np.sum(psf.size * np.log2(psf.shape))
    otf = np.real_if_close(otf, tol=n_ops)

    return otf
//  [python]
def psf2otf(psf, shape):
    """
    Convert point-spread function to optical transfer function.

    Compute the Fast Fourier Transform (FFT) of the point-spread
    function (PSF) array and creates the optical transfer function (OTF)
    array that is not influenced by the PSF off-centering.
    By default, the OTF array is the same size as the PSF array.

    To ensure that the OTF is not altered due to PSF off-centering, PSF2OTF
    post-pads the PSF array (down or to the right) with zeros to match
    dimensions specified in OUTSIZE, then circularly shifts the values of
    the PSF array up (or to the left) until the central pixel reaches (1,1)
    position.

    Parameters
    ----------
    psf : `numpy.ndarray`
        PSF array
    shape : int
        Output shape of the OTF array

    Returns
    -------
    otf : `numpy.ndarray`
        OTF array

    Notes
    -----
    Adapted from MATLAB psf2otf function

    """
    if np.all(psf == 0):
        return np.zeros_like(psf)

    inshape = psf.shape
    # Pad the PSF to outsize
    psf = zero_pad(psf, shape, position='corner')

    # Circularly shift OTF so that the 'center' of the PSF is
    # [0,0] element of the array
    for axis, axis_size in enumerate(inshape):
        psf = np.roll(psf, -int(axis_size / 2), axis=axis)

    # Compute the OTF
    otf = np.fft.fft2(psf)

    # Estimate the rough number of operations involved in the FFT
    # and discard the PSF imaginary part if within roundoff error
    # roundoff error  = machine epsilon = sys.float_info.epsilon
    # or np.finfo().eps
    n_ops = np.sum(psf.size * np.log2(psf.shape))
    otf = np.real_if_close(otf, tol=n_ops)

    return otf
//  [python]
def __eq__(self, other: Any) -> bool:
        assert isinstance(other, Matrix)
        return np.all(self.toArray() == other.toArray()).tolist()
//  [python]
def flip_until_delaunay(self):
        """Flip edges until the mesh is fully Delaunay."""
        # If all coedge/edge ratios are positive, all cells are Delaunay.
        if n.all(self.ce_ratios > 0):
            return False

        # If all _interior_ coedge/edge ratios are positive, all cells are Delaunay.
        if self.is_boundary_edge is None:
            self.mark_boundary()
        if n.all(self.ce_ratios[~self.is_boundary_edge] > 0):
            return False

        if self._edges_cells is None:
            self._compute_edges_cells()

        num_flip_steps = 0
        ce_ratios_per_interior_edge = self.ce_ratios_per_interior_edge
        while n.any(ce_ratios_per_interior_edge < 0.0):
            num_flip_steps += 1

            is_flip_interior_edge = ce_ratios_per_interior_edge < 0.0

            interior_edges_cells = self._edges_cells[2]
            adj_cells = interior_edges_cells[is_flip_interior_edge].T

            # Check if there are cells for which more than one edge needs to be flipped.
            # For those, only flip one edge, namely that with the smaller (more
            # negative) ce_ratio.
            cell_gids, num_flips_per_cell = n.unique(
                adj_cells, return_counts=True
            )
            critical_cell_gids = cell_gids[num_flips_per_cell > 1]
            while n.any(num_flips_per_cell > 1):
                for cell_gid in critical_cell_gids:
                    edge_gids = self.cells["edges"][cell_gid]
                    num_adj_cells, edge_id = self._edge_gid_to_edge_list[
                        edge_gids
                    ].T
                    edge_ids = edge_id[num_adj_cells == 2]
                    k = n.argmin(ce_ratios_per_interior_edge[edge_ids])
                    is_flip_interior_edge[edge_ids] = False
                    is_flip_interior_edge[edge_ids[k]] = True

                adj_cells = interior_edges_cells[is_flip_interior_edge].T
                cell_gids, num_flips_per_cell = n.unique(
                    adj_cells, return_counts=True
                )
                critical_cell_gids = cell_gids[num_flips_per_cell > 1]

            self.flip_interior_edges(is_flip_interior_edge)
            ce_ratios_per_interior_edge = self.ce_ratios_per_interior_edge

        return num_flip_steps > 1
//  [python]
def flip_until_delaunay(self):
        """Flip edges until the mesh is fully Delaunay."""
        # If all coedge/edge ratios are positive, all cells are Delaunay.
        if numpy.all(self.ce_ratios > 0):
            return False

        # If all _interior_ coedge/edge ratios are positive, all cells are Delaunay.
        if self.is_boundary_edge is None:
            self.mark_boundary()
        if numpy.all(self.ce_ratios[~self.is_boundary_edge] > 0):
            return False

        if self._edges_cells is None:
            self._compute_edges_cells()

        num_flip_steps = 0
        ce_ratios_per_interior_edge = self.ce_ratios_per_interior_edge
        while numpy.any(ce_ratios_per_interior_edge < 0.0):
            num_flip_steps += 1

            is_flip_interior_edge = ce_ratios_per_interior_edge < 0.0

            interior_edges_cells = self._edges_cells[2]
            adj_cells = interior_edges_cells[is_flip_interior_edge].T

            # Check if there are cells for which more than one edge needs to be flipped.
            # For those, only flip one edge, namely that with the smaller (more
            # negative) ce_ratio.
            cell_gids, num_flips_per_cell = numpy.unique(
                adj_cells, return_counts=True
            )
            critical_cell_gids = cell_gids[num_flips_per_cell > 1]
            while numpy.any(num_flips_per_cell > 1):
                for cell_gid in critical_cell_gids:
                    edge_gids = self.cells["edges"][cell_gid]
                    num_adj_cells, edge_id = self._edge_gid_to_edge_list[
                        edge_gids
                    ].T
                    edge_ids = edge_id[num_adj_cells == 2]
                    k = numpy.argmin(ce_ratios_per_interior_edge[edge_ids])
                    is_flip_interior_edge[edge_ids] = False
                    is_flip_interior_edge[edge_ids[k]] = True

                adj_cells = interior_edges_cells[is_flip_interior_edge].T
                cell_gids, num_flips_per_cell = numpy.unique(
                    adj_cells, return_counts=True
                )
                critical_cell_gids = cell_gids[num_flips_per_cell > 1]

            self.flip_interior_edges(is_flip_interior_edge)
            ce_ratios_per_interior_edge = self.ce_ratios_per_interior_edge

        return num_flip_steps > 1
//  [python]
def test_meanOmega(bovy14_setup):
    # Load the streamdf object
    sdf_bovy14 = bovy14_setup
    # Test that meanOmega is close to constant and the mean Omega close to the progenitor
    assert n.all(
        n.fabs(sdf_bovy14.meanOmega(0.1) - sdf_bovy14._progenitor_Omega)
        < 10.0**-2.0
    ), "meanOmega near progenitor not close to mean Omega for Bovy14 stream"
    assert n.all(
        n.fabs(sdf_bovy14.meanOmega(0.5) - sdf_bovy14._progenitor_Omega)
        < 10.0**-2.0
    ), "meanOmega near progenitor not close to mean Omega for Bovy14 stream"
    return None
//  [python]
def test_meanOmega(bovy14_setup):
    # Load the streamdf object
    sdf_bovy14 = bovy14_setup
    # Test that meanOmega is close to constant and the mean Omega close to the progenitor
    assert numpy.all(
        numpy.fabs(sdf_bovy14.meanOmega(0.1) - sdf_bovy14._progenitor_Omega)
        < 10.0**-2.0
    ), "meanOmega near progenitor not close to mean Omega for Bovy14 stream"
    assert numpy.all(
        numpy.fabs(sdf_bovy14.meanOmega(0.5) - sdf_bovy14._progenitor_Omega)
        < 10.0**-2.0
    ), "meanOmega near progenitor not close to mean Omega for Bovy14 stream"
    return None
//  [python]
def set_affine(self, affine):
        """Set the affine transform (operating in physical space).

        Also sets `self.affine_inv` - the inverse of `affine`, or None if
        there is no inverse.

        Parameters
        ----------
        affine : array, shape (dim + 1, dim + 1)
            the matrix representing the affine transform operating in
            physical space. The domain and co-domain information
            remains unchanged. If None, then `self` represents the identity
            transformation.

        """

        if affine is None:
            self.affine = None
            self.affine_inv = None
            return

        try:
            affine = np.array(affine)
        except Exception:
            raise TypeError("Input must be type ndarray, or be convertible"
                            " to one.")

        if len(affine.shape) != _number_dim_affine_matrix:
            raise AffineInversionError('Affine transform must be 2D')

        if not affine.shape[0] == affine.shape[1]:
            raise AffineInversionError("Affine transform must be a square "
                                       "matrix")

        if not np.isfinite(affine).all():
            raise AffineInvalidValuesError("Affine transform contains invalid"
                                           " elements")

        # checking on proper augmentation
        # First n-1 columns in last row in matrix contain non-zeros
        if not (affine[-1, :-1] == 0.0).all():
            raise AffineInvalidValuesError("First {n_1} columns in last row"
                                           " in matrix contain non-zeros!"
                                           .format(n_1=affine.shape[0] - 1))

        # Last row, last column in matrix must be 1.0!
        if affine[-1, -1] != 1.0:
            raise AffineInvalidValuesError("Last row, last column in matrix"
                                           " is not 1.0!")

        # making a copy to insulate it from changes outside object
        self.affine = affine.copy()

        try:
            self.affine_inv = npl.inv(affine)
        except npl.LinAlgError:
            raise AffineInversionError('Affine cannot be inverted')
//  [python]
def set_affine(self, affine):
        """Set the affine transform (operating in physical space).

        Also sets `self.affine_inv` - the inverse of `affine`, or None if
        there is no inverse.

        Parameters
        ----------
        affine : array, shape (dim + 1, dim + 1)
            the matrix representing the affine transform operating in
            physical space. The domain and co-domain information
            remains unchanged. If None, then `self` represents the identity
            transformation.

        """

        if affine is None:
            self.affine = None
            self.affine_inv = None
            return

        try:
            affine = np.array(affine)
        except Exception:
            raise TypeError("Input must be type ndarray, or be convertible"
                            " to one.")

        if len(affine.shape) != _number_dim_affine_matrix:
            raise AffineInversionError('Affine transform must be 2D')

        if not affine.shape[0] == affine.shape[1]:
            raise AffineInversionError("Affine transform must be a square "
                                       "matrix")

        if not np.all(np.isfinite(affine)):
            raise AffineInvalidValuesError("Affine transform contains invalid"
                                           " elements")

        # checking on proper augmentation
        # First n-1 columns in last row in matrix contain non-zeros
        if not np.all(affine[-1, :-1] == 0.0):
            raise AffineInvalidValuesError("First {n_1} columns in last row"
                                           " in matrix contain non-zeros!"
                                           .format(n_1=affine.shape[0] - 1))

        # Last row, last column in matrix must be 1.0!
        if affine[-1, -1] != 1.0:
            raise AffineInvalidValuesError("Last row, last column in matrix"
                                           " is not 1.0!")

        # making a copy to insulate it from changes outside object
        self.affine = affine.copy()

        try:
            self.affine_inv = npl.inv(affine)
        except npl.LinAlgError:
            raise AffineInversionError('Affine cannot be inverted')
//  [python]
def get_area_under_perf_diagram(success_ratio_by_threshold, pod_by_threshold):
    """Computes area under performance diagram.

    T = number of binarization thresholds

    :param success_ratio_by_threshold: length-T numpy array of success ratios.
    :param pod_by_threshold: length-T numpy array of corresponding POD values.
    :return: area_under_curve: Area under performance diagram.
    """

    error_checking.assert_is_numpy_array(
        success_ratio_by_threshold, num_dimensions=1)
    error_checking.assert_is_geq_numpy_array(
        success_ratio_by_threshold, 0., allow_nan=True)
    error_checking.assert_is_leq_numpy_array(
        success_ratio_by_threshold, 1., allow_nan=True)

    num_thresholds = len(success_ratio_by_threshold)
    expected_dim = numpy.array([num_thresholds], dtype=int)

    error_checking.assert_is_numpy_array(
        pod_by_threshold, exact_dimensions=expected_dim)
    error_checking.assert_is_geq_numpy_array(
        pod_by_threshold, 0., allow_nan=True)
    error_checking.assert_is_leq_numpy_array(
        pod_by_threshold, 1., allow_nan=True)

    sort_indices = numpy.argsort(success_ratio_by_threshold)
    success_ratio_by_threshold = success_ratio_by_threshold[sort_indices]
    pod_by_threshold = pod_by_threshold[sort_indices]

    nan_flags = numpy.logical_or(
        numpy.isnan(success_ratio_by_threshold),
        numpy.isnan(pod_by_threshold)
    )
    if numpy.alltrue(nan_flags):
        return numpy.nan

    real_indices = numpy.where(numpy.invert(nan_flags))[0]

    return sklearn.metrics.auc(
        success_ratio_by_threshold[real_indices],
        pod_by_threshold[real_indices]
    )
//  [python]
def get_area_under_perf_diagram(success_ratio_by_threshold, pod_by_threshold):
    """Computes area under performance diagram.

    T = number of binarization thresholds

    :param success_ratio_by_threshold: length-T numpy array of success ratios.
    :param pod_by_threshold: length-T numpy array of corresponding POD values.
    :return: area_under_curve: Area under performance diagram.
    """

    error_checking.assert_is_numpy_array(
        success_ratio_by_threshold, num_dimensions=1)
    error_checking.assert_is_geq_numpy_array(
        success_ratio_by_threshold, 0., allow_nan=True)
    error_checking.assert_is_leq_numpy_array(
        success_ratio_by_threshold, 1., allow_nan=True)

    num_thresholds = len(success_ratio_by_threshold)
    expected_dim = numpy.array([num_thresholds], dtype=int)

    error_checking.assert_is_numpy_array(
        pod_by_threshold, exact_dimensions=expected_dim)
    error_checking.assert_is_geq_numpy_array(
        pod_by_threshold, 0., allow_nan=True)
    error_checking.assert_is_leq_numpy_array(
        pod_by_threshold, 1., allow_nan=True)

    sort_indices = numpy.argsort(success_ratio_by_threshold)
    success_ratio_by_threshold = success_ratio_by_threshold[sort_indices]
    pod_by_threshold = pod_by_threshold[sort_indices]

    nan_flags = numpy.logical_or(
        numpy.isnan(success_ratio_by_threshold),
        numpy.isnan(pod_by_threshold)
    )
    if numpy.all(nan_flags):
        return numpy.nan

    real_indices = numpy.where(numpy.invert(nan_flags))[0]

    return sklearn.metrics.auc(
        success_ratio_by_threshold[real_indices],
        pod_by_threshold[real_indices]
    )
//  [python]
def test_bovy14_trackaa(bovy14_setup):
    # Load the streamdf object
    sdf_bovy14 = bovy14_setup
    # Test that the explicitly-calculated frequencies along the track are close to those that the track is based on (Fardal test, #194)
    from galpy.orbit import Orbit

    aastream = sdf_bovy14._ObsTrackAA  # freqs and angles that the track is based on
    RvR = sdf_bovy14._ObsTrack  # the track in R,vR,...
    aastream_expl = numpy.reshape(
        numpy.array(
            [sdf_bovy14._aA.actionsFreqsAngles(Orbit(trvr))[3:] for trvr in RvR]
        ),
        aastream.shape,
    )
    # frequencies, compare to offset between track and progenitor (spread in freq ~ 1/6 that diff, so as long as we're smaller than that we're fine)
    assert n.all(
        n.fabs(
            (aastream[:, :3] - aastream_expl[:, :3])
            / (aastream[0, :3] - sdf_bovy14._progenitor_Omega)
        )
        < 0.05
    ), "Explicitly calculated frequencies along the track do not agree with the frequencies on which the track is based for bovy14 setup"
    # angles
    assert n.all(
        n.fabs((aastream[:, 3:] - aastream_expl[:, 3:]) / 2.0 / numpy.pi) < 0.001
    ), "Explicitly calculated angles along the track do not agree with the angles on which the track is based for bovy14 setup"
    return None
//  [python]
def test_bovy14_trackaa(bovy14_setup):
    # Load the streamdf object
    sdf_bovy14 = bovy14_setup
    # Test that the explicitly-calculated frequencies along the track are close to those that the track is based on (Fardal test, #194)
    from galpy.orbit import Orbit

    aastream = sdf_bovy14._ObsTrackAA  # freqs and angles that the track is based on
    RvR = sdf_bovy14._ObsTrack  # the track in R,vR,...
    aastream_expl = numpy.reshape(
        numpy.array(
            [sdf_bovy14._aA.actionsFreqsAngles(Orbit(trvr))[3:] for trvr in RvR]
        ),
        aastream.shape,
    )
    # frequencies, compare to offset between track and progenitor (spread in freq ~ 1/6 that diff, so as long as we're smaller than that we're fine)
    assert numpy.all(
        numpy.fabs(
            (aastream[:, :3] - aastream_expl[:, :3])
            / (aastream[0, :3] - sdf_bovy14._progenitor_Omega)
        )
        < 0.05
    ), "Explicitly calculated frequencies along the track do not agree with the frequencies on which the track is based for bovy14 setup"
    # angles
    assert numpy.all(
        numpy.fabs((aastream[:, 3:] - aastream_expl[:, 3:]) / 2.0 / numpy.pi) < 0.001
    ), "Explicitly calculated angles along the track do not agree with the angles on which the track is based for bovy14 setup"
    return None
//  [python]
def test_orbitfit():
    from galpy.orbit import Orbit

    lp = potential.LogarithmicHaloPotential(normalize=1.0, q=0.9)
    o = Orbit([0.8, 0.3, 1.3, 0.4, 0.2, 2.0])
    ts = numpy.linspace(0.0, 1.0, 1001)
    o.integrate(ts, lp)
    # Create orbit points from this integrated orbit, each 100th point
    vxvv = o.getOrbit()[::100, :]
    # now fit
    of = Orbit.from_fit(o.vxvv[0], vxvv, pot=lp, tintJ=1.5)
    assert n.all(
        comp_orbfit(of, vxvv, numpy.linspace(0.0, 2.0, 1001), lp) < 10.0**-7.0
    ), "Orbit fit in configuration space does not work"
    return None