text:
import datetime
import os
import numpy as np
import re
import subprocess
import warnings

from collections import OrderedDict
from inspect import getfullargspec

from astra.utils import log, get_base_output_path

expand_path = lambda path: os.path.expandvars(os.path.expanduser(path))

_base_output_path = get_base_output_path()
def output_data_product_path(pk, **kwargs):
    """
    Temporary function to return an output data product path, until we have a data model
    for intermediate ASPCAP data products.
    """
    # Using warnings.warn here so this does not appear \infinity times in a loop
    warnings.warn(f"Using temporary function to return output data product path. This will be deprecated.")

    return os.path.join(
        _base_output_path,
        "ferre",
        "by_task_primary_key",
        f"{pk / 10_000:.0f}",
        f"{pk}.pkl"
    )

# TODO: Put elsewhere
def _grid_mid_point(headers):
    return np.mean(
        np.vstack([
            headers["LLIMITS"],
            headers["ULIMITS"]
        ]),
        axis=0
    )

def _parse_pca_chi(pca_chi, **kwargs):
    return dict(pcachi=int(pca_chi))

def _parse_pca_project(pca_project, **kwargs):
    return dict(pcaproject=int(pca_project))

def _parse_full_covariance(full_covariance, **kwargs):
    return dict(covprint=int(full_covariance))

def _parse_optimization_algorithm_flag(optimization_algorithm_flag, **kwargs):
    available = (1, 2, 3, 4)
    optimization_algorithm_flag = int(optimization_algorithm_flag)
    if optimization_algorithm_flag not in available:
        raise ValueError(f"optimization_algorithm_flag must be one of {available}")
    return dict(algor=optimization_algorithm_flag)

def _parse_interpolation_order(interpolation_order, **kwargs):
    available = {
        0: "nearest neighbour",
        1: "linear",
        2: "quadratic Bezier",
        3: "cubic Bezier",
        4: "cubic splines"
    }
    interpolation_order = int(interpolation_order)
    if interpolation_order not in available:
        raise ValueError(f"interpolation_order must be one of {tuple(list(available.keys()))}")
    return dict(inter=interpolation_order)

def _parse_error_algorithm_flag(error_algorithm_flag, **kwargs):
    available = {
        0: "adopt distance from the solution at which \chi^2 = min(\chi^2) + 1",
        1: "invert the curvature matrix",
        2: "perform numerical experiments injecting noise into the data"
    }
    error_algorithm_flag = int(error_algorithm_flag)
    if error_algorithm_flag not in available:
        raise ValueError(f"error_algorithm_flag must be one of {tuple(list(available.keys()))}")
    return dict(errbar=error_algorithm_flag)

def _parse_input_weights_path(input_weights_path, **kwargs):
    if input_weights_path is not None:
        input_weights_path = expand_path(input_weights_path)
        if not os.path.exists(input_weights_path):
            raise ValueError(f"the input_weights_path does not exist: {input_weights_path}")
        return dict(filterfile=input_weights_path)

def _parse_lsf_shape_flag_and_input_lsf_shape_path(input_lsf_shape_path, lsf_shape_flag, **kwargs):
    available = (0, 1, 2, 3, 4, 11, 12, 13, 14)
    if lsf_shape_flag is not None:
        lsf_shape_flag = int(lsf_shape_flag)
        if lsf_shape_flag not in available:
            raise ValueError(f"lsf_shape_flag must be one of {available}")

        if lsf_shape_flag != 0:
            if input_lsf_shape_path is None:
                raise ValueError(f"lsf_shape_flag is not 0, so an `input_lsf_shape_path` is needed")

            input_lsf_shape_path = expand_path(input_lsf_shape_path)
            if not os.path.exists(input_lsf_shape_path):
                raise ValueError(f"input_lsf_shape_path does not exist: {input_lsf_shape_path}")

            return dict(
                lsf=lsf_shape_flag,
                lsffile=input_lsf_shape_path
            )

    return None

def _parse_wavelength_interpolation_flag(wavelength_interpolation_flag, wavelength, **kwargs):
    available = (0, 1, 2)
    wavelength_interpolation_flag = int(wavelength_interpolation_flag or 0)
    if wavelength_interpolation_flag not in available:
        raise ValueError(f"wavelength_interpolation_flag must be one of {available}")

    if wavelength_interpolation_flag > 0:
        if wavelength is None:
            raise ValueError("if wavelength_interpolation_flag != 0 then wavelength must be given")
        return dict(winter=wavelength_interpolation_flag)

    return None

def _parse_header_path(header_path, **kwargs):
    header_path = expand_path(header_path)
    if not os.path.exists(header_path):
        raise ValueError(f"header_path does not exist: {header_path}")
    return dict([("synthfile(1)", header_path)])

def _parse_n_threads(n_threads, **kwargs):
    return dict(nthreads=int(n_threads))

def _parse_f_access(f_access, **kwargs):
    if f_access is None:
        # TODO: Figure out which is faster. For now default to load everything.
        f_access = False
    return dict(f_access=int(f_access))

def _parse_f_format(f_format, **kwargs):
    available = (0, 1)
    f_format = int(f_format)
    if f_format not in available:
        raise ValueError(f"f_format must be one of {available}")
    return dict(f_format=f_format)

def _parse_continuum_args(
        continuum_flag,
        continuum_order,
        continuum_segment,
        continuum_reject,
        continuum_observations_flag,
        **kwargs
    ):
    kwds = dict()
    if continuum_flag is not None:
        continuum_flag = int(continuum_flag)
        available = (0, 1, 2, 3)
        if continuum_flag not in available:
            raise ValueError(f"continuum_flag must be one of {available}")

        if continuum_flag == 1:
            # need continuum_order and continuum_reject
            if continuum_order is None:
                raise ValueError("continuum_order is required if continuum_flag == 1")
            if continuum_reject is None:
                raise ValueError("continuum_reject is required if continuum_flag == 1")

            kwds.update(
                cont=continuum_flag,
                rejectcont=float(continuum_reject),
                ncont=int(continuum_order)
            )

        elif continuum_flag in (2, 3):
            if continuum_segment is None:
                raise ValueError("continuum_segment is required if continuum_flag is 2 or 3")

            if continuum_flag == 2:
                continuum_segment = int(continuum_segment)
            elif continuum_flag == 3:
                continuum_segment = float(continuum_segment)

            if continuum_segment < 1:
                raise ValueError(f"continuum_segment must be a positive value")

            kwds.update(
                cont=continuum_flag,
                ncont=continuum_segment
            )

    if continuum_observations_flag is not None:
        kwds.update(obscont=int(continuum_observations_flag))

    return kwds

def _parse_ferre_kwargs(ferre_kwargs, **kwargs):
    return (ferre_kwargs or dict())

def _parse_wavelength_and_flux_and_sigma(
        wavelength,
        flux,
        sigma,
        wavelength_interpolation_flag,
        bad_pixel_value=1.0,
        bad_sigma_value=1e6,
        **kwargs
    ):

    # Check flux and sigma
    try:
        flux = np.atleast_2d(np.vstack(flux))
    except ValueError:
        # the number of pixels changes per observation.
        P_flux = list(map(len, flux))
        N_flux = len(flux)

    else:
        N_flux, P_flux = flux.shape
        P_flux = np.atleast_1d(P_flux)

    try:
        sigma = np.atleast_2d(np.vstack(sigma))
    except ValueError:
        P_sigma = list(map(len, sigma))
        N_sigma = len(sigma)
    else:
        N_sigma, P_sigma = sigma.shape
        P_sigma = np.atleast_1d(P_sigma)

    if N_flux != N_sigma:
        raise ValueError(f"the number of flux arrays and sigma arrays do not match ({N_flux} != {N_sigma})")

    if type(P_flux) != type(P_sigma):
        raise ValueError(f"the flux and sigma arrays are incompatible ({type(flux)} and {type(sigma)})")

    if not np.all(P_flux == P_sigma):
        N_mismatches = np.sum(P_flux != P_sigma)
        raise ValueError(f"different length arrays for flux and sigmas: {N_mismatches} mis-matches with {P_flux} and {P_sigma}")

    # TODO: This will cause problems when the number of pixels change per observation.
    bad = ~np.isfinite(flux) + ~np.isfinite(sigma) + (sigma == 0)
    flux[bad] = bad_pixel_value
    sigma[bad] = bad_sigma_value

    if wavelength_interpolation_flag > 0:
        # Check the wavelength values.
        raise NotImplementedError("check wavelength shapes match flux and sigma shapes")

    assert wavelength is not None
    wavelength = np.atleast_1d(wavelength)

    return (wavelength, flux, sigma)

def _default_ferre_kwds(**kwargs):
    return {
        "pfile": "parameters.input",
        "wfile": "wavelengths.input",
        "ffile": "flux.input",
        "erfile": "uncertainties.input",
        "opfile": "parameters.output",
        "offile": "flux.output",
        "sffile": "normalized_flux.output"
    }

def _parse_names_and_initial_and_frozen_parameters(
        names,
        initial_parameters,
        frozen_parameters,
        headers,
        flux,
        clip_initial_parameters_to_boundary_edges=True,
        clip_epsilon_percent=1,
        **kwargs
    ):

    # Read the labels from the first header path
    parameter_names = headers["LABEL"]

    # Need the number of spectra, which we will take from the flux array.
    N = len(flux)
    mid_point = _grid_mid_point(headers)
    parsed_initial_parameters = np.tile(mid_point, N).reshape((N, -1))

    log.debug(f"parsed initial parameters before {parsed_initial_parameters}")

    compare_parameter_names = list(map(sanitise_parameter_name, parameter_names))

    log.debug(f"Initial parameters passed for parsing {initial_parameters}")

    if initial_parameters is not None:
        log.debug(f"Comparison names {compare_parameter_names}")
        for i, (parameter_name, values) in enumerate(initial_parameters.items()):
            spn = sanitise_parameter_name(parameter_name)
            log.debug(f"{parameter_name} {values} {spn}")

            try:
                index = compare_parameter_names.index(spn)
            except ValueError:
                log.warning(f"Ignoring initial parameters for {parameter_name} as they are not in {parameter_names}")
                log.debug(f"Nothing matched for {spn} {parameter_name} {compare_parameter_names}")
            else:
                log.debug(f"Matched to index {index}")
                # Replace non-finite values with the mid point.
                finite = np.isfinite(values)
                if not np.all(finite):
                    log.warning(f"Missing or non-finite initial values given for {parameter_name}. Defaulting to the grid mid-point.")

                values = np.array(values)
                values[~finite] = mid_point[index]

                log.debug(f"values are {values} {type(values[0])} {finite}")
                parsed_initial_parameters[:, index] = values

    log.debug(f"parsed initial parameters after {parsed_initial_parameters}")

    kwds = dict()
    frozen_parameters = (frozen_parameters or dict())
    if frozen_parameters:
        # Ensure we have a dict-like thing.
        if isinstance(frozen_parameters, (list, tuple, np.ndarray)):
            frozen_parameters = { sanitise_parameter_name(k): True for k in frozen_parameters }
        elif isinstance(frozen_parameters, dict):
            # Exclude things that have boolean False.
            frozen_parameters = {
                sanitise_parameter_name(k): v for k, v in frozen_parameters.items() \
                if not (isinstance(v, bool) and not v)
            }
        else:
            raise TypeError(f"frozen_parameters must be list-like or dict-like")

        unknown_parameters = set(frozen_parameters).difference(compare_parameter_names)
        if unknown_parameters:
            raise ValueError(f"unknown parameter(s): {unknown_parameters} (available: {parameter_names})")

        indices = [
            i for i, pn in enumerate(compare_parameter_names, start=1) if pn not in frozen_parameters
        ]

        if len(indices) == 0:
            raise ValueError(f"all parameters frozen?!")

        # Over-ride initial values with the frozen ones if given.
        for parameter_name, value in frozen_parameters.items():
            if not isinstance(value, bool):
                log.debug(f"Over-writing initial values for {parameter_name} with frozen value of {value}")
                zero_index = compare_parameter_names.index(parameter_name)
                parsed_initial_parameters[:, zero_index] = value
    else:
        # No frozen parameters.
        indices = 1 + np.arange(len(parameter_names), dtype=int)

    # Build a frozen parameters dict for result metadata.
    parsed_frozen_parameters = { pn: (pn in frozen_parameters) for pn in compare_parameter_names }

    L = len(indices)
    kwds.update(
        ndim=headers["N_OF_DIM"],
        nov=L,
        indv=" ".join([f"{i:.0f}" for i in indices]),
        # We will always provide an initial guess, even if it is the grid mid point.
        init=0,
        indini=" ".join(["1"] * L)
    )

    # Now deal with names.
    if names is None:
        names = [f"{i:.0f}" for i in range(len(parsed_initial_parameters))]
    else:
        if len(names) != len(parsed_initial_parameters):
            raise ValueError(f"names and initial parameters does not match ({len(names)} != {len(parsed_initial_parameters)})")

    # Let's check the initial values are all within the grid boundaries.
    lower_limit, upper_limit = _get_grid_limits(headers)
    try:
        _check_initial_parameters_within_grid_limits(parsed_initial_parameters, lower_limit, upper_limit, parameter_names)
    except ValueError as e:
        log.exception(f"Exception when checking initial parameters within grid boundaries:")
        log.critical(e, exc_info=True)

        if clip_initial_parameters_to_boundary_edges:
            log.info(f"Clipping initial parameters to boundary edges (use clip_initial_parameters_to_boundary_edges=False to raise exception instead)")

            clip = clip_epsilon_percent * (upper_limit - lower_limit)/100.
            parsed_initial_parameters = np.round(np.clip(parsed_initial_parameters, lower_limit + clip, upper_limit - clip), 3)
        else:
            raise

    return (kwds, names, parsed_initial_parameters, parsed_frozen_parameters)

def _get_grid_limits(headers):
    """
    Return a two-length tuple that contains the lower limits of the grid and the upper limits of the grid.

    :param headers:
        The primary headers from a FERRE grid file.
    """
    return (
        headers["LLIMITS"],
        headers["LLIMITS"] + headers["STEPS"] * (headers["N_P"] - 1)
    )

def _check_initial_parameters_within_grid_limits(initial_parameters, lower_limit, upper_limit, parameter_names):
    """
    Check that the initial parameters are within the boundaries of the FERRE grid.

    :param initial_parameters:
        A 2D array of shape (N, L) where N is the number of spectra and L is the number of labels in the FERRE grid.

    :param headers:
        The primary headers from the FERRE grid file.

    :raise ValueError:
        If any initial parameter is outside the grid boundary.
    """

    in_limits = (upper_limit > initial_parameters) * (initial_parameters > lower_limit)

    if not np.all(in_limits):

        message = "Initial_parameters are not all within bounds of the grid. For example:\n"

        bad_parameter_indices = np.where(~np.all(in_limits, axis=0))[0]
        for j in bad_parameter_indices:
            i = np.where(~in_limits[:, j])[0]
            message += f"- {parameter_names[j]} has limits of ({lower_limit[j]:.2f}, {upper_limit[j]:.2f}) but {len(i)} indices ({i}) are outside this range: {initial_parameters[:, j][i]}\n"

        raise ValueError(message)

    return True

def parse_ferre_inputs(**kwargs):

    parsed_kwds = {}

    headers, *segment_headers = read_ferre_headers(expand_path(kwargs["header_path"]))

    # These two have different outputs than the rest.
    wavelength, flux, sigma = _parse_wavelength_and_flux_and_sigma(**kwargs)

    # Build wavelength arrays to generate a mask.
    model_wavelengths = tuple(map(wavelength_array_from_ferre_header, segment_headers))

    # Build from first wavelength array
    mask = np.zeros(wavelength.shape[1], dtype=bool)
    for model_wavelength in model_wavelengths:
        # TODO: Building wavelength mask off just the first wavelength array. Assuming all have the same wavelength array.
        s_index, e_index = wavelength[0].searchsorted(model_wavelength[[0, -1]])
        mask[s_index:e_index + 1] = True

    _kwds, names, initial_parameters, parsed_frozen_parameters = _parse_names_and_initial_and_frozen_parameters(
        headers=headers,
        **kwargs
    )
    parsed_kwds.update(_kwds)

    parsers = (
        _default_ferre_kwds,
        _parse_pca_chi,
        _parse_pca_project,
        _parse_full_covariance,
        _parse_optimization_algorithm_flag,
        _parse_interpolation_order,
        _parse_error_algorithm_flag,
        _parse_input_weights_path,
        _parse_lsf_shape_flag_and_input_lsf_shape_path,
        _parse_wavelength_interpolation_flag,
        _parse_header_path,
        _parse_n_threads,
        _parse_f_access,
        _parse_f_format,
        _parse_continuum_args,
        _parse_ferre_kwargs,
    )

    for parser in parsers:
        _kwds = parser(**kwargs)
        if _kwds: parsed_kwds.update(_kwds)

    # Create a metadata dict for results.
    meta = dict(
        parameter_names=headers["LABEL"],
        frozen_parameters=parsed_frozen_parameters,
        initial_parameters=initial_parameters
    )
    return (wavelength, flux, sigma, mask, names, initial_parameters, parsed_kwds, meta)

sanitise_parameter_name = lambda pn: pn.lower().strip().replace(" ", "_")
desanitise_parameter_name = lambda pn: pn.upper().replace("_", " ")

def line_count(filename):
    return int(subprocess.check_output(['wc', '-l', filename]).split()[0])

def wavelength_array_from_ferre_header(header):
    wave = header["WAVE"][0] + np.arange(header["NPIX"]) * header["WAVE"][1]
    if header["LOGW"] == 1:
        wave = 10**wave
    elif header["LOGW"] == 2:
        wave = np.exp(wave)

    return wave

def write_data_file(data, path, fmt="%.4e", footer="\n", **kwargs):
    return np.savetxt(
        path,
        data,
        fmt=fmt,
        footer=footer,
        **kwargs
    )

def safe_read_header(headers, keys):
    if not isinstance(keys, (tuple, list)):
        keys = [keys]

    for key in keys:
        try:
            return headers[key]
        except KeyError:
            continue

    else:
        raise KeyError(f"no header keyword found among {keys}")

def read_ferre_header(fp):
    """
    Read a FERRE library header into a dictionary.

    This functionality was originally written by Jon Holtzmann.

    :param fp:
        A file pointer to a FERRE header file.
    """

    header = dict()
    for i, line in enumerate(fp):
        if line.startswith(" /"):
            break

        if "=" in line:
            key, value = line.split(" = ")

            # Check values.
            if value.lstrip()[0] != "'":
                # Treat as numerical values.
                values = re.findall(
                    "[+|-|\d|\.|e|-]+",
                    value
                )

                dtype = float if "." in value else int

                if len(values) == 1:
                    value = dtype(values[0])
                else:
                    value = np.array(values, dtype=dtype)

            else:
                value = value.strip(" '\n")

            match = re.match(r'\s+(?P<key>\w+)\((?P<index>\d+)\)', key)
            if match:
                match = match.groupdict()
                key = match["key"]
                header.setdefault(key, [])
                # TODO: assuming header information is in order
                header[key].append(value)
            else:
                key = key.strip()
                header[key] = value

    # Put in upper grid limits as a 'value-added' header.
    header["ULIMITS"] = header["LLIMITS"] + header["STEPS"] * (header["N_P"] - 1)

    return header

def read_ferre_headers(path):
    """
    Read a full FERRE library header with multi-extensions.
    :param path:
        The path of a FERRE header file.
    Returns:
       libstr0, libstr : first header, then list of extension headers; headers returned as dictionaries
    """

    try:
        with open(path, "r") as fp:

            headers = [read_ferre_header(fp)]
            headers[0]["PATH"] = path

            for i in range(headers[0].get("MULTI", 0)):
                headers.append(read_ferre_header(fp))

    except:
        raise

    else:
        return headers

def get_grid_boundaries(ferre_grid_header_files, **kwargs):
    """
    Get the grid boundaries for the given FERRE grid header files.

    :param ferre_grid_header_files:
        A list of FERRE grid paths.

    :returns:
        A dictionary with FERRE grid paths as keys, and each value contains
        a dictionary of label names as keys and a tuple of (lower, upper)
        boundary values.
    """
    boundaries = {}

    for path in ferre_grid_header_files:
        expanded_path = os.path.expandvars(path)
        primary_header, *headers = read_ferre_headers(expanded_path)

        args = (
            primary_header["LABEL"],
            primary_header["LLIMITS"],
            primary_header["ULIMITS"]
        )
        boundaries[expanded_path] = {
            label_name: (lower, upper) for label_name, lower, upper in zip(*args)
        }

    return boundaries

def read_output_parameter_file(path, n_dimensions, full_covariance, **kwargs):
    """
    Three more columns follow, giving the fraction of photometric data
    points (useful when multiple grids combining spectroscopy and photometry are used),
    the average log(S/N)2
    for the spectrum, and the logarithm of the reduced χ^2 for the fit.

    Additional columns with the covariance matrix of the errors can be output setting to 1 the
    keyword COVPRINT.
    """

    names = np.loadtxt(path, usecols=(0, ), dtype=str)

    N_cols = 2 * n_dimensions + 3
    if full_covariance:
        raise NotImplementedError("needs testing")
        N_cols += 3 + n_dimensions**2

    results = np.atleast_2d(np.loadtxt(path, usecols=1 + np.arange(N_cols)))

    param = results[:, 0:n_dimensions]
    param_err = results[:, n_dimensions:2*n_dimensions]

    frac_phot_data_points, log_snr_sq, log_chisq_fit = results[:, 2*n_dimensions:2*n_dimensions + 3].T
    meta = dict(
        frac_phot_data_points=frac_phot_data_points,
        log_snr_sq=log_snr_sq,
        log_chisq_fit=log_chisq_fit
    #    cov=cov,
    )
    if full_covariance:
        cov = results[:, 2*n_dimensions+3:2*n_dimensions+3+n_dimensions**2]
        cov = cov.reshape((-1, n_dimensions, n_dimensions))
        meta.update(cov=cov)

    return (names, param, param_err, meta)

def format_ferre_control_keywords(ferre_kwds):
    r"""
    Format control keywords for FERRE to digest.

    :param ferre_kwds:
        A dictionary of FERRE-recognized keywords, in a suitable order.

    :returns:
        String contents for a FERRE input file.
    """

    preferred_order = (
        "ndim",
        "nov",
        "indv",
        "synthfile(1)",
        "pfile",
        "ffile",
        "erfile",
        "opfile",
        "offile",
        "obscont",
        "inter",
        "errbar",
        "nruns",
        "init",
        "indini",
        "winter",
        "algor",
        "lsf",
        "nthreads",
        "covprint",
        "pcachi",
        "f_format",
        "f_access",
    )

    contents = "&LISTA\n"
    remaining_keys = set(ferre_kwds).difference(preferred_order)
    keys = list(preferred_order) + list(remaining_keys)

    for key in keys:
        if key in ferre_kwds:
            value = ferre_kwds[key]
            if isinstance(value, str) and key not in ("indv", "indini"):
                value = f"'{value}'"
            else:
                value = f"{value}"
            contents += f"{key.upper()} = {value}\n"

    contents += "/\n"
    return contents

# Note: this old_format_ferre_input_parameters certainly will match the IPF files on the SDSS SAS,
#       but are giving me FERRE problems right now. Check format later on.
def format_ferre_input_parameters(*p, name="dummy"):
    r"""
    Format input parameters for FERRE to digest.

    :returns:
        String contents for a FERRE input parameter file.
    """
    contents = f"{name:40s}"
    for each in p:
        contents += f"{each:12.3f}"
    contents += "\n"
    return contents

def get_processing_times(stdout, n_threads):
    """
    Get the time taken to analyse spectra and estimate the initial load time.

    :param stdout: (optional)
        The standard output from FERRE.
    """

    if stdout is None or stdout == "":
        return None

    matches = re.findall('ellapsed time:\s+[{0-9}|.]+', stdout)
    load_time, *elapsed_time = [float(match.split()[-1]) for match in matches]

    # Offset load time.
    elapsed_time = np.array(elapsed_time) - load_time

    # Account for number of threads.
    O = n_threads - (elapsed_time.size % n_threads)
    A = np.hstack([elapsed_time, np.zeros(O)]).reshape((-1, n_threads))
    time_per_spectrum = np.hstack([elapsed_time[:n_threads], np.diff(A, axis=0).flatten()[:-O]])

    object_indices = [int(index.split()[-1]) for index in re.findall('next object #\s+[{0-9}]+', stdout)]

    # Sort.
    idx = np.argsort(object_indices)
    time_per_spectrum = time_per_spectrum[idx]

    return dict(
        load_time=load_time,
        elapsed_time=elapsed_time,
        indices=idx,
        time_per_ordered_spectrum=time_per_spectrum
    )

def get_abundance_keywords(element, header_label_names):
    """
    Return a dictionary of task parameters given a chemical element. These are adopted from DR16.

    :param element:
        The chemical element to measure.

    :param header_label_names:
        The list of label names in the FERRE header file.
    """

    # These can be inferred from running the following command on the SAS:
    # cd /uufs/chpc.utah.edu/common/home/sdss50/dr16/apogee/spectro/aspcap/r12/l33/apo25m/cal_all_apo25m007/ferre
    # egrep 'INDV|TIE|FILTERFILE' */input.nml

    controls = {
        "Al": {
            "INDV_LABEL": ('METALS', ),
            "TIES": [
                ('C', 0, -1),
                ('N', 0, -1),
                ('O Mg Si S Ca Ti', 0, -1)
            ]
        },
        "Ca": {
            "INDV_LABEL": ('O Mg Si S Ca Ti', ),
        },
        "Ce": {
            "INDV_LABEL": ('METALS', ),
        },
        "CI": {
            "INDV_LABEL": ('C', ),
        },
        "C": {
            "INDV_LABEL": ('C', ),
        },
        "CN": {
            "INDV_LABEL": ('C', 'O Mg Si S Ca Ti', ),
        },
        "Co": {
            "INDV_LABEL": ('METALS', ),
            "TIES": [
                ('C', 0, -1),
                ('N', 0, -1),
                ('O Mg Si S Ca Ti', 0, -1)
            ]
        },
        "Cr": {
            "INDV_LABEL": ('METALS', ),
            "TIES": [
                ('C', 0, -1),
                ('N', 0, -1),
                ('O Mg Si S Ca Ti', 0, -1)
            ]
        },
        "Cu": {
            "INDV_LABEL": ('METALS', ),
            "TIES": [
                ('C', 0, -1),
                ('N', 0, -1),
                ('O Mg Si S Ca Ti', 0, -1)
            ]
        },
        "Fe": {
            "INDV_LABEL": ('METALS', ),
            "TIES": [
                ('C', 0, -1),
                ('N', 0, -1),
                ('O Mg Si S Ca Ti', 0, -1)
            ]
        },
        "Ge": {
            "INDV_LABEL": ('METALS', ),
            "TIES": [
                ('C', 0, -1),
                ('N', 0, -1),
                ('O Mg Si S Ca Ti', 0, -1)
            ]
        },
        "K": {
            "INDV_LABEL": ('METALS', ),
            "TIES": [
                ('C', 0, -1),
                ('N', 0, -1),
                ('O Mg Si S Ca Ti', 0, -1)
            ]
        },
        "Mg": {
            "INDV_LABEL": ('O Mg Si S Ca Ti', ),
        },
        "Mn": {
            "INDV_LABEL": ('METALS', ),
            "TIES": [
                ('C', 0, -1),
                ('N', 0, -1),
                ('O Mg Si S Ca Ti', 0, -1)
            ]
        },
        "Na": {
            "INDV_LABEL": ('METALS', ),
            "TIES": [
                ('C', 0, -1),
                ('N', 0, -1),
                ('O Mg Si S Ca Ti', 0, -1)
            ]
        },
        "Nd": {
            "INDV_LABEL": ('METALS', ),
            "TIES": [
                ('C', 0, -1),
                ('N', 0, -1),
                ('O Mg Si S Ca Ti', 0, -1)
            ]
        },
        "Ni": {
            "INDV_LABEL": ('METALS', ),
            "TIES": [
                ('C', 0, -1),
                ('N', 0, -1),
                ('O Mg Si S Ca Ti', 0, -1)
            ]
        },
        "N": {
            "INDV_LABEL": ('N', ),
        },
        "O": {
            "INDV_LABEL": ('O Mg Si S Ca Ti', ),
        },
        "P": {
            "INDV_LABEL": ('METALS', ),
            "TIES": [
                ('C', 0, -1),
                ('N', 0, -1),
                ('O Mg Si S Ca Ti', 0, -1)
            ]
        },
        "Rb": {
            "INDV_LABEL": ('METALS', ),
            "TIES": [
                ('C', 0, -1),
                ('N', 0, -1),
                ('O Mg Si S Ca Ti', 0, -1)
            ]
        },
        "Si": {
            "INDV_LABEL": ('O Mg Si S Ca Ti', ),
        },
        "S": {
            "INDV_LABEL": ('O Mg Si S Ca Ti', ),
        },
        "TiII": {
            "INDV_LABEL": ('O Mg Si S Ca Ti', ),
        },
        "Ti": {
            "INDV_LABEL": ('O Mg Si S Ca Ti', ),
        },
        "V": {
            "INDV_LABEL": ('METALS', ),
            "TIES": [
                ('C', 0, -1),
                ('N', 0, -1),
                ('O Mg Si S Ca Ti', 0, -1)
            ]
        },
        "Yb": {
            "INDV_LABEL": ('METALS', ),
            "TIES": [
                ('C', 0, -1),
                ('N', 0, -1),
                ('O Mg Si S Ca Ti', 0, -1)
            ]
        }
    }

    def get_header_index(label):
        # FERRE uses 1-indexing and Python uses 0-indexing.
        return 1 + header_label_names.index(label)

    try:
        c = controls[element]
    except:
        raise ValueError(f"no abundance controls known for element '{element}' (available: {tuple(controls.keys())}")

    indv = [get_header_index(label) for label in c["INDV_LABEL"]]
    ties = c.get("TIES", [])

    ferre_kwds = {
        # We don't pass INDV here because this will be determined from the
        # 'frozen_<param>' arguments to the FerreGivenSDSSApStarFile tasks
        #"INDV": [get_header_index(label) for label in c["INDV_LABEL"]],
        "NTIE": len(ties),
        "TYPETIE": 1
    }
    for i, (tie_label, ttie0, ttie) in enumerate(ties, start=1):
        ferre_kwds.update({
            f"INDTIE({i:.0f})": get_header_index(tie_label),
            f"TTIE0({i:.0f})": ttie0,
            # TODO: What if we don't want to tie it back to first INDV element?
            f"TTIE({i:.0f},{indv[0]:.0f})": ttie
        })

    # Freeze all other labels.
    frozen_parameters = { hln: (hln not in c["INDV_LABEL"]) for hln in header_label_names }

    return (frozen_parameters, ferre_kwds)

max_stars_repo_name: sdss/astra

max_issues_repo_name: sdss/astra

max_forks_repo_name: sdss/astra


text:
import logging
from ror.Relation import Relation
from ror.Dataset import Dataset
from typing import List, Tuple
from ror.Constraint import Constraint, ConstraintVariable, ConstraintVariablesSet, ValueConstraintVariable
import numpy as np

# difference of 2 values greater than DIFF_EPS indicates that they are different
DIFF_EPS = 1e-10

def check_preconditions(data: Dataset) -> bool:
    if len(data.alternatives) < 3:
        logging.info('number of alternatives is lower than 3, skipping slope constraint')
        return False
    return True

def _create_slope_constraint(
        alternative_index: int,
        data: Dataset,
        criterion_name: str,
        relation: Relation,
        alternatives: List[str],
        alternative_scores: List[float]) -> Tuple[Constraint, Constraint]:
    '''
    Returns slope constraint or None if there would be division by 0 (in case when g_i(l) == g_i(l-1) or g_i(l-1) == g_i(l-2))
    Slope constraint is meeting the requirement | z - w | <= rho
    This constraint minimizes the differences between 2 consecutive characteristic points.
    This constraint requires partial utility function to be monotonic, non-decreasing
    '''
    first_diff = alternative_scores[alternative_index] - alternative_scores[alternative_index-1]
    # check if the 2 following points are not in the same place
    if abs(first_diff) < DIFF_EPS:
        logging.debug(
            f'Criterion {criterion_name} for alternative {alternatives[alternative_index]} has the same value ({alternative_scores[alternative_index-1]}) as alternative {alternatives[alternative_index-1]} on this criterion.')
        return None
    first_coeff = 1 / (first_diff)
    second_diff = alternative_scores[alternative_index-1] - alternative_scores[alternative_index-2]
    # check if the 2 following points are not in the same place
    if abs(second_diff) < DIFF_EPS:
        logging.debug(
            f'Criterion {criterion_name} for alternative {alternatives[alternative_index-1]} has the same value ({alternatives[alternative_index-2]}) as alternative {alternatives[alternative_index-2]} on this criterion.')
        return None
    second_coeff = 1 / (second_diff)

    delta_constraint = ConstraintVariable(
        "delta",
        -1.0
    ) if data.delta is None else ValueConstraintVariable(
        data.delta
    )

    # create constraint
    first_constraint = Constraint(ConstraintVariablesSet([
        ConstraintVariable(
            Constraint.create_variable_name(
                'u', criterion_name, alternatives[alternative_index]),
            first_coeff,
            alternatives[alternative_index]
        ),
        ConstraintVariable(
            Constraint.create_variable_name(
                'u', criterion_name, alternatives[alternative_index-1]),
            -first_coeff,
            alternatives
        ),
        ConstraintVariable(
            Constraint.create_variable_name(
                'u', criterion_name, alternatives[alternative_index-1]),
            -second_coeff,
            alternatives[alternative_index-1]
        ),
        ConstraintVariable(
            Constraint.create_variable_name(
                'u', criterion_name, alternatives[alternative_index-2]),
            second_coeff,
            alternatives[alternative_index-2]
        ),
        delta_constraint
    ]), relation, Constraint.create_variable_name("first_slope", criterion_name, alternative_index))

    second_constraint = Constraint(ConstraintVariablesSet([
        ConstraintVariable(
            Constraint.create_variable_name(
                'u', criterion_name, alternatives[alternative_index]),
            -first_coeff,
            alternatives[alternative_index]
        ),
        ConstraintVariable(
            Constraint.create_variable_name(
                'u', criterion_name, alternatives[alternative_index-1]),
            first_coeff,
            alternatives[alternative_index-1]
        ),
        ConstraintVariable(
            Constraint.create_variable_name(
                'u', criterion_name, alternatives[alternative_index-1]),
            second_coeff,
            alternatives[alternative_index-1]
        ),
        ConstraintVariable(
            Constraint.create_variable_name(
                'u', criterion_name, alternatives[alternative_index-2]),
            -second_coeff,
            alternatives[alternative_index-2]
        ),
        delta_constraint
    ]), relation, Constraint.create_variable_name("second_slope", criterion_name, alternative_index))

    return (first_constraint, second_constraint)

def create_slope_constraints(data: Dataset, relation: Relation = None) -> List[Constraint]:
    '''
    Returns slope constraints for all alternatives except the ones that have duplicated
    values in the criterion space.
    So the number of constraints will be
    2 x criteria + (m-2)*2
    where 'm' is the number of alternatives without duplicated data on each criterion
    and 'criteria' is the number of criteria in the data.
    '''
    if not check_preconditions(data):
        return []

    if relation is None:
        relation = Relation('<=')
    constraints = []
    for criterion_index, (criterion_name, _) in enumerate(data.criteria):
        alternative_score_on_criterion = data.matrix[:, criterion_index]

        for l in range(2, len(data.alternatives)):
            slope_constraints = _create_slope_constraint(
                l, data, criterion_name, relation, data.alternatives, alternative_score_on_criterion
            )
            if slope_constraints is not None:
                first_constraint, second_constraint = slope_constraints
                constraints.append(first_constraint)
                constraints.append(second_constraint)
    return constraints

ext: py

max_stars_repo_name: jakub-tomczak/ror

max_issues_repo_name: jakub-tomczak/ror

max_forks_repo_name: jakub-tomczak/ror

reason: import numpy


text:
import mbuild as mb
import numpy as np

class Isopropyl(mb.Compound):
    """ A isopropyl group. """
    def __init__(self):
        super(Isopropyl, self).__init__()

        mb.load('isopropyl.pdb', compound=self, relative_to_module=self.__module__)
        self.translate(-self[0].pos)

        self.add(mb.Port(anchor=self[0], orientation=[0, -1, 0], separation=0.07),
            'down')

if __name__ == '__main__':
    isopropyl = Isopropyl()
    isopropyl.save('isopropyl-test.mol2', overwrite=True)

max_stars_repo_name: dubosese/atools

max_issues_repo_name: dubosese/atools

max_forks_repo_name: dubosese/atools


text:
# Copyright 2021 The Brax Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Composer for environments.

ComponentEnv composes a scene from descriptions of the form below:

   composer = Composer(
    components=dict(
        ant1=dict(component='ant', pos=(0, 1, 0)),
        ant2=dict(component='ant', pos=(0, -1, 0)),
    ),
    edges=dict(ant1__ant2=dict(collide_type='full'),),
   )
   env = ComposeEnv(composer=composer)

 (more examples available in experimental/composer/env_descs.py)

During loading, it:
- creates components: loads and pieces together brax.Config()
    components defined in experimental/composer/components/
    such as ant.py or ground.py
  - support multiple instances of the same component through suffixes
  - each component requires: ROOT=root body, SYS_CONFIG=config in string form,
      TERM_FN=termination function of this component, COLLIDES=bodies that
      are allowed to collide, DEFAULT_OBSERVERS=a list of observers (
      see experimental/composer/observers.py for references)
- creates edges: automatically create necessary edge information
    between components, such as collide_include's in brax.Config()
  - optionally edge information can be supplied,
      e.g. `collide_type`={'full', 'root', None} specifying full collisons,
      collision only between roots, or no collision between two components
- sets termination as any(termination_fn of each component)
- sets observation to concatenation of observations of each component defined
    by each component's `observers` argument
"""
import collections
import copy
import functools
import itertools
from typing import Dict, Any, Callable

import brax
from brax import envs
from brax.envs import Env
from brax.envs import State
from brax.experimental.braxlines.common import sim_utils
from brax.experimental.composer import component_editor
from brax.experimental.composer import env_descs
from brax.experimental.composer import observers
import jax
from jax import numpy as jnp

MetaData = collections.namedtuple(
    'MetaData',
    ['components', 'edges', 'global_options', 'config_str', 'config_json'])

class Composer(object):
  """Compose a brax system."""

  def __init__(self,
               components: Dict[str, Dict[str, Any]],
               edges: Dict[str, Dict[str, Any]] = None,
               add_ground: bool = True,
               global_options: Dict[str, Any] = None):
    components = copy.deepcopy(components)
    edges = copy.deepcopy(edges or {})

    # load components
    if add_ground:
      components['ground'] = dict(component='ground')
    components = {
        name: component_editor.load_component(**value)
        for name, value in components.items()
    }
    component_keys = sorted(components.keys())
    components_ = collections.OrderedDict([
        (k, components[k]) for k in component_keys
    ])

    # set global
    v = dict(
        json=component_editor.json_global_options(**(global_options or {})))
    v['message_str'] = component_editor.json2message_str(v['json'])
    global_options_ = v

    for k, v in components_.items():

      # convert to json format for easy editing
      v['json'] = component_editor.message_str2json(v['message_str'])
      # add suffices
      suffix = v.get('suffix', k)
      if suffix:
        rename_fn = functools.partial(
            component_editor.json_add_suffix, suffix=suffix)
        v['json'] = rename_fn(v['json'])
        v['collides'] = rename_fn(v['collides'], force_add=True)
        v['root'] = rename_fn(v['root'], force_add=True)
      v['bodies'] = [b['name'] for b in v['json'].get('bodies', [])]
      v['joints'] = [b['name'] for b in v['json'].get('joints', [])]
      v['suffix'] = suffix
      # convert back to str
      v['message_str'] = component_editor.json2message_str(v['json'])

      # set transform or not
      if 'pos' in v or 'quat' in v:
        v['transform'] = True
        v['pos'] = jnp.array(v.get('pos', [0, 0, 0]), dtype='float')
        v['quat_origin'] = jnp.array(
            v.get('quat_origin', [0, 0, 0]), dtype='float')
        v['quat'] = jnp.array(v.get('quat', [1., 0., 0., 0.]), dtype='float')
      else:
        v['transform'] = False

    edges_ = {}
    for k1, k2 in itertools.combinations(list(components_.keys()), 2):
      if k1 == k2:
        continue
      k1, k2 = sorted([k1, k2])  # ensure the name is always sorted in order
      edge_name = f'{k1}__{k2}'
      v, new_v = edges.pop(edge_name, {}), {}
      v1, v2 = [components_[k] for k in [k1, k2]]
      collide_type = v.pop('collide_type', 'full')
      v_json = {}
      # add colliders
      if collide_type == 'full':
        v_json.update(
            component_editor.json_collides(v1['collides'], v2['collides']))
      elif collide_type == 'root':
        v_json.update(
            component_editor.json_collides([v1['root']], [v2['root']]))
      else:
        assert not collide_type, collide_type
      if v_json:
        # convert back to str
        new_v['message_str'] = component_editor.json2message_str(v_json)
      else:
        new_v['message_str'] = ''
      new_v['json'] = v_json
      assert not v, f'unused edges[{edge_name}]: {v}'
      edges_[edge_name] = new_v
    assert not edges, f'unused edges: {edges}'
    edge_keys = sorted(edges_.keys())
    edges_ = collections.OrderedDict([(k, edges_[k]) for k in edge_keys])

    # merge all message strs
    message_str = ''
    for _, v in sorted(components_.items()):
      message_str += v.get('message_str', '')
    for _, v in sorted(edges_.items()):
      message_str += v.get('message_str', '')
    message_str += global_options_.get('message_str', '')
    config_str = message_str
    config_json = component_editor.message_str2json(message_str)
    metadata = MetaData(
        components=components_,
        edges=edges_,
        global_options=global_options_,
        config_str=config_str,
        config_json=config_json,
    )
    config = component_editor.message_str2message(message_str)
    self.config, self.metadata = config, metadata

  def reset_fn(self, sys, qp: brax.QP):
    """Reset state."""
    # apply translations and rotations
    for _, v in sorted(self.metadata.components.items()):
      if v['transform']:
        _, _, mask = sim_utils.names2indices(sys.config, v['bodies'],
                                             'body')
        qp = sim_utils.transform_qp(qp, mask[..., None], v['quat'],
                                    v['quat_origin'], v['pos'])
    return qp

  def term_fn(self, done: jnp.ndarray, sys, qp: brax.QP, info: brax.Info):
    """Termination."""
    for k, v in self.metadata.components.items():
      term_fn = v['term_fn']
      if term_fn:
        done = term_fn(done, sys, qp, info, k)
    return done

  def obs_fn(self, sys, qp: brax.QP, info: brax.Info):
    """Return observation as OrderedDict."""
    obs_dict = collections.OrderedDict()
    for _, v in self.metadata.components.items():
      for observer in v['observers']:
        obs_dict_ = observers.get_obs_dict(sys, qp, info, observer, v)
        obs_dict = collections.OrderedDict(
            list(obs_dict.items()) + list(obs_dict_.items()))
    return obs_dict

class ComponentEnv(Env):
  """Make a brax Env fromc config/metadata for training and inference."""

  def __init__(self, composer: Composer, *args, **kwargs):
    self.observer_shapes = None
    self.composer = composer
    super().__init__(*args, config=self.composer.config, **kwargs)

  def reset(self, rng: jnp.ndarray) -> State:
    """Resets the environment to an initial state."""
    qp = self.sys.default_qp()
    qp = self.composer.reset_fn(self.sys, qp)
    info = self.sys.info(qp)
    obs = self._get_obs(qp, info)
    reward, done, steps = jnp.zeros(3)
    metrics = {}
    return State(rng, qp, info, obs, reward, done, steps, metrics)

  def step(self, state: State, action: jnp.ndarray) -> State:
    """Run one timestep of the environment's dynamics."""
    rng = state.rng
    # some codes, e.g. policy function, does not update state.rng.
    # this ensures that at each step, state.rng is updated for external use.
    rng, _ = jax.random.split(rng)
    qp, info = self.sys.step(state.qp, action)
    obs = self._get_obs(qp, info)
    reward = 0.0
    steps = state.steps + self.action_repeat
    done = jnp.where(steps >= self.episode_length, x=1.0, y=0.0)
    done = self.composer.term_fn(done, self.sys, qp, info)
    metrics = {}
    return State(rng, qp, info, obs, reward, done, steps, metrics)

  def _get_obs(
      self,
      qp: brax.QP,
      info: brax.Info,
  ) -> jnp.ndarray:
    """Observe."""
    obs_dict = self.composer.obs_fn(self.sys, qp, info)
    if self.observer_shapes is None:
      self.observer_shapes = observers.get_obs_dict_shape(obs_dict)
    return jnp.concatenate(list(obs_dict.values()))

def create(env_name: str = None,
           components: Dict[str, Dict[str, Any]] = None,
           edges: Dict[str, Dict[str, Any]] = None,
           add_ground: bool = True,
           global_options: Dict[str, Any] = None,
           **kwargs) -> Env:
  """Creates an Env with a specified brax system."""
  if env_name in env_descs.ENV_DESCS:
    composer = Composer(
        add_ground=add_ground,
        global_options=global_options,
        **env_descs.ENV_DESCS[env_name])
    return ComponentEnv(composer=composer, **kwargs)
  elif components:
    composer = Composer(
        components=components,
        edges=edges,
        add_ground=add_ground,
        global_options=global_options)
    return ComponentEnv(composer=composer, **kwargs)
  else:
    return envs.create(env_name, **kwargs)

def create_fn(env_name: str = None,
              components: Dict[str, Dict[str, Any]] = None,
              edges: Dict[str, Dict[str, Any]] = None,
              add_ground: bool = True,
              global_options: Dict[str, Any] = None,
              **kwargs) -> Callable[..., Env]:
  """Returns a function that when called, creates an Env."""
  return functools.partial(
      create,
      env_name=env_name,
      components=components,
      edges=edges,
      add_ground=add_ground,
      global_options=global_options,
      **kwargs)

max_stars_repo_name: lweitkamp/brax

max_issues_repo_name: lweitkamp/brax

max_forks_repo_name: lweitkamp/brax

reason: import numpy,import jax,from jax


text:
import os
import json
import pytest
import numpy as np
tf = pytest.importorskip('tensorflow')
from baseline.w2v import RandomInitVecModel
from baseline.tf.tagger.model import create_model, RNNTaggerModel
from baseline.utils import crf_mask as np_crf

os.environ['CUDA_VISIBLE_DEVICES'] = ''

HSZ = 100
WSZ = 30
S = '<GO>'
E = '<EOS>'
P = '<PAD>'
SPAN_TYPE="IOB2"
LOC = os.path.dirname(os.path.realpath(__file__))

@pytest.fixture
def label_vocab():
    LOC = os.path.dirname(os.path.realpath(__file__))
    vocab_loc = os.path.join(LOC, "test_data", "crf_vocab")
    return json.load(open(vocab_loc))

@pytest.fixture
def embeds():
    embeds = {}
    embeds['word'] = RandomInitVecModel(HSZ, {chr(i): i for i in range(100)})
    embeds['char'] = RandomInitVecModel(WSZ, {chr(i): i for i in range(100)})
    return embeds

@pytest.fixture
def save_file():
    base = 'del'
    path = os.path.join(LOC, base)
    yield path
    os.remove(path + "-char.vocab")
    os.remove(path + "-word.vocab")
    os.remove(os.path.join(LOC, 'checkpoint'))
    for file_name in os.listdir(LOC):
        if file_name.startswith(base):
            os.remove(os.path.join(LOC, file_name))

@pytest.fixture
def model(label_vocab, embeds):
    tf.reset_default_graph()
    model = create_model(
        label_vocab, embeds,
        crf=True, crf_mask=True, span_type=SPAN_TYPE,
        hsz=HSZ, cfiltsz=[3], wsz=WSZ,
        layers=2, rnntype="blstm"
    )
    model.create_loss()
    model.sess.run(tf.global_variables_initializer())
    return model

def test_mask_used(label_vocab, model):
    transition = model.sess.run(model.A)
    assert transition[label_vocab['O'], label_vocab[S]] == -1e4

def test_mask_is_transpose(label_vocab, model):
    transition = model.sess.run(model.mask)
    np_mask = np_crf(label_vocab, SPAN_TYPE, label_vocab[S], label_vocab[E], label_vocab[P])
    np.testing.assert_allclose(transition.T, np_mask)

def test_persists_save(model, save_file):
    model.save_using(tf.train.Saver())
    t1 = model.sess.run(model.A)
    model.save(save_file)
    m2 = RNNTaggerModel.load(save_file)
    t2 = model.sess.run(m2.A)
    np.testing.assert_allclose(t1, t2)

def test_skip_mask(label_vocab, embeds):
    tf.reset_default_graph()
    model = create_model(
        label_vocab, embeds,
        crf=True, span_type=SPAN_TYPE,
        hsz=HSZ, cfiltsz=[3], wsz=WSZ,
        layers=2, rnntype="blstm"
    )
    model.create_loss()
    model.sess.run(tf.global_variables_initializer())
    transition = model.sess.run(model.A)
    assert transition[label_vocab['O'], label_vocab[S]] != -1e4

def test_error_on_mask_and_no_span(label_vocab, embeds):
    tf.reset_default_graph()
    model = create_model(
        label_vocab, embeds,
        crf=True, crf_mask=True,
        hsz=HSZ, cfiltsz=[3], wsz=WSZ,
        layers=2, rnntype="blstm"
    )
    with pytest.raises(AssertionError):
        model.create_loss()

max_issues_repo_name: ZhenyueChin/baseline

max_forks_repo_name: ZhenyueChin/baseline


text:
import numpy as np
from .Solid import *
from .tools import *

class ConvexSolid(Solid):

    def __init__(self, name):

        super().__init__(name)

    ## Determine whether the ConvexSolid contains a given point
    def contains(self, p):

        return all([not(f.is_visible(p)) for f in self.faces])

    ## Add a vertex to the hull of this ConvexSolid
    def add_hull_vertex(self, vertex):

        if self.contains(vertex):
            return self

        cs = ConvexSolid(self.name)

        pv = np.asarray(vertex)
        visible_faces = [f for f in self.faces if f.is_visible(vertex)]
        horizon = Solid.boundary(visible_faces)
        horiz_length = len(horizon)

        for i in range(horiz_length):
            id1 = horizon[i]
            id2 = horizon[(i + 1) % horiz_length]
            v1 = self.vertices[id1]
            v2 = self.vertices[id2]
            adj_face = self.faces_with_edge(id1, id2)[1]
            if not adj_face.is_visible(pv, strict=False):
                cs.add_face([v1, v2, vertex])
            else:
                visible_faces.append(adj_face)
                new_vertices = []
                for id in adj_face.vertex_ids:
                    new_vertices.append(adj_face.get_coords(id))
                    if id == id2: new_vertices.append(pv)
                cs.add_face(new_vertices)

        for f in self.faces:
            if f not in visible_faces: cs.add_face(f.all_coords())

        self.overwrite(cs)
        return self

    ## Construct a tetrahedral ConvexSolid with 4 given (noncoplanar) vertices
    def tetrahedron(name, p1, p2, p3, p4):

        cs = ConvexSolid(name)

        pv1 = np.asarray(p1)
        pv2 = np.asarray(p2)
        pv3 = np.asarray(p3)
        pv4 = np.asarray(p4)
        cv = (pv1 + pv2 + pv3 + pv4)/4

        sv1 = pv1 - pv2
        sv2 = pv3 - pv2
        test_normal = np.cross(sv2, sv1)
        proj_sign = np.dot(test_normal, p4 - cv)

        if proj_sign > 0:
            cs.add_face([pv1, pv3, pv2])
            cs.add_face([pv2, pv3, pv4])
            cs.add_face([pv1, pv2, pv4])
            cs.add_face([pv1, pv4, pv3])
        else:
            cs.add_face([pv1, pv2, pv3])
            cs.add_face([pv2, pv4, pv3])
            cs.add_face([pv1, pv4, pv2])
            cs.add_face([pv1, pv3, pv4])

        return cs

    ## Construct a ConvexSolid as a convex hull of a given set of points
    def hull(name, pts, error=1e-7):

        pv1 = np.asarray(pts[0])
        pv2 = np.asarray(pts[1])
        pv3 = np.asarray(pts[2])
        cv = (pv1 + pv2 + pv3) / 3
        nv = np.cross(pv2-pv1, pv3-pv2)
        nv = nv / np.linalg.norm(nv)
        proj = 0
        id4 = 3
        pv4 = np.asarray(pts[id4])
        while abs(np.dot(nv, pv4-cv)) < error:
            id4 += 1
            pv4 = np.asarray(pts[id4])
        remaining_pts = [pts[i] for i in range(len(pts)) if i not in [0,1,2,id4]]

        cs = ConvexSolid.tetrahedron(name, pv1, pv2, pv3, pv4)
        for p in remaining_pts:
            cs.add_hull_vertex(p)

        return cs

max_stars_repo_name: Hand-and-Machine/polyhedra

max_issues_repo_name: Hand-and-Machine/polyhedra

max_forks_repo_name: Hand-and-Machine/polyhedra


text:
import torch
import itertools

import numpy as np
import scipy.sparse as sp

from typing import Any, Optional

import graphgallery as gg
from graphgallery import functional as gf

__all__ = ['asarray', 'index_to_mask',
           'repeat', 'get_length',
           'nx_graph_to_sparse_adj',
           'largest_indices', 'least_indices']

def asarray(x: Any, dtype: Optional[str] = None) -> np.ndarray:
    """Convert `x` to interger Numpy array.

    Parameters:
    ----------
    x: Tensor, Scipy sparse matrix,
        Numpy array-like, etc.

    Returns:
    ----------
    Integer Numpy array with dtype or `graphgallery.intx()`

    """
    if dtype is None:
        dtype = gg.intx()

    if gg.TF_ENABLED and gf.is_tensor(x, backend="tensorflow"):
        if x.dtype != dtype:
            import tensorflow as tf
            return tf.cast(x, dtype=dtype)
        else:
            return x

    if gf.is_tensor(x, backend="torch"):
        if x.dtype != dtype:
            return x.to(getattr(torch, dtype))
        else:
            return x

    if gg.is_intscalar(x):
        x = np.asarray([x], dtype=dtype)
    elif gg.is_listlike(x) or (isinstance(x, np.ndarray) and x.dtype != "O"):
        x = np.asarray(x, dtype=dtype)
    else:
        raise ValueError(
            f"Invalid input which should be either array-like or integer scalar, but got {type(x)}.")
    return x

def index_to_mask(indices: np.ndarray, shape: tuple) -> np.ndarray:
    mask = np.zeros(shape, dtype=gg.boolx())
    mask[indices] = True
    return mask

def repeat(src: Any, length: int = None) -> Any:
    if src == [] or src == ():
        return []
    if length is None:
        length = get_length(src)
    if any((gg.is_scalar(src), isinstance(src, str), src is None)):
        return list(itertools.repeat(src, length))
    if len(src) > length:
        return src[:length]
    if len(src) < length:
        return list(src) + list(itertools.repeat(src[-1], length - len(src)))
    return src

def get_length(obj: Any) -> int:
    if gg.is_iterable(obj):
        length = len(obj)
    else:
        length = 1
    return length

def nx_graph_to_sparse_adj(graph):
    num_nodes = graph.number_of_nodes()
    data = np.asarray(list(graph.edges().data('weight', default=1.0)))
    edge_index = data[:, :2].T.astype(np.int64)
    edge_weight = data[:, -1].T.astype(np.float32)
    adj_matrix = sp.csr_matrix((edge_weight, edge_index), shape=(num_nodes, num_nodes))
    return adj_matrix

def largest_indices(array: np.ndarray, n: int) -> tuple:
    """Returns the n largest indices from a numpy array.

    Parameters:
    ----------
        array {np.ndarray} -- data array
        n {int} -- number of elements to select

    Returns:
    --------
        tuple[np.ndarray] -- tuple of ndarray
        each ndarray is index
    """
    flat = array.ravel()
    indices = np.argpartition(flat, -n)[-n:]
    indices = indices[np.argsort(-flat[indices])]
    return np.unravel_index(indices, array.shape)

def least_indices(array: np.ndarray, n: int) -> tuple:
    """Returns the n least indices from a numpy array.

    Parameters:
    ----------
        array {np.ndarray} -- data array
        n {int} -- number of elements to select

    Returns:
    --------
        tuple[np.ndarray] -- tuple of ndarray
        each ndarray is index
    """
    flat = array.ravel()
    indices = np.argpartition(flat, n)[:n]
    indices = indices[np.argsort(flat[indices])]
    return np.unravel_index(indices, array.shape)

max_stars_repo_name: EdisonLeeeee/GraphGallery

max_issues_repo_name: EdisonLeeeee/GraphGallery

max_forks_repo_name: EdisonLeeeee/GraphGallery

reason: import numpy,import scipy


text:
# re-implement an idea of sparse hash mlp
# experiement: try multisize hash grid, and same size hashgrid

import torch
import torch.nn.functional as F

try:
    import tinycudann as tcnn
except:
    pass
from .tensoRF import TensorVMSplit, raw2alpha
import numpy as np

class SparseHashMLP(TensorVMSplit):
    """
    a drop-in replacement for for VM decompositon where we factor into sparse hashMLP instead
    """
    def __init__(self, *args, **kwargs):
        #self.hash_numfeatures = kwargs['hash_numfeatures']
        #self.sparsemlp_depth = [3,4,5,6,7,8,9,10,11]
        #self.sparsemlp_depth = [3,16,16,16,16]
        self.sparsemlp_depth = [3,4,8,16]
        self.sigma_activation = torch.relu
        print("Model: SparseHashMLP - big MLP from sparse hashgrid ")
        super(TensorVMSplit, self).__init__(*args, **kwargs)

    def init_svd_volume(self, res, device):
        #
        self.weight_hashs = []
        self.bias_hashs = []
        self.weight_inds = [0]
        self.bias_inds = [0]
        num_feature = 3 #define for x,y,z
        base_resolution = 1024
        log2_hashmap_size = 8
        for i in range(len(self.sparsemlp_depth)-1):
            #log2_size = self.hash_logsizes[i]
            #predict weight
            input_size = self.sparsemlp_depth[i]
            output_size = self.sparsemlp_depth[i+1]
            weights, bias = get_weight_and_bias(input_size, output_size, log2_hashmap_size=log2_hashmap_size, base_resolution=base_resolution, device=device)
            self.weight_inds += [self.weight_inds[-1] + len(weights)]
            self.bias_inds += [self.bias_inds[-1] + len(bias)]
            self.weight_hashs += weights
            self.bias_hashs += bias

        self.weight_hashs = torch.nn.ModuleList(self.weight_hashs)
        self.bias_hashs = torch.nn.ModuleList(self.bias_hashs)
        self.geo_feat_dim = self.sparsemlp_depth[-1] - 1

        # color network
        self.num_layers_color = 3
        self.hidden_dim_color = 64

        self.encoder_dir = tcnn.Encoding(
            n_input_dims=3,
            encoding_config={
                "otype": "SphericalHarmonics",
                "degree": 4,
            },
        )

        self.in_dim_color = self.encoder_dir.n_output_dims + self.geo_feat_dim
        self.color_net = tcnn.Network(
            n_input_dims=self.in_dim_color,
            n_output_dims=3,
            network_config={
                "otype": "CutlassMLP",
                "activation": "ReLU",
                "output_activation": "None",
                "n_neurons": self.hidden_dim_color,
                "n_hidden_layers": 3,
            },
        )
        """
        self.color_net = torch.nn.Sequential(
            torch.nn.Linear(self.encoder_dir.n_output_dims + self.geo_feat_dim, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, 3),
            torch.nn.Sigmoid()
        ).to(device)
        """
        """
        self.color_net = tcnn.Network(
            n_input_dims=self.in_dim_color,
            n_output_dims=3,
            network_config={
                "otype": "FullyFusedMLP",
                "activation": "ReLU",
                "output_activation": "None",
                "n_neurons": self.hidden_dim_color,
                "n_hidden_layers": self.num_layers_color - 1,
            },
        )
        """
        """
        self.encoder_dir = tcnn.Encoding(
            n_input_dims=3,
            encoding_config={
                "otype": "Frequency",
                "degree": 4,
            },
        )
        self.color_net = torch.nn.Sequential(
            torch.nn.Linear(self.encoder_dir.n_output_dims + self.geo_feat_dim, 128),
            torch.nn.ReLU(),
            torch.nn.Linear(128, 3),
            torch.nn.Sigmoid()
        ).to(device)
        """

    def get_optparam_groups(self, lr_init = 0.01, lr_basis = 0.001):
        # lr_init_spatialxyz = lr_init
        # lr_init_network = lr_basis
        grad_vars = [
            # HahsGrid parameters
            {'name': 'weight_hashs', 'params': self.weight_hashs.parameters(), 'lr': lr_init},
            {'name': 'bias_hashs', 'params': self.bias_hashs.parameters(), 'lr': lr_init},
            # Network parameter
            {'name': 'color_net', 'params': self.color_net.parameters(), 'lr': lr_basis, 'weight_decay': 1e-6}
        ]
        if isinstance(self.renderModule, torch.nn.Module):
            grad_vars += [{'params':self.renderModule.parameters(), 'lr':lr_init_network}]
        return grad_vars

    def feature2density(self, density_features):
        if self.fea2denseAct == "softplus":
            return F.softplus(density_features[...,0]+self.density_shift)
        elif self.fea2denseAct == "relu":
            val = density_features[...,0]
            val = F.relu(val)
            return val
        elif self.fea2denseAct == "abs":
            val = density_features[...,0]
            val = torch.abs(val)
            return val
        else:
            raise NotImplementedError()

    def compute_densityfeature(self, xyz_sampled):
        """
        @params xyz_sampled scale in [-1,1] #[num_ray, 3]
        @return sigma_feature #[num_ray]
        """
        xyz_sampled = (xyz_sampled + 1.0) / 2.0 #scale to [0,1]
        x = xyz_sampled # batch, 3
        for i in range(len(self.sparsemlp_depth)-1):
            num_ray = xyz_sampled.shape[0]
            input_size = self.sparsemlp_depth[i]
            output_size = self.sparsemlp_depth[i+1]
            w = get_hash(self.weight_hashs[self.weight_inds[i]:self.weight_inds[i+1]],xyz_sampled).type(xyz_sampled.dtype)
            w = w.view(xyz_sampled.shape[0],output_size, input_size).type(xyz_sampled.dtype)
            b = get_hash(self.bias_hashs[self.bias_inds[i]:self.bias_inds[i+1]],xyz_sampled)
            b = b.view(xyz_sampled.shape[0], output_size, 1)
            x = (w @ x[...,None]) + b
            x = x[..., 0]
            x = self.sigma_activation(x)
        return x.type(xyz_sampled.dtype)

    def compute_appfeature(self, xyz_sampled, density_features, viewdirs):
        #xyz_sampled = (xyz_sampled + 1.0) / 2.0 #scale to [0,1]
        #xyz_sampled = xyz_sampled[None] #[B, num_ray, 3]
        geo_feat = density_features[...,1:]  #[B, num_ray, 15]
        d = self.encoder_dir(viewdirs)
        h = torch.cat([d, geo_feat], dim=-1)
        h = self.color_net(h)
        color = torch.sigmoid(h)
        return color.type(viewdirs.dtype)

    @torch.no_grad()
    def upsample_volume_grid(self, res_target):
        # Hashgrid is not support UPsample size
        self.update_stepSize(res_target)
        print(f'upsamping to {res_target}')

    @torch.no_grad()
    def shrink(self, new_aabb):
        print("====> shrinking ...")
        xyz_min, xyz_max = new_aabb
        t_l, b_r = (xyz_min - self.aabb[0]) / self.units, (xyz_max - self.aabb[0]) / self.units
        t_l, b_r = torch.round(torch.round(t_l)).long(), torch.round(b_r).long() + 1
        b_r = torch.stack([b_r, self.gridSize]).amin(0)

        if not torch.all(self.alphaMask.gridSize == self.gridSize):
            t_l_r, b_r_r = t_l / (self.gridSize-1), (b_r-1) / (self.gridSize-1)
            correct_aabb = torch.zeros_like(new_aabb)
            correct_aabb[0] = (1-t_l_r)*self.aabb[0] + t_l_r*self.aabb[1]
            correct_aabb[1] = (1-b_r_r)*self.aabb[0] + b_r_r*self.aabb[1]
            print("aabb", new_aabb, "\ncorrect aabb", correct_aabb)
            new_aabb = correct_aabb

        newSize = b_r - t_l
        self.aabb = new_aabb
        self.update_stepSize((newSize[0], newSize[1], newSize[2]))

    def forward(self, rays_chunk, white_bg=True, is_train=False, ndc_ray=False, N_samples=-1):

        # sample points
        viewdirs = rays_chunk[:, 3:6]
        if ndc_ray:
            xyz_sampled, z_vals, ray_valid = self.sample_ray_ndc(rays_chunk[:, :3], viewdirs, is_train=is_train,N_samples=N_samples)
            dists = torch.cat((z_vals[:, 1:] - z_vals[:, :-1], torch.zeros_like(z_vals[:, :1])), dim=-1)
            rays_norm = torch.norm(viewdirs, dim=-1, keepdim=True)
            dists = dists * rays_norm
            viewdirs = viewdirs / rays_norm
        else:
            xyz_sampled, z_vals, ray_valid = self.sample_ray(rays_chunk[:, :3], viewdirs, is_train=is_train,N_samples=N_samples)
            dists = torch.cat((z_vals[:, 1:] - z_vals[:, :-1], torch.zeros_like(z_vals[:, :1])), dim=-1)
        viewdirs = viewdirs.view(-1, 1, 3).expand(xyz_sampled.shape)

        if self.alphaMask is not None:
            alphas = self.alphaMask.sample_alpha(xyz_sampled[ray_valid])
            alpha_mask = alphas > 0
            ray_invalid = ~ray_valid
            ray_invalid[ray_valid] |= (~alpha_mask)
            ray_valid = ~ray_invalid

        sigma = torch.zeros(xyz_sampled.shape[:-1], device=xyz_sampled.device)
        rgb = torch.zeros((*xyz_sampled.shape[:2], 3), device=xyz_sampled.device)

        if ray_valid.any():
            xyz_sampled = self.normalize_coord(xyz_sampled)
            sigma_feature = self.compute_densityfeature(xyz_sampled[ray_valid])
            validsigma = self.feature2density(sigma_feature)

            sigma[ray_valid] = validsigma
        #pdb.set_trace()

        alpha, weight, bg_weight = raw2alpha(sigma, dists * self.distance_scale)

        #app_mask = weight > self.rayMarch_weight_thres
        app_mask_ray = weight[ray_valid]  > self.rayMarch_weight_thres

        if app_mask_ray.any():
            xyz_r = xyz_sampled[ray_valid][app_mask_ray]
            viewdir_r = viewdirs[ray_valid][app_mask_ray]
            app_features = self.compute_appfeature(xyz_r, sigma_feature[app_mask_ray], viewdir_r)
            valid_rgbs = self.renderModule(xyz_r, viewdir_r, app_features)
            temp_rgb =  torch.zeros_like(rgb[ray_valid])
            temp_rgb[app_mask_ray] = valid_rgbs
            rgb[ray_valid] = temp_rgb

        acc_map = torch.sum(weight, -1)
        rgb_map = torch.sum(weight[..., None] * rgb, -2)

        if white_bg or (is_train and torch.rand((1,))<0.5):
            rgb_map = rgb_map + (1. - acc_map[..., None])

        rgb_map = rgb_map.clamp(0,1)

        with torch.no_grad():
            depth_map = torch.sum(weight * z_vals, -1)
            depth_map = depth_map + (1. - acc_map) * rays_chunk[..., -1]
        return rgb_map, depth_map # rgb, sigma, alpha, weight, bg_weight

    def density_L1(self):
        raise NotImplementedError("HashGrid3D shouldn't call this")

    def TV_loss_density(self, reg):
       raise NotImplementedError("HashGrid3D shouldn't call this")

    def TV_loss_app(self, reg):
       raise NotImplementedError("HashGrid3D shouldn't call this")

    def vectorDiffs(self, vector_comps):
        raise NotImplementedError("HashGrid3D shouldn't call this")

    def vector_comp_diffs(self):
        raise NotImplementedError("HashGrid3D shouldn't call this")

    @torch.no_grad()
    def up_sampling_VM(self, plane_coef, line_coef, res_target):
        raise NotImplementedError("HashGrid3D shouldn't call this")

def get_weight_and_bias(input_size, output_size, log2_hashmap_size=19, base_resolution=1024, device='cpu'):
    """
    generating weight and bias
    weight output size can be reshape to [input_size, output_size]
    bias output size can be reshape to [output_size, 1]
    """
    weight = []
    bias = []

    weight_pattern = get_weight_pattern(input_size * output_size)
    bias_pattern = get_weight_pattern(output_size)
    weight = get_weight_by_pattern(weight_pattern, log2_hashmap_size, base_resolution, device)
    bias = get_weight_by_pattern(bias_pattern, log2_hashmap_size, base_resolution, device)
    return weight, bias

def get_weight_by_pattern(patterns, log2_hashmap_size, base_resolution, device):
    weights = []
    for p in patterns:
        weights.append(tcnn.Encoding(
            n_input_dims=3,
            encoding_config={
               "otype": "HashGrid",
               "n_levels": p[0],
               "n_features_per_level": p[1],
               "log2_hashmap_size": log2_hashmap_size, #every table depth has the same size
               "base_resolution": base_resolution,
               "per_level_scale": 1,
            },
        ).to(device))
    return weights

def get_weight_pattern(weight_count):
    MAX_LEVEL = 32
    LATENT_SHAPE = [8,4,2,1]
    pattern = []
    i = MAX_LEVEL
    while weight_count > 0 and i > 0:
        jcnt = 0
        while jcnt < len(LATENT_SHAPE):
            j = LATENT_SHAPE[jcnt]
            if weight_count >= i*j:
                pattern.append([i,j])
                weight_count -= i*j
                continue
            jcnt += 1
        i -= 1
    return pattern

def get_hash(encoder_list,xyz_sampled):
    output = []
    num_ray = xyz_sampled.shape[0]
    for encoder in encoder_list:
        data = encoder(xyz_sampled)
        output.append(data.view(num_ray, -1))
    output = torch.cat(output, dim=-1)
    output = output.view(num_ray, -1)
    return output

"""
def get_hashgrid_pattern(weight_count)
    pattern = []
    while weight_count > 0:
        for b in [8,4,2,1]:
            if weight_count > b:
                pattenr.append(b)
                weight_count -= b
                break
    return pattern
"""

max_stars_repo_name: sWizad/TensoRF

max_issues_repo_name: sWizad/TensoRF

max_forks_repo_name: sWizad/TensoRF


text:
from __future__ import print_function
from quadruped_env import QuadrupedEnvironment
from ddpg import OUNoise, DDPG
import numpy as np

env = QuadrupedEnvironment()
state_shape = env.state_shape
action_shape = env.action_shape
agent = DDPG(state_shape,action_shape,batch_size=128,gamma=0.995,tau=0.001,
                                        actor_lr=0.0001, critic_lr=0.001, use_layer_norm=True)
print('DDPG agent configured')
max_episode = 10000
tot_rewards = []
print('env reset')
observation, done = env.reset()
action = agent.act(observation)
observation, reward, done = env.step(action)
noise_sigma = 0.15
save_cutoff = 1
cutoff_count = 0
save_count = 0
curr_highest_eps_reward = -1000.0
for i in xrange(max_episode):
    if i % 100 == 0 and noise_sigma>0.03:
        agent.noise = OUNoise(agent.nb_actions,sigma=noise_sigma)
        noise_sigma /= 2.0
    step_num = 0
    while done == False:
        step_num += 1
        action = agent.step(observation, reward, done)
        observation, reward, done = env.step(action)
        print('reward:',reward,'episode:', i, 'step:',step_num,'curr high eps reward:',curr_highest_eps_reward, 'saved:',save_count, 'cutoff count:', cutoff_count)
    action, eps_reward = agent.step(observation, reward, done)
    tot_rewards.append(eps_reward)
    if eps_reward > curr_highest_eps_reward:
        cutoff_count += 1
        curr_highest_eps_reward = eps_reward
    if cutoff_count >= save_cutoff:
        save_count += 1
        print('saving_model at episode:',i)
        agent.save_model()
        agent.save_memory()
        cutoff_count = 0
    observation, done = env.reset()
np.save('eps_rewards',tot_rewards)

import matplotlib.pyplot as plt
plt.plot(tot_rewards)


text:
import numpy as np
import matplotlib.pyplot as plt

from sklearn.datasets import load_iris
from sklearn.cross_validation import train_test_split

iris = load_iris()
X, y = iris.data, iris.target

print("Dataset size: %d  number of features: %d  number of classes: %d"
      % (X.shape[0], X.shape[1], len(np.unique(y))))

X_train, X_test, y_train, y_test = train_test_split(X, y)

plt.scatter(X_train[:, 0], X_train[:, 1], c=y_train)
plt.figure()
plt.scatter(X_train[:, 2], X_train[:, 3], c=y_train)


text:
import torch
import numpy as np

from queue import PriorityQueue

# simple bandit
class ThresholdBandit():
    def __init__(self):
        self.N = 100
        self.points = np.arange(self.N) / self.N
        self.active = np.zeros(self.N, dtype=np.bool)
        self.timer = 1

        self.cnt_pulls = np.zeros(self.N) + 1e-15
        self.sum_pulls = np.zeros(self.N)

        self.delayed_mean = np.zeros(self.N)
        self.gamma = 0.99

    def step(self, threshold, reward):
        self.timer += 1
        pos = int(threshold * self.N + 0.5)
        self.cnt_pulls[pos] += 1
        self.sum_pulls[pos] += reward
        self.delayed_mean[pos] = self.gamma * self.delayed_mean[pos] + (1 - self.gamma) * reward

        self.exploration_factor = np.sqrt(2 * np.log(self.timer) / (self.cnt_pulls + 1))
        self.value = self.delayed_mean + self.exploration_factor

    def choose_threshold(self):
        threshold_pos = np.argmax(self.value)
        return threshold_pos / self.N


text:
import c3py # import the python interface to the c3 library
import numpy as np
import matplotlib.pyplot as plt

np.random.seed(10)

def func(x):
    return x[:, 0]**2 + x[:, 1]**2 + x[:, 0] * x[:, 1] + np.sin(np.sum(x, axis=1)) \
        + x[:,0] * x[:,-1] + x[:, -1]**2.0 * np.cos(x[:, 0])

LB = -1.0                              # lower bounds of features
UB = 1.0                               # upper bounds of features
DIM = 20                               # number of features

Ntrain = 100
X = np.random.rand(Ntrain, 2)
Y = func(X)

Ntest = 1000
Xtest = np.random.rand(Ntest, 2)
Ytest = func(Xtest)

def build_ft_regress(xdata, ydata, nparam=2, init_rank=5, adaptrank=1, verbose=0):
    dim = xdata.shape[1]
    ndata = xdata.shape[0]

    ft = c3py.FunctionTrain(dim)
    for ii in range(dim):
        ft.set_dim_opts(ii, "legendre", LB, UB, nparam)

    ranks = [init_rank]*(dim+1)
    ranks[0] = 1
    ranks[dim] = 1
    ft.set_ranks(ranks)
    ft.build_data_model(ndata, xdata, ydata,
                        alg="AIO", obj="LS", adaptrank=adaptrank,
                        kickrank=1, roundtol=1e-10, verbose=verbose,
                        seed=10)
    return ft

print("X shape ", X.shape)
print("Y shape ", Y.shape)
ft_regress = build_ft_regress(X, Y, nparam=5, init_rank=1, adaptrank=1, verbose=0)

print("Xtest.shape ", Xtest.shape)
print("Ytest.shape ", Ytest.shape)
ft_evals = ft_regress.eval(Xtest)

err = np.linalg.norm(ft_evals - Ytest) / np.linalg.norm(Ytest)

print("Relmse = ", err)
# plt.figure()
# plt.plot(ft_evals, Ytest, 'o')
# plt.show()