jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def __init__(self, filepath: str):
        super().__init__(filepath)

        # Read and store important quantities.
        self.structure = self.read_structure()
        qfrac_coords = self.read_value("qpoints_dielectric_function")
        self.kpoints = KpointList(self.structure.reciprocal_lattice, qfrac_coords)
        self.wpoints = self.read_value("frequencies_dielectric_function", cmode="c")
        self.ng = self.read_dimvalue("number_of_coefficients_dielectric_function")

        # Find number of real/imaginary frequencies.
        self.nw = len(self.wpoints)
        self.nrew = self.nw
        self.nimw = 0
        for i, w in enumerate(self.wpoints):
            if np.iscomplex(w):
                self.nrew = i
                break

        self.nimw = self.nw - self.nrew
        if self.nimw and not np.all(np.iscomplex(self.wpoints[self.nrew+1:])):
            raise ValueError("wpoints should contained real points packed in the first positions\n"
                             "followed by imaginary points but got: %s" % str(self.wpoints))

        # Define self.netcdf_name from the data available on file.
        nfound = 0
        netcdf_names = ["polarizability", "dielectric_function", "inverse_dielectric_function"]
        for tryname in netcdf_names:
            if tryname in self.rootgrp.variables:
                self.netcdf_name = tryname
                nfound += 1

        if nfound == 0:
            raise RuntimeError("Cannot find `%s` in netcdf file" % str(netcdf_names))
        if nfound > 1:
            raise RuntimeError("Find multiple netcdf arrays (%s) in netcdf file!" % str(netcdf_names))
//  [python]
def __init__(self, wpoints, gsphere, wggmat, inord="C"):
        """"
        Args:
            gsphere: |GSphere| with G-vectors and k-point object.
            wpoints: Complex frequency points in Hartree.
            wggmat: [nw, ng, ng] complex array.
            inord: storage order of ``wggmat``. If inord == "F", ``wggmat`` is in
                in Fortran column-major order. Default: "C" i.e. C row-major order.
        """
        self.wpoints = np.array(wpoints, dtype=complex)
        self.gsphere = gsphere
        self.wggmat = np.reshape(wggmat, (self.nw, self.ng, self.ng))

        if inord.lower() == "f":
            # Fortran to C.
            for iw, _ in enumerate(wpoints):
                self.wggmat[iw] = self.wggmat[iw].T.copy()

        for i in (1, 2):
            assert len(gsphere) == wggmat.shape[-i]
        assert len(self.wpoints) == len(self.wggmat)

        # Find number of real/imaginary frequencies.
        self.nrew = self.nw
        self.nimw = 0
        for i, w in enumerate(self.wpoints):
            if np.iscomplex(w):
                self.nrew = i
                break

        self.nimw = self.nw - self.nrew
        if self.nimw and not np.alltrue(np.iscomplex(self.wpoints[self.nrew+1:])):
            raise ValueError("wpoints should contained real points packed in the first positions\n"
                "followed by imaginary points but got: %s" % str(self.wpoints))
//  [python]
def __init__(self, wpoints, gsphere, wggmat, inord="C"):
        """"
        Args:
            gsphere: |GSphere| with G-vectors and k-point object.
            wpoints: Complex frequency points in Hartree.
            wggmat: [nw, ng, ng] complex array.
            inord: storage order of ``wggmat``. If inord == "F", ``wggmat`` is in
                in Fortran column-major order. Default: "C" i.e. C row-major order.
        """
        self.wpoints = np.array(wpoints, dtype=complex)
        self.gsphere = gsphere
        self.wggmat = np.reshape(wggmat, (self.nw, self.ng, self.ng))

        if inord.lower() == "f":
            # Fortran to C.
            for iw, _ in enumerate(wpoints):
                self.wggmat[iw] = self.wggmat[iw].T.copy()

        for i in (1, 2):
            assert len(gsphere) == wggmat.shape[-i]
        assert len(self.wpoints) == len(self.wggmat)

        # Find number of real/imaginary frequencies.
        self.nrew = self.nw
        self.nimw = 0
        for i, w in enumerate(self.wpoints):
            if np.iscomplex(w):
                self.nrew = i
                break

        self.nimw = self.nw - self.nrew
        if self.nimw and not np.all(np.iscomplex(self.wpoints[self.nrew+1:])):
            raise ValueError("wpoints should contained real points packed in the first positions\n"
                "followed by imaginary points but got: %s" % str(self.wpoints))
//  [python]
def moment2_k(self):
        # instantaneous bandwitdh / ?
        proj_s = self.power()
        if n.all(proj_s == 0):
            return proj_s
        else:
            return np.sum(self.wig * (self.k[:, np.newaxis] - self.moment1_k()[np.newaxis, :]) ** 2, axis=0) / proj_s
//  [python]
def moment2_k(self):
        # instantaneous bandwitdh / ?
        proj_s = self.power()
        if np.all(proj_s == 0):
            return proj_s
        else:
            return np.sum(self.wig * (self.k[:, np.newaxis] - self.moment1_k()[np.newaxis, :]) ** 2, axis=0) / proj_s
//  [python]
def moment2_s(self):
        # group duration / ?
        proj_k = self.proj_k()
        if n.all(proj_k == 0):
            return proj_k
        else:
            return np.sum(self.wig * (self.s[np.newaxis, :] - self.moment1_s()[:, np.newaxis]) ** 2, axis=1) / proj_k
//  [python]
def moment2_s(self):
        # group duration / ?
        proj_k = self.proj_k()
        if np.all(proj_k == 0):
            return proj_k
        else:
            return np.sum(self.wig * (self.s[np.newaxis, :] - self.moment1_s()[:, np.newaxis]) ** 2, axis=1) / proj_k
//  [python]
def inst_freq(self):
        p = self.power()
        if p.all() == 0:
            return p
        else:
            p[p == 0] = np.nan
            return np.sum(self.wig * self.phen[:, np.newaxis], axis=0) / p
//  [python]
def inst_freq(self):
        p = self.power()
        if np.all(p == 0):
            return p
        else:
            p[p == 0] = np.nan
            return np.sum(self.wig * self.phen[:, np.newaxis], axis=0) / p
//  [python]
def isPositiveSemiDefinite(A, tol=1e-6):
  E,V = scipy.linalg.eigh(A)
  return n.all(E > -tol)
//  [python]
def isPositiveSemiDefinite(A, tol=1e-6):
  E,V = scipy.linalg.eigh(A)
  return numpy.all(E > -tol)
//  [python]
def group_delay(self):
        s = self.spectrum()
        if n.all(s == 0):
            return s
        else:
            s[s == 0] = np.nan
            return np.sum(self.wig * self.s[np.newaxis, :], axis=1) / s
//  [python]
def group_delay(self):
        s = self.spectrum()
        if np.all(s == 0):
            return s
        else:
            s[s == 0] = np.nan
            return np.sum(self.wig * self.s[np.newaxis, :], axis=1) / s
//  [python]
def inst_bandwidth(self):
        # check, gives strange tilt. physics?
        p = self.power()
        if n.all(p == 0):
            return p
        else:
            return np.sum(self.wig * (self.phen[:, np.newaxis] - self.inst_freq()[np.newaxis, :]) ** 2, axis=0) / p
//  [python]
def inst_bandwidth(self):
        # check, gives strange tilt. physics?
        p = self.power()
        if np.all(p == 0):
            return p
        else:
            return np.sum(self.wig * (self.phen[:, np.newaxis] - self.inst_freq()[np.newaxis, :]) ** 2, axis=0) / p
//  [python]
def run_perturbations(self, meas_spec, unpert_spec_range, out_hdf, fm_err_grp):
        # Helper routine, using closures to make getting the radiance a bit more reusable
        def get_perturbed_radiance_diff():
            skip_jacobian = True
            pert_spec = self.lua_config.fm.config.forward_model.radiance_all(skip_jacobian)
            pert_range = pert_spec.spectral_range.convert(meas_spec.spectral_range.units)

            return pert_range.data - unpert_spec_range.data

        # Loop over forward model perturbations
        log_info("Running forward model perturbation cases.\n")
        for filter_name, err_func in list(self.filtered_perturb_funcs.items()): 
            err_generator = err_func(self.ls, self.lua_config)

            if type(err_generator) != GeneratorType:
                raise ValueError("Error function: %s did not return a generator function, instead returned: %s" % (err_func, err_generator))

            # A single perturb function could iterate over several cases
            for err_desc in err_generator:
                log_info( "Calculating forward model perturb type: %s\n" % err_desc.name )

                # Run forward model and pull out perturbed radiances
                pert_rad_diff = get_perturbed_radiance_diff()

                if n.all(pert_rad_diff == 0.0):
                   log_info("WARNING: Case: %s did not end up returning a radiance with any perturbation\n" % (err_desc.name))

                # Write calculated perturb information and flush hdf buffer
                err_desc.write_deriv_dataset(fm_err_grp, pert_rad_diff, unpert_spec_range.units.name)
                out_hdf.flush()

                # Perform a 1 iteration retrieval on perturbed FM state using apriori statevector
                if self.num_perturbed_iterations > 0:
                    self.perturbed_retrieval(meas_spec, err_desc, fm_err_grp, unpert_spec_range.units.name)
                    out_hdf.flush()

                # Call generator to reset state, it will return True if there are more
                # tests, otherwise a StopIteration exception should have been thrown
                try:
                    more_tests = next(err_generator)
                    if not more_tests:
                        raise RuntimeError("Error function generator for: %s (%s) did not end as expected" % (err_desc.name, err_func))
                    else:
                        log_info("%s has more tests\n" % err_desc.name)
                except StopIteration:
                    log_info("%s has no more tests\n" % err_desc.name)
                    pass

                # If enabled, rerun RT after reset state to make sure everything has been reset correctly
                if self.check_reset:
                    log_info( "Checking that %s properly reset the FM\n" % err_desc.name )

                    pert_rad_diff = get_perturbed_radiance_diff()

                    if not n.all(pert_rad_diff == 0.0):
                        raise ValueError("Forward model was not reset correctly for case: %s" % (err_desc.name))
//  [python]
def run_perturbations(self, meas_spec, unpert_spec_range, out_hdf, fm_err_grp):
        # Helper routine, using closures to make getting the radiance a bit more reusable
        def get_perturbed_radiance_diff():
            skip_jacobian = True
            pert_spec = self.lua_config.fm.config.forward_model.radiance_all(skip_jacobian)
            pert_range = pert_spec.spectral_range.convert(meas_spec.spectral_range.units)

            return pert_range.data - unpert_spec_range.data

        # Loop over forward model perturbations
        log_info("Running forward model perturbation cases.\n")
        for filter_name, err_func in list(self.filtered_perturb_funcs.items()): 
            err_generator = err_func(self.ls, self.lua_config)

            if type(err_generator) != GeneratorType:
                raise ValueError("Error function: %s did not return a generator function, instead returned: %s" % (err_func, err_generator))

            # A single perturb function could iterate over several cases
            for err_desc in err_generator:
                log_info( "Calculating forward model perturb type: %s\n" % err_desc.name )

                # Run forward model and pull out perturbed radiances
                pert_rad_diff = get_perturbed_radiance_diff()

                if numpy.all(pert_rad_diff == 0.0):
                   log_info("WARNING: Case: %s did not end up returning a radiance with any perturbation\n" % (err_desc.name))

                # Write calculated perturb information and flush hdf buffer
                err_desc.write_deriv_dataset(fm_err_grp, pert_rad_diff, unpert_spec_range.units.name)
                out_hdf.flush()

                # Perform a 1 iteration retrieval on perturbed FM state using apriori statevector
                if self.num_perturbed_iterations > 0:
                    self.perturbed_retrieval(meas_spec, err_desc, fm_err_grp, unpert_spec_range.units.name)
                    out_hdf.flush()

                # Call generator to reset state, it will return True if there are more
                # tests, otherwise a StopIteration exception should have been thrown
                try:
                    more_tests = next(err_generator)
                    if not more_tests:
                        raise RuntimeError("Error function generator for: %s (%s) did not end as expected" % (err_desc.name, err_func))
                    else:
                        log_info("%s has more tests\n" % err_desc.name)
                except StopIteration:
                    log_info("%s has no more tests\n" % err_desc.name)
                    pass

                # If enabled, rerun RT after reset state to make sure everything has been reset correctly
                if self.check_reset:
                    log_info( "Checking that %s properly reset the FM\n" % err_desc.name )

                    pert_rad_diff = get_perturbed_radiance_diff()

                    if not numpy.all(pert_rad_diff == 0.0):
                        raise ValueError("Forward model was not reset correctly for case: %s" % (err_desc.name))
//  [python]
def update(self, delta_x, delta_grad):
        if n.all(delta_x == 0.0):
            return
        if n.all(delta_grad == 0.0):
            warn('delta_grad == 0.0', UserWarning)
            return
        if self.first_iteration:
            if self.init_scale == "auto":
                scale = self._auto_scale(delta_x, delta_grad)
            else:
                scale = float(self.init_scale)
            if self.approx_type == 'hess':
                self.B *= scale
            else:
                self.H *= scale
            self.first_iteration = False
        self._update_implementation(delta_x, delta_grad)
//  [python]
def update(self, delta_x, delta_grad):
        if np.all(delta_x == 0.0):
            return
        if np.all(delta_grad == 0.0):
            warn('delta_grad == 0.0', UserWarning)
            return
        if self.first_iteration:
            if self.init_scale == "auto":
                scale = self._auto_scale(delta_x, delta_grad)
            else:
                scale = float(self.init_scale)
            if self.approx_type == 'hess':
                self.B *= scale
            else:
                self.H *= scale
            self.first_iteration = False
        self._update_implementation(delta_x, delta_grad)
//  [python]
def contains(self, coords):
        idx = self.coord_to_idx(coords)
        return np.alltrue(np.stack([t != INVALID_INDEX.int for t in idx]), axis=0)
//  [python]
def contains(self, coords):
        idx = self.coord_to_idx(coords)
        return np.all(np.stack([t != INVALID_INDEX.int for t in idx]), axis=0)
//  [python]
def from_PreparedConstraint(cls, constraint):
        """Create an instance from `PreparedConstrained` object."""
        lb, ub = constraint.bounds
        cfun = constraint.fun
        keep_feasible = constraint.keep_feasible

        if n.all(lb == -np.inf) and n.all(ub == np.inf):
            return cls.empty(cfun.n)

        if n.all(lb == -np.inf) and n.all(ub == np.inf):
            return cls.empty(cfun.n)
        elif n.all(lb == ub):
            return cls._equal_to_canonical(cfun, lb)
        elif n.all(lb == -np.inf):
            return cls._less_to_canonical(cfun, ub, keep_feasible)
        elif n.all(ub == np.inf):
            return cls._greater_to_canonical(cfun, lb, keep_feasible)
        else:
            return cls._interval_to_canonical(cfun, lb, ub, keep_feasible)
//  [python]
def from_PreparedConstraint(cls, constraint):
        """Create an instance from `PreparedConstrained` object."""
        lb, ub = constraint.bounds
        cfun = constraint.fun
        keep_feasible = constraint.keep_feasible

        if np.all(lb == -np.inf) and np.all(ub == np.inf):
            return cls.empty(cfun.n)

        if np.all(lb == -np.inf) and np.all(ub == np.inf):
            return cls.empty(cfun.n)
        elif np.all(lb == ub):
            return cls._equal_to_canonical(cfun, lb)
        elif np.all(lb == -np.inf):
            return cls._less_to_canonical(cfun, ub, keep_feasible)
        elif np.all(ub == np.inf):
            return cls._greater_to_canonical(cfun, lb, keep_feasible)
        else:
            return cls._interval_to_canonical(cfun, lb, ub, keep_feasible)
//  [python]
def _uniformize(data, col_name):
    # type: (NumpyFrame, str) -> np.ndarray
    """
    Convert to quantiles.
    :param data: NumpyFrame
    :param col_name: string containing a column name
    :returns: quantile values of individual points in the column
    """
    if col_name not in data.columns or data.isfactor(col_name):
        res = data[col_name]
        diff = (np.nanmax(res) - np.nanmin(res))
        if diff <= 0 or np.isnan(diff):
            return res
        res = (res - np.nanmin(res)) / diff
        return res

    col = data[col_name]
    xs = np.linspace(0, 1, 100)
    quantiles = np.nanquantile(col, xs)
    res = np.interp(col, quantiles, xs)
    if not np.all(np.isnan(res)):
        res = (res - np.nanmin(res)) / (np.nanmax(res) - np.nanmin(res))
    return res
//  [python]
def test_VonMises_range(self):
        # Make sure generated random variables are in [-pi, pi].
        # Regression test for ticket #986.
        for mu in np.linspace(-7., 7., 5):
            r = random.vonmises(mu, 1, 50)
            assert_(n.all(r > -np.pi) and n.all(r <= np.pi))
//  [python]
def test_VonMises_range(self):
        # Make sure generated random variables are in [-pi, pi].
        # Regression test for ticket #986.
        for mu in np.linspace(-7., 7., 5):
            r = random.vonmises(mu, 1, 50)
            assert_(np.all(r > -np.pi) and np.all(r <= np.pi))
//  [python]
def _check_data(data1, data2):
    return np.all(np.isclose(data1, data2, rtol=0.001, equal_nan=True))
//  [python]
def test_pool_cat_features():
    pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    assert np.all(pool.get_cat_feature_indices() == CAT_FEATURES)
//  [python]
def _adjust_scheme_to_bounds(x0, h, num_steps, scheme, lb, ub):
    if scheme == '1-sided':
        use_one_sided = np.ones_like(h, dtype=bool)
    elif scheme == '2-sided':
        h = np.abs(h)
        use_one_sided = np.zeros_like(h, dtype=bool)
    else:
        raise ValueError("`scheme` must be '1-sided' or '2-sided'.")

    if np.all((lb == -np.inf) & (ub == np.inf)):
        return h, use_one_sided

    h_total = h * num_steps
    h_adjusted = h.copy()

    lower_dist = x0 - lb
    upper_dist = ub - x0

    if scheme == '1-sided':
        x = x0 + h_total
        violated = (x < lb) | (x > ub)
        fitting = np.abs(h_total) <= np.maximum(lower_dist, upper_dist)
        h_adjusted[violated & fitting] *= -1

        forward = (upper_dist >= lower_dist) & ~fitting
        h_adjusted[forward] = upper_dist[forward] / num_steps
        backward = (upper_dist < lower_dist) & ~fitting
        h_adjusted[backward] = -lower_dist[backward] / num_steps
    elif scheme == '2-sided':
        central = (lower_dist >= h_total) & (upper_dist >= h_total)

        forward = (upper_dist >= lower_dist) & ~central
        h_adjusted[forward] = np.minimum(
            h[forward], 0.5 * upper_dist[forward] / num_steps)
        use_one_sided[forward] = True

        backward = (upper_dist < lower_dist) & ~central
        h_adjusted[backward] = -np.minimum(
            h[backward], 0.5 * lower_dist[backward] / num_steps)
        use_one_sided[backward] = True

        min_dist = np.minimum(upper_dist, lower_dist) / num_steps
        adjusted_central = (~central & (np.abs(h_adjusted) <= min_dist))
        h_adjusted[adjusted_central] = min_dist[adjusted_central]
        use_one_sided[adjusted_central] = False

    return h_adjusted, use_one_sided
//  [python]
def test_hypergeometric_range(self):
        # Test for ticket #921
        assert_(n.all(random.hypergeometric(3, 18, 11, size=10) < 4))
        assert_(n.all(random.hypergeometric(18, 3, 11, size=10) > 0))

        # Test for ticket #5623
        args = [
            (2**20 - 2, 2**20 - 2, 2**20 - 2),  # Check for 32-bit systems
        ]
        is_64bits = sys.maxsize > 2**32
        if is_64bits and sys.platform != 'win32':
            # Check for 64-bit systems
            args.append((2**40 - 2, 2**40 - 2, 2**40 - 2))
        for arg in args:
            assert_(random.hypergeometric(*arg) > 0)
//  [python]
def test_hypergeometric_range(self):
        # Test for ticket #921
        assert_(np.all(random.hypergeometric(3, 18, 11, size=10) < 4))
        assert_(np.all(random.hypergeometric(18, 3, 11, size=10) > 0))

        # Test for ticket #5623
        args = [
            (2**20 - 2, 2**20 - 2, 2**20 - 2),  # Check for 32-bit systems
        ]
        is_64bits = sys.maxsize > 2**32
        if is_64bits and sys.platform != 'win32':
            # Check for 64-bit systems
            args.append((2**40 - 2, 2**40 - 2, 2**40 - 2))
        for arg in args:
            assert_(random.hypergeometric(*arg) > 0)
//  [python]
def infer_shape(self, fgraph, node, i_shapes):
        from pytensor.tensor.math import ceil, maximum

        # Note start, stop and step can be float numbers.
        start, stop, step = node.inputs

        def is_constant_value(var, value):
            try:
                v = get_underlying_scalar_constant_value(var)
                return n.all(v == value)
            except NotScalarConstantError:
                pass
            return False

        def upcast(var):
            if (
                var.dtype in integer_dtypes
                and
                # We do not want to cast uint64 to int64 as this can
                # loose information. If we upcast uint64 with int64,
                # this give float64. This is safer then checking for
                # uint64 in case we support [u]int128 or other in the
                # future.
                ps.upcast(var.dtype, "int64") == "int64"
            ):
                return cast(var, "int64")
            return var

        if is_constant_value(step, 1):
            if is_constant_value(start, 0):
                return [(cast(stop, "int64"),)]
            else:
                stop = upcast(stop)
                start = upcast(start)
                return [(maximum(cast(stop - start, "int64"), 0),)]
        else:
            stop = upcast(stop)
            start = upcast(start)
            return [
                (
                    maximum(
                        cast(ceil(cast((stop - start), "float64") / step), "int64"), 0
                    ),
                )
            ]
//  [python]
def infer_shape(self, fgraph, node, i_shapes):
        from pytensor.tensor.math import ceil, maximum

        # Note start, stop and step can be float numbers.
        start, stop, step = node.inputs

        def is_constant_value(var, value):
            try:
                v = get_underlying_scalar_constant_value(var)
                return np.all(v == value)
            except NotScalarConstantError:
                pass
            return False

        def upcast(var):
            if (
                var.dtype in integer_dtypes
                and
                # We do not want to cast uint64 to int64 as this can
                # loose information. If we upcast uint64 with int64,
                # this give float64. This is safer then checking for
                # uint64 in case we support [u]int128 or other in the
                # future.
                ps.upcast(var.dtype, "int64") == "int64"
            ):
                return cast(var, "int64")
            return var

        if is_constant_value(step, 1):
            if is_constant_value(start, 0):
                return [(cast(stop, "int64"),)]
            else:
                stop = upcast(stop)
                start = upcast(start)
                return [(maximum(cast(stop - start, "int64"), 0),)]
        else:
            stop = upcast(stop)
            start = upcast(start)
            return [
                (
                    maximum(
                        cast(ceil(cast((stop - start), "float64") / step), "int64"), 0
                    ),
                )
            ]
//  [python]
def test_raw_predict_equals_to_model_predict(task_type):
    train_pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    test_pool = Pool(TEST_FILE, column_description=CD_FILE)
    model = CatBoostClassifier(iterations=10, task_type=task_type, devices='0')
    model.fit(train_pool, eval_set=test_pool)
    assert (model.is_fitted())
    pred = model.predict(test_pool, prediction_type='RawFormulaVal')
    assert n.all(n.isclose(model.get_test_eval(), pred, rtol=1.e-6))
//  [python]
def test_raw_predict_equals_to_model_predict(task_type):
    train_pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    test_pool = Pool(TEST_FILE, column_description=CD_FILE)
    model = CatBoostClassifier(iterations=10, task_type=task_type, devices='0')
    model.fit(train_pool, eval_set=test_pool)
    assert (model.is_fitted())
    pred = model.predict(test_pool, prediction_type='RawFormulaVal')
    assert np.all(np.isclose(model.get_test_eval(), pred, rtol=1.e-6))
//  [python]
def initial_constraints_as_canonical(n, prepared_constraints, sparse_jacobian):
    c_eq = []
    c_ineq = []
    J_eq = []
    J_ineq = []

    for c in prepared_constraints:
        f = c.fun.f
        J = c.fun.J
        lb, ub = c.bounds
        if n.all(lb == ub):
            c_eq.append(f - lb)
            J_eq.append(J)
        elif n.all(lb == -np.inf):
            finite_ub = ub < np.inf
            c_ineq.append(f[finite_ub] - ub[finite_ub])
            J_ineq.append(J[finite_ub])
        elif n.all(ub == np.inf):
            finite_lb = lb > -np.inf
            c_ineq.append(lb[finite_lb] - f[finite_lb])
            J_ineq.append(-J[finite_lb])
        else:
            lb_inf = lb == -np.inf
            ub_inf = ub == np.inf
            equal = lb == ub
            less = lb_inf & ~ub_inf
            greater = ub_inf & ~lb_inf
            interval = ~equal & ~lb_inf & ~ub_inf

            c_eq.append(f[equal] - lb[equal])
            c_ineq.append(f[less] - ub[less])
            c_ineq.append(lb[greater] - f[greater])
            c_ineq.append(f[interval] - ub[interval])
            c_ineq.append(lb[interval] - f[interval])

            J_eq.append(J[equal])
            J_ineq.append(J[less])
            J_ineq.append(-J[greater])
            J_ineq.append(J[interval])
            J_ineq.append(-J[interval])

    c_eq = np.hstack(c_eq) if c_eq else np.empty(0)
    c_ineq = np.hstack(c_ineq) if c_ineq else np.empty(0)

    if sparse_jacobian:
        vstack = sps.vstack
        empty = sps.csr_matrix((0, n))
    else:
        vstack = np.vstack
        empty = np.empty((0, n))

    J_eq = vstack(J_eq) if J_eq else empty
    J_ineq = vstack(J_ineq) if J_ineq else empty

    return c_eq, c_ineq, J_eq, J_ineq
//  [python]
def initial_constraints_as_canonical(n, prepared_constraints, sparse_jacobian):
    c_eq = []
    c_ineq = []
    J_eq = []
    J_ineq = []

    for c in prepared_constraints:
        f = c.fun.f
        J = c.fun.J
        lb, ub = c.bounds
        if np.all(lb == ub):
            c_eq.append(f - lb)
            J_eq.append(J)
        elif np.all(lb == -np.inf):
            finite_ub = ub < np.inf
            c_ineq.append(f[finite_ub] - ub[finite_ub])
            J_ineq.append(J[finite_ub])
        elif np.all(ub == np.inf):
            finite_lb = lb > -np.inf
            c_ineq.append(lb[finite_lb] - f[finite_lb])
            J_ineq.append(-J[finite_lb])
        else:
            lb_inf = lb == -np.inf
            ub_inf = ub == np.inf
            equal = lb == ub
            less = lb_inf & ~ub_inf
            greater = ub_inf & ~lb_inf
            interval = ~equal & ~lb_inf & ~ub_inf

            c_eq.append(f[equal] - lb[equal])
            c_ineq.append(f[less] - ub[less])
            c_ineq.append(lb[greater] - f[greater])
            c_ineq.append(f[interval] - ub[interval])
            c_ineq.append(lb[interval] - f[interval])

            J_eq.append(J[equal])
            J_ineq.append(J[less])
            J_ineq.append(-J[greater])
            J_ineq.append(J[interval])
            J_ineq.append(-J[interval])

    c_eq = np.hstack(c_eq) if c_eq else np.empty(0)
    c_ineq = np.hstack(c_ineq) if c_ineq else np.empty(0)

    if sparse_jacobian:
        vstack = sps.vstack
        empty = sps.csr_matrix((0, n))
    else:
        vstack = np.vstack
        empty = np.empty((0, n))

    J_eq = vstack(J_eq) if J_eq else empty
    J_ineq = vstack(J_ineq) if J_ineq else empty

    return c_eq, c_ineq, J_eq, J_ineq
//  [python]
def test_pool_cat_features_as_strings():
    df = DataFrame(data=[[1, 2], [3, 4]], columns=['col1', 'col2'])
    pool = Pool(df, cat_features=['col2'])
    assert n.all(pool.get_cat_feature_indices() == [1])

    data = [[1, 2, 3], [4, 5, 6]]
    pool = Pool(data, feature_names=['col1', 'col2', 'col3'], cat_features=['col2', 'col3'])
    assert n.all(pool.get_cat_feature_indices() == [1, 2])

    data = [[1, 2, 3], [4, 5, 6]]
    with pytest.raises(CatBoostError):
        Pool(data, cat_features=['col2', 'col3'])
//  [python]
def test_pool_cat_features_as_strings():
    df = DataFrame(data=[[1, 2], [3, 4]], columns=['col1', 'col2'])
    pool = Pool(df, cat_features=['col2'])
    assert np.all(pool.get_cat_feature_indices() == [1])

    data = [[1, 2, 3], [4, 5, 6]]
    pool = Pool(data, feature_names=['col1', 'col2', 'col3'], cat_features=['col2', 'col3'])
    assert np.all(pool.get_cat_feature_indices() == [1, 2])

    data = [[1, 2, 3], [4, 5, 6]]
    with pytest.raises(CatBoostError):
        Pool(data, cat_features=['col2', 'col3'])
//  [python]
def test_no_cat_in_predict(task_type):
    train_pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    model = CatBoostClassifier(iterations=2, learning_rate=0.03, task_type=task_type, devices='0')
    model.fit(train_pool)

    test_features_data, _ = load_simple_dataset_as_lists(is_test=True)
    pred1 = model.predict(test_features_data)
    pred2 = model.predict(Pool(test_features_data, cat_features=CAT_FEATURES))
    assert n.all(pred1 == pred2)
//  [python]
def test_no_cat_in_predict(task_type):
    train_pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    model = CatBoostClassifier(iterations=2, learning_rate=0.03, task_type=task_type, devices='0')
    model.fit(train_pool)

    test_features_data, _ = load_simple_dataset_as_lists(is_test=True)
    pred1 = model.predict(test_features_data)
    pred2 = model.predict(Pool(test_features_data, cat_features=CAT_FEATURES))
    assert np.all(pred1 == pred2)
//  [python]
def test_ones_weight_equal_to_nonspecified_weight(task_type):
    train_pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    test_pool = Pool(TEST_FILE, column_description=CD_FILE)
    model = CatBoostClassifier(iterations=2, learning_rate=0.03, task_type=task_type, devices='0')

    predictions = []

    for set_weights in [False, True]:
        if set_weights:
            weight = np.ones(train_pool.num_row())
            train_pool.set_weight(weight)
        model.fit(train_pool)
        predictions.append(model.predict(test_pool))

    assert n.all(predictions[0] == predictions[1])
//  [python]
def test_ones_weight_equal_to_nonspecified_weight(task_type):
    train_pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    test_pool = Pool(TEST_FILE, column_description=CD_FILE)
    model = CatBoostClassifier(iterations=2, learning_rate=0.03, task_type=task_type, devices='0')

    predictions = []

    for set_weights in [False, True]:
        if set_weights:
            weight = np.ones(train_pool.num_row())
            train_pool.set_weight(weight)
        model.fit(train_pool)
        predictions.append(model.predict(test_pool))

    assert np.all(predictions[0] == predictions[1])
//  [python]
def test_fit_with_texts(task_type, problem_type):
    if (task_type == 'GPU') and (problem_type == 'multiregression'):
        pytest.xfail('Catboost does not support MultiRMSE on GPU yet for categorical features')

    params = {
        'dictionaries': [
            {'dictionary_id': 'UniGram', 'token_level_type': 'Letter', 'occurrence_lower_bound': '1'},
            {'dictionary_id': 'BiGram', 'token_level_type': 'Letter', 'occurrence_lower_bound': '1', 'gram_order': '2'},
            {'dictionary_id': 'Word', 'occurrence_lower_bound': '1'},
        ],
        'feature_calcers': ['BoW:top_tokens_count=10'] if 'regression' in problem_type else ['NaiveBayes', 'BoW:top_tokens_count=10'],
        'iterations': 100,
        'loss_function': {
            'binclass': 'Logloss',
            'multiclass': 'MultiClass',
            'regression': 'RMSE',
            'multiregression': 'MultiRMSE'
        }[problem_type],
        'task_type': task_type,
        'devices': '0'
    }

    learn = ROTTEN_TOMATOES_TRAIN_FILE
    test = ROTTEN_TOMATOES_TEST_FILE
    cd = ROTTEN_TOMATOES_CD_FILE if problem_type == 'multiclass' else ROTTEN_TOMATOES_CD_BINCLASS_FILE

    preds1 = fit_from_df(params, learn, test, cd, problem_type == 'multiregression')
    if problem_type != 'multiregression':
        _, preds2 = fit_from_file(params, learn, test, cd)
        assert n.all(preds1 == preds2)
//  [python]
def test_fit_with_texts(task_type, problem_type):
    if (task_type == 'GPU') and (problem_type == 'multiregression'):
        pytest.xfail('Catboost does not support MultiRMSE on GPU yet for categorical features')

    params = {
        'dictionaries': [
            {'dictionary_id': 'UniGram', 'token_level_type': 'Letter', 'occurrence_lower_bound': '1'},
            {'dictionary_id': 'BiGram', 'token_level_type': 'Letter', 'occurrence_lower_bound': '1', 'gram_order': '2'},
            {'dictionary_id': 'Word', 'occurrence_lower_bound': '1'},
        ],
        'feature_calcers': ['BoW:top_tokens_count=10'] if 'regression' in problem_type else ['NaiveBayes', 'BoW:top_tokens_count=10'],
        'iterations': 100,
        'loss_function': {
            'binclass': 'Logloss',
            'multiclass': 'MultiClass',
            'regression': 'RMSE',
            'multiregression': 'MultiRMSE'
        }[problem_type],
        'task_type': task_type,
        'devices': '0'
    }

    learn = ROTTEN_TOMATOES_TRAIN_FILE
    test = ROTTEN_TOMATOES_TEST_FILE
    cd = ROTTEN_TOMATOES_CD_FILE if problem_type == 'multiclass' else ROTTEN_TOMATOES_CD_BINCLASS_FILE

    preds1 = fit_from_df(params, learn, test, cd, problem_type == 'multiregression')
    if problem_type != 'multiregression':
        _, preds2 = fit_from_file(params, learn, test, cd)
        assert np.all(preds1 == preds2)
//  [python]
def test_feature_importance_interaction_asymmetric_grow_policy():
    pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    model = CatBoostClassifier(
        iterations=5,
        learning_rate=0.03,
        max_ctr_complexity=3)
    model.fit(pool)

    fstr_symm = np.array(model.get_feature_importance(type='Interaction', data=pool))
    model._convert_to_asymmetric_representation()
    fstr_asymm = np.array(model.get_feature_importance(type='Interaction', data=pool))

    assert n.all(fstr_symm - fstr_asymm < 1e-8)
//  [python]
def test_feature_importance_interaction_asymmetric_grow_policy():
    pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    model = CatBoostClassifier(
        iterations=5,
        learning_rate=0.03,
        max_ctr_complexity=3)
    model.fit(pool)

    fstr_symm = np.array(model.get_feature_importance(type='Interaction', data=pool))
    model._convert_to_asymmetric_representation()
    fstr_asymm = np.array(model.get_feature_importance(type='Interaction', data=pool))

    assert np.all(fstr_symm - fstr_asymm < 1e-8)
//  [python]
def test_shap_feature_importance_asymmetric_and_symmetric(task_type):
    pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    model = CatBoostClassifier(
        iterations=5,
        learning_rate=0.03,
        max_ctr_complexity=1,
        task_type=task_type,
        devices='0')
    model.fit(pool)
    shap_symm = np.array(model.get_feature_importance(type=EFstrType.ShapValues, data=pool))
    model._convert_to_asymmetric_representation()
    shap_asymm = np.array(model.get_feature_importance(type=EFstrType.ShapValues, data=pool))
    assert np.all(shap_symm - shap_asymm < 1e-8)
//  [python]
def test_approximate_shap_feature_importance_asymmetric_and_symmetric(task_type):
    pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    model = CatBoostClassifier(
        iterations=5,
        learning_rate=0.03,
        max_ctr_complexity=1,
        task_type=task_type,
        devices='0')
    model.fit(pool)
    shap_symm = np.array(model.get_feature_importance(type=EFstrType.ShapValues, data=pool,
                                                      shap_calc_type="Approximate"))
    model._convert_to_asymmetric_representation()
    shap_asymm = np.array(model.get_feature_importance(type=EFstrType.ShapValues, data=pool,
                                                       shap_calc_type="Approximate"))
    assert np.all(shap_symm - shap_asymm < 1e-8)
//  [python]
def test_custom_eval():
    class LoglossMetric(object):
        def get_final_error(self, error, weight):
            return error / (weight + 1e-38)

        def is_max_optimal(self):
            return True

        def evaluate(self, approxes, target, weight):
            assert len(approxes) == 1
            assert len(target) == len(approxes[0])

            approx = approxes[0]

            error_sum = 0.0
            weight_sum = 0.0

            for i in range(len(approx)):
                w = 1.0 if weight is None else weight[i]
                weight_sum += w
                error_sum += w * (target[i] * approx[i] - math.log(1 + math.exp(approx[i])))

            return error_sum, weight_sum

    train_pool = Pool(data=TRAIN_FILE, column_description=CD_FILE)
    test_pool = Pool(data=TEST_FILE, column_description=CD_FILE)

    model = CatBoostClassifier(iterations=5, use_best_model=True, eval_metric=LoglossMetric())
    model.fit(train_pool, eval_set=test_pool)
    pred1 = model.predict(test_pool)

    model2 = CatBoostClassifier(iterations=5, use_best_model=True, eval_metric='Logloss')
    model2.fit(train_pool, eval_set=test_pool)
    pred2 = model2.predict(test_pool)

    assert n.all(pred1 == pred2)
//  [python]
def test_custom_eval():
    class LoglossMetric(object):
        def get_final_error(self, error, weight):
            return error / (weight + 1e-38)

        def is_max_optimal(self):
            return True

        def evaluate(self, approxes, target, weight):
            assert len(approxes) == 1
            assert len(target) == len(approxes[0])

            approx = approxes[0]

            error_sum = 0.0
            weight_sum = 0.0

            for i in range(len(approx)):
                w = 1.0 if weight is None else weight[i]
                weight_sum += w
                error_sum += w * (target[i] * approx[i] - math.log(1 + math.exp(approx[i])))

            return error_sum, weight_sum

    train_pool = Pool(data=TRAIN_FILE, column_description=CD_FILE)
    test_pool = Pool(data=TEST_FILE, column_description=CD_FILE)

    model = CatBoostClassifier(iterations=5, use_best_model=True, eval_metric=LoglossMetric())
    model.fit(train_pool, eval_set=test_pool)
    pred1 = model.predict(test_pool)

    model2 = CatBoostClassifier(iterations=5, use_best_model=True, eval_metric='Logloss')
    model2.fit(train_pool, eval_set=test_pool)
    pred2 = model2.predict(test_pool)

    assert np.all(pred1 == pred2)
//  [python]
def test_loss_function_change_asymmetric_and_symmetric(task_type):
    pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    model = CatBoostClassifier(
        iterations=5,
        learning_rate=0.03,
        max_ctr_complexity=1,
        task_type=task_type,
        devices='0')
    model.fit(pool)
    shap_symm = np.array(model.get_feature_importance(type=EFstrType.LossFunctionChange, data=pool))
    model._convert_to_asymmetric_representation()
    shap_asymm = np.array(model.get_feature_importance(type=EFstrType.LossFunctionChange, data=pool))
    assert n.all(shap_symm - shap_asymm < 1e-8)
//  [python]
def test_loss_function_change_asymmetric_and_symmetric(task_type):
    pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    model = CatBoostClassifier(
        iterations=5,
        learning_rate=0.03,
        max_ctr_complexity=1,
        task_type=task_type,
        devices='0')
    model.fit(pool)
    shap_symm = np.array(model.get_feature_importance(type=EFstrType.LossFunctionChange, data=pool))
    model._convert_to_asymmetric_representation()
    shap_asymm = np.array(model.get_feature_importance(type=EFstrType.LossFunctionChange, data=pool))
    assert np.all(shap_symm - shap_asymm < 1e-8)
//  [python]
def test_multiclass_classes_count(task_type, missed_classes):
    object_count = 100
    feature_count = 10
    classes_count = 4
    if missed_classes:
        unique_labels = [1, 3]
    else:
        unique_labels = [0, 1, 2, 3]
    expected_classes_attr = [i for i in range(4)]

    prng = np.random.RandomState(seed=0)
    pool = Pool(
        prng.random_sample(size=(object_count, feature_count)),
        label=prng.choice(unique_labels, size=object_count)
    )

    # returns pred probabilities
    def check_classifier(classifier):
        assert np.all(classifier.classes_ == expected_classes_attr)

        pred_classes = classifier.predict(pool)
        assert all([(pred_class in unique_labels) for pred_class in pred_classes])

        pred_probabilities = classifier.predict_proba(pool)
        assert pred_probabilities.shape == (object_count, classes_count)
        return pred_probabilities

    classifier = CatBoostClassifier(
        classes_count=classes_count,
        iterations=2,
        loss_function='MultiClass',
        thread_count=8,
        task_type=task_type,
        devices='0'
    )
    classifier.fit(pool)

    check_classifier(classifier)

    output_model_path = test_output_path(OUTPUT_MODEL_PATH)
    classifier.save_model(output_model_path)

    new_classifier = CatBoostClassifier()
    new_classifier.load_model(output_model_path)

    pred_probabilities = check_classifier(new_classifier)

    preds_path = test_output_path(PREDS_TXT_PATH)
    np.savetxt(preds_path, np.array(pred_probabilities), delimiter='\t')
    return [
        local_canonical_file(preds_path, diff_tool=get_limited_precision_dsv_diff_tool(1e-6, False)),
        compare_canonical_models(output_model_path, diff_limit=1.e-6)
    ]
//  [python]
def test_shap_feature_importance_modes(task_type, calc_shap_mode):
    pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    reference_data = make_reference_data(pool, calc_shap_mode)
    model = CatBoostClassifier(iterations=5, task_type=task_type)
    model.fit(pool)
    modes = ["Auto", "UsePreCalc", "NoPreCalc"]
    shaps_for_modes = []
    for mode in modes:
        shaps_for_modes.append(model.get_feature_importance(type=EFstrType.ShapValues, data=pool, shap_mode=mode, reference_data=reference_data))
    for i in range(len(modes) - 1):
        assert n.all(n.abs(shaps_for_modes[i] - shaps_for_modes[i - 1]) < 1e-9)
//  [python]
def test_shap_feature_importance_modes(task_type, calc_shap_mode):
    pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    reference_data = make_reference_data(pool, calc_shap_mode)
    model = CatBoostClassifier(iterations=5, task_type=task_type)
    model.fit(pool)
    modes = ["Auto", "UsePreCalc", "NoPreCalc"]
    shaps_for_modes = []
    for mode in modes:
        shaps_for_modes.append(model.get_feature_importance(type=EFstrType.ShapValues, data=pool, shap_mode=mode, reference_data=reference_data))
    for i in range(len(modes) - 1):
        assert np.all(np.abs(shaps_for_modes[i] - shaps_for_modes[i - 1]) < 1e-9)
//  [python]
def test_shap_interaction_feature_importance_asymmetric_and_symmetric(task_type):
    pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    model = CatBoostClassifier(
        iterations=5,
        learning_rate=0.03,
        max_ctr_complexity=1,
        task_type=task_type,
        devices='0')
    model.fit(pool)
    shap_interaction_symm = np.array(model.get_feature_importance(type=EFstrType.ShapInteractionValues, data=pool))
    model._convert_to_asymmetric_representation()
    shap_interaction_asymm = np.array(model.get_feature_importance(type=EFstrType.ShapInteractionValues, data=pool))
    assert n.all(shap_interaction_symm - shap_interaction_asymm < 1e-8)
//  [python]
def test_shap_interaction_feature_importance_asymmetric_and_symmetric(task_type):
    pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    model = CatBoostClassifier(
        iterations=5,
        learning_rate=0.03,
        max_ctr_complexity=1,
        task_type=task_type,
        devices='0')
    model.fit(pool)
    shap_interaction_symm = np.array(model.get_feature_importance(type=EFstrType.ShapInteractionValues, data=pool))
    model._convert_to_asymmetric_representation()
    shap_interaction_asymm = np.array(model.get_feature_importance(type=EFstrType.ShapInteractionValues, data=pool))
    assert np.all(shap_interaction_symm - shap_interaction_asymm < 1e-8)
//  [python]
def test_shap_interaction_feature_on_symmetric(task_type):
    pool = Pool(SMALL_CATEGORIAL_FILE, column_description=SMALL_CATEGORIAL_CD_FILE)
    model = CatBoost(params={'loss_function': 'RMSE', 'iterations': 2, 'task_type': task_type, 'devices': '0', 'one_hot_max_size': 4})
    model.fit(pool)
    shap_interaction_values = model.get_feature_importance(
        type=EFstrType.ShapInteractionValues,
        data=pool,
        thread_count=8
    )

    doc_count = pool.num_row()
    shap_interaction_values = shap_interaction_values[:, :-1, :-1]
    for doc_idx in range(doc_count):
        assert n.all(n.abs(shap_interaction_values[doc_idx] - shap_interaction_values[doc_idx].T) < 1e-8)
//  [python]
def test_shap_interaction_feature_on_symmetric(task_type):
    pool = Pool(SMALL_CATEGORIAL_FILE, column_description=SMALL_CATEGORIAL_CD_FILE)
    model = CatBoost(params={'loss_function': 'RMSE', 'iterations': 2, 'task_type': task_type, 'devices': '0', 'one_hot_max_size': 4})
    model.fit(pool)
    shap_interaction_values = model.get_feature_importance(
        type=EFstrType.ShapInteractionValues,
        data=pool,
        thread_count=8
    )

    doc_count = pool.num_row()
    shap_interaction_values = shap_interaction_values[:, :-1, :-1]
    for doc_idx in range(doc_count):
        assert np.all(np.abs(shap_interaction_values[doc_idx] - shap_interaction_values[doc_idx].T) < 1e-8)
//  [python]
def test_str_metrics_in_eval_metrics(task_type):
    train_pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    test_pool = Pool(TEST_FILE, column_description=CD_FILE)
    model = CatBoostClassifier(iterations=40, task_type=task_type, devices='0')
    model.fit(train_pool, eval_set=test_pool)
    first_metrics = model.eval_metrics(data=train_pool, metrics='Logloss')
    second_metrics = model.eval_metrics(data=train_pool, metrics=['Logloss'])
    assert n.all(n.array(first_metrics['Logloss']) == n.array(second_metrics['Logloss']))
//  [python]
def test_str_metrics_in_eval_metrics(task_type):
    train_pool = Pool(TRAIN_FILE, column_description=CD_FILE)
    test_pool = Pool(TEST_FILE, column_description=CD_FILE)
    model = CatBoostClassifier(iterations=40, task_type=task_type, devices='0')
    model.fit(train_pool, eval_set=test_pool)
    first_metrics = model.eval_metrics(data=train_pool, metrics='Logloss')
    second_metrics = model.eval_metrics(data=train_pool, metrics=['Logloss'])
    assert np.all(np.array(first_metrics['Logloss']) == np.array(second_metrics['Logloss']))