<commit_before>from typing import Optional, Tuple

from silx.utils.proxy import docstring
from silx.gui import qt
from silx.gui.plot import items, PlotWindow

import os
import sys

sys.path.insert(0, os.path.dirname(__file__))
import silx_monkey_patch

# TODO pythonic slice accessor?, asynchronous
class LevelOfDetailImageData:
    """Multi-resolution image data accessor

     Provides access to a pyramid of images:
     Level of detail 0 gives access to the full resolution image.
     Level of detail 1 to data binned 2x2, and so on

    :param List[numpy.ndarray] data: Pyramid of image data.
    """

    def __init__(self, data):
        self.__lods = data

    def getMaxLevel(self) -> Optional[int]:
        """Returns the maximum available level of detail.

        :rtype: Union[int,None]
        """
        return len(self.__lods) - 1 if self.__lods else None

    def __checkLevel(self, level: int) -> None:
        """Check if provided level of detail is valid

        :param int level: Level of detail
        :raises ValueError: if requested level of detail is not available
        """
        maxLevel = self.getMaxLevel()
        if maxLevel is None or not 0 <= level <= maxLevel:
            raise ValueError("Requested level of detail is not available")

    def getData(self, level: int = 0, copy: bool = True):
        """Returns data array for selected level of detail

        :param int level: Level of detail
        :param bool copy: True to copy, False to get internal representation.
        :rtype: Union[numpy.ndarray,h5py.Dataset]
        :raises ValueError: if requested level of detail is not available
        """
        self.__checkLevel(level)
        data = self.__lods[level]
        return numpy.array(data, copy=True) if copy else data

    def __convertLevel0Slice(self, slice_, level) -> Tuple[slice]:
        """Convert Level of detail 0 slicing to another level.

        Conversion is done so that initial array subset is fully
        available in sub-sampled slicing.
        Converted slicing can cover a larger extent than the initial one,
        thus corresponding slicing in level 0 is also returned.

        :param slice_: Slice in level of detail 0
        :param level: The level of detail to convert slicing to
        :return: Converted slice, Corresponding level 0 slice
        :rtype: List[slice]
        :raises ValueError: if requested level of detail is not available
        """
        self.__checkLevel(level)
        if slice_.step not in (1, None):
            raise NotImplementedError("slice step must be 1 or None")

        factor = 2 ** level
        if slice_.start is None:
            start = None
        else:
            start = slice_.start // factor
        if slice_.stop is None:
            stop = None
        else:
            stop = int(numpy.ceil(slice_.stop / factor))
        return slice(start, stop), slice(
            None if start is None else start * factor,
            None if stop is None else stop * factor,
        )

    def getDataForLevel(
        self, rows: slice, cols: slice, level: int = 0, copy: bool = True
    ):
        """Returns data from the requested level and level 0 slicing.

        :param slice rows: Requested level of detail 0 slicing
        :param slice cols: Requested level of detail 0 slicing
        :param int level: Level of detail to convert slice to
        :return: Data at given level of detail for given slicing and
           corresponding level 0 slicing.
        :raises ValueError: if requested level of detail is not available
        """
        self.__checkLevel(level)
        lodrows, rows = self.__convertLevel0Slice(rows, level)
        lodcols, cols = self.__convertLevel0Slice(cols, level)
        data = self.__lods[level]
        return numpy.array(data[lodrows, lodcols], copy=copy), (rows, cols)

class Image(items.ImageData):

    sigVisibleSlicesChanged = qt.Signal()
    """Signal emitted when the visible slices of the array has changed."""

    def __init__(self):
        self.__loddata = LevelOfDetailImageData(
            [numpy.array((0, 0), dtype=numpy.float32)]
        )
        self.__previousLevelOfDetail = -1
        self.__previousVisibleSlices = slice(0), slice(0)
        self.__chunkShape = 1, 1
        super().__init__()
        self._sigVisibleBoundsChanged.connect(self.__visibleBoundsChanged)
        self._setVisibleBoundsTracking(True)

    @docstring(items.ImageData)
    def _setPlot(self, plot):
        previousPlot = self.getPlot()
        if previousPlot is not None:
            for axis in (plot.getXAxis(), plot.getYAxis()):
                axis.sigLimitsChanged.disconnect(self.__plotLimitsChanged)

        self.__previousLevelOfDetail = -1
        super()._setPlot(plot)

        if plot is not None:
            for axis in (plot.getXAxis(), plot.getYAxis()):
                axis.sigLimitsChanged.connect(self.__plotLimitsChanged)

    def __plotLimitsChanged(self, *args) -> None:
        level = self._getLevelOfDetail()
        if level != self.__previousLevelOfDetail:
            self.__previousLevelOfDetail = level
            self._updated()  # Dirty the

    def _getLevelOfDetail(self) -> int:
        """Returns current level of detail (0 for full resolution)"""
        plot = self.getPlot()
        if plot is None:
            return 0

        sx, sy = self.getScale()
        width, height = plot.getPixelSizeInData()
        elemPerPixel = max(width / sx, height / sy)
        level = int(numpy.ceil(numpy.log2(elemPerPixel))) if elemPerPixel > 1.0 else 0
        return numpy.clip(level, 0, self.__loddata.getMaxLevel())

    def getData(self, copy=True):
        """Returns the image data

        :param bool copy: True (Default) to get a copy,
                          False to use internal representation (do not modify!)
        :rtype: Union[numpy.ndarray,h5py.Dataset]
        """
        return self.__loddata.getData(copy=copy)

    def setData(self, data):
        """Set the level of detail image data

        :param List[numpy.ndarray] data:
        """
        self.__loddata = LevelOfDetailImageData(data)
        self._visibleBoundsChanged()
        self._updated(items.ItemChangedType.DATA)

    @docstring(items.ImageData)
    def _setVisibleBoundsTracking(self, enable: bool) -> None:
        if not enable:
            raise RuntimeError(
                "This item does not support disabling visible bounds tracking"
            )
        super()._setVisibleBoundsTracking(enable)

    def _addBackendRenderer(self, backend):
        """Update backend renderer"""
        plot = self.getPlot()
        assert plot is not None
        if not self._isPlotLinear(plot):
            # Do not render with non linear scales
            return None

        colormap = self.getColormap()
        if colormap.isAutoscale():
            # Avoid backend to compute autoscale: use item cache
            colormap = colormap.copy()
            colormap.setVRange(*colormap.getColormapRange(self))

        # Compute sub image and offset
        rowSlice, colSlice = self.getVisibleChunkSlices()
        emptySlice = slice(0)
        if rowSlice == emptySlice or colSlice == emptySlice:
            return None  # No data to display

        sx, sy = self.getScale()
        width, height = self.getPlot().getPixelSizeInData()
        elemPerPixel = max(width / sx, height / sy)
        level = int(numpy.ceil(numpy.log2(elemPerPixel))) if elemPerPixel > 1.0 else 0
        level = numpy.clip(level, 0, self.__loddata.getMaxLevel())
        dataToUse, (rowSlice, colSlice) = self.__loddata.getDataForLevel(
            rowSlice, colSlice, level, copy=False
        )
        if dataToUse.size == 0:
            return None  # No data to display

        # Offset origin to sub image
        ox, oy = self.getOrigin()
        origin = ox + sx * colSlice.start, oy + sy * rowSlice.start

        return backend.addImage(
            dataToUse,
            origin=origin,
            scale=(sx * 2**level, sy * 2**level),
            colormap=colormap if dataToUse.ndim == 2 else None,
            alpha=self.getAlpha(),
        )

    def __visibleBoundsChanged(self):
        """Emit sigVisibleSlicesChanged when slicing has changed."""
        slices = self.getVisibleChunkSlices()
        if slices != self.__previousVisibleSlices:
            self.__previousVisibleSlices = slices
            self._updated(checkVisibility=False)
            self.sigVisibleSlicesChanged.emit()

    def getVisibleSlices(self):
        """Returns the array slicing of the image part inside the plot area.

        This is inclusive in that partly visible array elements are included.

        :returns: (dim0 slice, dim1 slice)
        :rtype: List[slice]
        """
        bounds = self.getVisibleBounds()
        if bounds is None:
            return slice(0), slice(0)  # Empty slicing

        xmin, xmax, ymin, ymax = bounds

        ox, oy = self.getOrigin()
        sx, sy = self.getScale()

        return (
            slice(int((ymin - oy) / sy), int(numpy.ceil((ymax - oy) / sy))),
            slice(int((xmin - ox) / sx), int(numpy.ceil((xmax - ox) / sx))),
        )

    def getChunkShape(self) -> Tuple[float]:
        """Returns current chunk shape (rows, columns).

        :rtype: List[float]
        """
        return self.__chunkShape

    def setChunkShape(self, shape: Tuple[float]):
        """Set the chunk shape (rows, columns).

        :param List[float] shape:
        """
        if shape != self.__chunkShape:
            self.__chunkShape = shape
            self.__visibleBoundsChanged()

    def getVisibleChunkSlices(self):
        """Returns the array slicing of the image aligned to chunks.

        This is inclusive in that partly visible array elements are included.

        :returns: (dim0 slice, dim1 slice)
        :rtype: List[slice]
        """
        yslice, xslice = self.getVisibleSlices()

        emptySlice = slice(0)
        if yslice == emptySlice or xslice == emptySlice:
            return emptySlice, emptySlice  # Nothing to display

        height, width = self.getChunkShape()
        ystart, ystop = yslice.start, yslice.stop
        xstart, xstop = xslice.start, xslice.stop
        return (
            slice(
                height * (ystart // height), height * int(numpy.ceil(ystop / height))
            ),
            slice(width * (xstart // width), width * int(numpy.ceil((xstop / width)))),
        )

if __name__ == "__main__":
    from silx.gui import qt
    from silx.gui.plot.AlphaSlider import ActiveImageAlphaSlider
    import numpy
    import h5py
    import sys

    filename = sys.argv[1]
    bg_filename = sys.argv[2] if len(sys.argv) == 3 else None

    qt.QApplication.setAttribute(qt.Qt.AA_ShareOpenGLContexts, True)
    qt.QApplication.setAttribute(qt.Qt.AA_EnableHighDpiScaling, True)

    app = qt.QApplication([])
    plot = PlotWindow(backend="gl")
    #w.getYAxis().setInverted(True)
    #w.setAxesDisplayed(False)
    plot.setKeepDataAspectRatio(True)
    item = Image()
    colormap = plot.getDefaultColormap()
    colormap.setVRange(-0.1, 0.3)
    item.setColormap(colormap)

    f = h5py.File(filename, mode="r")
    lods = [f['level%d' % i] for i in range(8) if ('level%d' % i) in f.keys()]
    item.setData(lods)

    item.setChunkShape((256, 256))
    plot.addItem(item)

    if bg_filename is not None:
        bgfile = h5py.File(bg_filename, mode="r")
        lods = [bgfile['level%d' % i] for i in range(8) if ('level%d' % i) in bgfile.keys()]
        bgitem = Image()
        bgitem.setName("Background")
        bgitem.setZValue(-1)
        bgitem.setData(lods)
        bgitem.setChunkShape((256, 256))
        plot.addItem(bgitem)

    plot.setActiveImage(item.getName())
    plot.resetZoom()

    alphaSlider = ActiveImageAlphaSlider(plot=plot)
    alphaSlider.show()

    window = qt.QMainWindow()
    window.setCentralWidget(plot)
    alphaDock = qt.QDockWidget()
    alphaDock.setWidget(alphaSlider)
    window.addDockWidget(qt.Qt.LeftDockWidgetArea, alphaDock)
    window.show()
    app.exec_()
<commit_msg>add TODOs<commit_after>from typing import Optional, Tuple

from silx.utils.proxy import docstring
from silx.gui import qt
from silx.gui.plot import items, PlotWindow

import os
import sys

sys.path.insert(0, os.path.dirname(__file__))
import silx_monkey_patch

# TODO for live update of the HDF5 file:
# - make a h5py LevelOfDetailImageData that closes the file each time + no file locking
# - update at regular interval

# TODO pythonic slice accessor?, asynchronous
class LevelOfDetailImageData:
    """Multi-resolution image data accessor

     Provides access to a pyramid of images:
     Level of detail 0 gives access to the full resolution image.
     Level of detail 1 to data binned 2x2, and so on

    :param List[numpy.ndarray] data: Pyramid of image data.
    """

    def __init__(self, data):
        self.__lods = data

    def getMaxLevel(self) -> Optional[int]:
        """Returns the maximum available level of detail.

        :rtype: Union[int,None]
        """
        return len(self.__lods) - 1 if self.__lods else None

    def __checkLevel(self, level: int) -> None:
        """Check if provided level of detail is valid

        :param int level: Level of detail
        :raises ValueError: if requested level of detail is not available
        """
        maxLevel = self.getMaxLevel()
        if maxLevel is None or not 0 <= level <= maxLevel:
            raise ValueError("Requested level of detail is not available")

    def getData(self, level: int = 0, copy: bool = True):
        """Returns data array for selected level of detail

        :param int level: Level of detail
        :param bool copy: True to copy, False to get internal representation.
        :rtype: Union[numpy.ndarray,h5py.Dataset]
        :raises ValueError: if requested level of detail is not available
        """
        self.__checkLevel(level)
        data = self.__lods[level]
        return numpy.array(data, copy=True) if copy else data

    def __convertLevel0Slice(self, slice_, level) -> Tuple[slice]:
        """Convert Level of detail 0 slicing to another level.

        Conversion is done so that initial array subset is fully
        available in sub-sampled slicing.
        Converted slicing can cover a larger extent than the initial one,
        thus corresponding slicing in level 0 is also returned.

        :param slice_: Slice in level of detail 0
        :param level: The level of detail to convert slicing to
        :return: Converted slice, Corresponding level 0 slice
        :rtype: List[slice]
        :raises ValueError: if requested level of detail is not available
        """
        self.__checkLevel(level)
        if slice_.step not in (1, None):
            raise NotImplementedError("slice step must be 1 or None")

        factor = 2 ** level
        if slice_.start is None:
            start = None
        else:
            start = slice_.start // factor
        if slice_.stop is None:
            stop = None
        else:
            stop = int(numpy.ceil(slice_.stop / factor))
        return slice(start, stop), slice(
            None if start is None else start * factor,
            None if stop is None else stop * factor,
        )

    def getDataForLevel(
        self, rows: slice, cols: slice, level: int = 0, copy: bool = True
    ):
        """Returns data from the requested level and level 0 slicing.

        :param slice rows: Requested level of detail 0 slicing
        :param slice cols: Requested level of detail 0 slicing
        :param int level: Level of detail to convert slice to
        :return: Data at given level of detail for given slicing and
           corresponding level 0 slicing.
        :raises ValueError: if requested level of detail is not available
        """
        self.__checkLevel(level)
        lodrows, rows = self.__convertLevel0Slice(rows, level)
        lodcols, cols = self.__convertLevel0Slice(cols, level)
        data = self.__lods[level]
        return numpy.array(data[lodrows, lodcols], copy=copy), (rows, cols)

class Image(items.ImageData):

    sigVisibleSlicesChanged = qt.Signal()
    """Signal emitted when the visible slices of the array has changed."""

    def __init__(self):
        self.__loddata = LevelOfDetailImageData(
            [numpy.array((0, 0), dtype=numpy.float32)]
        )
        self.__previousLevelOfDetail = -1
        self.__previousVisibleSlices = slice(0), slice(0)
        self.__chunkShape = 1, 1
        super().__init__()
        self._sigVisibleBoundsChanged.connect(self.__visibleBoundsChanged)
        self._setVisibleBoundsTracking(True)

    @docstring(items.ImageData)
    def _setPlot(self, plot):
        previousPlot = self.getPlot()
        if previousPlot is not None:
            for axis in (plot.getXAxis(), plot.getYAxis()):
                axis.sigLimitsChanged.disconnect(self.__plotLimitsChanged)

        self.__previousLevelOfDetail = -1
        super()._setPlot(plot)

        if plot is not None:
            for axis in (plot.getXAxis(), plot.getYAxis()):
                axis.sigLimitsChanged.connect(self.__plotLimitsChanged)

    def __plotLimitsChanged(self, *args) -> None:
        level = self._getLevelOfDetail()
        if level != self.__previousLevelOfDetail:
            self.__previousLevelOfDetail = level
            self._updated()  # Dirty the

    def _getLevelOfDetail(self) -> int:
        """Returns current level of detail (0 for full resolution)"""
        plot = self.getPlot()
        if plot is None:
            return 0

        sx, sy = self.getScale()
        width, height = plot.getPixelSizeInData()
        elemPerPixel = max(width / sx, height / sy)
        level = int(numpy.ceil(numpy.log2(elemPerPixel))) if elemPerPixel > 1.0 else 0
        return numpy.clip(level, 0, self.__loddata.getMaxLevel())

    def getData(self, copy=True):
        """Returns the image data

        :param bool copy: True (Default) to get a copy,
                          False to use internal representation (do not modify!)
        :rtype: Union[numpy.ndarray,h5py.Dataset]
        """
        return self.__loddata.getData(copy=copy)

    def setData(self, data):
        """Set the level of detail image data

        :param List[numpy.ndarray] data:
        """
        self.__loddata = LevelOfDetailImageData(data)
        self._visibleBoundsChanged()
        self._updated(items.ItemChangedType.DATA)

    @docstring(items.ImageData)
    def _setVisibleBoundsTracking(self, enable: bool) -> None:
        if not enable:
            raise RuntimeError(
                "This item does not support disabling visible bounds tracking"
            )
        super()._setVisibleBoundsTracking(enable)

    def _addBackendRenderer(self, backend):
        """Update backend renderer"""
        plot = self.getPlot()
        assert plot is not None
        if not self._isPlotLinear(plot):
            # Do not render with non linear scales
            return None

        colormap = self.getColormap()
        if colormap.isAutoscale():
            # Avoid backend to compute autoscale: use item cache
            colormap = colormap.copy()
            colormap.setVRange(*colormap.getColormapRange(self))

        # Compute sub image and offset
        rowSlice, colSlice = self.getVisibleChunkSlices()
        emptySlice = slice(0)
        if rowSlice == emptySlice or colSlice == emptySlice:
            return None  # No data to display

        sx, sy = self.getScale()
        width, height = self.getPlot().getPixelSizeInData()
        elemPerPixel = max(width / sx, height / sy)
        level = int(numpy.ceil(numpy.log2(elemPerPixel))) if elemPerPixel > 1.0 else 0
        level = numpy.clip(level, 0, self.__loddata.getMaxLevel())
        dataToUse, (rowSlice, colSlice) = self.__loddata.getDataForLevel(
            rowSlice, colSlice, level, copy=False
        )
        if dataToUse.size == 0:
            return None  # No data to display

        # Offset origin to sub image
        ox, oy = self.getOrigin()
        origin = ox + sx * colSlice.start, oy + sy * rowSlice.start

        return backend.addImage(
            dataToUse,
            origin=origin,
            scale=(sx * 2**level, sy * 2**level),
            colormap=colormap if dataToUse.ndim == 2 else None,
            alpha=self.getAlpha(),
        )

    def __visibleBoundsChanged(self):
        """Emit sigVisibleSlicesChanged when slicing has changed."""
        slices = self.getVisibleChunkSlices()
        if slices != self.__previousVisibleSlices:
            self.__previousVisibleSlices = slices
            self._updated(checkVisibility=False)
            self.sigVisibleSlicesChanged.emit()

    def getVisibleSlices(self):
        """Returns the array slicing of the image part inside the plot area.

        This is inclusive in that partly visible array elements are included.

        :returns: (dim0 slice, dim1 slice)
        :rtype: List[slice]
        """
        bounds = self.getVisibleBounds()
        if bounds is None:
            return slice(0), slice(0)  # Empty slicing

        xmin, xmax, ymin, ymax = bounds

        ox, oy = self.getOrigin()
        sx, sy = self.getScale()

        return (
            slice(int((ymin - oy) / sy), int(numpy.ceil((ymax - oy) / sy))),
            slice(int((xmin - ox) / sx), int(numpy.ceil((xmax - ox) / sx))),
        )

    def getChunkShape(self) -> Tuple[float]:
        """Returns current chunk shape (rows, columns).

        :rtype: List[float]
        """
        return self.__chunkShape

    def setChunkShape(self, shape: Tuple[float]):
        """Set the chunk shape (rows, columns).

        :param List[float] shape:
        """
        if shape != self.__chunkShape:
            self.__chunkShape = shape
            self.__visibleBoundsChanged()

    def getVisibleChunkSlices(self):
        """Returns the array slicing of the image aligned to chunks.

        This is inclusive in that partly visible array elements are included.

        :returns: (dim0 slice, dim1 slice)
        :rtype: List[slice]
        """
        yslice, xslice = self.getVisibleSlices()

        emptySlice = slice(0)
        if yslice == emptySlice or xslice == emptySlice:
            return emptySlice, emptySlice  # Nothing to display

        height, width = self.getChunkShape()
        ystart, ystop = yslice.start, yslice.stop
        xstart, xstop = xslice.start, xslice.stop
        return (
            slice(
                height * (ystart // height), height * int(numpy.ceil(ystop / height))
            ),
            slice(width * (xstart // width), width * int(numpy.ceil((xstop / width)))),
        )

if __name__ == "__main__":
    from silx.gui import qt
    from silx.gui.plot.AlphaSlider import ActiveImageAlphaSlider
    import numpy
    import h5py
    import sys

    filename = sys.argv[1]
    bg_filename = sys.argv[2] if len(sys.argv) == 3 else None

    qt.QApplication.setAttribute(qt.Qt.AA_ShareOpenGLContexts, True)
    qt.QApplication.setAttribute(qt.Qt.AA_EnableHighDpiScaling, True)

    app = qt.QApplication([])
    plot = PlotWindow(backend="gl")
    #w.getYAxis().setInverted(True)
    #w.setAxesDisplayed(False)
    plot.setKeepDataAspectRatio(True)
    item = Image()
    colormap = plot.getDefaultColormap()
    colormap.setVRange(-0.1, 0.3)
    item.setColormap(colormap)

    f = h5py.File(filename, mode="r")
    lods = [f['level%d' % i] for i in range(8) if ('level%d' % i) in f.keys()]
    item.setData(lods)

    item.setChunkShape((256, 256))
    plot.addItem(item)

    if bg_filename is not None:
        bgfile = h5py.File(bg_filename, mode="r")
        lods = [bgfile['level%d' % i] for i in range(8) if ('level%d' % i) in bgfile.keys()]
        bgitem = Image()
        bgitem.setName("Background")
        bgitem.setZValue(-1)
        bgitem.setData(lods)
        bgitem.setChunkShape((256, 256))
        plot.addItem(bgitem)

    plot.setActiveImage(item.getName())
    plot.resetZoom()

    alphaSlider = ActiveImageAlphaSlider(plot=plot)
    alphaSlider.show()

    window = qt.QMainWindow()
    window.setCentralWidget(plot)
    alphaDock = qt.QDockWidget()
    alphaDock.setWidget(alphaSlider)
    window.addDockWidget(qt.Qt.LeftDockWidgetArea, alphaDock)
    window.show()
    app.exec_()
<|endoftext|>


<commit_before>from .datastructures import FilteredImage, SizedImage

class AlreadyRegistered(Exception):
    pass

class InvalidSizedImageSubclass(Exception):
    pass

class NotRegistered(Exception):
    pass

class VersatileImageFieldRegistry(object):
    """
    A VersatileImageFieldRegistry object allows new SizedImage & FilteredImage
    subclasses to be dynamically added to all SizedImageFileField instances
    at runtime. New SizedImage subclasses are registered with the
    register_sizer method. New ProcessedImage subclasses are registered
    with the register_filter method.
    """

    def __init__(self, name='versatileimage_registry'):
        self._sizedimage_registry = {}  # attr_name -> sizedimage_cls
        self._filter_registry = {}  # attr_name -> filter_cls
        self.name = name

    def register_sizer(self, attr_name, sizedimage_cls):
        """
        Registers a new SizedImage subclass (`sizedimage_cls`) to be used
        via the attribute (`attr_name`)
        """
        if not issubclass(sizedimage_cls, SizedImage):
            raise InvalidSizedImageSubclass(
                'Only subclasses of versatileimagefield.datastructures.SizedImage '
                'may be registered with register_sizer'
            )

        if attr_name in self._sizedimage_registry:
            raise AlreadyRegistered(
                'A SizedImage class is already registered to the %s attribute,'
                'If you would like to override this attribute, use the '
                'unregister method' % attr_name
            )
        else:
            self._sizedimage_registry[attr_name] = sizedimage_cls

    def unregister_sizer(self, attr_name):
        """
        Unregisters the SizedImage subclass currently assigned to `attr_name`.

        If a SizedImage subclass isn't already registered to `attr_name`
        NotRegistered will raise.
        """
        if attr_name not in self._sizedimage_registry:
            raise NotRegistered(
                'No SizedImage subclass is registered to %s' % attr_name
            )
        else:
            del self._sizedimage_registry[attr_name]

    def register_filter(self, attr_name, filterimage_cls):
        """
        Registers a new FilteredImage subclass (`filterimage_cls`) to be used
        via the attribute (filters.`attr_name`)
        """
        if not issubclass(filterimage_cls, FilteredImage):
            raise InvalidSizedImageSubclass(
                'Only subclasses of FilteredImage may be registered as filters'
                'with VersatileImageFieldRegistry')

        if attr_name in self._filter_registry:
            raise AlreadyRegistered(
                'A ProcessedImageMixIn class is already registered to the %s '
                'attribute. If you would like to override this attribute, use '
                'the unregister method' % attr_name
            )
        else:
            self._filter_registry[attr_name] = filterimage_cls

    def unregister_filter(self, attr_name):
        """
        Unregisters the FilteredImage subclass currently assigned to
        `attr_name`.

        If a FilteredImage subclass isn't already registered to filters.
        `attr_name` NotRegistered will raise.
        """
        if attr_name not in self._filter_registry:
            raise NotRegistered(
                'No FilteredImage subclass is registered to %s' % attr_name
            )
        else:
            del self._filter_registry[attr_name]

versatileimagefield_registry = VersatileImageFieldRegistry()

def autodiscover():
    """
    Auto-discover INSTALLED_APPS sizedimage.py modules and fail silently when
    not present. This forces an import on them to register any versatileimagefield
    bits they may want.

    This is a near 1-to-1 copy of how django's admin application registers
    models.
    """

    import copy
    from django.conf import settings
    from django.utils.importlib import import_module
    from django.utils.module_loading import module_has_submodule

    for app in settings.INSTALLED_APPS:
        mod = import_module(app)
        # Attempt to import the app's sizedimage module.
        try:
            before_import_sizedimage_registry = copy.copy(
                versatileimagefield_registry._sizedimage_registry
            )
            before_import_filter_registry = copy.copy(
                versatileimagefield_registry._filter_registry
            )
            import_module('%s.versatileimagefield' % app)
        except:
            # Reset the versatileimagefield_registry to the state before the last
            # import as this import will have to reoccur on the next request
            # and this could raise NotRegistered and AlreadyRegistered
            # exceptions (see django ticket #8245).
            versatileimagefield_registry._sizedimage_registry = before_import_sizedimage_registry
            versatileimagefield_registry._filter_registry = before_import_filter_registry

            # Decide whether to bubble up this error. If the app just
            # doesn't have a sizedimage module, we can ignore the error
            # attempting to import it, otherwise we want it to bubble up.
            if module_has_submodule(mod, 'versatileimagefield'):
                raise
<commit_msg>Adding better Filter/Sizer 'name' validation<commit_after>from .datastructures import FilteredImage, SizedImage

class AlreadyRegistered(Exception):
    pass

class InvalidSizedImageSubclass(Exception):
    pass

class NotRegistered(Exception):
    pass

class UnallowedSizerName(Exception):
    pass

class UnallowedFilterName(Exception):
    pass

class VersatileImageFieldRegistry(object):
    """
    A VersatileImageFieldRegistry object allows new SizedImage & FilteredImage
    subclasses to be dynamically added to all SizedImageFileField instances
    at runtime. New SizedImage subclasses are registered with the
    register_sizer method. New ProcessedImage subclasses are registered
    with the register_filter method.
    """
    unallowed_sizer_names = (
        'build_filters_and_sizers',
        'chunks',
        'close',
        'closed',
        'delete',
        'encoding',
        'field',
        'file',
        'fileno',
        'filters',
        'flush',
        'height',
        'instance',
        'isatty',
        'multiple_chunks',
        'name',
        'newlines',
        'open',
        'path',
        'ppoi',
        'read',
        'readinto',
        'readline',
        'readlines',
        'save',
        'seek',
        'size',
        'softspace',
        'storage',
        'tell',
        'truncate',
        'url',
        'validate_ppoi',
        'width',
        'write',
        'writelines',
        'xreadlines'
    )

    def __init__(self, name='versatileimage_registry'):
        self._sizedimage_registry = {}  # attr_name -> sizedimage_cls
        self._filter_registry = {}  # attr_name -> filter_cls
        self.name = name

    def register_sizer(self, attr_name, sizedimage_cls):
        """
        Registers a new SizedImage subclass (`sizedimage_cls`) to be used
        via the attribute (`attr_name`)
        """
        if attr_name.startswith('_') or attr_name in self.unallowed_sizer_names:
            raise UnallowedSizerName(
                "`%s` is an unallowed Sizer name. Sizer names cannot begin "
                "with an underscore or be named any of the "
                "following: %s." % (
                    attr_name,
                    ', '.join([
                        name
                        for name in self.unallowed_sizer_names
                    ])
                )
            )
        if not issubclass(sizedimage_cls, SizedImage):
            raise InvalidSizedImageSubclass(
                'Only subclasses of versatileimagefield.datastructures.'
                'SizedImage may be registered with register_sizer'
            )

        if attr_name in self._sizedimage_registry:
            raise AlreadyRegistered(
                'A SizedImage class is already registered to the %s attribute,'
                'If you would like to override this attribute, use the '
                'unregister method' % attr_name
            )
        else:
            self._sizedimage_registry[attr_name] = sizedimage_cls

    def unregister_sizer(self, attr_name):
        """
        Unregisters the SizedImage subclass currently assigned to `attr_name`.

        If a SizedImage subclass isn't already registered to `attr_name`
        NotRegistered will raise.
        """
        if attr_name not in self._sizedimage_registry:
            raise NotRegistered(
                'No SizedImage subclass is registered to %s' % attr_name
            )
        else:
            del self._sizedimage_registry[attr_name]

    def register_filter(self, attr_name, filterimage_cls):
        """
        Registers a new FilteredImage subclass (`filterimage_cls`) to be used
        via the attribute (filters.`attr_name`)
        """
        if attr_name.startswith('_'):
            raise UnallowedFilterName(
                '`%s` is an unallowed Filter name. Filter names cannot begin '
                'with an underscore.' % attr_name
            )
        if not issubclass(filterimage_cls, FilteredImage):
            raise InvalidSizedImageSubclass(
                'Only subclasses of FilteredImage may be registered as filters'
                'with VersatileImageFieldRegistry')

        if attr_name in self._filter_registry:
            raise AlreadyRegistered(
                'A ProcessedImageMixIn class is already registered to the %s '
                'attribute. If you would like to override this attribute, use '
                'the unregister method' % attr_name
            )
        else:
            self._filter_registry[attr_name] = filterimage_cls

    def unregister_filter(self, attr_name):
        """
        Unregisters the FilteredImage subclass currently assigned to
        `attr_name`.

        If a FilteredImage subclass isn't already registered to filters.
        `attr_name` NotRegistered will raise.
        """
        if attr_name not in self._filter_registry:
            raise NotRegistered(
                'No FilteredImage subclass is registered to %s' % attr_name
            )
        else:
            del self._filter_registry[attr_name]

versatileimagefield_registry = VersatileImageFieldRegistry()

def autodiscover():
    """
    Auto-discover INSTALLED_APPS sizedimage.py modules and fail silently when
    not present. This forces an import on them to register any versatileimagefield
    bits they may want.

    This is a near 1-to-1 copy of how django's admin application registers
    models.
    """

    import copy
    from django.conf import settings
    from django.utils.importlib import import_module
    from django.utils.module_loading import module_has_submodule

    for app in settings.INSTALLED_APPS:
        mod = import_module(app)
        # Attempt to import the app's sizedimage module.
        try:
            before_import_sizedimage_registry = copy.copy(
                versatileimagefield_registry._sizedimage_registry
            )
            before_import_filter_registry = copy.copy(
                versatileimagefield_registry._filter_registry
            )
            import_module('%s.versatileimagefield' % app)
        except:
            # Reset the versatileimagefield_registry to the state before the last
            # import as this import will have to reoccur on the next request
            # and this could raise NotRegistered and AlreadyRegistered
            # exceptions (see django ticket #8245).
            versatileimagefield_registry._sizedimage_registry = before_import_sizedimage_registry
            versatileimagefield_registry._filter_registry = before_import_filter_registry

            # Decide whether to bubble up this error. If the app just
            # doesn't have a sizedimage module, we can ignore the error
            # attempting to import it, otherwise we want it to bubble up.
            if module_has_submodule(mod, 'versatileimagefield'):
                raise
<|endoftext|>


<commit_before>keys = {
	'up': 'w',
	'down': 's',
	'left': 'a',
	'right': 'd',
	'done': '1'
}
<commit_msg>Convert tabs to spaces<commit_after>keys = {
    'up': 'w',
    'down': 's',
    'left': 'a',
    'right': 'd',
    'done': '1'
}
<|endoftext|>


<commit_before># Copyright 2016 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 library for loading 1B word benchmark dataset."""

import random

import numpy as np
import tensorflow as tf

class Vocabulary(object):
	"""Class that holds a vocabulary for the dataset."""

	def __init__(self, filename):
		"""Initialize vocabulary.

		Args:
			filename: Vocabulary file name.
		"""

		self._id_to_word = []
		self._word_to_id = {}
		self._unk = -1
		self._bos = -1
		self._eos = -1

		with tf.gfile.Open(filename) as f:
			idx = 0
			for line in f:
				word_name = line.strip()
				if word_name == '<S>':
					self._bos = idx
				elif word_name == '</S>':
					self._eos = idx
				elif word_name == '<UNK>':
					self._unk = idx
				if word_name == '!!!MAXTERMID':
					continue

				self._id_to_word.append(word_name)
				self._word_to_id[word_name] = idx
				idx += 1

	@property
	def bos(self):
		return self._bos

	@property
	def eos(self):
		return self._eos

	@property
	def unk(self):
		return self._unk

	@property
	def size(self):
		return len(self._id_to_word)

	def word_to_id(self, word):
		if word in self._word_to_id:
			return self._word_to_id[word]
		return self.unk

	def id_to_word(self, cur_id):
		if cur_id < self.size:
			return self._id_to_word[cur_id]
		return 'ERROR'

	def decode(self, cur_ids):
		"""Convert a list of ids to a sentence, with space inserted."""
		return ' '.join([self.id_to_word(cur_id) for cur_id in cur_ids])

	def encode(self, sentence):
		"""Convert a sentence to a list of ids, with special tokens added."""
		word_ids = [self.word_to_id(cur_word) for cur_word in sentence.split()]
		return np.array([self.bos] + word_ids + [self.eos], dtype=np.int32)

class CharsVocabulary(Vocabulary):
	"""Vocabulary containing character-level information."""

	def __init__(self, filename, max_word_length):
		super(CharsVocabulary, self).__init__(filename)
		self._max_word_length = max_word_length
		chars_set = set()

		for word in self._id_to_word:
			chars_set |= set(word)

		free_ids = []
		for i in range(256):
			if chr(i) in chars_set:
				continue
			free_ids.append(chr(i))

		if len(free_ids) < 5:
			raise ValueError('Not enough free char ids: %d' % len(free_ids))

		self.bos_char = free_ids[0]	# <begin sentence>
		self.eos_char = free_ids[1]	# <end sentence>
		self.bow_char = free_ids[2]	# <begin word>
		self.eow_char = free_ids[3]	# <end word>
		self.pad_char = free_ids[4]	# <padding>

		chars_set |= {self.bos_char, self.eos_char, self.bow_char, self.eow_char,
									self.pad_char}

		self._char_set = chars_set
		num_words = len(self._id_to_word)

		self._word_char_ids = np.zeros([num_words, max_word_length], dtype=np.int32)

		self.bos_chars = self._convert_word_to_char_ids(self.bos_char)
		self.eos_chars = self._convert_word_to_char_ids(self.eos_char)

		for i, word in enumerate(self._id_to_word):
			self._word_char_ids[i] = self._convert_word_to_char_ids(word)

	@property
	def word_char_ids(self):
		return self._word_char_ids

	@property
	def max_word_length(self):
		return self._max_word_length

	def _convert_word_to_char_ids(self, word):
		code = np.zeros([self.max_word_length], dtype=np.int32)
		code[:] = ord(self.pad_char)

		if len(word) > self.max_word_length - 2:
			word = word[:self.max_word_length-2]
		cur_word = self.bow_char + word + self.eow_char
		for j in range(len(cur_word)):
			code[j] = ord(cur_word[j])
		return code

	def word_to_char_ids(self, word):
		if word in self._word_to_id:
			return self._word_char_ids[self._word_to_id[word]]
		else:
			return self._convert_word_to_char_ids(word)

	def encode_chars(self, sentence):
		chars_ids = [self.word_to_char_ids(cur_word)
								 for cur_word in sentence.split()]
		return np.vstack([self.bos_chars] + chars_ids + [self.eos_chars])

def get_batch(generator, batch_size, num_steps, max_word_length, pad=False):
	"""Read batches of input."""
	cur_stream = [None] * batch_size

	inputs = np.zeros([batch_size, num_steps], np.int32)
	char_inputs = np.zeros([batch_size, num_steps, max_word_length], np.int32)
	global_word_ids = np.zeros([batch_size, num_steps], np.int32)
	targets = np.zeros([batch_size, num_steps], np.int32)
	weights = np.ones([batch_size, num_steps], np.float32)

	no_more_data = False
	while True:
		inputs[:] = 0
		char_inputs[:] = 0
		global_word_ids[:] = 0
		targets[:] = 0
		weights[:] = 0.0

		for i in range(batch_size):
			cur_pos = 0

			while cur_pos < num_steps:
				if cur_stream[i] is None or len(cur_stream[i][0]) <= 1:
					try:
						cur_stream[i] = list(generator.next())
					except StopIteration:
						# No more data, exhaust current streams and quit
						no_more_data = True
						break

				how_many = min(len(cur_stream[i][0]) - 1, num_steps - cur_pos)
				next_pos = cur_pos + how_many

				inputs[i, cur_pos:next_pos] = cur_stream[i][0][:how_many]
				char_inputs[i, cur_pos:next_pos] = cur_stream[i][1][:how_many]
				global_word_ids[i, cur_pos:next_pos] = cur_stream[i][2][:how_many]
				targets[i, cur_pos:next_pos] = cur_stream[i][0][1:how_many+1]
				weights[i, cur_pos:next_pos] = 1.0

				cur_pos = next_pos
				cur_stream[i][0] = cur_stream[i][0][how_many:]
				cur_stream[i][1] = cur_stream[i][1][how_many:]
				cur_stream[i][2] = cur_stream[i][2][how_many:]

				if pad:
					break

		if no_more_data and np.sum(weights) == 0:
			# There is no more data and this is an empty batch. Done!
			break
		yield inputs, char_inputs, global_word_ids, targets, weights

class LM1BDataset(object):
	"""Utility class for 1B word benchmark dataset.

	The current implementation reads the data from the tokenized text files.
	"""

	def __init__(self, filepattern, vocab):
		"""Initialize LM1BDataset reader.

		Args:
			filepattern: Dataset file pattern.
			vocab: Vocabulary.
		"""
		self._vocab = vocab
		self._all_shards = tf.gfile.Glob(filepattern)
		tf.logging.info('Found %d shards at %s', len(self._all_shards), filepattern)

	def _load_random_shard(self):
		"""Randomly select a file and read it."""
		return self._load_shard(random.choice(self._all_shards))

	def _load_shard(self, shard_name):
		"""Read one file and convert to ids.

		Args:
			shard_name: file path.

		Returns:
			list of (id, char_id, global_word_id) tuples.
		"""
		tf.logging.info('Loading data from: %s', shard_name)
		with tf.gfile.Open(shard_name) as f:
			sentences = f.readlines()
		chars_ids = [self.vocab.encode_chars(sentence) for sentence in sentences]
		ids = [self.vocab.encode(sentence) for sentence in sentences]

		global_word_ids = []
		current_idx = 0
		for word_ids in ids:
			current_size = len(word_ids) - 1	# without <BOS> symbol
			cur_ids = np.arange(current_idx, current_idx + current_size)
			global_word_ids.append(cur_ids)
			current_idx += current_size

		tf.logging.info('Loaded %d words.', current_idx)
		tf.logging.info('Finished loading')
		return zip(ids, chars_ids, global_word_ids)

	def _get_sentence(self, forever=True):
		while True:
			ids = self._load_random_shard()
			for current_ids in ids:
				yield current_ids
			if not forever:
				break

	def get_batch(self, batch_size, num_steps, pad=False, forever=True):
		return get_batch(self._get_sentence(forever), batch_size, num_steps,
										 self.vocab.max_word_length, pad=pad)

	@property
	def vocab(self):
		return self._vocab
<commit_msg>date util python3 support<commit_after>#!/usr/bin/python3

# Copyright 2016 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 library for loading 1B word benchmark dataset."""

import random

import numpy as np
import tensorflow as tf

class Vocabulary(object):
	"""Class that holds a vocabulary for the dataset."""

	def __init__(self, filename):
		"""Initialize vocabulary.

		Args:
			filename: Vocabulary file name.
		"""

		self._id_to_word = []
		self._word_to_id = {}
		self._unk = -1
		self._bos = -1
		self._eos = -1

		with tf.gfile.Open(filename) as f:
			idx = 0
			for line in f:
				word_name = line.strip()
				if word_name == '<S>':
					self._bos = idx
				elif word_name == '</S>':
					self._eos = idx
				elif word_name == '<UNK>':
					self._unk = idx
				if word_name == '!!!MAXTERMID':
					continue

				self._id_to_word.append(word_name)
				self._word_to_id[word_name] = idx
				idx += 1

	@property
	def bos(self):
		return self._bos

	@property
	def eos(self):
		return self._eos

	@property
	def unk(self):
		return self._unk

	@property
	def size(self):
		return len(self._id_to_word)

	def word_to_id(self, word):
		if word in self._word_to_id:
			return self._word_to_id[word]
		return self.unk

	def id_to_word(self, cur_id):
		if cur_id < self.size:
			return self._id_to_word[cur_id]
		return 'ERROR'

	def decode(self, cur_ids):
		"""Convert a list of ids to a sentence, with space inserted."""
		return ' '.join([self.id_to_word(cur_id) for cur_id in cur_ids])

	def encode(self, sentence):
		"""Convert a sentence to a list of ids, with special tokens added."""
		word_ids = [self.word_to_id(cur_word) for cur_word in sentence.split()]
		return np.array([self.bos] + word_ids + [self.eos], dtype=np.int32)

class CharsVocabulary(Vocabulary):
	"""Vocabulary containing character-level information."""

	def __init__(self, filename, max_word_length):
		super(CharsVocabulary, self).__init__(filename)
		self._max_word_length = max_word_length
		chars_set = set()

		for word in self._id_to_word:
			chars_set |= set(word)

		free_ids = []
		for i in range(256):
			if chr(i) in chars_set:
				continue
			free_ids.append(chr(i))

		if len(free_ids) < 5:
			raise ValueError('Not enough free char ids: %d' % len(free_ids))

		self.bos_char = free_ids[0]	# <begin sentence>
		self.eos_char = free_ids[1]	# <end sentence>
		self.bow_char = free_ids[2]	# <begin word>
		self.eow_char = free_ids[3]	# <end word>
		self.pad_char = free_ids[4]	# <padding>

		chars_set |= {self.bos_char, self.eos_char, self.bow_char, self.eow_char,
									self.pad_char}

		self._char_set = chars_set
		num_words = len(self._id_to_word)

		self._word_char_ids = np.zeros([num_words, max_word_length], dtype=np.int32)

		self.bos_chars = self._convert_word_to_char_ids(self.bos_char)
		self.eos_chars = self._convert_word_to_char_ids(self.eos_char)

		for i, word in enumerate(self._id_to_word):
			self._word_char_ids[i] = self._convert_word_to_char_ids(word)

	@property
	def word_char_ids(self):
		return self._word_char_ids

	@property
	def max_word_length(self):
		return self._max_word_length

	def _convert_word_to_char_ids(self, word):
		code = np.zeros([self.max_word_length], dtype=np.int32)
		code[:] = ord(self.pad_char)

		if len(word) > self.max_word_length - 2:
			word = word[:self.max_word_length-2]
		cur_word = self.bow_char + word + self.eow_char
		for j in range(len(cur_word)):
			code[j] = ord(cur_word[j])
		return code

	def word_to_char_ids(self, word):
		if word in self._word_to_id:
			return self._word_char_ids[self._word_to_id[word]]
		else:
			return self._convert_word_to_char_ids(word)

	def encode_chars(self, sentence):
		chars_ids = [self.word_to_char_ids(cur_word)
								 for cur_word in sentence.split()]
		return np.vstack([self.bos_chars] + chars_ids + [self.eos_chars])

def get_batch(generator, batch_size, num_steps, max_word_length, pad=False):
	"""Read batches of input."""
	cur_stream = [None] * batch_size

	inputs = np.zeros([batch_size, num_steps], np.int32)
	char_inputs = np.zeros([batch_size, num_steps, max_word_length], np.int32)
	global_word_ids = np.zeros([batch_size, num_steps], np.int32)
	targets = np.zeros([batch_size, num_steps], np.int32)
	weights = np.ones([batch_size, num_steps], np.float32)

	no_more_data = False
	while True:
		inputs[:] = 0
		char_inputs[:] = 0
		global_word_ids[:] = 0
		targets[:] = 0
		weights[:] = 0.0

		for i in range(batch_size):
			cur_pos = 0

			while cur_pos < num_steps:
				if cur_stream[i] is None or len(cur_stream[i][0]) <= 1:
					try:
						cur_stream[i] = list(generator.next())
					except StopIteration:
						# No more data, exhaust current streams and quit
						no_more_data = True
						break

				how_many = min(len(cur_stream[i][0]) - 1, num_steps - cur_pos)
				next_pos = cur_pos + how_many

				inputs[i, cur_pos:next_pos] = cur_stream[i][0][:how_many]
				char_inputs[i, cur_pos:next_pos] = cur_stream[i][1][:how_many]
				global_word_ids[i, cur_pos:next_pos] = cur_stream[i][2][:how_many]
				targets[i, cur_pos:next_pos] = cur_stream[i][0][1:how_many+1]
				weights[i, cur_pos:next_pos] = 1.0

				cur_pos = next_pos
				cur_stream[i][0] = cur_stream[i][0][how_many:]
				cur_stream[i][1] = cur_stream[i][1][how_many:]
				cur_stream[i][2] = cur_stream[i][2][how_many:]

				if pad:
					break

		if no_more_data and np.sum(weights) == 0:
			# There is no more data and this is an empty batch. Done!
			break
		yield inputs, char_inputs, global_word_ids, targets, weights

class LM1BDataset(object):
	"""Utility class for 1B word benchmark dataset.

	The current implementation reads the data from the tokenized text files.
	"""

	def __init__(self, filepattern, vocab):
		"""Initialize LM1BDataset reader.

		Args:
			filepattern: Dataset file pattern.
			vocab: Vocabulary.
		"""
		self._vocab = vocab
		self._all_shards = tf.gfile.Glob(filepattern)
		tf.logging.info('Found %d shards at %s', len(self._all_shards), filepattern)

	def _load_random_shard(self):
		"""Randomly select a file and read it."""
		return self._load_shard(random.choice(self._all_shards))

	def _load_shard(self, shard_name):
		"""Read one file and convert to ids.

		Args:
			shard_name: file path.

		Returns:
			list of (id, char_id, global_word_id) tuples.
		"""
		tf.logging.info('Loading data from: %s', shard_name)
		with tf.gfile.Open(shard_name) as f:
			sentences = f.readlines()
		chars_ids = [self.vocab.encode_chars(sentence) for sentence in sentences]
		ids = [self.vocab.encode(sentence) for sentence in sentences]

		global_word_ids = []
		current_idx = 0
		for word_ids in ids:
			current_size = len(word_ids) - 1	# without <BOS> symbol
			cur_ids = np.arange(current_idx, current_idx + current_size)
			global_word_ids.append(cur_ids)
			current_idx += current_size

		tf.logging.info('Loaded %d words.', current_idx)
		tf.logging.info('Finished loading')
		return zip(ids, chars_ids, global_word_ids)

	def _get_sentence(self, forever=True):
		while True:
			ids = self._load_random_shard()
			for current_ids in ids:
				yield current_ids
			if not forever:
				break

	def get_batch(self, batch_size, num_steps, pad=False, forever=True):
		return get_batch(self._get_sentence(forever), batch_size, num_steps,
										 self.vocab.max_word_length, pad=pad)

	@property
	def vocab(self):
		return self._vocab
<|endoftext|>


<commit_before>#!/usr/bin/env python
"""GuideImage represents a guide image

History:
2006-08-03 ROwen    Separated ImObj out of GuideWdg.py into this file
                    and renamed to GuideImage.
2006-09-12 ROwen    Modified GuideImage.getFITSObj to parse the header
                    and set various useful attributes.
2007-01-16 ROwen    Added commented-out code to print a traceback if file read fails.
2007-01-30 ROwen    Was not caching FITS header info (not setting parsedFITSHeader).
"""
import os
import pyfits
import sys
import TUI.HubModel
import SubFrame
import traceback

_DebugMem = False # print a message when a file is deleted from disk?

class BasicImage(object):
    """Information about an image.

    Inputs:
    - localBaseDir  root image directory on local machine
    - imageName path to image relative, specifically:
                if isLocal False, then a URL relative to the download host
                if isLocal True, then a local path relative to localBaseDir
    - imageName unix path to image, relative to host root directory
    - guideModel    guide model for this actor
    - fetchCallFunc function to call when image info changes state
    - isLocal   set True if image is local or already downloaded
    """
    Ready = "Ready to download"
    Downloading = "Downloading"
    Downloaded = "Downloaded"
    FileReadFailed = "Cannot read file"
    DownloadFailed = "Download failed"
    Expired = "Expired; file deleted"
    ErrorStates = (FileReadFailed, DownloadFailed, Expired)
    DoneStates = (Downloaded,) + ErrorStates

    def __init__(self,
        localBaseDir,
        imageName,
        downloadWdg = None,
        fetchCallFunc = None,
        isLocal = False,
    ):
        #print "%s localBaseDir=%r, imageName=%s" % (self.__class__.__name__, localBaseDir, imageName)
        self.localBaseDir = localBaseDir
        self.imageName = imageName
        self.downloadWdg = downloadWdg
        self.hubModel = TUI.HubModel.getModel()
        self.errMsg = None
        self.fetchCallFunc = fetchCallFunc
        self.isLocal = isLocal
        if not self.isLocal:
            self.state = self.Ready
        else:
            self.state = self.Downloaded
        self.isInSequence = not isLocal

        # set local path
        # this split suffices to separate the components because image names are simple
        if isLocal:
            self.localPath = os.path.join(self.localBaseDir, imageName)
        else:
            pathComponents = self.imageName.split("/")
            self.localPath = os.path.join(self.localBaseDir, *pathComponents)
        #print "GuideImage localPath=%r" % (self.localPath,)

    def didFail(self):
        """Return False if download failed or image expired"""
        return self.state in self.ErrorStates

    def expire(self):
        """Delete the file from disk and set state to expired.
        """
        if self.isLocal:
            if _DebugMem:
                print "Would delete %r, but is local" % (self.imageName,)
            return
        if self.state == self.Downloaded:
            # don't use _setState because no callback wanted
            # and _setState ignored new states once done
            self.state = self.Expired
            if os.path.exists(self.localPath):
                if _DebugMem:
                    print "Deleting %r" % (self.localPath,)
                os.remove(self.localPath)
            elif _DebugMem:
                print "Would delete %r, but not found on disk" % (self.imageName,)
        elif _DebugMem:
            print "Would delete %r, but state = %r is not 'downloaded'" % (self.imageName, self.state,)

    def fetchFile(self):
        """Start downloading the file."""
        #print "%s fetchFile; isLocal=%s" % (self, self.isLocal)
        if self.isLocal:
            self._setState(self.Downloaded)
            return

        host, hostRootDir = self.hubModel.httpRoot.get()[0]
        if None in (host, hostRootDir):
            self._setState(
                self.DownloadFailed,
                "Cannot download images; hub httpRoot keyword not available",
            )
            return

        self._setState(self.Downloading)

        fromURL = "".join(("http://", host, hostRootDir, self.imageName))
        self.downloadWdg.getFile(
            fromURL = fromURL,
            toPath = self.localPath,
            isBinary = True,
            overwrite = True,
            createDir = True,
            doneFunc = self._fetchDoneFunc,
            dispStr = self.imageName,
        )

    def getFITSObj(self):
        """If the file is available, return a pyfits object,
        else return None.
        """
        if self.state == self.Downloaded:
            try:
                fitsIm = pyfits.open(self.getLocalPath())
                if fitsIm:
                    return fitsIm

                self.state = self.FileReadFailed
                self.errMsg = "No image data found"
                return None
            except (SystemExit, KeyboardInterrupt):
                raise
            except Exception, e:
                self.state = self.FileReadFailed
                self.errMsg = str(e)
#               sys.stderr.write("Could not read file %r: %s\n" % (self.getLocalPath(), e))
#               traceback.print_exc(file=sys.stderr)
        return None

    def getLocalPath(self):
        """Return the full local path to the image."""
        return self.localPath

    def getStateStr(self):
        """Return a string describing the current state."""
        if self.errMsg:
            return "%s: %s" % (self.state, self.errMsg)
        return self.state

    def isDone(self):
        """Return True if download finished (successfully or otherwise)"""
        return self.state in self.DoneStates

    def _fetchDoneFunc(self, httpGet):
        """Called when image download ends.
        """
        if httpGet.getState() == httpGet.Done:
            self._setState(self.Downloaded)
        else:
            self._setState(self.DownloadFailed, httpGet.getErrMsg())
            #print "%s download failed: %s" % (self, self.errMsg)

    def _setState(self, state, errMsg=None):
        if self.isDone():
            return

        self.state = state
        if self.didFail():
            self.errMsg = errMsg

        if self.fetchCallFunc:
            self.fetchCallFunc(self)
        if self.isDone():
            self.fetchCallFunc = None

    def __str__(self):
        return "%s(%s)" % (self.__class__.__name__, self.imageName)

class GuideImage(BasicImage):
    def __init__(self,
        localBaseDir,
        imageName,
        downloadWdg = None,
        fetchCallFunc = None,
        defRadMult = None,
        defThresh = None,
        defGuideMode = None,
        isLocal = False,
    ):
        self.starDataDict = {} # dict of star type char: star keyword data
        self.defSelDataColor = None
        self.selDataColor = None
        self.guiderPredPos = None
        self.defRadMult = defRadMult
        self.defThresh = defThresh
        self.defGuideMode = defGuideMode
        self.currRadMult = None
        self.currThresh = None
        self.currGuideMode = None
        self.parsedFITSHeader = False
        self.binFac = None
        self.expTime = None
        self.subFrame = None

        BasicImage.__init__(self,
            localBaseDir = localBaseDir,
            imageName = imageName,
            downloadWdg = downloadWdg,
            fetchCallFunc = fetchCallFunc,
            isLocal = isLocal,
        )

    def getFITSObj(self):
        """Return the pyfits image object, or None if unavailable.

        Parse the FITS header, if not already done,
        and set the following attributes:
        - binFac: bin factor (a scalar; x = y)
        - expTime: exposure time (floating seconds)
        - subFrame: a SubFrame object
        - binSubBeg: subframe beginning (binned pixels; 0,0 is lower left corner)
        - binSubSize: subframe size (binned pixels)
        """
        fitsObj = BasicImage.getFITSObj(self)
        if fitsObj and not self.parsedFITSHeader:
            imHdr = fitsObj[0].header
            self.expTime = imHdr.get("EXPTIME")
            self.binFac = imHdr.get("BINX")
            try:
                self.subFrame = SubFrame.SubFrame.fromFITS(imHdr)
            except ValueError:
                pass
            self.parsedFITSHeader = True

        return fitsObj
<commit_msg>TUI/Guide/GuideImage.py - Minor tweak to history     <commit_after>#!/usr/bin/env python
"""GuideImage represents a guide image

History:
2006-08-03 ROwen    Separated ImObj out of GuideWdg.py into this file
                    and renamed to GuideImage.
2006-09-12 ROwen    Modified GuideImage.getFITSObj to parse the header
                    and set various useful attributes.
2007-01-16 ROwen    Added commented-out code to print a traceback if file read fails.
2007-01-30 ROwen    Was not caching FITS header info (despite code to do this).
"""
import os
import pyfits
import sys
import TUI.HubModel
import SubFrame
import traceback

_DebugMem = False # print a message when a file is deleted from disk?

class BasicImage(object):
    """Information about an image.

    Inputs:
    - localBaseDir  root image directory on local machine
    - imageName path to image relative, specifically:
                if isLocal False, then a URL relative to the download host
                if isLocal True, then a local path relative to localBaseDir
    - imageName unix path to image, relative to host root directory
    - guideModel    guide model for this actor
    - fetchCallFunc function to call when image info changes state
    - isLocal   set True if image is local or already downloaded
    """
    Ready = "Ready to download"
    Downloading = "Downloading"
    Downloaded = "Downloaded"
    FileReadFailed = "Cannot read file"
    DownloadFailed = "Download failed"
    Expired = "Expired; file deleted"
    ErrorStates = (FileReadFailed, DownloadFailed, Expired)
    DoneStates = (Downloaded,) + ErrorStates

    def __init__(self,
        localBaseDir,
        imageName,
        downloadWdg = None,
        fetchCallFunc = None,
        isLocal = False,
    ):
        #print "%s localBaseDir=%r, imageName=%s" % (self.__class__.__name__, localBaseDir, imageName)
        self.localBaseDir = localBaseDir
        self.imageName = imageName
        self.downloadWdg = downloadWdg
        self.hubModel = TUI.HubModel.getModel()
        self.errMsg = None
        self.fetchCallFunc = fetchCallFunc
        self.isLocal = isLocal
        if not self.isLocal:
            self.state = self.Ready
        else:
            self.state = self.Downloaded
        self.isInSequence = not isLocal

        # set local path
        # this split suffices to separate the components because image names are simple
        if isLocal:
            self.localPath = os.path.join(self.localBaseDir, imageName)
        else:
            pathComponents = self.imageName.split("/")
            self.localPath = os.path.join(self.localBaseDir, *pathComponents)
        #print "GuideImage localPath=%r" % (self.localPath,)

    def didFail(self):
        """Return False if download failed or image expired"""
        return self.state in self.ErrorStates

    def expire(self):
        """Delete the file from disk and set state to expired.
        """
        if self.isLocal:
            if _DebugMem:
                print "Would delete %r, but is local" % (self.imageName,)
            return
        if self.state == self.Downloaded:
            # don't use _setState because no callback wanted
            # and _setState ignored new states once done
            self.state = self.Expired
            if os.path.exists(self.localPath):
                if _DebugMem:
                    print "Deleting %r" % (self.localPath,)
                os.remove(self.localPath)
            elif _DebugMem:
                print "Would delete %r, but not found on disk" % (self.imageName,)
        elif _DebugMem:
            print "Would delete %r, but state = %r is not 'downloaded'" % (self.imageName, self.state,)

    def fetchFile(self):
        """Start downloading the file."""
        #print "%s fetchFile; isLocal=%s" % (self, self.isLocal)
        if self.isLocal:
            self._setState(self.Downloaded)
            return

        host, hostRootDir = self.hubModel.httpRoot.get()[0]
        if None in (host, hostRootDir):
            self._setState(
                self.DownloadFailed,
                "Cannot download images; hub httpRoot keyword not available",
            )
            return

        self._setState(self.Downloading)

        fromURL = "".join(("http://", host, hostRootDir, self.imageName))
        self.downloadWdg.getFile(
            fromURL = fromURL,
            toPath = self.localPath,
            isBinary = True,
            overwrite = True,
            createDir = True,
            doneFunc = self._fetchDoneFunc,
            dispStr = self.imageName,
        )

    def getFITSObj(self):
        """If the file is available, return a pyfits object,
        else return None.
        """
        if self.state == self.Downloaded:
            try:
                fitsIm = pyfits.open(self.getLocalPath())
                if fitsIm:
                    return fitsIm

                self.state = self.FileReadFailed
                self.errMsg = "No image data found"
                return None
            except (SystemExit, KeyboardInterrupt):
                raise
            except Exception, e:
                self.state = self.FileReadFailed
                self.errMsg = str(e)
#               sys.stderr.write("Could not read file %r: %s\n" % (self.getLocalPath(), e))
#               traceback.print_exc(file=sys.stderr)
        return None

    def getLocalPath(self):
        """Return the full local path to the image."""
        return self.localPath

    def getStateStr(self):
        """Return a string describing the current state."""
        if self.errMsg:
            return "%s: %s" % (self.state, self.errMsg)
        return self.state

    def isDone(self):
        """Return True if download finished (successfully or otherwise)"""
        return self.state in self.DoneStates

    def _fetchDoneFunc(self, httpGet):
        """Called when image download ends.
        """
        if httpGet.getState() == httpGet.Done:
            self._setState(self.Downloaded)
        else:
            self._setState(self.DownloadFailed, httpGet.getErrMsg())
            #print "%s download failed: %s" % (self, self.errMsg)

    def _setState(self, state, errMsg=None):
        if self.isDone():
            return

        self.state = state
        if self.didFail():
            self.errMsg = errMsg

        if self.fetchCallFunc:
            self.fetchCallFunc(self)
        if self.isDone():
            self.fetchCallFunc = None

    def __str__(self):
        return "%s(%s)" % (self.__class__.__name__, self.imageName)

class GuideImage(BasicImage):
    def __init__(self,
        localBaseDir,
        imageName,
        downloadWdg = None,
        fetchCallFunc = None,
        defRadMult = None,
        defThresh = None,
        defGuideMode = None,
        isLocal = False,
    ):
        self.starDataDict = {} # dict of star type char: star keyword data
        self.defSelDataColor = None
        self.selDataColor = None
        self.guiderPredPos = None
        self.defRadMult = defRadMult
        self.defThresh = defThresh
        self.defGuideMode = defGuideMode
        self.currRadMult = None
        self.currThresh = None
        self.currGuideMode = None
        self.parsedFITSHeader = False
        self.binFac = None
        self.expTime = None
        self.subFrame = None

        BasicImage.__init__(self,
            localBaseDir = localBaseDir,
            imageName = imageName,
            downloadWdg = downloadWdg,
            fetchCallFunc = fetchCallFunc,
            isLocal = isLocal,
        )

    def getFITSObj(self):
        """Return the pyfits image object, or None if unavailable.

        Parse the FITS header, if not already done,
        and set the following attributes:
        - binFac: bin factor (a scalar; x = y)
        - expTime: exposure time (floating seconds)
        - subFrame: a SubFrame object
        - binSubBeg: subframe beginning (binned pixels; 0,0 is lower left corner)
        - binSubSize: subframe size (binned pixels)
        """
        fitsObj = BasicImage.getFITSObj(self)
        if fitsObj and not self.parsedFITSHeader:
            imHdr = fitsObj[0].header
            self.expTime = imHdr.get("EXPTIME")
            self.binFac = imHdr.get("BINX")
            try:
                self.subFrame = SubFrame.SubFrame.fromFITS(imHdr)
            except ValueError:
                pass
            self.parsedFITSHeader = True

        return fitsObj
<|endoftext|>


<commit_before>from __future__ import division

def sum_partial_reflections(integrated_data, min_total_partiality=0.5):
  '''Sum partial reflections; weighted sum for summation integration; weighted
  average for profile fitted reflections. N.B. this will report total
  partiality for the summed reflection.'''

  return integrated_data

def scale_partial_reflections(integrated_data, min_partiality=0.5):
  '''Scale partial reflections (after summation) according to their estimated
  partiality - for profile fitted reflections this will result in no change,
  for summation integrated reflections will be scaled up by 1 / partiality
  with error accordingly scaled. N.B. this will report the scaled up partiality
  for the output reflection.'''

  return integrated_data

def export_mtz(integrated_data, experiment_list, hklout, ignore_panels=False,
               include_partials=False):
  '''Export data from integrated_data corresponding to experiment_list to an
  MTZ file hklout.'''
  from logging import info

  # for the moment assume (and assert) that we will convert data from exactly
  # one lattice...

  # FIXME allow for more than one experiment in here: this is fine just add
  # multiple MTZ data sets (DIALS1...DIALSN) and multiple batch headers: one
  # range of batches for each experiment

  assert(len(experiment_list) == 1)
  assert(min(integrated_data['id']) == max(integrated_data['id']) == 0)

  # strip out negative variance reflections: these should not really be there
  # FIXME Doing select on summation results. Should do on profile result if
  # present? Yes

  selection = integrated_data['intensity.sum.variance'] <= 0
  if selection.count(True) > 0:
    integrated_data.del_selected(selection)
    info('Removing %d reflections with negative variance' % \
          selection.count(True))

  if 'intensity.prf.variance' in integrated_data:
    selection = integrated_data['intensity.prf.variance'] <= 0
    if selection.count(True) > 0:
      integrated_data.del_selected(selection)
      info('Removing %d profile reflections with negative variance' % \
            selection.count(True))

  # FIXME in here work on including partial reflections => at this stage best
  # to split off the partial refections into a different selection & handle
  # gracefully... better to work on a short list as will need to "pop" them &
  # find matching parts to combine.

  if include_partials:
    integrated_data = sum_partial_reflections(integrated_data)
    integrated_data = scale_partial_reflections(integrated_data)

  if 'partiality' in integrated_data:
    selection = integrated_data['partiality'] < 0.99
    if selection.count(True) > 0:
      integrated_data.del_selected(selection)
      info('Removing %d incomplete reflections' % \
        selection.count(True))

  # FIXME TODO for more than one experiment into an MTZ file:
  #
  # - add an epoch (or recover an epoch) from the scan and add this as an extra
  #   column to the MTZ file for scaling, so we know that the two lattices were
  #   integrated at the same time
  # - decide a sensible BATCH increment to apply to the BATCH value between
  #   experiments and add this
  #
  # At the moment this is probably enough to be working on.

  experiment = experiment_list[0]

  # also only work with one panel(for the moment)

  if not ignore_panels:
    assert(len(experiment.detector) == 1)

  axis = experiment.goniometer.get_rotation_axis()
  s0 = experiment.beam.get_s0()
  wavelength = experiment.beam.get_wavelength()

  from scitbx import matrix

  panel = experiment.detector[0]
  origin = matrix.col(panel.get_origin())
  fast = matrix.col(panel.get_fast_axis())
  slow = matrix.col(panel.get_slow_axis())

  pixel_size = panel.get_pixel_size()

  fast *= pixel_size[0]
  slow *= pixel_size[1]

  cb_op_to_ref = experiment.crystal.get_space_group().info(
    ).change_of_basis_op_to_reference_setting()

  experiment.crystal = experiment.crystal.change_basis(cb_op_to_ref)

  U = experiment.crystal.get_U()
  unit_cell = experiment.crystal.get_unit_cell()

  from iotbx import mtz

  from scitbx.array_family import flex
  from math import floor, sqrt

  m = mtz.object()
  m.set_title('from dials.export_mtz')
  m.set_space_group_info(experiment.crystal.get_space_group().info())

  image_range = experiment.scan.get_image_range()

  for b in range(image_range[0], image_range[1] + 1):
    o = m.add_batch().set_num(b).set_nbsetid(1).set_ncryst(1)
    o.set_time1(0.0).set_time2(0.0).set_title('Batch %d' % b)
    o.set_ndet(1).set_theta(flex.float((0.0, 0.0))).set_lbmflg(0)
    o.set_alambd(wavelength).set_delamb(0.0).set_delcor(0.0)
    o.set_divhd(0.0).set_divvd(0.0)

    # FIXME hard-coded assumption on indealized beam vector below... this may be
    # broken when we come to process data from a non-imgCIF frame
    o.set_so(flex.float(s0)).set_source(flex.float((0, 0, -1)))

    # these are probably 0, 1 respectively, also flags for how many are set, sd
    o.set_bbfac(0.0).set_bscale(1.0)
    o.set_sdbfac(0.0).set_sdbscale(0.0).set_nbscal(0)

    # unit cell (this is fine) and the what-was-refined-flags FIXME hardcoded

    # take time-varying parameters from the *end of the frame* unlikely to
    # be much different at the end - however only exist if time-varying refinement
    # was used
    if experiment.crystal.num_scan_points:
      _unit_cell = experiment.crystal.get_unit_cell_at_scan_point(b)
      _U = experiment.crystal.get_U_at_scan_point(b)
    else:
      _unit_cell = unit_cell
      _U = U

    # FIXME need to get what was refined and what was constrained from the
    # crystal model
    o.set_cell(flex.float(_unit_cell.parameters()))
    o.set_lbcell(flex.int((-1, -1, -1, -1, -1, -1)))
    o.set_umat(flex.float(_U.transpose().elems))

    # get the mosaic spread though today it may not actually be set
    mosaic = experiment.crystal.get_mosaicity()
    o.set_crydat(flex.float([mosaic, 0.0, 0.0, 0.0, 0.0, 0.0,
                             0.0, 0.0, 0.0, 0.0, 0.0, 0.0]))

    o.set_lcrflg(0)
    o.set_datum(flex.float((0.0, 0.0, 0.0)))

    # detector size, distance
    o.set_detlm(flex.float([0.0, panel.get_image_size()[0],
                            0.0, panel.get_image_size()[1],
                            0, 0, 0, 0]))
    o.set_dx(flex.float([panel.get_distance(), 0.0]))

    # goniometer axes and names, and scan axis number, and number of axes, missets
    o.set_e1(flex.float(axis))
    o.set_e2(flex.float((0.0, 0.0, 0.0)))
    o.set_e3(flex.float((0.0, 0.0, 0.0)))
    o.set_gonlab(flex.std_string(('AXIS', '', '')))
    o.set_jsaxs(1)
    o.set_ngonax(1)
    o.set_phixyz(flex.float((0.0, 0.0, 0.0, 0.0, 0.0, 0.0)))

    # scan ranges, axis
    phi_start, phi_range = experiment.scan.get_image_oscillation(b)
    o.set_phistt(phi_start)
    o.set_phirange(phi_range)
    o.set_phiend(phi_start + phi_range)
    o.set_scanax(flex.float(axis))

    # number of misorientation angles
    o.set_misflg(0)

    # crystal axis closest to rotation axis (why do I want this?)
    o.set_jumpax(0)

    # type of data - 1; 2D, 2; 3D, 3; Laue
    o.set_ldtype(2)

  # now create the actual data structures - first keep a track of the columns
  # H K L M/ISYM BATCH I SIGI IPR SIGIPR FRACTIONCALC XDET YDET ROT WIDTH
  # LP MPART FLAG BGPKRATIOS

  from cctbx.array_family import flex as cflex # implicit import
  from cctbx.miller import map_to_asu_isym # implicit import

  # gather the required information for the reflection file

  nref = len(integrated_data['miller_index'])
  x_px, y_px, z_px = integrated_data['xyzcal.px'].parts()

  xdet = flex.double(x_px)
  ydet = flex.double(y_px)
  zdet = flex.double(z_px)

  # compute ROT values
  rot = flex.double([experiment.scan.get_angle_from_image_index(z)
                     for z in zdet])

  # compute BATCH values
  batch = flex.floor(zdet).iround() + 1

  # we're working with full reflections so...
  fractioncalc = flex.double(nref, 1.0)

  # now go for it and make an MTZ file...

  x = m.add_crystal('XTAL', 'DIALS', unit_cell.parameters())
  d = x.add_dataset('FROMDIALS', wavelength)

  # now add column information...

  # FIXME add DIALS_FLAG which can include e.g. was partial etc.

  type_table = {'IPR': 'J', 'BGPKRATIOS': 'R', 'WIDTH': 'R', 'I': 'J',
                'H': 'H', 'K': 'H', 'MPART': 'I', 'L': 'H', 'BATCH': 'B',
                'M_ISYM': 'Y', 'SIGI': 'Q', 'FLAG': 'I', 'XDET': 'R', 'LP': 'R',
                'YDET': 'R', 'SIGIPR': 'Q', 'FRACTIONCALC': 'R', 'ROT': 'R'}

  # derive index columns from original indices with
  #
  # from m.replace_original_index_miller_indices
  #
  # so all that is needed now is to make space for the reflections - fill with
  # zeros...

  m.adjust_column_array_sizes(nref)
  m.set_n_reflections(nref)

  # assign H, K, L, M_ISYM space
  for column in 'H', 'K', 'L', 'M_ISYM':
    d.add_column(column, type_table[column]).set_values(
      flex.double(nref, 0.0).as_float())

  m.replace_original_index_miller_indices(cb_op_to_ref.apply(
    integrated_data['miller_index']))

  d.add_column('BATCH', type_table['BATCH']).set_values(
    batch.as_double().as_float())

  lp = integrated_data['lp']
  I_profile = None
  V_profile = None
  I_sum = None
  V_sum = None
  if 'intensity.prf.value' in integrated_data:
    I_profile = integrated_data['intensity.prf.value'] * lp
    V_profile = integrated_data['intensity.prf.variance'] * lp
    d.add_column('IPR', type_table['I']).set_values(I_profile.as_float())
    d.add_column('SIGIPR', type_table['SIGI']).set_values(
      flex.sqrt(V_profile).as_float())
    # Trap negative variances
    assert V_profile.all_gt(0)
  if 'intensity.sum.value' in integrated_data:
    I_sum = integrated_data['intensity.sum.value'] * lp
    V_sum = integrated_data['intensity.sum.variance'] * lp
    d.add_column('I', type_table['I']).set_values(I_sum.as_float())
    d.add_column('SIGI', type_table['SIGI']).set_values(
      flex.sqrt(V_sum).as_float())
    # Trap negative variances
    assert V_sum.all_gt(0)

  d.add_column('FRACTIONCALC', type_table['FRACTIONCALC']).set_values(
    fractioncalc.as_float())

  d.add_column('XDET', type_table['XDET']).set_values(xdet.as_float())
  d.add_column('YDET', type_table['YDET']).set_values(ydet.as_float())
  d.add_column('ROT', type_table['ROT']).set_values(rot.as_float())
  d.add_column('LP', type_table['LP']).set_values(lp.as_float())

  m.write(hklout)

  return m
<commit_msg>Notes on handling partial reflections; will be removed but useful for future reference<commit_after>from __future__ import division

def sum_partial_reflections(integrated_data, min_total_partiality=0.5):
  '''Sum partial reflections; weighted sum for summation integration; weighted
  average for profile fitted reflections. N.B. this will report total
  partiality for the summed reflection.'''

  if not 'partiality' in integrated_data:
    return integrated_data

  # need to do something with these:
  #
  # background.mean     - not used in MTZ
  # background.mse      - not used in MTZ
  # profile.correlation - not used in MTZ
  # xyzobs.mm.value     - not used in MTZ
  # xyzobs.mm.variance  - not used in MTZ
  # xyzobs.px.value     - not used in MTZ
  # xyzobs.px.variance  - not used in MTZ
  #
  # rest of the columns are pretty well defined - these are all uniform for the
  # reflection so can just ignore other values:
  #
  # d
  # id
  # bbox (needs expanding...)
  # entering
  # flags
  # lp
  # miller_index
  # panel
  # partial_id
  # s1
  # xyzcal.mm
  # xyzcal.px
  # zeta
  #
  # Then finally these need a weighted average:
  #
  # intensity.prf.value
  # intensity.prf.variance
  # intensity.sum.value
  # intensity.sum.variance
  # partiality
  #
  # now just need to worry about those that I am actually outputting to the MTZ
  # file...

  trash = { }

  isel = (integrated_data['partiality'] < 0.99).iselection()

  for j in isel:
    p_id = integrated_data['partial_id'][j]
    if not p_id in trash:
      trash[p_id] = []
    trash[p_id].append(j)

  for p_id in sorted(trash):
    if len(trash[p_id]) == 1:
      continue
    p_tot = sum([integrated_data['partiality'][j] for j in trash[p_id]])
    if p_tot >= min_total_partiality:
      i_sum_tot = 0.0
      i_sum_var = 0.0
      i_prf_tot = 0.0
      i_prf_var = 0.0
      w_tot = 0.0
      for j in trash[p_id]:
        print p_id, integrated_data['zeta'][j]
        i_sum_tot += integrated_data['intensity.sum.value'][j]
        i_sum_var += integrated_data['intensity.sum.variance'][j]
        i_prf = integrated_data['intensity.prf.value'][j]
        i_prf_v = integrated_data['intensity.prf.variance'][j]
        p = integrated_data['partiality'][j]
        w = i_prf * i_prf / i_prf_v
        w_tot += w
        i_prf_tot += w * i_prf
        i_prf_var += w * i_prf_v
      i_prf_tot /= w_tot
      i_prf_var /= w_tot
      print p_id, p_tot, i_sum_tot, i_sum_var, i_prf_tot, i_prf_var

  print 1/0
  return integrated_data

def scale_partial_reflections(integrated_data, min_partiality=0.5):
  '''Scale partial reflections (after summation) according to their estimated
  partiality - for profile fitted reflections this will result in no change,
  for summation integrated reflections will be scaled up by 1 / partiality
  with error accordingly scaled. N.B. this will report the scaled up partiality
  for the output reflection.'''

  return integrated_data

def export_mtz(integrated_data, experiment_list, hklout, ignore_panels=False,
               include_partials=False):
  '''Export data from integrated_data corresponding to experiment_list to an
  MTZ file hklout.'''

  from dials.util import log
  from logging import info

  log.config(0, info='dials.export_mtz.log',
             debug='dials.export_mtz.debug.log')

  # for the moment assume (and assert) that we will convert data from exactly
  # one lattice...

  # FIXME allow for more than one experiment in here: this is fine just add
  # multiple MTZ data sets (DIALS1...DIALSN) and multiple batch headers: one
  # range of batches for each experiment

  assert(len(experiment_list) == 1)
  assert(min(integrated_data['id']) == max(integrated_data['id']) == 0)

  # strip out negative variance reflections: these should not really be there
  # FIXME Doing select on summation results. Should do on profile result if
  # present? Yes

  selection = integrated_data['intensity.sum.variance'] <= 0
  if selection.count(True) > 0:
    integrated_data.del_selected(selection)
    info('Removing %d reflections with negative variance' % \
          selection.count(True))

  if 'intensity.prf.variance' in integrated_data:
    selection = integrated_data['intensity.prf.variance'] <= 0
    if selection.count(True) > 0:
      integrated_data.del_selected(selection)
      info('Removing %d profile reflections with negative variance' % \
            selection.count(True))

  # FIXME in here work on including partial reflections => at this stage best
  # to split off the partial refections into a different selection & handle
  # gracefully... better to work on a short list as will need to "pop" them &
  # find matching parts to combine.

  if include_partials:
    integrated_data = sum_partial_reflections(integrated_data)
    integrated_data = scale_partial_reflections(integrated_data)

  if 'partiality' in integrated_data:
    selection = integrated_data['partiality'] < 0.99
    if selection.count(True) > 0:
      integrated_data.del_selected(selection)
      info('Removing %d incomplete reflections' % \
        selection.count(True))

  # FIXME TODO for more than one experiment into an MTZ file:
  #
  # - add an epoch (or recover an epoch) from the scan and add this as an extra
  #   column to the MTZ file for scaling, so we know that the two lattices were
  #   integrated at the same time
  # - decide a sensible BATCH increment to apply to the BATCH value between
  #   experiments and add this
  #
  # At the moment this is probably enough to be working on.

  experiment = experiment_list[0]

  # also only work with one panel(for the moment)

  if not ignore_panels:
    assert(len(experiment.detector) == 1)

  axis = experiment.goniometer.get_rotation_axis()
  s0 = experiment.beam.get_s0()
  wavelength = experiment.beam.get_wavelength()

  from scitbx import matrix

  panel = experiment.detector[0]
  origin = matrix.col(panel.get_origin())
  fast = matrix.col(panel.get_fast_axis())
  slow = matrix.col(panel.get_slow_axis())

  pixel_size = panel.get_pixel_size()

  fast *= pixel_size[0]
  slow *= pixel_size[1]

  cb_op_to_ref = experiment.crystal.get_space_group().info(
    ).change_of_basis_op_to_reference_setting()

  experiment.crystal = experiment.crystal.change_basis(cb_op_to_ref)

  U = experiment.crystal.get_U()
  unit_cell = experiment.crystal.get_unit_cell()

  from iotbx import mtz

  from scitbx.array_family import flex
  from math import floor, sqrt

  m = mtz.object()
  m.set_title('from dials.export_mtz')
  m.set_space_group_info(experiment.crystal.get_space_group().info())

  image_range = experiment.scan.get_image_range()

  for b in range(image_range[0], image_range[1] + 1):
    o = m.add_batch().set_num(b).set_nbsetid(1).set_ncryst(1)
    o.set_time1(0.0).set_time2(0.0).set_title('Batch %d' % b)
    o.set_ndet(1).set_theta(flex.float((0.0, 0.0))).set_lbmflg(0)
    o.set_alambd(wavelength).set_delamb(0.0).set_delcor(0.0)
    o.set_divhd(0.0).set_divvd(0.0)

    # FIXME hard-coded assumption on indealized beam vector below... this may be
    # broken when we come to process data from a non-imgCIF frame
    o.set_so(flex.float(s0)).set_source(flex.float((0, 0, -1)))

    # these are probably 0, 1 respectively, also flags for how many are set, sd
    o.set_bbfac(0.0).set_bscale(1.0)
    o.set_sdbfac(0.0).set_sdbscale(0.0).set_nbscal(0)

    # unit cell (this is fine) and the what-was-refined-flags FIXME hardcoded

    # take time-varying parameters from the *end of the frame* unlikely to
    # be much different at the end - however only exist if time-varying refinement
    # was used
    if experiment.crystal.num_scan_points:
      _unit_cell = experiment.crystal.get_unit_cell_at_scan_point(b)
      _U = experiment.crystal.get_U_at_scan_point(b)
    else:
      _unit_cell = unit_cell
      _U = U

    # FIXME need to get what was refined and what was constrained from the
    # crystal model
    o.set_cell(flex.float(_unit_cell.parameters()))
    o.set_lbcell(flex.int((-1, -1, -1, -1, -1, -1)))
    o.set_umat(flex.float(_U.transpose().elems))

    # get the mosaic spread though today it may not actually be set
    mosaic = experiment.crystal.get_mosaicity()
    o.set_crydat(flex.float([mosaic, 0.0, 0.0, 0.0, 0.0, 0.0,
                             0.0, 0.0, 0.0, 0.0, 0.0, 0.0]))

    o.set_lcrflg(0)
    o.set_datum(flex.float((0.0, 0.0, 0.0)))

    # detector size, distance
    o.set_detlm(flex.float([0.0, panel.get_image_size()[0],
                            0.0, panel.get_image_size()[1],
                            0, 0, 0, 0]))
    o.set_dx(flex.float([panel.get_distance(), 0.0]))

    # goniometer axes and names, and scan axis number, and number of axes, missets
    o.set_e1(flex.float(axis))
    o.set_e2(flex.float((0.0, 0.0, 0.0)))
    o.set_e3(flex.float((0.0, 0.0, 0.0)))
    o.set_gonlab(flex.std_string(('AXIS', '', '')))
    o.set_jsaxs(1)
    o.set_ngonax(1)
    o.set_phixyz(flex.float((0.0, 0.0, 0.0, 0.0, 0.0, 0.0)))

    # scan ranges, axis
    phi_start, phi_range = experiment.scan.get_image_oscillation(b)
    o.set_phistt(phi_start)
    o.set_phirange(phi_range)
    o.set_phiend(phi_start + phi_range)
    o.set_scanax(flex.float(axis))

    # number of misorientation angles
    o.set_misflg(0)

    # crystal axis closest to rotation axis (why do I want this?)
    o.set_jumpax(0)

    # type of data - 1; 2D, 2; 3D, 3; Laue
    o.set_ldtype(2)

  # now create the actual data structures - first keep a track of the columns
  # H K L M/ISYM BATCH I SIGI IPR SIGIPR FRACTIONCALC XDET YDET ROT WIDTH
  # LP MPART FLAG BGPKRATIOS

  from cctbx.array_family import flex as cflex # implicit import
  from cctbx.miller import map_to_asu_isym # implicit import

  # gather the required information for the reflection file

  nref = len(integrated_data['miller_index'])
  x_px, y_px, z_px = integrated_data['xyzcal.px'].parts()

  xdet = flex.double(x_px)
  ydet = flex.double(y_px)
  zdet = flex.double(z_px)

  # compute ROT values
  rot = flex.double([experiment.scan.get_angle_from_image_index(z)
                     for z in zdet])

  # compute BATCH values
  batch = flex.floor(zdet).iround() + 1

  # we're working with full reflections so...
  fractioncalc = flex.double(nref, 1.0)

  # now go for it and make an MTZ file...

  x = m.add_crystal('XTAL', 'DIALS', unit_cell.parameters())
  d = x.add_dataset('FROMDIALS', wavelength)

  # now add column information...

  # FIXME add DIALS_FLAG which can include e.g. was partial etc.

  type_table = {'IPR': 'J', 'BGPKRATIOS': 'R', 'WIDTH': 'R', 'I': 'J',
                'H': 'H', 'K': 'H', 'MPART': 'I', 'L': 'H', 'BATCH': 'B',
                'M_ISYM': 'Y', 'SIGI': 'Q', 'FLAG': 'I', 'XDET': 'R', 'LP': 'R',
                'YDET': 'R', 'SIGIPR': 'Q', 'FRACTIONCALC': 'R', 'ROT': 'R'}

  # derive index columns from original indices with
  #
  # from m.replace_original_index_miller_indices
  #
  # so all that is needed now is to make space for the reflections - fill with
  # zeros...

  m.adjust_column_array_sizes(nref)
  m.set_n_reflections(nref)

  # assign H, K, L, M_ISYM space
  for column in 'H', 'K', 'L', 'M_ISYM':
    d.add_column(column, type_table[column]).set_values(
      flex.double(nref, 0.0).as_float())

  m.replace_original_index_miller_indices(cb_op_to_ref.apply(
    integrated_data['miller_index']))

  d.add_column('BATCH', type_table['BATCH']).set_values(
    batch.as_double().as_float())

  lp = integrated_data['lp']
  I_profile = None
  V_profile = None
  I_sum = None
  V_sum = None
  if 'intensity.prf.value' in integrated_data:
    I_profile = integrated_data['intensity.prf.value'] * lp
    V_profile = integrated_data['intensity.prf.variance'] * lp
    d.add_column('IPR', type_table['I']).set_values(I_profile.as_float())
    d.add_column('SIGIPR', type_table['SIGI']).set_values(
      flex.sqrt(V_profile).as_float())
    # Trap negative variances
    assert V_profile.all_gt(0)
  if 'intensity.sum.value' in integrated_data:
    I_sum = integrated_data['intensity.sum.value'] * lp
    V_sum = integrated_data['intensity.sum.variance'] * lp
    d.add_column('I', type_table['I']).set_values(I_sum.as_float())
    d.add_column('SIGI', type_table['SIGI']).set_values(
      flex.sqrt(V_sum).as_float())
    # Trap negative variances
    assert V_sum.all_gt(0)

  d.add_column('FRACTIONCALC', type_table['FRACTIONCALC']).set_values(
    fractioncalc.as_float())

  d.add_column('XDET', type_table['XDET']).set_values(xdet.as_float())
  d.add_column('YDET', type_table['YDET']).set_values(ydet.as_float())
  d.add_column('ROT', type_table['ROT']).set_values(rot.as_float())
  d.add_column('LP', type_table['LP']).set_values(lp.as_float())

  m.write(hklout)

  return m
<|endoftext|>