text:
import itertools
import collections
from functools import partial

import pytest
import numpy as np

from jax import dtypes, jit, custom_jvp, lax, test_util as jtu

from fastar import test_util as tu
from fastar.core import Ones
from fastar.rules import conv_incounts, pad_dependency_rule

# This is borrowed from lax_tests.py in the JAX tests directory.
# TODO import directly from lax_tests.py

def supported_dtypes(dtypes):
  return [t for t in dtypes if t in jtu.supported_dtypes()]

float_dtypes = supported_dtypes([dtypes.bfloat16, np.float16, np.float32,
                                 np.float64])
complex_elem_dtypes = supported_dtypes([np.float32, np.float64])
complex_dtypes = supported_dtypes([np.complex64, np.complex128])
inexact_dtypes = float_dtypes + complex_dtypes
int_dtypes = supported_dtypes([np.int32, np.int64])
uint_dtypes = supported_dtypes([np.uint32, np.uint64])
bool_dtypes = [np.bool_]
default_dtypes = float_dtypes + int_dtypes
number_dtypes = default_dtypes + complex_dtypes
all_dtypes = number_dtypes + bool_dtypes

compatible_shapes = [[(3,)], [(3, 4), (3, 1), (1, 4)], [(2, 3, 4), (2, 1, 4)]]

CombosWithReplacement = itertools.combinations_with_replacement

OpRecord = collections.namedtuple(
    "OpRecord", ["op", "nargs", "dtypes", "rng_factory", "tol"])

def op_record(op, nargs, dtypes, rng_factory, tol=None):
  return OpRecord(op, nargs, dtypes, rng_factory, tol)
LAX_OPS = [
    op_record("neg", 1, default_dtypes + complex_dtypes, jtu.rand_small),
    op_record("sign", 1, default_dtypes + uint_dtypes, jtu.rand_small),
    op_record("floor", 1, float_dtypes, jtu.rand_small),
    op_record("ceil", 1, float_dtypes, jtu.rand_small),
    op_record("round", 1, float_dtypes, jtu.rand_default),
    op_record("nextafter", 2, [f for f in float_dtypes if f != dtypes.bfloat16],
              jtu.rand_default, tol=0),

    op_record("is_finite", 1, float_dtypes, jtu.rand_small),

    op_record("exp", 1, float_dtypes + complex_dtypes, jtu.rand_small),
    # TODO(b/142975473): on CPU, expm1 for float64 is only accurate to ~float32
    # precision.
    op_record("expm1", 1, float_dtypes + complex_dtypes, jtu.rand_small,
              {np.float64: 1e-8}),
    op_record("log", 1, float_dtypes + complex_dtypes, jtu.rand_positive),
    op_record("log1p", 1, float_dtypes + complex_dtypes, jtu.rand_positive),
    # TODO(b/142975473): on CPU, tanh for complex128 is only accurate to
    # ~float32 precision.
    # TODO(b/143135720): on GPU, tanh has only ~float32 precision.
    op_record("tanh", 1, float_dtypes + complex_dtypes, jtu.rand_small,
              {np.float64: 1e-9, np.complex128: 1e-7}),
    op_record("sin", 1, float_dtypes + complex_dtypes, jtu.rand_default),
    op_record("cos", 1, float_dtypes + complex_dtypes, jtu.rand_default),
    op_record("atan2", 2, float_dtypes, jtu.rand_default),

    op_record("sqrt", 1, float_dtypes + complex_dtypes, jtu.rand_positive),
    op_record("rsqrt", 1, float_dtypes + complex_dtypes, jtu.rand_positive),
    op_record("square", 1, float_dtypes + complex_dtypes, jtu.rand_default),
    op_record("reciprocal", 1, float_dtypes + complex_dtypes, jtu.rand_positive),
    op_record("tan", 1, float_dtypes, jtu.rand_default, {np.float32: 3e-5}),
    op_record("asin", 1, float_dtypes, jtu.rand_small),
    # TODO(j-towns) fix: op_record("acos", 1, float_dtypes, jtu.rand_small),
    op_record("atan", 1, float_dtypes, jtu.rand_small),
    op_record("asinh", 1, float_dtypes, jtu.rand_default),
    op_record("acosh", 1, float_dtypes, jtu.rand_positive),
    # TODO(b/155331781): atanh has only ~float precision
    op_record("atanh", 1, float_dtypes, jtu.rand_small, {np.float64: 1e-9}),
    op_record("sinh", 1, float_dtypes + complex_dtypes, jtu.rand_default),
    op_record("cosh", 1, float_dtypes + complex_dtypes, jtu.rand_default),
    op_record("lgamma", 1, float_dtypes, jtu.rand_positive,
              {np.float32: 1e-3 if jtu.device_under_test() == "tpu" else 1e-5,
               np.float64: 1e-14}),
    op_record("digamma", 1, float_dtypes, jtu.rand_positive,
              {np.float64: 1e-14}),
    op_record("betainc", 3, float_dtypes, jtu.rand_positive,
              {np.float64: 1e-14}),
    op_record("igamma", 2,
              [f for f in float_dtypes if f not in [dtypes.bfloat16, np.float16]],
              jtu.rand_positive, {np.float64: 1e-14}),
    op_record("igammac", 2,
              [f for f in float_dtypes if f not in [dtypes.bfloat16, np.float16]],
              jtu.rand_positive, {np.float64: 1e-14}),
    op_record("erf", 1, float_dtypes, jtu.rand_small),
    op_record("erfc", 1, float_dtypes, jtu.rand_small),
    # TODO(b/142976030): the approximation of erfinf used by XLA is only
    # accurate to float32 precision.
    op_record("erf_inv", 1, float_dtypes, jtu.rand_small,
              {np.float64: 1e-9}),
    op_record("bessel_i0e", 1, float_dtypes, jtu.rand_default),
    op_record("bessel_i1e", 1, float_dtypes, jtu.rand_default),

    op_record("real", 1, complex_dtypes, jtu.rand_default),
    op_record("imag", 1, complex_dtypes, jtu.rand_default),
    op_record("complex", 2, complex_elem_dtypes, jtu.rand_default),
    op_record("conj", 1, complex_elem_dtypes + complex_dtypes,
              jtu.rand_default),
    op_record("abs", 1, default_dtypes + complex_dtypes, jtu.rand_default),
    op_record("pow", 2, float_dtypes + complex_dtypes, jtu.rand_positive),

    op_record("bitwise_and", 2, bool_dtypes, jtu.rand_small),
    op_record("bitwise_not", 1, bool_dtypes, jtu.rand_small),
    op_record("bitwise_or", 2, bool_dtypes, jtu.rand_small),
    op_record("bitwise_xor", 2, bool_dtypes, jtu.rand_small),
    op_record("population_count", 1, uint_dtypes, partial(jtu.rand_int,
                                                          high=1 << 32)),

    op_record("add", 2, default_dtypes + complex_dtypes, jtu.rand_small),
    op_record("sub", 2, default_dtypes + complex_dtypes, jtu.rand_small),
    op_record("mul", 2, default_dtypes + complex_dtypes, jtu.rand_small),
    op_record("div", 2, default_dtypes + complex_dtypes, jtu.rand_nonzero),
    op_record("rem", 2, default_dtypes, jtu.rand_nonzero),

    op_record("max", 2, all_dtypes, jtu.rand_small),
    op_record("min", 2, all_dtypes, jtu.rand_small),

    op_record("eq", 2, all_dtypes, jtu.rand_some_equal),
    op_record("ne", 2, all_dtypes, jtu.rand_small),
    op_record("ge", 2, default_dtypes, jtu.rand_small),
    op_record("gt", 2, default_dtypes, jtu.rand_small),
    op_record("le", 2, default_dtypes, jtu.rand_small),
    op_record("lt", 2, default_dtypes, jtu.rand_small),
]

@pytest.mark.parametrize(
    'op_name,rng_factory,shapes,dtype,tol',
    [(rec.op, rec.rng_factory, shapes, dtype, rec.tol)
     for rec in LAX_OPS
     for shape_group in compatible_shapes
     for shapes in CombosWithReplacement(shape_group, rec.nargs)
     for dtype in rec.dtypes])
def test_nary(op_name, rng_factory, shapes, dtype, tol):
  rng = rng_factory(np.random)
  args = [rng(shape, dtype) for shape in shapes]
  tu.check_lazy_fun(getattr(lax, op_name), *args, atol=tol, rtol=tol)

LAX_REDUCE_OPS = [
  op_record("_reduce_sum", 1, number_dtypes, jtu.rand_default),
  op_record("_reduce_prod", 1, number_dtypes, jtu.rand_small_positive),
  op_record("_reduce_max", 1, all_dtypes, jtu.rand_default),
  op_record("_reduce_min", 1, all_dtypes, jtu.rand_default),
  op_record("_reduce_or", 1, bool_dtypes, jtu.rand_default),
  op_record("_reduce_and", 1, bool_dtypes, jtu.rand_default),
]

@pytest.mark.parametrize(
  'op_name,rng_factory,shape,axes,dtype,tol',
  [(rec.op, rec.rng_factory, shape, axes, dtype, rec.tol)
   for rec in LAX_REDUCE_OPS
   for (shape, axes) in [[(3, 4, 5), (0,)], [(3, 4, 5), (1, 2)],
                         [(3, 4, 5), (0, 2)], [(3, 4, 5), (0, 1, 2)]]
   for dtype in rec.dtypes])
def test_reduce(op_name, rng_factory, shape, axes, dtype, tol):
  rng = rng_factory(np.random)
  args = [rng(shape, dtype)]
  fun = partial(getattr(lax.lax, op_name), axes=axes)
  tu.check_lazy_fun(fun, *args, atol=tol, rtol=tol)

JAX_ARGMINMAX_RECORDS = [
  op_record("argmin", 1, default_dtypes, jtu.rand_some_equal),
  op_record("argmax", 1, default_dtypes, jtu.rand_some_equal),
]

@pytest.mark.parametrize(
  'op,rng_factory,shape,dtype,axis,index_dtype',
  [(rec.op, rec.rng_factory, shape, dtype, axis, index_dtype)
  for rec in JAX_ARGMINMAX_RECORDS
  for shape in [(1,), (4,), (2, 3), (1, 2, 1)]
  for dtype in rec.dtypes
  for index_dtype in int_dtypes
  for axis in range(len(shape))
  if dtype not in (np.complex128, dtypes.bfloat16)])
def test_argminmax(op, rng_factory, shape, dtype, axis, index_dtype):
  rng = rng_factory(np.random)
  def fun(x): return getattr(lax, op)(x, axis, index_dtype)
  args = [rng(shape, dtype)]
  tu.check_lazy_fun(fun, *args)

@pytest.mark.parametrize(
  'shape,dtype,dimensions,rng_factory',
  [(shape, dtype, dimensions, rng_factory)
   for dtype in default_dtypes
   for (shape, dimensions) in [
     [(1,), (0,)],
     [(1,), (-1,)],
     [(2, 1, 4), (1,)],
     [(2, 1, 3, 1), (1,)],
     [(2, 1, 3, 1), (1, 3)],
     [(2, 1, 3, 1), (3,)]]
   for rng_factory in [jtu.rand_default]])
def test_squeeze(shape, dtype, dimensions, rng_factory):
  rng = rng_factory(np.random)
  args = [rng(shape, dtype)]
  tu.check_lazy_fun(lambda x: lax.squeeze(x, dimensions), *args)

@pytest.mark.parametrize(
    'dim,base_shape,dtype,num_arrs,rng_factory',
    [(dim, base_shape, dtype, num_arrs, rng_factory)
     for num_arrs in [3]
     for dtype in default_dtypes
     for base_shape in [(4,), (3, 4), (2, 3, 4)]
     for dim in range(len(base_shape))
     for rng_factory in [jtu.rand_default]])
def test_concatenate(dim, base_shape, dtype, num_arrs, rng_factory):
  rng = rng_factory(np.random)
  shapes = [base_shape[:dim] + (size,) + base_shape[dim+1:]
            for size, _ in zip(itertools.cycle([3, 1, 4]), range(num_arrs))]
  args = [rng(shape, dtype) for shape in shapes]
  op = lambda *args: lax.concatenate(args, dim)
  tu.check_lazy_fun(op, *args)

@pytest.mark.parametrize(
  'shape,dtype,permutation,rng_factory',
  [(shape, dtype, permutation, rng_factory)
   for dtype in default_dtypes
   for shape, permutation in [((3, 4), (1, 0)),
                              ((3, 4, 5), (2, 1, 0)),
                              ((3, 4, 5), (1, 0, 2))]
   for rng_factory in [jtu.rand_default]])
def test_transpose(shape, dtype, permutation, rng_factory):
  rng = rng_factory(np.random)
  arg = rng(shape, dtype)
  tu.check_lazy_fun(lambda x: lax.transpose(x, permutation=permutation), arg)

@pytest.mark.parametrize(
    'shape,dtype,dimensions,rng_factory',
    [(shape, dtype, dimensions, rng_factory)
     for dtype in default_dtypes
     for shape, dimensions in [((4,), (0,)), ((3, 4), (1,)), ((2, 3, 4), (1, 2))]
     for rng_factory in [jtu.rand_default]])
def test_rev(shape, dtype, dimensions, rng_factory):
  rng = rng_factory(np.random)
  arg = rng(shape, dtype)
  tu.check_lazy_fun(lambda x: lax.rev(x, dimensions=dimensions), arg)

@pytest.mark.parametrize(
  'lhs_shape,rhs_shape,dtype,rng_factory',
  [(lhs_shape, rhs_shape, dtype, rng_factory)
   for lhs_shape in [(3,), (4, 3)]
   for rhs_shape in [(3,), (3, 6)]
   for dtype in float_dtypes
   for rng_factory in [jtu.rand_default]])
def test_dot(lhs_shape, rhs_shape, dtype, rng_factory):
  rng = rng_factory(np.random)
  args = [rng(lhs_shape, dtype), rng(rhs_shape, dtype)]
  tu.check_lazy_fun(lax.dot, *args)

@pytest.mark.parametrize(
  'lhs_shape,rhs_shape,dimension_numbers,dtype,rng_factory',
  [(lhs_shape, rhs_shape, dimension_numbers, dtype, rng_factory)
   for dtype in float_dtypes
   for lhs_shape, rhs_shape, dimension_numbers in
   [((3, 3, 2), (3, 2, 4), (([2], [1]), ([0], [0]))),
    ((3, 4, 2, 4), (3, 4, 3, 2), (([2], [3]), ([0, 1], [0, 1])))]
   for rng_factory in [jtu.rand_default]])
def test_dot_general_contract_and_batch(lhs_shape, rhs_shape, dimension_numbers, dtype, rng_factory):
  rng = rng_factory(np.random)
  args = [rng(lhs_shape, dtype), rng(rhs_shape, dtype)]
  tu.check_lazy_fun(partial(lax.dot_general, dimension_numbers=dimension_numbers), *args, atol=1e-5)

@pytest.mark.parametrize(
  'lhs_shape,rhs_shape,dtype,lhs_contracting,rhs_contracting,rng_factory',
  [(lhs_shape, rhs_shape, dtype, lhs_contracting, rhs_contracting, rng_factory)
   for dtype in float_dtypes
   for lhs_shape, rhs_shape, lhs_contracting, rhs_contracting in [
     [(3, 5), (2, 5), [1], [1]],
     [(5, 3), (5, 2), [0], [0]],
     [(5, 3, 2), (5, 2, 4), [0], [0]],
     [(5, 3, 2), (5, 2, 4), [0,2], [0,1]],
     [(1, 2, 2, 3), (1, 2, 3, 1), [1], [1]],
     [(3, 2), (2, 4), [1], [0]],
   ]
   for rng_factory in [jtu.rand_default]])
def test_dot_general_contract_only(
    lhs_shape, rhs_shape, dtype, lhs_contracting, rhs_contracting, rng_factory):
  rng = rng_factory(np.random)
  args = [rng(lhs_shape, dtype), rng(rhs_shape, dtype)]
  dimension_numbers = ((lhs_contracting, rhs_contracting), ((), ()))
  tu.check_lazy_fun(partial(lax.dot_general, dimension_numbers=dimension_numbers), *args, atol=1e-5)

@pytest.mark.parametrize(
  'shape,dtype,starts,limits,strides,rng_factory',
  [(shape, dtype, start_indices, limit_indices, strides, rng_factory)
   for shape, start_indices, limit_indices, strides in [
     [(3,), (1,), (2,), None],
     [(7,), (4,), (7,), None],
     [(5,), (1,), (5,), (2,)],
     [(8,), (1,), (6,), (2,)],
     [(5, 3), (1, 1), (3, 2), None],
     [(5, 3), (1, 1), (3, 1), None],
     [(7, 5, 3), (4, 0, 1), (7, 1, 3), None],
     [(5, 3), (1, 1), (2, 1), (1, 1)],
     [(5, 3), (1, 1), (5, 3), (2, 1)],
   ]
   for dtype in default_dtypes
   for rng_factory in [jtu.rand_default]])
def test_slice(shape, dtype, starts, limits, strides, rng_factory):
  rng = rng_factory(np.random)
  args = [rng(shape, dtype)]
  op = lambda x: lax.slice(x, starts, limits, strides)
  tu.check_lazy_fun(op, *args)

@pytest.mark.parametrize(
  "inshape,dtype,outshape,dimensions,rng_factory",
  [(inshape, dtype, outshape, broadcast_dimensions, rng_factory)
   for inshape, outshape, broadcast_dimensions in [
    ([2], [2, 2], [0]),
    ([2], [2, 2], [1]),
    ([2], [2, 3], [0]),
    ([], [2, 3], []),
    ([1], [2, 3], [1])]
  for dtype in default_dtypes
  for rng_factory in [jtu.rand_default]])
def test_broadcast_in_dim(inshape, dtype, outshape, dimensions, rng_factory):
  rng = rng_factory(np.random)
  args = [rng(inshape, dtype)]
  op = lambda x: lax.broadcast_in_dim(x, outshape, dimensions)
  tu.check_lazy_fun(op, *args)

@pytest.mark.parametrize(
  'shape,dtype,padding_config,rng_factory',
  [(shape, dtype, padding_config, jtu.rand_default)
   for shape in [(2, 3)]
   for padding_config in
   [
     [(0, 0, 0), (0, 0, 0)],  # no padding
     [(1, 1, 0), (2, 2, 0)],  # only positive edge padding
     [(1, 2, 1), (0, 1, 0)],  # edge padding and interior padding
     [(0, 0, 0), (-1, -1, 0)],  # negative padding
     [(0, 0, 0), (-2, -2, 4)],  # negative padding and interior padding
     [(0, 0, 0), (-2, -3, 1)],  # remove everything in one dimension
   ]
   for dtype in default_dtypes])
def test_pad(shape, dtype, padding_config, rng_factory):
  rng = rng_factory(np.random)
  args = [rng(shape, dtype), rng((), dtype)]
  op = lambda *args: lax.pad(*args, padding_config)
  tu.check_lazy_fun(op, *args)

@pytest.mark.parametrize(
  'lhs_shape,rhs_shape,dtype,strides,padding,rng_factory',
  [(lhs_shape, rhs_shape, dtype, strides, padding, rng_factory)
   for lhs_shape, rhs_shape in [
     ((b, i, 9, 10), (j, i, 4, 5))
     for b, i, j in itertools.product([2, 3], repeat=3)]
   for dtype in float_dtypes
   for strides in [(1, 1), (1, 2), (2, 1)]
   for padding in ["VALID", "SAME"]
   for rng_factory in [jtu.rand_small]])
def test_conv(lhs_shape, rhs_shape, dtype, strides, padding, rng_factory):
  rng = rng_factory(np.random)
  args = [rng(lhs_shape, dtype), rng(rhs_shape, dtype)]
  fun = lambda lhs, rhs: lax.conv(lhs, rhs, strides, padding)
  tu.check_lazy_fun(fun, *args)

@pytest.mark.parametrize(
  'lhs_shape,rhs_shape,dtype,strides,padding,lhs_dilation,rhs_dilation,'
  'feature_group_count,batch_group_count,dimension_numbers,perms,rng_factory',
  [(lhs_shape, rhs_shape, dtype, strides, padding, lhs_dilation, rhs_dilation,
    feature_group_count, batch_group_count, dimension_numbers, perms, rng_factory)
  for batch_group_count, feature_group_count in [(1, 1)] #, (2, 1), (1, 2)
  for lhs_shape, rhs_shape in [
    ((b * batch_group_count, i * feature_group_count, 9, w),
     (j * feature_group_count * batch_group_count, i, 4, 5))
    for w in [0, 10]
    for b, i, j in itertools.product([2, 3], repeat=3)]
  for dtype in float_dtypes
  for strides in [(1, 1), (2, 1)]
  for padding in [((1, 2), (2, 0)), ((10, 8), (7, 13))]
  for lhs_dilation in [(1, 1), (1, 2), (1, 4)]
  for rhs_dilation in [(1, 1)] # , (1, 2), (1, 4)
  for rng_factory in [jtu.rand_default]
  for dimension_numbers, perms in [
    (("NCHW", "OIHW", "NCHW"), ([0, 1, 2, 3], [0, 1, 2, 3])),
    (("NHWC", "HWIO", "NHWC"), ([0, 2, 3, 1], [2, 3, 1, 0])),
    (("NCHW", "HWIO", "NHWC"), ([0, 1, 2, 3], [2, 3, 1, 0])),
  ]])
def test_conv_general_dilated(lhs_shape, rhs_shape, dtype, strides,
                              padding, lhs_dilation, rhs_dilation,
                              feature_group_count, batch_group_count,
                              dimension_numbers, perms, rng_factory):
  rng = rng_factory(np.random)
  lhs_perm, rhs_perm = perms  # permute to compatible shapes
  args = [lax.transpose(rng(lhs_shape, dtype), lhs_perm),
          lax.transpose(rng(rhs_shape, dtype), rhs_perm)]
  def fun(lhs, rhs):
    return lax.conv_general_dilated(
      lhs, rhs, strides, padding, lhs_dilation, rhs_dilation,
      dimension_numbers, feature_group_count=feature_group_count,
      batch_group_count=batch_group_count)
  tu.check_lazy_fun(fun, *args, rtol=.005, atol=.2)

def test_conv_incounts():
  rhs_shape = (4, 1, 2, 3)
  lhs_count, rhs_count = conv_incounts((8, 1, 3, 7), rhs_shape, (1, 1))
  np.testing.assert_array_equal(4 * np.array(8 * [[[
    [1, 2, 3, 3, 3, 2, 1],
    [2, 4, 6, 6, 6, 4, 2],
    [1, 2, 3, 3, 3, 2, 1]
  ]]]), lhs_count)
  np.testing.assert_array_equal(np.full(rhs_shape, 8), rhs_count)

def test_conv_incounts_strided():
  rhs_shape = (1, 1, 7)
  lhs_count, rhs_count = conv_incounts((1, 1, 21), rhs_shape, (3,))
  np.testing.assert_array_equal(
    [[[1, 1, 1, 2, 2, 2, 3, 2, 2, 3, 2, 2, 3, 2, 2, 2, 1, 1, 1, 0, 0]]], lhs_count)
  np.testing.assert_array_equal(np.ones(rhs_shape), rhs_count)

@pytest.mark.parametrize(
  'outstart,outcount,expected_incount',
  [((7,), np.arange(7), [0, 3]), # outstart aligned with padded element
   ((7,), np.arange(1), [0]), # outstart + outstop aligned with padded element
   ((7,), Ones((1,)), Ones((1,))), # outcount Ones
   ((6,), np.arange(8), [1, 4]), # outstart in interior padding
   ((6,), np.arange(3), [1]), # outstart, outstop in interior padding
   ((0,), np.arange(4), None), # outstart in lo, no elements
   ((1,), np.arange(4), [3]), # outstart in lo, including padded element
   ((15,), np.arange(1), None), # outstart in hi
  ])
def test_pad_dependency_rule(outstart, outcount, expected_incount):
  (instart, _), (incount, _), outslice = pad_dependency_rule(
    outstart, outcount, np.arange(3), 0, ((4, 4, 2),))
  np.testing.assert_array_equal(expected_incount, incount)
  inslice = None if incount is None else np.ones(incount.shape, int)
  assert outslice(inslice, 0).shape == outcount.shape

@pytest.mark.parametrize(
    'pred_shape,arg_shape,arg_dtype,rng_factory',
    [(pred_shape, arg_shape, arg_dtype, rng_factory)
     for arg_shape in [(), (3,), (2, 3)]
     for pred_shape in ([(), arg_shape] if arg_shape else [()])
     for arg_dtype in default_dtypes
     for rng_factory in [jtu.rand_default]])
def test_select(pred_shape, arg_shape, arg_dtype, rng_factory):
  rng = rng_factory(np.random)
  args = [rng(pred_shape, np.bool_), rng(arg_shape, arg_dtype),
          rng(arg_shape, arg_dtype)]
  return tu.check_lazy_fun(lax.select, *args)

def test_custom_jvp():
  @custom_jvp
  def f(x):
    return x ** 2

  f.defjvp(lambda x: 2 * x)
  rng = jtu.rand_small(np.random)
  tu.check_lazy_fun(f, rng((1,), 'float32'))

def test_jit():
  rng = jtu.rand_small(np.random)
  tu.check_lazy_fun(jit(lambda x: x * 2), rng((1,), int))

def test_jit_freevar():
  rng = jtu.rand_small(np.random)
  tu.check_lazy_fun(lambda x, y: jit(lambda x: x * y)(x), rng((1,), int), rng((1,), int))

max_stars_repo_name: j-towns/fastar

max_issues_repo_name: j-towns/fastar

max_forks_repo_name: j-towns/fastar

reason: import numpy,from jax


text:
"""
Processing Utilities
"""
import numpy as np

def smooth_curve(curve, length):
    """
    Smooth a curve with a boxcar of length n elements
    """
    boxcar = np.ones(length) / length
    return np.convolve(boxcar, curve, mode="same")

ext: py

max_stars_repo_name: lawrencegwm/petpy

max_issues_repo_name: lawrencegwm/petpy

max_forks_repo_name: lawrencegwm/petpy

reason: import numpy


text:
'''
Deimos longslit data reduction

Dependencies: the usual astropy suite (numpy, scipy, matplotlib, astropy)
along with astroscrappy for cosmic ray subtraction

Development notes:
The DEIMOS chip layout is shown in the engineering drawings found here:
http://www.ucolick.org/~sla/fits/mosaic/d0307j.pdf

After getting a single array with the (as currently named)
multiextension_to_array function, the dispersion axis will be in the y-axis
(columns) and the spatial axis will be in x (rows).  Wavelength increases with
y indices.

Much of the inspiration for this comes from the PyDIS package, found here:
https://github.com/jradavenport/pydis

History:
1/25/2016 initial working (somewhat) version
'''

from __future__ import (absolute_import, division,
                        print_function, unicode_literals,
                        with_statement)
from builtins import (bytes, dict, int, list, object, range, str,
                      ascii, chr, hex, input, next, oct, open,
                      pow, round, super,
                      filter, map, zip)

import sys, os, glob
import numpy as np
from matplotlib import pyplot as plt
from scipy.optimize import curve_fit, minimize
from astropy.convolution import convolve, Box1DKernel
from scipy.interpolate import UnivariateSpline, interp1d
from scipy.integrate import quad
from scipy.signal import medfilt
from astropy.io import fits
from astropy.wcs import WCS
from astroscrappy import detect_cosmics

from utils import *

def multiextension_to_array(fitsname=None, data_arrays=None, headers=None):
    '''
    Takes the eight DEIMOS ccd arrays and returns a single array oriented
    correctly.  Supply the list of headers to insure the correct orientation.
    '''
    rows = 2
    columns = 4
    nimages = rows * columns

    if fitsname is not None:
        hdulist = fits.open(fitsname)
        primaryHDU = hdulist[0]
        headers = [hdu.header for hdu in hdulist[1: nimages + 1]]
        arrays = [hdu.data for hdu in hdulist[1: nimages + 1]]
    elif data_arrays is not None and headers is not None:
        arrays = data_arrays
    else:
        raise KeyError('Need either fitsname or both arrays and headers!')
    wcslist = [WCS(header) for header in headers]

    # get the mosaic image size from any of the headers
    det_x0, det_x1, det_y0, det_y1 = get_indices(headers[0]['DETSIZE'])
    mosaic = np.empty(shape=(det_y1, det_x1))
    mosaic[:] = np.nan

    for i, data in enumerate(arrays):
        '''
        I'm assuming the wcs transformation is perfectly linear, so I can just
        specify the start and end pixels.
        '''
        # select the data which isn't part of the overscan region
        x0, x1, y0, y1 = get_indices(headers[i]['DATASEC'])
        # get the direction of pixel increments
        dx = float(headers[i]['CD1_1'])
        dy = float(headers[i]['CD2_2'])

        # subtract by 1 to convert from FITS format 1-based indexing to 0-based
        indices = np.index_exp[y0 - 1: y1, x0 - 1: x1]
        start = [x0 - 1, y0 - 1]
        end = [x1, y1]
        old_coords = np.array([start, end])

        # subtract 0.5 to convert from DEIMOS plane wcs definition to indices
        new_coords = wcslist[i].all_pix2world(old_coords, 0) - 0.5
        new_start, new_end = map(list, new_coords)
        # need to handle going backwards in the slice
        new_end[0] = None if new_end[0] < 0 else new_end[0]
        new_end[1] = None if new_end[1] < 0 else new_end[1]

        # remember to switch from x, y in FITS to y, x in numpy
        new_indices = np.index_exp[new_start[1]: new_end[1]: dy,
                                   new_start[0]: new_end[0]: dx]
        mosaic[new_indices] = data[indices]

    return mosaic

def make_masterbias(output="masterbias.fits", biasdir="bias/", remove_cr=True):
    '''
    Creates a median bias file from fits files found in biasdir.
    '''
    bias_files = glob.glob(biasdir + "*.fits")
    # arbitrarily taking the first exposure as the header info
    hdulist = fits.open(bias_files[0])
    primary_header = hdulist[0].header
    headers = [hdu.header for hdu in hdulist[1:9]]

    bias_data = []
    # axis 0 is different exposures, axis 1 is different CCDs
    # axis 2, 3 are y, x
    for i, f in enumerate(bias_files):
        bias_data.append(np.array([hdu.data for hdu in fits.open(f)[1:9]]))
    bias_data = np.array(bias_data)

    # if we have more bias frames, we should cr subtract before medianing biases
    bias = np.nanmedian(bias_data, axis=0)
    if remove_cr:
        masks = []
        cr_cleaned = []
        for frame in bias:
            mask, cleaned = detect_cosmics(frame, verbose=True,
                                           cleantype='medmask',
                                           sigclip=0.2)
            masks.append(mask)
            cr_cleaned.append(cleaned)
        bias = np.array(cr_cleaned)

    bias = multiextension_to_array(data_arrays=bias, headers=headers)
    if output is not None:
        header = primary_header
        header.extend(deimos_cards(bias.shape))
        hdu = fits.PrimaryHDU(data=bias,
                              header=header)
        hdu.writeto(output, clobber=True)
    return bias

def make_masterflat(output="masterflat.fits", flatdir="flats/", bias="masterbias.fits"):
    '''
    Creates a mean flat field from fits files found in flatdir.
    '''
    if isinstance(bias, str):
        bias = fits.getdata(bias)

    flat_files = glob.glob(flatdir + "*.fits")
    flat_data = np.array([multiextension_to_array(f) for f in flat_files])

    # arbitrarily taking the first exposure as the header info
    hdulist = fits.open(flat_files[0])
    headers = [hdu.header for hdu in hdulist]

    # divide by median after subtracting bias
    for i, data in enumerate(flat_data):
        flat_data[i] -= bias
        flat_data[i] /= np.nanmedian(flat_data[i])
    # median across all frames
    flat = np.nanmedian(flat_data, axis=0)

    # normalize to lamp response
    # response technique from PyDIS (https://github.com/jradavenport/pydis)
    spectral_size, spatial_size = flat.shape
    spectral_pixels = np.arange(spectral_size)
    spatial_pixels = np.arange(spatial_size)
    flat_1d = np.log10(convolve(flat.sum(axis=1), Box1DKernel(5)))
    spline = UnivariateSpline(spectral_pixels, flat_1d, k=2, s=0.001)
    flat_curve = 10.0 ** spline(spectral_pixels)
    # tile back up to shape of flat file
    flat_curve = np.tile(np.split(flat_curve, flat_curve.size, axis=0), (1, spatial_size))
    flat /= flat_curve

    if output is not None:
        header = headers[0]
        header.extend(deimos_cards(flat.shape))
        hdu = fits.PrimaryHDU(data=flat,
                              header=header)
        hdu.writeto(output, clobber=True)
    return flat

def normalize(fitsname, output=None, cr_remove=True, multiextension=True,
              masterflat="masterflat.fits", masterbias="masterbias.fits",
              cr_options=None):
    '''
    Flat field, bias subtraction, and cosmic ray removal with astroscrappy.

    Parameters:
    -----------
        fitsname: either string or 2d array, if string then name of fits file,
                  if 2d array, then the output of arrange_fits
        output: string, name of file to write out to (optional)
                if None, then just returns the array
        cr_remove: boolean, if true, the subtract cosmics
        multiextension: boolean, if true, then
        masterflat: str, name of fits file with master flat (optional)
                    if None, then make a master flat from files in flatdir
        masterbias: str, name of fits file with master bias (optional)
    if None, then make a master bias from files in biasdir
        :cr_options: dict, keys should be keywords of detect_cosmics from
                     astroscrappy, overrides defaults below
    Returns:
    --------
        :normalized_data: 2d numpy array, rows are spectral, columns are spatial
    '''
    if masterbias is None:
        bias = make_masterbias()
    else:
        bias = fits.getdata(masterbias)
    if masterflat is None:
        flat = make_masterflat()
    else:
        flat = fits.getdata(masterflat)

    if multiextension:
        raw_data = multiextension_to_array(fitsname)
        header = fits.open(fitsname)[0].header
        header.extend(deimos_cards(raw_data.shape))
    else:
        raw_data = fits.getdata(fitsname)
        header = fits.getheader(fitsname)

    if cr_remove:
        flat_mask = get_slitmask(flat)
        kwargs = {'verbose': True, 'inmask': flat_mask, 'cleantype': 'medmask',
                  'sigclip': 0.5, 'sigfrac': 0.1, 'niter': 4}
        if cr_options is not None:
            for key, value in cr_options.items():
                kwargs[key] = value
        mask, cleaned = detect_cosmics(raw_data, **kwargs)
        normed_data = (cleaned - bias) / flat
    else:
        normed_data = (raw_data - bias) / flat

    if output is not None:
        fits.writeto(output, data=normed_data, header=header, clobber=True)

    return normed_data

def ap_trace(data, initial_guess=None, masterflat="masterflat.fits",
             nbins=20, nsigma=15, seeing=1., arcsec_per_pix=0.1185, frac_med=1.5, slit_width=1.):
    '''
    Traces the spatial apeture of a gaussian point source.

    Parameters:
    -----------
    data: 2d numpy array
    initial_guess: int, initial guess for the spatial position (here assumed as
                   x-axis) of the source
    masterflat: str, name of flat file to use
    nbins: int, number of spectral bins to fit trace
    nsigma: float, how far away from the peak should the trace look, in stdev
    seeing: float, in arcsec, used to get an initial sigma guess for the trace
    arcsec_per_pix: float, float, should get from header, default for DEIMOS
    frac_med: float, fraction of median image above which to mask (e.g., for sky lines)
    slit_width: float, arcsec, width of slit

    Returns:
    --------
    trace: function from spectral index (y) to spatial index (x)
    sigma: function from spectral index (y) to width in spatial index
    '''

    # mask maker mask maker make me a mask
    # True for bad pixels
    flat_mask = get_slitmask(masterflat=masterflat)
    # sky_mask = data > np.nanmedian(data) * frac_med
    # mask = np.logical_or(flat_mask, sky_mask)
    mask = flat_mask

    y_size, x_size = data.shape
    # try to do binning on smaller scale based on slit width
    nbins = int(y_size * arcsec_per_pix / slit_width / 10)
    spectral_bin_edges = np.linspace(0, y_size, nbins + 1).astype(int)
    bin_sizes = y_size / nbins
    spectral_bin_centers = np.linspace(bin_sizes * 0.5,
                                       bin_sizes * (nbins - 0.5) ,
                                       nbins).astype(int)
    box_size = round(seeing / arcsec_per_pix * nsigma)

    if initial_guess is None:
        # make a rough guess
        spatial = np.sum(data, axis=0)
        smooth = medfilt(spatial, kernel_size=5)
        # roughly captures how far we have to go into the chip for constant illumination
        frac = 0.05
        x = np.arange(smooth.size)
        spatial_mask = np.logical_or(x < x.size * frac, x > x.size * (1 - frac))
        smooth[spatial_mask] = np.nan
        initial_guess = np.nanargmax(smooth)

    x_low = initial_guess - box_size
    x_high = initial_guess + box_size

    crop = data[:, x_low: x_high]
    mask = mask[:, x_low: x_high]
    crop[mask] = np.nan

    specbin_arrays = np.array_split(crop, spectral_bin_edges[1:-1])
    specbins = np.empty(shape=(nbins, x_high - x_low))
    for i, array in enumerate(specbin_arrays):
        specbins[i] = np.nanmean(array, axis=0)

    fit_centers = np.empty(specbins.shape[0])
    fit_sigmas = np.empty(specbins.shape[0])

    x = np.arange(specbins.shape[1])
    sigma = seeing / arcsec_per_pix
    sky_limit = round(3 * sigma)
    for i, specbin in enumerate(specbins):
        mask = ~np.isnan(specbin)
        a = np.nanmax(specbin[mask])
        x0 = np.nanargmax(specbin)
        sky = np.concatenate((specbin[:x0 - sky_limit], specbin[x0 + sky_limit:]))
        b = np.nanmedian(sky)
        params = [a, b, x0, sigma]
        popt, pcov = curve_fit(gaussian, x[mask], specbin[mask], p0=params)
        # plot for sanity check
        # space = np.linspace(0, x.max(), 1000)
        # plt.clf()
        # plt.plot(space, gaussian(space, *popt))
        # plt.plot(x[mask], specbin[mask], 'ko')
        # plt.show()

        # if err > 10**2, reject fit
        perr = np.sqrt(np.diag(pcov))
        # fit to gaussian and remember to add back cropped out spatial
        if perr[2] < 10**2:
            fit_centers[i] = popt[2] + x_low
        else:
            fit_centers[i] = np.nan
        fit_sigmas[i] = popt[3]

    mask = ~np.isnan(fit_centers)
    sigma_median = np.nanmedian(fit_sigmas)
    sigma_range = seeing / arcsec_per_pix
    mask = mask & (fit_sigmas < sigma_median + sigma_range)
    mask = mask & (fit_sigmas > sigma_median - sigma_range)

    # polynomial fit to fit centers and sigmas
    poly_center = np.polyfit(spectral_bin_centers[mask].astype(float),
                             fit_centers[mask].astype(float), deg=2,
                             w=1 / fit_sigmas[mask].astype(float)**2)
    poly_sigma = np.polyfit(spectral_bin_centers[mask].astype(float),
                            fit_sigmas[mask].astype(float), deg=2)

    trace = lambda y: np.polyval(poly_center, y)
    sigma = lambda y: np.polyval(poly_sigma, y)
    return trace, sigma

    # spline interpolation to get in between the binned spectral data
    # trace_spline = UnivariateSpline(spectral_bin_centers[mask], fit_centers[mask],
    #                                 k=3, s=1)
    # sigma_spline = UnivariateSpline(spectral_bin_centers[mask], fit_sigmas[mask],
    #                                 k=3, s=1)
    # return trace_spline, sigma_spline

def ap_extract(data, trace_spl, sigma_spl,
               apwidth=2, skysep=1, skywidth=2, skydeg=0, sky_subtract=True):
    '''
    Extract the spectrum using a specified trace.
    Data is the 2d array, trace_spl, sigma_spl are the splines from ap_trace.

    Parameters
    -----------
    :apwidth: in factors of sigma, corresponds to aperature radius
    :skysep: in factors of sigma, corresponds to the separation between sky and aperature
    :skywidth: in factors of sigma, corresponds to the width of sky windows on either side
    :skydeg: degree of polynomial fit for the sky spatial profile at each wavelength

    Returns
    -------
    ap
    sky
    uncertainty
    '''

    y_size, x_size = data.shape
    y_indices, x_indices = np.indices(data.shape)
    y_bins = np.arange(y_size)
    x_bins = np.arange(x_size)
    x_centers = trace_spl(y_bins)
    x_sigmas = sigma_spl(y_bins)

    # specify aperature as a function of spectral position
    ap_lows = x_centers - apwidth * x_sigmas
    ap_highs = x_centers + apwidth * x_sigmas

    right_sky_lows = x_centers - (apwidth + skysep + skywidth) * x_sigmas
    right_sky_highs = x_centers - (apwidth + skysep) * x_sigmas
    left_sky_lows = x_centers + (apwidth + skysep) * x_sigmas
    left_sky_highs = x_centers + (apwidth + skysep + skywidth) * x_sigmas

    ap_pixels = np.logical_and(ap_lows < x_indices, x_indices < ap_highs)
    right_sky_pixels = np.logical_and(right_sky_lows < x_indices, x_indices < right_sky_highs)
    left_sky_pixels = np.logical_and(left_sky_lows < x_indices, x_indices < left_sky_highs)
    sky_pixels = np.logical_or(right_sky_pixels, left_sky_pixels)

    aperture = np.empty(y_size)
    sky = np.empty(y_size)
    unc = np.empty(y_size)

    if not sky_subtract:
        # just return aperture sum
        for i in range(y_size):
            aperture[i] = np.nansum(data[i, ap_pixels[i]])
        return aperture

    for i in range(y_size):
        data_slice = data[i]
        ap_slice = data_slice[ap_pixels[i]]
        aperture[i] = np.nansum(ap_slice)
        sky_slice = data_slice[sky_pixels[i]]
        x_sky = x_bins[sky_pixels[i]]
        x_ap = x_bins[ap_pixels[i]]
        if skydeg > 0:
            pfit = np.polyfit(x_sky, sky_slice, skydeg)
            sky[i] = np.nansum(np.polyval(pfit, x_ap))
        elif skydeg == 0:
            sky[i] = np.nanmean(sky_slice) * (apwidth * x_sigmas[i] * 2 + 1)

        # for now...
        coaddN = 1
        # uncertainty, as done by PyDIS
        sigB = np.std(sky_slice) # stddev in the background data
        N_B = len(x_sky) # number of bkgd pixels
        N_A = apwidth * x_sigmas[i] * 2. + 1 # number of aperture pixels
        # based on aperture phot err description by F. Masci, Caltech:
        # http://wise2.ipac.caltech.edu/staff/fmasci/ApPhotUncert.pdf
        unc[i] = np.sqrt(np.sum((aperture[i] - sky[i]) / coaddN) +
                         (N_A + N_A**2. / N_B) * (sigB**2.))

    return aperture, sky, unc

def sky_wavesol(sky, header, deg=4):
    '''
    Calculate the wavelength solution from the sky calibration.

    Parameters
    ----------
    sky: float array, sky spectrum from ap_extract
    header: header from data file
    deg: int, degree of fit for poly mode

    Returns
    -------
    wfunc: function from trace center to wavelength, in Angstroms
    '''

    p0 = get_initial_wavesol(header)
    if len(p0) < deg + 1:
        # assuming that p0 refers to the lower degrees only
        p0 = np.append(p0, np.zeros(deg + 1 - len(p0)))
    grating_name = header['GRATENAM']

    sky_wave, sky_flux = get_sky_spectrum(grating_name)
    # function from wavelength in angstroms to sky flux
    f_ref = interp1d(sky_wave, sky_flux, bounds_error=False, fill_value=np.nan)

    # map pixel space to sky flux
    # I'm using lower to higher terms for fit, but polyval
    # uses higher to lower
    def poly_wave(y, p): return np.polyval(p[::-1], y)
    def fit_function(y, *p): return f_ref(poly_wave(y, p))
    return fit_function

    y = np.arange(sky.size)
    y0 = y.size / 2.

    # normalize to height
    sky_normed = sky * np.amax(sky_flux) / np.amax(sky)
    # w_init = poly_wave(y - y0, p0)
    # w_low = w_init[0]
    # w_high = w_init[-1]
    # ref_area = quad(f_ref, w_low, w_high)[0]
    # sky_area = np.sum(sky)
    # sky_normed = sky * ref_area / sky_area

    # subtract y pixels by zeros point at center
    popt, pcov = curve_fit(fit_function, y - y0, sky_normed, p0)
    def wfunc(y): return poly_wave(y - y0, popt)
    return wfunc

def wavelength_solution(trace_spl, sigma_spl, arc_name, lines="spec2d_lamp_NIST.dat",
                        mode='poly', deg=4, method='Nelder-Mead', linear_no_fit=False):
    '''
    Calculate the wavelength solution, given an arcline calibration frame.
    I'm assuming that the initial guess is very good (i.e., that we can identify
    lines based on which are the closest) and that there are more lines in
    arc_lines than list_lines.

    Parameters
    ----------
    trace_spl: trace spline (i.e., from ap_trace)
    sigma_spl: sigma spline (from ap_trace)
    arc_name: str, name of fits file with arc frame
    lines: str, name of file with lines (in Angstroms) as the first column
    mode: str, poly for polynomial fit
    deg: int, degree of fit for poly mode
    method: str, optimization method to use
    linear_no_fit: bool, if true, then just return the inital guess from header

    Returns
    -------
    wfunc: function from trace center to wavelength, in Angstroms
    '''

    header = fits.getheader(arc_name)
    arc_data = fits.getdata(arc_name)
    grating_name = header['GRATENAM']
    grating_position = header['GRATEPOS']
    # this should be the nominal wavelength at the center of the detector
    # if I'm reading the info at the page below correctly
    # http://www2.keck.hawaii.edu/inst/deimos/grating-info.html
    w0 = header['G' + str(grating_position).strip() + 'TLTWAV']
    w_per_pix = grating_to_disp(grating_name)
    y0 = arc_data.shape[0] / 2

    if linear_no_fit:
        def wfunc(y): return w0 + w_per_pix * (y - y0)
        return wfunc

    list_lines, list_heights = good_lines(lines)

    # get pixel values of lines
    ap_lines = ap_extract(arc_data, trace_spl, sigma_spl, sky_subtract=False)

    # list of arc lines
    arc_lines, arc_heights = get_lines(ap_lines)

    # array of zeroth degree and linear terms
    p0 = np.array([w0, w_per_pix])
    # append zero terms for higher order fits
    if len(p0) < deg + 1:
        # assuming that p0 refers to the lower degrees only
        p0 = np.append(p0, np.zeros(deg + 1 - len(p0)))

    if mode == 'poly':
        # I'm using lower to higher terms for fit, but polyval
        # uses higher to lower
        def fit_function(y, p): return np.polyval(p[::-1], y)
    else:
        raise NotImplementedError("Only mode='poly' is currently implemented.")

    # subtract y pixels by zeros point at center
    def min_function(p): return premetric(fit_function(arc_lines - y0, p), list_lines)
    res = minimize(min_function, p0, method=method, options={'disp': True})

    def wfunc(y): return fit_function(y - y.shape[0]/2., res.x)
    return wfunc

def reduce_files(files, save=True, plot=True):

    for f in files:

        print("Cleaning", f)
        clean_name = f.split('.')[0] + '.clean.fits'
        if os.path.exists(clean_name):
            clean = fits.getdata(clean_name)
        else:
            clean = normalize(f, output=clean_name)
        continue
        print("Tracing", f)
        trace_spl, sigma_spl = ap_trace(clean)

        print("Extracting aperture and sky subtracting", f)
        ap, sky, unc = ap_extract(clean, trace_spl, sigma_spl)

        print("Fitting wavelength solution", f)
        wfunc = wavelength_solution(trace_spl, sigma_spl, "normed_NeArKrXe.fits")

        y = np.arange(ap.size)
        w = wfunc(y)
        if save:
            header = "Spectrum extracted from " + f + "\nAngstroms\tFlux"
            write_array = list(zip(w, ap - sky, unc))
            np.savetxt(f.split('.')[0] + ".dat", write_array, header=header)
        if plot:
            plt.clf()
            plt.plot(wfunc(y), ap - sky, 'k-')
            plt.xlabel("Wavelength (Angstroms)")
            plt.ylabel("Relatived flux (arbitrary units)")
            plt.show()

if __name__ == "__main__":
    reduce_files(sys.argv[1:])

max_stars_repo_name: adwasser/deimos-longslit

max_issues_repo_name: adwasser/deimos-longslit

max_forks_repo_name: adwasser/deimos-longslit

reason: import numpy,from scipy,from astropy


text:
# Copyright 2020 DeepMind Technologies Limited.
#
# 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.

"""ActionSpace implementations."""

import re
from typing import Generator, Iterable, List, Optional, Tuple, Sequence, Union

from absl import logging
from dm_env import specs
from dm_robotics.agentflow import core
from dm_robotics.agentflow import spec_utils
import numpy as np

# Internal profiling

class _Deslicer(core.ActionSpace):
  """Creates a value by projecting the given value into a larger one."""

  def __init__(
      self,
      name: str,
      spec: core.Spec,
      output_shape: Tuple[int],  # 1-d
      output_dtype: np.dtype,
      mask: Sequence[bool],
      default_value: Optional[np.floating] = None):
    super().__init__()
    self._name = name
    self._input_spec = spec
    self._output_shape = output_shape
    self._output_dtype = output_dtype
    self._mask = mask
    if not np.any(mask):
      logging.warning('Deslicer passed a mask with no valid indices to write'
                      'to.')
    self._default = np.nan if default_value is None else default_value

  @property
  def name(self) -> str:
    return self._name

  def spec(self) -> core.Spec:
    return self._input_spec

  # Profiling for .wrap('Deslicer.project')
  def project(self, action: np.ndarray) -> np.ndarray:
    output = np.full(
        shape=self._output_shape,
        fill_value=self._default,
        dtype=self._output_dtype)
    if np.any(self._mask):
      output[self._mask] = action
    return output

def prefix_slicer(
    spec: core.Spec,
    prefix: str,
    default_value: Optional[float] = None) -> core.ActionSpace[core.Spec]:
  """An ActionSpace for actions starting with prefix.

  The spec name is split on tab, and it's expected that the number of names
  from this split matches the shape of the spec, I.e. that each component has
  a name.
  The returned ActionSpace will project from actions with the prefix to the
  given spec, inserting default_value (or NaN, if missing).
  I.e, given a spec with seven components that start with 'robot_0' and
  some that do not, this function will return an ActionSpace of size seven
  that when projected will have the same size as the input spec did, but with
  NaNs for all the components that don't start with 'robot_0'.

  Args:
    spec: The (primitive) action spec to 'slice'.
    prefix: A regular expression for components to select. Note that we use
      regular expression comparison, so you can use exclusion patterns too.
    default_value: The default value used by the Desclicer.

  Returns:
    An ActionSpace.
      Given a value that conforms to the (currently implicit) input spec,
      return a value that conforms to the larger spec given to this function.

  Raises:
    ValueError: If a non-primitive spec is given or the names in the spec don't
      split as expected.
  """
  # Special case an empty input spec.
  if np.prod(spec.shape) == 0:
    return core.IdentityActionSpace(spec)

  names = np.asarray(spec.name.split('\t'))
  prefix_expr = re.compile(prefix)
  indices: List[bool] = [
      re.match(prefix_expr, name) is not None for name in names
  ]

  if len(names) != spec.shape[0]:
    raise ValueError('Expected {} names, got {}.  Name: {}'.format(
        spec.shape[0], len(names), names))

  if isinstance(spec, specs.DiscreteArray):
    raise ValueError('Support for DiscreteArray not implemented, yet.')
  elif isinstance(spec, specs.BoundedArray):
    input_spec = specs.BoundedArray(
        shape=(np.count_nonzero(indices),),
        dtype=spec.dtype,
        minimum=spec.minimum[indices],
        maximum=spec.maximum[indices],
        name='\t'.join(names[indices]))
  elif isinstance(spec, specs.Array):
    input_spec = specs.Array(
        shape=(np.count_nonzero(indices),),
        dtype=spec.dtype,
        name='\t'.join(names[indices]))
  else:
    raise ValueError('unknown spec type: {}'.format(type(spec)))

  return _Deslicer(
      name=prefix,
      spec=input_spec,
      output_shape=spec.shape,
      output_dtype=spec.dtype,
      mask=indices,
      default_value=default_value)

class CastActionSpace(core.ActionSpace):
  """Casts actions to the appropriate dtype for the provided spec."""

  def __init__(self,
               spec: core.Spec,
               ignore_nan: Optional[bool] = None,
               name: str = 'cast'):
    """Initializes SequentialActionSpace.

    Note: ShrinkToFitActionSpace also casts, so this should only be used if
    scaling is not desired.

    Args:
      spec: Specification for value to cast to.
      ignore_nan: If True, NaN values will not fail validation. If None, this is
        determined by the shape of `value`, so that large arrays (e.g. images)
        are not checked  (for performance reasons).
      name: A name for action space.
    """
    if np.issubdtype(spec.dtype, np.integer):
      logging.warning('Casting to %s will fail for NaN', spec.dtype)

    self._spec = spec
    self._ignore_nan = ignore_nan
    self._name = name

  @property
  def name(self) -> str:
    return self._name

  def spec(self) -> core.Spec:
    return self._spec

  def project(self, action: np.ndarray) -> np.ndarray:
    cast_action = action.astype(self._spec.dtype)
    spec_utils.validate(self._spec, cast_action, self._ignore_nan)
    return cast_action

class ShrinkToFitActionSpace(core.ActionSpace[specs.BoundedArray]):
  """Action space that scales an action if any component falls out of bounds."""

  def __init__(self,
               spec: specs.BoundedArray,
               ignore_nan: Optional[bool] = None,
               name: str = 'shrink_to_fit'):
    """Action space that scales the value towards zero to fit within spec.

    This action-space also casts the value to the dtype of the provided spec.

    Args:
      spec: Specification for value to scale and clip.
      ignore_nan: If True, NaN values will not fail validation. If None, this is
        determined by the shape of `value`, so that large arrays (e.g. images)
        are not checked  (for performance reasons).
      name: A name for action space.
    """
    self._spec = spec
    self._ignore_nan = ignore_nan
    self._name = name

  @property
  def name(self) -> str:
    return self._name

  def spec(self) -> specs.BoundedArray:
    return self._spec

  def project(self, action: np.ndarray) -> np.ndarray:
    return spec_utils.shrink_to_fit(
        value=action, spec=self._spec, ignore_nan=self._ignore_nan)

class FixedActionSpace(core.ActionSpace):
  """Like a partial function application, for an action space."""

  def __init__(self, action_space: core.ActionSpace, fixed_value: np.ndarray):
    self._action_space = action_space
    self._fixed_value = fixed_value
    space_shape = action_space.spec().shape
    value_shape = fixed_value.shape
    if space_shape != value_shape:
      raise ValueError('Shape mismatch. Spec: {} ({}), Value: {} ({})'.format(
          action_space, space_shape, fixed_value, value_shape))

  @property
  def name(self) -> str:
    return 'Fixed'

  def spec(self) -> core.Spec:
    return specs.BoundedArray((0,), np.float32, minimum=[], maximum=[], name='')

  def project(self, action: np.ndarray) -> np.ndarray:
    return self._action_space.project(self._fixed_value)

class SequentialActionSpace(core.ActionSpace):
  """Apply a sequence of ActionSpaces iteratively.

  This allows users to compose action transformations, or apply a transformation
  to a subset of a larger action space, e.g. one sliced out by `prefix_slicer`.
  """

  def __init__(self,
               action_spaces: Sequence[core.ActionSpace[core.Spec]],
               name: Optional[str] = None):
    """Initialize SequentialActionSpace.

    Args:
      action_spaces: A sequence of action spaces to apply in order.
      name: Optional name. Defaults to 0th action space name.
    """
    self._action_spaces = action_spaces
    self._name = name or action_spaces[0].name

  @property
  def name(self) -> str:
    return self._name

  def spec(self) -> core.Spec:
    return self._action_spaces[0].spec()

  def project(self, action: np.ndarray) -> np.ndarray:
    """Projects the action iteratively through the sequence."""
    for action_space in self._action_spaces:
      action = action_space.project(action)
    return action

class CompositeActionSpace(core.ActionSpace[specs.BoundedArray]):
  """Composite Action Space consisting of other action spaces.

  Assumptions (which are verified):
  1. The input action spaces all `project` to the same shape space. I.e. the
     output of project for each space has the same shape.
  2. The input action spaces are all one dimensional.
  3. The input action spaces are all ActionSpace[BoundedArray] instances.

  Behaviour:
  The outputs of the composed action spaces are merged not by concatenation -
  they all output arrays of the same shape - but by using NaN as a sentinel, as
  as missing value.  An example merging:
    input1: [1.0, 2.0, NaN, NaN, Nan]
    input2: [NaN, NaN, NaN, 4.0, 5.0]
    merged: [1.0, 2.0, NaN, 4.0, 5.0]

  It is invalid if there is a position where both inputs have a non-NaN value.
  """

  def __init__(self,
               action_spaces: Iterable[core.ActionSpace[specs.BoundedArray]],
               name: Optional[str] = None):
    if not action_spaces:
      logging.warning('CompositeActionSpace created with no action_spaces.')
    sub_specs = [space.spec() for space in action_spaces]
    for spec in sub_specs:
      if not isinstance(spec, specs.BoundedArray):
        raise ValueError('spec {} is not a BoundedArray, (type: {})'.format(
            spec, type(spec)))

    sizes = [spec.shape[0] for spec in sub_specs]
    minimums = [spec.minimum for spec in sub_specs]
    maximums = [spec.maximum for spec in sub_specs]

    if sub_specs:
      spec_dtype = np.find_common_type([spec.dtype for spec in sub_specs], [])
      min_dtype = np.find_common_type([lim.dtype for lim in minimums], [])
      max_dtype = np.find_common_type([lim.dtype for lim in maximums], [])
      minimum = np.concatenate(minimums).astype(min_dtype)
      maximum = np.concatenate(maximums).astype(max_dtype)
    else:
      # No input spaces; we have to default the data type.
      spec_dtype = np.float32
      minimum = np.asarray([], dtype=spec_dtype)
      maximum = np.asarray([], dtype=spec_dtype)

    self._component_action_spaces = action_spaces
    self._composite_action_spec = specs.BoundedArray(
        shape=(sum(sizes),),
        dtype=spec_dtype,
        minimum=minimum,
        maximum=maximum,
        name='\t'.join([spec.name for spec in sub_specs if spec.name]))

    self._name = name or '_'.join(
        [space.name for space in action_spaces if space.name])

  @property
  def name(self) -> str:
    return self._name

  def spec(self) -> specs.BoundedArray:
    return self._composite_action_spec

  # Profiling for .wrap('CompositeActionSpace.project')
  def project(self, action: np.ndarray) -> np.ndarray:
    if not self._component_action_spaces:
      return np.asarray([], dtype=self.spec().dtype)

    # Check input value has correct shape (and legal values).
    spec_utils.validate(self._composite_action_spec, action, ignore_nan=True)

    cur_action = None  # type: np.ndarray
    for action_space, action_component in zip(self._component_action_spaces,
                                              self._action_components(action)):
      projection = action_space.project(action_component)
      if cur_action is None:
        cur_action = np.full(
            projection.shape, fill_value=np.nan, dtype=projection.dtype)
      elif not np.all(cur_action.shape == projection.shape):
        raise ValueError(f'Projected actions differ in shape! cur_action: '
                         f'{cur_action.shape}, projection: {projection.shape}')
      cur_empty_indices = np.isnan(cur_action)
      proj_empty_indices = np.isnan(projection)
      assert np.all(
          np.logical_or(proj_empty_indices, cur_empty_indices)
      ), 'The projection and current action empty indices do not align'
      proj_valid_indices = np.logical_not(proj_empty_indices)
      cur_action[proj_valid_indices] = projection[proj_valid_indices]

    assert cur_action is not None, 'Program error, no action created!'

    return cur_action

  def _action_components(
      self, action: np.ndarray) -> Generator[np.ndarray, None, None]:
    start_index = 0
    for action_space in self._component_action_spaces:
      input_length = action_space.spec().shape[0]
      end_index = start_index + input_length

      assert end_index <= self._composite_action_spec.shape[0]
      action_component = action[start_index:end_index]
      start_index = end_index

      yield action_component
    assert start_index == self._composite_action_spec.shape[0]

def constrained_action_spec(minimum: Union[float, Sequence[float]],
                            maximum: Union[float, Sequence[float]],
                            base: specs.BoundedArray) -> specs.BoundedArray:
  """Returns action spec bounds constrained by the given minimum and maximum.

  Args:
    minimum: The new minimum spec constraint.
    maximum: The new maximum spec constraint.
    base: The base action space spec.
  """

  minimum = np.array(minimum, dtype=base.dtype)
  if minimum.shape != base.minimum.shape:
    raise ValueError('minimum not compatible with base shape')
  minimum = np.maximum(base.minimum, minimum)

  maximum = np.array(maximum, dtype=base.dtype)
  if maximum.shape != base.maximum.shape:
    raise ValueError('maximum not compatible with base shape')
  maximum = np.minimum(base.maximum, maximum)

  # Check that mins and maxs are non intersecting.
  if np.any(minimum > maximum):
    raise ValueError('minimum and maximum bounds intersect')

  return specs.BoundedArray(
      shape=base.shape,
      dtype=base.dtype,
      minimum=minimum,
      maximum=maximum,
      name=base.name)

def constrained_action_space(
    minimum: Union[float, Sequence[float]],
    maximum: Union[float, Sequence[float]],
    base: core.ActionSpace[specs.BoundedArray],
    name: Optional[str] = None) -> core.ActionSpace[specs.BoundedArray]:
  """Returns an action space that is a constrained version of the base space."""

  spec = constrained_action_spec(minimum, maximum, base.spec())
  constrained_space = core.IdentityActionSpace(spec)
  return SequentialActionSpace([constrained_space, base], name)


text:
"""
NSLS-II insertion device commissioning/operation

copyright (C) 2014, Yongjun Li, Yoshi Hidaka, Lingyun Yang
"""

import aphla as ap
import itertools
import numpy as np
import re
import h5py

_params = {
    "dw100g1c08u":
        {"unitsys": "phy",
         "gap": (119.0, 147.0, 30, 0.1),
         "cch": ("cch0", "cch1", "cch2", "cch3", "cch4", "cch5"),
         "background": {"gap": 147.0},
         "Imin": 2.0, # mA
         "Tmin": 2.0, # hour
         "timeout": 150, },
    "dw100g1c08d":
        {"unitsys": "phy",
         "gap": (119.0, 147.0, 30, 0.1),
         "cch": ("cch0", "cch1", "cch2", "cch3", "cch4", "cch5"),
         "background": {"gap": 147.0},
         "Imin": 2.0, # mA
         "Tmin": 2.0, # hour
         "timeout": 150, },
    "dw100g1c18u":
        {"unitsys": "phy",
         "gap": (119.0, 147.0, 30, 0.1),
         "cch": ("cch0", "cch1", "cch2", "cch3", "cch4", "cch5"),
         "background": {"gap": 147.0},
         "Imin": 2.0, # mA
         "Tmin": 2.0, # hour
         "timeout": 150, },
    "dw100g1c18d":
        {"unitsys": "phy",
         "gap": (119.0, 147.0, 30, 0.1),
         "cch": ("cch0", "cch1", "cch2", "cch3", "cch4", "cch5"),
         "background": {"gap": 147.0},
         "Imin": 2.0, # mA
         "Tmin": 2.0, # hour
         "timeout": 150, },
    "dw100g1c28u":
        {"unitsys": "phy",
         "gap": (119.0, 147.0, 30, 0.1),
         "cch": ("cch0", "cch1", "cch2", "cch3", "cch4", "cch5"),
         "background": {"gap": 147.0},
         "Imin": 2.0, # mA
         "Tmin": 2.0, # hour
         "timeout": 150, },
    "dw100g1c28d":
        {"unitsys": "phy",
         "gap": (119.0, 147.0, 30, 0.1),
         "cch": ("cch0", "cch1", "cch2", "cch3", "cch4", "cch5"),
         "background": {"gap": 147.0},
         "Imin": 2.0, # mA
         "Tmin": 2.0, # hour
         "timeout": 150, },
    }

def getBrho(E_GeV):
    """
    """

    import scipy.constants as const

    m_e_GeV = const.m_e*const.c*const.c/const.e/1e9
    gamma = E_GeV / m_e_GeV

    Brho = const.m_e * const.c * gamma / const.e # magnetic rigidity [T*m]

    return Brho

def putPar(ID, parList, **kwargs):
    """
    Put (write) a set of parameters (list) on an ID while the hardware
    itself (motor control) checks whether the target state is reached or not.

    inputs:
     ID: aphla ID instance

     parList: 2d parameter list in the format of [name, value, tolerance]
     [['gap',15,1e-4],['phase',12,1e-4]]

     timeout: Maximum time the motor control should wait for each "put"
     in the unit of seconds.

     verbose: integer larger means more details.
     throw: raise exception if True, otherwise return False

    returns: True if success, otherwise False
    """
    timeout = kwargs.get("timeout", _params[ID.name].get("timeout", 150))
    unitsys = kwargs.get("unitsys", _params[ID.name].get("unitsys", 'phy'))
    throw = kwargs.get("throw", True)
    verbose = kwargs.get("verbose", 0)

    agree = True
    for par in parList:
        ID.put(par[0], par[1], timeout=timeout, unitsys=unitsys, trig=1)
        p0 = ID.get(par[0], unitsys=unitsys)
        if abs(p0-par[1]) <= par[2]:
            continue
        # error handling
        agree = False
        if verbose:
            print 'For "{0}" of {1}:'.format(par[0], ID.name)
            print 'Target SP = {0:.9g}, Current RB = {1:.9g}, Tol = {2:.9g}'.\
                format(par[1], p0, par[2])
        if throw:
            raise RuntimeError('Failed to set device within tolerance.')
        else:
            break
    return agree

def createCorrectorField(ID):
    return [(ID, fld) for fld in _params[ID.name].get("cch", [])]

def createParList(ID, parScale):
    """
    create parameter list based on the paraneter range, spaced type
     parRange: 2d parameter range in the format of
     [[name, spacedType, start, end, step, tolerance],...]
     example:[['gap','log',150,15,21,0.1]]
     scan table will cover 15~150 with 21 steps, tolerance is 0.1,
     spacedType: log or linear

    return parameter list for communicating with hardware, table for data
    archive
    """
    nlist,vlist,tlist = [],[],[] #name, value and tolerance list
    for fld, scale in parScale:
        if not _params[ID.name].get(fld, None): continue
        nlist.append(fld)
        vmin, vmax, vstep, vtol = _params[ID.name][fld]
        if scale == 'linear':
            vlist.append(list(np.linspace(vmin, vmax, int(vstep))))
        elif scale == 'log':
            if vmin<=0 or vmax<=0:
                raise RuntimeError('negative boundary can not be spaced Logarithmically')
            else:
                vlist.append(list(np.logspace(np.log10(vmin),np.log10(vmax),int(vstep))))
        else:
            raise RuntimeError('unknown spaced pattern: %s'%p[1])
        tlist.append(vtol)
    valueList = itertools.product(*vlist)
    parList = []
    for v in valueList:
        tmp = []
        for i,n in enumerate(nlist):
            tmp.append([n,v[i],tlist[i]])
        parList.append(tmp)
    valueList = itertools.product(*vlist)
    table = np.array([vi for vi in valueList])
    return parList, nlist, table

def putParHardCheck(ID, parList, timeout=30, throw=True, unitsys='phy'):
    '''
    Put (write) a set of parameters (list) on an ID while the hardware
    itself (motor control) checks whether the target state is reached or not.

    ID: aphla ID instance

    parList: 2d parameter list in the format of [name, value, tolerance]
    [['gap',15,1e-4],['phase',12,1e-4]]

    timeout: Maximum time the motor control should wait for each "put"
    in the unit of seconds.

    return: True if success, otherwise throws an exception.
    '''

    agree = True
    for par in parList:
        ID.put(par[0], par[1], timeout=timeout, unitsys=unitsys, trig=1)
        # raw unit for "gap"   = [um]
        # raw unit for "phase" = [um?]

        p0 = ID.get(par[0], unitsys=unitsys)

        if abs(p0-par[1]) <= par[2]: # TODO: readback & setpoint unit may be different! Check it!
            continue # print "Agree: ", p0, par[1], "eps=", par[2]
        # error handling
        agree = False
        print 'For "{0}" of {1}:'.format(par[0], ID.name)
        print 'Target Setpoint = {0:.9g}, Current Readback = {1:.9g}, Tolerance = {2:.9g}'.format(
            par[1], p0, par[2])
        if throw:
            raise RuntimeError('Failed to set device within tolerance.')
        else:
            break
    return agree

# <codecell>

def putParSoftCheck(ID, parList, timeout=30, online=False):
    '''
    Put (write) a set of parameters (list) on an ID while this function
    checks whether the target state is reached or not through readbacks
    for given tolerances.

    ID: aphla ID instance

    parList: 2d parameter list in the format of [name, value, tolerance]
    [['gap',15,1e-4],['phase',12,1e-4]]

    timeout: Maximum time the motor control should wait for each "put"
    in the unit of seconds.

    return: True if success, otherwise throws an exception.
    '''

    if not online: return True # TODO: To be reomved once we are allowed to move ID motors

    for par in parList:
        t0 = datetime.now()
        converged = False

        try:
            ID.put(par[0], par[1], unitsys=None) # raw unit
        except:
            print 'Failed to set the setpoint for {0} to {1}'.format(par[0], par[1])
            raise

        # TODO: remove hardcoding
        ap.caput("SR:C28-ID:G1{DW100:2}ManG:Go_.PROC", 1, wait=False)

        while not converged:
            p0 = ID.get(par[0], unitsys=None)
            if abs(p0-par[1]) <= par[2]: # TODO: readback & setpoint unit may be different! Check it!
                # print "Agree: ", p0, par[1], "eps=", par[2]
                converged = True
                break
            t1 = datetime.now()
            if (t1-t0).total_seconds() > timeout:
                break
            time.sleep(0.5)
        if not converged:
            raise RuntimeError("timeout at setting {0}={1} (epsilon={2})".format(par[0], par[1], par[2]))

    return True

def putBackground(ID, **kwargs):
    """
    put ID to passive status,
    gap to max, phase to 0 if apply, all correction cch to zeros
    """
    gapMin, gapMax, gapStep, gapTol = kwargs.get("gap",
                                                 _params[ID.name]["gap"])
    phaseMin, phaseMax, phaseStep, phaseTol = \
        kwargs.get("phase",  _params[ID.name].get("phase", (None, None, None, None)))
    zeroPhase = 0.0
    timeout = kwargs.get("timeout", 150)
    throw   = kwargs.get("throw", True)
    unitsys = kwargs.get("unitsys", 'phy')
    verbose = kwargs.get("verbose", 0)

    flds = ID.fields()
    parList = []
    if 'gap' in flds:
        parList.append(['gap',gapMax,gapTol])
    if 'phase' in flds:
        parList.append(['phase',zeroPhase,phaseTol])

    if putPar(ID, parList, timeout=timeout,
              throw=throw, unitsys=unitsys, verbose=verbose):
        # put correcting coils to zeros
        for i in range(len(ID.cch)):
            ID.put('cch'+str(i), 0.0, unitsys=None)
        return True
    else:
        return False

def checkBeam(Imin=2, Tmin=2, online=False):
    """
    check beam life time and current
    if beam lifetime is less than Tmin [hr], 2hrs by default,
    or current is less then Imin [mA], 2mA by default
    return False, otherwise True
    """
    tau, Ib = ap.getLifetimeCurrent()
    if Ib < Imin:
        print 'Beam current too low ({0} < {1})'.format(Ib, Imin)
        return False
    if tau < Tmin:
        print 'Beam lifetime too short ({0} < {1})'.format(tau, Tmin)
        return False
    return True

def checkGapPhase(ID, **kwargs):
    """
    check ID gap, phase
    return True if success, otherwise False
    """
    gapMin, gapMax, gapStep, gapTol = kwargs.get("gap",
                                                 _params[ID.name]["gap"])
    phaseMin, phaseMax, phaseStep, phaseTol = \
        kwargs.get("phase",  _params[ID.name].get("phase", (None, None, None, None)))
    timeout = kwargs.get("timeout", 150)
    throw   = kwargs.get("throw", True)
    unitsys = kwargs.get("unitsys", _params[ID.name]["unitsys"])
    verbose = kwargs.get("verbose", 0)
    gapStep = kwargs.get("gapStep", gapStep)
    phaseStep = kwargs.get("phaseStep", phaseStep)

    flds = ID.fields()
    if 'gap' in flds:
        for gap in np.linspace(gapMin, gapMax, gapStep):
            gapList = [['gap',gap, gapTol]]
            gapStatus = putPar(ID,gapList,timeout=timeout,
                               throw=throw,unitsys=unitsys,verbose=verbose)
            if not gapStatus:
                return False
    if 'phase' in flds:
        for phase in np.linspace(phaseMin,phaseMax,phaseStep):
            phaseList = [['phase',phase,phaseTol]]
            phaseStatus = putPar(ID,phaseList,timeout=timeout,
                               throw=throw,unitsys=unitsys,verbose=verbose)
            if not phaseStatus:
                return False
    return True

def switchFeedback(fftable = "off"):
    """
    switchFeedback("on") or "off"
    """
    if fftable not in ["on", "off"]:
        raise RuntimeError("invalid feed forward table state: ('on'|'off')")

    for dw in ap.getGroupMembers(["DW",], op="union"):
        if "gap" not in dw.fields():
            print "WARNING: no 'gap' field in {0}".format(dw.name)
            continue
        pv = dw.pv(field="gap", handle="setpoint")[0]
        m = re.match(r"([^\{\}]+)\{(.+)\}", pv)
        if not m:
            print "WARNING: inconsistent naming '{0}'".format(pv)
        pvffwd = "{0}{{{1}}}MPS:Lookup_.INPA".format(m.group(1), m.group(2))
        pvffwd_pref = "{0}{{{1}-Mtr:Gap}}.RBV ".format(m.group(1), m.group(2))

        pvffwd_val = {"on": pvffwd_pref + "CP NM",
                      "off": pvffwd_pref + "NPP N"}
        print "set {0}='{1}'".format(pvffwd, pvffwd_val[fftable])
        ap.caput(pvffwd, pvffwd_val[fftable])

    # fast/slow co
    # all ID feed forward
    # weixing Bunch by Bunch

def initFile(ID, fieldList, parTable):
    """initilize file name with path, save parameter table to hdf5"""
    fileName = ap.outputFileName("ID", ID.name+"_")
    fid = h5py.File(fileName)
    grp = fid.require_group(ID.name)
    grp.attrs["__FORMAT__"] = 1
    # setup parameters
    subg = grp.require_group("parameters")
    subg["scanTable"] = parTable #
    subg["scanTable"].attrs["columns"] = fieldList
    #for p in nameList:
    #    subg["scanTable"].attrs[p] = []
    bkg = _params[ID.name]["background"]
    # like one row of scanTable, same columns
    subg["background"] = [bkg[fld] for fld in fieldList]
    subg["background"].attrs["columns"] = fieldList
    # timestamp ISO "2007-03-01 13:00:00"
    subg["minCurrent"] = _params[ID.name]["Imin"]
    subg["minCurrent"].attrs["unit"] = "mA"
    subg["minLifetime"] = _params[ID.name]["Tmin"]
    subg["minLifetime"].attrs["unit"] = "hr"
    fid.close()

    return fileName

def chooseBpmCor(ID, userBpm=False):
    """
    choose bpm and corrector
    """
    bpms = ap.getElements('BPM')
    if userBpm:
        bpms += ap.getElements('UBPM')
    bpmfields = []
    for bpm in bpms:
        bpmflds.append([bpm,'x'])
        bpmflds.append([bpm,'y'])

    corfields = []
    for i in range(len(ID.cch)):
        corflds.append([ID,'cch'+'%i'%i])

    return bpmFields, corFields

def saveToDB(fileName):
    print "save to file (Guobao's DB)"
    pass

def measBackground(ID, output, iiter):
    """measure the background and return saved group name"""
    if not nsls2id.putBackground(ID):
        print "Failed at setting {0} to background mode".format(ID)
        return None
    # create background subgroup with index
    fid = h5py.File(output)
    prefix = "iter_"
    iterindex = max([int(g[len(prefix):]) for g in fid[ID.name].keys()
                     if g.startswith(prefix)] + [-1]) + 1
    bkgGroup = "iter_{0:04d}".format(iterindex)
    grp = fid[ID.name].create_group(bkgGroup)
    orb0 = ap.getOrbit(spos=True)
    grp["orbit"] = orb0
    tau, I = ap.getLifetimeCurrent()
    grp["lifetime"] = tau
    grp["current"] = I
    grp.attrs["iter"] = iiter
    fid.close()
    return bkgGroup

def virtKicks2FldInt(virtK1, virtK2, idLen, idKickOffset1, idKickOffset2, E_GeV):
    """
    Calculate the 1st and 2nd field integrals of an insertion device (ID)
    from the given upstream/downstream virtual kicks.

    Parameters
    ----------
    virtK1, virtK2 : float
       Virtual kick values [rad] at the upsteam and downstream of the ID,
       respectively.

    idLen : float
       Length of the ID [m].

    idKickOffset1, idKickOffset2 : float
       Position offset [m] of virtual kicks with respect to the undulator
       extremeties. `idKickOffset1` == 0 means that the upstream virtual kick
       is exactly located at the upstream entrance of the ID. If `idKickOffset1`
       is a positive value, then the virtual kick is inside of the ID by the
       amount `idKickOffset1`. If negative, the virtual kick is outside of the
       ID by the absolute value of `idKickOffset1`. The same is true for the
       downstream side.

    E_GeV : float
       Electron beam energy [GeV].

    Returns
    -------
    I1 : float
       First field integral [G*m].

    I2 : float
       Second field integral [G*(m^2)].
    """

    Brho = getBrho(E_GeV) # magnetic rigidity [T*m]

    common = Brho * 1e4
    I1 = common * (virtK1 + virtK2) # [G*m]
    I2 = common * ((idLen-idKickOffset1)*virtK1 + idKickOffset2*virtK2) # [G*(m^2)]

    return I1, I2

# <codecell>

def fldInt2VirtKicks(I1, I2, idLen, idKickOffset1, idKickOffset2, E_GeV):
    """
    Calculate upstream/downstream virtual kicks from the given 1st and 2nd field
    integrals of an insertion device (ID).

    Parameters
    ----------
    I1 : float
       First field integral [G*m].

    I2 : float
       Second field integral [G*(m^2)].

    idLen : float
       Length of the ID [m].

    idKickOffset1, idKickOffset2 : float
       Position offset [m] of virtual kicks with respect to the undulator
       extremeties. `idKickOffset1` == 0 means that the upstream virtual kick
       is exactly located at the upstream entrance of the ID. If `idKickOffset1`
       is a positive value, then the virtual kick is inside of the ID by the
       amount `idKickOffset1`. If negative, the virtual kick is outside of the
       ID by the absolute value of `idKickOffset1`. The same is true for the
       downstream side.

    E_GeV : float
       Electron beam energy [GeV].

    Returns
    -------
    virtK1, virtK2 : float
       Virtual kick values [rad] at the upsteam and downstream of the ID,
       respectively.
    """

    Brho = getBrho(E_GeV) # magnetic rigidity [T*m]

    common = 1e-4 / Brho / (idLen-idKickOffset1-idKickOffset2)
    virtK1 = common * (I2 - I1 * idKickOffset2)         # [rad]
    virtK2 = common * (I1 * (idLen-idKickOffset1) - I2) # [rad]

    return virtK1, virtK2

def save1DFeedFowardTable(filepath, table, fmt='%.16e'):
    """
    Save a valid 1-D Stepped Feedforward table (NSLS-II format) to a text file.
    """

    np.savetxt(filepath, table, fmt=fmt, delimiter=', ', newline='\n')

def get1DFeedForwardTable(centers, half_widths, dI_array,
                          I0_array=None, fmt='%.16e'):
    """
    Get a valid 1-D Stepped Feedforward table (NSLS-II format)
    """

    if I0_array is None:
        I_array = dI_array
    else:
        I_array = I0_array + dI_array

    table = np.hstack((np.array(centers).reshape((-1,1)),
                       np.array(half_widths).reshape((-1,1)),
                       I_array))

    return table

def getZeroed1DFeedForwardTable(parDict, nIDCor):
    """
    Get a valid 1-D Stepped Feedforward table (NSLS-II format) with
    all ID correctors being set to zero for all the entire range of
    ID property specified in "parDict".
    """

    try:
        scanVectors = parDict['vectors']
        bkgList = parDict['bkgTable'].flatten().tolist()

        assert len(scanVectors) == len(bkgList) == 1
    except:
        print 'len(scanVectors) = {0:d}'.format(len(scanVectors))
        print 'len(bkgList) = {0:d}'.format(len(bkgList))
        print 'This function is only for 1D feedforward table.'
        raise RuntimeError(('Lengths of "scanVectors" and "bkgList" must be 1.'))

    array = scanVectors[0] + [bkgList[0]]
    minVal, maxVal = np.min(array), np.max(array)

    centers = [(minVal + maxVal) / 2.0]
    half_widths = [(maxVal - minVal) / 2.0 * 1.01] # Extra margin of 1% added
    dI_array = np.array([0.0]*nIDCor).reshape((1,-1))

    return get1DFeedForwardTable(centers, half_widths, dI_array,
                                 I0_array=None, fmt='%.16e')

def create1DFeedForwardTable(centers, half_widths, dI_array, I0_array=None):
    """
    Create a valid 1-D Stepped Feedforward table (NSLS-II format)
    """

    if I0_array is None:
        I_array = dI_array
    else:
        I_array = I0_array + dI_array

    table = np.hstack((np.array(centers).reshape((-1,1)),
                       np.array(half_widths).reshape((-1,1)),
                       I_array))

    return table

def calc1DFeedForwardColumns(
    ID_filepath, n_interp_pts=None, interp_step_size=None, step_size_unit=None,
    cor_inds_ignored=None, bpm_inds_ignored=None, nsv=None):
    """
    """

    # TODO: Make sure all the units are correct in the generated table
    # Gap & interval are in microns => [um]
    # Currents in ppm of 10 Amps => [10uA]

    compIterInds = getCompletedIterIndexes(ID_filepath)
    nCompletedIter = len(compIterInds)

    f = h5py.File(ID_filepath, 'r')

    ID_name = f.keys()[0]
    grp = f[ID_name]

    meas_state_1d_array = grp['parameters']['scanTable'].value
    state_unitsymb = grp['parameters']['scanTable'].attrs['unit'] # TODO: need unit conversion
    nIter, ndim = meas_state_1d_array.shape
    if ndim != 1:
        f.close()
        raise NotImplementedError('Only 1-D scan has been implemented.')
    if nCompletedIter != nIter:
        print '# of completed scan states:', nCompletedIter
        print '# of requested scan states:', nIter
        f.close()
        raise RuntimeError('You have not scanned all specified states.')
    meas_state_1d_array = meas_state_1d_array.flatten()
    state_min = np.min(meas_state_1d_array)
    state_max = np.max(meas_state_1d_array)

    if (n_interp_pts is not None) and (interp_step_size is not None):
        f.close()
        raise ValueError(('You can only specify either one of "n_interp_pts" '
                          'or "interp_step_size", not both.'))
    elif n_interp_pts is not None:
        interp_state_1d_array = np.linspace(state_min, state_max, n_interp_pts)
    elif interp_step_size is not None:
        interp_state_1d_array = np.arange(state_min, state_max, interp_step_size)
        if interp_state_1d_array[-1] != state_max:
            interp_state_1d_array = np.array(interp_state_1d_array.tolist()+
                                             [state_max])
    else:
        interp_state_1d_array = meas_state_1d_array

    M_list        = [None]*nIter
    diff_orb_list = [None]*nIter
    for k in grp.keys():
        if k.startswith('iter_'):

            iIter = grp[k].attrs['iteration']

            M_list[iIter] = grp[k]['orm']['m'].value

            orb = grp[k]['orbit'].value

            bkgGroup = grp[k].attrs['background']
            orb0 = grp[bkgGroup]['orbit'].value[:,:-1] # Ignore s-pos column

            diff_orb_list[iIter] = orb - orb0

    f.close()

    interp_state_1d_array = np.sort(interp_state_1d_array)

    center_list = interp_state_1d_array.tolist()

    half_width_list = (np.diff(interp_state_1d_array)/2.0).tolist()
    half_width_list.append(center_list[-1]-center_list[-2]-half_width_list[-1])

    dI_list = []
    for M, diff_orb in zip(M_list, diff_orb_list):

        TF = np.ones(diff_orb.shape)
        if bpm_inds_ignored is not None:
            for i in bpm_inds_ignored:
                TF[i,:] = 0
        TF = TF.astype(bool)

        diff_orb_trunc = diff_orb[TF].reshape((-1,2))

        # Reverse sign to get desired orbit change
        dObs = (-1.0)*diff_orb_trunc.T.flatten().reshape((-1,1))

        TF = np.ones(M.shape)
        if cor_inds_ignored is not None:
            for i in cor_inds_ignored:
                TF[:,i] = 0
        if bpm_inds_ignored is not None:
            nBPM = M.shape[0]/2
            try:
                assert nBPM*2 == M.shape[0]
            except:
                raise ValueError('Number of rows for response matrix must be 2*nBPM.')
            for i in bpm_inds_ignored:
                TF[i     ,:] = 0
                TF[i+nBPM,:] = 0
        TF = TF.astype(bool)

        M_trunc = M[TF].reshape((dObs.size,-1))

        U, sv, V = np.linalg.svd(M_trunc, full_matrices=0, compute_uv=1)

        S_inv = np.linalg.inv(np.diag(sv))
        if nsv is not None:
            S_inv[nsv:, nsv:] = 0.0

        dI = V.T.dot(S_inv.dot(U.T.dot(dObs))).flatten().tolist()

        if cor_inds_ignored is not None:
            for i in cor_inds_ignored:
                # Set 0 Amp for unused correctors
                dI.insert(i, 0.0)

        dI_list.append(dI)

    dI_array = np.array(dI_list)
    nCor = dI_array.shape[1]
    interp_dI_array = np.zeros((interp_state_1d_array.size, nCor))
    for i in range(nCor):
        interp_dI_array[:,i] = np.interp(
            interp_state_1d_array, meas_state_1d_array, dI_array[:,i])

    return {'centers': np.array(center_list),
            'half_widths': np.array(half_width_list),
            'raw_dIs': dI_array, 'interp_dIs': interp_dI_array}

#----------------------------------------------------------------------
def getCompletedIterIndexes(ID_filepath):
    """
    """

    f = h5py.File(ID_filepath, 'r')

    ID_name = f.keys()[0]
    grp = f[ID_name]

    completed_iter_indexes = []
    for k in grp.keys():
        if k.startswith('iter_') and grp[k].attrs.has_key('completed'):
            completed_iter_indexes.append(grp[k].attrs['iteration'].value)

    f.close()

    if completed_iter_indexes != []:
        if not np.all(np.diff(completed_iter_indexes) == 1):
            raise RuntimeError(
                'List of completed iteration indexes has some missing indexes.')
        if np.min(completed_iter_indexes) != 0:
            raise RuntimeError('List of completed iteration indexes does not start from 0.')

    return completed_iter_indexes

if __name__ == '__main__':

    ID_filepath = '/epics/data/aphla/SR/2014_09/ID/dw100g1c08u_2014_09_24_142644.hdf5'

    d = calc1DFeedForwardColumns(ID_filepath, interp_step_size=1.0,
                                 cor_inds_ignored=[2,3],
                                 bpm_inds_ignored=None, nsv=None)
    table = create1DFeedForwardTable(d['centers'], d['half_widths'],
                                     d['interp_dIs'])
    save1DFeedFowardTable('test_ff.txt', table, fmt='%.16e')

max_stars_repo_name: NSLS-II/aphla

max_issues_repo_name: NSLS-II/aphla

max_forks_repo_name: NSLS-II/aphla

reason: import numpy,import scipy


text:
import numpy as np

# https://github.com/rbgirshick/fast-rcnn/blob/b612190f279da3c11dd8b1396dd5e72779f8e463/lib/fast_rcnn/test.py#L306
def nms(dets, thresh):
    x1 = dets[:, 0]
    y1 = dets[:, 1]
    x2 = dets[:, 2]
    y2 = dets[:, 3]
    scores = dets[:, 4]

    areas = (x2 - x1 + 1) * (y2 - y1 + 1)
    order = scores.argsort()[::-1]

    keep = []
    while order.size > 0:
        i = order[0]
        keep.append(i)
        xx1 = np.maximum(x1[i], x1[order[1:]])
        yy1 = np.maximum(y1[i], y1[order[1:]])
        xx2 = np.minimum(x2[i], x2[order[1:]])
        yy2 = np.minimum(y2[i], y2[order[1:]])

        w = np.maximum(0.0, xx2 - xx1 + 1)
        h = np.maximum(0.0, yy2 - yy1 + 1)
        inter = w * h
        ovr = inter / (areas[i] + areas[order[1:]] - inter)

        inds = np.where(ovr <= thresh)[0]
        order = order[inds + 1]

    return keep

# https://towardsdatascience.com/non-maximum-suppression-nms-93ce178e177c
# Malisiewicz et al.

def non_max_suppression_fast(boxes, overlapThresh):
	# if there are no boxes, return an empty list
	if len(boxes) == 0:
		return []
	# if the bounding boxes integers, convert them to floats --
	# this is important since we'll be doing a bunch of divisions
	if boxes.dtype.kind == "i":
		boxes = boxes.astype("float")
	# initialize the list of picked indexes
	pick = []
	# grab the coordinates of the bounding boxes
	x1 = boxes[:,0]
	y1 = boxes[:,1]
	x2 = boxes[:,2]
	y2 = boxes[:,3]
	# compute the area of the bounding boxes and sort the bounding
	# boxes by the bottom-right y-coordinate of the bounding box
	area = (x2 - x1 + 1) * (y2 - y1 + 1)
	idxs = np.argsort(y2)
	# keep looping while some indexes still remain in the indexes
	# list
	while len(idxs) > 0:
		# grab the last index in the indexes list and add the
		# index value to the list of picked indexes
		last = len(idxs) - 1
		i = idxs[last]
		pick.append(i)
		# find the largest (x, y) coordinates for the start of
		# the bounding box and the smallest (x, y) coordinates
		# for the end of the bounding box
		xx1 = np.maximum(x1[i], x1[idxs[:last]])
		yy1 = np.maximum(y1[i], y1[idxs[:last]])
		xx2 = np.minimum(x2[i], x2[idxs[:last]])
		yy2 = np.minimum(y2[i], y2[idxs[:last]])
		# compute the width and height of the bounding box
		w = np.maximum(0, xx2 - xx1 + 1)
		h = np.maximum(0, yy2 - yy1 + 1)
		# compute the ratio of overlap
		overlap = (w * h) / area[idxs[:last]]
		# delete all indexes from the index list that have
		idxs = np.delete(idxs, np.concatenate(([last],
			np.where(overlap > overlapThresh)[0])))
	# return only the bounding boxes that were picked using the
	# integer data type
	return boxes[pick].astype("int")

max_stars_repo_name: gregbugaj/form-processor

max_issues_repo_name: gregbugaj/form-processor

max_forks_repo_name: gregbugaj/form-processor


text:
import pathlib
import json
import numpy as np

from registry import registry
from eval_settings.eval_setting_base import EvalSetting, StandardDataset

def validate_vid_robust(logits, targets, image_paths, py_model, merge_op='max'):
    if 'smoothing' in py_model.name:
        merge_op = 'sum'

    with open((pathlib.Path(__file__).parent / 'imagenet-vid-robust_metadata/class_idx_map.json').resolve()) as f:
        CLASS_IDX_LIST = sorted(list(map(int, json.loads(f.read()).keys())))

    with open((pathlib.Path(__file__).parent / 'imagenet-vid-robust_metadata/rev_class_idx_map.json').resolve()) as f:
        REV_CLASS_IDX_MAP = json.loads(f.read())

    with open((pathlib.Path(__file__).parent / 'imagenet-vid-robust_metadata/labels.json').resolve()) as f:
        labels = json.load(f)

    with open((pathlib.Path(__file__).parent / 'imagenet-vid-robust_metadata/pmsets.json').resolve()) as f:
        pmsets = json.loads(f.read())

    def project_imagenet_predictions_imagenet_vid(preds, merge_op='max'):
        new_preds = np.zeros((preds.shape[0], 30))
        for k,v in REV_CLASS_IDX_MAP.items():
            v = list(map(lambda x: CLASS_IDX_LIST.index(x), v))
            if (merge_op == 'mean'):
                new_preds[:, int(k)] = np.mean(preds[:, v], axis=1).squeeze()
            elif (merge_op == 'median'):
                new_preds[:, int(k)] = np.median(preds[:, v], axis=1).squeeze()
            elif (merge_op == 'max'):
                new_preds[:, int(k)] = np.max(preds[:, v], axis=1).squeeze()
            elif (merge_op == 'sum'):
                new_preds[:, int(k)] = np.sum(preds[:, v], axis=1)
            else:
                raise Exception(f'unsupported merge operation {merge_op} not allowed')
        return new_preds

    def score_predictions(preds, pmsets, labels):
        correct_anchor = 0
        correct_pmk = 0
        N = len(pmsets)
        wrong_map = {}
        for anchor, pmset in pmsets.items():
            pmset_correct = 0
            wrongs = []
            for elem in pmset:
                if np.argmax(preds[elem]) in labels[elem]:
                    pmset_correct += 1
                else:
                    wrongs.append(elem)

            if np.argmax(preds[anchor]) in labels[anchor]:
                correct_anchor  += 1
                pmset_correct += 1
                if len(wrongs) > 0:
                    wrong_map[anchor] = wrongs[-1]

            if pmset_correct == len(pmset) + 1:
                correct_pmk += 1

        return correct_anchor/N, correct_pmk/N

    preds_dict = {}
    logits_projected = project_imagenet_predictions_imagenet_vid(logits.numpy(), merge_op)
    for i, img_name in enumerate(image_paths):
        preds_dict['val/' + img_name] =  logits_projected[i]

    benign,pmk = score_predictions(preds_dict, pmsets, labels)
    metrics_dict = {}
    metrics_dict['pm0'] = benign
    metrics_dict['pm10'] = pmk
    metrics_dict['merge_op'] = merge_op
    return metrics_dict

def validate_val_on_vid_robust_classes(logits, targets, image_paths, py_model, merge_op='max'):
    if 'smoothing' in py_model.name:
        merge_op = 'sum'

    with open((pathlib.Path(__file__).parent / 'imagenet-vid-robust_metadata/class_idx_map.json').resolve()) as f:
        CLASS_IDX_LIST = sorted(list(map(int, json.loads(f.read()).keys())))

    with open((pathlib.Path(__file__).parent / 'imagenet-vid-robust_metadata/class_idx_map.json').resolve()) as f:
        CLASS_IDX_MAP = {int(k): v for k, v in json.loads(f.read()).items()}

    with open((pathlib.Path(__file__).parent / 'imagenet-vid-robust_metadata/rev_class_idx_map.json').resolve()) as f:
        REV_CLASS_IDX_MAP = json.loads(f.read())

    with open((pathlib.Path(__file__).parent / 'imagenet-vid-robust_metadata/labels.json').resolve()) as f:
        labels = json.load(f)

    def project_imagenet_predictions_imagenet_vid(preds, merge_op='max'):
        new_preds = np.zeros((preds.shape[0], 30))
        for k,v in REV_CLASS_IDX_MAP.items():
            v = list(map(lambda x: CLASS_IDX_LIST.index(x), v))
            if (merge_op == 'mean'):
                new_preds[:, int(k)] = np.mean(preds[:, v], axis=1).squeeze()
            elif (merge_op == 'median'):
                new_preds[:, int(k)] = np.median(preds[:, v], axis=1).squeeze()
            elif (merge_op == 'max'):
                new_preds[:, int(k)] = np.max(preds[:, v], axis=1).squeeze()
            elif (merge_op == 'sum'):
                new_preds[:, int(k)] = np.sum(preds[:, v], axis=1)
            else:
                raise Exception(f'unsupported merge operation {merge_op} not allowed')
        return new_preds

    def score_predictions(logits_projected, targets):
        mask = np.isin(targets, CLASS_IDX_LIST)
        logits_projected, targets = logits_projected[mask], targets[mask]
        targets = [CLASS_IDX_MAP[x] for x in targets]
        preds = logits_projected.argmax(axis=1)

        acc = np.mean(np.equal(preds, targets))

        weights = np.zeros((30,))
        for _, target in labels.items():
            weights[target] += 1
        weights /= weights.sum()

        indv_acc = np.zeros((30,))
        for i in range(30):
            indv_acc[i] = np.mean(np.equal(preds[np.equal(targets, i)], i))

        return {'top1': acc * 100,
                'top1_uniform_class_weight': indv_acc.mean() * 100,
                'top1_vid_robust_balanced_class_weight': np.sum(indv_acc * weights) * 100}

    logits_projected = project_imagenet_predictions_imagenet_vid(logits.numpy(), merge_op)

    metrics_dict = score_predictions(logits_projected, targets.numpy())
    metrics_dict['merge_op'] = merge_op
    return metrics_dict

with open((pathlib.Path(__file__).parent / 'imagenet-vid-robust_metadata/class_idx_map.json').resolve()) as f:
	class_sublist = sorted(list(map(int, json.loads(f.read()).keys())))

registry.add_eval_setting(
    EvalSetting(
        name = 'imagenet-vid-robust',
        dataset = StandardDataset(name='imagenet-vid-robust'),
        size = 1109,
        metrics_fn = validate_vid_robust,
        class_sublist = class_sublist,
    )
)

idx_subsample_list = [range(x*50, (x+1)*50) for x in class_sublist]
idx_subsample_list = sorted([item for sublist in idx_subsample_list for item in sublist])

registry.add_eval_setting(
    EvalSetting(
        name = 'val-on-vid-robust-classes',
        dataset = StandardDataset(name='val'),
        size = 14400,
        metrics_fn = validate_val_on_vid_robust_classes,
        class_sublist = class_sublist,
        idx_subsample_list = idx_subsample_list,
    )
)

max_stars_repo_name: saarimrahman/imagenet-testbed

max_issues_repo_name: saarimrahman/imagenet-testbed

max_forks_repo_name: saarimrahman/imagenet-testbed


text:
#!/usr/bin/env python3

from fuzzconfig import *
import numpy as np
import os

os.system("rm -rf work_binop")
os.mkdir("work_binop")

for idx in range(num):
    with open("work_binop/binop_%02d.v" % idx, "w") as f:
        print("module top(input a, b, output y);", file=f)
        print("  assign y = a%sb;" % np.random.choice([" ^ ", " ^ ~", " & ", " & ~", " | ", " | ~"]), file=f)
        print("endmodule", file=f)
    with open("work_binop/binop_%02d.pcf" % idx, "w") as f:
        p = np.random.permutation(pins)
        print("set_io a %s" % p[0], file=f)
        print("set_io b %s" % p[1], file=f)
        print("set_io y %s" % p[2], file=f)

with open("work_binop/Makefile", "w") as f:
    print("all: %s" % " ".join(["binop_%02d.bin" % i for i in range(num)]), file=f)
    for i in range(num):
        print("binop_%02d.bin:" % i, file=f)
        print("\t-bash ../icecube.sh binop_%02d > binop_%02d.log 2>&1 && rm -rf binop_%02d.tmp || tail binop_%02d.log" % (i, i, i, i), file=f)


text:
#!/usr/bin/env python3

import tensorflow as tf
import numpy as np

def main():
	''' Create data '''
	x = np.random.rand(100).astype(np.float32)
	y = x * 0.7 + 0.5

	''' Create TensorFlow model '''
	# Define Weights and biases
	Weights = tf.Variable(tf.random_uniform([1], -1.0, 1.0))
	biases = tf.Variable(tf.zeros([1]))

	# Define training function
	y_train = Weights * x + biases

	# Define loss function and optimizer
	loss = tf.reduce_mean(tf.square(y_train - y))
	optimizer = tf.train.GradientDescentOptimizer(0.5)
	train = optimizer.minimize(loss)

	# Initialize all variables in TensorFlow
	init = tf.initialize_all_variables()

	''' Start training '''
	sess = tf.Session()
	sess.run(init)

	# Print the Weights and biases for every 20 steps
	for step in range(201):
	    sess.run(train)
	    if step % 20 == 0:
		    print('Step %3d:' % step, sess.run(Weights), sess.run(biases))

	sess.close()

''' ENTRY POINT '''
if __name__ == "__main__":
    main()

max_stars_repo_name: yungshenglu/tensorflow-practice

max_issues_repo_name: yungshenglu/tensorflow-practice

max_forks_repo_name: yungshenglu/tensorflow-practice


text:
import numpy as np
import pkg_resources
import torch
import torch.nn.functional as F

_CACHE_LOCATION = '.cache'

def bitget(byteval, idx):
    return ((byteval & (1 << idx)) != 0)

def cache_location():
    return pkg_resources.resource_filename(__name__, _CACHE_LOCATION)

def compute_cm(gt, pred, N, with_cuda=False):
    '''Computes confusion matrix between prediction and gt
    Args:
      gt (Tensor) : gt.
      pred (Tensor) : predictions.
      N (int) : number of classes in matrix.
    Returns:
      Confusion matrix
      (Tensor of size (N, N)).
    '''
    # flatten gt and preds
    pred = pred.view([-1])
    gt = gt.view([-1])

    # check for valid labels
    mask_1 = gt >= 0
    mask_2 = gt < N
    mask = torch.logical_and(mask_1, mask_2)

    # compute scatter idx
    pred = pred[mask]
    gt = gt[mask]
    gt *= N
    idx = gt + pred

    cm = torch.zeros([N * N], dtype=torch.float32)
    if with_cuda:
        cm = cm.cuda()
    samples = torch.ones_like(gt, dtype=torch.float)

    # compute confusion matrix and reshape
    cm = cm.scatter_add_(0, idx, samples)
    cm = cm.view([N, N]).type(torch.int64)

    return cm

def compute_iu(cm):
    '''Compute IU from confusion matrix.
    Args:
      cm (ndarray) : square confusion matrix.
    Returns:
      IU vector (ndarray).
    '''
    N = cm.shape[0]
    IU = np.ones(N)

    for i in range(N):
        fn = np.sum(cm[:, i])
        fp = np.sum(cm[i, :])
        tp = cm[i, i]
        total = fn + fp - tp
        if total > 0:
            IU[i] = tp / total
    return IU

def forward_multi_scale(model, img):
    # sample multi scale images
    height = img.shape[-2]
    width = img.shape[-1]

    # scales
    scales = [0.4, 0.6, 0.8, 1, 1.2]

    # predict at multiple scales
    predictions = []
    for scale in scales:

        # compute scaled height and width
        scaled_height = int(scale * height)
        scaled_width = int(scale * width)

        # interpolate image to scale
        scaled_img = F.interpolate(img, (scaled_height, scaled_width),
                                   mode='bilinear')

        # forward pass through
        logits = model(scaled_img)

        # interpolate logits back to original image size
        prediction = F.softmax(logits, dim=1)
        prediction = F.interpolate(prediction, (height, width),
                                   mode='bilinear')
        predictions.append(prediction)

    prediction = torch.cat(predictions, dim=0)
    prediction = torch.mean(prediction, dim=0)

    # average across images
    prediction = torch.argmax(prediction, dim=0)
    prediction = torch.squeeze(prediction)

    return prediction

def forward_single_scale(model, img):
    # forward pass through
    logits = model(img)

    # interpolate logits back to original image size
    prediction = F.softmax(logits, dim=1)
    prediction = F.interpolate(prediction, (img.shape[-2], img.shape[-1]),
                               mode='bilinear')
    prediction = torch.argmax(prediction, dim=1)
    prediction = torch.squeeze(prediction)

    return prediction

class ColourMap(object):

    def __init__(self, N=256, normalised=False, dataset='voc'):
        dtype = 'float32' if normalised else 'uint8'

        if dataset == 'voc' or dataset == 'nyu':
            self.map = np.zeros((N, 3), dtype=dtype)
            for i in range(N):
                r = g = b = 0
                c = i
                for j in range(8):
                    r = r | (bitget(c, 0) << 7 - j)
                    g = g | (bitget(c, 1) << 7 - j)
                    b = b | (bitget(c, 2) << 7 - j)
                    c = c >> 3

                self.map[i] = np.array([r, g, b])
        else:
            self.map = np.array(
                [[128, 64, 128], [244, 35, 232], [70, 70, 70], [102, 102, 156],
                 [190, 153, 153], [153, 153, 153], [250, 170, 30], [
                     220, 220, 0
                 ], [107, 142, 35], [152, 251, 152], [70, 130, 180],
                 [220, 20, 60], [255, 0, 0], [0, 0, 142], [0, 0, 70],
                 [0, 60, 100], [0, 80, 100], [0, 0, 230], [119, 11, 32]],
                dtype=dtype)

        self.map = self.map / 255 if normalised else self.map

    def colourise(self, prediction):
        prediction = self.map[prediction]
        return prediction

max_stars_repo_name: Best-of-ACRV/pytorch-refinenet

max_issues_repo_name: Best-of-ACRV/pytorch-refinenet

max_forks_repo_name: Best-of-ACRV/pytorch-refinenet