<commit_before>"""Describe TestBlock class."""
# pylint: disable=dangerous-default-value,too-many-arguments
from itertools import count

from rotest.common.config import ROTEST_WORK_DIR
from rotest.core.flow_component import (AbstractFlowComponent, MODE_OPTIONAL,
                                        MODE_FINALLY, MODE_CRITICAL,
                                        BlockInput, BlockOutput)

assert MODE_FINALLY
assert MODE_CRITICAL
assert MODE_OPTIONAL

class TestBlock(AbstractFlowComponent):
    """Define TestBlock, which is a part of a test.

    Defines tests that are parts of a greater test. The block is dependent on
    the containing :class:`rotest.core.flow.TestFlow` to lock and pass
    resources to it.

    Test blocks can only be skipped because of test related issues
    (e.g. no resources, test-flow stopped due to failure, intentional skip)
    and can't be skipped on account of 'run_delta' (passed in previous runs),
    tags filtering, etc.

    Declaring 'inputs': assign class fields to instances of BlockInput to
    ask for values for the block (values are passed via common, parametrize,
    previous blocks passing them as outputs, or as requested resources).
    You can pass a default value to BlockInput to assign if non is supplied
    (making it an optional input).

    Declaring 'outputs': assign class fields to instances of BlockOutput to
    share values from the instance (self) to the parent and siblings.
    the block automatically shares the declared outputs after teardown.

    In case the blocks under a flow don't 'connect' properly (a block doesn't
    have its declared output in self.__dict__ or a block doesn't get all its
    inputs from) an error would be raised before the tests start.

    Test authors should subclass TestBlock for their own tests and override
    'inputs' tuple with the names of the fields required for the run of the
    block, and override 'mode' to state the type of the block.

    Attributes:
        resources (tuple): list of the required resources. each item is a
            tuple of (resource_name, resource type, parameters dictionary),
            you can use :func:`rotest.core..request` to create the tuple.
        identifier (number): unique id of the test.
        data (rotest.core.models._data.Data): contain information
            about a test  run.
        logger (logging.Logger): test logger.
        save_state (bool): a flag to determine if storing the states of
            resources is required.
        force_initialize (bool): a flag to determine if the resources will be
            initialized even if their validation succeeds.
        config (AttrDict): dictionary of configurations.
        enable_debug (bool): whether to enable entering ipdb debugging mode
            upon any exception in a test statement.
        resource_manager (ClientResourceManager): client resource manager.
        skip_init (bool): True to skip resources initialize and validation.
        mode (number): running mode code. available modes are:
            CRITICAL: stop test flow on failure or error.
            FINALLY: always run this block, regardless of the others' result.
            OPTIONAL: don't stop test flow on failure (but do so on error),
                failure in this type of block still fails the test-flow.
        TAGS (list): list of tags by which the test may be filtered.
        IS_COMPLEX (bool): if this test is complex (may contain sub-tests).
    """
    IS_COMPLEX = False

    def __init__(self, indexer=count(), base_work_dir=ROTEST_WORK_DIR,
                 save_state=True, force_initialize=False, config=None,
                 parent=None, run_data=None, enable_debug=True,
                 resource_manager=None, skip_init=False, is_main=True,
                 parameters={}):

        super(TestBlock, self).__init__(parent=parent,
                                        config=config,
                                        indexer=indexer,
                                        is_main=is_main,
                                        run_data=run_data,
                                        skip_init=skip_init,
                                        parameters=parameters,
                                        save_state=save_state,
                                        enable_debug=enable_debug,
                                        base_work_dir=base_work_dir,
                                        force_initialize=force_initialize,
                                        resource_manager=resource_manager)

        self.addCleanup(self._share_outputs)

        for input_name, value in self.get_inputs().iteritems():
            if value.is_optional():
                setattr(self, input_name, value.default)

    @classmethod
    def get_name(cls, **parameters):
        """Return test name.

        This method gets gets instantiation arguments that are passed to the
        block via 'parametrize' call, and can be overridden to give unique
        names to blocks.

        Returns:
            str. test name.
        """
        class_name = parameters.get(cls.COMPONENT_NAME_PARAMETER,
                                    cls.__name__)

        method_name = cls.get_test_method_name()
        return '.'.join((class_name, method_name))

    @classmethod
    def get_inputs(cls):
        """Return a list of all the input instances of this block.

        Returns:
            dict. block's inputs (name: input placeholder instance).
        """
        all_inputs = {}
        checked_class = cls
        while checked_class is not TestBlock:
            all_inputs.update({key: value for (key, value) in
                               checked_class.__dict__.iteritems()
                               if isinstance(value, BlockInput)})

            checked_class = checked_class.__bases__[0]

        return all_inputs

    @classmethod
    def get_outputs(cls):
        """Return a list of all the input instances of this block.

        Returns:
            dict. block's inputs (name: input placeholder instance).
        """
        all_outputs = {}
        checked_class = cls
        while checked_class is not TestBlock:
            all_outputs.update({key: value for (key, value) in
                                checked_class.__dict__.iteritems()
                                if isinstance(value, BlockOutput)})

            checked_class = checked_class.__bases__[0]

        return all_outputs

    def _share_outputs(self):
        """Share all the declared outputs of the block."""
        outputs_dict = {}
        for output_name in self.get_outputs():
            if output_name not in self.__dict__:
                raise RuntimeError("Block %r didn't create output %r" %
                                   (self.data.name, output_name))

            outputs_dict[output_name] = getattr(self, output_name)

        self.share_data(**outputs_dict)

    def _validate_inputs(self, extra_inputs=[]):
        """Validate that all the required inputs of the blocks were passed.

        All names under the 'inputs' list must be attributes of the test-block
        when it begins to run, otherwise the block would raise an exception.

        Args:
            extra_inputs (list): fields the component would get from its parent
                or siblings.

        Raises:
            AttributeError: not all inputs were passed to the block.
        """
        required_inputs = [name for (name, value) in self.get_inputs().iteritems()
                           if not value.is_optional()]

        required_inputs.extend(self._pipes.itervalues())

        missing_inputs = [input_name for input_name in required_inputs
                          if (input_name not in self.__dict__ and
                              input_name not in extra_inputs and
                              input_name not in self._pipes)]

        if len(missing_inputs) > 0:
            raise AttributeError("Block %r under %r is missing mandatory "
                                 "inputs %s" %
                                 (self.data.name, self.parent, missing_inputs))
<commit_msg>fixed by pylint<commit_after>"""Describe TestBlock class."""
# pylint: disable=dangerous-default-value,too-many-arguments
from itertools import count

from rotest.common.config import ROTEST_WORK_DIR
from rotest.core.flow_component import (AbstractFlowComponent, MODE_OPTIONAL,
                                        MODE_FINALLY, MODE_CRITICAL,
                                        BlockInput, BlockOutput)

assert MODE_FINALLY
assert MODE_CRITICAL
assert MODE_OPTIONAL

class TestBlock(AbstractFlowComponent):
    """Define TestBlock, which is a part of a test.

    Defines tests that are parts of a greater test. The block is dependent on
    the containing :class:`rotest.core.flow.TestFlow` to lock and pass
    resources to it.

    Test blocks can only be skipped because of test related issues
    (e.g. no resources, test-flow stopped due to failure, intentional skip)
    and can't be skipped on account of 'run_delta' (passed in previous runs),
    tags filtering, etc.

    Declaring 'inputs': assign class fields to instances of BlockInput to
    ask for values for the block (values are passed via common, parametrize,
    previous blocks passing them as outputs, or as requested resources).
    You can pass a default value to BlockInput to assign if non is supplied
    (making it an optional input).

    Declaring 'outputs': assign class fields to instances of BlockOutput to
    share values from the instance (self) to the parent and siblings.
    the block automatically shares the declared outputs after teardown.

    In case the blocks under a flow don't 'connect' properly (a block doesn't
    have its declared output in self.__dict__ or a block doesn't get all its
    inputs from) an error would be raised before the tests start.

    Test authors should subclass TestBlock for their own tests and override
    'inputs' tuple with the names of the fields required for the run of the
    block, and override 'mode' to state the type of the block.

    Attributes:
        resources (tuple): list of the required resources. each item is a
            tuple of (resource_name, resource type, parameters dictionary),
            you can use :func:`rotest.core..request` to create the tuple.
        identifier (number): unique id of the test.
        data (rotest.core.models._data.Data): contain information
            about a test  run.
        logger (logging.Logger): test logger.
        save_state (bool): a flag to determine if storing the states of
            resources is required.
        force_initialize (bool): a flag to determine if the resources will be
            initialized even if their validation succeeds.
        config (AttrDict): dictionary of configurations.
        enable_debug (bool): whether to enable entering ipdb debugging mode
            upon any exception in a test statement.
        resource_manager (ClientResourceManager): client resource manager.
        skip_init (bool): True to skip resources initialize and validation.
        mode (number): running mode code. available modes are:
            CRITICAL: stop test flow on failure or error.
            FINALLY: always run this block, regardless of the others' result.
            OPTIONAL: don't stop test flow on failure (but do so on error),
                failure in this type of block still fails the test-flow.
        TAGS (list): list of tags by which the test may be filtered.
        IS_COMPLEX (bool): if this test is complex (may contain sub-tests).
    """
    IS_COMPLEX = False

    def __init__(self, indexer=count(), base_work_dir=ROTEST_WORK_DIR,
                 save_state=True, force_initialize=False, config=None,
                 parent=None, run_data=None, enable_debug=True,
                 resource_manager=None, skip_init=False, is_main=True,
                 parameters={}):

        super(TestBlock, self).__init__(parent=parent,
                                        config=config,
                                        indexer=indexer,
                                        is_main=is_main,
                                        run_data=run_data,
                                        skip_init=skip_init,
                                        parameters=parameters,
                                        save_state=save_state,
                                        enable_debug=enable_debug,
                                        base_work_dir=base_work_dir,
                                        force_initialize=force_initialize,
                                        resource_manager=resource_manager)

        self.addCleanup(self._share_outputs)

        for input_name, value in self.get_inputs().iteritems():
            if value.is_optional():
                setattr(self, input_name, value.default)

    @classmethod
    def get_name(cls, **parameters):
        """Return test name.

        This method gets gets instantiation arguments that are passed to the
        block via 'parametrize' call, and can be overridden to give unique
        names to blocks.

        Returns:
            str. test name.
        """
        class_name = parameters.get(cls.COMPONENT_NAME_PARAMETER,
                                    cls.__name__)

        method_name = cls.get_test_method_name()
        return '.'.join((class_name, method_name))

    @classmethod
    def get_inputs(cls):
        """Return a list of all the input instances of this block.

        Returns:
            dict. block's inputs (name: input placeholder instance).
        """
        all_inputs = {}
        checked_class = cls
        while checked_class is not TestBlock:
            all_inputs.update({key: value for (key, value) in
                               checked_class.__dict__.iteritems()
                               if isinstance(value, BlockInput)})

            checked_class = checked_class.__bases__[0]

        return all_inputs

    @classmethod
    def get_outputs(cls):
        """Return a list of all the input instances of this block.

        Returns:
            dict. block's inputs (name: input placeholder instance).
        """
        all_outputs = {}
        checked_class = cls
        while checked_class is not TestBlock:
            all_outputs.update({key: value for (key, value) in
                                checked_class.__dict__.iteritems()
                                if isinstance(value, BlockOutput)})

            checked_class = checked_class.__bases__[0]

        return all_outputs

    def _share_outputs(self):
        """Share all the declared outputs of the block."""
        outputs_dict = {}
        for output_name in self.get_outputs():
            if output_name not in self.__dict__:
                raise RuntimeError("Block %r didn't create output %r" %
                                   (self.data.name, output_name))

            outputs_dict[output_name] = getattr(self, output_name)

        self.share_data(**outputs_dict)

    def _validate_inputs(self, extra_inputs=[]):
        """Validate that all the required inputs of the blocks were passed.

        All names under the 'inputs' list must be attributes of the test-block
        when it begins to run, otherwise the block would raise an exception.

        Args:
            extra_inputs (list): fields the component would get from its parent
                or siblings.

        Raises:
            AttributeError: not all inputs were passed to the block.
        """
        required_inputs = [name
                           for (name, value) in self.get_inputs().iteritems()
                           if not value.is_optional()]

        required_inputs.extend(self._pipes.itervalues())

        missing_inputs = [input_name for input_name in required_inputs
                          if (input_name not in self.__dict__ and
                              input_name not in extra_inputs and
                              input_name not in self._pipes)]

        if len(missing_inputs) > 0:
            raise AttributeError("Block %r under %r is missing mandatory "
                                 "inputs %s" %
                                 (self.data.name, self.parent, missing_inputs))
<|endoftext|>


<commit_before>#This file is part of QuTIP.
#
#    QuTIP is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#    QuTIP is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with QuTIP.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright (C) 2011-2012, Paul D. Nation & Robert J. Johansson
#
###########################################################################

from pylab import *
from matplotlib import pyplot, mpl,cm
from mpl_toolkits.mplot3d import Axes3D

from qutip.istests import *

from qutip.Qobj import Qobj

#
# A collection of various visalization functions.
#

# Adopted from the SciPy Cookbook.
def _blob(x, y, w, w_max, area):
    """
    Draws a square-shaped blob with the given area (< 1) at
    the given coordinates.
    """
    hs = sqrt(area) / 2
    xcorners = array([x - hs, x + hs, x + hs, x - hs])
    ycorners = array([y - hs, y - hs, y + hs, y + hs])

    fill(xcorners, ycorners, color=cm.jet(int((w+w_max) * 256 / (2*w_max))))

# Adopted from the SciPy Cookbook.
def hinton(rho):
    """Draws a Hinton diagram for visualizing a density matrix.

    Parameters
    ----------
    rho : qobj
        Input density matrix.

    Raises
    ------
    ValueError
        Input argument is not a quantum object.

    """
    if not isoper(rho):
        raise ValueError("argument must be a quantum operator")

    W = rho.full()

    clf()
    height, width = W.shape

    w_max = 1.25 * max(abs(diag(matrix(W))))
    if w_max <= 0.0:
        w_max = 1.0

    fill(array([0,width,width,0]),array([0,0,height,height]), color=cm.jet(128))
    axis('off')
    axis('equal')
    for x in range(width):
        for y in range(height):
            _x = x+1
            _y = y+1
            if real(W[x,y]) < 0.0:
                _blob(_x - 0.5, height - _y + 0.5,  abs(W[x,y]), w_max, min(1,abs(W[x,y])/w_max))
            else:
                _blob(_x - 0.5, height - _y + 0.5, -abs(W[x,y]), w_max, min(1,abs(W[x,y])/w_max))

def matrix_histogram(M, xlabels, ylabels, title, limits=None, ax=None):
    """
    Draw a histogram for the matrix M, with the given x and y labels and title.

    Parameters
    ----------
    M : Matrix of Qobj
        The matrix to visualize

    xlabels : list of strings
        list of x labels

    ylabels : list of strings
        list of y labels

    title : string
        title of the plot

    limits : list/array with two float numbers
        The z-axis limits [min, max] (optional)

    ax : a matplotlib axes instance
        The axes context in which the plot will be drawn.

    Returns
    -------

        An matplotlib axes instance for the plot.

    Raises
    ------
    ValueError
        Input argument is not valid.

    """

    if isinstance(M, Qobj):
        # extract matrix data from Qobj
        M = M.full()

    n=size(M)
    xpos,ypos=meshgrid(range(M.shape[0]),range(M.shape[1]))
    xpos=xpos.T.flatten()-0.5
    ypos=ypos.T.flatten()-0.5
    zpos = zeros(n)
    dx = dy = 0.8 * ones(n)
    dz = real(M.flatten())

    if limits: # check that limits is a list type
        z_min = limits[0]
        z_max = limits[1]
    else:
        z_min = min(dz)
        z_max = max(dz)

    norm=mpl.colors.Normalize(z_min, z_max)
    cmap=get_cmap('jet') # Spectral
    colors=cmap(norm(dz))

    if ax == None:
        fig = plt.figure()
        ax = Axes3D(fig, azim=-35, elev=35)

    ax.bar3d(xpos, ypos, zpos, dx, dy, dz, color=colors)
    plt.title(title)

    # x axis
    ax.axes.w_xaxis.set_major_locator(IndexLocator(1,-0.5))
    ax.set_xticklabels(xlabels)
    ax.tick_params(axis='x', labelsize=14)

    # y axis
    ax.axes.w_yaxis.set_major_locator(IndexLocator(1,-0.5))
    ax.set_yticklabels(ylabels)
    ax.tick_params(axis='y', labelsize=14)

    # z axis
    ax.axes.w_zaxis.set_major_locator(IndexLocator(1,0.5))
    ax.set_zlim3d([z_min, z_max])

    # color axis
    cax, kw = mpl.colorbar.make_axes(ax, shrink=.75, pad=.02)
    cb1 = mpl.colorbar.ColorbarBase(cax, cmap=cmap, norm=norm)

    return ax

def matrix_histogram_complex(M, xlabels, ylabels, title, limits=None, ax=None):
    """
    Draw a histogram for the amplitudes of matrix M, using the argument of each element
    for coloring the bars, with the given x and y labels and title.

    Parameters
    ----------
    M : Matrix of Qobj
        The matrix to visualize

    xlabels : list of strings
        list of x labels

    ylabels : list of strings
        list of y labels

    title : string
        title of the plot

    limits : list/array with two float numbers
        The z-axis limits [min, max] (optional)

    ax : a matplotlib axes instance
        The axes context in which the plot will be drawn.

    Returns
    -------

        An matplotlib axes instance for the plot.

    Raises
    ------
    ValueError
        Input argument is not valid.

    """

    if isinstance(M, Qobj):
        # extract matrix data from Qobj
        M = M.full()

    n=size(M)
    xpos,ypos=meshgrid(range(M.shape[0]),range(M.shape[1]))
    xpos=xpos.T.flatten()-0.5
    ypos=ypos.T.flatten()-0.5
    zpos = zeros(n)
    dx = dy = 0.8 * ones(n)
    Mvec = M.flatten()
    dz = abs(Mvec)

    if limits: # check that limits is a list type
        z_min = limits[0]
        z_max = limits[1]
    else:
        z_min = -pi
        z_max = pi

    norm=mpl.colors.Normalize(z_min, z_max)
    cmap=get_cmap('jet') # Spectral

    # cdict = {'red':  ((0.0, 0.0, 0.0),
    #                   (0.5, 0.5, 0.5),
    #                   (1.0, 0.0, 0.0)),
    #         'green': ((0.0, 0.0, 0.0),
    #                   (0.5, 0.5, 0.5),
    #                   (1.0, 0.0, 0.0)),
    #         'blue':  ((0.0, 0.0, 0.0),
    #                   (0.5, 0.5, 0.5),
    #                   (1.0, 0.0, 0.0))}
    # cmap = matplotlib.colors.LinearSegmentedColormap('my_colormap',cdict, 256)

    colors=cmap(angle(Mvec))

    if ax == None:
        fig = plt.figure()
        ax = Axes3D(fig, azim=-35, elev=35)

    ax.bar3d(xpos, ypos, zpos, dx, dy, dz, color=colors)
    plt.title(title)

    # x axis
    ax.axes.w_xaxis.set_major_locator(IndexLocator(1,-0.5))
    ax.set_xticklabels(xlabels)
    ax.tick_params(axis='x', labelsize=14)

    # y axis
    ax.axes.w_yaxis.set_major_locator(IndexLocator(1,-0.5))
    ax.set_yticklabels(ylabels)
    ax.tick_params(axis='y', labelsize=14)

    # z axis
    ax.axes.w_zaxis.set_major_locator(IndexLocator(1,0.5))
    ax.set_zlim3d([z_min, z_max])

    # color axis
    cax, kw = mpl.colorbar.make_axes(ax, shrink=.75, pad=.02)
    cb1 = mpl.colorbar.ColorbarBase(cax, cmap=cmap, norm=norm)

    return ax<commit_msg>added cylic colormap for the phases in the plotting function for complex matrix<commit_after>#This file is part of QuTIP.
#
#    QuTIP is free software: you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation, either version 3 of the License, or
#   (at your option) any later version.
#
#    QuTIP is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with QuTIP.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright (C) 2011-2012, Paul D. Nation & Robert J. Johansson
#
###########################################################################

from pylab import *
from matplotlib import pyplot, mpl,cm
from mpl_toolkits.mplot3d import Axes3D

from qutip.istests import *

from qutip.Qobj import Qobj

#
# A collection of various visalization functions.
#

# Adopted from the SciPy Cookbook.
def _blob(x, y, w, w_max, area):
    """
    Draws a square-shaped blob with the given area (< 1) at
    the given coordinates.
    """
    hs = sqrt(area) / 2
    xcorners = array([x - hs, x + hs, x + hs, x - hs])
    ycorners = array([y - hs, y - hs, y + hs, y + hs])

    fill(xcorners, ycorners, color=cm.jet(int((w+w_max) * 256 / (2*w_max))))

# Adopted from the SciPy Cookbook.
def hinton(rho):
    """Draws a Hinton diagram for visualizing a density matrix.

    Parameters
    ----------
    rho : qobj
        Input density matrix.

    Raises
    ------
    ValueError
        Input argument is not a quantum object.

    """
    if not isoper(rho):
        raise ValueError("argument must be a quantum operator")

    W = rho.full()

    clf()
    height, width = W.shape

    w_max = 1.25 * max(abs(diag(matrix(W))))
    if w_max <= 0.0:
        w_max = 1.0

    fill(array([0,width,width,0]),array([0,0,height,height]), color=cm.jet(128))
    axis('off')
    axis('equal')
    for x in range(width):
        for y in range(height):
            _x = x+1
            _y = y+1
            if real(W[x,y]) < 0.0:
                _blob(_x - 0.5, height - _y + 0.5,  abs(W[x,y]), w_max, min(1,abs(W[x,y])/w_max))
            else:
                _blob(_x - 0.5, height - _y + 0.5, -abs(W[x,y]), w_max, min(1,abs(W[x,y])/w_max))

def matrix_histogram(M, xlabels, ylabels, title, limits=None, ax=None):
    """
    Draw a histogram for the matrix M, with the given x and y labels and title.

    Parameters
    ----------
    M : Matrix of Qobj
        The matrix to visualize

    xlabels : list of strings
        list of x labels

    ylabels : list of strings
        list of y labels

    title : string
        title of the plot

    limits : list/array with two float numbers
        The z-axis limits [min, max] (optional)

    ax : a matplotlib axes instance
        The axes context in which the plot will be drawn.

    Returns
    -------

        An matplotlib axes instance for the plot.

    Raises
    ------
    ValueError
        Input argument is not valid.

    """

    if isinstance(M, Qobj):
        # extract matrix data from Qobj
        M = M.full()

    n=size(M)
    xpos,ypos=meshgrid(range(M.shape[0]),range(M.shape[1]))
    xpos=xpos.T.flatten()-0.5
    ypos=ypos.T.flatten()-0.5
    zpos = zeros(n)
    dx = dy = 0.8 * ones(n)
    dz = real(M.flatten())

    if limits: # check that limits is a list type
        z_min = limits[0]
        z_max = limits[1]
    else:
        z_min = min(dz)
        z_max = max(dz)

    norm=mpl.colors.Normalize(z_min, z_max)
    cmap=get_cmap('jet') # Spectral
    colors=cmap(norm(dz))

    if ax == None:
        fig = plt.figure()
        ax = Axes3D(fig, azim=-35, elev=35)

    ax.bar3d(xpos, ypos, zpos, dx, dy, dz, color=colors)
    plt.title(title)

    # x axis
    ax.axes.w_xaxis.set_major_locator(IndexLocator(1,-0.5))
    ax.set_xticklabels(xlabels)
    ax.tick_params(axis='x', labelsize=14)

    # y axis
    ax.axes.w_yaxis.set_major_locator(IndexLocator(1,-0.5))
    ax.set_yticklabels(ylabels)
    ax.tick_params(axis='y', labelsize=14)

    # z axis
    ax.axes.w_zaxis.set_major_locator(IndexLocator(1,0.5))
    ax.set_zlim3d([z_min, z_max])

    # color axis
    cax, kw = mpl.colorbar.make_axes(ax, shrink=.75, pad=.02)
    cb1 = mpl.colorbar.ColorbarBase(cax, cmap=cmap, norm=norm)

    return ax

def matrix_histogram_complex(M, xlabels, ylabels, title, limits=None, ax=None):
    """
    Draw a histogram for the amplitudes of matrix M, using the argument of each element
    for coloring the bars, with the given x and y labels and title.

    Parameters
    ----------
    M : Matrix of Qobj
        The matrix to visualize

    xlabels : list of strings
        list of x labels

    ylabels : list of strings
        list of y labels

    title : string
        title of the plot

    limits : list/array with two float numbers
        The z-axis limits [min, max] (optional)

    ax : a matplotlib axes instance
        The axes context in which the plot will be drawn.

    Returns
    -------

        An matplotlib axes instance for the plot.

    Raises
    ------
    ValueError
        Input argument is not valid.

    """

    if isinstance(M, Qobj):
        # extract matrix data from Qobj
        M = M.full()

    n=size(M)
    xpos,ypos=meshgrid(range(M.shape[0]),range(M.shape[1]))
    xpos=xpos.T.flatten()-0.5
    ypos=ypos.T.flatten()-0.5
    zpos = zeros(n)
    dx = dy = 0.8 * ones(n)
    Mvec = M.flatten()
    dz = abs(Mvec)

    # make small numbers real, to avoid random colors
    idx, = where(abs(Mvec) < 0.001)
    Mvec[idx] = abs(Mvec[idx])

    if limits: # check that limits is a list type
        phase_min = limits[0]
        phase_max = limits[1]
    else:
        phase_min = -pi
        phase_max = pi

    norm=mpl.colors.Normalize(phase_min, phase_max)

    # create a cyclic colormap
    cdict = {'blue': ((0.00, 0.0, 0.0),
                      (0.25, 0.0, 0.0),
                      (0.50, 1.0, 1.0),
                      (0.75, 1.0, 1.0),
                      (1.00, 0.0, 0.0)),
            'green': ((0.00, 0.0, 0.0),
                      (0.25, 1.0, 1.0),
                      (0.50, 0.0, 0.0),
                      (0.75, 1.0, 1.0),
                      (1.00, 0.0, 0.0)),
            'red':   ((0.00, 1.0, 1.0),
                      (0.25, 0.5, 0.5),
                      (0.50, 0.0, 0.0),
                      (0.75, 0.0, 0.0),
                      (1.00, 1.0, 1.0))}
    cmap = matplotlib.colors.LinearSegmentedColormap('phase_colormap', cdict, 256)

    colors = cmap(norm(angle(Mvec)))

    if ax == None:
        fig = plt.figure()
        ax = Axes3D(fig, azim=-35, elev=35)

    ax.bar3d(xpos, ypos, zpos, dx, dy, dz, color=colors)
    plt.title(title)

    # x axis
    ax.axes.w_xaxis.set_major_locator(IndexLocator(1,-0.5))
    ax.set_xticklabels(xlabels)
    ax.tick_params(axis='x', labelsize=14)

    # y axis
    ax.axes.w_yaxis.set_major_locator(IndexLocator(1,-0.5))
    ax.set_yticklabels(ylabels)
    ax.tick_params(axis='y', labelsize=14)

    # z axis
    #ax.axes.w_zaxis.set_major_locator(IndexLocator(1,0.5))
    ax.set_zlim3d([0, 1])

    # color axis
    cax, kw = mpl.colorbar.make_axes(ax, shrink=.75, pad=.02)
    cb = mpl.colorbar.ColorbarBase(cax, cmap=cmap, norm=norm)
    cb.set_ticks([-pi, -pi/2, 0, pi/2, pi])
    cb.set_ticklabels((r'$-\pi$',r'$-\pi/2$',r'$0$',r'$\pi/2$',r'$\pi$'))

    return ax<|endoftext|>


<commit_before>"""
Created on 3 May 2021

@author: Bruno Beloff (bruno.beloff@southcoastscience.com)

https://packaging.python.org/guides/single-sourcing-package-version/
"""

__version__ = '1.0.32'
<commit_msg>Updated version number<commit_after>"""
Created on 3 May 2021

@author: Bruno Beloff (bruno.beloff@southcoastscience.com)

https://packaging.python.org/guides/single-sourcing-package-version/
"""

__version__ = '1.0.33'
<|endoftext|>


<commit_before>"""
Created on 3 May 2021

@author: Bruno Beloff (bruno.beloff@southcoastscience.com)

https://packaging.python.org/guides/single-sourcing-package-version/
"""

__version__ = '1.1.8'
<commit_msg>Updated version number<commit_after>"""
Created on 3 May 2021

@author: Bruno Beloff (bruno.beloff@southcoastscience.com)

https://packaging.python.org/guides/single-sourcing-package-version/
"""

__version__ = '1.1.9'
<|endoftext|>


<commit_before>import sys
import math
import wave
import struct
import pyaudio
import matplotlib.pyplot as plt

# PyAudio object variable
pa = pyaudio.PyAudio()

# Size of each read-in chunk
CHUNK = 1
# Amount of channels of the live recording
CHANNELS = 2
# Sample width of the live recording
WIDTH = 2
# Sample rate in Hz of the live recording
SAMPLE_RATE = 44100

def main():
    # Mode the program was run in
    mode = sys.argv[1]

    # Execute the chosen mode
    if mode == '--file' or mode == '-f':
        filemode()
    elif mode == '--live' or mode == '-l':
        livemode()
    else:
        print('Please either choose file-mode or live-mode')

def filemode():
    # Give some feedback
    print('Now noise-cancelling the file')

    # Read in the given file
    (waveform, stream) = readin(sys.argv[2])

    # Collecting the volume levels in decibels in a list
    decibel_levels = []

    # Counting the iterations of the while-loop
    iteration = 0

    # Read a first chunk and continue to do so for as long as there is a stream to read in
    original = waveform.readframes(CHUNK)
    while original != b'':
        # Invert the original audio
        inverted = invert(original)

        # Play back both audios
        stream.write(original)
        stream.write(inverted)

        # On every 1000th iteration append the difference between the level of the source audio and the inverted one
        if iteration % 1000 == 0:
            # Calculate the difference of the source and the inverted audio
            difference = calculate_decibel(original) - calculate_decibel(inverted)
            # Print the current difference
            print(difference)
            # Append the difference to the list used for the plot
            decibel_levels.append(difference)

        # Read in the next chunk of data
        original = waveform.readframes(CHUNK)

        # Add up one to the iterations
        iteration += 1

    # Outputting feedback regarding the end of the file
    print('Finished noise-cancelling the file')

    # Stop the stream after there is no more data to read
    stream.stop_stream()
    stream.close()

    # Plot the results
    plot_results(decibel_levels)

    # Terminate PyAudio as well as the program
    pa.terminate()
    sys.exit()

def livemode():
    # Start live recording
    print('Now noise-cancelling live')

    # Create a new PyAudio object using the preset constants
    stream = pa.open(
        format=pa.get_format_from_width(WIDTH),
        channels=CHANNELS,
        rate=SAMPLE_RATE,
        frames_per_buffer=CHUNK,
        input=True,
        output=True
    )

    # Collecting the volume levels in decibels in a list
    decibel_levels = []

    # Grab a chunk of data in iterations according to the preset constants
    try:
        for i in range(0, int(SAMPLE_RATE / CHUNK * sys.maxunicode)):
            # Read in a chunk of live audio on each iteration
            original = stream.read(CHUNK)

            # Invert the original audio
            inverted = invert(original)

            # Play back the inverted audio
            stream.write(inverted, CHUNK)

            # On every 1000th iteration append the difference between the level of the source audio and the inverted one
            if i % 3000 == 0:
                # Calculate the difference of the source and the inverted audio
                difference = calculate_decibel(original) - calculate_decibel(inverted)
                # Print the current difference
                print(difference)
                # Append the difference to the list used for the plot
                decibel_levels.append(difference)
    except (KeyboardInterrupt, SystemExit):
        # Outputting feedback regarding the end of the file
        print('Finished noise-cancelling the file')
        # Plot the results
        plot_results(decibel_levels)
        # Terminate the program
        sys.exit()

def readin(file):
    """
    Reads in the given wave file and returns a new PyAudio stream object from it.

    :param file: The path to the file to read in
    :return (waveform, stream): (The actual audio data as a waveform, the PyAudio object for said data)
    """

    # Open the waveform from the command argument
    waveform = wave.open(file, 'r')

    # Load PyAudio and create a useable waveform object
    stream = pa.open(
        format=pa.get_format_from_width(waveform.getsampwidth()),
        channels=waveform.getnchannels(),
        rate=waveform.getframerate(),
        output=True
    )

    # Return the waveform as well as the generated PyAudio stream object
    return waveform, stream

def invert(data):
    """
    Inverts the byte data it received utilizing an XOR operation.

    :param data: A chunk of byte data
    :return inverted: The same size of chunked data inverted bitwise
    """

    # Convert the bytestring into an integer
    intwave = int.from_bytes(data, byteorder='big')
    # Invert the integer
    intwave ^= 2147483647
    # Convert the integer back into a bytestring
    inverted = intwave.to_bytes(4, byteorder='big')
    # Return the inverted audio data
    return inverted

def calculate_decibel(data):
    """
    Calculates the volume level in decibel of the given data

    :param data: A bytestring used to calculate the decibel level
    :return db: The calculated volume level in decibel
    """

    count = len(data) / 2
    form = "%dh" % count
    shorts = struct.unpack(form, data)
    sum_squares = 0.0
    for sample in shorts:
        n = sample * (1.0/32768)
        sum_squares += n * n
    rms = math.sqrt(sum_squares / count) + 0.0001
    db = 20 * math.log10(rms)
    return db

def plot_results(data):
    """
    Plots the list it receives and cuts off the first ten entries to circumvent the plotting of initial silence

    :param data: A list of data to be plotted
    """

    plt.plot(data[10:])
    plt.xlabel('Time (every 1000th byte)')
    plt.ylabel('Volume level (in dB)')
    plt.show()

main()
<commit_msg>The program now expects the second argument of its call to be a number that defines on which iterations plot data will be collected<commit_after>import sys
import math
import wave
import struct
import pyaudio
import matplotlib.pyplot as plt

# PyAudio object variable
pa = pyaudio.PyAudio()

# Size of each read-in chunk
CHUNK = 1
# Amount of channels of the live recording
CHANNELS = 2
# Sample width of the live recording
WIDTH = 2
# Sample rate in Hz of the live recording
SAMPLE_RATE = 44100
# Set how often plot data will be saved (every nth CHUNK) - the lower the number, the more precise the result
NTH_ITERATION = int(sys.argv[2])

def main():
    # Mode the program was run in
    mode = sys.argv[1]

    # Execute the chosen mode
    if mode == '--file' or mode == '-f':
        filemode()
    elif mode == '--live' or mode == '-l':
        livemode()
    else:
        print('Please either choose file-mode or live-mode')

def filemode():
    # Give some feedback
    print('Now noise-cancelling the file')

    # Read in the given file
    (waveform, stream) = readin(sys.argv[3])

    # Collecting the volume levels in decibels in a list
    decibel_levels = []

    # Counting the iterations of the while-loop
    iteration = 0

    # Read a first chunk and continue to do so for as long as there is a stream to read in
    original = waveform.readframes(CHUNK)
    while original != b'':
        # Invert the original audio
        inverted = invert(original)

        # Play back both audios
        stream.write(original)
        stream.write(inverted)

        # On every nth iteration append the difference between the level of the source audio and the inverted one
        if iteration % NTH_ITERATION == 0:
            # Calculate the difference of the source and the inverted audio
            difference = calculate_decibel(original) - calculate_decibel(inverted)
            # Print the current difference
            print(difference)
            # Append the difference to the list used for the plot
            decibel_levels.append(difference)

        # Read in the next chunk of data
        original = waveform.readframes(CHUNK)

        # Add up one to the iterations
        iteration += 1

    # Outputting feedback regarding the end of the file
    print('Finished noise-cancelling the file')

    # Stop the stream after there is no more data to read
    stream.stop_stream()
    stream.close()

    # Plot the results
    plot_results(decibel_levels, NTH_ITERATION)

    # Terminate PyAudio as well as the program
    pa.terminate()
    sys.exit()

def livemode():
    # Start live recording
    print('Now noise-cancelling live')

    # Create a new PyAudio object using the preset constants
    stream = pa.open(
        format=pa.get_format_from_width(WIDTH),
        channels=CHANNELS,
        rate=SAMPLE_RATE,
        frames_per_buffer=CHUNK,
        input=True,
        output=True
    )

    # Collecting the volume levels in decibels in a list
    decibel_levels = []

    # Grab a chunk of data in iterations according to the preset constants
    try:
        for i in range(0, int(SAMPLE_RATE / CHUNK * sys.maxunicode)):
            # Read in a chunk of live audio on each iteration
            original = stream.read(CHUNK)

            # Invert the original audio
            inverted = invert(original)

            # Play back the inverted audio
            stream.write(inverted, CHUNK)

            # On every nth iteration append the difference between the level of the source audio and the inverted one
            if i % NTH_ITERATION == 0:
                # Calculate the difference of the source and the inverted audio
                difference = calculate_decibel(original) - calculate_decibel(inverted)
                # Print the current difference
                print(difference)
                # Append the difference to the list used for the plot
                decibel_levels.append(difference)
    except (KeyboardInterrupt, SystemExit):
        # Outputting feedback regarding the end of the file
        print('Finished noise-cancelling the file')
        # Plot the results
        plot_results(decibel_levels, NTH_ITERATION)
        # Terminate the program
        sys.exit()

def readin(file):
    """
    Reads in the given wave file and returns a new PyAudio stream object from it.

    :param file: The path to the file to read in
    :return (waveform, stream): (The actual audio data as a waveform, the PyAudio object for said data)
    """

    # Open the waveform from the command argument
    waveform = wave.open(file, 'r')

    # Load PyAudio and create a useable waveform object
    stream = pa.open(
        format=pa.get_format_from_width(waveform.getsampwidth()),
        channels=waveform.getnchannels(),
        rate=waveform.getframerate(),
        output=True
    )

    # Return the waveform as well as the generated PyAudio stream object
    return waveform, stream

def invert(data):
    """
    Inverts the byte data it received utilizing an XOR operation.

    :param data: A chunk of byte data
    :return inverted: The same size of chunked data inverted bitwise
    """

    # Convert the bytestring into an integer
    intwave = int.from_bytes(data, byteorder='big')
    # Invert the integer
    intwave ^= 2147483647
    # Convert the integer back into a bytestring
    inverted = intwave.to_bytes(4, byteorder='big')
    # Return the inverted audio data
    return inverted

def calculate_decibel(data):
    """
    Calculates the volume level in decibel of the given data

    :param data: A bytestring used to calculate the decibel level
    :return db: The calculated volume level in decibel
    """

    count = len(data) / 2
    form = "%dh" % count
    shorts = struct.unpack(form, data)
    sum_squares = 0.0
    for sample in shorts:
        n = sample * (1.0/32768)
        sum_squares += n * n
    rms = math.sqrt(sum_squares / count) + 0.0001
    db = 20 * math.log10(rms)
    return db

def plot_results(data, nth_iteration):
    """
    Plots the list it receives and cuts off the first ten entries to circumvent the plotting of initial silence

    :param data: A list of data to be plotted
    :param nth_iteration: Used for the label of the x axis
    """

    plt.plot(data[10:])
    plt.xlabel('Time (every {}th {} byte)'.format(nth_iteration, CHUNK))
    plt.ylabel('Volume level (in dB)')
    plt.show()

main()
<|endoftext|>


<commit_before>import sentry_top

from collections import defaultdict
from nydus.db import create_cluster
from time import time

from django.conf import settings
from django.core.exceptions import ImproperlyConfigured

from sentry.models import Project
from sentry.plugins.base import Plugin

if not getattr(settings, 'SENTRY_TOP', None):
    raise ImproperlyConfigured('You need to configure SENTRY_TOP')

def get_cluster(hosts=None, router='nydus.db.routers.keyvalue.PartitionRouter'):
    if hosts is None:
        hosts = {
            0: {}  # localhost / default
        }

    return create_cluster({
        'engine': 'nydus.db.backends.redis.Redis',
        'router': router,
        'hosts': hosts,
    })

redis = get_cluster(**settings.SENTRY_TOP['redis'])

MINUTES = settings.SENTRY_TOP.get('total_minutes', 15)

class TopPlugin(Plugin):
    author = 'Sentry Team'
    author_url = 'https://github.com/getsentry/sentry-top'
    version = sentry_top.VERSION
    description = 'Tracks active projects ala `top`'
    resource_links = [
        ('Bug Tracker', 'https://github.com/getsentry/sentry-top/issues'),
        ('Source', 'https://github.com/getsentry/sentry-top'),
    ]

    slug = 'top'
    title = 'Top'
    conf_title = title
    conf_key = 'top'

    def add_event(self, project, client=redis):
        minute = int(time() / 60)
        keys = [
            # 'stop:e:{0}:{1}'.format(event.group_id),
            'stop:p:{0}'.format(minute),
        ]
        with client.map() as conn:
            for key in keys:
                conn.zincrby(key, project.id)
                conn.expire(key, (MINUTES + 1) * 60)

    def top_projects(self, minutes=15, num=100, client=redis):
        now = int(time() / 60)
        keys = []
        for minute in xrange(minutes):
            keys.append('stop:p:{0}'.format(now - minute))

        counts = []
        with client.map() as conn:
            for key in keys:
                counts.append(conn.zrevrange(key, 0, num, withscores=True))

        results = defaultdict(int)
        for countset in counts:
            for project_id, count in countset:
                results[int(project_id)] += int(count)

        sorted_results = sorted(
            results.items(), key=lambda x: x[1], reverse=True)[:num]

        project_map = dict(
            (p.id, p) for p in Project.objects.filter(id__in=[
                p_id for p_id, _ in sorted_results
            ]).select_related('team')
        )

        return [
            (project_map[p_id], c)
            for (p_id, c) in sorted_results
            if p_id in project_map
        ]

        return sorted_results[:num]

    def is_rate_limited(self, project):
        # TODO(dcramer): we need a way to hook into Sentry at event input
        # that guarantees this stat
        self.add_event(project)
<commit_msg>Cannot enable for projects<commit_after>import sentry_top

from collections import defaultdict
from nydus.db import create_cluster
from time import time

from django.conf import settings
from django.core.exceptions import ImproperlyConfigured

from sentry.models import Project
from sentry.plugins.base import Plugin

if not getattr(settings, 'SENTRY_TOP', None):
    raise ImproperlyConfigured('You need to configure SENTRY_TOP')

def get_cluster(hosts=None, router='nydus.db.routers.keyvalue.PartitionRouter'):
    if hosts is None:
        hosts = {
            0: {}  # localhost / default
        }

    return create_cluster({
        'engine': 'nydus.db.backends.redis.Redis',
        'router': router,
        'hosts': hosts,
    })

redis = get_cluster(**settings.SENTRY_TOP['redis'])

MINUTES = settings.SENTRY_TOP.get('total_minutes', 15)

class TopPlugin(Plugin):
    author = 'Sentry Team'
    author_url = 'https://github.com/getsentry/sentry-top'
    version = sentry_top.VERSION
    description = 'Tracks active projects ala `top`'
    resource_links = [
        ('Bug Tracker', 'https://github.com/getsentry/sentry-top/issues'),
        ('Source', 'https://github.com/getsentry/sentry-top'),
    ]

    slug = 'top'
    title = 'Top'
    conf_title = title
    conf_key = 'top'

    def can_enable_for_projects(self):
        return False

    def add_event(self, project, client=redis):
        minute = int(time() / 60)
        keys = [
            # 'stop:e:{0}:{1}'.format(event.group_id),
            'stop:p:{0}'.format(minute),
        ]
        with client.map() as conn:
            for key in keys:
                conn.zincrby(key, project.id)
                conn.expire(key, (MINUTES + 1) * 60)

    def top_projects(self, minutes=15, num=100, client=redis):
        now = int(time() / 60)
        keys = []
        for minute in xrange(minutes):
            keys.append('stop:p:{0}'.format(now - minute))

        counts = []
        with client.map() as conn:
            for key in keys:
                counts.append(conn.zrevrange(key, 0, num, withscores=True))

        results = defaultdict(int)
        for countset in counts:
            for project_id, count in countset:
                results[int(project_id)] += int(count)

        sorted_results = sorted(
            results.items(), key=lambda x: x[1], reverse=True)[:num]

        project_map = dict(
            (p.id, p) for p in Project.objects.filter(id__in=[
                p_id for p_id, _ in sorted_results
            ]).select_related('team')
        )

        return [
            (project_map[p_id], c)
            for (p_id, c) in sorted_results
            if p_id in project_map
        ]

        return sorted_results[:num]

    def is_rate_limited(self, project):
        # TODO(dcramer): we need a way to hook into Sentry at event input
        # that guarantees this stat
        self.add_event(project)
<|endoftext|>


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

# Copyright (c) 2015-2016 MIT Probabilistic Computing Project

# 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.

import matplotlib.cm
import matplotlib.colors
import matplotlib.pyplot as plt
import numpy as np
import sklearn.datasets

from cgpm.pca.ppca import PPCA
from cgpm.utils import general as gu
from cgpm.utils import config as cu

from scipy.stats import multivariate_normal

rng = gu.gen_rng(12)

# Number of datapoints.
N = 100

# First principal component vector.
W = np.asarray([
    [3.,],
    [20.],
    [-4.],])

# First principal component scores.
X = rng.normal(size=(1,N))

# Mean vector.
m = np.asarray([
    [1],
    [2],
    [-3],])
M = np.repeat(m, repeats=N, axis=1)

# Noise vector
sigma = 1.9
e = rng.normal(scale=sigma, size=(3,N))

# Observations.
Y = np.dot(W,X) + M + e

assert Y.shape == (3, N)

fig, ax = plt.subplots()

from sklearn.decomposition import FactorAnalysis
fa = FactorAnalysis(n_components=2)
fa.fit(Y.T)

# Hard code posterior.
def transform(fa, x):
    dot, inv = np.dot, np.linalg.inv
    L, D = fa.components_.shape
    mu = fa.mean_
    assert mu.shape == (D,)
    Phi = np.diag(fa.noise_variance_)
    assert Phi.shape == (D, D)
    W = fa.components_.T
    assert W.shape == (D, L)
    I = np.eye(L)
    # Compute using Murphy.
    S1 = inv((I + dot(W.T, dot(inv(Phi), W))))
    m1 = dot(S1, dot(W.T, dot(inv(Phi), (x-mu))))
    # Compute using the Schur complement.
    S2 = I - dot(dot(W.T, inv(dot(W,W.T) + Phi)), W)
    m2 = dot(dot(W.T, inv(dot(W, W.T)+Phi)), (x-mu))
    assert np.allclose(m1, m2)
    assert np.allclose(S1, S2)
    print m1, m2
    return m1

def logp(fa, x):
    dot, inv = np.dot, np.linalg.inv
    L, D = fa.components_.shape
    mu = fa.mean_
    assert mu.shape == (D,)
    Phi = np.diag(fa.noise_variance_)
    assert Phi.shape == (D, D)
    W = fa.components_.T
    assert W.shape == (D, L)
    return multivariate_normal.logpdf(x, mu, Phi + np.dot(W, W.T))

for x in Y.T:
    print x
    # Compute the latent variables z.
    z_a = fa.transform(np.asarray([x]))
    z_b = transform(fa, x)
    assert np.allclose(z_a, z_b)
    # Compute the probability p(x), marginalizing over the latent variables.
    lp_a = fa.score(np.asarray([x]))
    lp_b = logp(fa, x)
    assert np.allclose(lp_a, lp_b)

def mvn_marginalize(mu, cov, query, evidence):
    # Q = [0]
    # E = [1,3,4]
    # QE = Q+E
    mu = np.random.rand(5)
    cov = np.outer(mu, mu)
    # Query and evidence indices.
    Q = query
    E = evidence
    # QE = Q+E
    # muQ = mu[Q]
    # muE = mu[E]
    # muQE = mu[Q+E]
    covQ = cov[Q][:,Q]
    covE = cov[E][:,E]
    covJ = cov[Q][:,E]
    top = np.column_stack((covQ, covJ))
    bottom = np.column_stack((covJ.T, covE))
    covQE = np.row_stack((top, bottom))
    assert np.allclose(covQE, covQE.T)
    return mu[Q], mu[E], covQ, covE, covJ

def mvn_condition(mu, cov, query, evidence):
    # assert isinstance(query, list)
    # assert isinstance(evidence, dict)
    # assert len(mu) == cov.shape[0] == cov.shape[1]
    # assert len(query) + len(evidence) <= len(mu)

    Q, E = sorted(query), sorted(evidence.keys())
    Ev = np.asarray([evidence[e] for e in E])
    muQ, muE, covQ, covE, covJ = mvn_marginalize(mu, cov, Q, E)

    P = np.dot(covJ, np.linalg.inv(covE))

    muZ = muQ + np.dot(P, Ev - muE)
    covZ = covQ - np.dot(P, covJ.T)
<commit_msg>Add some plotting to line.py before convering to test.<commit_after># -*- coding: utf-8 -*-

# Copyright (c) 2015-2016 MIT Probabilistic Computing Project

# 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.

import matplotlib.cm
import matplotlib.colors
import matplotlib.pyplot as plt
import numpy as np
import sklearn.datasets

from cgpm.pca.ppca import PPCA
from cgpm.utils import general as gu
from cgpm.utils import config as cu

from scipy.stats import multivariate_normal

rng = gu.gen_rng(12)

# Number of datapoints.
N = 100
# First principal component vector.
W = np.asarray([
    [3.,],
    [20.],
    [-4.],])
# First principal component scores.
X = rng.normal(size=(1,N))
# Mean vector.
m = np.asarray([
    [1],
    [2],
    [-3],])
M = np.repeat(m, repeats=N, axis=1)
# Noise vector
sigma = 1.9
e = rng.normal(scale=sigma, size=(3,N))
# Observations.
Y = np.dot(W,X) + M + e
assert Y.shape == (3, N)

from sklearn.decomposition import FactorAnalysis
fa = FactorAnalysis(n_components=1)
fa.fit(Y.T)

# Hard code posterior.
def transform(fa, x):
    dot, inv = np.dot, np.linalg.inv
    L, D = fa.components_.shape
    mu = fa.mean_
    assert mu.shape == (D,)
    Phi = np.diag(fa.noise_variance_)
    assert Phi.shape == (D, D)
    W = fa.components_.T
    assert W.shape == (D, L)
    I = np.eye(L)
    # Compute using Murphy.
    S1 = inv((I + dot(W.T, dot(inv(Phi), W))))
    m1 = dot(S1, dot(W.T, dot(inv(Phi), (x-mu))))
    # Compute using the Schur complement.
    S2 = I - dot(dot(W.T, inv(dot(W,W.T) + Phi)), W)
    m2 = dot(dot(W.T, inv(dot(W, W.T)+Phi)), (x-mu))
    assert np.allclose(m1, m2)
    assert np.allclose(S1, S2)
    print m1, m2
    return m1

def logp(fa, x):
    dot, inv = np.dot, np.linalg.inv
    L, D = fa.components_.shape
    mu = fa.mean_
    assert mu.shape == (D,)
    Phi = np.diag(fa.noise_variance_)
    assert Phi.shape == (D, D)
    W = fa.components_.T
    assert W.shape == (D, L)
    return multivariate_normal.logpdf(x, mu, Phi + np.dot(W, W.T))

for x in Y.T:
    print x
    # Compute the latent variables z.
    z_a = fa.transform(np.asarray([x]))
    z_b = transform(fa, x)
    assert np.allclose(z_a, z_b)
    # Compute the probability p(x), marginalizing over the latent variables.
    lp_a = fa.score(np.asarray([x]))
    lp_b = logp(fa, x)
    assert np.allclose(lp_a, lp_b)

def mvn_marginalize(mu, cov, query, evidence):
    # Q = [0]
    # E = [1,3,4]
    # QE = Q+E
    mu = np.random.rand(5)
    cov = np.outer(mu, mu)
    # Query and evidence indices.
    Q = query
    E = evidence
    # QE = Q+E
    # muQ = mu[Q]
    # muE = mu[E]
    # muQE = mu[Q+E]
    covQ = cov[Q][:,Q]
    covE = cov[E][:,E]
    covJ = cov[Q][:,E]
    top = np.column_stack((covQ, covJ))
    bottom = np.column_stack((covJ.T, covE))
    covQE = np.row_stack((top, bottom))
    assert np.allclose(covQE, covQE.T)
    return mu[Q], mu[E], covQ, covE, covJ

def mvn_condition(mu, cov, query, evidence):
    # assert isinstance(query, list)
    # assert isinstance(evidence, dict)
    # assert len(mu) == cov.shape[0] == cov.shape[1]
    # assert len(query) + len(evidence) <= len(mu)
    Q, E = sorted(query), sorted(evidence.keys())
    Ev = np.asarray([evidence[e] for e in E])
    muQ, muE, covQ, covE, covJ = mvn_marginalize(mu, cov, Q, E)
    P = np.dot(covJ, np.linalg.inv(covE))
    muZ = muQ + np.dot(P, Ev - muE)
    covZ = covQ - np.dot(P, covJ.T)

# Number of datapoints.
N = 100
# First principal component vector.
W = np.asarray([
    [3.,],
    [20.]])
# First principal component scores.
X = rng.normal(size=(1,N))
# Mean vector.
m = np.asarray([
    [1],
    [2]])
M = np.repeat(m, repeats=N, axis=1)
# Noise vector
e = np.column_stack((
    rng.normal(scale=1.9, size=N), rng.normal(scale=3, size=N))).T
# Observations.
Y = np.dot(W,X) + M + e
assert Y.shape == (2, N)

# Plot the raw data.
fig, ax = plt.subplots()
ax.scatter(Y.T[:,0], Y.T[:,1], color='b')
ax.grid()

from sklearn.decomposition import FactorAnalysis
fa = FactorAnalysis(n_components=1)
fa.fit(Y.T)

# Plot the components
xs = np.asarray(ax.get_xlim())
w0, w1 = fa.components_[0]
ys = w1/w0*xs
ax.plot(xs, ys, color='r')
<|endoftext|>


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

#Author: Mathias Hauser
#Date: 14.12.2014

import netCDF4
import numpy as np
from glob import glob

from collections import OrderedDict
# SUBCLASS Dataset and Variable
# ---------------------------------------------------------------------------

def _select(self, item):

    if len(item) > 3:
        raise IndexError("selector for dimension %i (%s) has \
            too many elements" % ("XXXX", "XXX"))

    data = self.variable[:]

    if data is None:
        raise RuntimeError("dimension XYZ has no data")

    start = wherenearest(data, item[0])
    stop = start if len(item) < 2 else wherenearest(data, item[1])
    step = 1 if len(item) < 3 else item[2]

    # reverse start stop if dimen is sorted the other way round
    if start > stop:
        start, stop = stop, start
    # python slice notation: stop is the first that is NOT selected
    sel = slice(start, stop + 1, step)

    vals = data[sel]
    verbose(self.name, item, vals[0], stop_val=vals[-1])
    return sel

# ---------------------------------------------------------------------------

def verbose(name, item, start_val, stop_val):
    """print item and resulting selection"""
    string = "Select '{0}': ".format(name)

    if len(item) == 1:
        msg = "{0}{1} selects: {2:5.3f}".format(string, item[0], start_val)
    elif len(item) == 2:
        msg = "{0}{1}...{2} selects: {3:5.3f}...{4:5.3f}".format(string,
         item[0], item[1], start_val, stop_val)

    print(msg)

# ============================================================================

class Select(object):
    """docstring for select"""
    def __init__(self, name, dimension, variable):
        super(Select, self).__init__()

        self.name = name
        self.dimension = dimension
        self.variable = variable

        # se
        self.selections = OrderedDict()

    def __getitem__(self, item):
        # check if this particular case was already selected
        sel = self.selections.get(item, None)

        if sel is not None:
            return sel
        else:
            # get the new selection
            sel = _select(self, item)

            # assign the new selection to the dict
            self.selections[item] = sel
            return sel

    # selections can not be assigned
    def __setitem__(self, item, value):
        raise RuntimeError("Coordinate Subscripts can not be set manually")

# ============================================================================

def wherenearest(grid, pos):
    """return index nearest to pos"""
    val = np.abs(grid - pos)
    IDX = np.where(val == np.min(val))
    return IDX[0][0]

# ============================================================================
# subclass the netCDF4.Dataset class in order to (1) use the new Variable
# class and (2) the Select class

class Dataset(netCDF4.Dataset):
    """subclass of netCDF4.Dataset that uses the """

    #select = dict()

    def __init__(self, *arg, **kwargs):

        try:
            super(Dataset, self).__init__(*arg, **kwargs)
        except RuntimeError:
            raise RuntimeError("No such file or directory '%s'" % arg[0])

        select = kwargs.pop('select', None)
        if select is None:
            self.__dict__['select'] = OrderedDict()
        else:
            self.__dict__['select'] = select

        for var in self.variables.keys():
            ncv = self.variables[var]
            # unfortunately we have to reassign "variables" in order to
            # use the new version
            self.variables[var] = Variable(ncv.group(),
                                           ncv._name,
                                           ncv.datatype,
                                           ncv.dimensions,
                                           id=ncv._varid,
                                           select_parent=self.select
                                           )

            # add the new Select class to the Dataset
            self.select[var] = Select(var,
                                      self.dimensions.get(var, [None]),
                                      self.variables.get(var, [None])
                                      )

# ============================================================================
# subclass the netCDF4.MFDataset class in order to (1) use the new Variable
# class and (2) the Select class

class MFDataset(netCDF4.MFDataset):
    """subclass of netCDF4.Dataset that uses the """

#    select = dict()

    def __init__(self, *arg, **kwargs):

        try:
            super(MFDataset, self).__init__(*arg, **kwargs)
        except RuntimeError:
            raise RuntimeError("No such file or directory '%s'" % arg[0])

        self.__dict__['select'] = OrderedDict()

        for var in self.variables.keys():
            ncv = self.variables[var]
            # unfortunately we have to reassign "variables" in order to
            # use the new version

            if isinstance(ncv, netCDF4.Variable):
                self._vars[var] = Variable(ncv.group(),
                                           ncv._name,
                                           ncv.datatype,
                                           ncv.dimensions,
                                           id=ncv._varid,
                                           select_parent=self.select
                                           )
            else:
                self._vars[var] = _Variable(self,
                                            ncv._name,
                                            ncv._mastervar,
                                            ncv._recdimname,
                                            select_parent=self.select
                                            )

            # add the new Select class to the Dataset
            self.select[var] = Select(var,
                                      self.dimensions.get(var, [None]),
                                      self.variables.get(var, [None])
                                      )

# # ============================================================================

class MSFDataset(object):
    """docstring for MSFDataset"""
    def __init__(self, files, check=False, aggdim=None, exclude=[]):
        super(MSFDataset, self).__init__()

        # Open the master file in the base class, so that the CDFMF instance
        # can be used like a CDF instance.
        if isinstance(files, str):
            if files.startswith('http'):
                msg='cannot using file globbing for remote (OPeNDAP) datasets'
                raise ValueError(msg)
            else:
                files = sorted(glob(files))

        master = files[0]
        print(master)
        print(files)

        cdfm = Dataset(master)
        # copy attributes from master.
        for name, value in cdfm.__dict__.items():
            self.__dict__[name] = value

        self._files = files

        ncfs = []
        for _file in files[2:]:
            ncfs.append(Dataset(_file, select=self.select))

        self._ncfs = ncfs

    def __getitem__(self, key):
        return self._ncfs[key]

    def __iter__(self):
        return self._ncfs.__iter__()

    def __len__(self):
        return self._ncfs.__len__()

# ============================================================================

# subclass netCDF4.Variable to alter __getitem__

def __expand_elem__(elem, ndim):
    """parse the slice input"""

    # PARSE ELEM
    # we need to be sure of the position of all elem

    # (1) only 1 elem is given
    if not isinstance(elem, tuple):
        elem = (elem, )
    # if type(elem) is list or isinstance(elem, set):
    #     elem = (elem,)

    # if not np.iterable(elem):
    #     elem = (elem, )

    # search for named elements (dicts)
    had_dict = False

    _elem = ()
    _dict = dict()

    for i, el in enumerate(elem):
        if isinstance(el, dict):
            had_dict = True

            for key in el.keys():
                _dict[key] = el[key]

        elif had_dict:
            raise KeyError('dict argument can not be followed by positional argument')

        else:
            _elem = _elem + (el, )

    elem = _elem

    # (2) make sure there are not too many dimensions in slice.
    if len(elem) > ndim:
        raise ValueError("slicing expression exceeds the number of dimensions \
            of the variable")

    # (2) if less elements are given then there are dimensions:
    #     fill them with Ellipsis

    missing_dim = ndim - len(elem)

    if missing_dim > 0:
        # the first occurence of Ellipsis gets expanded
        if Ellipsis in elem:
            # first occurence of Ellipsis
            i = next(el[0] for el in enumerate(elem) if el[1] == Ellipsis)
            # expand
            elem = elem[:i+1] + (slice(None), ) * missing_dim + elem[i+1:]
        else:
            elem = elem + (slice(None), ) * missing_dim

    return elem, _dict

# ----------------------------------------------------------------------------

def __parse_el__(self, elem, _dict):
    """find slices that have to be selected"""

    sel_elem = list()

    for i, el in enumerate(elem):

        # coordinate subsetting happens for sets
        if isinstance(el, set):
            # convert to tuple (must be hashable)
            el = tuple(el)

            if len(el) != 2 and len(el) != 1:
                raise IndexError("Length of coordinate subsetting must be 1 or 2")

            # name of the dimension
            dim = self.dimensions[i]

            # need "Select" of the parent of the Variable
            # get the slice for this specific selection
            sel_elem.append(self._select_parent[dim][el])
        else:
            sel_elem.append(el)

        for i, dim in enumerate(self.dimensions):
            el = _dict.pop(dim, None)

            if el is not None:
                sel_elem[i] = self._select_parent[dim][el]

        for key in _dict:
            print('ignored key: {key}'.format(key=key))

    return tuple(sel_elem)

# ----------------------------------------------------------------------------

class Variable(netCDF4.Variable):
    """subclass netCDF4 to alter __getitem__"""
    def __init__(self, *arg, **kwargs):
        super(Variable, self).__init__(*arg, **kwargs)
        self.__dict__['_select_parent'] = kwargs.pop('select_parent')

    def __getitem__(self, elem):

        # add Ellipsis if elem has less members than ndim
        elem, _dict = __expand_elem__(elem, self.ndim)

        # find slices that have to be selected and select
        sel_elem = __parse_el__(self, elem, _dict)

        data = super(Variable, self).__getitem__(sel_elem)

        return data

# ============================================================================

class _Variable(netCDF4._Variable):
    """subclass netCDF4 to alter __getitem__"""

    def __init__(self, *arg, **kwargs):
        self._select_parent = kwargs.pop('select_parent')
        super(_Variable, self).__init__(*arg, **kwargs)

    def __getitem__(self, elem):

        # add Ellipsis if elem has less members than ndim
        elem, _dict = __expand_elem__(elem, self.ndim)

        # find slices that have to be selected and select
        sel_elem = __parse_el__(self, elem, _dict)

        data = super(_Variable, self).__getitem__(sel_elem)

        return data

# ============================================================================

fN = '/net/exo/landclim/mathause/cesm_data/f.e121.FC5.f19_g16.CTRL_2000-io384.001/lnd/hist/f.e121.FC5.f19_g16.CTRL_2000-io384.001.clm2.h0.0001-01.nc'

ncf = Dataset(fN)

#ncf = ncp.Dataset(fN)

#print(ncf.variables['SOILLIQ'][Ellipsis, 1:3, {'lat' : (1, 3)}].shape)

print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.select['lat'][(3, 5)])
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][0, 5, slice(None), slice(None)].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][0, 5].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][:, [0, 5]].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][:].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][:].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][0].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][Ellipsis].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][0, {0, 0.1}].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][0, {0, 0.1}, {0, 30}, ...].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILICE'][0, {0, 0.1}, {0, 30}, ...].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][{0}].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print('??????????????')
print(ncf.variables['lat'][{3}, {'lat' : (5, 10), 'lat' : (20, 30), 'aglkn' : (1, 10)}])
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
# print(ncf.variables['lat'][{3}, 7, {'lat' : (5, 10)}])
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')

fN = '/net/exo/landclim/mathause/cesm_data/f.e121.FC5.f19_g16.CTRL_2000-io384.001/lnd/hist/f.e121.FC5.f19_g16.CTRL_2000-io384.001.clm2.h0.0005-01.nc'
ncf = Dataset(fN)

print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][0, {0, 0.1}, {0, 30}, ...].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')

fN = '/net/exo/landclim/mathause/cesm_data/f.e121.FC5.f19_g16.CTRL_2000-io384.001/lnd/hist/f.e121.FC5.f19_g16.CTRL_2000-io384.001.clm2.h0.000?-01.nc'
ncf = MFDataset(fN)

print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][:, {0, 0.1}, {0, 30}, ...].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')

# print(ncf.variables['SOILLIQ'][{'lat' : (3, 15)}].shape)
# print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')

ncf = MSFDataset(fN)<commit_msg>better verbose messages<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

#Author: Mathias Hauser
#Date: 14.12.2014

import netCDF4
import numpy as np
from glob import glob

from collections import OrderedDict
# SUBCLASS Dataset and Variable
# ---------------------------------------------------------------------------

def _select(self, item):

    if len(item) > 3:
        raise IndexError("selector for dimension %i (%s) has \
            too many elements" % ("XXXX", "XXX"))

    data = self.variable[:]

    if data is None:
        raise RuntimeError("dimension XYZ has no data")

    start = wherenearest(data, item[0])
    stop = start if len(item) < 2 else wherenearest(data, item[1])
    step = 1 if len(item) < 3 else item[2]

    # reverse start stop if dimen is sorted the other way round
    if start > stop:
        start, stop = stop, start
    # python slice notation: stop is the first that is NOT selected
    sel = slice(start, stop + 1, step)

    vals = data[sel]
    print_selection(self.name, item, vals[0], stop_val=vals[-1])
    return sel

# ---------------------------------------------------------------------------

def print_selection(name, item, start_val, stop_val):
    """print item and resulting selection"""
    string = "Select '{0}': ".format(name)

    if len(item) == 1:
        msg = "{0}{1} selects: {2:5.3f}".format(string, item[0], start_val)
    elif len(item) == 2:
        msg = "{0}{1}...{2} selects: {3:5.3f}...{4:5.3f}".format(string,
         item[0], item[1], start_val, stop_val)

    print(msg)

# ============================================================================

class Select(object):
    """docstring for select"""
    def __init__(self, name, dimension, variable):
        super(Select, self).__init__()

        self.name = name
        self.dimension = dimension
        self.variable = variable

        # se
        self.selections = OrderedDict()

    def __getitem__(self, item):
        # check if this particular case was already selected
        sel = self.selections.get(item, None)

        if sel is not None:
            return sel
        else:
            # get the new selection
            sel = _select(self, item)

            # assign the new selection to the dict
            self.selections[item] = sel
            return sel

    # selections can not be assigned
    def __setitem__(self, item, value):
        raise RuntimeError("Coordinate Subscripts can not be set manually")

# ============================================================================

def wherenearest(grid, pos):
    """return index nearest to pos"""
    val = np.abs(grid - pos)
    IDX = np.where(val == np.min(val))
    return IDX[0][0]

# ============================================================================
# subclass the netCDF4.Dataset class in order to (1) use the new Variable
# class and (2) the Select class

class Dataset(netCDF4.Dataset):
    """subclass of netCDF4.Dataset that uses the """

    #select = dict()

    def __init__(self, *arg, **kwargs):

        try:
            super(Dataset, self).__init__(*arg, **kwargs)
        except RuntimeError:
            raise RuntimeError("No such file or directory '%s'" % arg[0])

        select = kwargs.pop('select', None)
        if select is None:
            self.__dict__['select'] = OrderedDict()
        else:
            self.__dict__['select'] = select

        for var in self.variables.keys():
            ncv = self.variables[var]
            # unfortunately we have to reassign "variables" in order to
            # use the new version
            self.variables[var] = Variable(ncv.group(),
                                           ncv._name,
                                           ncv.datatype,
                                           ncv.dimensions,
                                           id=ncv._varid,
                                           select_parent=self.select
                                           )

            # add the new Select class to the Dataset
            self.select[var] = Select(var,
                                      self.dimensions.get(var, [None]),
                                      self.variables.get(var, [None])
                                      )

# ============================================================================
# subclass the netCDF4.MFDataset class in order to (1) use the new Variable
# class and (2) the Select class

class MFDataset(netCDF4.MFDataset):
    """subclass of netCDF4.Dataset that uses the """

#    select = dict()

    def __init__(self, *arg, **kwargs):

        try:
            super(MFDataset, self).__init__(*arg, **kwargs)
        except RuntimeError:
            raise RuntimeError("No such file or directory '%s'" % arg[0])

        self.__dict__['select'] = OrderedDict()

        for var in self.variables.keys():
            ncv = self.variables[var]
            # unfortunately we have to reassign "variables" in order to
            # use the new version

            if isinstance(ncv, netCDF4.Variable):
                self._vars[var] = Variable(ncv.group(),
                                           ncv._name,
                                           ncv.datatype,
                                           ncv.dimensions,
                                           id=ncv._varid,
                                           select_parent=self.select
                                           )
            else:
                self._vars[var] = _Variable(self,
                                            ncv._name,
                                            ncv._mastervar,
                                            ncv._recdimname,
                                            select_parent=self.select
                                            )

            # add the new Select class to the Dataset
            self.select[var] = Select(var,
                                      self.dimensions.get(var, [None]),
                                      self.variables.get(var, [None])
                                      )

# # ============================================================================

class MSFDataset(object):
    """docstring for MSFDataset"""
    def __init__(self, files, check=False, aggdim=None, exclude=[]):
        super(MSFDataset, self).__init__()

        # Open the master file in the base class, so that the CDFMF instance
        # can be used like a CDF instance.
        if isinstance(files, str):
            if files.startswith('http'):
                msg='cannot using file globbing for remote (OPeNDAP) datasets'
                raise ValueError(msg)
            else:
                files = sorted(glob(files))

        master = files[0]
        print(master)
        print(files)

        cdfm = Dataset(master)
        # copy attributes from master.
        for name, value in cdfm.__dict__.items():
            self.__dict__[name] = value

        self._files = files

        ncfs = []
        for _file in files[2:]:
            ncfs.append(Dataset(_file, select=self.select))

        self._ncfs = ncfs

    def __getitem__(self, key):
        return self._ncfs[key]

    def __iter__(self):
        return self._ncfs.__iter__()

    def __len__(self):
        return self._ncfs.__len__()

# ============================================================================

# subclass netCDF4.Variable to alter __getitem__

def __expand_elem__(elem, ndim):
    """parse the slice input"""

    # PARSE ELEM
    # we need to be sure of the position of all elem

    # (1) only 1 elem is given
    if not isinstance(elem, tuple):
        elem = (elem, )
    # if type(elem) is list or isinstance(elem, set):
    #     elem = (elem,)

    # if not np.iterable(elem):
    #     elem = (elem, )

    # search for named elements (dicts)
    had_dict = False

    _elem = ()
    _dict = dict()

    for i, el in enumerate(elem):
        if isinstance(el, dict):
            had_dict = True

            for key in el.keys():
                _dict[key] = el[key]

        elif had_dict:
            raise KeyError('dict argument can not be followed by positional argument')

        else:
            _elem = _elem + (el, )

    elem = _elem

    # (2) make sure there are not too many dimensions in slice.
    if len(elem) > ndim:
        raise ValueError("slicing expression exceeds the number of dimensions \
            of the variable")

    # (2) if less elements are given then there are dimensions:
    #     fill them with Ellipsis

    missing_dim = ndim - len(elem)

    if missing_dim > 0:
        # the first occurence of Ellipsis gets expanded
        if Ellipsis in elem:
            # first occurence of Ellipsis
            i = next(el[0] for el in enumerate(elem) if el[1] == Ellipsis)
            # expand
            elem = elem[:i+1] + (slice(None), ) * missing_dim + elem[i+1:]
        else:
            elem = elem + (slice(None), ) * missing_dim

    return elem, _dict

# ----------------------------------------------------------------------------

def __parse_el__(self, elem, _dict):
    """find slices that have to be selected"""

    sel_elem = list()

    for i, el in enumerate(elem):

        # coordinate subsetting happens for sets
        if isinstance(el, set):
            # convert to tuple (must be hashable)
            el = tuple(el)

            if len(el) != 2 and len(el) != 1:
                raise IndexError("Length of coordinate subsetting must be 1 or 2")

            # name of the dimension
            dim = self.dimensions[i]

            # need "Select" of the parent of the Variable
            # get the slice for this specific selection
            sel_elem.append(self._select_parent[dim][el])
        else:
            sel_elem.append(el)

        for i, dim in enumerate(self.dimensions):
            el = _dict.pop(dim, None)

            if el is not None:
                sel_elem[i] = self._select_parent[dim][el]

        for key in _dict:
            print("Select: ignored named subsetting key: '{key}' (not a dimension of the variable)".format(key=key))

    return tuple(sel_elem)

# ----------------------------------------------------------------------------

class Variable(netCDF4.Variable):
    """subclass netCDF4 to alter __getitem__"""
    def __init__(self, *arg, **kwargs):
        super(Variable, self).__init__(*arg, **kwargs)
        self.__dict__['_select_parent'] = kwargs.pop('select_parent')

    def __getitem__(self, elem):

        # add Ellipsis if elem has less members than ndim
        elem, _dict = __expand_elem__(elem, self.ndim)

        # find slices that have to be selected and select
        sel_elem = __parse_el__(self, elem, _dict)

        data = super(Variable, self).__getitem__(sel_elem)

        return data

# ============================================================================

class _Variable(netCDF4._Variable):
    """subclass netCDF4 to alter __getitem__"""

    def __init__(self, *arg, **kwargs):
        self._select_parent = kwargs.pop('select_parent')
        super(_Variable, self).__init__(*arg, **kwargs)

    def __getitem__(self, elem):

        # add Ellipsis if elem has less members than ndim
        elem, _dict = __expand_elem__(elem, self.ndim)

        # find slices that have to be selected and select
        sel_elem = __parse_el__(self, elem, _dict)

        data = super(_Variable, self).__getitem__(sel_elem)

        return data

# ============================================================================

fN = '/net/exo/landclim/mathause/cesm_data/f.e121.FC5.f19_g16.CTRL_2000-io384.001/lnd/hist/f.e121.FC5.f19_g16.CTRL_2000-io384.001.clm2.h0.0001-01.nc'

ncf = Dataset(fN)

#ncf = ncp.Dataset(fN)

#print(ncf.variables['SOILLIQ'][Ellipsis, 1:3, {'lat' : (1, 3)}].shape)

print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.select['lat'][(3, 5)])
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][0, 5, slice(None), slice(None)].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][0, 5].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][:, [0, 5]].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][:].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][:].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][0].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][Ellipsis].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][0, {0, 0.1}].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][0, {0, 0.1}, {0, 30}, ...].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILICE'][0, {0, 0.1}, {0, 30}, ...].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][{0}].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print('??????????????')
print(ncf.variables['lat'][{3}, {'lat' : (5, 10), 'lat' : (20, 30), 'aglkn' : (1, 10)}])
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
# print(ncf.variables['lat'][{3}, 7, {'lat' : (5, 10)}])
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')

fN = '/net/exo/landclim/mathause/cesm_data/f.e121.FC5.f19_g16.CTRL_2000-io384.001/lnd/hist/f.e121.FC5.f19_g16.CTRL_2000-io384.001.clm2.h0.0005-01.nc'
ncf = Dataset(fN)

print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][0, {0, 0.1}, {0, 30}, ...].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')

fN = '/net/exo/landclim/mathause/cesm_data/f.e121.FC5.f19_g16.CTRL_2000-io384.001/lnd/hist/f.e121.FC5.f19_g16.CTRL_2000-io384.001.clm2.h0.000?-01.nc'
ncf = MFDataset(fN)

print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')
print(ncf.variables['SOILLIQ'][:, {0, 0.1}, {0, 30}, ...].shape)
print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')

# print(ncf.variables['SOILLIQ'][{'lat' : (3, 15)}].shape)
# print('~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~')

ncf = MSFDataset(fN)<|endoftext|>