text:
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Dec 5, 2018
@author: Quentin Lutz <qlutz@enst.fr>
Nathan de Lara <ndelara@enst.fr>
"""
import warnings
from csv import reader
from typing import Optional, List, Tuple, Union
from xml.etree import ElementTree

import numpy as np
from scipy import sparse

from sknetwork.utils import Bunch
from sknetwork.utils.format import directed2undirected

def load_edge_list(file: str, directed: bool = False, bipartite: bool = False, weighted: Optional[bool] = None,
                   named: Optional[bool] = None, comment: str = '%#', delimiter: str = None, reindex: bool = True,
                   fast_format: bool = True) -> Bunch:
    """Parse Tabulation-Separated, Comma-Separated or Space-Separated (or other) Values datasets in the form of
    edge lists.

    Parameters
    ----------
    file : str
        The path to the dataset in TSV format
    directed : bool
        If ``True``, considers the graph as directed.
    bipartite : bool
        If ``True``, returns a biadjacency matrix of shape (n1, n2).
    weighted : Optional[bool]
        Retrieves the weights in the third field of the file. None makes a guess based on the first lines.
    named : Optional[bool]
        Retrieves the names given to the nodes and renumbers them. Returns an additional array. None makes a guess
        based on the first lines.
    comment : str
        Set of characters denoting lines to ignore.
    delimiter : str
        delimiter used in the file. None makes a guess
    reindex : bool
        If True and the graph nodes have numeric values, the size of the returned adjacency will be determined by the
        maximum of those values. Does not work for bipartite graphs.
    fast_format : bool
        If True, assumes that the file is well-formatted:

        * no comments except for the header
        * only 2 or 3 columns
        * only int or float values

    Returns
    -------
    graph: :class:`Bunch`
    """
    header_len, guess_delimiter, guess_weighted, guess_named, guess_string_present, guess_type = scan_header(file,
                                                                                comment)

    if weighted is None:
        weighted = guess_weighted
    if named is None:
        named = guess_named
    if delimiter is None:
        delimiter = guess_delimiter

    with open(file, 'r', encoding='utf-8') as f:
        for i in range(header_len):
            f.readline()
        if fast_format and not guess_string_present:
            # fromfile raises a DeprecationWarning on fail. This should be changed to ValueError in the future.
            warnings.filterwarnings("error")
            try:
                parsed = np.fromfile(f, sep=guess_delimiter, dtype=guess_type)
            except (DeprecationWarning, ValueError):
                raise ValueError('File not suitable for fast parsing. Set fast_format to False.')
            warnings.filterwarnings("default")
            n_entries = len(parsed)
            if weighted:
                parsed.resize((n_entries//3, 3))
                row, col, data = parsed[:, 0], parsed[:, 1], parsed[:, 2]
            else:
                parsed.resize((n_entries//2, 2))
                row, col = parsed[:, 0], parsed[:, 1]
                data = np.ones(row.shape[0], dtype=bool)
        else:
            row, col, data = [], [], []
            csv_reader = reader(f, delimiter=delimiter)
            for line in csv_reader:
                if line[0] not in comment:
                    if named:
                        row.append(line[0])
                        col.append(line[1])
                    else:
                        row.append(int(line[0]))
                        col.append(int(line[1]))
                    if weighted:
                        data.append(float(line[2]))
            row, col, data = np.array(row), np.array(col), np.array(data)

    return from_edge_list(row=row, col=col, data=data, directed=directed, bipartite=bipartite,
                          reindex=reindex, named=named)

def convert_edge_list(edge_list: Union[np.ndarray, List[Tuple], List[List]], directed: bool = False,
                      bipartite: bool = False, reindex: bool = True, named: Optional[bool] = None) -> Bunch:
    """Turn an edge list into a :class:`Bunch`.

    Parameters
    ----------
    edge_list : Union[np.ndarray, List[Tuple], List[List]]
        The edge list to convert, given as a NumPy array of size (n, 2) or (n, 3) or a list of either lists or tuples of
        length 2 or 3.
    directed : bool
        If ``True``, considers the graph as directed.
    bipartite : bool
        If ``True``, returns a biadjacency matrix of shape (n1, n2).
    reindex : bool
        If True and the graph nodes have numeric values, the size of the returned adjacency will be determined by the
        maximum of those values. Does not work for bipartite graphs.
    named : Optional[bool]
        Retrieves the names given to the nodes and renumbers them. Returns an additional array. None makes a guess
        based on the first lines.

    Returns
    -------
    graph: :class:`Bunch`
    """
    if isinstance(edge_list, list):
        if edge_list:
            if isinstance(edge_list[0], tuple) or isinstance(edge_list[0], list):
                edge_list = np.array(edge_list)
    if isinstance(edge_list, np.ndarray):
        if edge_list.ndim == 2:
            if edge_list.shape[1] == 2:
                row, col, data = edge_list[:, 0], edge_list[:, 1], np.array([])
            elif edge_list.shape[1] == 3:
                row, col, data = edge_list[:, 0], edge_list[:, 1], edge_list[:, 2].astype(float)
            else:
                raise ValueError('Edges must be given as couples or triplets.')
        else:
            raise ValueError('Too many dimensions.')
    else:
        raise TypeError('Edge lists must be given as NumPy arrays or lists of lists or lists of tuples.')
    return from_edge_list(row=row, col=col, data=data, directed=directed, bipartite=bipartite,
                          reindex=reindex, named=named)

def from_edge_list(row: np.ndarray, col: np.ndarray, data: np.ndarray, directed: bool = False, bipartite: bool = False,
                   reindex: bool = True, named: Optional[bool] = None) -> Bunch:
    """Turn an edge list given as a triplet of NumPy arrays into a :class:`Bunch`.

    Parameters
    ----------
    row : np.ndarray
        The array of sources in the graph.
    col : np.ndarray
        The array of targets in the graph.
    data : np.ndarray
        The array of weights in the graph. Pass an empty array for unweighted graphs.
    directed : bool
        If ``True``, considers the graph as directed.
    bipartite : bool
        If ``True``, returns a biadjacency matrix of shape (n1, n2).
    reindex : bool
        If True and the graph nodes have numeric values, the size of the returned adjacency will be determined by the
        maximum of those values. Does not work for bipartite graphs.
    named : Optional[bool]
        Retrieves the names given to the nodes and renumbers them. Returns an additional array. None makes a guess
        based on the first lines.

    Returns
    -------
    graph: :class:`Bunch`
    """
    reindexed = False
    if named is None:
        named = (row.dtype != int) or (col.dtype != int)
    weighted = bool(len(data))
    n_edges = len(row)
    graph = Bunch()
    if bipartite:
        names_row, row = np.unique(row, return_inverse=True)
        names_col, col = np.unique(col, return_inverse=True)
        if not reindex:
            n_row = names_row.max() + 1
            n_col = names_col.max() + 1
        else:
            n_row = len(names_row)
            n_col = len(names_col)
        if not weighted:
            data = np.ones(n_edges, dtype=bool)
        biadjacency = sparse.csr_matrix((data, (row, col)), shape=(n_row, n_col))
        graph.biadjacency = biadjacency
        if named or reindex:
            graph.names = names_row
            graph.names_row = names_row
            graph.names_col = names_col
    else:
        nodes = np.concatenate((row, col), axis=None)
        names, new_nodes = np.unique(nodes, return_inverse=True)
        if not reindex:
            n_nodes = names.max() + 1
        else:
            n_nodes = len(names)
        if named:
            row = new_nodes[:n_edges]
            col = new_nodes[n_edges:]
        else:
            should_reindex = not (names[0] == 0 and names[-1] == n_nodes - 1)
            if should_reindex and reindex:
                reindexed = True
                row = new_nodes[:n_edges]
                col = new_nodes[n_edges:]
        if not weighted:
            data = np.ones(n_edges, dtype=bool)
        adjacency = sparse.csr_matrix((data, (row, col)), shape=(n_nodes, n_nodes))
        if not directed:
            adjacency = directed2undirected(adjacency, weighted=weighted)
        graph.adjacency = adjacency
        if named or reindexed:
            graph.names = names

    return graph

def load_adjacency_list(file: str, bipartite: bool = False, comment: str = '%#',
                        delimiter: str = None, ) -> Bunch:
    """Parse Tabulation-Separated, Comma-Separated or Space-Separated (or other) Values datasets in the form of
    adjacency lists.

    Parameters
    ----------
    file : str
        The path to the dataset in TSV format
    bipartite : bool
        If ``True``, returns a biadjacency matrix of shape (n1, n2).
    comment : str
        Set of characters denoting lines to ignore.
    delimiter : str
        delimiter used in the file. None makes a guess

    Returns
    -------
    graph: :class:`Bunch`
    """
    header_len, guess_delimiter, _, _, _, _ = scan_header(file, comment)
    if delimiter is None:
        delimiter = guess_delimiter
    indptr, indices = [0], []
    with open(file, 'r', encoding='utf-8') as f:
        for i in range(header_len):
            f.readline()
        for row in f:
            neighbors = [int(el) for el in row.split(delimiter)]
            indices += neighbors
            indptr.append(indptr[-1] + len(neighbors))
    indices = np.array(indices)
    n_rows = len(indptr) - 1
    min_index = indices.min()
    n_cols = indices.max() + 1 - min_index
    indices -= min_index
    graph = Bunch()
    if not bipartite:
        max_dim = max(n_rows, n_cols)
        new_indptr = np.full(max_dim + 1, indptr[-1])
        new_indptr[:len(indptr)] = indptr
        graph.adjacency = sparse.csr_matrix((np.ones_like(indices, dtype=bool), indices, new_indptr),
                                            shape=(max_dim, max_dim))
    else:
        indptr = np.array(indptr)
        graph.biadjacency = sparse.csr_matrix((np.ones_like(indices, dtype=bool), indices, indptr),
                                              shape=(n_rows, n_cols))
    return graph

def scan_header(file: str, comment: str):
    """Infer some properties of the graph in a TSV file from the first few lines."""
    header_len = -1
    possible_delimiters = ['\t', ',', ' ']
    del_count = np.zeros(3, dtype=int)
    lines = []
    row = comment
    with open(file, 'r', encoding='utf-8') as f:
        while row[0] in comment:
            row = f.readline()
            header_len += 1
        for line in range(3):
            for i, poss_del in enumerate(possible_delimiters):
                if poss_del in row:
                    del_count[i] += 1
            lines.append(row.rstrip())
            row = f.readline()
        lines = [line for line in lines if line != '']
        guess_delimiter = possible_delimiters[int(np.argmax(del_count))]
        guess_weighted = bool(min([line.count(guess_delimiter) for line in lines]) - 1)
        guess_named = not all([all([el.strip().isdigit() for el in line.split(guess_delimiter)][0:2])
                               for line in lines])
        guess_string_present = not all([all([isnumber(el.strip()) for el in line.split(guess_delimiter)][0:2])
                                        for line in lines])
        if not guess_named:
            guess_type = np.int32
        else:
            guess_type = np.float32
    return header_len, guess_delimiter, guess_weighted, guess_named, guess_string_present, guess_type

def load_labels(file: str) -> np.ndarray:
    """Parser for files with a single entry on each row.

    Parameters
    ----------
    file : str
        The path to the dataset

    Returns
    -------
    labels: np.ndarray
        Labels.
    """
    rows = []
    with open(file, 'r', encoding='utf-8') as f:
        for row in f:
            rows.append(row.strip())
    return np.array(rows)

def load_header(file: str):
    """Check if the graph is directed, bipartite, weighted."""
    directed, bipartite, weighted = False, False, True
    with open(file, 'r', encoding='utf-8') as f:
        row = f.readline()
        if 'bip' in row:
            bipartite = True
        if 'unweighted' in row:
            weighted = False
        if 'asym' in row:
            directed = True
    return directed, bipartite, weighted

def load_metadata(file: str, delimiter: str = ': ') -> Bunch:
    """Extract metadata from the file."""
    metadata = Bunch()
    with open(file, 'r', encoding='utf-8') as f:
        for row in f:
            parts = row.split(delimiter)
            key, value = parts[0], ': '.join(parts[1:]).strip('\n')
            metadata[key] = value
    return metadata

def load_graphml(file: str, weight_key: str = 'weight', max_string_size: int = 512) -> Bunch:
    """Parse GraphML datasets.

    Hyperedges and nested graphs are not supported.

    Parameters
    ----------
    file: str
        The path to the dataset
    weight_key: str
        The key to be used as a value for edge weights
    max_string_size: int
        The maximum size for string features of the data

    Returns
    -------
    data: :class:`Bunch`
        The dataset in a bunch with the adjacency as a CSR matrix.
    """
    # see http://graphml.graphdrawing.org/primer/graphml-primer.html
    # and http://graphml.graphdrawing.org/specification/dtd.html#top
    tree = ElementTree.parse(file)
    n_nodes = 0
    n_edges = 0
    symmetrize = None
    naming_nodes = True
    default_weight = 1
    weight_type = bool
    weight_id = None
    # indices in the graph tree
    node_indices = []
    edge_indices = []
    data = Bunch()
    graph = None
    file_description = None
    attribute_descriptions = Bunch()
    attribute_descriptions.node = Bunch()
    attribute_descriptions.edge = Bunch()
    keys = {}
    for file_element in tree.getroot():
        if file_element.tag.endswith('graph'):
            graph = file_element
            symmetrize = (graph.attrib['edgedefault'] == 'undirected')
            for index, element in enumerate(graph):
                if element.tag.endswith('node'):
                    node_indices.append(index)
                    n_nodes += 1
                elif element.tag.endswith('edge'):
                    edge_indices.append(index)
                    if 'directed' in element.attrib:
                        if element.attrib['directed'] == 'true':
                            n_edges += 1
                        else:
                            n_edges += 2
                    elif symmetrize:
                        n_edges += 2
                    else:
                        n_edges += 1
            if 'parse.nodeids' in graph.attrib:
                naming_nodes = not (graph.attrib['parse.nodeids'] == 'canonical')
    for file_element in tree.getroot():
        if file_element.tag.endswith('key'):
            attribute_name = file_element.attrib['attr.name']
            attribute_type = java_type_to_python_type(file_element.attrib['attr.type'])
            if attribute_name == weight_key:
                weight_type = java_type_to_python_type(file_element.attrib['attr.type'])
                weight_id = file_element.attrib['id']
                for key_element in file_element:
                    if key_element.tag == 'default':
                        default_weight = attribute_type(key_element.text)
            else:
                default_value = None
                if file_element.attrib['for'] == 'node':
                    size = n_nodes
                    if 'node_attribute' not in data:
                        data.node_attribute = Bunch()
                    for key_element in file_element:
                        if key_element.tag.endswith('desc'):
                            attribute_descriptions.node[attribute_name] = key_element.text
                        elif key_element.tag.endswith('default'):
                            default_value = attribute_type(key_element.text)
                    if attribute_type == str:
                        local_type = '<U' + str(max_string_size)
                    else:
                        local_type = attribute_type
                    if default_value:
                        data.node_attribute[attribute_name] = np.full(size, default_value, dtype=local_type)
                    else:
                        data.node_attribute[attribute_name] = np.zeros(size, dtype=local_type)
                elif file_element.attrib['for'] == 'edge':
                    size = n_edges
                    if 'edge_attribute' not in data:
                        data.edge_attribute = Bunch()
                    for key_element in file_element:
                        if key_element.tag.endswith('desc'):
                            attribute_descriptions.edge[attribute_name] = key_element.text
                        elif key_element.tag.endswith('default'):
                            default_value = attribute_type(key_element.text)
                    if attribute_type == str:
                        local_type = '<U' + str(max_string_size)
                    else:
                        local_type = attribute_type
                    if default_value:
                        data.edge_attribute[attribute_name] = np.full(size, default_value, dtype=local_type)
                    else:
                        data.edge_attribute[attribute_name] = np.zeros(size, dtype=local_type)
                keys[file_element.attrib['id']] = [attribute_name, attribute_type]
        elif file_element.tag.endswith('desc'):
            file_description = file_element.text
    if file_description or attribute_descriptions.node or attribute_descriptions.edge:
        data.meta = Bunch()
        if file_description:
            data.meta['description'] = file_description
        if attribute_descriptions.node or attribute_descriptions.edge:
            data.meta['attributes'] = attribute_descriptions
    if graph is not None:
        row = np.zeros(n_edges, dtype=int)
        col = np.zeros(n_edges, dtype=int)
        dat = np.full(n_edges, default_weight, dtype=weight_type)
        data.names = None
        if naming_nodes:
            data.names = np.zeros(n_nodes, dtype='<U512')

        node_map = {}
        # deal with nodes first
        for number, index in enumerate(node_indices):
            node = graph[index]
            if naming_nodes:
                name = node.attrib['id']
                data.names[number] = name
                node_map[name] = number
            for node_attribute in node:
                if node_attribute.tag.endswith('data'):
                    data.node_attribute[keys[node_attribute.attrib['key']][0]][number] = \
                        keys[node_attribute.attrib['key']][1](node_attribute.text)
        # deal with edges
        edge_index = -1
        for index in edge_indices:
            edge_index += 1
            duplicate = False
            edge = graph[index]
            if naming_nodes:
                node1 = node_map[edge.attrib['source']]
                node2 = node_map[edge.attrib['target']]
            else:
                node1 = int(edge.attrib['source'][1:])
                node2 = int(edge.attrib['target'][1:])
            row[edge_index] = node1
            col[edge_index] = node2
            for edge_attribute in edge:
                if edge_attribute.tag.endswith('data'):
                    if edge_attribute.attrib['key'] == weight_id:
                        dat[edge_index] = weight_type(edge_attribute.text)
                    else:
                        data.edge_attribute[keys[edge_attribute.attrib['key']][0]][edge_index] = \
                            keys[edge_attribute.attrib['key']][1](edge_attribute.text)
            if 'directed' in edge.attrib:
                if edge.attrib['directed'] != 'true':
                    duplicate = True
            elif symmetrize:
                duplicate = True
            if duplicate:
                edge_index += 1
                row[edge_index] = node2
                col[edge_index] = node1
                for edge_attribute in edge:
                    if edge_attribute.tag.endswith('data'):
                        if edge_attribute.attrib['key'] == weight_id:
                            dat[edge_index] = weight_type(edge_attribute.text)
                        else:
                            data.edge_attribute[keys[edge_attribute.attrib['key']][0]][edge_index] = \
                                keys[edge_attribute.attrib['key']][1](edge_attribute.text)
        data.adjacency = sparse.csr_matrix((dat, (row, col)), shape=(n_nodes, n_nodes))
        if data.names is None:
            data.pop('names')
        return data
    else:
        raise ValueError(f'No graph defined in {file}.')

def java_type_to_python_type(value: str) -> type:
    if value == 'boolean':
        return bool
    elif value == 'int':
        return int
    elif value == 'string':
        return str
    elif value in ('long', 'float', 'double'):
        return float

def isnumber(s):
    try:
        float(s)
        return True
    except ValueError:
        return False

max_stars_repo_name: altana-tech/scikit-network

max_issues_repo_name: altana-tech/scikit-network

max_forks_repo_name: altana-tech/scikit-network


text:
from .DATA import *
from PIL import Image
from numpy import array

def get_normalize_size_shape_method(img, max_length):
    [ height, width, channels ] = img.shape
    if height >= width:
        longerSize = height
        shorterSize = width
    else:
        longerSize = width
        shorterSize = height

    scale = float(max_length) / float(longerSize)

    outputHeight = int(round(height*scale))
    outputWidth = int(round(width*scale))
    return outputHeight, outputWidth

def cpu_normalize_image(img, max_length):
    outputHeight, outputWidth = get_normalize_size_shape_method(img, max_length)
    outputImg = Image.fromarray(img)
    outputImg = outputImg.resize((outputWidth, outputHeight), Image.ANTIALIAS)
    outputImg = array(outputImg)
    return outputImg

ext: py

max_stars_repo_name: pnbao/deep-photo-enhance

reason: from numpy


text:
import datetime
import json

import matplotlib.gridspec
import matplotlib.lines
import matplotlib.patches
import matplotlib.pyplot as plt
import matplotlib.style
import numpy as np
import pandas as pd

from ..plot import DEPARTMENT_GRAND_EST_COLOR, plot_int

data_source = "combined_icubam_public"

def plot(data):
  fig, (ax1, ax2) = plt.subplots(2, figsize=(20, 10), sharex=True)
  date_range_idx = np.arange(len(data.date.unique()))
  for dept, dg in data.groupby("department"):
    dg = dg.sort_values(by="date")
    plot_int(
      x=date_range_idx,
      y=dg.n_hospitalised_patients / dg.department_pop * 100e3,
      ax=ax1,
      marker=None,
      color=DEPARTMENT_GRAND_EST_COLOR[dept],
      label=dept,
      lw=2,
    )
    y1 = dg.n_icu_patients_icubam / dg.department_pop * 100e3
    plot_int(
      x=date_range_idx,
      y=y1,
      ax=ax2,
      marker=None,
      color=DEPARTMENT_GRAND_EST_COLOR[dept],
      ls="dashed",
      lw=2,
    )
    y2 = dg.n_icu_patients_public / dg.department_pop * 100e3
    plot_int(
      x=date_range_idx,
      y=y2,
      ax=ax2,
      marker=None,
      color=DEPARTMENT_GRAND_EST_COLOR[dept],
      ls="solid",
      label=dept,
      lw=2,
    )
    ax2.fill_between(
      x=date_range_idx,
      y1=y1,
      y2=y2,
      facecolor=DEPARTMENT_GRAND_EST_COLOR[dept],
      alpha=0.3,
    )

  ax1.set_ylabel("Patients hospitalisés / 100,000")
  ax1.set_title("Évolution du nombre de patients hospitalisés / 100,000 hab.")
  ax1.legend(ncol=2)
  ax2.set_ylabel("Patients en réanimation / 100,000")
  ax2.set_title(
    "Évolution du nombre de patients en réanimation / 100,000 hab."
  )

  ax2.set_xticks(date_range_idx)
  ax2.set_xticklabels(
    [date.strftime("%d-%m") for date in sorted(data.date.unique())],
    rotation=45,
  )
  ax2.legend(
    [
      matplotlib.lines.Line2D([0], [0], color="k", lw=2, ls="solid"),
      matplotlib.lines.Line2D([0], [0], color="k", lw=2, ls="dashed"),
    ],
    ["Donnée publique", "Donnée ICUBAM"],
    ncol=1,
  )
  ax2.set_xlim(0, len(date_range_idx) - 1)
  tikzplotlib_kwargs = dict(
    axis_width="12cm",
    axis_height="8cm",
  )
  return fig, tikzplotlib_kwargs

max_stars_repo_name: maximedenes/icubam

max_issues_repo_name: maximedenes/icubam

max_forks_repo_name: maximedenes/icubam

reason: import numpy


text:
"""
Grady Williams
May 2013

This program computes all of the integer points inside an n-dimensional box
of radius R.
"""

import time
import numpy as np
import scipy.linalg as la
import pycuda.driver as cuda
import pycuda.autoinit
from pycuda.compiler import SourceModule
from pycuda import gpuarray
from riemanntheta_cuda import RiemannThetaCuda

def riemanntheta_high_dim(X, Yinv, T, z, g, rad, max_points = 10000000):
    parRiemann = RiemannThetaCuda(1,512)
    #initialize parRiemann
    parRiemann.compile(g)
    parRiemann.cache_omega_real(X)
    parRiemann.cache_omega_imag(Yinv,T)
    #compile the box_points program
    point_finder = func1()
    R = get_rad(T, rad)
    print R
    num_int_points = (2*R + 1)**g
    num_partitions = num_int_points//max_points
    num_final_partition = num_int_points - num_partitions*max_points
    osc_part = 0 + 0*1.j
    if (num_partitions > 0):
        S = gpuarray.zeros(np.int(max_points * g), dtype=np.double)
    print "Required number of iterations"
    print num_partitions
    print
    for p in range(num_partitions):
        print p
        print
        S = box_points(point_finder, max_points*p, max_points*(p+1),g,R, S)
        parRiemann.cache_intpoints(S, gpu_already=True)
        osc_part += parRiemann.compute_v_without_derivs(np.array([z]))
    S = gpuarray.zeros(np.int((num_int_points - num_partitions*max_points)*g), dtype = np.double)
    print num_partitions*max_points,num_int_points
    S = box_points(point_finder, num_partitions*max_points, num_int_points, g, R,S)
    parRiemann.cache_intpoints(S,gpu_already = True)
    osc_part += parRiemann.compute_v_without_derivs(np.array([z]))
    print osc_part
    return osc_part

def get_rad(T, rad):
    ell = np.dot(T, T.transpose())*rad
    eigs = la.eigvals(ell).real
    eigs = np.absolute(eigs)
    R = np.ceil(max(eigs))
    return np.int64(R)

def box_points(point_finder, min_p,max_p,g,r, S):
    num_points = max_p - min_p
    min_p_d = gpuarray.to_gpu(np.array([min_p], dtype = np.int64))
    max_p_d = gpuarray.to_gpu(np.array([max_p], dtype = np.int64))
    g_d = gpuarray.to_gpu(np.array([g], dtype = np.int64))
    r_d = gpuarray.to_gpu(np.array([r], dtype = np.int64))
    blocksize = (10, 1, 1)
    gridsize = (1000, 1000, 1)
    cuda.Context.synchronize()
    point_finder(S, min_p_d,max_p_d, g_d, r_d, block=blocksize, grid=gridsize)
    cuda.Context.synchronize()
    return S

def checklist(S):
    B = S.reshape(len(S)/2,2)
    print B[10000000:]

def func1():
    template = """

#include <stdlib.h>
#include <stdio.h>
#include <math.h>

__global__ void point_finder(double* S, int* start, int* max_p, int* g, int* r)
{
   int tdx = threadIdx.x;
   int bdx = blockIdx.x;
   int bdy = blockIdx.y;
   int index = 1000*(10)*bdy + 10*bdx + tdx + start[0];
   if (index < max_p[0]) {
     int i = 0;
     for (i = 0; i < g[0]; i++) {
       int p = 1;
       int j = 0;
       for (j = 0; j < i; j++){
         p *= 2*r[0] + 1;
       }
       S[g[0]*(index - start[0]) + i] = ((index/p) % (2*r[0] + 1)) - r[0];
     }
   }
}
"""
    mod = SourceModule(template)
    func = mod.get_function("point_finder")
    return func

if __name__=="__main__":
    print "Testing GPU capacity"
    rad = 100
    for i in range(1,32):
        box_points(10**i, 0)
        print i

max_stars_repo_name: RParini/abelfunctions

max_issues_repo_name: RParini/abelfunctions

max_forks_repo_name: RParini/abelfunctions

reason: import numpy,import scipy,import pycuda,from pycuda


text:
#!/usr/bin/env python3

"""
 Copyright (c) 2018, salesforce.com, inc.
 All rights reserved.
 SPDX-License-Identifier: BSD-3-Clause
 For full license text, see the LICENSE file in the repo root or https://opensource.org/licenses/BSD-3-Clause

 Experiment Portal.
"""

import copy
import itertools
import numpy as np
import os, sys
import random
import logging

import torch

from src.parse_args import parser
from src.parse_args import args
import src.data_utils as data_utils
import src.eval
from src.knowledge_graph import KnowledgeGraph
from src.emb.fact_network import ComplEx, ConvE, DistMult
from src.emb.fact_network import get_conve_kg_state_dict, get_complex_kg_state_dict, get_distmult_kg_state_dict
from src.emb.emb import EmbeddingBasedMethod
from src.rl.graph_search.pn import GraphSearchPolicy
from src.rl.graph_search.pg import PolicyGradient
from src.rl.graph_search.rs_pg import RewardShapingPolicyGradient

torch.cuda.set_device(args.gpu)

torch.manual_seed(args.seed)
torch.cuda.manual_seed_all(args.seed)

def process_data():
    data_dir = args.data_dir
    raw_kb_path = os.path.join(data_dir, 'raw.kb')
    train_path = data_utils.get_train_path(args)
    dev_path = os.path.join(data_dir, 'dev.triples')
    test_path = os.path.join(data_dir, 'test.triples')
    data_utils.prepare_kb_envrioment(raw_kb_path, train_path, dev_path, test_path, args.test,
                                     args.add_reverse_relations)

def initialize_model_directory(args, random_seed=None):
    # add model parameter info to model directory
    model_root_dir = args.model_root_dir
    dataset = os.path.basename(os.path.normpath(args.data_dir))

    reverse_edge_tag = '-RV' if args.add_reversed_training_edges else ''
    entire_graph_tag = '-EG' if args.train_entire_graph else ''
    if args.xavier_initialization:
        initialization_tag = '-xavier'
    elif args.uniform_entity_initialization:
        initialization_tag = '-uniform'
    else:
        initialization_tag = ''

    # Hyperparameter signature

    if args.model in ['rule']:
        print("rule")
        hyperparam_sig = '{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}'.format(
            args.baseline,
            args.entity_dim,
            args.relation_dim,
            args.history_num_layers,
            args.learning_rate,
            args.emb_dropout_rate,
            args.ff_dropout_rate,
            args.action_dropout_rate,
            args.bandwidth,
            args.beta,
            args.rl_module,
            args.beam_size_high,
            args.beam_size_low
        )
    elif args.model.startswith('point'):
        print("point")
        if args.baseline == 'avg_reward':
            logging.info('* Policy Gradient Baseline: average reward')
        elif args.baseline == 'avg_reward_normalized':
            logging.info('* Policy Gradient Baseline: average reward baseline plus normalization')
        else:
            logging.info('* Policy Gradient Baseline: None')
        if args.action_dropout_anneal_interval < 1000:
            hyperparam_sig = '{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}'.format(
                args.baseline,
                args.entity_dim,
                args.relation_dim,
                args.history_num_layers,
                args.learning_rate,
                args.emb_dropout_rate,
                args.ff_dropout_rate,
                args.action_dropout_rate,
                args.action_dropout_anneal_factor,
                args.action_dropout_anneal_interval,
                args.bandwidth,
                args.beta,
                args.gamma,
                args.num_rollout_steps,
                args.num_rollouts,
                args.rl_module,
                args.beam_size_high,
                args.beam_size_low,
                args.beam_size,
                args.relation_dropout_rate,
                args.tailentity_dropout_rate,
                args.train_batch_size,
                args.dev_batch_size,
                args.additional_tailentity_size
            )
            if args.mu != 1.0:
                hyperparam_sig += '-{}'.format(args.mu)
        else:
            hyperparam_sig = '{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}-{}'.format(
                args.baseline,
                args.entity_dim,
                args.relation_dim,
                args.history_num_layers,
                args.learning_rate,
                args.emb_dropout_rate,
                args.ff_dropout_rate,
                args.action_dropout_rate,
                args.bandwidth,
                args.beta,
                args.gamma,
                args.num_rollout_steps,
                args.num_rollouts,
                args.rl_module,
                args.beam_size_high,
                args.beam_size_low,
                args.beam_size,
                args.relation_dropout_rate,
                args.tailentity_dropout_rate,
                args.train_batch_size,
                args.dev_batch_size,
                args.additional_tailentity_size
            )
        if args.reward_shaping_threshold > 0:
            hyperparam_sig += '-{}'.format(args.reward_shaping_threshold)
    elif args.model == 'distmult':
        hyperparam_sig = '{}-{}-{}-{}-{}'.format(
            args.entity_dim,
            args.relation_dim,
            args.learning_rate,
            args.emb_dropout_rate,
            args.label_smoothing_epsilon
        )
    elif args.model == 'complex':
        hyperparam_sig = '{}-{}-{}-{}-{}'.format(
            args.entity_dim,
            args.relation_dim,
            args.learning_rate,
            args.emb_dropout_rate,
            args.label_smoothing_epsilon
        )
    elif args.model in ['conve', 'hypere', 'triplee']:
        hyperparam_sig = '{}-{}-{}-{}-{}-{}-{}-{}-{}'.format(
            args.entity_dim,
            args.relation_dim,
            args.learning_rate,
            args.num_out_channels,
            args.kernel_size,
            args.emb_dropout_rate,
            args.hidden_dropout_rate,
            args.feat_dropout_rate,
            args.label_smoothing_epsilon
        )
    else:
        raise NotImplementedError
    model_sub_dir = '{}-{}{}{}{}-{}'.format(
        dataset,
        args.model,
        reverse_edge_tag,
        entire_graph_tag,
        initialization_tag,
        hyperparam_sig
    )

    if args.model == 'set':
        model_sub_dir += '-{}'.format(args.beam_size)
        model_sub_dir += '-{}'.format(args.num_paths_per_entity)
    if args.relation_only:
        model_sub_dir += '-ro'
    elif args.relation_only_in_path:
        model_sub_dir += '-rpo'
    elif args.type_only:
        model_sub_dir += '-to'

    if args.test:
        model_sub_dir += '-test'

    if random_seed:
        model_sub_dir += '.{}'.format(random_seed)

    model_dir = os.path.join(model_root_dir, 'test/' + model_sub_dir)

    if not os.path.exists(model_dir):
        os.makedirs(model_dir)
        logging.info('Model directory created: {}'.format(model_dir))
    else:
        logging.info('Model directory exists: {}'.format(model_dir))

    args.model_dir = model_dir

def construct_model(args):
    """
    Construct NN graph.
    """
    kg = KnowledgeGraph(args)
    if args.model.endswith('.gc'):
        kg.load_fuzzy_facts()

    if args.model in ['point', 'point.gc']:
        # pn => policy network
        pn = GraphSearchPolicy(args)
        # lf => learn framework
        lf = PolicyGradient(args, kg, pn)
    elif args.model.startswith('point.rs'):
        fn_model = args.model.split('.')[2]
        fn_args = copy.deepcopy(args)
        fn_args.model = fn_model
        fn_args.relation_only = False
        if fn_model == 'complex':
            fn = ComplEx(fn_args)
            fn_kg = KnowledgeGraph(fn_args)
        elif fn_model == 'distmult':
            fn = DistMult(fn_args)
            fn_kg = KnowledgeGraph(fn_args)
        elif fn_model == 'conve':
            fn = ConvE(fn_args, kg.num_entities)
            fn_kg = KnowledgeGraph(fn_args)
        pn = GraphSearchPolicy(args, fn, fn_kg)
        lf = RewardShapingPolicyGradient(args, kg, pn, fn_kg, fn)
    elif args.model == 'complex':
        fn = ComplEx(args)
        lf = EmbeddingBasedMethod(args, kg, fn)
    elif args.model == 'distmult':
        fn = DistMult(args)
        lf = EmbeddingBasedMethod(args, kg, fn)
    elif args.model == 'conve':
        fn = ConvE(args, kg.num_entities)
        lf = EmbeddingBasedMethod(args, kg, fn)
    else:
        raise NotImplementedError
    return lf

def train(lf):
    train_path = data_utils.get_train_path(args)
    dev_path = os.path.join(args.data_dir, 'dev.triples')
    test_path = os.path.join(args.data_dir, 'test.triples')

    entity_index_path = os.path.join(args.data_dir, 'entity2id.txt')
    relation_index_path = os.path.join(args.data_dir, 'relation2id.txt')

    if 'NELL' in args.data_dir:
        adj_list_path = os.path.join(args.data_dir, 'adj_list.pkl')
        seen_entities = data_utils.load_seen_entities(adj_list_path, entity_index_path)
    else:
        seen_entities = set()

    train_data = data_utils.load_triples(
        train_path, entity_index_path, relation_index_path, group_examples_by_query=args.group_examples_by_query,
        add_reverse_relations=args.add_reversed_training_edges)
    dev_data = data_utils.load_triples(dev_path, entity_index_path, relation_index_path,
                                       seen_entities=seen_entities)

    if args.checkpoint_path is not None:
        lf.load_checkpoint(args.checkpoint_path)
    lf.run_train(train_data, dev_data)

def inference(lf):
    lf.batch_size = args.dev_batch_size
    lf.eval()
    if args.model == 'hypere':
        conve_kg_state_dict = get_conve_kg_state_dict(torch.load(args.conve_state_dict_path))
        lf.kg.load_state_dict(conve_kg_state_dict)
        secondary_kg_state_dict = get_complex_kg_state_dict(torch.load(args.complex_state_dict_path))
        lf.secondary_kg.load_state_dict(secondary_kg_state_dict)
    elif args.model == 'triplee':
        conve_kg_state_dict = get_conve_kg_state_dict(torch.load(args.conve_state_dict_path))
        lf.kg.load_state_dict(conve_kg_state_dict)
        complex_kg_state_dict = get_complex_kg_state_dict(torch.load(args.complex_state_dict_path))
        lf.secondary_kg.load_state_dict(complex_kg_state_dict)
        distmult_kg_state_dict = get_distmult_kg_state_dict(torch.load(args.distmult_state_dict_path))
        lf.tertiary_kg.load_state_dict(distmult_kg_state_dict)
    else:
        lf.load_checkpoint(get_checkpoint_path(args))
    entity_index_path = os.path.join(args.data_dir, 'entity2id.txt')
    relation_index_path = os.path.join(args.data_dir, 'relation2id.txt')
    if 'NELL' in args.data_dir:
        adj_list_path = os.path.join(args.data_dir, 'adj_list.pkl')
        seen_entities = data_utils.load_seen_entities(adj_list_path, entity_index_path)
    else:
        seen_entities = set()

    eval_metrics = {
        'dev': {},
        'test': {}
    }

    dev_path = os.path.join(args.data_dir, 'dev.triples')
    test_path = os.path.join(args.data_dir, 'test.triples')
    dev_data = data_utils.load_triples(
        dev_path, entity_index_path, relation_index_path, seen_entities=seen_entities, verbose=False)
    test_data = data_utils.load_triples(
        test_path, entity_index_path, relation_index_path, seen_entities=seen_entities, verbose=False)
    logging.info('Dev set performance:')
    pred_scores = lf.forward(dev_data, verbose=False)
    dev_metrics = src.eval.hits_and_ranks(dev_data, pred_scores, lf.kg.dev_objects, verbose=True)
    eval_metrics['dev'] = {}
    eval_metrics['dev']['hits_at_1'] = dev_metrics[0]
    eval_metrics['dev']['hits_at_3'] = dev_metrics[1]
    eval_metrics['dev']['hits_at_5'] = dev_metrics[2]
    eval_metrics['dev']['hits_at_10'] = dev_metrics[3]
    eval_metrics['dev']['mrr'] = dev_metrics[4]
    src.eval.hits_and_ranks(dev_data, pred_scores, lf.kg.all_objects, verbose=True)
    logging.info('Test set performance:')
    pred_scores = lf.forward(test_data, verbose=args.save_beam_search_paths)
    test_metrics = src.eval.hits_and_ranks(test_data, pred_scores, lf.kg.all_objects, verbose=True)
    eval_metrics['test']['hits_at_1'] = test_metrics[0]
    eval_metrics['test']['hits_at_3'] = test_metrics[1]
    eval_metrics['test']['hits_at_5'] = test_metrics[2]
    eval_metrics['test']['hits_at_10'] = test_metrics[3]
    eval_metrics['test']['mrr'] = test_metrics[4]

    return eval_metrics

def run_ablation_studies(args):
    """
    Run the ablation study experiments reported in the paper.
    """

    def set_up_lf_for_inference(args):
        initialize_model_directory(args)
        lf = construct_model(args)
        lf.cuda()
        lf.batch_size = args.dev_batch_size
        lf.load_checkpoint(get_checkpoint_path(args))
        lf.eval()
        return lf

    def rel_change(metrics, ab_system, kg_portion):
        ab_system_metrics = metrics[ab_system][kg_portion]
        base_metrics = metrics['ours'][kg_portion]
        return int(np.round((ab_system_metrics - base_metrics) / base_metrics * 100))

    entity_index_path = os.path.join(args.data_dir, 'entity2id.txt')
    relation_index_path = os.path.join(args.data_dir, 'relation2id.txt')
    if 'NELL' in args.data_dir:
        adj_list_path = os.path.join(args.data_dir, 'adj_list.pkl')
        seen_entities = data_utils.load_seen_entities(adj_list_path, entity_index_path)
    else:
        seen_entities = set()
    dataset = os.path.basename(args.data_dir)
    dev_path = os.path.join(args.data_dir, 'dev.triples')
    dev_data = data_utils.load_triples(
        dev_path, entity_index_path, relation_index_path, seen_entities=seen_entities, verbose=False)
    to_m_rels, to_1_rels, (to_m_ratio, to_1_ratio) = data_utils.get_relations_by_type(args.data_dir,
                                                                                relation_index_path)
    relation_by_types = (to_m_rels, to_1_rels)
    to_m_ratio *= 100
    to_1_ratio *= 100
    seen_queries, (seen_ratio, unseen_ratio) = data_utils.get_seen_queries(args.data_dir, entity_index_path,
                                                                           relation_index_path)
    seen_ratio *= 100
    unseen_ratio *= 100

    systems = ['ours', '-ad', '-rs']
    mrrs, to_m_mrrs, to_1_mrrs, seen_mrrs, unseen_mrrs = {}, {}, {}, {}, {}
    for system in systems:
        logging.info('** Evaluating {} system **'.format(system))
        if system == '-ad':
            args.action_dropout_rate = 0.0
            if dataset == 'umls':
                # adjust dropout hyperparameters
                args.emb_dropout_rate = 0.3
                args.ff_dropout_rate = 0.1
        elif system == '-rs':
            config_path = os.path.join('configs', '{}.sh'.format(dataset.lower()))
            args = parser.parse_args()
            args = data_utils.load_configs(args, config_path)

        lf = set_up_lf_for_inference(args)
        pred_scores = lf.forward(dev_data, verbose=False)
        _, _, _, _, mrr = src.eval.hits_and_ranks(dev_data, pred_scores, lf.kg.dev_objects, verbose=True)
        if to_1_ratio == 0:
            to_m_mrr = mrr
            to_1_mrr = -1
        else:
            to_m_mrr, to_1_mrr = src.eval.hits_and_ranks_by_relation_type(
                dev_data, pred_scores, lf.kg.dev_objects, relation_by_types, verbose=True)
        seen_mrr, unseen_mrr = src.eval.hits_and_ranks_by_seen_queries(
            dev_data, pred_scores, lf.kg.dev_objects, seen_queries, verbose=True)
        mrrs[system] = {'': mrr * 100}
        to_m_mrrs[system] = {'': to_m_mrr * 100}
        to_1_mrrs[system] = {'': to_1_mrr * 100}
        seen_mrrs[system] = {'': seen_mrr * 100}
        unseen_mrrs[system] = {'': unseen_mrr * 100}
        _, _, _, _, mrr_full_kg = src.eval.hits_and_ranks(dev_data, pred_scores, lf.kg.all_objects, verbose=True)
        if to_1_ratio == 0:
            to_m_mrr_full_kg = mrr_full_kg
            to_1_mrr_full_kg = -1
        else:
            to_m_mrr_full_kg, to_1_mrr_full_kg = src.eval.hits_and_ranks_by_relation_type(
                dev_data, pred_scores, lf.kg.all_objects, relation_by_types, verbose=True)
        seen_mrr_full_kg, unseen_mrr_full_kg = src.eval.hits_and_ranks_by_seen_queries(
            dev_data, pred_scores, lf.kg.all_objects, seen_queries, verbose=True)
        mrrs[system]['full_kg'] = mrr_full_kg * 100
        to_m_mrrs[system]['full_kg'] = to_m_mrr_full_kg * 100
        to_1_mrrs[system]['full_kg'] = to_1_mrr_full_kg * 100
        seen_mrrs[system]['full_kg'] = seen_mrr_full_kg * 100
        unseen_mrrs[system]['full_kg'] = unseen_mrr_full_kg * 100

    # overall system comparison (table 3)
    logging.info('Partial graph evaluation')
    logging.info('--------------------------')
    logging.info('Overall system performance')
    logging.info('Ours(ConvE)\t-RS\t-AD')
    logging.info('{:.1f}\t{:.1f}\t{:.1f}'.format(mrrs['ours'][''], mrrs['-rs'][''], mrrs['-ad']['']))
    logging.info('--------------------------')
    # performance w.r.t. relation types (table 4, 6)
    logging.info('Performance w.r.t. relation types')
    logging.info('\tTo-many\t\t\t\tTo-one\t\t')
    logging.info('%\tOurs\t-RS\t-AD\t%\tOurs\t-RS\t-AD')
    logging.info('{:.1f}\t{:.1f}\t{:.1f} ({:d})\t{:.1f} ({:d})\t{:.1f}\t{:.1f}\t{:.1f} ({:d})\t{:.1f} ({:d})'.format(
        to_m_ratio, to_m_mrrs['ours'][''], to_m_mrrs['-rs'][''], rel_change(to_m_mrrs, '-rs', ''), to_m_mrrs['-ad'][''],
        rel_change(to_m_mrrs, '-ad', ''),
        to_1_ratio, to_1_mrrs['ours'][''], to_1_mrrs['-rs'][''], rel_change(to_1_mrrs, '-rs', ''), to_1_mrrs['-ad'][''],
        rel_change(to_1_mrrs, '-ad', '')))
    logging.info('--------------------------')
    # performance w.r.t. seen queries (table 5, 7)
    logging.info('Performance w.r.t. seen/unseen queries')
    logging.info('\tSeen\t\t\t\tUnseen\t\t')
    logging.info('%\tOurs\t-RS\t-AD\t%\tOurs\t-RS\t-AD')
    logging.info('{:.1f}\t{:.1f}\t{:.1f} ({:d})\t{:.1f} ({:d})\t{:.1f}\t{:.1f}\t{:.1f} ({:d})\t{:.1f} ({:d})'.format(
        seen_ratio, seen_mrrs['ours'][''], seen_mrrs['-rs'][''], rel_change(seen_mrrs, '-rs', ''), seen_mrrs['-ad'][''],
        rel_change(seen_mrrs, '-ad', ''),
        unseen_ratio, unseen_mrrs['ours'][''], unseen_mrrs['-rs'][''], rel_change(unseen_mrrs, '-rs', ''),
        unseen_mrrs['-ad'][''], rel_change(unseen_mrrs, '-ad', '')))
    logging.info()
    logging.info('Full graph evaluation')
    logging.info('--------------------------')
    logging.info('Overall system performance')
    logging.info('Ours(ConvE)\t-RS\t-AD')
    logging.info(
        '{:.1f}\t{:.1f}\t{:.1f}'.format(mrrs['ours']['full_kg'], mrrs['-rs']['full_kg'], mrrs['-ad']['full_kg']))
    logging.info('--------------------------')
    logging.info('Performance w.r.t. relation types')
    logging.info('\tTo-many\t\t\t\tTo-one\t\t')
    logging.info('%\tOurs\t-RS\t-AD\t%\tOurs\t-RS\t-AD')
    logging.info('{:.1f}\t{:.1f}\t{:.1f} ({:d})\t{:.1f} ({:d})\t{:.1f}\t{:.1f}\t{:.1f} ({:d})\t{:.1f} ({:d})'.format(
        to_m_ratio, to_m_mrrs['ours']['full_kg'], to_m_mrrs['-rs']['full_kg'], rel_change(to_m_mrrs, '-rs', 'full_kg'),
        to_m_mrrs['-ad']['full_kg'], rel_change(to_m_mrrs, '-ad', 'full_kg'),
        to_1_ratio, to_1_mrrs['ours']['full_kg'], to_1_mrrs['-rs']['full_kg'], rel_change(to_1_mrrs, '-rs', 'full_kg'),
        to_1_mrrs['-ad']['full_kg'], rel_change(to_1_mrrs, '-ad', 'full_kg')))
    logging.info('--------------------------')
    logging.info('Performance w.r.t. seen/unseen queries')
    logging.info('\tSeen\t\t\t\tUnseen\t\t')
    logging.info('%\tOurs\t-RS\t-AD\t%\tOurs\t-RS\t-AD')
    logging.info('{:.1f}\t{:.1f}\t{:.1f} ({:d})\t{:.1f} ({:d})\t{:.1f}\t{:.1f}\t{:.1f} ({:d})\t{:.1f} ({:d})'.format(
        seen_ratio, seen_mrrs['ours']['full_kg'], seen_mrrs['-rs']['full_kg'], rel_change(seen_mrrs, '-rs', 'full_kg'),
        seen_mrrs['-ad']['full_kg'], rel_change(seen_mrrs, '-ad', 'full_kg'),
        unseen_ratio, unseen_mrrs['ours']['full_kg'], unseen_mrrs['-rs']['full_kg'],
        rel_change(unseen_mrrs, '-rs', 'full_kg'), unseen_mrrs['-ad']['full_kg'],
        rel_change(unseen_mrrs, '-ad', 'full_kg')))

def export_to_embedding_projector(lf):
    lf.load_checkpoint(get_checkpoint_path(args))
    lf.export_to_embedding_projector()

def export_reward_shaping_parameters(lf):
    lf.load_checkpoint(get_checkpoint_path(args))
    lf.export_reward_shaping_parameters()

def export_fuzzy_facts(lf):
    lf.load_checkpoint(get_checkpoint_path(args))
    lf.export_fuzzy_facts()

def export_error_cases(lf):
    lf.load_checkpoint(get_checkpoint_path(args))
    lf.batch_size = args.dev_batch_size
    lf.eval()
    entity_index_path = os.path.join(args.data_dir, 'entity2id.txt')
    relation_index_path = os.path.join(args.data_dir, 'relation2id.txt')
    dev_path = os.path.join(args.data_dir, 'dev.triples')
    dev_data = data_utils.load_triples(dev_path, entity_index_path, relation_index_path)
    lf.load_checkpoint(get_checkpoint_path(args))
    logging.info('Dev set performance:')
    pred_scores = lf.forward(dev_data, verbose=False)
    src.eval.hits_and_ranks(dev_data, pred_scores, lf.kg.dev_objects, verbose=True)
    src.eval.export_error_cases(dev_data, pred_scores, lf.kg.dev_objects, os.path.join(lf.model_dir, 'error_cases.pkl'))

def compute_fact_scores(lf):
    data_dir = args.data_dir
    train_path = os.path.join(data_dir, 'train.triples')
    dev_path = os.path.join(data_dir, 'dev.triples')
    test_path = os.path.join(data_dir, 'test.triples')
    entity_index_path = os.path.join(args.data_dir, 'entity2id.txt')
    relation_index_path = os.path.join(args.data_dir, 'relation2id.txt')
    train_data = data_utils.load_triples(train_path, entity_index_path, relation_index_path)
    dev_data = data_utils.load_triples(dev_path, entity_index_path, relation_index_path)
    test_data = data_utils.load_triples(test_path, entity_index_path, relation_index_path)
    lf.eval()
    lf.load_checkpoint(get_checkpoint_path(args))
    train_scores = lf.forward_fact(train_data)
    dev_scores = lf.forward_fact(dev_data)
    test_scores = lf.forward_fact(test_data)

    logging.info('Train set average fact score: {}'.format(float(train_scores.mean())))
    logging.info('Dev set average fact score: {}'.format(float(dev_scores.mean())))
    logging.info('Test set average fact score: {}'.format(float(test_scores.mean())))

def get_checkpoint_path(args):
    if not args.checkpoint_path:
        return os.path.join(args.model_dir, 'model_best.tar')
    else:
        return args.checkpoint_path

def load_configs(config_path):
    with open(config_path) as f:
        logging.info('loading configuration file {}'.format(config_path))
        for line in f:
            if not '=' in line:
                continue
            arg_name, arg_value = line.strip().split('=')
            if arg_value.startswith('"') and arg_value.endswith('"'):
                arg_value = arg_value[1:-1]
            if hasattr(args, arg_name):
                logging.info('{} = {}'.format(arg_name, arg_value))
                arg_value2 = getattr(args, arg_name)
                if type(arg_value2) is str:
                    setattr(args, arg_name, arg_value)
                elif type(arg_value2) is bool:
                    if arg_value == 'True':
                        setattr(args, arg_name, True)
                    elif arg_value == 'False':
                        setattr(args, arg_name, False)
                    else:
                        raise ValueError('Unrecognized boolean value description: {}'.format(arg_value))
                elif type(arg_value2) is int:
                    setattr(args, arg_name, int(arg_value))
                elif type(arg_value2) is float:
                    setattr(args, arg_name, float(arg_value))
                else:
                    raise ValueError('Unrecognized attribute type: {}: {}'.format(arg_name, type(arg_value2)))
            else:
                raise ValueError('Unrecognized argument: {}'.format(arg_name))
    return args

def run_experiment(args):
    if args.test:
        if 'NELL' in args.data_dir:
            dataset = os.path.basename(args.data_dir)
            args.distmult_state_dict_path = data_utils.change_to_test_model_path(dataset, args.distmult_state_dict_path)
            args.complex_state_dict_path = data_utils.change_to_test_model_path(dataset, args.complex_state_dict_path)
            args.conve_state_dict_path = data_utils.change_to_test_model_path(dataset, args.conve_state_dict_path)
        args.data_dir += '.test'

    if args.process_data:

        # Process knowledge graph data

        process_data()
    else:
        with torch.set_grad_enabled(args.train or args.search_random_seed or args.grid_search):
            initialize_model_directory(args)
            lf = construct_model(args)
            lf.cuda()

            if args.train:
                train(lf)
            elif args.inference:
                inference(lf)
            elif args.eval_by_relation_type:
                inference(lf)
            elif args.eval_by_seen_queries:
                inference(lf)
            elif args.export_to_embedding_projector:
                export_to_embedding_projector(lf)
            elif args.export_reward_shaping_parameters:
                export_reward_shaping_parameters(lf)
            elif args.compute_fact_scores:
                compute_fact_scores(lf)
            elif args.export_fuzzy_facts:
                export_fuzzy_facts(lf)
            elif args.export_error_cases:
                export_error_cases(lf)

def set_logger(args):
    """
    Write logs to checkpoint and console
    """
    dataset = args.data_dir
    dataset = dataset.split('/')[-1]
    if args.train:
        log_name = dataset + '-' + args.model + '-' + str(args.rl_module) + \
                   '-num_rollout_steps' + str(args.num_rollout_steps) + \
                   '-num_rollouts' + str(args.num_rollouts) + \
                   '-gamma' + str(args.gamma) + '-beamsize_high' + str(args.beam_size_high) + \
                   '-beamsize-low' + str(args.beam_size_low) + '-beamsize' + str(args.beam_size) + \
                   '-ff_dropout_rate' + str(args.ff_dropout_rate) + \
                   '-r_dropout_rate' + str(args.relation_dropout_rate) + \
                   '-te_dropout_rate' + str(args.tailentity_dropout_rate) + \
                   '-train_bs' + str(args.train_batch_size) + \
                   '-dev_bs' + str(args.dev_batch_size) + \
                   '-add_tail_size' + str(args.additional_tailentity_size)
        log_name += '-addtailentity'

    elif args.inference:
        log_name = dataset + '-' + args.model + '-' + 'inference'
    else:
        log_name = dataset + '-' + args.model + '-' + 'analysis'

    log_file = os.path.join(args.log_dir, 'test/' + log_name + '.log')
    print("Log file name => {}".format(log_file))

    # logging module
    logger = logging.getLogger()
    # clear default handlers for avioding the mutiple output infos
    logger.handlers = []
    # logger.setLevel(logging.INFO)  # Log等级总开关
    formatter = logging.Formatter('%(asctime)s %(levelname)-8s %(message)s')
    ## log to file
    fh = logging.FileHandler(log_file)
    fh.setFormatter(formatter)
    fh.setLevel(logging.INFO)
    logger.addHandler(fh)
    # log to console
    ch = logging.StreamHandler()
    ch.setLevel(logging.INFO)  # 输出到console的log等级的开关
    ch.setFormatter(formatter)
    logger.addHandler(ch)

if __name__ == '__main__':
    set_logger(args)
    run_experiment(args)


text:
import numpy as np
import matplotlib.pyplot as plt
import keras
from keras.layers import Activation, Add, BatchNormalization, LeakyReLU
from keras.layers import Input, Conv2D, Conv2DTranspose, Dense, Flatten, Concatenate
from keras.models import Model
from keras.initializers import RandomNormal
from img_io import Dataset

# GAN model class, it contains all sub models and adversarial model itself.
# Model networks and inputs can be editable from model schemes
class GAN:
    def __init__(self,
                 input_shape: tuple,
                 loss: str,
                 metrics: list,
                 optimizer: keras.optimizers.Optimizer):
        # Model input information
        self.input_shape = input_shape
        self.input_width = input_shape[0]
        self.input_height = input_shape[1]
        self.input_channel = input_shape[2]

        # Model compile informations
        self.batch_size = None
        self.optimizer = optimizer
        self.loss = loss
        self.metrics = metrics

        # Sub-Models and Adversarial Model
        self.generator = None
        self.discriminator = None
        self.adversarial = None

        # Private network schemes
        self.__generator_network = None
        self.__discriminator_network = None
        self.__generator_outputs = None
        self.__discriminator_outputs = None

    # Creates generator sub-model of adversarial network
    def create_generator(self):
        inputs = Input(shape=(self.input_width, self.input_height, self.input_channel),
                       name='generator/input')

        outputs = self._generator_network(inputs)

        model = Model(inputs=inputs, outputs=outputs, name='Generator_Model')
        self.generator = model

    # Generator network scheme
    @staticmethod
    def _generator_network(inputs):
        x = Conv2D(64, kernel_size=(4, 4), strides=(1, 1), padding='same', name='generator/b1/conv2d')(inputs)
        x_add1 = Activation('relu', name='generator/b1/relu')(x)

        x = Conv2D(64, kernel_size=(3, 3), strides=(1, 1), padding='same', name='generator/b2/conv2d')(x_add1)
        x = BatchNormalization(name='generator/b2/batch_norm')(x)
        x = Activation('relu', name='generator/b2/relu')(x)

        x = Conv2D(64, kernel_size=(3, 3), strides=(1, 1), padding='same', name='generator/b3/conv2d')(x)
        x = BatchNormalization(name='generator/b3/batch_norm')(x)
        x_add2 = Activation('relu', name='generator/b3/relu')(x)

        x = Concatenate(name='generator/b4/skip_con1')([x_add1, x_add2])

        x = Conv2DTranspose(3, kernel_size=(4, 4), strides=(1, 1), padding='same', name='generator/b4/deconv2d')(x)
        outputs = Activation('tanh', name='generator/b4/tanh')(x)
        return outputs

    # Creates discriminator sub-model of adversarial network
    def create_discriminator(self, is_trainable: bool = False):
        inputs_gen = Input(shape=(self.input_width, self.input_height, self.input_channel),
                           name='discriminator/input_gen')
        inputs_low = Input(batch_shape=(self.batch_size, self.input_width, self.input_height, self.input_channel),
                           name='discriminator/input_low')

        inputs = Concatenate(name='discriminator/inputs')([inputs_gen, inputs_low])

        outputs = self._discriminator_network(inputs)

        model = Model(inputs=[inputs_gen, inputs_low], outputs=outputs, name='Discriminator_Model')
        model.compile(optimizer=self.optimizer, loss=self.loss, metrics=self.metrics)

        self.discriminator = model
        self.discriminator.trainable = is_trainable

    # Discriminator network scheme
    @staticmethod
    def _discriminator_network(inputs):
        init = RandomNormal(stddev=0.02)
        x = Conv2D(64, kernel_size=(4, 4), strides=(2, 2), padding='same', kernel_initializer=init, name='discriminator/b1/conv2d')(inputs)
        x = LeakyReLU(alpha=0.2, name='discriminator/b1/leaky')(x)

        x = Conv2D(128, kernel_size=(4, 4), strides=(2, 2), padding='same', kernel_initializer=init, name='discriminator/b2/conv2d')(x)
        x = BatchNormalization(name='discriminator/b2/batch_norm')(x)
        x = LeakyReLU(alpha=0.2, name='discriminator/b2/leaky')(x)

        x = Conv2D(256, kernel_size=(4, 4), strides=(2, 2), padding='same', kernel_initializer=init, name='discriminator/b3/conv2d')(x)
        x = BatchNormalization(name='discriminator/b3/batch_norm')(x)
        x = LeakyReLU(alpha=0.2, name='discriminator/b3/leaky')(x)

        x = Conv2D(512, kernel_size=(4, 4), strides=(1, 1), padding='same', kernel_initializer=init, name='discriminator/b4/conv2d')(x)
        x = BatchNormalization(name='discriminator/b4/batch_norm')(x)
        x = LeakyReLU(alpha=0.2, name='discriminator/b4/leaky')(x)

        x = Conv2D(1, kernel_size=(4, 4), strides=(1, 1), padding='same', kernel_initializer=init, name='discriminator/b5/conv2d')(x)
        x = Flatten(name='discriminator/b5/flatten')(x)
        x = Dense(1, kernel_initializer=init, name='discriminator/b5/dense')(x)
        outputs = Activation('sigmoid', name='discriminator/b5/sigmoid')(x)

        return outputs

    # Creates combined network from generative and discriminator sub-models to create adversarial model
    def create_adversarial(self):
        low_input = Input(batch_shape=(self.batch_size, self.input_width, self.input_height, self.input_channel),
                          name='adversarial/input_low')
        # high_input = Input(batch_shape=(self.batch_size, self.input_width, self.input_height, self.input_channel),
        #                    name='adversarial/input_high')

        gen_output = self.generator(low_input)
        disc_output = self.discriminator([gen_output, low_input])

        model = Model(inputs=low_input, outputs=disc_output, name='Adversarial_Model')
        model.compile(optimizer=self.optimizer, loss=self.loss, )

        self.adversarial = model

# Train given model with given dataset.
# Epochs and batch sizes can be editable
class Train:
    def __init__(self, model: GAN, dataset: Dataset, epoch: int, batch_size: int):
        # Models itself
        self.model = model
        self.model.batch_size = batch_size

        # Networks inside main GAN model
        self.adversarial: keras.models.Model = model.adversarial
        self.generator: keras.models.Model = model.generator
        self.discriminator: keras.models.Model = model.discriminator

        # Total image number in dataset.
        # Not sum of high and low
        self.total_image: int = dataset.total_image

        # Dataset object and selected load batch size
        # Batch size need to be same with models batch size
        self.dataset = dataset
        self.dataset.batch_size = batch_size

        # Training parameters
        self.epoch: int = epoch
        self.batch_size: int = batch_size
        self.steps_per_epoch: int = int(self.dataset.total_image / self.dataset.batch_size)

    # Trains model with batches
    def train_on_batch(self):
        y_disc_real = np.ones((self.batch_size, 1))
        y_disc_real[:, 0] = 0.9
        y_disc_fake = np.zeros((self.batch_size, 1))
        y_gen = np.ones(self.batch_size)

        for epoch in range(self.epoch):
            for batch in range(self.steps_per_epoch):
                high_input, low_input = self.dataset.load_batch(batch)
                fake_output = self.generator.predict(low_input)

                disc_loss_real = self.discriminator.train_on_batch([high_input, low_input], y_disc_real)
                disc_loss_fake = self.discriminator.train_on_batch([fake_output, low_input], y_disc_fake)
                disc_loss = 0.5 * np.add(disc_loss_real, disc_loss_fake)

                gen_loss = self.adversarial.train_on_batch(low_input, y_gen)

                self.train_logger('batch', epoch, batch, disc_loss, gen_loss)
        # noinspection PyUnboundLocalVariable
        self.train_logger('epoch', epoch, batch, disc_loss, gen_loss)

    def train_logger(self, log_select, epoch, batch, disc_loss, gen_loss):
        if log_select == 'batch' and batch % 300 == 0 and batch != 0:
            print(f'Step: {batch}/{self.steps_per_epoch}\n'
                  f'Discriminator Loss: {disc_loss[0]:2.5f}\n'
                  f'Accuracy: %{disc_loss[1] * 100:3.2f}\n'
                  f'Generator Loss: {float(gen_loss):2.5f}\n'
                  f'Sample images saved in: generated/IE-CGAN_epoch{epoch}_batch{batch}.png\n')
            self.save_sample_images(epoch, batch)
        elif log_select == 'batch' and batch % 100 == 0 and batch != 0:
            print(f'Step: {batch}/{self.steps_per_epoch}\n'
                  f'Discriminator Loss: {disc_loss[0]:2.5f}\n'
                  f'Accuracy: %{disc_loss[1] * 100:3.2f}\n'
                  f'Generator Loss: {float(gen_loss):2.5f}\n')

        if log_select == 'epoch':
            print(f'\tEpoch: {epoch}/{self.epoch}'
                  f'\tAvg. Discriminator Loss: {disc_loss[0]:2.5f}\n'
                  f'\tAvg. Accuracy: %{disc_loss[1] * 100:3.2f}\n'
                  f'\tAvg. Generator Loss: {float(gen_loss):2.5f}\n')

    def save_sample_images(self, epoch, step):
        gen_imgs = self.generator.predict_on_batch(self.dataset.sample_images)

        gen_imgs = 0.5 * gen_imgs + 0.5
        gen_imgs = np.array(gen_imgs * 255).astype(np.uint8)

        r, c = 2, 3
        fig, axs = plt.subplots(r, c)
        cnt = 0
        for i in range(r):
            for j in range(c):
                axs[i, j].imshow(gen_imgs[cnt, :, :, 0], cmap='gray')
                axs[i, j].axis('off')
                cnt += 1
        fig.savefig(f"generated/IE-CGAN_epoch{epoch}_batch{step}.png", dpi=300)
        plt.close()

max_stars_repo_name: omerferhatt/gan-ir-image-enhance

max_issues_repo_name: omerferhatt/gan-ir-image-enhance

max_forks_repo_name: omerferhatt/gan-ir-image-enhance


text:
from collections import deque

import arche.tools.schema as schema_tools
import numpy as np
import pytest

schema = {
    "$schema": "http://json-schema.org/draft-07/schema#",
    "definitions": {
        "url": {
            "pattern": (
                r"^https?://(www\.)?[a-z0-9.-]*\.[a-z]{2,}"
                r"([^<>%\x20\x00-\x1f\x7F]|%[0-9a-fA-F]{2})*$"
            )
        }
    },
    "type": "object",
    "properties": {
        "url": {"type": "string"},
        "id": {"type": "integer"},
        "race": {
            "type": "object",
            "properties": {"t": {"type": "string"}},
            "additionalProperties": False,
            "required": ["t"],
        },
        "a": {
            "items": {
                "additionalProperties": False,
                "properties": {"k": {"type": "string"}},
                "required": ["k"],
                "type": "object",
            },
            "type": "array",
            "uniqueItems": True,
        },
    },
    "additionalProperties": False,
    "required": ["a", "id", "race", "url"],
}

def test_set_item_no():
    assert schema_tools.set_item_no(4) == [0, 1, 2, 3]
    assert schema_tools.set_item_no(1) == [0]
    assert len(schema_tools.set_item_no(5)) == 4
    assert len(schema_tools.set_item_no(124112414)) == 4

def test_infer_schema():
    item1 = {
        "url": "https://example.com",
        "id": 0,
        "race": {"t": "21k"},
        "a": [{"k": "v"}],
    }
    item2 = {
        "url": "https://example.com",
        "id": 1,
        "race": {"t": "10k"},
        "a": [{"k": "v"}],
    }

    assert schema_tools.infer_schema([item1, item2]) == schema
    assert schema_tools.infer_schema([item2]) == schema

def test_basic_json_schema(mocker):
    mocked_create_js = mocker.patch(
        "arche.tools.schema.create_json_schema", return_value=schema, autospec=True
    )
    assert schema_tools.basic_json_schema("235801/1/15", [0, 5]).raw == schema
    mocked_create_js.assert_called_once_with("235801/1/15", [0, 5])

def test_create_json_schema(mocker, get_job, get_raw_items):
    mocker.patch("arche.tools.api.get_job", return_value=get_job, autospec=True)
    mocker.patch("arche.tools.api.get_items", return_value=get_raw_items, autospec=True)
    schema_tools.create_json_schema(get_job.key, [2])
    assert schema_tools.create_json_schema(get_job.key, [2, 3]) == {
        "$schema": "http://json-schema.org/draft-07/schema#",
        "definitions": {
            "url": {
                "pattern": (
                    r"^https?://(www\.)?[a-z0-9.-]*\.[a-z]{2,}"
                    r"([^<>%\x20\x00-\x1f\x7F]|%[0-9a-fA-F]{2})*$"
                )
            }
        },
        "additionalProperties": False,
        "type": "object",
        "properties": {"name": {"type": "string"}, "price": {"type": "integer"}},
        "required": ["name", "price"],
    }

@pytest.mark.parametrize(
    "error_msg, path, schema_path, validator, formatted_msg",
    [
        (
            "is not of type 'string'",
            ["root_field", 0, "deep", "deeper", 100],
            [],
            "type",
            "root_field/deep/deeper is not of type 'string'",
        ),
        ("other error", ["field"], [], "", "field - other error"),
        (
            "'biography' is a required property",
            [],
            [],
            "required",
            "'biography' is a required property",
        ),
        (
            "'s' does not match any of the regexes: '^spec(_\\w.+){1,}$'",
            [],
            [],
            "pattern",
            "'s' does not match any of the regexes: '^spec(_\\w.+){1,}$'",
        ),
        (
            "Additional properties are not allowed ('price_was', 'stock' were unexpected)",
            [],
            [],
            "additionalProperties",
            "Additional properties are not allowed ('price_was', 'stock' were unexpected)",
        ),
        (
            "https://sto' is not a 'date-time'",
            ["url"],
            [],
            "format",
            "url is not a 'date-time'",
        ),
        ("a very long error message", [], ["anyOf"], "anyOf", "'schema/anyOf' failed"),
        (
            "a very very long error message",
            ["options", 0],
            ["properties", "options", "items", "anyOf"],
            "anyOf",
            "'options' does not satisfy 'schema/properties/options/items/anyOf'",
        ),
    ],
)
def test_format_validation_message(
    error_msg, path, schema_path, validator, formatted_msg
):
    assert (
        schema_tools.format_validation_message(
            error_msg, deque(path), deque(schema_path), validator
        )
        == formatted_msg
    )

def test_full_validate_fails():
    errors = schema_tools.full_validate(
        {"properties": {"NAME": {"type": "string"}}},
        np.array([{"NAME": None, "_key": "0"}, {"NAME": None}]),
        keys=[0, 1],
    )
    assert errors == {"NAME is not of type 'string'": {0, 1}}

def test_full_validate_type_key():
    errors = schema_tools.full_validate(
        {"properties": {"A": {"type": "number"}}, "additionalProperties": False},
        np.array([{"A": 0, "_key": "0", "_type": "Some"}, {"A": 1}]),
        [0, 1],
    )
    assert not errors

def test_fast_validate_fails():
    errors = schema_tools.fast_validate(
        {"type": "object", "properties": {"v": {"type": "number"}}},
        [{"v": "0"}, {"v": "1"}],
        keys=[0, 1],
    )
    assert errors == {"data.v must be number": {0, 1}}

def test_fast_validate(get_schema, get_raw_items):
    errors = schema_tools.fast_validate(
        get_schema, get_raw_items, keys=list(range(len(get_raw_items)))
    )
    assert not errors

max_stars_repo_name: WinterComes/arche

max_issues_repo_name: WinterComes/arche

max_forks_repo_name: WinterComes/arche


text:
from fv3fit._shared import (
    Predictor,
    io,
    match_prediction_to_input_coords,
    SAMPLE_DIM_NAME,
    stack,
)
from fv3fit.keras.adapters import (
    ensure_dict_output,
    rename_dict_input,
    rename_dict_output,
    ensure_tuple_model,
)
import tensorflow as tf
from typing import Any, Dict, Hashable, Iterable, Sequence, Mapping
import xarray as xr
import os
from ...._shared import get_dir, put_dir, InputSensitivity
import yaml
import numpy as np
from .halos import append_halos, append_halos_using_mpi
from fv3fit.keras.jacobian import compute_jacobians, nondimensionalize_jacobians

@io.register("all-keras")
class PureKerasModel(Predictor):
    """Model which uses Keras for packing and normalization.

    Assumes wrapped keras model accepts [sample, feature] arrays.
    """

    _MODEL_FILENAME = "model.tf"
    _CONFIG_FILENAME = "config.yaml"
    custom_objects: Dict[str, Any] = {}

    def __init__(
        self,
        input_variables: Iterable[Hashable],
        output_variables: Iterable[Hashable],
        model: tf.keras.Model,
        unstacked_dims: Sequence[str],
        n_halo: int = 0,
    ):
        """Initialize the predictor

        Args:
            input_variables: names of input variables
            output_variables: names of output variables
            model: keras model to wrap
            unstacked_dims: non-sample dimensions of model output
            n_halo: number of halo points required in input data
        """
        super().__init__(input_variables, output_variables)
        self.input_variables = input_variables
        self.output_variables = output_variables
        self.model = model
        self._n_halo = n_halo
        self._unstacked_dims = unstacked_dims

    @classmethod
    def load(cls, path: str) -> "PureKerasModel":
        """Load a serialized model from a directory."""
        with get_dir(path) as path:
            model_filename = os.path.join(path, cls._MODEL_FILENAME)
            model = tf.keras.models.load_model(
                model_filename, custom_objects=cls.custom_objects
            )
            with open(os.path.join(path, cls._CONFIG_FILENAME), "r") as f:
                config = yaml.load(f, Loader=yaml.Loader)
            obj = cls(
                config["input_variables"],
                config["output_variables"],
                model,
                unstacked_dims=config.get("unstacked_dims", None),
                n_halo=config.get("n_halo", 0),
            )
            return obj

    def _array_prediction_to_dataset(
        self, names, outputs, stacked_coords
    ) -> xr.Dataset:
        ds = xr.Dataset()
        for name, output in zip(names, outputs):
            dims = [SAMPLE_DIM_NAME] + list(self._unstacked_dims)
            scalar_singleton_dim = (
                len(output.shape) == len(dims) and output.shape[-1] == 1
            )
            if scalar_singleton_dim:  # remove singleton dimension
                output = output[..., 0]
                dims = dims[:-1]
            da = xr.DataArray(
                data=output, dims=dims, coords={SAMPLE_DIM_NAME: stacked_coords},
            ).unstack(SAMPLE_DIM_NAME)
            dim_order = [dim for dim in self._unstacked_dims if dim in da.dims]
            ds[name] = da.transpose(*dim_order, ...)
        return ds

    def predict(self, X: xr.Dataset) -> xr.Dataset:
        """Predict an output xarray dataset from an input xarray dataset."""
        if self._n_halo > 0:
            if "tile" not in X.dims:
                try:
                    X = append_halos_using_mpi(ds=X, n_halo=self._n_halo)
                except RuntimeError as err:
                    raise ValueError(
                        "either dataset must have tile dimension or MPI must be present"
                    ) from err
            else:
                X = append_halos(ds=X, n_halo=self._n_halo)
        X_stacked = stack(X, unstacked_dims=self._unstacked_dims)
        inputs = [X_stacked[name].values for name in self.input_variables]
        outputs = self.model.predict(inputs)
        if isinstance(outputs, np.ndarray):
            outputs = [outputs]
        return_ds = self._array_prediction_to_dataset(
            self.output_variables, outputs, X_stacked.coords[SAMPLE_DIM_NAME],
        )
        return match_prediction_to_input_coords(X, return_ds)

    def dump(self, path: str) -> None:
        with put_dir(path) as path:
            if self.model is not None:
                model_filename = os.path.join(path, self._MODEL_FILENAME)
                self.model.save(model_filename)
            with open(os.path.join(path, self._CONFIG_FILENAME), "w") as f:
                f.write(
                    yaml.dump(
                        {
                            "input_variables": self.input_variables,
                            "output_variables": self.output_variables,
                            "unstacked_dims": self._unstacked_dims,
                            "n_halo": self._n_halo,
                        }
                    )
                )

    def get_dict_compatible_model(self) -> tf.keras.Model:
        """
        Rename keras model inputs/outputs to match input_variables and
        output_variables, and ensure model accepts and outputs dicts.
        """
        renamed_inputs: Mapping[str, str] = {
            str(layer_name): str(input_var)
            for layer_name, input_var in zip(
                self.model.input_names, self.input_variables
            )
        }
        renamed_outputs: Mapping[str, str] = {
            str(layer_name): str(output_var)
            for layer_name, output_var in zip(
                self.model.output_names, self.output_variables
            )
        }

        dict_compatible_model = ensure_dict_output(self.model)

        model_renamed_inputs = rename_dict_input(dict_compatible_model, renamed_inputs)
        return rename_dict_output(model_renamed_inputs, renamed_outputs)

    def input_sensitivity(self, stacked_sample: xr.Dataset) -> InputSensitivity:
        # Jacobian functions take in dict input
        data_dict = {}
        for var in stacked_sample:
            values = stacked_sample[var].values
            if len(stacked_sample[var].dims) == 1:
                values = values.reshape(-1, 1)
            data_dict[var] = values

        jacobians = compute_jacobians(
            self.get_dict_compatible_model(),  # type: ignore
            data_dict,
            self.input_variables,
        )
        # normalize factors so sensitivities are comparable but still
        # preserve level-relative magnitudes
        std_factors = {name: np.std(data, axis=0) for name, data in data_dict.items()}
        jacobians_std = nondimensionalize_jacobians(jacobians, std_factors)
        return InputSensitivity(jacobians=jacobians_std)

@io.register("all-keras-dict")
class PureKerasDictPredictor(Predictor):
    """Save a keras model encoding the input/output names into a saved_model
    artifact

    The output names of the keras model are assumed to be correct. To do this
    requires including the following logic when building the model::

        out_with_correct_name = tf.keras.layers.Lambda(
            lambda x: x, name="desired_name")

    Note:
        This class relies on PureKerasModel to implement Mhe Predictor methods.
        It might be worth refactoring the implementation here instead since the
        dict-in/out format has more structure.

    """

    _MODEL_FILENAME = "model.tf"
    _CONFIG_FILENAME = "config.yaml"

    def __init__(self, model: tf.keras.Model, passthrough: Any = None):
        """Initialize the predictor.

        Args:
            model: a keras model. Tuple in/out models are okay, but will be cast
                to a dict-output model before saving.
            passthrough: any data! Is NOT serialized, but can be
                used by other parts of the training workflow. Mostly exists for
                backwards compatibility and should be avoided in new code.
        """
        dict_model = ensure_dict_output(model)
        tuple_model = ensure_tuple_model(dict_model)
        super().__init__(tuple_model.input_names, tuple_model.output_names)
        self._model = tuple_model
        self._pure_keras_predictor = PureKerasModel(
            self.input_variables,
            self.output_variables,
            tuple_model,
            # unstacked dims is hardcoded most places in fv3fit.keras, but it
            # seems like useful metadata, but more of an aspect of the training
            # data pipeline since it is the piece of code that removes the
            # dimension names...would be nice to remove this concept...e.g. pass
            # unstacked_dims to the .dump rather than in the constructor
            # TODO consider this refactor
            unstacked_dims=["z"],
        )
        self.passthrough = passthrough

    def predict(self, X: xr.Dataset) -> xr.Dataset:
        """Predict an output xarray dataset from an input xarray dataset."""
        return self._pure_keras_predictor.predict(X)

    def dump(self, path: str) -> None:
        """Serialize to a directory."""
        with put_dir(path) as local_path:
            self._model.save(self._model_filename(local_path))
            path

    @classmethod
    def _model_filename(cls, path: str) -> str:
        return os.path.join(path, cls._MODEL_FILENAME)

    @classmethod
    def load(cls, path: str) -> "Predictor":
        """Load a serialized model from a directory."""
        model = tf.keras.models.load_model(cls._model_filename(path))
        return cls(model)

    def input_sensitivity(self, stacked_sample: xr.Dataset) -> InputSensitivity:
        """Calculate sensitivity to input features."""
        raise NotImplementedError(
            "input_sensitivity is not implemented...but could be pretty "
            "easily using jacobian routines."
        )


text:
# -*- coding: utf-8 -*-
"""Some utility functions."""
# Authors: Alexandre Gramfort <alexandre.gramfort@inria.fr>
#
# License: BSD (3-clause)

from collections.abc import Iterable
import os
import os.path as op
import logging
import tempfile
from threading import Thread
import time

import numpy as np

from ._logging import logger

class ProgressBar(object):
    """Generate a command-line progressbar.

    Parameters
    ----------
    iterable : iterable | int | None
        The iterable to use. Can also be an int for backward compatibility
        (acts like ``max_value``).
    initial_value : int
        Initial value of process, useful when resuming process from a specific
        value, defaults to 0.
    mesg : str
        Message to include at end of progress bar.
    max_total_width : int | str
        Maximum total message width. Can use "auto" (default) to try to set
        a sane value based on the current terminal width.
    max_value : int | None
        The max value. If None, the length of ``iterable`` will be used.
    **kwargs : dict
        Additional keyword arguments for tqdm.
    """

    def __init__(self, iterable=None, initial_value=0, mesg=None,
                 max_total_width='auto', max_value=None,
                 **kwargs):  # noqa: D102
        from ..externals.tqdm import auto
        tqdm = auto.tqdm
        defaults = dict(
            leave=True, mininterval=0.016, miniters=0, smoothing=0.9,
            bar_format='{percentage:3.0f}%|{bar}| {desc} : {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt:>11}{postfix}]',  # noqa: E501
        )
        for key, val in defaults.items():
            if key not in kwargs:
                kwargs.update({key: val})
        if isinstance(iterable, Iterable):
            self.iterable = iterable
            if max_value is None:
                self.max_value = len(iterable)
            else:
                self.max_value = max_value
        else:  # ignore max_value then
            self.max_value = int(iterable)
            self.iterable = None
        if max_total_width == 'auto':
            max_total_width = None  # tqdm's auto
        with tempfile.NamedTemporaryFile('wb', prefix='tmp_mne_prog') as tf:
            self._mmap_fname = tf.name
        del tf  # should remove the file
        self._mmap = None
        disable = logger.level > logging.INFO
        self._tqdm = tqdm(
            iterable=self.iterable, desc=mesg, total=self.max_value,
            initial=initial_value, ncols=max_total_width,
            disable=disable, **kwargs)

    def update(self, cur_value):
        """Update progressbar with current value of process.

        Parameters
        ----------
        cur_value : number
            Current value of process.  Should be <= max_value (but this is not
            enforced).  The percent of the progressbar will be computed as
            ``(cur_value / max_value) * 100``.
        """
        self.update_with_increment_value(cur_value - self._tqdm.n)

    def update_with_increment_value(self, increment_value):
        """Update progressbar with an increment.

        Parameters
        ----------
        increment_value : int
            Value of the increment of process.  The percent of the progressbar
            will be computed as
            ``(self.cur_value + increment_value / max_value) * 100``.
        """
        self._tqdm.update(increment_value)

    def __iter__(self):
        """Iterate to auto-increment the pbar with 1."""
        for x in self._tqdm:
            yield x

    def subset(self, idx):
        """Make a joblib-friendly index subset updater.

        Parameters
        ----------
        idx : ndarray
            List of indices for this subset.

        Returns
        -------
        updater : instance of PBSubsetUpdater
            Class with a ``.update(ii)`` method.
        """
        return _PBSubsetUpdater(self, idx)

    def __enter__(self):  # noqa: D105
        # This should only be used with pb.subset and parallelization
        if op.isfile(self._mmap_fname):
            os.remove(self._mmap_fname)
        # prevent corner cases where self.max_value == 0
        self._mmap = np.memmap(self._mmap_fname, bool, 'w+',
                               shape=max(self.max_value, 1))
        self.update(0)  # must be zero as we just created the memmap

        # We need to control how the pickled bars exit: remove print statements
        self._thread = _UpdateThread(self)
        self._thread.start()
        return self

    def __exit__(self, type_, value, traceback):  # noqa: D105
        # Restore exit behavior for our one from the main thread
        self.update(self._mmap.sum())
        self._tqdm.close()
        self._thread._mne_run = False
        self._thread.join()
        self._mmap = None
        if op.isfile(self._mmap_fname):
            os.remove(self._mmap_fname)

    def __del__(self):
        """Ensure output completes."""
        if getattr(self, '_tqdm') is not None:
            self._tqdm.close()

class _UpdateThread(Thread):

    def __init__(self, pb):
        super(_UpdateThread, self).__init__(daemon=True)
        self._mne_run = True
        self._mne_pb = pb

    def run(self):
        while self._mne_run:
            self._mne_pb.update(self._mne_pb._mmap.sum())
            time.sleep(1. / 30.)  # 30 Hz refresh is plenty

class _PBSubsetUpdater(object):

    def __init__(self, pb, idx):
        self.mmap = pb._mmap
        self.idx = idx

    def update(self, ii):
        self.mmap[self.idx[ii - 1]] = True

max_stars_repo_name: mshader/mne-python

max_issues_repo_name: mshader/mne-python

max_forks_repo_name: mshader/mne-python