text:
"""V-trace 테스트 모듈."""

import numpy as np
import torch

import vtrace
from learner import calc_loss

def _shaped_arange(*shape):
    """주어진 모양에 맞게 arange 행렬 생성.

    >>> _shaped_arange([2, 3])
    array([[0., 1., 2.],
           [3., 4., 5.]], dtype=float32)
    """
    return np.arange(np.prod(shape), dtype=np.float32).reshape(*shape)

def _softmax(logits):
    """Apply softmax non-linearity on inputs."""
    return np.exp(logits) / np.sum(np.exp(logits), axis=-1, keepdims=True)

def _ground_truth_calculation(discounts, log_rhos, rewards, values,
                              bootstrap_value, clip_rho_threshold,
                              clip_pg_rho_threshold):
    """Python/Numpy로 V-trace 참값 계산.

    Args:
        discounts: 감쇄율. [T, B] 형태
        log_rhos: 로그 IS 가중치. [T, B] 형태.
        rewards: 행위 정책에 따라 생성된 리워드. [T, B] 형태
        values: 대상 정책에 대한 가치 함수 추정. [T, B] 형태
        bootstrap_value, 시간 T에서 가치 함수 추정. [B] 형태
        clip_rho_threshold: IS(rho) 클리핑 임계치. 논문에서 rho bar
        clip_pg_rho_threshold: V-trace AC에서 정책 모수 갱신 경사식 rho s
    """
    vs = []
    seq_len = len(discounts)
    rhos = np.exp(log_rhos)     # 선형 -> 지수
    cs = np.minimum(rhos, 1.0)  # 지수가 1을 넘지 않
    clipped_rhos = rhos
    if clip_rho_threshold:
        clipped_rhos = np.minimum(rhos, clip_rho_threshold)  # <= 3.7
    clipped_pg_rhos = rhos
    if clip_pg_rho_threshold:
        clipped_pg_rhos = np.minimum(rhos, clip_pg_rho_threshold)  # <= 2.2

    # 이것은 매우 비효율적인 V-trace 참값 계산이지만, 논문의 식1의 수학적 표기에 가깝게 구현.
    # V-trace.
    # v_s = V(x_s)
    #       + \sum^{T-1}_{t=s} \gamma^{t-s}
    #         * \prod_{i=s}^{t-1} c_i
    #         * \rho_t (r_t + \gamma V(x_{t+1}) - V(x_t))
    # 주: 논문의 t-1은 포함되기에 코드에서 s:t로 표기, np.prod([]) == 1

    # t + 1의 value
    values_t_plus_1 = np.concatenate([values, bootstrap_value[None, :]],
                                     axis=0)
    # 전개(unroll) 횟수 만큼
    for s in range(seq_len):
        v_s = np.copy(values[s])  # 중요한 복사
        for t in range(s, seq_len):  # s + n - 1
            delta = clipped_rhos[t] * (rewards[t] + discounts[t] *
                                       values_t_plus_1[t + 1] - values[t])
            # 식에서 감쇄는 power로 되어 있으나 결국 같음.
            v_s += (np.prod(discounts[s:t], axis=0) *
                    np.prod(cs[s:t], axis=0) * delta)
        vs.append(v_s)
    vs = np.stack(vs, axis=0)
    pg_advantages = (
        clipped_pg_rhos * (rewards + discounts * np.concatenate(
            [vs[1:], bootstrap_value[None, :]], axis=0) - values))

    return vtrace.VTraceReturns(vs=vs, pg_advantages=pg_advantages)

def test_vtrace_from_iw():
    """V-trace 중요도 가중치 테스트."""
    batch_size = 1  # 2
    seq_len = 5
    log_rhos = _shaped_arange(seq_len, batch_size) / (batch_size * seq_len)
    log_rhos = 5 * (log_rhos - 0.5)  # [0.0, 1.0) -> [-2.5, 2.5).
    values = {
        'log_rhos': log_rhos,
        # T, B where B_i: [0.9 / (i+1)] * T
        'discounts':
            np.array([[0.9 / (b + 1)
                       for b in range(batch_size)]
                      for _ in range(seq_len)]),
        'rewards':
            _shaped_arange(seq_len, batch_size),
        'values':
            _shaped_arange(seq_len, batch_size) / batch_size,
        'bootstrap_value':
            _shaped_arange(batch_size) + 1.0,
        'clip_rho_threshold':
            3.7,
        'clip_pg_rho_threshold':
            2.2,
    }
    output = vtrace.from_importance_weights(**values)
    ground_truth = _ground_truth_calculation(**values)

    for g, o in zip(ground_truth, output):
        assert np.allclose(g, o.data.tolist())

def test_log_probs_from_logits_and_actions():
    """로짓과 동작에서 로그 확률 얻기 테스트."""
    seq_len = 7
    num_actions = 3
    batch_size = 2
    policy_logits = _shaped_arange(seq_len, batch_size, num_actions) + 10
    np.random.seed(0)

    actions = np.random.randint(
        0, num_actions, size=(seq_len, batch_size), dtype=np.int32)
    action_log_probs = vtrace.log_probs_from_logits_and_actions(
        policy_logits, actions)

    # Ground Truth
    # Using broadcasting to create a mask that indexes action logits
    action_index_mask = actions[..., None] == np.arange(num_actions)

    def index_with_mask(array, mask):
        return array[mask].reshape(*array.shape[:-1])

    # Note: Normally log(softmax) is not a good idea because it's not
    # numerically stable. However, in this test we have well-behaved values.
    ground_truth = index_with_mask(
        np.log(_softmax(policy_logits)), action_index_mask)

    for g, o in zip(ground_truth, action_log_probs):
        assert np.allclose(g, o.data.tolist())

def test_vtrace_from_logit():
    """V-trace를 로짓에서 계산 테스트."""
    seq_len = 5  # n-step
    num_actions = 3
    batch_size = 2
    clip_rho_threshold = None  # No clipping.
    clip_pg_rho_threshold = None  # No clipping.

    np.random.seed(0)
    values = {
        'behavior_policy_logits':
            _shaped_arange(seq_len, batch_size, num_actions),
        'target_policy_logits':
            _shaped_arange(seq_len, batch_size, num_actions),
        'actions':
            np.random.randint(0, num_actions - 1, size=(seq_len, batch_size)),
        'discounts':
            np.array(  # T, B where B_i: [0.9 / (i+1)] * T
                [[0.9 / (b + 1)
                  for b in range(batch_size)]
                 for _ in range(seq_len)]),
        'rewards':
            _shaped_arange(seq_len, batch_size),
        'values':
            _shaped_arange(seq_len, batch_size) / batch_size,
        'bootstrap_value':
            _shaped_arange(batch_size) + 1.0,  # B
    }

    from_logit_output = vtrace.from_logits(
        clip_rho_threshold=clip_rho_threshold,
        clip_pg_rho_threshold=clip_pg_rho_threshold,
        **values)

    ground_truth_target_log_probs = vtrace.log_probs_from_logits_and_actions(
        values['target_policy_logits'], values['actions'])
    ground_truth_behavior_log_probs = vtrace.log_probs_from_logits_and_actions(
        values['behavior_policy_logits'], values['actions'])
    ground_truth_log_rhos = ground_truth_target_log_probs - \
        ground_truth_behavior_log_probs

    from_iw = vtrace.from_importance_weights(
        log_rhos=ground_truth_log_rhos,
        discounts=values['discounts'],
        rewards=values['rewards'],
        values=values['values'],
        bootstrap_value=values['bootstrap_value'],
        clip_rho_threshold=clip_rho_threshold,
        clip_pg_rho_threshold=clip_pg_rho_threshold
    )

    # 중요도 가중치 결과 == 로짓 결과 == ground truth
    for g, o in zip(from_iw.vs, from_logit_output.vs):
        assert np.allclose(g, o.data.tolist())
    for g, o in zip(from_iw.pg_advantages, from_logit_output.pg_advantages):
        assert np.allclose(g, o.data.tolist())
    for g, o in zip(ground_truth_behavior_log_probs,
                    from_logit_output.behavior_action_log_probs):
        assert np.allclose(g, o.data.tolist())
    for g, o in zip(ground_truth_target_log_probs,
                    from_logit_output.target_action_log_probs):
        assert np.allclose(g, o.data.tolist())
    for g, o in zip(ground_truth_log_rhos, from_logit_output.log_rhos):
        assert np.allclose(g, o.data.tolist())

    logits = torch.Tensor(values['behavior_policy_logits'])
    actions = torch.LongTensor(values['actions'])
    advantages = from_iw.pg_advantages
    import pdb; pdb.set_trace()  # breakpoint fd504776 //
    loss = calc_loss(logits, actions, advantages)
    pass


text:
"""
Create forcing file of constant sea ice thickness at 2 m

Notes
-----
    Reference : Kay et al. [2014]
    Author : Zachary Labe
    Date   : 16 August 2017
"""

### Import modules
import numpy as np
from netCDF4 import Dataset
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
import nclcmaps as ncm
import datetime

### Define directories
directorydata = '/surtsey/zlabe/LENS/ForcingPerturb/'
directoryfigure = '/home/zlabe/Desktop/'

### Define time
now = datetime.datetime.now()
currentmn = str(now.month)
currentdy = str(now.day)
currentyr = str(now.year)
currenttime = currentmn + '_' + currentdy + '_' + currentyr
titletime = currentmn + '/' + currentdy + '/' + currentyr
print('\n' '----Calculate forcing file SIT constant - %s----' % titletime)

### Set all values to 2 m
### Used NCO by:
# ncap2 -s 'where(ice_thick>0) ice_thick=2;' SST-SIT_lens_CTL.nc test.nc
# ncap2 -s 'where(ice_thick>0) ice_thick=2;' SST-SIC-SIT_lens_2051-2080_FICT.nc  SST-SIC-SIT_lens_2051-2080_FIC.nc

### Read in data
data = Dataset(directorydata + 'SST-SIC-SIT_lens_2051-2080_FIC.nc')
lon = data.variables['lon'][:]
lat = data.variables['lat'][:]
sit = data.variables['ice_thick'][:]
data.close()

lons,lats = np.meshgrid(lon,lat)

#### Test Data
plt.rc('text',usetex=True)
plt.rc('font',**{'family':'sans-serif','sans-serif':['Avant Garde']})

fig = plt.figure()
ax = plt.subplot(111)

m = Basemap(projection='ortho',lon_0=300,lat_0=90,resolution='l')

var = sit[2]

m.drawmapboundary(fill_color='white')
m.drawcoastlines(color='dimgrey',linewidth=0.3)
parallels = np.arange(-90,90,30)
meridians = np.arange(-180,180,60)
m.drawparallels(parallels,labels=[True,True,True,True],
                linewidth=0.3,color='k',fontsize=6)
m.drawmeridians(meridians,labels=[True,True,True,True],
                linewidth=0.3,color='k',fontsize=6)

cs = m.contourf(lons,lats,var,np.arange(0,3,0.1),latlon=True,extend='both')

m.fillcontinents(color='dimgrey')

cs.set_cmap('cubehelix')

cbar = plt.colorbar(cs,extend='both')

cbar.set_label(r'\textbf{SIT (m)}')
ticks = np.arange(0,5,1)
cbar.set_ticks(ticks)
cbar.set_ticklabels(list(map(str,ticks)))

plt.savefig(directoryfigure + 'test_sitconstant.png',dpi=300)

print('Completed: Script done!')

ext: py

max_stars_repo_name: zmlabe/ThicknessSensitivity

max_issues_repo_name: zmlabe/ThicknessSensitivity

max_forks_repo_name: zmlabe/ThicknessSensitivity

reason: import numpy


text:
import numpy
from scipy.special import hankel1

PI = numpy.pi
TWOPI = 2. * PI
FOURPI = 2. * TWOPI

def incident(kvec, point):
    """
    Incident wave

    @param kvec incident wave vector
    @param point target point
    @return complex number
    """
    return numpy.exp(1j*kvec.dot(point))

def gradIncident(nvec, kvec, point):
    """
    Normal gradient of the incident wave, assumes incident wave is exp(1j * kvec.x)

    @param nvec normal vector pointing inwards
    @param kvec incident wave vector
    @param point (source) point
    @return complex number
    """
    return 1j*nvec.dot(kvec)*incident(kvec, point)

def computeScatteredWaveElement(kvec, p0, p1, point):
    """
    Scattered wave contribution from a single segment
    @param kvec incident wave vector
    @param p0 starting point of the segment
    @param p1 end point of the segment
    @param point observer point
    @return complex value
    """

    # xdot is anticlockwise
    xdot = p1 - p0

    # mid point of the segment
    pmid = 0.5*(p0 + p1)

    # segment length
    dsdt = numpy.sqrt(xdot.dot(xdot))

    # normal vector, pointintg inwards and normalised
    nvec = numpy.array([-xdot[1], xdot[0],])
    nvec /= numpy.sqrt(nvec.dot(nvec))

    # from segment mid-point to observer
    rvec = point - pmid
    r = numpy.sqrt(rvec.dot(rvec))

    kmod = numpy.sqrt(kvec.dot(kvec))
    kr = kmod * r

    # Green functions and normal derivatives
    g = (1j/4.) * hankel1(0, kr)
    dgdn = (-1j/4.) * hankel1(1, kr) * kmod * nvec.dot(rvec) / r

    # contribution from the gradient of the incident wave on the surface
    # of the obstacle. The normal derivative of the scattered wave is
    # - normal derivative of the incident wave.
    scattered_wave = - dsdt * g * gradIncident(nvec, kvec, pmid)

    # shadow side: total wave is nearly zero
    #              => scattered wave amplitude = -incident wave ampl.
    #
    # illuminated side:
    #              => scattered wave amplitude = +incident wave ampl.
    shadow = 2*((nvec.dot(kvec) > 0.) - 0.5) # +1 on the shadow side, -1 on the illuminated side
    scattered_wave += shadow * dsdt * dgdn * incident(kvec, pmid)

    return scattered_wave

def computeScatteredWave(kvec, xc, yc, point):
    """
    Total scattered wave response, summing up
    contributions from each segment

    @param kvec incident wave vector
    @param xc list of x coordinates representing the contour, must close
    @param yc list of y coordinates representing the contour, must close
    @param point observer point
    @return complex value
    """
    res = 0j
    n = len(xc)
    for i0 in range(n - 1):
        p0 = numpy.array([xc[i0], yc[i0],])
        i1 = i0 + 1
        p1 = numpy.array([xc[i1], yc[i1],])
        res += computeScatteredWaveElement(kvec, p0, p1, point)
    return res

max_stars_repo_name: pletzer/scatter

max_issues_repo_name: pletzer/scatter

max_forks_repo_name: pletzer/scatter

reason: import numpy,from scipy


text:
# Open3D: www.open3d.org
# The MIT License (MIT)
# See license file or visit www.open3d.org for details

# examples/Python/Basic/mesh_subdivision.py

import numpy as np
import open3d as o3d

def create_mesh_triangle():
    mesh = o3d.geometry.TriangleMesh()
    mesh.vertices = o3d.utility.Vector3dVector(
        np.array([(np.sqrt(8 / 9), 0, -1 / 3),
                  (-np.sqrt(2 / 9), np.sqrt(2 / 3), -1 / 3),
                  (-np.sqrt(2 / 9), -np.sqrt(2 / 3), -1 / 3)],
                 dtype=np.float32))
    mesh.triangles = o3d.utility.Vector3iVector(np.array([[0, 1, 2]]))
    return mesh

def create_mesh_plane():
    mesh = o3d.geometry.TriangleMesh()
    mesh.vertices = o3d.utility.Vector3dVector(
        np.array([[0, 0, 0], [0, 1, 0], [1, 1, 0], [1, 0, 0]],
                 dtype=np.float32))
    mesh.triangles = o3d.utility.Vector3iVector(np.array([[0, 2, 1], [2, 0,
                                                                      3]]))
    return mesh

def mesh_generator():
    yield create_mesh_triangle()
    yield create_mesh_plane()
    yield o3d.geometry.create_mesh_tetrahedron()
    yield o3d.geometry.create_mesh_box()
    yield o3d.geometry.create_mesh_octahedron()
    yield o3d.geometry.create_mesh_icosahedron()
    yield o3d.geometry.create_mesh_sphere()
    yield o3d.geometry.create_mesh_cone()
    yield o3d.geometry.create_mesh_cylinder()
    yield o3d.io.read_triangle_mesh("../../TestData/knot.ply")
    yield o3d.io.read_triangle_mesh("../../TestData/bathtub_0154.ply")

if __name__ == "__main__":
    np.random.seed(42)

    number_of_iterations = 3

    for mesh in mesh_generator():
        mesh.compute_vertex_normals()
        n_verts = np.asarray(mesh.vertices).shape[0]
        colors = np.random.uniform(0, 1, size=(n_verts, 3))
        mesh.vertex_colors = o3d.utility.Vector3dVector(colors)

        print("original mesh has %d triangles and %d vertices" % (np.asarray(
            mesh.triangles).shape[0], np.asarray(mesh.vertices).shape[0]))
        o3d.visualization.draw_geometries([mesh])

        mesh_up = o3d.geometry.subdivide_midpoint(
            mesh, number_of_iterations=number_of_iterations)
        print("midpoint upsampled mesh has %d triangles and %d vertices" %
              (np.asarray(mesh_up.triangles).shape[0],
               np.asarray(mesh_up.vertices).shape[0]))
        o3d.visualization.draw_geometries([mesh_up])

        mesh_up = o3d.geometry.subdivide_loop(
            mesh, number_of_iterations=number_of_iterations)
        print("loop upsampled mesh has %d triangles and %d vertices" %
              (np.asarray(mesh_up.triangles).shape[0],
               np.asarray(mesh_up.vertices).shape[0]))
        o3d.visualization.draw_geometries([mesh_up])


text:
# general imports
import pytest
from pathlib import Path
import numpy as np
from scipy.interpolate import interp1d

# AHA imports
import magma as m

# msdsl imports
from ..common import *
from msdsl import VerilogGenerator
from msdsl.templates.channel import ChannelModel
from msdsl.rf import s4p_to_step

THIS_DIR = Path(__file__).resolve().parent
BUILD_DIR = THIS_DIR / 'build'
TOP_DIR = THIS_DIR.parent.parent
S4P_FILE = 'peters_01_0605_B1_thru.s4p'
NPTS = 4

def pytest_generate_tests(metafunc):
    pytest_sim_params(metafunc)
    pytest_real_type_params(metafunc, [RealType.FloatReal, RealType.FixedPoint])

@pytest.mark.parametrize(
    'test_pts,chan_type,err_lim', [
        (200, 'exp', 1e-4),
        (200, 's4p', 8e-3)
    ]
)
def test_chan_interp(simulator, real_type, test_pts, err_lim, chan_type,
                     ui=62.5e-12, ui_var=0.25, func_numel=512):
    # make sure test is repeatable
    np.random.seed(0)

    # generate step response data
    if chan_type in {'exp'}:
        t_dur=1e-9
        t_step = np.linspace(0, t_dur, func_numel)
        v_step = 1.0-np.exp(-0.5*t_step/ui)
    elif chan_type in {'s4p'}:
        t_orig, v_orig = s4p_to_step(TOP_DIR/S4P_FILE, dt=0.1e-12, T=10e-9)
        t_step = np.linspace(2e-9, 6e-9, 40001)  # i.e., 0.1 ps spacing
        v_step = interp1d(t_orig, v_orig)(t_step)
        t_step -= t_step[0]
        t_dur = t_step[-1] - t_step[0]
    else:
        raise Exception(f'Unknown chan_type: {chan_type}')

    # determine number of terms for the history
    num_terms = int(round(t_dur/ui))

    # generate model
    model = ChannelModel(
        t_step = t_step,
        v_step = v_step,
        module_name='model',
        dtmax=ui,
        num_spline=NPTS,
        num_terms=num_terms,
        func_numel=func_numel,
        build_dir=BUILD_DIR,
        clk='clk',
        rst='rst',
        real_type=real_type
    )
    model_file = model.compile_to_file(VerilogGenerator())

    # create IO dictionary
    ios = {f'out_{k}': fault.RealOut for k in range(NPTS)}

    # declare circuit
    class dut(m.Circuit):
        name = 'test_chan_interp'
        io = m.IO(
            dt=fault.RealIn,
            in_=fault.RealIn,
            clk=m.In(m.Clock),
            rst=m.BitIn,
            **ios
        )

    # create the tester
    tester = MsdslTester(dut, dut.clk)

    # convenience functions for reading/writing splines
    def get_output_spline():
        retval = []
        for k in range(NPTS):
            retval.append(tester.get_value(getattr(dut, f'out_{k}')))
        return retval

    # define input segments
    inputs = 2*(np.random.randint(0, 2, test_pts)-0.5)
    widths = np.random.uniform(1-ui_var, 1+ui_var, test_pts)*ui
    times = np.concatenate(([0], np.cumsum(widths)))

    # initialize
    tester.poke(dut.dt, 0)
    tester.poke(dut.in_, 0)
    tester.poke(dut.clk, 0)
    tester.poke(dut.rst, 1)
    tester.eval()

    # apply reset
    tester.step(2)

    # clear reset
    tester.poke(dut.rst, 0)
    tester.step(2)

    # run test
    outputs = []
    for width, in_ in zip(widths, inputs):
        tester.poke(dut.dt, width)
        tester.poke(dut.in_, in_)
        tester.eval()
        outputs.append(get_output_spline())
        tester.step(2)

    # run the simulation
    parameters = {
        'in_range': 2.5,
        'out_range': 2.5
    }
    tester.compile_and_run(
        directory=BUILD_DIR,
        simulator=simulator,
        ext_srcs=[model_file, get_file('chan_interp/test_chan_interp.sv')],
        parameters=parameters,
        real_type=real_type,
        #dump_waveforms=True
    )

    # convert measurements to arrays
    meas = [np.array([elem.value for elem in spline]) for spline in outputs]

    # generate expected output
    func = interp1d(t_step, v_step, bounds_error=False, fill_value=(v_step[0], v_step[-1]))
    expt = []
    svec = np.linspace(0, 1, NPTS) * ui
    for i in range(len(inputs)):
        expt.append(np.zeros((NPTS,), dtype=float))
        for j in range(i+1):
            if i==j:
                expt[-1] += inputs[j]*func(svec)
            else:
                expt[-1] += inputs[j]*(func(svec+times[i]-times[j])-func(svec+times[i]-times[j+1]))

    # calculate errors
    errs = [vm-ve for vm, ve in zip(meas, expt)]

    # check errors
    max_err = max([np.max(elem) for elem in errs])
    print(f'max_err: {max_err}')
    assert max_err < err_lim

max_stars_repo_name: sgherbst/msdsl

max_issues_repo_name: sgherbst/msdsl

max_forks_repo_name: sgherbst/msdsl


text:
from numpy import e

def sigmoid(z):
    """computes the sigmoid of z."""

# ====================== YOUR CODE HERE ======================
# Instructions: Compute the sigmoid of each value of z (z can be a matrix,
#               vector or scalar).

# =============================================================
    return g

max_stars_repo_name: cyberbaker/CSML

max_issues_repo_name: cyberbaker/CSML

max_forks_repo_name: cyberbaker/CSML

reason: from numpy


text:
# -*- coding: utf-8 -*-
"""Functions for testing permutation test routines.

Last modified: 22th March 2018.
Author: Tuomas Puoliväli
Email: tuomas.puolivali@helsinki.fi
License: Revised 3-clause BSD
"""

import sys
sys.path.append('../')

import numpy as np

import unittest

from multipy.permutation import _cluster_by_adjacency

class TestPermutation(unittest.TestCase):

    def test_cluster_by_adjacency_one_cluster(self):
        """Test clustering by adjacency with only one cluster."""
        x = np.array([1, 1, 1, 1, 1])
        clusters = _cluster_by_adjacency(x)
        self.assertTrue(np.all(clusters == 1))

    def test_cluster_by_adjacency_no_clusters(self):
        """Test clustering by adjacency with no clusters."""
        x = np.zeros([5, 1])
        clusters = _cluster_by_adjacency(x)
        self.assertTrue(np.all(clusters == 0))

    def test_cluster_by_adjacency_general(self):
        """Test clustering by adjacency for a simple general test case."""
        x = np.array([1, 1, 1, 0, 0, 1, 1, 0, 1])
        clusters = _cluster_by_adjacency(x)
        # Check that the numbering is correct for each cluster.
        self.assertTrue(np.all(clusters[0:3] == 1))
        self.assertTrue(np.all(clusters[5:7] == 2))
        self.assertTrue(clusters[8] == 3)

if __name__ == '__main__':
    unittest.main()

max_stars_repo_name: OrionisBio/multipy

max_issues_repo_name: OrionisBio/multipy

max_forks_repo_name: OrionisBio/multipy


text:
# coding=utf-8
# Copyright 2021 The Tensor2Robot 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.

# Lint as python3
"""TFModel abstract subclasses."""

import abc
import os
from typing import Any, Callable, Dict, Optional, Sequence, Text, Tuple, Union

from absl import flags
from absl import logging
import gin
from gin.tf import utils as gin_utils
import numpy as np
import six
from tensor2robot.models import model_interface
from tensor2robot.models import optimizers
from tensor2robot.preprocessors import abstract_preprocessor
from tensor2robot.preprocessors import noop_preprocessor
from tensor2robot.utils import tensorspec_utils
import tensorflow.compat.v1 as tf
from tensorflow.contrib import framework as contrib_framework
from tensorflow.contrib import tpu as contrib_tpu
from tensorflow.contrib import training as contrib_training

FLAGS = flags.FLAGS
TRAIN = tf.estimator.ModeKeys.TRAIN
EVAL = tf.estimator.ModeKeys.EVAL
PREDICT = tf.estimator.ModeKeys.PREDICT

RunConfigType = Optional[Union[tf.estimator.RunConfig, contrib_tpu.RunConfig]]
ParamsType = Optional[Dict[Text, Any]]

DictOrSpec = Union[Dict[Text, tf.Tensor], tensorspec_utils.TensorSpecStruct]

ModelTrainOutputType = Union[tf.Tensor, Tuple[tf.Tensor, DictOrSpec]]
ExportOutputType = Union[Dict[Text, tf.Tensor], Tuple[
    Dict[Text, tf.Tensor], Dict[Text, tf.estimator.export.PredictOutput]]]
InferenceNetworkOutputsType = Union[DictOrSpec,
                                    Tuple[DictOrSpec,
                                          Optional[Sequence[tf.Tensor]]]]
TransformGradsType = Sequence[Tuple[tf.Tensor, tf.Variable]]

try:
  flags.DEFINE_string('master', '', 'Master for TPU RunConfig')
except flags.DuplicateFlagError:
  pass

DEVICE_TYPE_CPU = 'cpu'
DEVICE_TYPE_GPU = 'gpu'
DEVICE_TYPE_TPU = 'tpu'

gin_configurable_run_config_cls = gin.external_configurable(
    tf.estimator.RunConfig,
    name='tf.estimator.RunConfig',
    denylist=['model_dir'])

gin_configurable_tpu_run_config_cls = gin.external_configurable(
    contrib_tpu.RunConfig,
    name='tf.contrib.tpu.RunConfig',
    denylist=['model_dir', 'tpu_config'])

gin_configurable_tpu_config_cls = gin.external_configurable(
    contrib_tpu.TPUConfig, name='tf.contrib.tpu.TPUConfig')

# Expose the tf.train.Saver to gin.
gin_configurable_saver = gin.external_configurable(
    tf.train.Saver,
    name='tf.train.Saver',
    allowlist=['save_relative_paths', 'var_list', 'allow_empty'])

@gin.configurable
def default_init_from_checkpoint_fn(
    checkpoint,
    allow_partial_restore = False,
    filter_restorables_fn = None):
  """init_from_checkpoint_fn that can be used to init a model from a checkpoint.

  Args:
    checkpoint: String pointing to path of TF checkpoint.
    allow_partial_restore: If True, we allow partial restore, otherwise we raise
      an error if a variable cannot be restored.
    filter_restorables_fn: (Optional) A function that takes a restorable
      TensorFlow variable and returns whether it should be restored or not.
      By default, all restorable variables are updated. Note that
      allow_partial_restore is about how to handle variables are in the
      checkpoint, but not in the graph. The filter_restorables_fn argument
      is about variables that are in the checkpoint and the graph, which we
      don't want to restore into the graph.

  Raises:
    A ValueError if a variable(s) is missing and partial restore is not
    explicitly enabled.
  """
  logging.info('Initializing model weights from %s', checkpoint)
  reader = tf.train.load_checkpoint(checkpoint)
  variables_to_restore = contrib_framework.get_variables()
  assignment_map = {}
  for v in variables_to_restore:
    if filter_restorables_fn is not None and not filter_restorables_fn(v):
      continue
    op_name = v.op.name
    if reader.has_tensor(op_name):
      logging.info('Loading variable %s from checkpoint', op_name)
      assignment_map[op_name] = v
    elif allow_partial_restore:
      logging.warning('Variable %s is not in the checkpoint, skipping.',
                      op_name)
    else:
      raise ValueError('Attempting to restore variable {} which is '
                       'not in the checkpoint.'.format(op_name))
  tf.train.init_from_checkpoint(checkpoint, assignment_map)

class V2SummaryInitHook(tf.estimator.SessionRunHook):
  """Runs v2 summary init op.

  When running code that creates v2 summaries in TF 1.x graph mode, a v2
  summary writer must be created in the same graph, before calling any code
  that would add v2 summaries. You cannot create the V2 summary writer
  in a SessionRunHook, since this runs after the model code.

  When running tf.estimator.train_and_evaluate, the provided model_fn in the
  estimator is called twice, once in mode TRAIN and once in mode EVAL. These
  calls are done in separate graphs. These graphs are not the default graph and
  are created by tf.estimator. You cannot create the V2 summary writer before
  the model_fn, because doing so adds the writer to the default graph, instead
  of the graph tf.estimator creates.

  Therefore, the model must create the summary writer itself, at the same time
  as the model code. We expect the model to expose
  the writer init_op to this SessionRunHook, which initializes the summary
  writer in time for the rest of the model code.

  This is run through a hook instead of a custom init_op because the default
  init_op is not exposed well in tf.estimator. Using a hook makes it easier to
  guarantee both the default init_op and summary writer init_op get run.
  """

  def __init__(self, init_op):
    self.init_op = init_op

  def after_create_session(self, session=None, coord=None):
    session.run(self.init_op)

@gin.configurable
class AbstractT2RModel(
    six.with_metaclass(abc.ABCMeta, model_interface.ModelInterface)):
  """Base class encapsulating a model_fn and metadata about input/output sizes.

  The `T2RModel` abstraction defines a `model_fn` method that can be constructed
  using inputs coming from a tf.data.Dataset, or placeholders. We generate
  these automatically for tf.Example using an input_generator and preprocessor
  if no custom version is provided. The model_fn should not do any
  preprocessing. If any preprocessing besides the raw input tensors is necessary
  please use a custom preprocessor.

  """

  def __init__(self,
               preprocessor_cls=None,
               create_optimizer_fn=optimizers.default_create_optimizer_fn,
               device_type = DEVICE_TYPE_CPU,
               summarize_gradients = True,
               use_sync_replicas_optimizer = False,
               use_avg_model_params = False,
               init_from_checkpoint_fn=None):
    """Base constructor to be used by subclass.

    Args:
      preprocessor_cls: (Optional) A class derived from
        preprocessors.AbstractPreprocessor.
      create_optimizer_fn: A callable function which returns an instance of a
        subclass of tf.train.Optimizer. We couldn't take an optimizer instance
        here because some optimizer's constructor may need to access the
        variables in the graph, which will be created by Estimator when calling
        model_fn. More precisely we will only create an instance during training
        (mode == ModeKeys.TRAIN) within the _optimizer property call which will
        wrap the optimizer instance for GPU towers or TPUs if necessary. The
        _optimizer property is only used within create_train_op function.
      device_type: The device type this model will be deployed on (
        DEVICE_TYPE_CPU, DEVICE_TYPE_GPU, DEVICE_TYPE_TPU).
      summarize_gradients: If True summaries for the gradients produced by the
        train_op will be created. Note, we will automatically disable these
        summaries in case of DEVICE_TYPE_TPU.
      use_sync_replicas_optimizer: If True, synchronize gradient updates from
        the different replicas. (GPU-only, since TPUs are already synchronous).
      use_avg_model_params: During training use a MovingAverageOptimizer and
        swapping saver to compute a running average of the model variables for
        inference.
      init_from_checkpoint_fn: A function that calls
        tf.train.init_from_checkpoint.
    """
    self._preprocessor_cls = preprocessor_cls
    self._create_optimizer_fn = create_optimizer_fn
    self._device_type = device_type
    self._summarize_gradients = summarize_gradients
    self._use_sync_replicas_optimizer = use_sync_replicas_optimizer
    self._sync_replicas_optimizer = None
    self._use_avg_model_params = use_avg_model_params
    self._init_from_checkpoint_fn = init_from_checkpoint_fn
    self._optimizer = None  # type: Optional[tf.train.Optimizer]
    self._scaffold_fn = tf.train.Scaffold

  def create_pack_features(
      self, feature_spec,
      label_spec
  ):
    """Creates a callable function which packs features for model inference.

    Note, it is important that this function is very much independent of the
    model intrinsics. To achieve hermetic models using exported saved
    models this function should not depend on parameters such that previously
    exported models break or have to be aware of such parameters. Pack features
    is required to map from other data representations, e.g. python objects and
    protos to the saved model inputs. In essence, pack features allows to
    extract a feed_dict which is fed to the loaded saved model.

    Args:
      feature_spec: tensorspec_utils.TensorSpecStruct input feature
        specifications of the model.
      label_spec: tensorspec_utils.TensorSpecStruct input label specifications
        of the model.
    """
    raise NotImplementedError()

  def get_feature_specification_for_packing(
      self, mode):
    """Returns the feature_spec that create_pack_features expects."""
    return self.preprocessor.get_in_feature_specification(mode)

  def get_label_specification_for_packing(
      self, mode):
    """Returns the label_spec that create_pack_features expects."""
    return self.preprocessor.get_in_label_specification(mode)

  @property
  def default_preprocessor_cls(self):
    """The default preprocessor class if the user has not provided another one.

    We use this property to define the preprocessor class in order to avoid
    dependencies between the super constructor invocation.

    Returns:
      A preprocessor class.
    """
    return noop_preprocessor.NoOpPreprocessor

  @property
  def preprocessor(self):
    feature_specification_fn = self.get_feature_specification
    label_specification_fn = self.get_label_specification

    preprocessor_cls = self._preprocessor_cls
    if preprocessor_cls is None:
      preprocessor_cls = self.default_preprocessor_cls

    return preprocessor_cls(
        model_feature_specification_fn=feature_specification_fn,
        model_label_specification_fn=label_specification_fn,
        is_model_device_tpu=self.is_device_tpu)

  @property
  def scaffold_fn(self):
    """Returns a scaffold function object for model loading."""
    return self._scaffold_fn

  #############################################################################
  # START DEPRECATED functions which will be removed soon.
  #############################################################################
  def get_eval_hooks(self, config, params):
    """Get eval_hooks to be passed to estimator spec."""
    logging.warning('This function is deprecated and will be replaced.')
    hooks = []
    summary_op = tf.summary.merge_all()
    if summary_op is not None:
      eval_name = 'eval'
      if params is not None:
        eval_name = params.get('eval_name', eval_name)
      hooks = [
          tf.train.SummarySaverHook(
              output_dir=os.path.join(config.model_dir, eval_name),
              save_steps=config.save_summary_steps,
              summary_op=summary_op),
      ]
    return hooks

  #############################################################################
  # END DEPRECATED functions which will be removed soon.
  #############################################################################

  @abc.abstractmethod
  def get_feature_specification(
      self, mode):
    """Required features for the model_fn/model_inference_fn.

    Note, the model_fn might use additional features for debugging/development
    purposes. The create_export_outputs_fn will however only require the
    specified required features. Only this subset of features will be used to
    generate automatic tf.Example extractors and numpy placeholders for the
    serving models.

    Args:
      mode: The mode for feature specifications
    """

  @abc.abstractmethod
  def get_label_specification(
      self, mode):
    """Required labels for the model_fn/model_train_fn/model_eval_fn.

    Note, the model_fn might use additional labels for debugging/development
    purposes.

    Args:
      mode: The mode for feature specifications
    """

  @gin.configurable
  def create_train_op(
      self,
      loss,
      optimizer,
      update_ops = None,
      train_outputs = None,
      filter_trainables_fn = None,
      **kwargs):
    """Create the train_op of from the loss obtained from model_train_fn.

    Args:
      loss: The loss we compute within model_train_fn.
      optimizer: An instance of `tf.train.Optimizer`.
      update_ops: List of update ops to execute alongside the training op.
      train_outputs: (Optional) A dict with additional tensors the training
        model generates.
      filter_trainables_fn: (Optional) A function that takes a trainable
        TensorFlow variable and returns whether it should be updated or not.
        By default, all trainable variables are updated.
      **kwargs: (Optional) Other keyword arguments passed directly to the
        underlying create_train_op function.

    Returns:
      train_op: Op for the training step.
    """
    summarize_gradients = self._summarize_gradients
    if self.is_device_tpu:
      # TPUs don't support summaries up until now. Hence, we overwrite the user
      # provided summarize_gradients option to False.
      if self._summarize_gradients:
        logging.info('We cannot use summarize_gradients on TPUs.')
      summarize_gradients = False
    variables_to_train = None
    if filter_trainables_fn is not None:
      logging.info('Filtering trainable variables')
      variables_to_train = [
          var for var in tf.trainable_variables() if filter_trainables_fn(var)]
      logging.info('Only updating the following trainables:')
      for var in variables_to_train:
        logging.info('  %s', var.name)
    return contrib_training.create_train_op(
        loss,
        optimizer,
        summarize_gradients=summarize_gradients,
        update_ops=update_ops,
        variables_to_train=variables_to_train,
        **kwargs)

  def maybe_init_from_checkpoint(self):
    """Optionally initialize the model from a checkpoint.

    We only automatically initialize from a checkpoint other than the model_dir
    if this function is overloaded with an actual model specific implementation.
    The recommended way to initialize a model from a checkpoint is done via
    tf.train.init_from_checkpoint(ckpt_dir_or_file, assignment_map).
    """
    if self._init_from_checkpoint_fn:
      self._init_from_checkpoint_fn()

  def raise_no_tpu_support(self):
    """Convenience function to raise on tpu request for an unsupported model.

    Raises:
      ValueError: If the model should run on a TPU.
    """
    if self.is_device_tpu:
      raise ValueError('This model {} does not support TPUs'.format(
          self.__name__))

  @abc.abstractmethod
  def inference_network_fn(
      self,
      features,
      labels,
      mode,
      config = None,
      params = None):
    """The inference network implementation.

    This creates the main network based on features.
    Optionally (mode=ModeKeys.TRAIN or ModeKeys.EVAL) the model can do
    additional processing on labels, however, it has to be ensured that this is
    optional and the graph is fully operational without labels. At inference
    time we will have no access to labels. Tensors which are required for loss
    computation or debugging must be put into the inference_outputs dict.
    Having a dedicated inference_network_fn allows to compose new networks by
    using other TFModels.

    Please, use the following pattern to add not supported tpu model components
    such as tf.summary.*
    if self.use_summaries(params):
      # Do operations which are not supported on tpus.

    If your model does not support TPUs at all, please call the following
    function.
    self.raise_no_tpu_support()

    Args:
      features: This is the first item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      labels: This is the second item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      mode: (ModeKeys) Specifies if this is training, evaluation or prediction.
      config: (Optional tf.estimator.RunConfig or contrib_tpu.RunConfig) Will
        receive what is passed to Estimator in config parameter, or the default
        config (tf.estimator.RunConfig). Allows updating things in your model_fn
        based on  configuration such as num_ps_replicas, or model_dir.
      params: An optional dict of hyper parameters that will be passed into
        input_fn and model_fn. Keys are names of parameters, values are basic
        python types. There are reserved keys for TPUEstimator, including
        'batch_size'.

    Returns:
      inference_outputs: A dict with output tensors.
    """

  @abc.abstractmethod
  def model_train_fn(self,
                     features,
                     labels,
                     inference_outputs,
                     mode,
                     config = None,
                     params = None):
    """The training model implementation.

    This model_fn should add the loss computation based on the inference_outputs
    and labels. For better debugging we also provide access to the input
    features. Note, no new variables should be generated in this model_fn since
    the model_inference_fn and the maybe_init_from_checkpoint function would
    not have access to these variables. We output the final loss (scalar) and
    a dict of optional train_outputs which might be useful for the
    model_eval_fn.

    Please, use the following pattern to add not supported tpu model components
    such as tf.summary.*
    if self.use_summaries(params):
      # Do operations which are not supported on tpus.

    If your model does not support TPUs at all, please call the following
    function.
    self.raise_no_tpu_support()

    Args:
      features: This is the first item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      labels: This is the second item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      inference_outputs: A dict containing the output tensors of
        model_inference_fn.
      mode: (ModeKeys) Specifies if this is training, evaluation or prediction.
      config: (Optional tf.estimator.RunConfig or contrib_tpu.RunConfig) Will
        receive what is passed to Estimator in config parameter, or the default
        config (tf.estimator.RunConfig). Allows updating things in your model_fn
        based on  configuration such as num_ps_replicas, or model_dir.
      params: An optional dict of hyper parameters that will be passed into
        input_fn and model_fn. Keys are names of parameters, values are basic
        python types. There are reserved keys for TPUEstimator, including
        'batch_size'.

    Returns:
      loss: The loss we will optimize.
      train_outputs: (Optional) A dict with additional tensors the training
        model generates. We output these tensors such that model_eval_fn could
        introspect these tensors.
    """

  def model_eval_fn(self,
                    features,
                    labels,
                    inference_outputs,
                    train_loss,
                    train_outputs,
                    mode,
                    config = None,
                    params = None):
    """The eval model implementation, by default we report the loss for eval.

    This function should add the eval_metrics computation based on the
    inference_outputs, labels and the train_loss. For better debugging we also
    provide access to the input features and the train_outputs. Note, no new
    variables should be generated in this model_fn since the model_inference_fn
    and the maybe_init_from_checkpoint function would not have access to these
    variables.

    Please, use the following pattern to add not supported tpu model components
    such as tf.summary.*
    if self.use_summaries(params):
      # Do operations which are not supported on tpus.

    If your model does not support TPUs at all, please call the following
    function.
    self.raise_no_tpu_support()

    Args:
      features: This is the first item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      labels: This is the second item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      inference_outputs: A dict containing the output tensors of
        model_inference_fn.
      train_loss: The final loss from model_train_fn.
      train_outputs: A dict containing the output tensors (dict) of
        model_train_fn.
      mode: (ModeKeys) Specifies if this is training, evaluation or prediction.
      config: (Optional tf.estimator.RunConfig or contrib_tpu.RunConfig) Will
        receive what is passed to Estimator in config parameter, or the default
        config (tf.estimator.RunConfig). Allows updating things in your model_fn
        based on  configuration such as num_ps_replicas, or model_dir.
      params: An optional dict of hyper parameters that will be passed into
        input_fn and model_fn. Keys are names of parameters, values are basic
        python types. There are reserved keys for TPUEstimator, including
        'batch_size'.

    Returns:
      eval_metrics: A tuple of (metric_fn, metric_fn_inputs) where metric_fn
        is a dict with {metric_description: tf.metrics.*}.
    """
    del features, labels, inference_outputs, train_loss, train_outputs
    del mode, config, params
    # By default we don't have any eval_metrics. The loss computation used
    # to optimize the model_fn will be reported for the model_eval_fn as well.
    # Hence, by default the EVAL mode can be used to determine the loss
    # performance on the eval dataset or even a larger train dataset.
    return None

  def add_summaries(self,
                    features,
                    labels,
                    inference_outputs,
                    train_loss,
                    train_outputs,
                    mode,
                    config = None,
                    params = None):
    """Add summaries to the graph.

    Having a central place to add all summaries to the graph is helpful in order
    to compose models. For example, if an inference_network_fn is used within
    a while loop no summaries can be added. This function will allow to add
    summaries after the while loop has been processed.

    Args:
      features: This is the first item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      labels: This is the second item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      inference_outputs: A dict containing the output tensors of
        model_inference_fn.
      train_loss: The final loss from model_train_fn.
      train_outputs: A dict containing the output tensors (dict) of
        model_train_fn.
      mode: (ModeKeys) Specifies if this is training, evaluation or prediction.
      config: (Optional tf.estimator.RunConfig or contrib_tpu.RunConfig) Will
        receive what is passed to Estimator in config parameter, or the default
        config (tf.estimator.RunConfig). Allows updating things in your model_fn
        based on  configuration such as num_ps_replicas, or model_dir.
      params: An optional dict of hyper parameters that will be passed into
        input_fn and model_fn. Keys are names of parameters, values are basic
        python types. There are reserved keys for TPUEstimator, including
        'batch_size'.
    """
    del features, labels, inference_outputs, train_loss, train_outputs, mode
    del config
    if not self.use_summaries(params):
      return

  def create_export_outputs_fn(self,
                               features,
                               inference_outputs,
                               mode,
                               config = None,
                               params = None):
    """We export the final output used for model inference.

    This model_fn should create the optional export_outputs, see
    tf.estimator.EstimatorSpec for a more in depth description, and the
    required predictions dict. Note, the predictions dict should more often
    than not be a small subset of the inference_outputs.

    Please, use the following pattern to add not supported tpu model components
    such as tf.summary.*
    if self.use_summaries(params):
      # Do operations which are not supported on tpus.

    If your model does not support TPUs at all, please call the following
    function.
    self.raise_no_tpu_support()

    Args:
      features: This is the first item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      inference_outputs: A dict containing the output tensors of
        model_inference_fn.
      mode: (ModeKeys) Specifies if this is training, evaluation or prediction.
      config: (Optional tf.estimator.RunConfig or contrib_tpu.RunConfig) Will
        receive what is passed to Estimator in config parameter, or the default
        config (tf.estimator.RunConfig). Allows updating things in your model_fn
        based on  configuration such as num_ps_replicas, or model_dir.
      params: An optional dict of hyper parameters that will be passed into
        input_fn and model_fn. Keys are names of parameters, values are basic
        python types. There are reserved keys for TPUEstimator, including
        'batch_size'.

    Returns:
      predictions: A dict of tensors.
      export_outputs: (Optional) A dict containing an arbitrary name for the
        output and tf.estimator.export.PredictOutput(output_dict) as value.
        The output dict is a {name: tensor} mapping. If None, the default
        mapping for predictions is generated. The export_outputs are used
        for the serving model. Multi-headed models should have one name
        per head.
    """
    del features, mode, config, params
    # By default we will export all outputs generated by the
    # inference_network_fn.
    return inference_outputs

  def model_fn(self,
               features,
               labels,
               mode,
               config = None,
               params = None):
    """Estimator model_fn.

    Args:
      features: This is the first item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      labels: This is the second item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      mode: (ModeKeys) Specifies if this is training, evaluation or prediction.
      config: (Optional tf.estimator.RunConfig or contrib_tpu.RunConfig) Will
        receive what is passed to Estimator in config parameter, or the default
        config (tf.estimator.RunConfig). Allows updating things in your model_fn
        based on  configuration such as num_ps_replicas, or model_dir.
      params: An optional dict of hyper parameters that will be passed into
        input_fn and model_fn. Keys are names of parameters, values are basic
        python types. There are reserved keys for TPUEstimator, including
        'batch_size'.

    Raises:
      ValueError: If the mode key is not supported, not in [PREDICT, TRAIN,
        EVAL].

    Returns:
      An EstimatorSpec.
    """
    features = tensorspec_utils.validate_and_pack(
        expected_spec=self.get_feature_specification(mode),
        actual_tensors_or_spec=features,
        ignore_batch=True)
    if labels:
      labels = tensorspec_utils.validate_and_pack(
          expected_spec=self.get_label_specification(mode),
          actual_tensors_or_spec=labels,
          ignore_batch=True)
    inference_outputs = self.inference_network_fn(features, labels, mode,
                                                  config, params)
    update_ops = None
    if isinstance(inference_outputs, tuple):
      if len(inference_outputs) != 2:
        raise ValueError('Unknown output of inference_network_fn: '
                         'tuple of length %d' % len(inference_outputs))
      outputs = inference_outputs[0]
      update_ops = inference_outputs[1]
      inference_outputs = outputs

    if mode == tf.estimator.ModeKeys.PREDICT:
      model_fn_results = self.create_export_outputs_fn(features,
                                                       inference_outputs, mode,
                                                       config, params)
      export_outputs = None
      if isinstance(model_fn_results, tuple):
        predictions = model_fn_results[0]
        export_outputs = model_fn_results[1]
      elif isinstance(model_fn_results, dict):
        export_outputs = {}
        if len(model_fn_results) == 1:
          name, output = list(model_fn_results.items())[0]
          export_outputs[name] = tf.estimator.export.RegressionOutput(output)
        export_outputs[tf.saved_model.signature_constants
                       .DEFAULT_SERVING_SIGNATURE_DEF_KEY] = (
                           tf.estimator.export.PredictOutput(model_fn_results))
        predictions = model_fn_results
      else:
        raise ValueError('The create_export_outputs_fn should return a '
                         'tuple(predictions, export_outputs) or predictions.')

      return tf.estimator.EstimatorSpec(
          mode=mode, predictions=predictions, export_outputs=export_outputs)

    train_fn_result = self.model_train_fn(features, labels, inference_outputs,
                                          mode, config, params)
    if isinstance(train_fn_result, tf.Tensor):
      train_loss = train_fn_result
      train_outputs = {}
    elif isinstance(train_fn_result, tuple):
      train_loss = train_fn_result[0]
      train_outputs = train_fn_result[1]
    else:
      raise ValueError('The model_train_fn should return a '
                       'tuple(loss, train_outputs) or loss.')

    if mode == tf.estimator.ModeKeys.TRAIN:
      # Create the tf.train.Optimizer.
      optimizer = self.create_optimizer(params)

      train_op = self.create_train_op(train_loss, optimizer, update_ops,
                                      train_outputs)

      self.add_summaries(features, labels, inference_outputs, train_loss,
                         train_outputs, mode, config, params)

      # Now the optimizer has been created, therefore, the checkpoint could be
      # initialized.
      # No new variables are allowed to be added, otherwise
      # we would not initialize these variables.
      # Note, this feature is only available for train to bootstrap a model
      # (partially) from a different model. As soon as this checkpoint is
      # written all other modes will use the local checkpoint within model_dir.
      self.maybe_init_from_checkpoint()
      training_hooks = []

      # EstimatorSpec has training_chief_hooks, but TPUEstimatorSpec does not,
      # so we have to use training_hooks here and check is_chief.
      if config and config.is_chief:  # pytype: disable=attribute-error
        training_hooks.append(
            gin_utils.GinConfigSaverHook(
                config.model_dir, summarize_config=True))
        if hasattr(self, 'writer_init_ops'):
          training_hooks.append(V2SummaryInitHook(self.writer_init_ops[mode]))

      # `SyncReplicasOptimizer` needs to attach a training hook.
      if self._sync_replicas_optimizer:
        training_hooks.append(
            self._sync_replicas_optimizer.make_session_run_hook(
                config.is_chief))  # pytype: disable=attribute-error

      # Return the value of the property first since it might be changed.
      scaffold_fn = self.scaffold_fn
      scaffold = scaffold_fn()

      # In order to export asynchronously the saver has to be registered
      # in the graph collection. The scaffold function might register a
      # saver already which is why it is checked here and a saver only
      # added it has none has been added.
      if not tf.get_collection(tf.GraphKeys.SAVERS):
        # TODO(T2R_CONTRIBUTORS): Switch to using gin config for all saver params.
        keep_checkpoint_every_n_hours = None
        max_to_keep = None
        if config is not None:
          keep_checkpoint_every_n_hours = config.keep_checkpoint_every_n_hours
          max_to_keep = config.keep_checkpoint_max
        saver = gin_configurable_saver(
            keep_checkpoint_every_n_hours=keep_checkpoint_every_n_hours,
            max_to_keep=max_to_keep,
        )
        tf.add_to_collection(tf.GraphKeys.SAVERS, saver)
      return tf.estimator.EstimatorSpec(
          mode=mode,
          loss=train_loss,
          train_op=train_op,
          training_hooks=training_hooks,
          scaffold=scaffold)

    if mode == tf.estimator.ModeKeys.EVAL:
      self.add_summaries(features, labels, inference_outputs, train_loss,
                         train_outputs, mode, config, params)

      eval_metrics = self.model_eval_fn(  # pylint: disable=assignment-from-none
          features, labels, inference_outputs, train_loss, train_outputs, mode,
          config, params)
      evaluation_hooks = self.get_eval_hooks(config, params)
      if config and config.is_chief:  # pytype: disable=attribute-error
        eval_name = params.get('eval_name', 'eval')  # pytype: disable=attribute-error
        evaluation_hooks.append(
            gin_utils.GinConfigSaverHook(
                os.path.join(config.model_dir, eval_name),
                summarize_config=True))
        if hasattr(self, 'writer_init_ops'):
          evaluation_hooks.append(V2SummaryInitHook(self.writer_init_ops[mode]))
      return tf.estimator.EstimatorSpec(
          mode=mode,
          loss=train_loss,
          eval_metric_ops=eval_metrics,
          evaluation_hooks=evaluation_hooks)

    raise ValueError('The mode {} is not supported yet.'.format(mode))

  def create_optimizer(self, params):
    """Create the optimizer used for training.

    This function optionally wraps the base optimizer with SyncReplicasOptimizer
    (aggregrate gradients across devices).

    Args:
      params: An optional dict of hyper parameters that will be passed into
        input_fn and model_fn. Keys are names of parameters, values are basic
        python types. There are reserved keys for TPUEstimator, including
        'batch_size'.

    Returns:
      An instance of `tf.train.Optimizer`.
    """

    config = self.get_run_config()
    optimizer = self._create_optimizer_fn(
        use_summaries=self.use_summaries(params))
    if self._use_avg_model_params:
      optimizer = optimizers.create_moving_average_optimizer(optimizer)

      def create_swapping_saver_scaffold(saver=None):
        saver = optimizers.create_swapping_saver(optimizer)
        tf.add_to_collection(tf.GraphKeys.SAVERS, saver)
        return tf.train.Scaffold(saver=saver)

      self._scaffold_fn = create_swapping_saver_scaffold
    if (self._use_sync_replicas_optimizer and (not self.is_device_tpu) and
        config is not None and config.num_worker_replicas > 1):
      optimizer = tf.train.SyncReplicasOptimizer(
          optimizer,
          replicas_to_aggregate=config.num_worker_replicas - 1,
          total_num_replicas=config.num_worker_replicas)
      self._sync_replicas_optimizer = optimizer
    return optimizer

  def host_call_fn(
      self,
      features,
      labels,
      inference_outputs,
      train_loss,
      train_outputs,
      mode,
      config = None,
      params = None
  ):
    """Host call for TPU Estimator model.

    Default to None. Can be passed into TPUEstimatorSpec which will be useful
    to write summries when train on TPU.

    Args:
      features: This is the first item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      labels: This is the second item returned from the input_fn and parsed by
        tensorspec_utils.validate_and_pack. A spec_structure which fulfills the
        requirements of the self.get_feature_specification.
      inference_outputs: A dict containing the output tensors of
        model_inference_fn.
      train_loss: The final loss from model_train_fn.
      train_outputs: A dict containing the output tensors (dict) of
        model_train_fn.
      mode: (ModeKeys) Specifies if this is training, evaluation or prediction.
      config: (Optional tf.estimator.RunConfig or contrib_tpu.RunConfig) Will
        receive what is passed to Estimator in config parameter, or the default
        config (tf.estimator.RunConfig). Allows updating things in your model_fn
        based on  configuration such as num_ps_replicas, or model_dir.
      params: An optional dict of hyper parameters that will be passed into
        input_fn and model_fn. Keys are names of parameters, values are basic
        python types. There are reserved keys for TPUEstimator, including
        'batch_size'.
    Returns:
      None, or the host_call which is a tuple of a function and
      a list or dictionary of tensors
    """
    return None

  def use_summaries(self, params = None):
    """Determine whether or not summaries should be used within this model.

    Note, we cannot simply use members for this operation since we reuse
    models, therefore we have to use a input/output data structure.

    Args:
      params: An optional dict of hyper parameters that will be passed into
        input_fn and model_fn. The key 'use_summaries' will be checked within
        this function and allows models to partially disable summaries during
        execution.

    Returns:
      True if summaries should be used and we are not running on TPUs otherwise
      False.
    """
    if self.is_device_tpu:
      return False
    if params is not None and not params.get('use_summaries', True):
      return False
    return True

  def get_run_config(self):
    """Get the RunConfig for Estimator model.

    Returns:
      tf.estimator.RunConfig() for this model.
    """
    return gin_configurable_run_config_cls(
        session_config=self.get_session_config())

  def get_tpu_run_config(self):
    """Get the TPU RunConfig for Estimator model.

    Returns:
      contrib_tpu.RunConfig() for this model.
    """
    return gin_configurable_tpu_run_config_cls(
        master=FLAGS.master, tpu_config=gin_configurable_tpu_config_cls())

  def get_session_config(self):
    """Get the session config for Estimator model.

    Defaults to None which tells tf.Estimator to use its default session config.
    Not used in TPU jobs at the moment.

    Returns:
      None, or the desired session config.
    """
    return None

  @property
  def is_device_tpu(self):
    return self._device_type == DEVICE_TYPE_TPU

  @property
  def is_device_gpu(self):
    return self._device_type == DEVICE_TYPE_GPU

  @property
  def device_type(self):
    return self._device_type

  @device_type.setter
  def device_type(self, device_type):
    self._device_type = device_type


text:
# Copyright (C) H.R. Oosterhuis 2021.
# Distributed under the MIT License (see the accompanying README.md and LICENSE files).

import argparse
import numpy as np
import time
import json
import pickle

import utils.dataset as dataset
import utils.pretrained_models as prtr
import utils.clicks as clk
import utils.dcg_ips as ips
import utils.ranking as rnk
import utils.evaluate as evl

parser = argparse.ArgumentParser()
parser.add_argument("model_file", type=str,
                    help="Model file output from pretrained model.")
parser.add_argument("input_file", type=str,
                    help="Path to click input.")
parser.add_argument("output_path", type=str,
                    help="Path to output model.")
parser.add_argument("--loss", type=str,
                    help="Loss to optimize.",
                    default='dcg2')
parser.add_argument("--dataset_info_path", type=str,
                    default="datasets_info.txt",
                    help="Path to dataset info file.")

args = parser.parse_args()

print('Reading clicks from: %s' % args.input_file)
with open(args.input_file, 'rb') as f:
  loaded_clicks = pickle.load(f)

train_clicks = loaded_clicks['train']
select_clicks = loaded_clicks['select']
validation_clicks = loaded_clicks['validation']

combined_clicks = clk.add_clicks([loaded_clicks['train'],
                              loaded_clicks['select']],
                              safe_copy=True)

click_model = loaded_clicks['click_model']
eta = loaded_clicks['eta']
dataset_name = loaded_clicks['dataset_name']
fold_id = loaded_clicks['dataset_fold']
num_clicks = loaded_clicks['target_num_clicks']

binarize_labels = 'binarized' in click_model
num_proc = 0

data = dataset.get_dataset_from_json_info(
                  dataset_name,
                  args.dataset_info_path,
                  shared_resource = False,
                )
data = data.get_data_folds()[fold_id]

start = time.time()
data.read_data()
print('Time past for reading data: %d seconds' % (time.time() - start))

pretrain_model = prtr.read_model(args.model_file, data, 1.0)
train_ranking = rnk.data_split_model_rank_and_invert(
                      pretrain_model,
                      data.train
                    )[1]
validation_ranking = rnk.data_split_model_rank_and_invert(
                          pretrain_model,
                          data.validation
                        )[1]

(train_weights,
 train_prop) = clk.compute_weights(
                        train_clicks,
                        data.train,
                        eta)
(_,
 select_prop) = clk.compute_weights(
                        select_clicks,
                        data.train,
                        eta)
(validation_weights,
 validation_prop) = clk.compute_weights(
                        validation_clicks,
                        data.validation,
                        eta)

(combined_weights,
 combined_prop) = clk.compute_weights(
                        combined_clicks,
                        data.train,
                        eta)

epsilon_thres=0.000001
click_order = np.log(np.sum(combined_clicks['num_clicks']))/np.log(10)
if dataset_name == 'Webscope_C14_Set1':
  lr = 10.
  tries = 30+int(10*click_order)
elif dataset_name == 'MSLR-WEB30k':
  lr = 50.
  tries = 30+int(10*click_order)
elif dataset_name == 'istella':
  # Istella
  lr = .01
  tries = 20+int(2*click_order)

print('Training model on training subset for deciding activations.')
train_result = ips.optimize_dcg(
                    args.loss,
                    data,
                    train_weights,
                    validation_weights,
                    learning_rate=lr,
                    learning_rate_decay=1.,
                    trial_epochs=tries,
                    # max_epochs=1000,
                    max_epochs=2,
                    epsilon_thres=epsilon_thres,
                    cutoff=0,
                   )

print('Training model on complete training set for deployment.')
combined_result = ips.optimize_dcg(
                    args.loss,
                    data,
                    combined_weights,
                    validation_weights,
                    learning_rate=lr,
                    learning_rate_decay=1.,
                    trial_epochs=tries,
                    # max_epochs=1000,
                    max_epochs=2,
                    epsilon_thres=epsilon_thres,
                    cutoff=0,
                   )
print('Finished training models.')

general_model = train_result['model']
combined_general_model = combined_result['model']

general_ranking = rnk.data_split_model_rank_and_invert(
                      general_model,
                      data.train
                    )[1]
specialized_ranking = rnk.data_split_rank_and_invert(
                            train_weights,
                            data.train
                          )[1]
combined_general_ranking = rnk.data_split_model_rank_and_invert(
                      combined_general_model,
                      data.train
                    )[1]
combined_specialized_ranking = rnk.data_split_rank_and_invert(
                            combined_weights,
                            data.train
                          )[1]
ideal_ranking = rnk.data_split_rank_and_invert(
                            data.train.label_vector,
                            data.train
                          )[1]

pretrain_docrewards = evl.dcg_reward_per_doc(train_ranking)
general_docrewards = evl.dcg_reward_per_doc(general_ranking)
specialized_docrewards = evl.dcg_reward_per_doc(specialized_ranking)
combined_general_docrewards = evl.dcg_reward_per_doc(combined_general_ranking)
combined_specialized_docrewards = evl.dcg_reward_per_doc(combined_specialized_ranking)

(pretrain_mean,
 pretrain_variance) = evl.mean_variance_reward(
                            pretrain_docrewards,
                            select_prop,
                            select_clicks
                          )

(general_mean,
 general_variance) = evl.mean_variance_reward(
                            general_docrewards,
                            select_prop,
                            select_clicks
                          )

(gendiff_mean,
 gendiff_variance) = evl.mean_variance_reward(
                            general_docrewards-pretrain_docrewards,
                            select_prop,
                            select_clicks
                          )

(pretrain_mean_q,
 pretrain_variance_q) = evl.mean_variance_reward_per_query(
                            pretrain_docrewards,
                            select_prop,
                            select_clicks,
                            data.train,
                          )

(general_mean_q,
 general_variance_q) = evl.mean_variance_reward_per_query(
                            general_docrewards,
                            select_prop,
                            select_clicks,
                            data.train,
                          )

(specialized_mean_q,
 specialized_variance_q) = evl.mean_variance_reward_per_query(
                            specialized_docrewards,
                            select_prop,
                            select_clicks,
                            data.train,
                          )
(specpre_diff_mean_q,
 specpre_diff_variance_q) = evl.mean_variance_reward_per_query(
                            specialized_docrewards - pretrain_docrewards,
                            select_prop,
                            select_clicks,
                            data.train,
                          )
(specgen_diff_mean_q,
 specgen_diff_variance_q) = evl.mean_variance_reward_per_query(
                            specialized_docrewards - general_docrewards,
                            select_prop,
                            select_clicks,
                            data.train,
                          )

true_pretrain_rewards = evl.true_reward_per_query(
                              train_ranking,
                              data.train,
                              binarize=binarize_labels,
                            )
true_general_rewards = evl.true_reward_per_query(
                              general_ranking,
                              data.train,
                              binarize=binarize_labels,
                            )
true_specialized_rewards = evl.true_reward_per_query(
                              specialized_ranking,
                              data.train,
                              binarize=binarize_labels,
                            )
true_combined_general_rewards = evl.true_reward_per_query(
                              combined_general_ranking,
                              data.train,
                              binarize=binarize_labels,
                            )
true_combined_specialized_rewards = evl.true_reward_per_query(
                              combined_specialized_ranking,
                              data.train,
                              binarize=binarize_labels,
                            )
ideal_rewards = evl.true_reward_per_query(
                              ideal_ranking,
                              data.train,
                              binarize=binarize_labels,
                            )

test_ranking = rnk.data_split_model_rank_and_invert(
                          pretrain_model,
                          data.test
                        )[1]

test_general_ranking = rnk.data_split_model_rank_and_invert(
                      combined_general_model,
                      data.test
                    )[1]
test_ideal_ranking = rnk.data_split_rank_and_invert(
                            data.test.label_vector,
                            data.test
                          )[1]
true_pretrain_test_rewards = evl.true_reward_per_query(
                                  test_ranking,
                                  data.test,
                                  binarize=binarize_labels,
                                )
true_general_test_rewards = evl.true_reward_per_query(
                                  test_general_ranking,
                                  data.test,
                                  binarize=binarize_labels,
                                )
ideal_test_rewards = evl.true_reward_per_query(
                                  test_ideal_ranking,
                                  data.test,
                                  binarize=binarize_labels,
                                )
ideal_test_rewards[np.equal(ideal_test_rewards, 0)] = 1.

def mean_ndcg(rewards, ideal_rewards):
  mask = np.greater(ideal_rewards, 0)
  return np.mean(rewards[mask]/ideal_rewards[mask])

def weighted_mean_ndcg(rewards, ideal_rewards, query_weights):
  mask = np.greater(ideal_rewards, 0)
  return np.sum(query_weights[mask]*(rewards[mask]/ideal_rewards[mask]))

def perc_act(activations, ideal_rewards):
  mask = np.greater(ideal_rewards, 0)
  return np.sum(activations[mask])/float(np.sum(mask))

results = {
    'simulation': {
      'loss function': args.loss,
      'click model': click_model,
      'number of clicks': num_clicks,
      'eta': eta,
      'cutoff': 0,
      'dataset': dataset_name,
      'number of train queries': data.train.num_queries(),
      'number of validation queries': data.validation.num_queries(),
    },
    'common results': {
      'pretrained ndcg mean':  mean_ndcg(true_pretrain_rewards, ideal_rewards),
      'general ndcg mean':     mean_ndcg(true_general_rewards, ideal_rewards),
      'specialized ndcg mean': mean_ndcg(true_specialized_rewards, ideal_rewards),
      'combined general ndcg mean':     mean_ndcg(true_combined_general_rewards, ideal_rewards),
      'combined specialized ndcg mean': mean_ndcg(true_combined_specialized_rewards, ideal_rewards),
      'pretrained unseen ndcg mean':  mean_ndcg(true_pretrain_test_rewards, ideal_test_rewards),
      'general unseen ndcg mean':     mean_ndcg(true_general_test_rewards, ideal_test_rewards),
    },
    'SEA confidence selected': {},
    'GENSPEC confidence selected': {},
  }

print('COMMON STATISTICS')
print('Total number of clicks:', train_clicks['num_clicks'])
print('Mean clicks per query:', train_clicks['num_clicks']/float(data.train.num_queries()))
print('pretrain:', mean_ndcg(true_pretrain_rewards, ideal_rewards))
print('general:', mean_ndcg(true_general_rewards, ideal_rewards))
print('specialized:', mean_ndcg(true_specialized_rewards, ideal_rewards))
print('combined general:', mean_ndcg(true_combined_general_rewards, ideal_rewards))
print('combined specialized:', mean_ndcg(true_combined_specialized_rewards, ideal_rewards))
print('unseen pretrain:', mean_ndcg(true_pretrain_test_rewards, ideal_test_rewards))
print('unseen general:', mean_ndcg(true_general_test_rewards, ideal_test_rewards))

min_prop = 1./np.amax(data.train.query_sizes())
for bound_prob in ([0.0, 0.01, 0.5, 0.75, 0.95]):
  pretrain_cb = evl.confidence_bound(
                             bound_prob,
                             pretrain_variance,
                             1.,
                             min_prop,
                             select_clicks
                            )
  general_cb = evl.confidence_bound(bound_prob,
                             general_variance,
                             1.,
                             min_prop,
                             select_clicks
                            )
  gendiff_cb = evl.confidence_bound(bound_prob,
                             gendiff_variance,
                             1.,
                             min_prop,
                             select_clicks
                            )
  pretrain_cb_q = evl.confidence_bound_per_query(
                        bound_prob,
                        pretrain_variance_q,
                        1.,
                        train_prop,
                        select_clicks,
                        data.train,
                      )

  general_cb_q = evl.confidence_bound_per_query(
                        bound_prob,
                        general_variance_q,
                        1.,
                        train_prop,
                        select_clicks,
                        data.train,
                      )
  specialized_cb_q = evl.confidence_bound_per_query(
                        bound_prob,
                        specialized_variance_q,
                        1.,
                        train_prop,
                        select_clicks,
                        data.train,
                      )

  if general_mean - general_cb > pretrain_mean + general_cb:
    general_model_activated = True
    selected_reward = true_combined_general_rewards
    selected_query_mean = general_mean_q
    selected_query_cb = general_cb_q
    true_selected_rewards = true_combined_general_rewards.copy()
    true_selected_test_rewards = true_general_test_rewards
  else:
    general_model_activated = False
    selected_reward = true_pretrain_rewards
    selected_query_mean = pretrain_mean_q
    selected_query_cb = pretrain_cb_q
    true_selected_rewards = true_pretrain_rewards.copy()
    true_selected_test_rewards = true_pretrain_test_rewards

  selected_upper = selected_query_mean + selected_query_cb
  specialized_lower = specialized_mean_q - specialized_cb_q

  spec_act = np.greater(specialized_lower, selected_upper)
  true_selected_rewards[spec_act] = true_combined_specialized_rewards[spec_act]

  print('CONFIDENCE', bound_prob)
  print('SEA BOUNDS')
  print('confidence bound:', bound_prob)
  print('General model activated:', general_model_activated)
  print('Num queries specialized:', np.sum(spec_act))
  print('Perc. queries specialized:', np.sum(spec_act)/float(spec_act.size))
  print('Cor. Perc. queries specialized:', perc_act(spec_act, ideal_rewards))
  print('selected:', mean_ndcg(true_selected_rewards, ideal_rewards))
  print('unseen selected:', mean_ndcg(true_selected_test_rewards, ideal_test_rewards))

  results['SEA confidence selected'][bound_prob] = {
    'confidence': bound_prob,
    'general activated': general_model_activated,
    'number of activations': int(np.sum(spec_act)),
    'percentage of activations': np.sum(spec_act)/float(data.train.num_queries()),
    'corrected percentage of activations': perc_act(spec_act, ideal_rewards),
    'selected ndcg mean': mean_ndcg(true_selected_rewards, ideal_rewards),
    'selected unseen ndcg mean': mean_ndcg(true_selected_test_rewards, ideal_test_rewards),
  }
  if gendiff_mean - gendiff_cb > 0:
    general_model_activated = True
    selected_reward = true_combined_general_rewards
    selected_query_mean = specgen_diff_mean_q
    selected_query_cb = evl.confidence_bound_per_query(
                              bound_prob,
                              specgen_diff_variance_q,
                              1.,
                              train_prop,
                              select_clicks,
                              data.train,
                            )
    true_selected_rewards = true_combined_general_rewards.copy()
    true_selected_test_rewards = true_general_test_rewards

  else:
    general_model_activated = False
    selected_reward = true_pretrain_rewards
    selected_query_mean = specpre_diff_mean_q
    selected_query_cb = evl.confidence_bound_per_query(
                              bound_prob,
                              specpre_diff_variance_q,
                              1.,
                              train_prop,
                              select_clicks,
                              data.train,
                            )
    true_selected_rewards = true_pretrain_rewards.copy()
    true_selected_test_rewards = true_pretrain_test_rewards

  spec_act = np.greater(selected_query_mean - selected_query_cb, 0)
  true_selected_rewards[spec_act] = true_combined_specialized_rewards[spec_act]

  print('GENSPEC BOUND')
  print('confidence bound:', bound_prob)
  print('General model activated:', general_model_activated)
  print('Num queries specialized:', np.sum(spec_act))
  print('Perc. queries specialized:', np.sum(spec_act)/float(spec_act.size))
  print('Cor. Perc. queries specialized:', perc_act(spec_act, ideal_rewards))
  print('selected:', mean_ndcg(true_selected_rewards, ideal_rewards))
  print('unseen selected:', mean_ndcg(true_selected_test_rewards, ideal_test_rewards))

  results['GENSPEC confidence selected'][bound_prob] = {
    'confidence': bound_prob,
    'general activated': general_model_activated,
    'number of activations': int(np.sum(spec_act)),
    'percentage of activations': np.sum(spec_act)/float(data.train.num_queries()),
    'corrected percentage of activations': perc_act(spec_act, ideal_rewards),
    'selected ndcg mean': mean_ndcg(true_selected_rewards, ideal_rewards),
    'selected unseen ndcg mean': mean_ndcg(true_selected_test_rewards, ideal_test_rewards),
  }

print('Writing results to %s' % args.output_path)
with open(args.output_path, 'w') as f:
  json.dump(results, f)


text:
import cv2
import numpy as np
import os
# from sys import argv
import time

# main direction (program)
def get_parent_dir(n=1):
    """returns the n-th parent directory of the current
    working directory"""
    current_path = os.path.dirname(os.path.abspath(__file__))
    for _ in range(n):
        current_path = os.path.dirname(current_path)
    return current_path

def clear_coordinates():
    coordinates_files = open("coordinates.txt", "w")
    coordinates_files.write("")
    coordinates_files.close()

def output_coordinates(x, y, centre_x, centre_y):
    coordinate_file = open("coordinates.txt", "a+")
    coordinate_file.writelines([str(x),"\t", str(y),"\t", str(centre_x), "\t",str(centre_y), "\n"])
    coordinate_file.close()

clear_coordinates()

number = 0
# initialising directories
detection_folder = os.path.join(get_parent_dir(1), "1_Detection")
camera_inputs = os.path.join(detection_folder, "Camera_Inputs")
annotated_images = os.path.join(detection_folder, "Annotated_Images")
trained_yolo = os.path.join(get_parent_dir(1), "Trained_YOLO")
# annotated_images = os.path.join(detection_folder, "Annotated_Images")

# Load Yolo
weights = os.path.join(trained_yolo, "yolov4-custom.weights")
cfg = os.path.join(trained_yolo, "custom-yolov4-detector.cfg")
net = cv2.dnn.readNet(weights, cfg)
# data_classes = os.path.join(trained_yolo, "coco.names")

# labels = []
# with open(data_classes, "r") as f:
#     labels = [line.strip() for line in f.readlines()]
labels = ["dice", "dice face"]
layer_names = net.getLayerNames()
output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]
colors = np.random.uniform(0, 255, size=(len(labels), 3))
# colors = np.random.uniform(0, 255, size = (36, 3))

print(labels)
# Load Images from input_images folder
# folder  = "input_images"
# filename = "dice_2.jpg"
# img = cv2.imread(filename)
# img = cv2.resize(img, None, fx=0.6, fy=0.6)
# height, width, channels = img.shape

# This is for webcam
webcam = cv2.VideoCapture(0)

while True:
    # This is for webcam
    time.sleep(0.5)
    net, img = webcam.read()
    height, width, channels = img.shape

    # Detecting objects
    blob = cv2.dnn.blobFromImage(img, 0.00392, (416, 416), (0, 0, 0), True, crop=False)
    net.setInput(blob)
    outs = net.forward(output_layers)

    # Showing information on the screen
    class_ids = []
    confidences = []
    boxes = []
    centers_x = []
    centers_y = []
    xs = []
    ys = []
    dice_iteration = 0
    for out in outs:
        for detection in out:
            scores = detection[5:]
            class_id = np.argmax(scores)
            confidence = scores[class_id]
            if confidence > 0.5:
                # Object detected
                center_x = int(detection[0] * width)
                center_y = int(detection[1] * height)
                w = int(detection[2] * width)
                h = int(detection[3] * height)
                # Rectangle coordinates
                x = int(center_x - w / 2)
                y = int(center_y - h / 2)
                boxes.append([x, y, w, h])
                confidences.append(float(confidence))
                class_ids.append(class_id)
                # output_coordinates(x,y,center_x,center_y)
                # coordinates = [ x, y, center_x, center_y]
                centers_x.append(center_x)
                centers_y.append(center_y)
                xs.append(x)
                ys.append(y)

    indexes = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)
    font = cv2.FONT_HERSHEY_SIMPLEX
    for i in range(len(boxes)):
        if i in indexes:
            x, y, w, h = boxes[i]
            label = str(labels[class_ids[i]])
            # label = labels
            # color = (0,0,0)
            if label == "dice":
                color = (255,0,0)
                output_coordinates(xs[i], ys[i], centers_x[i], centers_y[i])
            else:
                color = (0,255,0)
            cv2.rectangle(img, (x, y), (x + w, y + h), color, 2)
            cv2.rectangle(img, (x, y), (x + 50*len(label), y-30), color, -1)
            cv2.putText(img, label +",{:.2f}".format(confidences[i]), (x, y-5), font, 1, (255,255,255), 1)
            print(label + ": {:.2f}".format(confidences[i]))
            number = number + 1
            outfile = os.path.join(camera_inputs, 'dice_%s.jpg' % str(number))
            cv2.imwrite(outfile, img)
            if label == "dice":
                dice_iteration = dice_iteration + 1
                print("im here")
                output_file = os.path.join(annotated_images, "dice_%s.jpg" % str(dice_iteration))
                image = cv2.imread("dice_2.jpg")
                crop_img = image[y:y+h, x:x + w]
                cv2.imwrite(output_file, crop_img)
                # cv2.imshow("crop", crop_img)

    # clear_coordinates()
    # folder = "output_images"
    # if cv2.waitKey(1) & 0xFF == ord('p'):

    # for webcam
    if cv2.waitKey(1) & 0xFF == ord('q'):
        webcam.release()
        cv2.destroyAllWindows()

    cv2.imshow("Image", img)
    cv2.waitKey(0)
    # for images
    # if cv2.waitKey(1) & 0xFF == ord('l'):
        # img.release()
        # cv2.destroyAllWindows()


text:
# coding: utf-8

# In[30]:

from numpy import *
import pickle

# In[3]:

import sys
sys.path.append('../')
sys.path.append('../support/')
sys.path.append('../lung_segmentation/')
import os
from preprocessing import *
from ct_reader import *
from scipy.ndimage import morphology
from tqdm import tqdm
import time
from os.path import join, basename, isfile
from multiprocessing import Pool
from scipy.ndimage import label
import scipy.ndimage.filters as filters
from glob import glob
from paths import *
import functools

from scipy.linalg import norm
from scipy.ndimage.filters import gaussian_filter, laplace
import pandas as pd

# In[4]:

def read_ct(path, ret_ct_scan=False, ret_xy_spacing=False):
    patient = read_ct_scan(path)
    if ret_ct_scan:
        return get_pixels_hu(patient), patient
    if ret_xy_spacing:
        return get_pixels_hu(patient), patient.GetSpacing()[0]

    return get_pixels_hu(patient)

# In[5]:

SPACING = array([1., 1., 1.])
ISOLATED_THRESHOLD = -600
DOT_ENHANCED_THRESHOLD = 6
OUT_LUNGS = -9000
BOTTOM = 4
BORDER = 32
TOP = 16
CPU = 6
FILTERS_AMOUNT = 6
ISOLATED_MIN_VOLUME = 9
ISOLATED_MAX_VOLUME = 500
JUXTAVASCULAR_MIN_VOLUME = 9
JUXTAPLEURAL_MIN_VALUME = 1

# In[6]:

def get_scales(bottom=BOTTOM, top=TOP,
               amount=FILTERS_AMOUNT):
    radius = (top / bottom) ** (1. / (amount - 1))
    sigmas = [bottom / 4.]
    for i in range(amount - 1):
        sigmas.append(sigmas[0] * (radius ** i + 1))
    return sigmas

# In[7]:

def load_data(patient_path, mask_path):
    ct_scan = read_ct_scan(patient_path)
    mask = load(mask_path)
    patient = get_pixels_hu(ct_scan)
    patient, spacing = resample(patient, ct_scan, SPACING)
    mask, spacing = resample(mask, ct_scan, SPACING)

    mask = morphology.binary_fill_holes(
        morphology.binary_dilation(
            morphology.binary_fill_holes(mask > 0),
            iterations=4)
    )

    return patient, mask

# In[ ]:

paths = PATH['STAGE_ENHANCED']
data_path = PATH['STAGE_DATA']

patients = list(sorted(glob(join(PATH['STAGE_DATA'], '*')), reverse=True))
processed = [basename(path).split('amounts.npy')[0]
             for path in glob(join(paths,'*amounts.npy'))]

patients = [patient
            for patient in patients
           if basename(patient) not in processed]

erroneus = pickle.load(open(join(PATH['STAGE_MASKS'], 'erroneus.pkl'), 'rb'))

patients = [patient
            for patient in patients
           if basename(patient) not in erroneus]

erroneus = set(erroneus)
erroneus = erroneus.difference(pickle.load(open(join(PATH['WEIGHTS'],
                                                    'still_erroneus'), 'rb')))
patients = [data_path + fixed for fixed in list(erroneus)]

def region_growing(img, seed, minthr, maxthr, structure=None):
    """code was taken from:
    https://github.com/loli/medpy/wiki/Basic-image-manipulation
    """
    img[seed] = minthr
    thrimg = (img <= maxthr) & (img >= minthr)
    lmap, _ = label(thrimg, structure=structure)
    lids = unique(lmap[seed])
    region = zeros(img.shape, bool)
    for lid in lids:
        region |= lmap == lid
    return region

def find_closest_label(lung, enh):
    center = zeros(lung.shape)
    center[lung.shape[0] // 2, lung.shape[1] // 2, lung.shape[2] // 2] = 1
    ans = (enh * center).max()
    while ans == 0:
        center = morphology.binary_dilation(center)
        ans = (enh * center).max()
    return ans

def segment_nodule(btch, mask):
    batch = btch.copy()
    interval = [64, 64]
    grown = [mask == find_closest_label(batch,mask)] * 2
    minval = median(batch[grown[-1]])
    maxval = median(batch[grown[-1]])
    sums = [grown[0].sum()] * 2
    while True: #unique((mask == mask[32][32]) == grown).shape[0] == 1:
        grown.append(region_growing(batch,
                                    grown[-1],
                                    minval,
                                    maxval))
        minval -= interval[0]
        maxval += interval[1]
        grown.pop(0)
        sums.append(grown[-1].sum())
        sums.pop(0)
        if 2 * sums[0] < sums[1]:
            interval[0] = interval[0] // 2
            interval[1] = interval[1] // 2
            grown = [grown[0]] * 2
            sums = [grown[0].sum()] * 2
            if interval[0] == 1:
                break
        if grown[0].sum() > 27000:
            break
    return grown[0]

def overlap(lung, mask):
    labeled, num = label(mask)
    coords = list()
    for colour in range(1, labeled.max() + 1):
        coords.append(where(labeled == colour))

    coords = array([[int(coord[0].mean())
                     for coord in coords],
                    [int(coord[1].mean())
                     for coord in coords],
                    [int(coord[2].mean())
                     for coord in coords]])

    pads = ((BORDER, BORDER),
            (BORDER, BORDER),
            (BORDER, BORDER))

    lung = pad(lung, pads,
               mode='constant', constant_values=OUT_LUNGS)
    res = zeros(lung.shape)
    labeled = pad(labeled, pads,
                  mode='constant', constant_values=OUT_LUNGS)
    patches = list()
    masks = list()
    for coord in coords.T:
        res[coord[0]: coord[0] + 2 * BORDER,
            coord[1]: coord[1] + 2 * BORDER,
            coord[2]: coord[2] + 2 * BORDER] += \
            segment_nodule(lung[coord[0]: coord[0] + 2 * BORDER,
                                coord[1]: coord[1] + 2 * BORDER,
                                coord[2]: coord[2] + 2 * BORDER],
                           labeled[coord[0]: coord[0] + 2 * BORDER,
                                   coord[1]: coord[1] + 2 * BORDER,
                                   coord[2]: coord[2] + 2 * BORDER])

    return res[BORDER: -BORDER,
               BORDER: -BORDER,
               BORDER: -BORDER]

print(len(patients))

for i, patient_path in enumerate(patients):
    mask_path = join(PATH['STAGE_MASKS'], basename(patient_path) + '.npy')
    patient, mask = load_data(patient_path, mask_path)

    nodules = load(join(paths, basename(patient_path) + '.npy'))
    patient += OUT_LUNGS * (mask == 0)
    save(mask_path, mask + 16 * overlap(patient, nodules))

max_stars_repo_name: vessemer/LungCancerDetection

max_issues_repo_name: vessemer/LungCancerDetection

max_forks_repo_name: vessemer/LungCancerDetection

reason: from numpy,import scipy,from scipy


text:
from serial import Serial, EIGHTBITS, PARITY_NONE, STOPBITS_ONE
from time import sleep
import sys
import cv2
import matplotlib.pyplot as plt
import numpy as np
from skimage import exposure
import os
import json
import tensorflow as tf
from keras.models import Model
from sklearn.metrics.pairwise import cosine_similarity
from model_hercules import Center_mFCN3_relu_max2_WID_fz1_2
import time
COM_PORT = 'COM3'
#COM_PORT = '/dev/serial/by-id/usb-Digilent_Digilent_USB_Device_251642542476-if00-port0'
#BAUD_RATES = 115200
#BAUD_RATES = int(sys.argv[1])
model_path = 'lcfr.hdf5'
input_shape = (64, 64, 1)
model = Center_mFCN3_relu_max2_WID_fz1_2(input_shape=input_shape, classes=10582, bn_axis=-1, isCenterloss=False)
model.load_weights(model_path, by_name=True)
FC1_vector = Model(model.input, model.get_layer('new_FV').output)

#preprae database
with open('database.json') as f:
	cha_arr = json.load(f)
s = Serial(
		port='COM3',
		baudrate=1000000,
		bytesize=EIGHTBITS,
		parity=PARITY_NONE,
		stopbits=STOPBITS_ONE,
		xonxoff=False,
		rtscts=False
	)
print("Configured")
while True:
	#mcu_feedback = s.readline().decode()
	# mcu_feedback = ser.read(1)
	b1, b2, b3 = b'\x00', b'\x00', b'\x00'
	RGB = True

	img = np.zeros((64, 64, 3)).astype(np.uint8)

	while True:
		b1, b2, b3 = b2, b3, s.read(1)

		if b1 == b'R' and b2 == b'D' and b3 == b'Y':
			start = time.time();
			print("GET")
			cnt = 0
			while True:
				if(not RGB):
					b1 = s.read(1)
					b2 = s.read(1)
					b1i = int.from_bytes(b1, byteorder='little')
					b2i = int.from_bytes(b2, byteorder='little')
					R = b1i // 8
					tmp = b1i % 8
					G = tmp * 8 + b2i // 32
					B = b2i % 32
				else:
					b1 = s.read(1)
					b2 = s.read(1)
					b3 = s.read(1)
					R = int.from_bytes(b1, byteorder='little')
					G = int.from_bytes(b2, byteorder='little')
					B = int.from_bytes(b3, byteorder='little')

				#print(R, G, B)
                # print(b1, b2);

				img[cnt // 64][cnt % 64][2] = np.clip(R * 8, 0, 255)
				img[cnt // 64][cnt % 64][1] = np.clip((R+B)*4, 0, 255)
				img[cnt // 64][cnt % 64][0] = np.clip(B * 8, 0, 255)

				cnt = cnt +1
				if cnt == 64 * 64:
					break
			imgs = []
			img3 = cv2.cvtColor(img,cv2.COLOR_BGR2GRAY)
			cv2.imwrite('color_img.jpg', img)

			#img2 = exposure.equalize_adapthist(img, clip_limit=0.03)
			img2 = cv2.resize(img, (240, 240))
			#cv2.imshow("Image", img2)

			b = s.read(1)
			bi = int.from_bytes(b, byteorder='little')
			if(bi>0):
				print("\n--------------------Face detected------------------------\n")
				img3 = (img3-127.5)/127.5
				img3 = np.reshape(img3, [64,64,1])
				imgs.append(img3)
				imgs = np.array(imgs)
				vec0 = FC1_vector.predict(imgs)
				vec0 = vec0[0][0][0].tolist()

				#recognition
				best_score = 0.5
				best_name = 'nobody'
				for key, value in cha_arr.items():
					vec1 = value
					score = cosine_similarity([vec0], [vec1])
					#print(key + ' ' + str(score))
					if score > best_score:
						best_score = score
						best_name = key
				#print("inferencce time:"+str(end-start)+"\n")
				print('\n')
				print('+------------------------------------------------->')
				print('|   Hi ' + best_name + ', confidence is ' + str(best_score))

				print('+------------------------------------------------->')
				cv2.putText(img2, "Hello "+best_name, (80,210), cv2.FONT_HERSHEY_DUPLEX, 0.7, (0, 255, 255), 1, cv2.LINE_AA)
			else:
				print("\n--------------------No face-----------------------\n")
				cv2.putText(img2, "No Face", (80,210), cv2.FONT_HERSHEY_DUPLEX, 0.7, (0, 255, 255), 1, cv2.LINE_AA)
			cv2.imshow("Image", img2)
			cv2.waitKey(1)
			#print("frame done \n")
			'''
			b = s.read(1)
			bi = int.from_bytes(b, byteorder='little')
			if(bi>0):
				print("Face detected\n")
			else:
				print("no face\n")
			'''
	#print(mcu_feedback,)
if __name__ == "__main__":
	main()


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

cap = cv2.VideoCapture("v_CricketShot_g04_c01.avi")
plt.ion()
ret, frame1 = cap.read()
prvs = cv2.cvtColor(frame1,cv2.COLOR_BGR2GRAY)
hsv = np.zeros_like(frame1)
p_flow = cv2.DualTVL1OpticalFlow_create()
gif = []
while True:
	ret, frame2 = cap.read()
	if frame2 is None:
		break
	next = cv2.cvtColor(frame2,cv2.COLOR_BGR2GRAY)

	flow = p_flow.calc(prvs, next, None)
	flow[flow > 20] = 20
	flow[flow < -20] = -20
	flow = flow / 20.
	flow = np.c_[flow, np.zeros((240, 320, 1))] + 0.5
	gif.append(flow)
	prvs = next

cap.release()

imageio.mimsave('1.gif', gif)

max_stars_repo_name: dythebs/repo

max_issues_repo_name: dythebs/repo

max_forks_repo_name: dythebs/repo