<commit_before># Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""A Python interface for creating dataset servers."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

# pylint: disable=invalid-import-order,g-bad-import-order, unused-import
from tensorflow.core.protobuf.data.experimental import service_config_pb2
from tensorflow.python import pywrap_tensorflow
from tensorflow.python.data.experimental.service import _pywrap_server_lib
from tensorflow.python.util.tf_export import tf_export

@tf_export("data.experimental.service.DispatcherConfig")
class DispatcherConfig(
    collections.namedtuple(
        "DispatcherConfig",
        ["port", "protocol", "work_dir", "fault_tolerant_mode"])):
  """Configuration class for tf.data service dispatchers.

  Fields:
    port: Specifies the port to bind to. A value of 0 indicates that the server
      may bind to any available port.
    protocol: The protocol to use for communicating with the tf.data service.
      Acceptable values include `"grpc" and "grpc+local"`.
    work_dir: A directory to store dispatcher state in. This
      argument is required for the dispatcher to be able to recover from
      restarts.
    fault_tolerant_mode: Whether the dispatcher should write its state to a
      journal so that it can recover from restarts. Dispatcher state, including
      registered datasets and created jobs, is synchronously written to the
      journal before responding to RPCs. If `True`, `work_dir` must also be
      specified.
  """

  def __new__(cls,
              port=0,
              protocol="grpc",
              work_dir=None,
              fault_tolerant_mode=False):
    return super(DispatcherConfig, cls).__new__(cls, port, protocol, work_dir,
                                                fault_tolerant_mode)

@tf_export("data.experimental.service.DispatchServer", v1=[])
class DispatchServer(object):
  """An in-process tf.data service dispatch server.

  A `tf.data.experimental.service.DispatchServer` coordinates a cluster of
  `tf.data.experimental.service.WorkerServer`s. When the workers start, they
  register themselves with the dispatcher.

  >>> dispatcher = tf.data.experimental.service.DispatchServer()
  >>> dispatcher_address = dispatcher.target.split("://")[1]
  >>> worker = tf.data.experimental.service.WorkerServer(WorkerConfig(
  ...     dispatcher_address=dispatcher_address))
  >>> dataset = tf.data.Dataset.range(10)
  >>> dataset = dataset.apply(tf.data.experimental.service.distribute(
  ...     processing_mode="parallel_epochs", service=dispatcher.target))
  >>> print(list(dataset.as_numpy_iterator()))
  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

  When starting a dedicated tf.data dispatch process, use join() to block
  indefinitely after starting up the server.

  ```
  dispatcher = tf.data.experimental.service.DispatchServer(
      tf.data.experimental.service.DispatcherConfig(port=5050))
  dispatcher.join()
  ```

  To start a `DispatchServer` in fault-tolerant mode, set `work_dir` and
  `fault_tolerant_mode` like below:

  ```
  dispatcher = tf.data.experimental.service.DispatchServer(
      tf.data.experimental.service.DispatcherConfig(
          port=5050,
          work_dir="gs://my-bucket/dispatcher/work_dir",
          fault_tolerant_mode=True))
  ```
  """

  def __init__(self, config=None, start=True):
    """Creates a new dispatch server.

    Args:
      config: (Optional.) A `tf.data.experimental.service.DispatcherConfig`
        configration. If `None`, the dispatcher will be use default
        configuration values.
      start: (Optional.) Boolean, indicating whether to start the server after
        creating it.
    """
    config = config or DispatcherConfig()
    if config.fault_tolerant_mode and not config.work_dir:
      raise ValueError(
          "Cannot enable fault tolerant mode without configuring a work_dir")
    self._config = config
    config_proto = service_config_pb2.DispatcherConfig(
        port=config.port,
        protocol=config.protocol,
        work_dir=config.work_dir,
        fault_tolerant_mode=config.fault_tolerant_mode)
    self._server = _pywrap_server_lib.TF_DATA_NewDispatchServer(
        config_proto.SerializeToString())
    if start:
      self._server.start()

  def start(self):
    """Starts this server.

    >>> dispatcher = tf.data.experimental.service.DispatchServer(start=False)
    >>> dispatcher.start()

    Raises:
      tf.errors.OpError: Or one of its subclasses if an error occurs while
        starting the server.
    """
    self._server.start()

  def join(self):
    """Blocks until the server has shut down.

    This is useful when starting a dedicated dispatch process.

    ```
    dispatcher = tf.data.experimental.service.DispatchServer(
        tf.data.experimental.service.DispatcherConfig(port=5050))
    dispatcher.join()
    ```

    Raises:
      tf.errors.OpError: Or one of its subclasses if an error occurs while
        joining the server.
    """
    self._server.join()

  @property
  def target(self):
    """Returns a target that can be used to connect to the server.

    >>> dispatcher = tf.data.experimental.service.DispatchServer()
    >>> dataset = tf.data.Dataset.range(10)
    >>> dataset = dataset.apply(tf.data.experimental.service.distribute(
    ...     processing_mode="parallel_epochs", service=dispatcher.target))

    The returned string will be in the form protocol://address, e.g.
    "grpc://localhost:5050".
    """
    return "{0}://localhost:{1}".format(self._config.protocol,
                                        self._server.bound_port())

  def _stop(self):
    """Stops the server.

    Raises:
      tf.errors.OpError: Or one of its subclasses if an error occurs while
        stopping the server.
    """
    self._server.stop()

  def __del__(self):
    self._stop()

  @property
  def _address(self):
    """Returns the address of the server.

    The returned string will be in the form address:port, e.g. "localhost:1000".
    """
    return "localhost:{0}".format(self._server.bound_port())

  def _num_workers(self):
    """Returns the number of workers registered with the dispatcher."""
    return self._server.num_workers()

@tf_export("data.experimental.service.WorkerConfig")
class WorkerConfig(
    collections.namedtuple(
        "WorkerConfig",
        ["dispatcher_address", "worker_address", "port", "protocol"])):
  """Configuration class for tf.data service dispatchers.

  Fields:
    dispatcher_address: Specifies the address of the dispatcher.
    worker_address: Specifies the address of the worker server. This address is
      passed to the dispatcher so that the dispatcher can tell clients how to
      connect to this worker.
    port: Specifies the port to bind to. A value of 0 indicates that the worker
      can bind to any available port.
    protocol: (Optional.) Specifies the protocol to be used by the server.
      Acceptable values include `"grpc" and "grpc+local"`.
  """

  def __new__(cls,
              dispatcher_address,
              worker_address=None,
              port=0,
              protocol="grpc"):
    worker_address = ("localhost:%port%"
                      if worker_address is None else worker_address)
    return super(WorkerConfig, cls).__new__(cls, dispatcher_address,
                                            worker_address, port, protocol)

@tf_export("data.experimental.service.WorkerServer", v1=[])
class WorkerServer(object):
  """An in-process tf.data service worker server.

  A `tf.data.experimental.service.WorkerServer` performs `tf.data.Dataset`
  processing for user-defined datasets, and provides the resulting elements over
  RPC. A worker is associated with a single
  `tf.data.experimental.service.DispatchServer`.

  >>> dispatcher = tf.data.experimental.service.DispatchServer()
  >>> dispatcher_address = dispatcher.target.split("://")[1]
  >>> worker = tf.data.experimental.service.WorkerServer(
  ...     tf.data.experimental.service.WorkerConfig(
  ...         dispatcher_address=dispatcher_address))
  >>> dataset = tf.data.Dataset.range(10)
  >>> dataset = dataset.apply(tf.data.experimental.service.distribute(
  ...     processing_mode="parallel_epochs", service=dispatcher.target))
  >>> print(list(dataset.as_numpy_iterator()))
  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

  When starting a dedicated tf.data worker process, use join() to block
  indefinitely after starting up the server.

  ```
  worker = tf.data.experimental.service.WorkerServer(
      port=5051, dispatcher_address="grpc://localhost:5050")
  worker.join()
  ```
  """

  def __init__(self, config, start=True):
    """Creates a new worker server.

    Args:
      config: A `tf.data.experimental.service.WorkerConfig` configration.
      start: (Optional.) Boolean, indicating whether to start the server after
        creating it.
    """
    if config.dispatcher_address is None:
      raise ValueError("must specify a dispatcher_address")
    config_proto = service_config_pb2.WorkerConfig(
        dispatcher_address=config.dispatcher_address,
        worker_address=config.worker_address,
        port=config.port,
        protocol=config.protocol)
    self._server = _pywrap_server_lib.TF_DATA_NewWorkerServer(
        config_proto.SerializeToString())
    if start:
      self._server.start()

  def start(self):
    """Starts this server.

    Raises:
      tf.errors.OpError: Or one of its subclasses if an error occurs while
        starting the server.
    """
    self._server.start()

  def join(self):
    """Blocks until the server has shut down.

    This is useful when starting a dedicated worker process.

    ```
    worker_server = tf.data.experimental.service.WorkerServer(
        port=5051, dispatcher_address="grpc://localhost:5050")
    worker_server.join()
    ```

    This method currently blocks forever.

    Raises:
      tf.errors.OpError: Or one of its subclasses if an error occurs while
        joining the server.
    """
    self._server.join()

  def _stop(self):
    """Stops the server.

    Raises:
      tf.errors.OpError: Or one of its subclasses if an error occurs while
        stopping the server.
    """
    self._server.stop()

  def __del__(self):
    self._stop()

  @property
  def _address(self):
    """Returns the address of the server.

    The returned string will be in the form address:port, e.g. "localhost:1000".
    """
    return "localhost:{0}".format(self._server.bound_port())
<commit_msg>supporting protocol changes<commit_after># Copyright 2020 The TensorFlow Authors. All Rights Reserved.
#
# 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.
# ==============================================================================
"""A Python interface for creating dataset servers."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import collections

# pylint: disable=invalid-import-order,g-bad-import-order, unused-import
from tensorflow.core.protobuf.data.experimental import service_config_pb2
from tensorflow.python import pywrap_tensorflow
from tensorflow.python.data.experimental.service import _pywrap_server_lib
from tensorflow.python.util.tf_export import tf_export

@tf_export("data.experimental.service.DispatcherConfig")
class DispatcherConfig(
    collections.namedtuple(
        "DispatcherConfig",
        ["port", "protocol", "work_dir", "fault_tolerant_mode"])):
  """Configuration class for tf.data service dispatchers.

  Fields:
    port: Specifies the port to bind to. A value of 0 indicates that the server
      may bind to any available port.
    protocol: The protocol to use for communicating with the tf.data service.
      Defaults to the only acceptable value of `"grpc"`.
    work_dir: A directory to store dispatcher state in. This
      argument is required for the dispatcher to be able to recover from
      restarts.
    fault_tolerant_mode: Whether the dispatcher should write its state to a
      journal so that it can recover from restarts. Dispatcher state, including
      registered datasets and created jobs, is synchronously written to the
      journal before responding to RPCs. If `True`, `work_dir` must also be
      specified.
  """

  def __new__(cls,
              port=0,
              protocol="grpc",
              work_dir=None,
              fault_tolerant_mode=False):
    return super(DispatcherConfig, cls).__new__(cls, port, protocol, work_dir,
                                                fault_tolerant_mode)

@tf_export("data.experimental.service.DispatchServer", v1=[])
class DispatchServer(object):
  """An in-process tf.data service dispatch server.

  A `tf.data.experimental.service.DispatchServer` coordinates a cluster of
  `tf.data.experimental.service.WorkerServer`s. When the workers start, they
  register themselves with the dispatcher.

  >>> dispatcher = tf.data.experimental.service.DispatchServer()
  >>> dispatcher_address = dispatcher.target.split("://")[1]
  >>> worker = tf.data.experimental.service.WorkerServer(WorkerConfig(
  ...     dispatcher_address=dispatcher_address))
  >>> dataset = tf.data.Dataset.range(10)
  >>> dataset = dataset.apply(tf.data.experimental.service.distribute(
  ...     processing_mode="parallel_epochs", service=dispatcher.target))
  >>> print(list(dataset.as_numpy_iterator()))
  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

  When starting a dedicated tf.data dispatch process, use join() to block
  indefinitely after starting up the server.

  ```
  dispatcher = tf.data.experimental.service.DispatchServer(
      tf.data.experimental.service.DispatcherConfig(port=5050))
  dispatcher.join()
  ```

  To start a `DispatchServer` in fault-tolerant mode, set `work_dir` and
  `fault_tolerant_mode` like below:

  ```
  dispatcher = tf.data.experimental.service.DispatchServer(
      tf.data.experimental.service.DispatcherConfig(
          port=5050,
          work_dir="gs://my-bucket/dispatcher/work_dir",
          fault_tolerant_mode=True))
  ```
  """

  def __init__(self, config=None, start=True):
    """Creates a new dispatch server.

    Args:
      config: (Optional.) A `tf.data.experimental.service.DispatcherConfig`
        configration. If `None`, the dispatcher will be use default
        configuration values.
      start: (Optional.) Boolean, indicating whether to start the server after
        creating it.
    """
    config = config or DispatcherConfig()
    if config.fault_tolerant_mode and not config.work_dir:
      raise ValueError(
          "Cannot enable fault tolerant mode without configuring a work_dir")
    self._config = config
    config_proto = service_config_pb2.DispatcherConfig(
        port=config.port,
        protocol=config.protocol,
        work_dir=config.work_dir,
        fault_tolerant_mode=config.fault_tolerant_mode)
    self._server = _pywrap_server_lib.TF_DATA_NewDispatchServer(
        config_proto.SerializeToString())
    if start:
      self._server.start()

  def start(self):
    """Starts this server.

    >>> dispatcher = tf.data.experimental.service.DispatchServer(start=False)
    >>> dispatcher.start()

    Raises:
      tf.errors.OpError: Or one of its subclasses if an error occurs while
        starting the server.
    """
    self._server.start()

  def join(self):
    """Blocks until the server has shut down.

    This is useful when starting a dedicated dispatch process.

    ```
    dispatcher = tf.data.experimental.service.DispatchServer(
        tf.data.experimental.service.DispatcherConfig(port=5050))
    dispatcher.join()
    ```

    Raises:
      tf.errors.OpError: Or one of its subclasses if an error occurs while
        joining the server.
    """
    self._server.join()

  @property
  def target(self):
    """Returns a target that can be used to connect to the server.

    >>> dispatcher = tf.data.experimental.service.DispatchServer()
    >>> dataset = tf.data.Dataset.range(10)
    >>> dataset = dataset.apply(tf.data.experimental.service.distribute(
    ...     processing_mode="parallel_epochs", service=dispatcher.target))

    The returned string will be in the form protocol://address, e.g.
    "grpc://localhost:5050".
    """
    return "{0}://localhost:{1}".format(self._config.protocol,
                                        self._server.bound_port())

  def _stop(self):
    """Stops the server.

    Raises:
      tf.errors.OpError: Or one of its subclasses if an error occurs while
        stopping the server.
    """
    self._server.stop()

  def __del__(self):
    self._stop()

  @property
  def _address(self):
    """Returns the address of the server.

    The returned string will be in the form address:port, e.g. "localhost:1000".
    """
    return "localhost:{0}".format(self._server.bound_port())

  def _num_workers(self):
    """Returns the number of workers registered with the dispatcher."""
    return self._server.num_workers()

@tf_export("data.experimental.service.WorkerConfig")
class WorkerConfig(
    collections.namedtuple(
        "WorkerConfig",
        ["dispatcher_address", "worker_address", "port", "protocol"])):
  """Configuration class for tf.data service dispatchers.

  Fields:
    dispatcher_address: Specifies the address of the dispatcher.
    worker_address: Specifies the address of the worker server. This address is
      passed to the dispatcher so that the dispatcher can tell clients how to
      connect to this worker.
    port: Specifies the port to bind to. A value of 0 indicates that the worker
      can bind to any available port.
    protocol: (Optional.) Specifies the protocol to be used by the server.
      Defaults to the only acceptable value of `"grpc"`.
  """

  def __new__(cls,
              dispatcher_address,
              worker_address=None,
              port=0,
              protocol="grpc"):
    worker_address = ("localhost:%port%"
                      if worker_address is None else worker_address)
    return super(WorkerConfig, cls).__new__(cls, dispatcher_address,
                                            worker_address, port, protocol)

@tf_export("data.experimental.service.WorkerServer", v1=[])
class WorkerServer(object):
  """An in-process tf.data service worker server.

  A `tf.data.experimental.service.WorkerServer` performs `tf.data.Dataset`
  processing for user-defined datasets, and provides the resulting elements over
  RPC. A worker is associated with a single
  `tf.data.experimental.service.DispatchServer`.

  >>> dispatcher = tf.data.experimental.service.DispatchServer()
  >>> dispatcher_address = dispatcher.target.split("://")[1]
  >>> worker = tf.data.experimental.service.WorkerServer(
  ...     tf.data.experimental.service.WorkerConfig(
  ...         dispatcher_address=dispatcher_address))
  >>> dataset = tf.data.Dataset.range(10)
  >>> dataset = dataset.apply(tf.data.experimental.service.distribute(
  ...     processing_mode="parallel_epochs", service=dispatcher.target))
  >>> print(list(dataset.as_numpy_iterator()))
  [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]

  When starting a dedicated tf.data worker process, use join() to block
  indefinitely after starting up the server.

  ```
  worker = tf.data.experimental.service.WorkerServer(
      port=5051, dispatcher_address="grpc://localhost:5050")
  worker.join()
  ```
  """

  def __init__(self, config, start=True):
    """Creates a new worker server.

    Args:
      config: A `tf.data.experimental.service.WorkerConfig` configration.
      start: (Optional.) Boolean, indicating whether to start the server after
        creating it.
    """
    if config.dispatcher_address is None:
      raise ValueError("must specify a dispatcher_address")
    config_proto = service_config_pb2.WorkerConfig(
        dispatcher_address=config.dispatcher_address,
        worker_address=config.worker_address,
        port=config.port,
        protocol=config.protocol)
    self._server = _pywrap_server_lib.TF_DATA_NewWorkerServer(
        config_proto.SerializeToString())
    if start:
      self._server.start()

  def start(self):
    """Starts this server.

    Raises:
      tf.errors.OpError: Or one of its subclasses if an error occurs while
        starting the server.
    """
    self._server.start()

  def join(self):
    """Blocks until the server has shut down.

    This is useful when starting a dedicated worker process.

    ```
    worker_server = tf.data.experimental.service.WorkerServer(
        port=5051, dispatcher_address="grpc://localhost:5050")
    worker_server.join()
    ```

    This method currently blocks forever.

    Raises:
      tf.errors.OpError: Or one of its subclasses if an error occurs while
        joining the server.
    """
    self._server.join()

  def _stop(self):
    """Stops the server.

    Raises:
      tf.errors.OpError: Or one of its subclasses if an error occurs while
        stopping the server.
    """
    self._server.stop()

  def __del__(self):
    self._stop()

  @property
  def _address(self):
    """Returns the address of the server.

    The returned string will be in the form address:port, e.g. "localhost:1000".
    """
    return "localhost:{0}".format(self._server.bound_port())
<|endoftext|>


<commit_before><commit_msg>Set background_median and background_std to None when no background ROI's present<commit_after><|endoftext|>


<commit_before># emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
"""
Intensity-based image registration
"""

import numpy as np

from ...core.image.image_spaces import (make_xyz_image,
                                        as_xyz_image,
                                        xyz_affine)

from .optimizer import configure_optimizer
from .affine import inverse_affine, subgrid_affine, affine_transforms
from .chain_transform import ChainTransform
from .similarity_measures import similarity_measures as _sms
from ._registration import _joint_histogram

MAX_INT = np.iinfo(np.intp).max

# Module globals
VERBOSE = True  # enables online print statements
OPTIMIZER = 'powell'
XTOL = 1e-2
FTOL = 1e-2
GTOL = 1e-3
MAXITER = 25
MAXFUN = None
CLAMP_DTYPE = 'short'  # do not edit
NPOINTS = 64 ** 3

# Dictionary of interpolation methods (partial volume, trilinear,
# random)
interp_methods = {'pv': 0, 'tri': 1, 'rand': -1}

class HistogramRegistration(object):
    """
    A class to reprensent a generic intensity-based image registration
    algorithm.
    """
    def __init__(self, from_img, to_img,
                 from_bins=256, to_bins=None,
                 from_mask=None, to_mask=None,
                 similarity='crl1', interp='pv',
                 normalize=True, dist=None):
        """
        Creates a new histogram registration object.

        Parameters
        ----------
        from_img : nipy-like image
          `From` image
        to_img : nipy-like image
          `To` image
        from_bins : integer
          Number of histogram bins to represent the `from` image
        to_bins : integer
          Number of histogram bins to represent the `to` image
        from_mask : array-like
          Mask to apply to the `from` image
        to_mask : array-like
          Mask to apply to the `to` image
        similarity : str or callable
          Cost-function for assessing image similarity. If a string,
          one of 'cc': correlation coefficient, 'cr': correlation
          ratio, 'crl1': L1-norm based correlation ratio, 'mi': mutual
          information, 'nmi': normalized mutual information, 'slr':
          supervised log-likelihood ratio. If a callable, it should
          take a two-dimensional array representing the image joint
          histogram as an input and return a float.
       interp : str
         Interpolation method.  One of 'pv': Partial volume, 'tri':
         Trilinear, 'rand': Random interpolation.  See ``joint_histogram.c``
        """
        # Function assumes xyx_affine for inputs
        from_img = as_xyz_image(from_img)
        to_img = as_xyz_image(to_img)

        # Binning sizes
        if to_bins is None:
            to_bins = from_bins

        # Clamping of the `from` image. The number of bins may be
        # overriden if unnecessarily large.
        data, from_bins = clamp(from_img.get_data(), from_bins,
                                mask=from_mask)
        self._from_img = make_xyz_image(data, xyz_affine(from_img), 'scanner')
        # Set field of view in the `from` image with potential
        # subsampling for faster similarity evaluation. This also sets
        # the _from_data and _vox_coords attributes
        if from_mask is None:
            self.subsample(npoints=NPOINTS)
        else:
            corner, size = smallest_bounding_box(from_mask)
            self.set_fov(corner=corner, size=size, npoints=NPOINTS)

        # Clamping of the `to` image including padding with -1
        data, to_bins = clamp(to_img.get_data(), to_bins,
                              mask=to_mask)
        self._to_data = -np.ones(np.array(to_img.shape) + 2, dtype=CLAMP_DTYPE)
        self._to_data[1:-1, 1:-1, 1:-1] = data
        self._to_inv_affine = inverse_affine(xyz_affine(to_img))

        # Joint histogram: must be double contiguous as it will be
        # passed to C routines which assume so
        self._joint_hist = np.zeros([from_bins, to_bins], dtype='double')

        # Set default registration parameters
        self._set_interp(interp)
        self._set_similarity((similarity, normalize, dist))

    def _get_interp(self):
        return interp_methods.keys()\
            [interp_methods.values().index(self._interp)]

    def _set_interp(self, interp):
        self._interp = interp_methods[interp]

    interp = property(_get_interp, _set_interp)

    def set_fov(self, spacing=None, corner=(0, 0, 0), size=None,
                npoints=None):
        """
        Defines a subset of the `from` image to restrict joint
        histogram computation.

        Parameters
        ----------
        spacing : sequence (3,) of positive integers
          Subsampling of image in voxels, where None (default) results
          in the subsampling to be automatically adjusted to roughly
          match a cubic grid with `npoints` voxels
        corner : sequence (3,) of positive integers
          Bounding box origin in voxel coordinates
        size : sequence (3,) of positive integers
          Desired bounding box size
        npoints : positive integer
          Desired number of voxels in the bounding box. If a `spacing`
          argument is provided, then `npoints` is ignored.
        """
        if spacing is None and npoints is None:
            spacing = [1, 1, 1]
        if size is None:
            size = self._from_img.shape
        slicer = lambda c, s, sp:\
            tuple([slice(c[i], s[i] + c[i], sp[i]) for i in range(3)])
        # Adjust spacing to match desired field of view size
        if not spacing is None:
            fov_data = self._from_img.get_data()[slicer(corner, size, spacing)]
        else:
            fov_data = self._from_img.get_data()[slicer(corner, size, [1, 1, 1])]
            spacing = ideal_spacing(fov_data, npoints=npoints)
            fov_data = self._from_img.get_data()[slicer(corner, size, spacing)]
        self._from_data = fov_data
        self._from_npoints = (fov_data >= 0).sum()
        self._from_affine = subgrid_affine(xyz_affine(self._from_img),
                                           slicer(corner, size, spacing))
        # We cache the voxel coordinates of the clamped image
        self._vox_coords =\
            np.indices(self._from_data.shape).transpose((1, 2, 3, 0))

    def subsample(self, spacing=None, npoints=None):
        self.set_fov(spacing=spacing, npoints=npoints)

    def _set_similarity(self, arg):
        normalize = True
        dist = None
        if not isinstance(arg, (tuple, list)):
            similarity = arg
        else:
            if len(arg) == 1:
                similarity = arg[0]
            elif len(arg) == 2:
                similarity, normalize = arg
            elif len(arg) == 3:
                similarity, normalize, dist = arg
            else:
                raise ValueError('cannot understand input argument')
        if similarity in _sms:
            self._similarity = similarity
            self._similarity_call =\
                _sms[similarity](self._joint_hist.shape, normalize, dist)
        else:
            if not hasattr(similarity, '__call__'):
                raise ValueError('similarity should be callable')
            self._similarity = 'custom'
            self._similarity_call = similarity

    def _get_similarity(self):
        return self._similarity

    similarity = property(_get_similarity, _set_similarity)

    def eval(self, T):
        """
        Evaluate similarity function given a world-to-world transform.

        Parameters
        ----------
        T : Transform
            Transform object implementing ``apply`` method
        """
        Tv = ChainTransform(T, pre=self._from_affine, post=self._to_inv_affine)
        return self._eval(Tv)

    def _eval(self, Tv):
        """
        Evaluate similarity function given a voxel-to-voxel transform.

        Parameters
        ----------
        Tv : Transform
             Transform object implementing ``apply`` method
             Should map voxel space to voxel space
        """
        # trans_vox_coords needs be C-contiguous
        trans_vox_coords = Tv.apply(self._vox_coords)
        interp = self._interp
        if self._interp < 0:
            interp = - np.random.randint(MAX_INT)
        _joint_histogram(self._joint_hist,
                         self._from_data.flat,  # array iterator
                         self._to_data,
                         trans_vox_coords,
                         interp)
        return self._similarity_call(self._joint_hist)

    def optimize(self, T, optimizer=OPTIMIZER, **kwargs):
        """ Optimize transform `T` with respect to similarity measure.

        The input object `T` will change as a result of the optimization.

        Parameters
        ----------
        T : object or str
          An object representing a transformation that should
          implement ``apply`` method and ``param`` attribute or
          property. If a string, one of 'rigid', 'similarity', or
          'affine'. The corresponding transformation class is then
          initialized by default.
        optimizer : str
          Name of optimization function (one of 'powell', 'steepest',
          'cg', 'bfgs', 'simplex')
        **kwargs : dict
          keyword arguments to pass to optimizer
        """
        # Replace T if a string is passed
        if T in affine_transforms:
            T = affine_transforms[T]()

        # Pull callback out of keyword arguments, if present
        callback = kwargs.pop('callback', None)

        # Create transform chain object with T generating params
        Tv = ChainTransform(T, pre=self._from_affine, post=self._to_inv_affine)
        tc0 = Tv.param

        # Cost function to minimize
        def cost(tc):
            # This is where the similarity function is calculcated
            Tv.param = tc
            return -self._eval(Tv)

        # Callback during optimization
        if callback is None and VERBOSE:

            def callback(tc):
                Tv.param = tc
                print(Tv.optimizable)
                print(str(self.similarity) + ' = %s' % self._eval(Tv))
                print('')

        # Switching to the appropriate optimizer
        if VERBOSE:
            print('Initial guess...')
            print(Tv.optimizable)

        kwargs.setdefault('xtol', XTOL)
        kwargs.setdefault('ftol', FTOL)
        kwargs.setdefault('gtol', GTOL)
        kwargs.setdefault('maxiter', MAXITER)
        kwargs.setdefault('maxfun', MAXFUN)

        fmin, args, kwargs = configure_optimizer(optimizer,
                                                 fprime=None,
                                                 fhess=None,
                                                 **kwargs)

        # Output
        if VERBOSE:
            print ('Optimizing using %s' % fmin.__name__)
        kwargs['callback'] = callback
        Tv.param = fmin(cost, tc0, *args, **kwargs)
        return Tv.optimizable

    def explore(self, T0, *args):
        """
        Evaluate the similarity at the transformations specified by
        sequences of parameter values.

        For instance:

        explore(T0, (0, [-1,0,1]), (4, [-2.,2]))
        """
        nparams = T0.param.size
        deltas = [[0] for i in range(nparams)]
        for a in args:
            deltas[a[0]] = a[1]
        grids = np.mgrid[[slice(0, len(d)) for d in deltas]]
        ntrials = np.prod(grids.shape[1:])
        Deltas = [np.asarray(deltas[i])[grids[i, :]].ravel()\
                      for i in range(nparams)]
        simis = np.zeros(ntrials)
        params = np.zeros([nparams, ntrials])

        Tv = ChainTransform(T0, pre=self._from_affine,
                            post=self._to_inv_affine)
        param0 = Tv.param
        for i in range(ntrials):
            param = param0 + np.array([D[i] for D in Deltas])
            Tv.param = param
            simis[i] = self._eval(Tv)
            params[:, i] = param

        return simis, params

def _clamp(x, y, bins):

    # Threshold
    dmaxmax = 2 ** (8 * y.dtype.itemsize - 1) - 1
    dmax = bins - 1  # default output maximum value
    if dmax > dmaxmax:
        raise ValueError('Excess number of bins')
    xmin = float(x.min())
    xmax = float(x.max())
    d = xmax - xmin

    """
    If the image dynamic is small, no need for compression: just
    downshift image values and re-estimate the dynamic range (hence
    xmax is translated to xmax-tth casted to the appropriate
    dtype. Otherwise, compress after downshifting image values (values
    equal to the threshold are reset to zero).
    """
    if issubclass(x.dtype.type, np.integer) and d <= dmax:
        y[:] = x - xmin
        bins = int(d) + 1
    else:
        a = dmax / d
        y[:] = np.round(a * (x - xmin))

    return y, bins

def clamp(x, bins, mask=None):
    """
    Clamp array values that fall within a given mask in the range
    [0..bins-1] and reset masked values to -1.

    Parameters
    ----------
    x : ndarray
      The input array

    bins : number
      Desired number of bins

    mask : ndarray, tuple or slice
      Anything such that x[mask] is an array.

    Returns
    -------
    y : ndarray
      Clamped array, masked items are assigned -1

    bins : number
      Adjusted number of bins

    """
    if bins > np.iinfo(np.short).max:
        raise ValueError('Too large a bin size')
    y = -np.ones(x.shape, dtype=CLAMP_DTYPE)
    if mask is None:
        y, bins = _clamp(x, y, bins)
    else:
        ym = y[mask]
        xm = x[mask]
        ym, bins = _clamp(xm, ym, bins)
        y[mask] = ym
    return y, bins

def ideal_spacing(data, npoints):
    """
    Tune spacing factors so that the number of voxels in the
    output block matches a given number.

    Parameters
    ----------
    data : ndarray or sequence
      Data image to subsample

    npoints : number
      Target number of voxels (negative values will be ignored)

    Returns
    -------
    spacing: ndarray
      Spacing factors

    """
    dims = data.shape
    actual_npoints = (data >= 0).sum()
    spacing = np.ones(3, dtype='uint')

    while actual_npoints > npoints:

        # Subsample the direction with the highest number of samples
        ddims = dims / spacing
        if ddims[0] >= ddims[1] and ddims[0] >= ddims[2]:
            dir = 0
        elif ddims[1] > ddims[0] and ddims[1] >= ddims[2]:
            dir = 1
        else:
            dir = 2
        spacing[dir] += 1
        subdata = data[::spacing[0], ::spacing[1], ::spacing[2]]
        actual_npoints = (subdata >= 0).sum()

    return spacing

def smallest_bounding_box(msk):
    """
    Extract the smallest bounding box from a mask
    """
    x, y, z = np.where(msk > 0)
    corner = [x.min(), y.min(), z.min()]
    size = [x.max() + 1, y.max() + 1, z.max() + 1]
    return corner, size
<commit_msg>avoid modifying input arg in explore method of HistogramRegistration<commit_after># emacs: -*- mode: python; py-indent-offset: 4; indent-tabs-mode: nil -*-
# vi: set ft=python sts=4 ts=4 sw=4 et:
"""
Intensity-based image registration
"""

import numpy as np

from ...core.image.image_spaces import (make_xyz_image,
                                        as_xyz_image,
                                        xyz_affine)

from .optimizer import configure_optimizer
from .affine import inverse_affine, subgrid_affine, affine_transforms
from .chain_transform import ChainTransform
from .similarity_measures import similarity_measures as _sms
from ._registration import _joint_histogram

MAX_INT = np.iinfo(np.intp).max

# Module globals
VERBOSE = True  # enables online print statements
OPTIMIZER = 'powell'
XTOL = 1e-2
FTOL = 1e-2
GTOL = 1e-3
MAXITER = 25
MAXFUN = None
CLAMP_DTYPE = 'short'  # do not edit
NPOINTS = 64 ** 3

# Dictionary of interpolation methods (partial volume, trilinear,
# random)
interp_methods = {'pv': 0, 'tri': 1, 'rand': -1}

class HistogramRegistration(object):
    """
    A class to reprensent a generic intensity-based image registration
    algorithm.
    """
    def __init__(self, from_img, to_img,
                 from_bins=256, to_bins=None,
                 from_mask=None, to_mask=None,
                 similarity='crl1', interp='pv',
                 normalize=True, dist=None):
        """
        Creates a new histogram registration object.

        Parameters
        ----------
        from_img : nipy-like image
          `From` image
        to_img : nipy-like image
          `To` image
        from_bins : integer
          Number of histogram bins to represent the `from` image
        to_bins : integer
          Number of histogram bins to represent the `to` image
        from_mask : array-like
          Mask to apply to the `from` image
        to_mask : array-like
          Mask to apply to the `to` image
        similarity : str or callable
          Cost-function for assessing image similarity. If a string,
          one of 'cc': correlation coefficient, 'cr': correlation
          ratio, 'crl1': L1-norm based correlation ratio, 'mi': mutual
          information, 'nmi': normalized mutual information, 'slr':
          supervised log-likelihood ratio. If a callable, it should
          take a two-dimensional array representing the image joint
          histogram as an input and return a float.
       interp : str
         Interpolation method.  One of 'pv': Partial volume, 'tri':
         Trilinear, 'rand': Random interpolation.  See ``joint_histogram.c``
        """
        # Function assumes xyx_affine for inputs
        from_img = as_xyz_image(from_img)
        to_img = as_xyz_image(to_img)

        # Binning sizes
        if to_bins is None:
            to_bins = from_bins

        # Clamping of the `from` image. The number of bins may be
        # overriden if unnecessarily large.
        data, from_bins = clamp(from_img.get_data(), from_bins,
                                mask=from_mask)
        self._from_img = make_xyz_image(data, xyz_affine(from_img), 'scanner')
        # Set field of view in the `from` image with potential
        # subsampling for faster similarity evaluation. This also sets
        # the _from_data and _vox_coords attributes
        if from_mask is None:
            self.subsample(npoints=NPOINTS)
        else:
            corner, size = smallest_bounding_box(from_mask)
            self.set_fov(corner=corner, size=size, npoints=NPOINTS)

        # Clamping of the `to` image including padding with -1
        data, to_bins = clamp(to_img.get_data(), to_bins,
                              mask=to_mask)
        self._to_data = -np.ones(np.array(to_img.shape) + 2, dtype=CLAMP_DTYPE)
        self._to_data[1:-1, 1:-1, 1:-1] = data
        self._to_inv_affine = inverse_affine(xyz_affine(to_img))

        # Joint histogram: must be double contiguous as it will be
        # passed to C routines which assume so
        self._joint_hist = np.zeros([from_bins, to_bins], dtype='double')

        # Set default registration parameters
        self._set_interp(interp)
        self._set_similarity((similarity, normalize, dist))

    def _get_interp(self):
        return interp_methods.keys()\
            [interp_methods.values().index(self._interp)]

    def _set_interp(self, interp):
        self._interp = interp_methods[interp]

    interp = property(_get_interp, _set_interp)

    def set_fov(self, spacing=None, corner=(0, 0, 0), size=None,
                npoints=None):
        """
        Defines a subset of the `from` image to restrict joint
        histogram computation.

        Parameters
        ----------
        spacing : sequence (3,) of positive integers
          Subsampling of image in voxels, where None (default) results
          in the subsampling to be automatically adjusted to roughly
          match a cubic grid with `npoints` voxels
        corner : sequence (3,) of positive integers
          Bounding box origin in voxel coordinates
        size : sequence (3,) of positive integers
          Desired bounding box size
        npoints : positive integer
          Desired number of voxels in the bounding box. If a `spacing`
          argument is provided, then `npoints` is ignored.
        """
        if spacing is None and npoints is None:
            spacing = [1, 1, 1]
        if size is None:
            size = self._from_img.shape
        slicer = lambda c, s, sp:\
            tuple([slice(c[i], s[i] + c[i], sp[i]) for i in range(3)])
        # Adjust spacing to match desired field of view size
        if not spacing is None:
            fov_data = self._from_img.get_data()[slicer(corner, size, spacing)]
        else:
            fov_data = self._from_img.get_data()[slicer(corner, size, [1, 1, 1])]
            spacing = ideal_spacing(fov_data, npoints=npoints)
            fov_data = self._from_img.get_data()[slicer(corner, size, spacing)]
        self._from_data = fov_data
        self._from_npoints = (fov_data >= 0).sum()
        self._from_affine = subgrid_affine(xyz_affine(self._from_img),
                                           slicer(corner, size, spacing))
        # We cache the voxel coordinates of the clamped image
        self._vox_coords =\
            np.indices(self._from_data.shape).transpose((1, 2, 3, 0))

    def subsample(self, spacing=None, npoints=None):
        self.set_fov(spacing=spacing, npoints=npoints)

    def _set_similarity(self, arg):
        normalize = True
        dist = None
        if not isinstance(arg, (tuple, list)):
            similarity = arg
        else:
            if len(arg) == 1:
                similarity = arg[0]
            elif len(arg) == 2:
                similarity, normalize = arg
            elif len(arg) == 3:
                similarity, normalize, dist = arg
            else:
                raise ValueError('cannot understand input argument')
        if similarity in _sms:
            self._similarity = similarity
            self._similarity_call =\
                _sms[similarity](self._joint_hist.shape, normalize, dist)
        else:
            if not hasattr(similarity, '__call__'):
                raise ValueError('similarity should be callable')
            self._similarity = 'custom'
            self._similarity_call = similarity

    def _get_similarity(self):
        return self._similarity

    similarity = property(_get_similarity, _set_similarity)

    def eval(self, T):
        """
        Evaluate similarity function given a world-to-world transform.

        Parameters
        ----------
        T : Transform
            Transform object implementing ``apply`` method
        """
        Tv = ChainTransform(T, pre=self._from_affine, post=self._to_inv_affine)
        return self._eval(Tv)

    def _eval(self, Tv):
        """
        Evaluate similarity function given a voxel-to-voxel transform.

        Parameters
        ----------
        Tv : Transform
             Transform object implementing ``apply`` method
             Should map voxel space to voxel space
        """
        # trans_vox_coords needs be C-contiguous
        trans_vox_coords = Tv.apply(self._vox_coords)
        interp = self._interp
        if self._interp < 0:
            interp = - np.random.randint(MAX_INT)
        _joint_histogram(self._joint_hist,
                         self._from_data.flat,  # array iterator
                         self._to_data,
                         trans_vox_coords,
                         interp)
        return self._similarity_call(self._joint_hist)

    def optimize(self, T, optimizer=OPTIMIZER, **kwargs):
        """ Optimize transform `T` with respect to similarity measure.

        The input object `T` will change as a result of the optimization.

        Parameters
        ----------
        T : object or str
          An object representing a transformation that should
          implement ``apply`` method and ``param`` attribute or
          property. If a string, one of 'rigid', 'similarity', or
          'affine'. The corresponding transformation class is then
          initialized by default.
        optimizer : str
          Name of optimization function (one of 'powell', 'steepest',
          'cg', 'bfgs', 'simplex')
        **kwargs : dict
          keyword arguments to pass to optimizer
        """
        # Replace T if a string is passed
        if T in affine_transforms:
            T = affine_transforms[T]()

        # Pull callback out of keyword arguments, if present
        callback = kwargs.pop('callback', None)

        # Create transform chain object with T generating params
        Tv = ChainTransform(T, pre=self._from_affine, post=self._to_inv_affine)
        tc0 = Tv.param

        # Cost function to minimize
        def cost(tc):
            # This is where the similarity function is calculcated
            Tv.param = tc
            return -self._eval(Tv)

        # Callback during optimization
        if callback is None and VERBOSE:

            def callback(tc):
                Tv.param = tc
                print(Tv.optimizable)
                print(str(self.similarity) + ' = %s' % self._eval(Tv))
                print('')

        # Switching to the appropriate optimizer
        if VERBOSE:
            print('Initial guess...')
            print(Tv.optimizable)

        kwargs.setdefault('xtol', XTOL)
        kwargs.setdefault('ftol', FTOL)
        kwargs.setdefault('gtol', GTOL)
        kwargs.setdefault('maxiter', MAXITER)
        kwargs.setdefault('maxfun', MAXFUN)

        fmin, args, kwargs = configure_optimizer(optimizer,
                                                 fprime=None,
                                                 fhess=None,
                                                 **kwargs)

        # Output
        if VERBOSE:
            print ('Optimizing using %s' % fmin.__name__)
        kwargs['callback'] = callback
        Tv.param = fmin(cost, tc0, *args, **kwargs)
        return Tv.optimizable

    def explore(self, T0, *args):
        """
        Evaluate the similarity at the transformations specified by
        sequences of parameter values.

        For instance:

        explore(T0, (0, [-1,0,1]), (4, [-2.,2]))
        """
        nparams = T0.param.size
        T0 = T0.copy()
        deltas = [[0] for i in range(nparams)]
        for a in args:
            deltas[a[0]] = a[1]
        grids = np.mgrid[[slice(0, len(d)) for d in deltas]]
        ntrials = np.prod(grids.shape[1:])
        Deltas = [np.asarray(deltas[i])[grids[i, :]].ravel()\
                      for i in range(nparams)]
        simis = np.zeros(ntrials)
        params = np.zeros([nparams, ntrials])

        Tv = ChainTransform(T0, pre=self._from_affine,
                            post=self._to_inv_affine)
        param0 = Tv.param
        for i in range(ntrials):
            param = param0 + np.array([D[i] for D in Deltas])
            Tv.param = param
            simis[i] = self._eval(Tv)
            params[:, i] = param

        return simis, params

def _clamp(x, y, bins):

    # Threshold
    dmaxmax = 2 ** (8 * y.dtype.itemsize - 1) - 1
    dmax = bins - 1  # default output maximum value
    if dmax > dmaxmax:
        raise ValueError('Excess number of bins')
    xmin = float(x.min())
    xmax = float(x.max())
    d = xmax - xmin

    """
    If the image dynamic is small, no need for compression: just
    downshift image values and re-estimate the dynamic range (hence
    xmax is translated to xmax-tth casted to the appropriate
    dtype. Otherwise, compress after downshifting image values (values
    equal to the threshold are reset to zero).
    """
    if issubclass(x.dtype.type, np.integer) and d <= dmax:
        y[:] = x - xmin
        bins = int(d) + 1
    else:
        a = dmax / d
        y[:] = np.round(a * (x - xmin))

    return y, bins

def clamp(x, bins, mask=None):
    """
    Clamp array values that fall within a given mask in the range
    [0..bins-1] and reset masked values to -1.

    Parameters
    ----------
    x : ndarray
      The input array

    bins : number
      Desired number of bins

    mask : ndarray, tuple or slice
      Anything such that x[mask] is an array.

    Returns
    -------
    y : ndarray
      Clamped array, masked items are assigned -1

    bins : number
      Adjusted number of bins

    """
    if bins > np.iinfo(np.short).max:
        raise ValueError('Too large a bin size')
    y = -np.ones(x.shape, dtype=CLAMP_DTYPE)
    if mask is None:
        y, bins = _clamp(x, y, bins)
    else:
        ym = y[mask]
        xm = x[mask]
        ym, bins = _clamp(xm, ym, bins)
        y[mask] = ym
    return y, bins

def ideal_spacing(data, npoints):
    """
    Tune spacing factors so that the number of voxels in the
    output block matches a given number.

    Parameters
    ----------
    data : ndarray or sequence
      Data image to subsample

    npoints : number
      Target number of voxels (negative values will be ignored)

    Returns
    -------
    spacing: ndarray
      Spacing factors

    """
    dims = data.shape
    actual_npoints = (data >= 0).sum()
    spacing = np.ones(3, dtype='uint')

    while actual_npoints > npoints:

        # Subsample the direction with the highest number of samples
        ddims = dims / spacing
        if ddims[0] >= ddims[1] and ddims[0] >= ddims[2]:
            dir = 0
        elif ddims[1] > ddims[0] and ddims[1] >= ddims[2]:
            dir = 1
        else:
            dir = 2
        spacing[dir] += 1
        subdata = data[::spacing[0], ::spacing[1], ::spacing[2]]
        actual_npoints = (subdata >= 0).sum()

    return spacing

def smallest_bounding_box(msk):
    """
    Extract the smallest bounding box from a mask
    """
    x, y, z = np.where(msk > 0)
    corner = [x.min(), y.min(), z.min()]
    size = [x.max() + 1, y.max() + 1, z.max() + 1]
    return corner, size
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""Tests for completerlib.

"""

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------

import os
import shutil
import sys
import tempfile
import unittest
from os.path import join

import nose.tools as nt

from IPython.core.completerlib import magic_run_completer, module_completion, try_import
from IPython.utils.tempdir import TemporaryDirectory
from IPython.testing.decorators import onlyif_unicode_paths

class MockEvent(object):
    def __init__(self, line):
        self.line = line

#-----------------------------------------------------------------------------
# Test functions begin
#-----------------------------------------------------------------------------
class Test_magic_run_completer(unittest.TestCase):
    files = [u"aao.py", u"a.py", u"b.py", u"aao.txt"]
    dirs = [u"adir/", "bdir/"]

    def setUp(self):
        self.BASETESTDIR = tempfile.mkdtemp()
        for fil in self.files:
            with open(join(self.BASETESTDIR, fil), "w") as sfile:
                sfile.write("pass\n")
        for d in self.dirs:
            os.mkdir(join(self.BASETESTDIR, d))

        self.oldpath = os.getcwd()
        os.chdir(self.BASETESTDIR)

    def tearDown(self):
        os.chdir(self.oldpath)
        shutil.rmtree(self.BASETESTDIR)

    def test_1(self):
        """Test magic_run_completer, should match two alternatives
        """
        event = MockEvent(u"%run a")
        mockself = None
        match = set(magic_run_completer(mockself, event))
        self.assertEqual(match, {u"a.py", u"aao.py", u"adir/"})

    def test_2(self):
        """Test magic_run_completer, should match one alternative
        """
        event = MockEvent(u"%run aa")
        mockself = None
        match = set(magic_run_completer(mockself, event))
        self.assertEqual(match, {u"aao.py"})

    def test_3(self):
        """Test magic_run_completer with unterminated " """
        event = MockEvent(u'%run "a')
        mockself = None
        match = set(magic_run_completer(mockself, event))
        self.assertEqual(match, {u"a.py", u"aao.py", u"adir/"})

    def test_completion_more_args(self):
        event = MockEvent(u'%run a.py ')
        match = set(magic_run_completer(None, event))
        self.assertEqual(match, set(self.files + self.dirs))

    def test_completion_in_dir(self):
        # Github issue #3459
        event = MockEvent(u'%run a.py {}'.format(join(self.BASETESTDIR, 'a')))
        print(repr(event.line))
        match = set(magic_run_completer(None, event))
        # We specifically use replace here rather than normpath, because
        # at one point there were duplicates 'adir' and 'adir/', and normpath
        # would hide the failure for that.
        self.assertEqual(match, {join(self.BASETESTDIR, f).replace('\\','/')
                            for f in (u'a.py', u'aao.py', u'aao.txt', u'adir/')})

class Test_magic_run_completer_nonascii(unittest.TestCase):
    @onlyif_unicode_paths
    def setUp(self):
        self.BASETESTDIR = tempfile.mkdtemp()
        for fil in [u"aaø.py", u"a.py", u"b.py"]:
            with open(join(self.BASETESTDIR, fil), "w") as sfile:
                sfile.write("pass\n")
        self.oldpath = os.getcwd()
        os.chdir(self.BASETESTDIR)

    def tearDown(self):
        os.chdir(self.oldpath)
        shutil.rmtree(self.BASETESTDIR)

    @onlyif_unicode_paths
    def test_1(self):
        """Test magic_run_completer, should match two alternatives
        """
        event = MockEvent(u"%run a")
        mockself = None
        match = set(magic_run_completer(mockself, event))
        self.assertEqual(match, {u"a.py", u"aaø.py"})

    @onlyif_unicode_paths
    def test_2(self):
        """Test magic_run_completer, should match one alternative
        """
        event = MockEvent(u"%run aa")
        mockself = None
        match = set(magic_run_completer(mockself, event))
        self.assertEqual(match, {u"aaø.py"})

    @onlyif_unicode_paths
    def test_3(self):
        """Test magic_run_completer with unterminated " """
        event = MockEvent(u'%run "a')
        mockself = None
        match = set(magic_run_completer(mockself, event))
        self.assertEqual(match, {u"a.py", u"aaø.py"})

# module_completer:

def test_import_invalid_module():
    """Testing of issue https://github.com/ipython/ipython/issues/1107"""
    invalid_module_names = {'foo-bar', 'foo:bar', '10foo'}
    valid_module_names = {'foobar'}
    with TemporaryDirectory() as tmpdir:
        sys.path.insert( 0, tmpdir )
        for name in invalid_module_names | valid_module_names:
            filename = os.path.join(tmpdir, name + '.py')
            open(filename, 'w').close()

        s = set( module_completion('import foo') )
        intersection = s.intersection(invalid_module_names)
        nt.assert_equal(intersection, set())

        assert valid_module_names.issubset(s), valid_module_names.intersection(s)

def test_bad_module_all():
    """Test module with invalid __all__

    https://github.com/ipython/ipython/issues/9678
    """
    testsdir = os.path.dirname(__file__)
    sys.path.insert(0, testsdir)
    try:
        results = module_completion('from bad_all import ')
        nt.assert_in('puppies', results)
        for r in results:
            nt.assert_is_instance(r, str)

        # bad_all doesn't contain submodules, but this completion
        # should finish without raising an exception:
        results = module_completion("import bad_all.")
        nt.assert_equal(results, [])
    finally:
        sys.path.remove(testsdir)

def test_module_without_init():
    """
    Test module without __init__.py.

    https://github.com/ipython/ipython/issues/11226
    """
    fake_module_name = "foo"
    with TemporaryDirectory() as tmpdir:
        sys.path.insert(0, tmpdir)
        try:
            os.makedirs(os.path.join(tmpdir, fake_module_name))
            s = try_import(mod=fake_module_name)
            assert s == []
        finally:
            sys.path.remove(tmpdir)

def test_valid_exported_submodules():
    """
    Test checking exported (__all__) objects are submodules
    """
    results = module_completion("import os.pa")
    # ensure we get a valid submodule:
    nt.assert_in("os.path", results)
    # ensure we don't get objects that aren't submodules:
    nt.assert_not_in("os.pathconf", results)
<commit_msg>[core][tests][completerlib] Remove nose<commit_after># -*- coding: utf-8 -*-
"""Tests for completerlib.

"""

#-----------------------------------------------------------------------------
# Imports
#-----------------------------------------------------------------------------

import os
import shutil
import sys
import tempfile
import unittest
from os.path import join

from IPython.core.completerlib import magic_run_completer, module_completion, try_import
from IPython.utils.tempdir import TemporaryDirectory
from IPython.testing.decorators import onlyif_unicode_paths

class MockEvent(object):
    def __init__(self, line):
        self.line = line

#-----------------------------------------------------------------------------
# Test functions begin
#-----------------------------------------------------------------------------
class Test_magic_run_completer(unittest.TestCase):
    files = [u"aao.py", u"a.py", u"b.py", u"aao.txt"]
    dirs = [u"adir/", "bdir/"]

    def setUp(self):
        self.BASETESTDIR = tempfile.mkdtemp()
        for fil in self.files:
            with open(join(self.BASETESTDIR, fil), "w") as sfile:
                sfile.write("pass\n")
        for d in self.dirs:
            os.mkdir(join(self.BASETESTDIR, d))

        self.oldpath = os.getcwd()
        os.chdir(self.BASETESTDIR)

    def tearDown(self):
        os.chdir(self.oldpath)
        shutil.rmtree(self.BASETESTDIR)

    def test_1(self):
        """Test magic_run_completer, should match two alternatives
        """
        event = MockEvent(u"%run a")
        mockself = None
        match = set(magic_run_completer(mockself, event))
        self.assertEqual(match, {u"a.py", u"aao.py", u"adir/"})

    def test_2(self):
        """Test magic_run_completer, should match one alternative
        """
        event = MockEvent(u"%run aa")
        mockself = None
        match = set(magic_run_completer(mockself, event))
        self.assertEqual(match, {u"aao.py"})

    def test_3(self):
        """Test magic_run_completer with unterminated " """
        event = MockEvent(u'%run "a')
        mockself = None
        match = set(magic_run_completer(mockself, event))
        self.assertEqual(match, {u"a.py", u"aao.py", u"adir/"})

    def test_completion_more_args(self):
        event = MockEvent(u'%run a.py ')
        match = set(magic_run_completer(None, event))
        self.assertEqual(match, set(self.files + self.dirs))

    def test_completion_in_dir(self):
        # Github issue #3459
        event = MockEvent(u'%run a.py {}'.format(join(self.BASETESTDIR, 'a')))
        print(repr(event.line))
        match = set(magic_run_completer(None, event))
        # We specifically use replace here rather than normpath, because
        # at one point there were duplicates 'adir' and 'adir/', and normpath
        # would hide the failure for that.
        self.assertEqual(match, {join(self.BASETESTDIR, f).replace('\\','/')
                            for f in (u'a.py', u'aao.py', u'aao.txt', u'adir/')})

class Test_magic_run_completer_nonascii(unittest.TestCase):
    @onlyif_unicode_paths
    def setUp(self):
        self.BASETESTDIR = tempfile.mkdtemp()
        for fil in [u"aaø.py", u"a.py", u"b.py"]:
            with open(join(self.BASETESTDIR, fil), "w") as sfile:
                sfile.write("pass\n")
        self.oldpath = os.getcwd()
        os.chdir(self.BASETESTDIR)

    def tearDown(self):
        os.chdir(self.oldpath)
        shutil.rmtree(self.BASETESTDIR)

    @onlyif_unicode_paths
    def test_1(self):
        """Test magic_run_completer, should match two alternatives
        """
        event = MockEvent(u"%run a")
        mockself = None
        match = set(magic_run_completer(mockself, event))
        self.assertEqual(match, {u"a.py", u"aaø.py"})

    @onlyif_unicode_paths
    def test_2(self):
        """Test magic_run_completer, should match one alternative
        """
        event = MockEvent(u"%run aa")
        mockself = None
        match = set(magic_run_completer(mockself, event))
        self.assertEqual(match, {u"aaø.py"})

    @onlyif_unicode_paths
    def test_3(self):
        """Test magic_run_completer with unterminated " """
        event = MockEvent(u'%run "a')
        mockself = None
        match = set(magic_run_completer(mockself, event))
        self.assertEqual(match, {u"a.py", u"aaø.py"})

# module_completer:

def test_import_invalid_module():
    """Testing of issue https://github.com/ipython/ipython/issues/1107"""
    invalid_module_names = {'foo-bar', 'foo:bar', '10foo'}
    valid_module_names = {'foobar'}
    with TemporaryDirectory() as tmpdir:
        sys.path.insert( 0, tmpdir )
        for name in invalid_module_names | valid_module_names:
            filename = os.path.join(tmpdir, name + '.py')
            open(filename, 'w').close()

        s = set( module_completion('import foo') )
        intersection = s.intersection(invalid_module_names)
        assert intersection == set()

        assert valid_module_names.issubset(s), valid_module_names.intersection(s)

def test_bad_module_all():
    """Test module with invalid __all__

    https://github.com/ipython/ipython/issues/9678
    """
    testsdir = os.path.dirname(__file__)
    sys.path.insert(0, testsdir)
    try:
        results = module_completion("from bad_all import ")
        assert "puppies" in results
        for r in results:
            assert isinstance(r, str)

        # bad_all doesn't contain submodules, but this completion
        # should finish without raising an exception:
        results = module_completion("import bad_all.")
        assert results == []
    finally:
        sys.path.remove(testsdir)

def test_module_without_init():
    """
    Test module without __init__.py.

    https://github.com/ipython/ipython/issues/11226
    """
    fake_module_name = "foo"
    with TemporaryDirectory() as tmpdir:
        sys.path.insert(0, tmpdir)
        try:
            os.makedirs(os.path.join(tmpdir, fake_module_name))
            s = try_import(mod=fake_module_name)
            assert s == []
        finally:
            sys.path.remove(tmpdir)

def test_valid_exported_submodules():
    """
    Test checking exported (__all__) objects are submodules
    """
    results = module_completion("import os.pa")
    # ensure we get a valid submodule:
    assert "os.path" in results
    # ensure we don't get objects that aren't submodules:
    assert "os.pathconf" not in results
<|endoftext|>


<commit_before>#!/usr/bin/env python
# -*- encoding: utf-8 -*-

# Copyright 2011-2016, Nigel Small
#
# 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 py2neo.database import cypher_escape
from py2neo.types import Node, remote
from py2neo.util import label_case, relationship_case

class Property(object):

    def __init__(self, key=None):
        self.key = key

    def __get__(self, instance, owner):
        return instance.__db_node__[self.key]

    def __set__(self, instance, value):
        instance.__db_node__[self.key] = value

class Label(object):

    def __init__(self, name=None):
        self.name = name

    def __get__(self, instance, owner):
        return instance.__db_node__.has_label(self.name)

    def __set__(self, instance, value):
        if value:
            instance.__db_node__.add_label(self.name)
        else:
            instance.__db_node__.remove_label(self.name)

class Related(object):

    __instance__ = None

    def __init__(self, type=None):
        self.type = type

    @property
    def items(self):
        return self.__instance__.__db_relationships__.setdefault(self.type, set())

    def add(self, value):
        if isinstance(value, GraphObject):
            self.items.add(value)
        else:
            raise TypeError("Related items must be GraphObject instances")

    def remove(self, value):
        if isinstance(value, GraphObject):
            self.items.remove(value)
        else:
            raise TypeError("Related items must be GraphObject instances")

class GraphObjectMeta(type):

    def __new__(mcs, name, bases, attributes):
        related = []
        for attr_name, attr in list(attributes.items()):
            if isinstance(attr, Property):
                attributes[attr_name] = Property(attr.key or attr_name)
            elif isinstance(attr, Label):
                attributes[attr_name] = Label(attr.name or label_case(attr_name))
            elif isinstance(attr, Related):
                attributes[attr_name] = Related(attr.type or relationship_case(attr_name))
                related.append(attributes[attr_name])
        attributes.setdefault("__primary_label__", name)
        attributes.setdefault("__primary_key__", "__id__")
        instance = super().__new__(mcs, name, bases, attributes)
        for attr in related:
            attr.__instance__ = instance
        return instance

class GraphObject(metaclass=GraphObjectMeta):
    __primary_label__ = None
    __primary_key__ = None

    __db_node = None
    __db_relationships = None

    @classmethod
    def load(cls, graph, primary_value):
        # Label:key=value
        primary_key = cls.__primary_key__
        if primary_key == "__id__":
            node = graph.evaluate("MATCH (a:%s) WHERE id(a)={x}" %
                                  cypher_escape(cls.__primary_label__), x=primary_value)
        else:
            node = graph.find_one(cls.__primary_label__, primary_key, primary_value)
        if node is None:
            raise LookupError("Cannot load object")
        inst = GraphObject()
        inst.__db_node = node
        print(list(graph.match(node)))
        # TODO: load relationships
        inst.__class__ = cls
        return inst

    @property
    def __primary_value__(self):
        node = self.__db_node__
        primary_key = self.__primary_key__
        if primary_key == "__id__":
            remote_node = remote(node)
            if remote_node:
                return remote_node._id
            else:
                return None
        else:
            return node[primary_key]

    @property
    def __db_node__(self):
        if self.__db_node is None:
            self.__db_node = Node(self.__primary_label__)
        return self.__db_node

    @property
    def __db_relationships__(self):
        if self.__db_relationships is None:
            self.__db_relationships = {}
        return self.__db_relationships

    @property
    def __remote__(self):
        return self.__db_node__.__remote__

    def __db_create__(self, tx):
        tx.merge(self.__db_node__, self.__primary_label__, self.__primary_key__)

    def __db_delete__(self, tx):
        # TODO: delete if not bound
        tx.delete(self.__db_node__)

    def __db_pull__(self, graph):
        graph.pull(self.__db_node__)

    def __db_push__(self, graph):
        graph.push(self.__db_node__)

class Person(GraphObject):
    __primary_key__ = "name"

    name = Property()
    year_of_birth = Property(key="born")

    acted_in = Related("Movie")
    directed = Related("Movie")
    produced = Related("Movie")

    def __init__(self, name, year_of_birth):
        self.name = name
        self.year_of_birth = year_of_birth

class Movie(GraphObject):
    __primary_key__ = "title"

    title = Property()
    year_of_release = Property(key="released")

    def __init__(self, title, year_of_release):
        self.title = title
        self.year_of_release = year_of_release

def main():
    keanu = Person("Keanu Reeves", 1968)
    the_matrix = Movie("The Matrix", 1999)
    keanu.acted_in.add(the_matrix)
<commit_msg>Working but complicated<commit_after>#!/usr/bin/env python
# -*- encoding: utf-8 -*-

# Copyright 2011-2016, Nigel Small
#
# 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 py2neo.database import cypher_escape
from py2neo.types import Node, remote
from py2neo.util import label_case, relationship_case

class Property(object):

    def __init__(self, key=None):
        self.key = key

    def __get__(self, instance, owner):
        return instance.__db_node__[self.key]

    def __set__(self, instance, value):
        instance.__db_node__[self.key] = value

class Label(object):

    def __init__(self, name=None):
        self.name = name

    def __get__(self, instance, owner):
        return instance.__db_node__.has_label(self.name)

    def __set__(self, instance, value):
        if value:
            instance.__db_node__.add_label(self.name)
        else:
            instance.__db_node__.remove_label(self.name)

class Related(object):

    __instance__ = None

    def __init__(self, related_node_type, relationship_type=None):
        self.related_node_type = related_node_type
        self.relationship_type = relationship_type

    def __get__(self, instance, owner):
        return instance.__db_relationships__.setdefault(self.relationship_type, set())

    def add(self, item):
        pass

    def remove(self, item):
        pass

class RelatedItems(object):

    pass

class GraphObjectMeta(type):

    def __new__(mcs, name, bases, attributes):
        for attr_name, attr in list(attributes.items()):
            if isinstance(attr, Property):
                attributes[attr_name] = Property(attr.key or attr_name)
            elif isinstance(attr, Label):
                attributes[attr_name] = Label(attr.name or label_case(attr_name))
            elif isinstance(attr, Related):
                relationship_type = attr.relationship_type or relationship_case(attr_name)
                attributes[attr_name] = Related(attr.related_node_type, relationship_type)
        attributes.setdefault("__primary_label__", name)
        attributes.setdefault("__primary_key__", "__id__")
        return super().__new__(mcs, name, bases, attributes)

class GraphObject(metaclass=GraphObjectMeta):
    __graph__ = None
    __primary_label__ = None
    __primary_key__ = None

    __db_node = None
    __db_relationships = None

    @classmethod
    def load(cls, primary_value):
        graph = cls.__graph__
        # Label:key=value
        primary_key = cls.__primary_key__
        if primary_key == "__id__":
            node = graph.evaluate("MATCH (a:%s) WHERE id(a)={x}" %
                                  cypher_escape(cls.__primary_label__), x=primary_value)
        else:
            node = graph.find_one(cls.__primary_label__, primary_key, primary_value)
        if node is None:
            raise LookupError("Cannot load object")
        return cls.wrap(node)

    @classmethod
    def wrap(cls, node):
        graph = cls.__graph__
        inst = GraphObject()
        inst.__db_node = node
        inst.__db_relationships = {}
        for relationship in graph.match(node):
            related_nodes = inst.__db_relationships.setdefault(relationship.type(), set())
            related_nodes.add(relationship.end_node())
        inst.__class__ = cls
        return inst

    def __eq__(self, other):
        try:
            return self.__db_node__ == other.__db_node__
        except AttributeError:
            return False

    def __ne__(self, other):
        return not self.__eq__(other)

    def __hash__(self):
        return hash(self.__db_node__)

    @property
    def __db_node__(self):
        if self.__db_node is None:
            self.__db_node = Node(self.__primary_label__)
        return self.__db_node

    @property
    def __db_relationships__(self):
        if self.__db_relationships is None:
            self.__db_relationships = {}
        return self.__db_relationships

    @property
    def __primary_value__(self):
        node = self.__db_node__
        primary_key = self.__primary_key__
        if primary_key == "__id__":
            remote_node = remote(node)
            if remote_node:
                return remote_node._id
            else:
                return None
        else:
            return node[primary_key]

    @property
    def __remote__(self):
        return self.__db_node__.__remote__

    def __db_create__(self, tx):
        tx.merge(self.__db_node__, self.__primary_label__, self.__primary_key__)

    def __db_delete__(self, tx):
        # TODO: delete if not bound
        tx.delete(self.__db_node__)

    def __db_pull__(self, graph):
        graph.pull(self.__db_node__)

    def __db_push__(self, graph):
        graph.push(self.__db_node__)

class Person(GraphObject):
    __primary_key__ = "name"

    name = Property()
    year_of_birth = Property(key="born")

    acted_in = Related("Movie")
    directed = Related("Movie")
    produced = Related("Movie")

class Movie(GraphObject):
    __primary_key__ = "title"

    title = Property()
    tag_line = Property(key="tagline")
    year_of_release = Property(key="released")

def main():
    from py2neo import Graph
    GraphObject.__graph__ = Graph()
    keanu = Person.load("Keanu Reeves")
    the_matrix = Movie.load("The Matrix")
    keanu.acted_in.add(the_matrix)
<|endoftext|>