<commit_before># -*- coding: utf-8 -*-

#  Copyright 2015 Matthias Fuchs
#
#  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.
#

from stromx import register, runtime
import time
import unittest

class DeadlockObserver(runtime.ConnectorObserver):
    interrupt = False

    def observe(self, connector, oldData, newData, thread):
        # consider outputs only
        if connector.type() == runtime.Connector.INPUT:
            return

        # consider only 'sendOutputData()' calls -> leads to a deadlock when
        # acquiring the write access below
        if newData.empty():
            return

        try:
            with runtime.ReadAccess(newData) as access:
                print access.get().get()
        except runtime.Exception as e:
            print e
            self.interrupt = True

class OutputObserver(runtime.ConnectorObserver):
    data = None

    def observe(self, connector, oldData, newData, thread):
        # consider outputs only
        if connector.type() == runtime.Connector.INPUT:
            return

        # consider only 'clearOutputData()' calls
        if oldData.empty():
            return

        with runtime.ReadAccess(oldData) as access:
            print access.get().get()
            self.data = access.get().get()

class ObserversTest(unittest.TestCase):
    def setUp(self):
        # set up an operator factory
        factory = runtime.Factory()
        register("libstromx_runtime.so", factory)
        register("libstromx_cvsupport.so", factory)
        register("libstromx_test.so", factory)

        # get the operator kernels
        counterKernel = factory.newOperator('runtime', 'Counter')
        dumpKernel = factory.newOperator('runtime', 'Dump')
        deadlockKernel = factory.newOperator('test', 'DeadlockOperator')

        # create the stream operator
        stream = runtime.Stream()
        stream.setDelay(100)
        counter = stream.addOperator(counterKernel)
        deadlock = stream.addOperator(deadlockKernel)
        dump = stream.addOperator(dumpKernel)

        stream.initializeOperator(counter)
        stream.initializeOperator(deadlock)
        stream.initializeOperator(dump)

        # obtain write access in deadlock operator
        deadlock.setParameter(3, runtime.Bool(True))

        stream.connect(counter, 0, deadlock, 0)
        stream.connect(deadlock, 1, dump, 0)

        thread = stream.addThread()
        thread.addInput(deadlock, 0)
        thread.addInput(dump, 0)

        self.stream = stream
        self.deadlock = deadlock

    def tearDown(self):
        self.stream = None

    def testInterrupt(self):
        observer = DeadlockObserver()
        self.deadlock.addObserver(observer)
        self.stream.start()

        time.sleep(0.3)

        self.stream.stop()
        self.stream.join()

        self.assertTrue(observer.interrupt)

    def testObserverOutput(self):
        observer = OutputObserver()
        self.deadlock.addObserver(observer)
        self.stream.start()

        time.sleep(0.3)

        self.stream.stop()
        self.stream.join()

        self.assertNotEqual(None, observer.data)

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

<commit_msg>Fix python observer test<commit_after># -*- coding: utf-8 -*-

#  Copyright 2015 Matthias Fuchs
#
#  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.
#

from stromx import register, runtime
import time
import unittest

class DeadlockObserver(runtime.ConnectorObserver):
    interrupt = False

    def observe(self, connector, oldData, newData, thread):
        # consider outputs only
        if connector.type() == runtime.Connector.INPUT:
            return

        # consider only 'sendOutputData()' calls -> leads to a deadlock when
        # acquiring the write access below
        if newData.empty():
            return

        try:
            with runtime.ReadAccess(newData) as access:
                print access.get().get()
        except runtime.Exception as e:
            print e
            self.interrupt = True

class OutputObserver(runtime.ConnectorObserver):
    data = None

    def observe(self, connector, oldData, newData, thread):
        # consider outputs only
        if connector.type() == runtime.Connector.INPUT:
            return

        # consider only 'clearOutputData()' calls
        if oldData.empty():
            return

        with runtime.ReadAccess(oldData) as access:
            print access.get().get()
            self.data = access.get().get()

class ObserversTest(unittest.TestCase):
    def setUp(self):
        # set up an operator factory
        factory = runtime.Factory()
        register("libstromx_runtime.so", factory)
        register("libstromx_cvsupport.so", factory)
        register("libstromx_test.so", factory)

        # get the operator kernels
        counterKernel = factory.newOperator('runtime', 'Counter')
        dumpKernel = factory.newOperator('runtime', 'Dump')
        deadlockKernel = factory.newOperator('test', 'DeadlockOperator')

        # create the stream operator
        stream = runtime.Stream()
        stream.setDelay(100)
        counter = stream.addOperator(counterKernel)
        deadlock = stream.addOperator(deadlockKernel)
        dump = stream.addOperator(dumpKernel)

        stream.initializeOperator(counter)
        stream.initializeOperator(deadlock)
        stream.initializeOperator(dump)

        # obtain write access in deadlock operator
        deadlock.setParameter(5, runtime.Bool(True))

        stream.connect(counter, 0, deadlock, 0)
        stream.connect(deadlock, 1, dump, 0)

        thread = stream.addThread()
        thread.addInput(deadlock, 0)
        thread.addInput(dump, 0)

        self.stream = stream
        self.deadlock = deadlock

    def tearDown(self):
        self.stream = None

    def testInterrupt(self):
        observer = DeadlockObserver()
        self.deadlock.addObserver(observer)
        self.stream.start()

        time.sleep(0.3)

        self.stream.stop()
        self.stream.join()

        self.assertTrue(observer.interrupt)

    def testObserverOutput(self):
        observer = OutputObserver()
        self.deadlock.addObserver(observer)
        self.stream.start()

        time.sleep(0.3)

        self.stream.stop()
        self.stream.join()

        self.assertNotEqual(None, observer.data)

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

<|endoftext|>


<commit_before># import for cache
import os
import tempfile
import shutil
import hashlib
import sysconfig
import sys
import weakref
import contextlib
import io
import torch.utils.cpp_extension
# import for just-in-time compilation
import distutils
import setuptools.command.build_ext
import setuptools
# triton
import triton.frameworks as fw
import triton.utils
import triton._C.libtriton as libtriton
import os
import time
import platform

@contextlib.contextmanager
def quiet():
    old_stdout, old_stderr = sys.stdout, sys.stderr
    sys.stdout, sys.stderr = io.StringIO(), io.StringIO()
    try:
        yield
    finally:
        sys.stdout, sys.stderr = old_stdout, old_stderr

def _build(src, path, name):
  ccdir = os.path.join(libtriton.__file__, os.path.pardir)
  ccdir = os.path.realpath(ccdir)
  # include / libraries
  include_dirs = [os.path.join(ccdir, 'include')]
  library_dirs = [ccdir]
  libraries = ['triton']
  # create extension module
  abi = fw.torch._C._GLIBCXX_USE_CXX11_ABI
  extra_compile_args = ['-fPIC', '-Wno-deprecated-declarations', '-D_GLIBCXX_USE_CXX11_ABI={abi}'.format(abi=abi)]
  extra_compile_args += ['-DTORCH_EXTENSION_NAME={}'.format(name)]
  extra_compile_args += ['-DTORCH_API_INCLUDE_EXTENSION_H']

  ext = torch.utils.cpp_extension.CUDAExtension(
      name = name,
      language = 'c++',
      sources = [src],
      include_dirs = include_dirs,
      library_dirs = library_dirs,
      libraries = libraries,
      extra_compile_args = extra_compile_args,
      depends = [os.path.realpath(libtriton.__file__)]
  )
  # build extension module
  args = ['build_ext']
  tmp = tempfile.mkdtemp()
  args.append('--build-temp=' + tmp)
  args.append('--build-lib=' + path)
  args.append('-q')
  args = dict(
      name = name,
      ext_modules = [ext],
      script_args = args,
  )
  with quiet():
    setuptools.setup(**args)
  shutil.rmtree(tmp)

def _cvt_to_def_str(obj):
  # bool
  if isinstance(obj, bool):
    return str(int(obj))
  # torch type
  if fw.has_torch():
    if isinstance(obj, fw.torch.dtype):
      return {fw.torch.int8: 'char',
              fw.torch.int16: 'short',
              fw.torch.int32: 'int',
              fw.torch.int64: 'long',
              fw.torch.float16: 'half',
              fw.torch.float32: 'float',
              fw.torch.float64: 'double'}[obj]
  else:
    assert False
  # default
  return str(obj)

def _encode(arg_types):
  codes = {
    libtriton.arg_type.int1:   'i1',
    libtriton.arg_type.int8:   'i8',
    libtriton.arg_type.int32:  'i32',
    libtriton.arg_type.int64:  'i64',
    libtriton.arg_type.half:   'f16',
    libtriton.arg_type.float:  'f32',
    libtriton.arg_type.double: 'f64',
    libtriton.arg_type.buffer: 'buf'
  }
  ret = '_'.join(map(codes.get, arg_types))
  return ret

def _make_framework_op(arg_types):
  name = _encode(arg_types)
  # path of .cpp and .so file
  home = os.path.expanduser('~')
  root = os.path.join(home, '.triton', 'torch', name)
  try:
    os.makedirs(root)
  except FileExistsError:
    pass
  suffix = sysconfig.get_config_var('EXT_SUFFIX')
  so = os.path.join(root, f'{name}.so')
  cpp = os.path.join(root, f'op.cpp')
  # handle cached .so file
  if os.path.exists(so) and os.stat(so).st_size > 0:
    tt_mtime = os.stat(os.path.realpath(libtriton.__file__)).st_mtime
    so_mtime = os.stat(so).st_mtime
    # can use cached if libtriton is older than the .so
    if tt_mtime < so_mtime:
      fw.torch.ops.load_library(so)
      return getattr(fw.torch.ops.triton, name)
  # create torch source code
  print('[TRITON] Compiling op...')
  baton = torch.utils.file_baton.FileBaton(os.path.join(root, 'lock'))
  if baton.try_acquire():
    try:
      src, _ = libtriton.make_torch_src(name, arg_types)
      with open(cpp, 'w+') as handle:
        handle.writelines(src)
      ccdir = os.path.join(libtriton.__file__, os.path.pardir)
      ccdir = os.path.realpath(ccdir)
      #include_dirs = [os.path.join(ccdir, 'include')]
      #library_dirs = [ccdir]
      #_build(cpp, root, 'op')
      #libraries = ['triton']
      machine = platform.machine()
      torch.utils.cpp_extension._write_ninja_file_and_build(
                    name=name,
                    sources=[cpp],
                    extra_cflags=['-std=gnu++11'] if machine == 'ppc64le' else [],
                    extra_cuda_cflags=[],
                    extra_ldflags=[f'-L{ccdir}', '-ltriton'],
                    extra_include_paths=[os.path.join(ccdir, 'include')],
                    build_directory=root,
                    verbose=False,
                    with_cuda=True)
    finally:
      baton.release()
  else:
    baton.wait()
  print('[TRITON] Done compiling...')
  fw.torch.ops.load_library(so)
  return getattr(fw.torch.ops.triton, name)

class kernel:

  def __init__(self, src, defines = dict(), num_warps = [2, 4, 8]):
    self.src = src
    # create constants
    self.cst = dict()
    # create triton op
    macros = []
    for k, v in defines.items():
      cvt = lambda x: _cvt_to_def_str(x)
      if(isinstance(v, list)):
        values = list(map(cvt, v))
      else:
        values = [cvt(v)]
      macros.append((k, values))
    opt = libtriton.options_space()
    opt.defines = macros
    opt.num_warps = num_warps
    self.op_id = libtriton.make_op_id()
    self.opt = opt
    self.registered = set()
    # create pytorch hook
    arg_types = libtriton.get_fn_signature(self.src, opt)
    self.fw_op = _make_framework_op(arg_types)

  def set_constant(self, name, value):
    libtriton.register_cst(self.op_id, name, value)

  def __call__(self, *args, **kwargs):
    for x in args:
      if isinstance(x, fw.torch.Tensor):
        device = x.device.index
        break
    # lazily register function for device
    if device not in self.registered:
      self.registered.add(device)
      libtriton.register_fn((self.op_id, device), self.src, self.opt, os.path.realpath(libtriton.__file__))
    # launch options
    bench = kwargs['bench']         if 'bench'     in kwargs else 0
    bench_id = libtriton.make_scalar_id() if bench > 0 else -1
    # launch grid
    if 'grid' not in kwargs:
      raise RuntimeError('Must provide grid for kernel launch')
    grid = kwargs['grid']
    libtriton.register_grid((self.op_id, device), grid)
    # launch
    self.fw_op(self.op_id, device, bench, bench_id, *args)
    if bench > 0:
      return libtriton.retrieve_scalar(bench_id)<commit_msg>[TRITON][PYTHON] Reverted back to distutils<commit_after># import for cache
import os
import tempfile
import shutil
import hashlib
import sysconfig
import sys
import weakref
import contextlib
import io
import torch.utils.cpp_extension
# import for just-in-time compilation
import distutils
import setuptools.command.build_ext
import setuptools
# triton
import triton.frameworks as fw
import triton.utils
import triton._C.libtriton as libtriton
import os
import time
import platform

@contextlib.contextmanager
def quiet():
    old_stdout, old_stderr = sys.stdout, sys.stderr
    sys.stdout, sys.stderr = io.StringIO(), io.StringIO()
    try:
        yield
    finally:
        sys.stdout, sys.stderr = old_stdout, old_stderr

def _build(src, path, name):
  ccdir = os.path.join(libtriton.__file__, os.path.pardir)
  ccdir = os.path.realpath(ccdir)
  # include / libraries
  include_dirs = [os.path.join(ccdir, 'include')]
  library_dirs = [ccdir]
  libraries = ['triton']
  # create extension module
  abi = fw.torch._C._GLIBCXX_USE_CXX11_ABI
  extra_compile_args = ['-fPIC', '-Wno-deprecated-declarations', f'-D_GLIBCXX_USE_CXX11_ABI={str(int(abi))}']
  extra_compile_args += ['-DTORCH_EXTENSION_NAME={}'.format(name)]
  extra_compile_args += ['-DTORCH_API_INCLUDE_EXTENSION_H']

  ext = torch.utils.cpp_extension.CUDAExtension(
      name = name,
      language = 'c++',
      sources = [src],
      include_dirs = include_dirs,
      library_dirs = library_dirs,
      libraries = libraries,
      extra_compile_args = extra_compile_args,
      depends = [os.path.realpath(libtriton.__file__)]
  )
  # build extension module
  args = ['build_ext']
  tmp = tempfile.mkdtemp()
  args.append('--build-temp=' + tmp)
  args.append('--build-lib=' + path)
  args.append('-q')
  args = dict(
      name = name,
      ext_modules = [ext],
      script_args = args,
  )
  with quiet():
    setuptools.setup(**args)
  shutil.rmtree(tmp)

def _cvt_to_def_str(obj):
  # bool
  if isinstance(obj, bool):
    return str(int(obj))
  # torch type
  if fw.has_torch():
    if isinstance(obj, fw.torch.dtype):
      return {fw.torch.int8: 'char',
              fw.torch.int16: 'short',
              fw.torch.int32: 'int',
              fw.torch.int64: 'long',
              fw.torch.float16: 'half',
              fw.torch.float32: 'float',
              fw.torch.float64: 'double'}[obj]
  else:
    assert False
  # default
  return str(obj)

def _encode(arg_types):
  codes = {
    libtriton.arg_type.int1:   'i1',
    libtriton.arg_type.int8:   'i8',
    libtriton.arg_type.int32:  'i32',
    libtriton.arg_type.int64:  'i64',
    libtriton.arg_type.half:   'f16',
    libtriton.arg_type.float:  'f32',
    libtriton.arg_type.double: 'f64',
    libtriton.arg_type.buffer: 'buf'
  }
  ret = '_'.join(map(codes.get, arg_types))
  return ret

def _make_framework_op(arg_types):
  name = _encode(arg_types)
  # path of .cpp and .so file
  home = os.path.expanduser('~')
  root = os.path.join(home, '.triton', 'torch', name)
  try:
    os.makedirs(root)
  except FileExistsError:
    pass
  suffix = sysconfig.get_config_var('EXT_SUFFIX')
  so = os.path.join(root, f'op{suffix}')
  cpp = os.path.join(root, f'op.cpp')
  # handle cached .so file
  if os.path.exists(so) and os.stat(so).st_size > 0:
    tt_mtime = os.stat(os.path.realpath(libtriton.__file__)).st_mtime
    so_mtime = os.stat(so).st_mtime
    # can use cached if libtriton is older than the .so
    if tt_mtime < so_mtime:
      fw.torch.ops.load_library(so)
      return getattr(fw.torch.ops.triton, name)
  # create torch source code
  print('[TRITON] Compiling op...')
  baton = torch.utils.file_baton.FileBaton(os.path.join(root, 'lock'))
  if baton.try_acquire():
    try:
      src, _ = libtriton.make_torch_src(name, arg_types)
      with open(cpp, 'w+') as handle:
        handle.writelines(src)
      ccdir = os.path.join(libtriton.__file__, os.path.pardir)
      ccdir = os.path.realpath(ccdir)
      _build(cpp, root, 'op')
    finally:
      baton.release()
  else:
    baton.wait()
  print('[TRITON] Done compiling...')
  fw.torch.ops.load_library(so)
  return getattr(fw.torch.ops.triton, name)

class kernel:

  def __init__(self, src, defines = dict(), num_warps = [2, 4, 8]):
    self.src = src
    # create constants
    self.cst = dict()
    # create triton op
    macros = []
    for k, v in defines.items():
      cvt = lambda x: _cvt_to_def_str(x)
      if(isinstance(v, list)):
        values = list(map(cvt, v))
      else:
        values = [cvt(v)]
      macros.append((k, values))
    opt = libtriton.options_space()
    opt.defines = macros
    opt.num_warps = num_warps
    self.op_id = libtriton.make_op_id()
    self.opt = opt
    self.registered = set()
    # create pytorch hook
    arg_types = libtriton.get_fn_signature(self.src, opt)
    self.fw_op = _make_framework_op(arg_types)

  def set_constant(self, name, value):
    libtriton.register_cst(self.op_id, name, value)

  def __call__(self, *args, **kwargs):
    for x in args:
      if isinstance(x, fw.torch.Tensor):
        device = x.device.index
        break
    # lazily register function for device
    if device not in self.registered:
      self.registered.add(device)
      libtriton.register_fn((self.op_id, device), self.src, self.opt, os.path.realpath(libtriton.__file__))
    # launch options
    bench = kwargs['bench']         if 'bench'     in kwargs else 0
    bench_id = libtriton.make_scalar_id() if bench > 0 else -1
    # launch grid
    if 'grid' not in kwargs:
      raise RuntimeError('Must provide grid for kernel launch')
    grid = kwargs['grid']
    libtriton.register_grid((self.op_id, device), grid)
    # launch
    self.fw_op(self.op_id, device, bench, bench_id, *args)
    if bench > 0:
      return libtriton.retrieve_scalar(bench_id)<|endoftext|>


<commit_before># Copyright (C) 2015, Carlo de Franchis <carlo.de-franchis@cmla.ens-cachan.fr>
# Copyright (C) 2015, Gabriele Facciolo <facciolo@cmla.ens-cachan.fr>
# Copyright (C) 2015, Enric Meinhardt <enric.meinhardt@cmla.ens-cachan.fr>
# Copyright (C) 2015, Julien Michel <julien.michel@cnes.fr>

import numpy as np
import os

import piio
import common
import estimation
import rpc_model
import rpc_utils
import pointing_accuracy

def plot_line(im, x1, y1, x2, y2, colour):
    """
    Plots a line on a rgb image stored as a numpy array

    Args:
        im: 3D numpy array containing the image values. It may be stored as
            uint8 or float32.
        x1, y1, x2, y2: integer coordinates of the line endpoints
        colour: list of length 3 giving the colour used for the plotted line
            (ie [r, g, b])

    Returns:
        a copy of the input numpy array, with the plotted lines on it. It means
        that the intensities of pixels located on the plotted line are changed.
    """
    # colour points of the line pixel by pixel. Loop over x or y, depending on
    # the biggest dimension.
    if np.abs(x2 - x1) >= np.abs(y2 - y1):
        n = np.abs(x2 - x1)
        for i in range(int(n+1)):
            x = x1 + i * (x2 - x1) / n
            y = np.round(y1 + i * (y2 - y1) / n)
            im[y, x] = colour
    else:
        n = np.abs(y2 - y1)
        for i in range(int(n+1)):
            y = y1 + i * (y2 - y1) / n
            x = np.round(x1 + i * (x2 - x1) / n)
            im[y, x] = colour

    return im

def plot_matches(im1, im2, matches):
    """
    Displays two images side by side with matches highlighted

    Args:
        im1, im2: paths to the two input images
        matches: 2D numpy array of size 4xN containing a list of matches (a
            list of pairs of points, each pair being represented by x1, y1, x2,
            y2)

    Returns:
        path to the resulting image, to be displayed
    """
    # load images
    img1 = piio.read(im1).astype(np.uint8)
    img2 = piio.read(im2).astype(np.uint8)

    # if images have more than 3 channels, keep only the first 3
    if img1.shape[2] > 3:
        img1 = img1[:, :, 0:3]
    if img2.shape[2] > 3:
        img2 = img2[:, :, 0:3]

    # build the output image
    h1, w1 = img1.shape[:2]
    h2, w2 = img2.shape[:2]
    w = w1 + w2
    h = max(h1, h2)
    out = np.zeros((h, w, 3), np.uint8)
    out[:h1, :w1] = img1
    out[:h2, w1:w] = img2

    # define colors, according to min/max intensity values
    out_min = min(np.nanmin(img1), np.nanmin(img2))
    out_max = max(np.nanmax(img1), np.nanmax(img2))
    green = [out_min, out_max, out_min]
    blue = [out_min, out_min, out_max]

    # plot the matches
    for i in range(len(matches)):
        x1 = matches[i, 0]
        y1 = matches[i, 1]
        x2 = matches[i, 2] + w1
        y2 = matches[i, 3]
        # convert endpoints to int (nn interpolation)
        x1, y1, x2, y2 = np.round([x1, y1, x2, y2])
        plot_line(out, x1, y1, x2, y2, blue)
        out[y1, x1] = green
        out[y2, x2] = green

    # save the output image, and return its path
    outfile = common.tmpfile('.png')
    piio.write(outfile, out)
    return outfile

def plot_matches_pleiades(im1, im2, rpc1, rpc2, matches, x=None, y=None, w=None, h=None,
                 outfile=None):
    """
    Plot matches on Pleiades images

    Args:
        im1, im2: paths to full Pleiades images
        rpc1, rpc2: paths to xml files containing the rpc coefficients
        matches: 2D numpy array of size 4xN containing a list of matches (a
            list of pairs of points, each pair being represented by x1, y1, x2,
            y2). The coordinates are given in the frame of the full images.
        x, y, w, h (optional, default is None): ROI in the reference image
        outfile (optional, default is None): path to the output file. If None,
            the file image is displayed using the pvflip viewer

    Returns:
        path to the displayed output
    """
    # if no matches, no plot
    if not matches.size:
        print "visualisation.plot_matches: nothing to plot"
        return

    # read rpcs
    r1 = rpc_model.RPCModel(rpc1)
    r2 = rpc_model.RPCModel(rpc2)

    # determine regions to crop in im1 and im2
    if x is not None:
        x1 = x
    else:
        x1 = np.min(matches[:, 0])

    if y is not None:
        y1 = y
    else:
        y1 = np.min(matches[:, 1])

    if w is not None:
        w1 = w
    else:
        w1 = np.max(matches[:, 0]) - x1

    if h is not None:
        h1 = h
    else:
        h1 = np.max(matches[:, 1]) - y1

    x2, y2, w2, h2 = rpc_utils.corresponding_roi(r1, r2, x1, y1, w1, h1)
    # x2 = np.min(matches[:, 2])
    # w2 = np.max(matches[:, 2]) - x2
    # y2 = np.min(matches[:, 3])
    # h2 = np.max(matches[:, 3]) - y2

    # # add 20 pixels offset and round. The image_crop_tif function will round
    # # off the coordinates before it does the crops.
    # x1 -= 20; x1 = np.round(x1)
    # y1 -= 20; y1 = np.round(y1)
    # x2 -= 20; x2 = np.round(x2)
    # y2 -= 20; y2 = np.round(y2)
    # w1 += 40; w1 = np.round(w1)
    # h1 += 40; h1 = np.round(h1)
    # w2 += 40; w2 = np.round(w2)
    # h2 += 40; h2 = np.round(h2)

    # do the crops
    crop1 = common.image_qauto(common.image_crop_tif(im1, x1, y1, w1, h1))
    crop2 = common.image_qauto(common.image_crop_tif(im2, x2, y2, w2, h2))

    # compute matches coordinates in the cropped images
    pts1 = matches[:, 0:2] - [x1, y1]
    pts2 = matches[:, 2:4] - [x2, y2]

    # plot the matches on the two crops
    to_display = plot_matches(crop1, crop2, np.hstack((pts1, pts2)))
    if outfile is None:
        os.system('v %s &' % (to_display))
    else:
        common.run('cp %s %s' % (to_display, outfile))

    return

def plot_vectors(p, v, x, y, w, h, f=1, out_file=None):
    """
    Plots vectors on an image, using gnuplot

    Args:
        p: points (origins of vectors),represented as a numpy Nx2 array
        v: vectors, represented as a numpy Nx2 array
        x, y, w, h: rectangular ROI
        f: (optional, default is 1) exageration factor
        out_file: (optional, default is None) path to the output file

    Returns:
        nothing, but opens a display or write a png file
    """
    tmp = common.tmpfile('.txt')
    data = np.hstack((p, v))
    np.savetxt(tmp, data, fmt='%6f')
    gp_string = 'set term png size %d,%d;unset key;unset tics;plot [%d:%d] [%d:%d] "%s" u($1):($2):(%d*$3):(%d*$4) w vectors head filled' % (w, h, x, x+w, y, y+h, tmp, f, f)

    if out_file is None:
        out_file = common.tmpfile('.png')

    common.run("gnuplot -p -e '%s' > %s" % (gp_string, out_file))
    print out_file

    if out_file is None:
        os.system("v %s &" % out_file)

def plot_pointing_error_tile(im1, im2, rpc1, rpc2, x, y, w, h,
        matches_sift=None, f=100, out_files_pattern=None):
    """
    Args:
        im1, im2: path to full images
        rpc1, rcp2: path to associated rpc xml files
        x, y, w, h: four integers defining the rectangular tile in the reference
            image. (x, y) is the top-left corner, and (w, h) are the dimensions
            of the tile.
        f (optional, default is 100): exageration factor for the error vectors
        out_files_pattern (optional, default is None): pattern used to name the
            two output files (plots of the pointing error)

    Returns:
        nothing, but opens a display
    """
    # read rpcs
    r1 = rpc_model.RPCModel(rpc1)
    r2 = rpc_model.RPCModel(rpc2)

    # compute sift matches
    if not matches_sift:
        matches_sift = sift.matches_on_rpc_roi(im1, im2, r1, r2, x, y, w, h)

    # compute rpc matches
    matches_rpc = rpc_utils.matches_from_rpc(r1, r2, x, y, w, h, 5)

    # estimate affine fundamental matrix
    F = estimation.affine_fundamental_matrix(matches_rpc)

    # compute error vectors
    e = pointing_accuracy.error_vectors(matches_sift, F, 'ref')

    A = pointing_accuracy.local_translation(r1, r2, x, y, w, h, matches_sift)
    p = matches_sift[:, 0:2]
    q = matches_sift[:, 2:4]
    qq = common.points_apply_homography(A, q)
    ee = pointing_accuracy.error_vectors(np.hstack((p, qq)), F, 'ref')
    print pointing_accuracy.evaluation_from_estimated_F(im1, im2, r1, r2, x, y, w, h, None, matches_sift)
    print pointing_accuracy.evaluation_from_estimated_F(im1, im2, r1, r2, x, y, w, h, A, matches_sift)

    # plot the vectors: they go from the point x to the line (F.T)x'
    plot_vectors(p, -e, x, y, w, h, f, out_file='%s_before.png' % out_files_pattern)
    plot_vectors(p, -ee, x, y, w, h, f, out_file='%s_after.png' % out_files_pattern)
<commit_msg>fix name clash bug in visualisation module<commit_after># Copyright (C) 2015, Carlo de Franchis <carlo.de-franchis@cmla.ens-cachan.fr>
# Copyright (C) 2015, Gabriele Facciolo <facciolo@cmla.ens-cachan.fr>
# Copyright (C) 2015, Enric Meinhardt <enric.meinhardt@cmla.ens-cachan.fr>
# Copyright (C) 2015, Julien Michel <julien.michel@cnes.fr>

import numpy as np
import os

import piio
import common
import estimation
import rpc_model
import rpc_utils
import pointing_accuracy

def plot_line(im, x1, y1, x2, y2, colour):
    """
    Plots a line on a rgb image stored as a numpy array

    Args:
        im: 3D numpy array containing the image values. It may be stored as
            uint8 or float32.
        x1, y1, x2, y2: integer coordinates of the line endpoints
        colour: list of length 3 giving the colour used for the plotted line
            (ie [r, g, b])

    Returns:
        a copy of the input numpy array, with the plotted lines on it. It means
        that the intensities of pixels located on the plotted line are changed.
    """
    # colour points of the line pixel by pixel. Loop over x or y, depending on
    # the biggest dimension.
    if np.abs(x2 - x1) >= np.abs(y2 - y1):
        n = np.abs(x2 - x1)
        for i in range(int(n+1)):
            x = x1 + i * (x2 - x1) / n
            y = np.round(y1 + i * (y2 - y1) / n)
            im[y, x] = colour
    else:
        n = np.abs(y2 - y1)
        for i in range(int(n+1)):
            y = y1 + i * (y2 - y1) / n
            x = np.round(x1 + i * (x2 - x1) / n)
            im[y, x] = colour

    return im

def plot_matches_low_level(im1, im2, matches):
    """
    Displays two images side by side with matches highlighted

    Args:
        im1, im2: paths to the two input images
        matches: 2D numpy array of size 4xN containing a list of matches (a
            list of pairs of points, each pair being represented by x1, y1, x2,
            y2)

    Returns:
        path to the resulting image, to be displayed
    """
    # load images
    img1 = piio.read(im1).astype(np.uint8)
    img2 = piio.read(im2).astype(np.uint8)

    # if images have more than 3 channels, keep only the first 3
    if img1.shape[2] > 3:
        img1 = img1[:, :, 0:3]
    if img2.shape[2] > 3:
        img2 = img2[:, :, 0:3]

    # build the output image
    h1, w1 = img1.shape[:2]
    h2, w2 = img2.shape[:2]
    w = w1 + w2
    h = max(h1, h2)
    out = np.zeros((h, w, 3), np.uint8)
    out[:h1, :w1] = img1
    out[:h2, w1:w] = img2

    # define colors, according to min/max intensity values
    out_min = min(np.nanmin(img1), np.nanmin(img2))
    out_max = max(np.nanmax(img1), np.nanmax(img2))
    green = [out_min, out_max, out_min]
    blue = [out_min, out_min, out_max]

    # plot the matches
    for i in range(len(matches)):
        x1 = matches[i, 0]
        y1 = matches[i, 1]
        x2 = matches[i, 2] + w1
        y2 = matches[i, 3]
        # convert endpoints to int (nn interpolation)
        x1, y1, x2, y2 = np.round([x1, y1, x2, y2])
        plot_line(out, x1, y1, x2, y2, blue)
        out[y1, x1] = green
        out[y2, x2] = green

    # save the output image, and return its path
    outfile = common.tmpfile('.png')
    piio.write(outfile, out)
    return outfile

def plot_matches(im1, im2, rpc1, rpc2, matches, x=None, y=None, w=None, h=None,
                 outfile=None):
    """
    Plot matches on Pleiades images

    Args:
        im1, im2: paths to full Pleiades images
        rpc1, rpc2: paths to xml files containing the rpc coefficients
        matches: 2D numpy array of size 4xN containing a list of matches (a
            list of pairs of points, each pair being represented by x1, y1, x2,
            y2). The coordinates are given in the frame of the full images.
        x, y, w, h (optional, default is None): ROI in the reference image
        outfile (optional, default is None): path to the output file. If None,
            the file image is displayed using the pvflip viewer

    Returns:
        path to the displayed output
    """
    # if no matches, no plot
    if not matches.size:
        print "visualisation.plot_matches: nothing to plot"
        return

    # read rpcs
    r1 = rpc_model.RPCModel(rpc1)
    r2 = rpc_model.RPCModel(rpc2)

    # determine regions to crop in im1 and im2
    if x is not None:
        x1 = x
    else:
        x1 = np.min(matches[:, 0])

    if y is not None:
        y1 = y
    else:
        y1 = np.min(matches[:, 1])

    if w is not None:
        w1 = w
    else:
        w1 = np.max(matches[:, 0]) - x1

    if h is not None:
        h1 = h
    else:
        h1 = np.max(matches[:, 1]) - y1

    x2, y2, w2, h2 = rpc_utils.corresponding_roi(r1, r2, x1, y1, w1, h1)
    # x2 = np.min(matches[:, 2])
    # w2 = np.max(matches[:, 2]) - x2
    # y2 = np.min(matches[:, 3])
    # h2 = np.max(matches[:, 3]) - y2

    # # add 20 pixels offset and round. The image_crop_tif function will round
    # # off the coordinates before it does the crops.
    # x1 -= 20; x1 = np.round(x1)
    # y1 -= 20; y1 = np.round(y1)
    # x2 -= 20; x2 = np.round(x2)
    # y2 -= 20; y2 = np.round(y2)
    # w1 += 40; w1 = np.round(w1)
    # h1 += 40; h1 = np.round(h1)
    # w2 += 40; w2 = np.round(w2)
    # h2 += 40; h2 = np.round(h2)

    # do the crops
    crop1 = common.image_qauto(common.image_crop_tif(im1, x1, y1, w1, h1))
    crop2 = common.image_qauto(common.image_crop_tif(im2, x2, y2, w2, h2))

    # compute matches coordinates in the cropped images
    pts1 = matches[:, 0:2] - [x1, y1]
    pts2 = matches[:, 2:4] - [x2, y2]

    # plot the matches on the two crops
    to_display = plot_matches_low_level(crop1, crop2, np.hstack((pts1, pts2)))
    if outfile is None:
        os.system('v %s &' % (to_display))
    else:
        common.run('cp %s %s' % (to_display, outfile))

    return

def plot_vectors(p, v, x, y, w, h, f=1, out_file=None):
    """
    Plots vectors on an image, using gnuplot

    Args:
        p: points (origins of vectors),represented as a numpy Nx2 array
        v: vectors, represented as a numpy Nx2 array
        x, y, w, h: rectangular ROI
        f: (optional, default is 1) exageration factor
        out_file: (optional, default is None) path to the output file

    Returns:
        nothing, but opens a display or write a png file
    """
    tmp = common.tmpfile('.txt')
    data = np.hstack((p, v))
    np.savetxt(tmp, data, fmt='%6f')
    gp_string = 'set term png size %d,%d;unset key;unset tics;plot [%d:%d] [%d:%d] "%s" u($1):($2):(%d*$3):(%d*$4) w vectors head filled' % (w, h, x, x+w, y, y+h, tmp, f, f)

    if out_file is None:
        out_file = common.tmpfile('.png')

    common.run("gnuplot -p -e '%s' > %s" % (gp_string, out_file))
    print out_file

    if out_file is None:
        os.system("v %s &" % out_file)

def plot_pointing_error_tile(im1, im2, rpc1, rpc2, x, y, w, h,
        matches_sift=None, f=100, out_files_pattern=None):
    """
    Args:
        im1, im2: path to full images
        rpc1, rcp2: path to associated rpc xml files
        x, y, w, h: four integers defining the rectangular tile in the reference
            image. (x, y) is the top-left corner, and (w, h) are the dimensions
            of the tile.
        f (optional, default is 100): exageration factor for the error vectors
        out_files_pattern (optional, default is None): pattern used to name the
            two output files (plots of the pointing error)

    Returns:
        nothing, but opens a display
    """
    # read rpcs
    r1 = rpc_model.RPCModel(rpc1)
    r2 = rpc_model.RPCModel(rpc2)

    # compute sift matches
    if not matches_sift:
        matches_sift = sift.matches_on_rpc_roi(im1, im2, r1, r2, x, y, w, h)

    # compute rpc matches
    matches_rpc = rpc_utils.matches_from_rpc(r1, r2, x, y, w, h, 5)

    # estimate affine fundamental matrix
    F = estimation.affine_fundamental_matrix(matches_rpc)

    # compute error vectors
    e = pointing_accuracy.error_vectors(matches_sift, F, 'ref')

    A = pointing_accuracy.local_translation(r1, r2, x, y, w, h, matches_sift)
    p = matches_sift[:, 0:2]
    q = matches_sift[:, 2:4]
    qq = common.points_apply_homography(A, q)
    ee = pointing_accuracy.error_vectors(np.hstack((p, qq)), F, 'ref')
    print pointing_accuracy.evaluation_from_estimated_F(im1, im2, r1, r2, x, y, w, h, None, matches_sift)
    print pointing_accuracy.evaluation_from_estimated_F(im1, im2, r1, r2, x, y, w, h, A, matches_sift)

    # plot the vectors: they go from the point x to the line (F.T)x'
    plot_vectors(p, -e, x, y, w, h, f, out_file='%s_before.png' % out_files_pattern)
    plot_vectors(p, -ee, x, y, w, h, f, out_file='%s_after.png' % out_files_pattern)
<|endoftext|>


<commit_before>
from operator import attrgetter
import logging
import blocks
logger = logging.getLogger(__name__)

class HashChainTask(object):
    """
    - get hashes chain until we see a known block hash
    """

    NUM_HASHES_PER_REQUEST = 2000
    GENESIS_HASH = blocks.genesis().hash

    def __init__(self, chain_manager, peer, block_hash):
        self.chain_manager = chain_manager
        self.peer = peer
        self.hash_chain = [] # [youngest, ..., oldest]
        self.request(block_hash)

    def request(self, block_hash):
        logger.debug('%r requesting block_hashes starting from %r', self.peer, block_hash.encode('hex'))
        self.peer.send_GetBlockHashes(block_hash, self.NUM_HASHES_PER_REQUEST)

    def received_block_hashes(self, block_hashes):
        logger.warn('HashChainTask.received_block_hashes %d', len(block_hashes))
        if self.GENESIS_HASH in block_hashes:
            logger.warn('%r has different chain starting from genesis', self.peer)
        for bh in block_hashes:
            if bh in self.chain_manager or bh == self.GENESIS_HASH:
                logger.debug('%r matching block hash found %r', self.peer, bh.encode('hex'))
                return list(reversed(self.hash_chain))
            self.hash_chain.append(bh)
            #logger.debug('hash_chain.append(%r) %d', bh.encode('hex'), len(self.hash_chain))
        self.request(bh)

class SynchronizationTask(object):
    """
    Created if we receive a unknown block w/o known parent. Possibly from a different branch.

    - get hashes chain until we see a known block hash
    - request missing blocks

    - once synced
        - rerequest blocks that lacked a reference before syncing
    """
    NUM_BLOCKS_PER_REQUEST = 200

    def __init__(self, chain_manager, peer, block_hash):
        self.chain_manager = chain_manager
        self.peer = peer
        self.hash_chain = [] # [oldest to youngest]
        logger.debug('%r syncing %r', self.peer, block_hash.encode('hex'))
        self.hash_chain_task = HashChainTask(self.chain_manager, self.peer, block_hash)

    def received_block_hashes(self, block_hashes):
        res = self.hash_chain_task.received_block_hashes(block_hashes)
        if res:
            self.hash_chain = res
            logger.debug('%r hash chain with %d hashes for missing blocks', self.peer, len(self.hash_chain))
            self.request_blocks()

    def received_blocks(self, transient_blocks):
        logger.debug('%r received %d of %d missing blocks', self.peer, len(transient_blocks), len(self.hash_chain))
        for tb in transient_blocks:
            if len(self.hash_chain) and self.hash_chain[0] == tb.hash:
                self.hash_chain.pop(0)
            else:
                logger.debug('%r received unexpected block %r', self.peer, tb)
                return False
        if self.hash_chain:
            # still blocks to fetch
            logger.debug('%r still missing %d blocks', self.peer, len(self.hash_chain))
            self.request_blocks()
        else: # done
            return True

    def request_blocks(self):
        logger.debug('%r requesting %d of %d missing blocks', self.peer, self.NUM_BLOCKS_PER_REQUEST, len(self.hash_chain))
        self.peer.send_GetBlocks(self.hash_chain[:self.NUM_BLOCKS_PER_REQUEST])

class Synchronizer(object):
    """"
    Cases:
        on "recv_Status": received unknown head_hash w/ sufficient difficulty
        on "recv_Blocks": received block w/o parent (new block mined, competing chain discovered)

    Naive Strategy:
        assert we see a block for which we have no parent
        assume that the sending peer knows the parent
        if we have not yet syncer for this unknown block:
            create new syncer
            sync direction genesis until we see known block_hash
            sync also (re)requests the block we missed, so it can be added on top of the synced chain
        else
            do nothing
            syncing (if finished) will be started with the next broadcasted block w/ missing parent
    """

    def __init__(self, chain_manager):
        self.chain_manager = chain_manager
        self.synchronization_tasks = {} # peer > syncer # syncer.unknown_hash as marker for task

    def stop_synchronization(self, peer):
        logger.debug('%r sync stopped', peer)
        if peer in self.synchronization_tasks:
            del self.synchronization_tasks[peer]

    def synchronize_unknown_block(self, peer, block_hash, force=False):
        "Case: block with unknown parent. Fetches unknown ancestors and this block"
        logger.debug('%r sync %r', peer, block_hash.encode('hex'))
        assert block_hash not in self.chain_manager
        if peer and (not peer in self.synchronization_tasks) or force:
            logger.debug('%r new sync task', peer)
            self.synchronization_tasks[peer] = SynchronizationTask(self.chain_manager, peer, block_hash)
        else:
            logger.debug('%r already has a synctask, sorry', peer)

    def synchronize_status(self, peer, block_hash, total_difficulty):
        "Case: unknown head with sufficient difficulty"
        logger.debug('%r status  with %r %d', peer,  block_hash.encode('hex'), total_difficulty)
        assert block_hash not in self.chain_manager
        # guesstimate the max difficulty difference possible for a sucessfully competing chain
        # worst case if skip it: we are on a stale chain until the other catched up
        # assume difficulty is constant
        num_blocks_behind = 7
        avg_uncles_per_block = 4
        max_diff = self.chain_manager.head.difficulty * num_blocks_behind * (1 + avg_uncles_per_block)
        if total_difficulty + max_diff > self.chain_manager.head.difficulty:
            logger.debug('%r sufficient difficulty, syncing', peer)
            self.synchronize_unknown_block(peer, block_hash)
        else:
            logger.debug('%r insufficient difficulty, not syncing', peer)

    def received_block_hashes(self, peer, block_hashes):
        if peer in self.synchronization_tasks:
            logger.debug("Synchronizer.received_block_hashes %d for: %r", len(block_hashes), peer)
            self.synchronization_tasks[peer].received_block_hashes(block_hashes)

    def received_blocks(self, peer, transient_blocks):
        if peer in self.synchronization_tasks:
            res = self.synchronization_tasks[peer].received_blocks(transient_blocks)
            if res is True:
                logger.debug("Synchronizer.received_blocks: chain w %r synced", peer)
                del self.synchronization_tasks[peer]
<commit_msg>ethereum protocol 31<commit_after>
from operator import attrgetter
import logging
import blocks
logger = logging.getLogger(__name__)

class HashChainTask(object):
    """
    - get hashes chain until we see a known block hash
    """

    NUM_HASHES_PER_REQUEST = 2000
    GENESIS_HASH = blocks.genesis().hash

    def __init__(self, chain_manager, peer, block_hash):
        self.chain_manager = chain_manager
        self.peer = peer
        self.hash_chain = [] # [youngest, ..., oldest]
        self.request(block_hash)

    def request(self, block_hash):
        logger.debug('%r requesting block_hashes starting from %r', self.peer, block_hash.encode('hex'))
        self.peer.send_GetBlockHashes(block_hash, self.NUM_HASHES_PER_REQUEST)

    def received_block_hashes(self, block_hashes):
        logger.warn('HashChainTask.received_block_hashes %d', len(block_hashes))
        if self.GENESIS_HASH in block_hashes:
            logger.warn('%r has different chain starting from genesis', self.peer)
        for bh in block_hashes:
            if bh in self.chain_manager or bh == self.GENESIS_HASH:
                logger.debug('%r matching block hash found %r', self.peer, bh.encode('hex'))
                return list(reversed(self.hash_chain))
            self.hash_chain.append(bh)
            #logger.debug('hash_chain.append(%r) %d', bh.encode('hex'), len(self.hash_chain))
        self.request(bh)

class SynchronizationTask(object):
    """
    Created if we receive a unknown block w/o known parent. Possibly from a different branch.

    - get hashes chain until we see a known block hash
    - request missing blocks

    - once synced
        - rerequest blocks that lacked a reference before syncing
    """
    NUM_BLOCKS_PER_REQUEST = 200

    def __init__(self, chain_manager, peer, block_hash):
        self.chain_manager = chain_manager
        self.peer = peer
        self.hash_chain = [] # [oldest to youngest]
        logger.debug('%r syncing %r', self.peer, block_hash.encode('hex'))
        self.hash_chain_task = HashChainTask(self.chain_manager, self.peer, block_hash)

    def received_block_hashes(self, block_hashes):
        res = self.hash_chain_task.received_block_hashes(block_hashes)
        if res:
            self.hash_chain = res
            logger.debug('%r hash chain with %d hashes for missing blocks', self.peer, len(self.hash_chain))
            self.request_blocks()

    def received_blocks(self, transient_blocks):
        logger.debug('%r received %d of %d missing blocks', self.peer, len(transient_blocks), len(self.hash_chain))
        for tb in transient_blocks:
            if len(self.hash_chain) and self.hash_chain[0] == tb.hash:
                self.hash_chain.pop(0)
            else:
                logger.debug('%r received unexpected block %r', self.peer, tb)
                return False
        if self.hash_chain:
            # still blocks to fetch
            logger.debug('%r still missing %d blocks', self.peer, len(self.hash_chain))
            self.request_blocks()
        else: # done
            return True

    def request_blocks(self):
        logger.debug('%r requesting %d of %d missing blocks', self.peer, self.NUM_BLOCKS_PER_REQUEST, len(self.hash_chain))
        self.peer.send_GetBlocks(self.hash_chain[:self.NUM_BLOCKS_PER_REQUEST])

class Synchronizer(object):
    """"
    Cases:
        on "recv_Status": received unknown head_hash w/ sufficient difficulty
        on "recv_Blocks": received block w/o parent (new block mined, competing chain discovered)

    Naive Strategy:
        assert we see a block for which we have no parent
        assume that the sending peer knows the parent
        if we have not yet syncer for this unknown block:
            create new syncer
            sync direction genesis until we see known block_hash
            sync also (re)requests the block we missed, so it can be added on top of the synced chain
        else
            do nothing
            syncing (if finished) will be started with the next broadcasted block w/ missing parent
    """

    def __init__(self, chain_manager):
        self.chain_manager = chain_manager
        self.synchronization_tasks = {} # peer > syncer # syncer.unknown_hash as marker for task

    def stop_synchronization(self, peer):
        logger.debug('%r sync stopped', peer)
        if peer in self.synchronization_tasks:
            del self.synchronization_tasks[peer]

    def synchronize_unknown_block(self, peer, block_hash, force=False):
        "Case: block with unknown parent. Fetches unknown ancestors and this block"
        logger.debug('%r sync %r', peer, block_hash.encode('hex'))
        assert block_hash not in self.chain_manager
        if peer and (not peer in self.synchronization_tasks) or force:
            logger.debug('%r new sync task', peer)
            self.synchronization_tasks[peer] = SynchronizationTask(self.chain_manager, peer, block_hash)
        else:
            logger.debug('%r already has a synctask, sorry', peer)

    def synchronize_status(self, peer, block_hash, total_difficulty):
        "Case: unknown head with sufficient difficulty"
        logger.debug('%r status  with %r %d', peer,  block_hash.encode('hex'), total_difficulty)
        if block_hash == blocks.genesis().hash:
            logger.debug('%r head == genesis, skipping', peer)
            return
        assert block_hash not in self.chain_manager

        # guesstimate the max difficulty difference possible for a sucessfully competing chain
        # worst case if skip it: we are on a stale chain until the other catched up
        # assume difficulty is constant
        num_blocks_behind = 7
        avg_uncles_per_block = 4
        max_diff = self.chain_manager.head.difficulty * num_blocks_behind * (1 + avg_uncles_per_block)
        if total_difficulty + max_diff > self.chain_manager.head.difficulty:
            logger.debug('%r sufficient difficulty, syncing', peer)
            self.synchronize_unknown_block(peer, block_hash)
        else:
            logger.debug('%r insufficient difficulty, not syncing', peer)

    def received_block_hashes(self, peer, block_hashes):
        if peer in self.synchronization_tasks:
            logger.debug("Synchronizer.received_block_hashes %d for: %r", len(block_hashes), peer)
            self.synchronization_tasks[peer].received_block_hashes(block_hashes)

    def received_blocks(self, peer, transient_blocks):
        if peer in self.synchronization_tasks:
            res = self.synchronization_tasks[peer].received_blocks(transient_blocks)
            if res is True:
                logger.debug("Synchronizer.received_blocks: chain w %r synced", peer)
                del self.synchronization_tasks[peer]
<|endoftext|>


<commit_before>#
# Sublime Text plugin to support Corona Editor
#
# Copyright (c) 2013 Corona Labs Inc. A mobile development software company. All rights reserved.
#
# MIT License - see https://raw.github.com/coronalabs/CoronaSDK-SublimeText/master/LICENSE

# Note: the leading underscore on the name of this module forces Sublime Text to load it first which
# prevents it being loaded more than once (once when it's first imported, e.g. by about.py, and then
# again when it comes up in the load order)

import sublime
import os
import re
import threading
import io
import subprocess
import sys

# Define "check_output" for Python <= 2.6 (ST2 uses this)
if "check_output" not in dir( subprocess ):
  def f(*popenargs, **kwargs):
    if 'stdout' in kwargs:
      raise ValueError('stdout argument not allowed, it will be overridden.')
    process = subprocess.Popen(stdout=subprocess.PIPE, *popenargs, **kwargs)
    output, unused_err = process.communicate()
    retcode = process.poll()
    if retcode:
      cmd = kwargs.get("args")
      if cmd is None:
        cmd = popenargs[0]
      raise subprocess.CalledProcessError(retcode, cmd)
    return output
  subprocess.check_output = f

SUBLIME_VERSION = "not set"
PLUGIN_PATH = "not set"
PACKAGE_NAME = "not set"
PACKAGE_DIR = "not set"
PACKAGE_USER_DIR = "not set"
ST_PACKAGE_PATH = "not set"
_corona_sdk_debug = False

# In Sublime Text 3 most APIs are unavailable until a module level function is called (fortunately
# sublime.version() is available so we can correctly fake things in Sublime Text 2; see about.py)
SUBLIME_VERSION = 3000 if sublime.version() == '' else int(sublime.version())

def GetSetting(key,default=None):
  # repeated calls to load_settings return same object without further disk reads
  debug("GetSetting: " + str(key) + " (default: " + str(default) + ")")
  settings = sublime.load_settings('Corona Editor.sublime-settings')
  debug("GetSetting: value from CE settings: " + str(settings.get(key, default)))
  # If we don't have a value for this preference in the Corona Editor settings, look in the view for a value
  # (this happens if the preference is set in the main Sublime Text preference file instead of the Corona Editor file)
  if not settings.get(key, default) and sublime and sublime.active_window() and sublime.active_window().active_view():
    debug("GetSetting: getting value from view: " + str(sublime.active_window().active_view()))
    settings = sublime.active_window().active_view().settings()

  value = settings.get(key, default)

  if not value:
    debug("GetSetting: no value for preference '" + str(key) + "' found (using default '"+ str(default) + "')")

  return value

def debug(*args):
  global _corona_sdk_debug
  if _corona_sdk_debug:
    print("Corona Editor: ", '\t'.join(map(str, args)))

InitializedEvent = threading.Event()

def Init():
  global SUBLIME_VERSION
  global PLUGIN_PATH
  global PACKAGE_NAME
  global PACKAGE_DIR
  global PACKAGE_USER_DIR
  global ST_PACKAGE_PATH
  global _corona_sdk_debug

  # Always look for "corona_sdk_debug" in 'Corona Editor.sublime-settings'
  settings = sublime.load_settings('Corona Editor.sublime-settings')
  _corona_sdk_debug = settings.get("corona_sdk_debug", False)

  print("Corona Editor: Init")
  debug("Python: " + str(sys.version))

  PLUGIN_PATH = os.path.dirname(os.path.realpath(__file__))
  if PLUGIN_PATH.lower().endswith('coronasdk-sublimetext'):
    if SUBLIME_VERSION < 3000:
      PLUGIN_PATH = os.path.expanduser('~/Library/Application Support/Sublime Text 2/Packages/Corona Editor')
    else:
      PLUGIN_PATH = os.path.expanduser('~/Library/Application Support/Sublime Text 3/Packages/Corona Editor')
  debug("PLUGIN_PATH: " + PLUGIN_PATH)

  PACKAGE_NAME = os.path.basename(PLUGIN_PATH) if SUBLIME_VERSION < 3000 else os.path.basename(PLUGIN_PATH).replace(".sublime-package", "")
  debug("PACKAGE_NAME: " + PACKAGE_NAME)

  # This is the actual package dir on ST2 and the "unpacked" package dir on ST3 ("~/Library/Application Support/Sublime Text 2/Packages")
  PACKAGE_DIR = os.path.join(sublime.packages_path(), PACKAGE_NAME)
  debug("PACKAGE_DIR: " + PACKAGE_DIR)

  # This is the user dir on ST
  PACKAGE_USER_DIR = os.path.join(sublime.packages_path(), 'User', PACKAGE_NAME)
  debug("PACKAGE_USER_DIR: " + PACKAGE_USER_DIR)

  # This is the faux path used by various Sublime Text functions ("Packages/CoronaSDK/")
  ST_PACKAGE_PATH = "Packages/" + PACKAGE_NAME + "/"
  debug("ST_PACKAGE_PATH: " + ST_PACKAGE_PATH)

  # This allows other threads to wait until this code has run
  InitializedEvent.set()

# This is run by ST3 after the plugins have loaded (and the ST API is ready for calls)
def plugin_loaded():
  Init()

# This fakes the above functionality on ST2
if SUBLIME_VERSION < 3000:
  Init()

# Return the path to the Simulator for the current project and platform
# First we look in any build.settings file in the project, then we look
# for a user preference, then we pick a platform appropriate default
def GetSimulatorCmd(mainlua=None, debug=False):
  platform = sublime.platform()
  arch = sublime.arch()
  view = sublime.active_window().active_view()

  simulator_path = ""
  simulator_flags = []
  simulator_version = GetSetting("corona_sdk_version", None)

  if mainlua is not None:
    simulator_path = GetSimulatorPathFromBuildSettings(mainlua)
    if simulator_path is None:
      simulator_path = GetSetting("corona_sdk_simulator_path", None)

  if platform == "osx":
    if simulator_path is None:
      simulator_path = GetLatestDailyBuildSimulatorPath()
    if simulator_path is None:
      simulator_path = "/Applications/CoronaSDK/Corona Simulator.app"
    if simulator_path.endswith(".app"):
      simulator_path += "/Contents/MacOS/Corona Simulator"
    simulator_flags = ["-singleton", "1", "-no-console", "1"]
    if debug:
      simulator_flags.append("-debug")
      simulator_flags.append("1")
      simulator_flags.append("-project")
    if os.path.isfile(simulator_path) and os.access(simulator_path, os.X_OK):
      simulator_version = str(subprocess.check_output(['defaults', 'read', os.path.join(os.path.dirname(simulator_path), '../Info.plist'), 'CFBundleVersion'], stderr=subprocess.STDOUT).decode('utf8').strip())
  elif platform == 'windows':
    if simulator_path is None:
      if arch == "x64":
        simulator_path = "C:\\Program Files (x86)\\Corona Labs\\Corona SDK\\Corona Simulator.exe"
      else:
        simulator_path = "C:\\Program Files\\Corona Labs\\Corona SDK\\Corona Simulator.exe"
    simulator_flags = ["/singleton", "/no-console"]
    if debug:
      simulator_flags.append("/debug")

  # Can we find an executable file at the path
  if not os.path.isfile(simulator_path) or not os.access(simulator_path, os.X_OK):
    sublime.error_message("Cannot find executable Corona Simulator at path '{0}'\n\nYou can set the user preference 'corona_sdk_simulator_path' to the location of the Simulator.".format(simulator_path))
    return None, None, None

  return simulator_path, simulator_flags, simulator_version

# Find the highest numbered CoronaSDK directory in /Applications
def GetLatestDailyBuildSimulatorPath():
  dirs = os.listdir("/Applications")

  # Find the highest numbered Daily Build
  maxBuildNum = 0
  buildNum = 0
  latestSDK = ""
  for entry in dirs:
    if entry.startswith("CoronaSDK-"):
      result = re.findall(r'\d+', entry)
      if len(result) > 0:
        buildNum = int(re.findall(r'\d+', entry)[0])
      if buildNum > maxBuildNum:
        maxBuildNum = buildNum
        latestSDK = "/Applications/" + entry

  if len(latestSDK) > 0:
    return  latestSDK + "/Corona Simulator.app"
  else:
    return None

# Given a path to a main.lua file, see if there's a "corona_sdk_simulator_path" setting in
# the corresponding build.settings file
def GetSimulatorPathFromBuildSettings(mainlua):
  simulator_path = None
  project_path = os.path.dirname(mainlua)
  build_settings = os.path.join(project_path, "build.settings")
  bs_contents = None
  if os.path.isfile(build_settings):
    try:
      try: # see if we're in P3
        with open(build_settings, "r", encoding="utf-8") as bs_fd:
          bs_contents = bs_fd.read()
      except: # nope, try a P2 idiom
        with io.open(build_settings, "r", encoding="utf-8") as bs_fd:
          bs_contents = bs_fd.read()
    except IOError:
      pass  # we don't care if the file doesn't exist

  # debug("bs_contents: ", str(bs_contents))
  if bs_contents is not None:
    # Remove comments
    bs_contents = re.sub(r'--.*', '', bs_contents)
    # Note we can't use a Python r'' string here because we then can't escape the single quotes
    bs_matches = re.findall('corona_sdk_simulator_path\s*=\s*["\'](.*)["\']', bs_contents)
    if bs_matches is not None and len(bs_matches) > 0:
      # Last one wins
      simulator_path = bs_matches[-1]
      print("Corona Editor: corona_sdk_simulator_path set from "+str(build_settings))

  return simulator_path

# Given an existing file path or directory, find the likely "main.lua" for this project
def ResolveMainLua(path):
  # debug("ResolveMainLua: path: "+str(path))
  path = os.path.abspath(path)
  if not os.path.isdir(path):
    path = os.path.dirname(path)

  mainlua = os.path.join(path, "main.lua")
  if mainlua == "/main.lua" or mainlua == "\\main.lua":
    return None
  elif os.path.isfile(mainlua):
    return mainlua
  else:
    return ResolveMainLua(os.path.join(path, ".."))

<commit_msg>Add option to show Corona Simulator Console when running a project<commit_after>#
# Sublime Text plugin to support Corona Editor
#
# Copyright (c) 2013 Corona Labs Inc. A mobile development software company. All rights reserved.
#
# MIT License - see https://raw.github.com/coronalabs/CoronaSDK-SublimeText/master/LICENSE

# Note: the leading underscore on the name of this module forces Sublime Text to load it first which
# prevents it being loaded more than once (once when it's first imported, e.g. by about.py, and then
# again when it comes up in the load order)

import sublime
import os
import re
import threading
import io
import subprocess
import sys

# Define "check_output" for Python <= 2.6 (ST2 uses this)
if "check_output" not in dir( subprocess ):
  def f(*popenargs, **kwargs):
    if 'stdout' in kwargs:
      raise ValueError('stdout argument not allowed, it will be overridden.')
    process = subprocess.Popen(stdout=subprocess.PIPE, *popenargs, **kwargs)
    output, unused_err = process.communicate()
    retcode = process.poll()
    if retcode:
      cmd = kwargs.get("args")
      if cmd is None:
        cmd = popenargs[0]
      raise subprocess.CalledProcessError(retcode, cmd)
    return output
  subprocess.check_output = f

SUBLIME_VERSION = "not set"
PLUGIN_PATH = "not set"
PACKAGE_NAME = "not set"
PACKAGE_DIR = "not set"
PACKAGE_USER_DIR = "not set"
ST_PACKAGE_PATH = "not set"
_corona_sdk_debug = False

# In Sublime Text 3 most APIs are unavailable until a module level function is called (fortunately
# sublime.version() is available so we can correctly fake things in Sublime Text 2; see about.py)
SUBLIME_VERSION = 3000 if sublime.version() == '' else int(sublime.version())

def GetSetting(key,default=None):
  # repeated calls to load_settings return same object without further disk reads
  debug("GetSetting: " + str(key) + " (default: " + str(default) + ")")
  settings = sublime.load_settings('Corona Editor.sublime-settings')
  debug("GetSetting: value from CE settings: " + str(settings.get(key, default)))
  # If we don't have a value for this preference in the Corona Editor settings, look in the view for a value
  # (this happens if the preference is set in the main Sublime Text preference file instead of the Corona Editor file)
  if not settings.get(key, default) and sublime and sublime.active_window() and sublime.active_window().active_view():
    debug("GetSetting: getting value from view: " + str(sublime.active_window().active_view()))
    settings = sublime.active_window().active_view().settings()

  value = settings.get(key, default)

  if not value:
    debug("GetSetting: no value for preference '" + str(key) + "' found (using default '"+ str(default) + "')")

  return value

def debug(*args):
  global _corona_sdk_debug
  if _corona_sdk_debug:
    print("Corona Editor: ", '\t'.join(map(str, args)))

InitializedEvent = threading.Event()

def Init():
  global SUBLIME_VERSION
  global PLUGIN_PATH
  global PACKAGE_NAME
  global PACKAGE_DIR
  global PACKAGE_USER_DIR
  global ST_PACKAGE_PATH
  global _corona_sdk_debug

  # Always look for "corona_sdk_debug" in 'Corona Editor.sublime-settings'
  settings = sublime.load_settings('Corona Editor.sublime-settings')
  _corona_sdk_debug = settings.get("corona_sdk_debug", False)

  print("Corona Editor: Init")
  debug("Python: " + str(sys.version))

  PLUGIN_PATH = os.path.dirname(os.path.realpath(__file__))
  if PLUGIN_PATH.lower().endswith('coronasdk-sublimetext'):
    if SUBLIME_VERSION < 3000:
      PLUGIN_PATH = os.path.expanduser('~/Library/Application Support/Sublime Text 2/Packages/Corona Editor')
    else:
      PLUGIN_PATH = os.path.expanduser('~/Library/Application Support/Sublime Text 3/Packages/Corona Editor')
  debug("PLUGIN_PATH: " + PLUGIN_PATH)

  PACKAGE_NAME = os.path.basename(PLUGIN_PATH) if SUBLIME_VERSION < 3000 else os.path.basename(PLUGIN_PATH).replace(".sublime-package", "")
  debug("PACKAGE_NAME: " + PACKAGE_NAME)

  # This is the actual package dir on ST2 and the "unpacked" package dir on ST3 ("~/Library/Application Support/Sublime Text 2/Packages")
  PACKAGE_DIR = os.path.join(sublime.packages_path(), PACKAGE_NAME)
  debug("PACKAGE_DIR: " + PACKAGE_DIR)

  # This is the user dir on ST
  PACKAGE_USER_DIR = os.path.join(sublime.packages_path(), 'User', PACKAGE_NAME)
  debug("PACKAGE_USER_DIR: " + PACKAGE_USER_DIR)

  # This is the faux path used by various Sublime Text functions ("Packages/CoronaSDK/")
  ST_PACKAGE_PATH = "Packages/" + PACKAGE_NAME + "/"
  debug("ST_PACKAGE_PATH: " + ST_PACKAGE_PATH)

  # This allows other threads to wait until this code has run
  InitializedEvent.set()

# This is run by ST3 after the plugins have loaded (and the ST API is ready for calls)
def plugin_loaded():
  Init()

# This fakes the above functionality on ST2
if SUBLIME_VERSION < 3000:
  Init()

# Return the path to the Simulator for the current project and platform
# First we look in any build.settings file in the project, then we look
# for a user preference, then we pick a platform appropriate default
def GetSimulatorCmd(mainlua=None, debug=False):
  platform = sublime.platform()
  arch = sublime.arch()
  view = sublime.active_window().active_view()

  simulator_path = ""
  simulator_flags = []
  simulator_version = GetSetting("corona_sdk_version", None)

  if mainlua is not None:
    simulator_path = GetSimulatorPathFromBuildSettings(mainlua)
    if simulator_path is None:
      simulator_path = GetSetting("corona_sdk_simulator_path", None)

  if platform == "osx":
    if simulator_path is None:
      simulator_path = GetLatestDailyBuildSimulatorPath()
    if simulator_path is None:
      simulator_path = "/Applications/CoronaSDK/Corona Simulator.app"
    if simulator_path.endswith(".app"):
      simulator_path += "/Contents/MacOS/Corona Simulator"
    simulator_flags = ["-singleton", "1"]
    if not GetSetting("corona_sdk_simulator_show_console", False):
      simulator_flags += ["-no-console", "1"]
    if debug:
      simulator_flags.append("-debug")
      simulator_flags.append("1")
      simulator_flags.append("-project")
    if os.path.isfile(simulator_path) and os.access(simulator_path, os.X_OK):
      simulator_version = str(subprocess.check_output(['defaults', 'read', os.path.join(os.path.dirname(simulator_path), '../Info.plist'), 'CFBundleVersion'], stderr=subprocess.STDOUT).decode('utf8').strip())
  elif platform == 'windows':
    if simulator_path is None:
      if arch == "x64":
        simulator_path = "C:\\Program Files (x86)\\Corona Labs\\Corona SDK\\Corona Simulator.exe"
      else:
        simulator_path = "C:\\Program Files\\Corona Labs\\Corona SDK\\Corona Simulator.exe"
    simulator_flags = ["/singleton"]
    if not GetSetting("corona_sdk_simulator_show_console", False):
      simulator_flags += ["/no-console"]
    if debug:
      simulator_flags.append("/debug")

  # Can we find an executable file at the path
  if not os.path.isfile(simulator_path) or not os.access(simulator_path, os.X_OK):
    sublime.error_message("Cannot find executable Corona Simulator at path '{0}'\n\nYou can set the user preference 'corona_sdk_simulator_path' to the location of the Simulator.".format(simulator_path))
    return None, None, None

  return simulator_path, simulator_flags, simulator_version

# Find the highest numbered CoronaSDK directory in /Applications
def GetLatestDailyBuildSimulatorPath():
  dirs = os.listdir("/Applications")

  # Find the highest numbered Daily Build
  maxBuildNum = 0
  buildNum = 0
  latestSDK = ""
  for entry in dirs:
    if entry.startswith("CoronaSDK-"):
      result = re.findall(r'\d+', entry)
      if len(result) > 0:
        buildNum = int(re.findall(r'\d+', entry)[0])
      if buildNum > maxBuildNum:
        maxBuildNum = buildNum
        latestSDK = "/Applications/" + entry

  if len(latestSDK) > 0:
    return  latestSDK + "/Corona Simulator.app"
  else:
    return None

# Given a path to a main.lua file, see if there's a "corona_sdk_simulator_path" setting in
# the corresponding build.settings file
def GetSimulatorPathFromBuildSettings(mainlua):
  simulator_path = None
  project_path = os.path.dirname(mainlua)
  build_settings = os.path.join(project_path, "build.settings")
  bs_contents = None
  if os.path.isfile(build_settings):
    try:
      try: # see if we're in P3
        with open(build_settings, "r", encoding="utf-8") as bs_fd:
          bs_contents = bs_fd.read()
      except: # nope, try a P2 idiom
        with io.open(build_settings, "r", encoding="utf-8") as bs_fd:
          bs_contents = bs_fd.read()
    except IOError:
      pass  # we don't care if the file doesn't exist

  # debug("bs_contents: ", str(bs_contents))
  if bs_contents is not None:
    # Remove comments
    bs_contents = re.sub(r'--.*', '', bs_contents)
    # Note we can't use a Python r'' string here because we then can't escape the single quotes
    bs_matches = re.findall('corona_sdk_simulator_path\s*=\s*["\'](.*)["\']', bs_contents)
    if bs_matches is not None and len(bs_matches) > 0:
      # Last one wins
      simulator_path = bs_matches[-1]
      print("Corona Editor: corona_sdk_simulator_path set from "+str(build_settings))

  return simulator_path

# Given an existing file path or directory, find the likely "main.lua" for this project
def ResolveMainLua(path):
  # debug("ResolveMainLua: path: "+str(path))
  path = os.path.abspath(path)
  if not os.path.isdir(path):
    path = os.path.dirname(path)

  mainlua = os.path.join(path, "main.lua")
  if mainlua == "/main.lua" or mainlua == "\\main.lua":
    return None
  elif os.path.isfile(mainlua):
    return mainlua
  else:
    return ResolveMainLua(os.path.join(path, ".."))

<|endoftext|>


<commit_before>from utilUdpCmd import recvCommand, procCommand
from utilSetting import CSetting
import socket

# TODO: 1> something does not sit right concerning setting passing (g_setting)

g_setting = CSetting()

def main():
	cmdip = ""
	g_setting = CSetting()
	cmdport = g_setting.getCmdPort()
	cmdsock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
	cmdsock.bind((cmdip, cmdport))
	cmdsock.setblocking(0)

	print "cmdport=", cmdport

	rcvcmd = ""

	while 1:
		rcvcmd,rcvd = recvCommand(cmdsock, rcvcmd)
		if rcvd == True and "\n" in rcvcmd:
			procCommand(rcvcmd)
			rcvcmd = ""

if __name__ == '__main__':
	main()
<commit_msg>linemonitor: fix bug<commit_after>from utilUdpCmd import recvCommand, procCommand, getSetting
from utilSetting import CSetting
import socket

# HACKME: something does not sit right concerning setting passing (g_setting)

g_setting = CSetting()

def main():
	cmdip = ""
	g_setting = getSetting() # HACKME:
	cmdport = g_setting.getCmdPort()
	cmdsock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
	cmdsock.bind((cmdip, cmdport))
	cmdsock.setblocking(0)

	print "cmdport=", cmdport

	rcvcmd = ""

	while 1:
		rcvcmd,rcvd = recvCommand(cmdsock, rcvcmd)
		if rcvd == True and "\n" in rcvcmd:
			procCommand(rcvcmd)
			rcvcmd = ""

if __name__ == '__main__':
	main()
<|endoftext|>


<commit_before>#!/usr/bin/env python
import math
from dataclasses import dataclass
from functools import cached_property
from typing import Optional

import torch
from torch import nn

class TextUnembedding(nn.Module):
    def __init__(self, cfg, dtype=None, device=None):
        super().__init__()

        self.final_ln = nn.LayerNorm(cfg.d_model, dtype=dtype, device=device)
        self.unembedding = nn.Linear(
            cfg.d_model, cfg.n_vocab, dtype=dtype, device=device
        )

    def init_weights(self, cfg):
        self.final_ln.reset_parameters()
        self.unembedding.weight.data.normal_(mean=0.0, std=cfg.initializer_range)

    def forward(self, acts):
        acts = self.final_ln(acts)
        return self.unembedding(acts)

@dataclass
class Config:
    n_vocab: int
    n_layer: int
    d_model: int
    d_head: int
    n_ctx: int = 1024
    initializer_range: float = 0.02
    positional_encoding: str = "sin"

    output_head: type = TextUnembedding

    autoregressive_mask: bool = True

    @cached_property
    def d_mlp(self):
        return 4 * self.d_model

    @cached_property
    def n_head(self):
        assert self.d_model % self.d_head == 0
        return self.d_model // self.d_head

    @cached_property
    def n_parameters(self):
        return self.n_layer * (2 * self.d_mlp * self.d_model + 4 * self.d_model**2)

class Resblock(nn.Module):
    def __init__(self, cfg: Config, dtype=None, device=None):
        super().__init__()
        self.attn_ln = nn.LayerNorm(cfg.d_model, dtype=dtype, device=device)
        self.attn = nn.MultiheadAttention(
            cfg.d_model, cfg.n_head, batch_first=True, dtype=dtype, device=device
        )

        self.mlp_ln = nn.LayerNorm(cfg.d_model, dtype=dtype, device=device)
        self.mlp_up = nn.Linear(cfg.d_model, cfg.d_mlp, dtype=dtype, device=device)
        self.mlp_act = nn.ReLU()
        self.mlp_down = nn.Linear(cfg.d_mlp, cfg.d_model, dtype=dtype, device=device)

    def init_weights(self, cfg: Config):
        std = cfg.initializer_range / math.sqrt(cfg.n_layer)
        self.attn_ln.reset_parameters()
        self.attn.in_proj_weight.data.normal_(mean=0, std=std)
        self.attn.out_proj.weight.data.normal_(mean=0, std=std)
        self.mlp_ln.reset_parameters()
        self.mlp_up.weight.data.normal_(mean=0, std=std)
        self.mlp_down.weight.data.normal_(mean=0, std=std)

    def forward(
        self,
        resid,
        attn_mask: Optional[torch.Tensor] = None,
        padding_mask: Optional[torch.Tensor] = None,
    ):
        n_batch, n_ctx, d_model = resid.shape

        attn_ln = self.attn_ln(resid)

        attn_out, attn_pattern = self.attn(
            attn_ln,
            attn_ln,
            attn_ln,
            attn_mask=attn_mask,
            key_padding_mask=padding_mask,
        )
        resid = resid + attn_out

        mlp_ln = self.mlp_ln(resid)
        mlp_out = self.mlp_down(self.mlp_act(self.mlp_up(mlp_ln)))
        return resid + mlp_out

def ar_mask(n_ctx: int, dtype: torch.dtype, device: torch.device):
    return torch.triu(
        torch.ones((n_ctx, n_ctx), dtype=torch.bool, device=device), diagonal=1
    )

class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_n_ctx: int = 2048, device=None, dtype=None):
        super().__init__()

        position = torch.arange(max_n_ctx, device=device, dtype=dtype).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2, device=device, dtype=dtype)
            * (-math.log(10000.0) / d_model)
        )
        pe = torch.zeros((max_n_ctx, d_model), dtype=dtype, device=device)
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: Tensor, shape [batch_size, n_ctx, d_model]
        """
        x = x + self.pe[: x.size(1)]
        return x

class LearnedPositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_n_ctx: int = 2048, device=None, dtype=None):
        super().__init__()

        self.pe = nn.Parameter(
            torch.empty((max_n_ctx, d_model), dtype=dtype, device=device)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: Tensor, shape [batch_size, n_ctx, d_model]
        """
        x = x + self.pe[: x.size(1)]
        return x

class Torso(nn.Module):
    autoregressive_mask: torch.Tensor

    def __init__(self, cfg: Config, dtype=None, device=None):
        super().__init__()

        if cfg.autoregressive_mask:
            self.autoregressive_mask = ar_mask(cfg.n_ctx, dtype=dtype, device=device)

        self.layers = nn.ModuleList(
            [Resblock(cfg, dtype=dtype, device=device) for _ in range(cfg.n_layer)]
        )

    def forward(self, acts, padding_mask=None):
        if hasattr(self, "autoregressive_mask"):
            n_ctx = acts.size(0)
            attn_mask = self.autoregressive_mask[:n_ctx, :n_ctx]
        else:
            attn_mask = None

        for layer in self.layers:
            acts = layer(acts, attn_mask, padding_mask)
        return acts

    def init_weights(self, cfg):
        for layer in self.layers:
            layer.init_weights(cfg)

class TextEmbedding(nn.Module):
    def __init__(self, cfg, dtype=None, device=None):
        super().__init__()

        self.embedding = nn.Embedding(
            cfg.n_vocab, cfg.d_model, dtype=dtype, device=device
        )
        if cfg.positional_encoding == "sin":
            self.positional_encoding = PositionalEncoding(
                d_model=cfg.d_model, max_n_ctx=cfg.n_ctx, dtype=dtype, device=device
            )
        elif cfg.positional_encoding == "learned":
            self.positional_encoding = LearnedPositionalEncoding(
                d_model=cfg.d_model, max_n_ctx=cfg.n_ctx, dtype=dtype, device=device
            )
        elif cfg.positional_encoding == "none":
            self.positional_encoding = lambda acts: acts
        else:
            raise ValueError(
                f"Unknown positional encoding type: {cfg.positional_encoding!r}"
            )

    def init_weights(self, cfg):
        self.embedding.weight.data.normal_(
            mean=0.0, std=cfg.initializer_range * math.sqrt(cfg.d_model)
        )
        if isinstance(self.positional_encoding, LearnedPositionalEncoding):
            self.positional_encoding.pe.data.normal_(
                mean=0.0, std=cfg.initializer_range
            )

    def forward(self, tokens):
        acts = self.embedding(tokens)
        acts = self.positional_encoding(acts)
        return acts

class Transformer(nn.Module):
    def __init__(self, cfg, dtype=None, device=None):
        super().__init__()
        self.cfg = cfg

        self.embedding = TextEmbedding(cfg, dtype=dtype, device=device)
        self.torso = Torso(cfg, dtype=dtype, device=device)
        self.unembedding = cfg.output_head(cfg, dtype=dtype, device=device)

    def init_weights(self):
        self.embedding.init_weights(self.cfg)
        self.torso.init_weights(self.cfg)
        if hasattr(self.unembedding, "init_weights"):
            self.unembedding.init_weights(self.cfg)

    def forward(self, input, padding_mask: torch.Tensor = None):
        acts = self.embedding(input)
        acts = self.torso(acts, padding_mask)
        return self.unembedding(acts)
<commit_msg>slice properly<commit_after>#!/usr/bin/env python
import math
from dataclasses import dataclass
from functools import cached_property
from typing import Optional

import torch
from torch import nn

class TextUnembedding(nn.Module):
    def __init__(self, cfg, dtype=None, device=None):
        super().__init__()

        self.final_ln = nn.LayerNorm(cfg.d_model, dtype=dtype, device=device)
        self.unembedding = nn.Linear(
            cfg.d_model, cfg.n_vocab, dtype=dtype, device=device
        )

    def init_weights(self, cfg):
        self.final_ln.reset_parameters()
        self.unembedding.weight.data.normal_(mean=0.0, std=cfg.initializer_range)

    def forward(self, acts):
        acts = self.final_ln(acts)
        return self.unembedding(acts)

@dataclass
class Config:
    n_vocab: int
    n_layer: int
    d_model: int
    d_head: int
    n_ctx: int = 1024
    initializer_range: float = 0.02
    positional_encoding: str = "sin"

    output_head: type = TextUnembedding

    autoregressive_mask: bool = True

    @cached_property
    def d_mlp(self):
        return 4 * self.d_model

    @cached_property
    def n_head(self):
        assert self.d_model % self.d_head == 0
        return self.d_model // self.d_head

    @cached_property
    def n_parameters(self):
        return self.n_layer * (2 * self.d_mlp * self.d_model + 4 * self.d_model**2)

class Resblock(nn.Module):
    def __init__(self, cfg: Config, dtype=None, device=None):
        super().__init__()
        self.attn_ln = nn.LayerNorm(cfg.d_model, dtype=dtype, device=device)
        self.attn = nn.MultiheadAttention(
            cfg.d_model, cfg.n_head, batch_first=True, dtype=dtype, device=device
        )

        self.mlp_ln = nn.LayerNorm(cfg.d_model, dtype=dtype, device=device)
        self.mlp_up = nn.Linear(cfg.d_model, cfg.d_mlp, dtype=dtype, device=device)
        self.mlp_act = nn.ReLU()
        self.mlp_down = nn.Linear(cfg.d_mlp, cfg.d_model, dtype=dtype, device=device)

    def init_weights(self, cfg: Config):
        std = cfg.initializer_range / math.sqrt(cfg.n_layer)
        self.attn_ln.reset_parameters()
        self.attn.in_proj_weight.data.normal_(mean=0, std=std)
        self.attn.out_proj.weight.data.normal_(mean=0, std=std)
        self.mlp_ln.reset_parameters()
        self.mlp_up.weight.data.normal_(mean=0, std=std)
        self.mlp_down.weight.data.normal_(mean=0, std=std)

    def forward(
        self,
        resid,
        attn_mask: Optional[torch.Tensor] = None,
        padding_mask: Optional[torch.Tensor] = None,
    ):
        n_batch, n_ctx, d_model = resid.shape

        attn_ln = self.attn_ln(resid)

        attn_out, attn_pattern = self.attn(
            attn_ln,
            attn_ln,
            attn_ln,
            attn_mask=attn_mask,
            key_padding_mask=padding_mask,
        )
        resid = resid + attn_out

        mlp_ln = self.mlp_ln(resid)
        mlp_out = self.mlp_down(self.mlp_act(self.mlp_up(mlp_ln)))
        return resid + mlp_out

def ar_mask(n_ctx: int, dtype: torch.dtype, device: torch.device):
    return torch.triu(
        torch.ones((n_ctx, n_ctx), dtype=torch.bool, device=device), diagonal=1
    )

class PositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_n_ctx: int = 2048, device=None, dtype=None):
        super().__init__()

        position = torch.arange(max_n_ctx, device=device, dtype=dtype).unsqueeze(1)
        div_term = torch.exp(
            torch.arange(0, d_model, 2, device=device, dtype=dtype)
            * (-math.log(10000.0) / d_model)
        )
        pe = torch.zeros((max_n_ctx, d_model), dtype=dtype, device=device)
        pe[:, 0::2] = torch.sin(position * div_term)
        pe[:, 1::2] = torch.cos(position * div_term)
        self.register_buffer("pe", pe)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: Tensor, shape [batch_size, n_ctx, d_model]
        """
        x = x + self.pe[: x.size(1)]
        return x

class LearnedPositionalEncoding(nn.Module):
    def __init__(self, d_model: int, max_n_ctx: int = 2048, device=None, dtype=None):
        super().__init__()

        self.pe = nn.Parameter(
            torch.empty((max_n_ctx, d_model), dtype=dtype, device=device)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Args:
            x: Tensor, shape [batch_size, n_ctx, d_model]
        """
        x = x + self.pe[: x.size(1)]
        return x

class Torso(nn.Module):
    autoregressive_mask: torch.Tensor

    def __init__(self, cfg: Config, dtype=None, device=None):
        super().__init__()

        if cfg.autoregressive_mask:
            self.autoregressive_mask = ar_mask(cfg.n_ctx, dtype=dtype, device=device)

        self.layers = nn.ModuleList(
            [Resblock(cfg, dtype=dtype, device=device) for _ in range(cfg.n_layer)]
        )

    def forward(self, acts, padding_mask=None):
        if hasattr(self, "autoregressive_mask"):
            n_ctx = acts.size(1)
            attn_mask = self.autoregressive_mask[:n_ctx, :n_ctx]
        else:
            attn_mask = None

        for layer in self.layers:
            acts = layer(acts, attn_mask, padding_mask)
        return acts

    def init_weights(self, cfg):
        for layer in self.layers:
            layer.init_weights(cfg)

class TextEmbedding(nn.Module):
    def __init__(self, cfg, dtype=None, device=None):
        super().__init__()

        self.embedding = nn.Embedding(
            cfg.n_vocab, cfg.d_model, dtype=dtype, device=device
        )
        if cfg.positional_encoding == "sin":
            self.positional_encoding = PositionalEncoding(
                d_model=cfg.d_model, max_n_ctx=cfg.n_ctx, dtype=dtype, device=device
            )
        elif cfg.positional_encoding == "learned":
            self.positional_encoding = LearnedPositionalEncoding(
                d_model=cfg.d_model, max_n_ctx=cfg.n_ctx, dtype=dtype, device=device
            )
        elif cfg.positional_encoding == "none":
            self.positional_encoding = lambda acts: acts
        else:
            raise ValueError(
                f"Unknown positional encoding type: {cfg.positional_encoding!r}"
            )

    def init_weights(self, cfg):
        self.embedding.weight.data.normal_(
            mean=0.0, std=cfg.initializer_range * math.sqrt(cfg.d_model)
        )
        if isinstance(self.positional_encoding, LearnedPositionalEncoding):
            self.positional_encoding.pe.data.normal_(
                mean=0.0, std=cfg.initializer_range
            )

    def forward(self, tokens):
        acts = self.embedding(tokens)
        acts = self.positional_encoding(acts)
        return acts

class Transformer(nn.Module):
    def __init__(self, cfg, dtype=None, device=None):
        super().__init__()
        self.cfg = cfg

        self.embedding = TextEmbedding(cfg, dtype=dtype, device=device)
        self.torso = Torso(cfg, dtype=dtype, device=device)
        self.unembedding = cfg.output_head(cfg, dtype=dtype, device=device)

    def init_weights(self):
        self.embedding.init_weights(self.cfg)
        self.torso.init_weights(self.cfg)
        if hasattr(self.unembedding, "init_weights"):
            self.unembedding.init_weights(self.cfg)

    def forward(self, input, padding_mask: torch.Tensor = None):
        acts = self.embedding(input)
        acts = self.torso(acts, padding_mask)
        return self.unembedding(acts)
<|endoftext|>