Fix issues with plotting IQazimuthal with wedges

            _check_parallel(intensity, error, qx, qy)

        elif len(intensity.shape) == 2:

            if len(qx.shape) == 1:

                # Qx and Qy are given in 1D array (not meshed)

                _check_parallel(intensity, error)

                if intensity.shape[0] != qx.shape[0]:

                    raise TypeError(

                        )

                    )

            elif len(qx.shape) == 2:

                # Qx and Qy are given in meshed 2D

                _check_parallel(intensity, error, qx, qy)

            else:

                raise TypeError(

        # make the qx and qy have the same shape as the data

        if len(intensity.shape) == 2 and len(qx.shape) == 1 and len(qy.shape) == 1:

            qy_length = qy.shape[0]

            qx_length = qx.shape[0]

            qx = np.tile(qx, (qy_length, 1))

            qy = np.tile(qy, (qx_length, 1)).transpose()

            # Using meshgrid to construct the Qx and Qy 2D arrays.  This is consistent with the algorithm

            # that is used in bin_iq_2d()

            qx, qy = np.meshgrid(qx, qy, indexing='ij')

            # Sanity check

            assert qx.shape == intensity.shape, f'qx and intensity must have same shapes.  ' \

                                                f'It is not now: {qx.shape} vs {intensity.shape}'

            assert qy.shape == intensity.shape, f'qy and intensity must have same shapes.  ' \

                                                f'It is not now: {qy.shape} vs {intensity.shape}'

        # pass everything to namedtuple

        return super(IQazimuthal, cls).__new__(

# https://code.ornl.gov/sns-hfir-scse/sans/sans-backend/blob/next/drtsans/dataobjects.py

# https://code.ornl.gov/sns-hfir-scse/sans/sans-backend/blob/next/docs/drtsans/dataobjects.rst

import numpy

from drtsans.dataobjects import (

    DataType,

    getDataType,

    return wedge_i_of_q

def _toQmodAndAzimuthal(data):

def _toQmodAndAzimuthal(data: IQazimuthal) -> Tuple[numpy.ndarray, numpy.ndarray,

                                                    numpy.ndarray, numpy.ndarray, numpy.ndarray]:

    """This function returns the values of qmod and azimuthal that are parallel

    to the original data array. It requires that the data is IQazimuthal

            return Backend[mode.upper()]

def _saveFile(figure, filename, backend, show=False):

def _save_file(figure, filename, backend, show=False):

    """Convenience method for the common bits of saving the file based on

    the selected backend.

            raise RuntimeError('Do not know how to plot type="{}"'.format(datatype))

    fig, ax = plt.subplots()

    handles = []

    for n, workspace in enumerate(workspaces):

        eb, _, _ = ax.errorbar(

            workspace.mod_q, workspace.intensity, yerr=workspace.error

    if kwargs:

        plt.setp(ax, **kwargs)

    _saveFile(fig, filename, backend)

    _save_file(fig, filename, backend)

def _create_ring_roi(iq2d, q_min, q_max, input_roi) -> np.ndarray:

    """Create a mask or ROI by q range and result in a ring (or circle)

    Returns

    -------

    np.ndarray

        Boolean numpy array with same shape as input iq2d.  False for masking, True for ROI

    """

    # create region of interest overlay

    output_roi = input_roi

    # Larger than min Q

    if q_min is not None:

        # roi = np.logical_and(roi, np.square(iq2d.qx) + np.square(iq2d.qy) > np.square(q_min))

        q_min_roi = np.square(iq2d.qx) + np.square(iq2d.qy) > np.square(q_min)

        # act on output

        output_roi = np.logical_and(output_roi, q_min_roi)

    # Less than max Q

    if q_max is not None:

        q_max_roi = np.square(iq2d.qx) + np.square(iq2d.qy) < np.square(q_max)

        # act on output

        output_roi = np.logical_and(output_roi, q_max_roi)

    return output_roi

def _create_wedge_roi(iq2d, wedges, symmetric_wedges: bool,

                      input_roi: np.ndarray, qx_limit: float = 1e10) -> np.ndarray:

    """Create ROI matrix with

    Parameters

    ----------

    iq2d

    wedges

    symmetric_wedges

    input_roi

    qx_limit: float

        Upper limit of possible Qx

    Returns

    -------

    np.ndarray

        boolean array with same shape as intensity.  True for ROI.

    """

    # set output

    output_roi = input_roi

    # Check: wedge must be specified

    if wedges is None:

        return output_roi

    # create bool array selecting nothing

    roi_wedges = np.zeros(iq2d.intensity.shape).astype(bool)

    # expand the supplied variables into an easier form

    # get validated wedge in groups and flatten it to list of wedge angles

    wedge_angles = validate_wedges_groups(wedges, symmetric_wedges)

    wedge_angles = [

        wedge_angle for wedges_group in wedge_angles for wedge_angle in wedges_group

    ]

    # create the individual selections and combine with 'or'

    # Note: qx is in [[qx0, qx0, qx0, ...], [qx1, qx1, qx1, ...], ...]

    #       qy is in [[qy0, qy1, qy2, ...], [qy0, qy1, qy2, ...], ...]

    # this is transposed comparing to how Qx and Qy is plotted for the output

    azimuthal = np.rad2deg(

        np.arctan2(iq2d.qy, iq2d.qx)

    )

    # Try 1azimuthal = np.rad2deg(np.arctan2(workspace.qx, workspace.qy))

    azimuthal[azimuthal <= -90.0] += 360.0

    for lower_boundary_angle, upper_boundary_angle in wedge_angles:

        wedge = np.logical_and((azimuthal > lower_boundary_angle), (azimuthal < upper_boundary_angle))

        roi_wedges = np.logical_or(roi_wedges, wedge)

    # combine with existing roi

    output_roi = np.logical_and(output_roi, roi_wedges)

    return output_roi

def _require_transpose_intensity(iq2d) -> bool:

    """Check whether the intensity/ROI in IQazimuthal shall be transposed to plot

    as Qx in horizontal and Qy in vertical

    """

    # Determine whether intensity matrix shall be inverted or not

    qx2d = iq2d.qx

    qy2d = iq2d.qy

    # set up the flag to transpose ROI if I(Qx, Qy) is to be tranposed

    transpose_flag = False

    if len(qx2d.shape) == 1:

        # No need to transpose if Qx and Qy are given in 1-dim

        transpose_flag = True

    if len(qx2d.shape) == 2 and qx2d.shape[0] > 1 and np.sum(qx2d[0] == qx2d[1]) == qx2d.shape[1]:

        # Input Qx and Qy are 2-dim and

        # I(Qx, Qy) is of same order as meshgrid(Qx, Qy)

        # Qx have identical among rows:

        if qy2d.shape[1] > 1:

            # sanity check

            assert (np.sum(qy2d[:, 0] == qy2d[:, 1]) == qy2d.shape[0]), "Qy shall have identical columns"

    else:

        # I(Qx, Qy) is transposed to meshgrid(Qx, Qy)

        transpose_flag = True

    return transpose_flag

def plot_IQazimuthal(

    kwargs: ~dict

        Additional key word arguments for :py:obj:`matplotlib.axes.Axes`

    """

    # Set up backend and verify data type

    backend = Backend.getMode(backend)

    datatype = getDataType(workspace)

    if datatype != DataType.IQ_AZIMUTHAL:

        raise RuntimeError('Do not know how to plot type="{}"'.format(datatype))

    qxmin = workspace.qx.min()

    qxmax = workspace.qx.max()

    qymin = workspace.qy.min()

    qymax = workspace.qy.max()

    # create region of interest overlay

    roi = None

    if qmin is not None:

        # no need to check for exsting roi

        roi = np.square(workspace.qx) + np.square(workspace.qy) > np.square(qmin)

    if qmax is not None:

        roi_qmax = np.square(workspace.qx) + np.square(workspace.qy) < np.square(qmax)

        if roi is None:

            roi = roi_qmax

        else:

            roi = np.logical_and(roi, roi_qmax)

    if wedges is not None:

        # create bool array selecting nothing

        roi_wedges = np.logical_not(workspace.qx < 1000.0)

        # expand the supplied variables into an easier form

        # get validated wedge in groups and flatten it to list of wedge angles

        wedge_angles = validate_wedges_groups(wedges, symmetric_wedges)

        wedge_angles = [

            wedge_angle for wedges_group in wedge_angles for wedge_angle in wedges_group

        ]

        # create the individual selections and combine with 'or'

        # Note: qx is in [[qx0, qx0, qx0, ...], [qx1, qx1, qx1, ...], ...]

        #       qy is in [[qy0, qy1, qy2, ...], [qy0, qy1, qy2, ...], ...]

        # this is transposed comparing to how Qx and Qy is plotted for the output

        azimuthal = np.rad2deg(

            np.arctan2(workspace.qy.transpose(), workspace.qx.transpose())

        )

        # Try 1azimuthal = np.rad2deg(np.arctan2(workspace.qx, workspace.qy))

        azimuthal[azimuthal <= -90.0] += 360.0

        for left, right in wedge_angles:

            wedge = np.logical_and((azimuthal > left), (azimuthal < right))

            # Try 0 wedge = np.logical_and((azimuthal < right), (azimuthal > left))

            roi_wedges = np.logical_or(roi_wedges, wedge)

    # Set up ROI/mask

    roi = (np.zeros(workspace.intensity.shape) + 1).astype(bool)

    # ROI as a ring (qmin, qmax)

    roi = _create_ring_roi(workspace, qmin, qmax, roi)

    # wedge

    roi = _create_wedge_roi(workspace, wedges, symmetric_wedges, roi)

        # combine with existing roi

        if roi is None:

            roi = roi_wedges

    # Make sure the orientation of intensity array can be shown correctly with imshow

    # imshow thinks the bottom right corner is (0, n-1) while it is intensity(qx_max, qy_min)

    transpose_flag = _require_transpose_intensity(workspace)

    if transpose_flag:

        # transpose both intensity and ROI

        intensity = workspace.intensity.T

        roi = roi.T

    else:

            roi = np.logical_and(roi, roi_wedges)

        intensity = workspace.intensity

    # convert ROI to masks

    roi = np.ma.masked_where(roi, roi.astype(int))

    # put together the plot

    fig, ax = plt.subplots()

    current_cmap = matplotlib.cm.get_cmap()

    current_cmap.set_bad(color="grey")

    # Determine whether intensity matrix shall be inverted or not

    qx2d = workspace.qx

    qy2d = workspace.qy

    # set up the flag to transpose ROI if I(Qx, Qy) is to be tranposed

    transpose_flag = False

    if qx2d.shape[0] > 1 and np.sum(qx2d[0] == qx2d[1]) == qx2d.shape[1]:

        # I(Qx, Qy) is of same order as meshgrid(Qx, Qy)

        # Qx have identical among rows:

        if qy2d.shape[1] > 1:

            # sanity check

            assert (

                np.sum(qy2d[:, 0] == qy2d[:, 1]) == qy2d.shape[0]

            ), "Qy shall have identical columns"

        intensity = workspace.intensity

    else:

        # I(Qx, Qy) is tranposed to meshgrid(Qx, Qy)

        intensity = workspace.intensity.T

        transpose_flag = True

    qxmin = workspace.qx.min()

    qxmax = workspace.qx.max()

    qymin = workspace.qy.min()

    qymax = workspace.qy.max()

    pcm = ax.imshow(

        intensity,

        extent=(qxmin, qxmax, qymin, qymax),

    )

    # add calculated region of interest

    if roi is not None:

        roi = np.ma.masked_where(roi, roi.astype(int))

        # transpose ROI

        if transpose_flag:

            roi = roi.T

    ax.imshow(

        roi,

        alpha=mask_alpha,

    if kwargs:

        plt.setp(ax, **kwargs)

    _saveFile(fig, filename, backend)

    _save_file(fig, filename, backend)

def plot_detector(

    fig.tight_layout()

    if filename is not None:

        _saveFile(fig, filename, backend)

        _save_file(fig, filename, backend)

            [12.2, 13.8, 14.8, 13.4, 13.8, 16.7, 14.8, 17.0, 14.8, 14.8, 12.2],

        ]

    )

    # IQazimuthal's constructor is corrected in

    # https://code.ornl.gov/sns-hfir-scse/sans/sans-backend/-/merge_requests/951

    # Thus the test input data shall be tranposed accordingly

    intensities = intensities.T

    errors = errors.T

    data2d = IQazimuthal(

        intensity=intensities,

        error=errors,

    assert delta_qmod is None

    assert azimuthal.shape == intensity.shape

    from drtsans.plots.api import plot_IQazimuthal

    plot_IQazimuthal(data2d, 'input_wedge.png', backend='mpl')

    # numbers taken from the spreadsheet

    q_exp = np.array(

        [

    )

    np.testing.assert_allclose(qmod, q_exp.ravel(), atol=0.005)

    np.testing.assert_allclose(azimuthal, azimuthal_exp.ravel(), atol=0.5)

    # IQazimuthal's constructor is corrected in

    # https://code.ornl.gov/sns-hfir-scse/sans/sans-backend/-/merge_requests/951

    # Thus the expected test result shall be tranposed accordingly

    np.testing.assert_allclose(azimuthal, azimuthal_exp.T.ravel(), atol=0.5)

def test_bin_into_q_and_azimuthal():

        azimuthal_delta=azimuthal_delta,

    )

    # from drtsans.plots.api import plot_IQazimuthal

    # plot_IQazimuthal(data2d, 'wedge_int_test_debug.png', backend='mpl')

    # assert False, f'Wedges: {wedges}'

    # tests

    assert len(wedges) == 2

    for wedge in wedges:

import os

import pytest

import matplotlib.pyplot as plt

from matplotlib.pyplot import imread

from typing import Tuple, Any

def verify_images(test_png: str, gold_png):

    """Verify the image output from test is same as what is expected (gold)

    AssertionError will be raised if they do not match

    """

    # check

    assert os.path.exists(gold_png), f'Gold/reference data file {gold_png} cannot be located'

    assert os.path.exists(test_png), f'Test result data file {test_png} cannot be located'

    # Import PNG to 2d array

    tested_image = imread(test_png)

    gold_image = imread(gold_png)

    # Verify

    np.testing.assert_allclose(tested_image, gold_image, err_msg=f'Testing result {tested_image} does not match '

                                                                 f'the expected result {gold_image}')

def fileCheckAndRemove(filename, remove=True):

    """Convienience function for doing simple checs that the file was created.

    """Convenience function for doing simple checks that the file was created.

    The ``remove`` option is available to make debugging new tests easier."""

    assert os.path.exists(filename), '"{}" does not exist'.format(filename)

    assert os.path.getsize(filename) > 100, '"{}" is too small'.format(filename)

    fileCheckAndRemove(filename)

@pytest.fixture

def test_iq2d_data() -> Tuple[Any, Any, Any, Any]:

    """Generate Qx, Qy, I(qx, qy), Error(qx, qy) data

    """

    # Qx: 60 values, Qy: 40 values

    x = np.linspace(0.0, 4 * np.pi, 60) - 3.

    y = np.linspace(0.5 * np.pi, 4.5 * np.pi, 40) - 3.

    # Calculate intensity and error

    mesh_x, mesh_y = np.meshgrid(x, y, sparse=False, indexing="ij")

    intensity = np.abs(np.sin(mesh_x) + np.cos(mesh_y))

    error = np.sqrt(intensity)

    # Transfer to correct orientation

    mesh_x = mesh_x.T

    mesh_y = mesh_y.T

    intensity = intensity.T

    error = error.T

    # now try to find zero...

    x_zero_index = np.argmin(np.abs(x))

    y_zero_index = np.argmin(np.abs(y))

    # mask center

    for x_index in [-1, 0, 1]:

        for y_index in [-1, 0, 1]:

            center_x_index = x_zero_index + x_index

            center_y_index = y_zero_index + y_index

            intensity[center_y_index][center_x_index] = np.nan

            error[center_y_index][center_x_index] = np.nan

    # mark lower right corner

    intensity[1, 59] = 8.

    assert intensity.shape == (40, 60), f'Expected intensity is 40 row (Qy) and 60 column (Qx) ' \

                                        f'but not {intensity.shape}'

    return mesh_x, mesh_y, intensity, error

@pytest.mark.parametrize(

    "backend, filename",

    [("mpl", "test_IQazimuthal_1d.png"), ("d3", "test_IQazimuthal_1d.json")],

    ids=["mpl", "d3"],

    "backend, filename, reference_name",

    [

        ("mpl", "test_IQazimuthal_1d.png", 'tests/unit/references/gold_IQazimuthal_2d_T.png'),

        ("d3", "test_IQazimuthal_1d.json", None)

    ],

    ids=["mpl_test", "d3"],

)

def test_IQazimuthal_1d(backend, filename):

def test_IQazimuthal_1d(backend, filename, reference_name, test_iq2d_data):

    """Test plotting IQazimuthal with 1d Qx and Qy arrays"""

    x = np.linspace(0.0, 4 * np.pi, 50)

    y = np.linspace(0.5 * np.pi, 4.5 * np.pi, 50)

    error = np.zeros((50, 50))

    data = error[:]

    for i in range(x.size):

        for j in range(y.size):

            data[i, j] = np.sin(x[i]) + np.cos(y[j])

    data = IQazimuthal(intensity=data, error=error, qx=x, qy=y)

    # Generate input arrays

    mesh_x, mesh_y, intensity, error = test_iq2d_data

    x = mesh_x[0]

    assert x.min() < x.max()

    y = mesh_y[:, 0]

    assert y.min() < y.max()

    # construct IQazimuthal: following bin_iq_2d() routine, i.e., intensity is (num_qx, num_qy)

    assert intensity.T.shape[0] == len(x)

    assert intensity.T.shape[1] == len(y)

    data = IQazimuthal(intensity=intensity.T, error=error.T, qx=x, qy=y)

    # plot

    plot_IQazimuthal(data, filename=filename, backend=backend)

    plt.close()

    fileCheckAndRemove(filename)

    # verify

    if reference_name:

        fileCheckAndRemove(filename, remove=False)

        verify_images(filename, reference_name)

    fileCheckAndRemove(filename, remove=False)

@pytest.mark.parametrize(

    "backend, filename",

    [("mpl", "test_IQazimuthal_2d.png"), ("d3", "test_IQazimuthal_2d.json")],

    "backend, filename, reference_name",

    [("mpl", "test_IQazimuthal_2d.png", 'tests/unit/references/gold_IQazimuthal_2d_T.png'),

     ("d3", "test_IQazimuthal_2d.json", None)],

    ids=["mpl", "d3"],

)

def test_IQazimuthal_2d(backend, filename):

def test_IQazimuthal_2d(backend, filename, reference_name, test_iq2d_data):

    """Test plotting IQazimuthal with 2d Qx and Qy arrays"""

    x, y = np.meshgrid(

        np.linspace(0.0, 4 * np.pi, 50),

        np.linspace(0.5 * np.pi, 4.5 * np.pi, 50),

        sparse=False,

        indexing="ij",

    # construct IQazimuthal

    x, y, intensity, error = test_iq2d_data

    data = IQazimuthal(intensity=intensity, error=error, qx=x, qy=y)

    # plot

    plot_IQazimuthal(data, filename=filename, backend=backend)

    plt.close()

    # verify

    if reference_name:

        fileCheckAndRemove(filename, remove=False)

        verify_images(filename, reference_name)

    fileCheckAndRemove(filename, remove=True)

@pytest.mark.parametrize(

    "backend, filename, reference_name",

    [("mpl", "test_IQazimuthal_2d_T.png", 'tests/unit/references/gold_IQazimuthal_2d_T.png'),

     ("d3", "test_IQazimuthal_2d_T.json", None)],

    ids=["mpl", "d3"],

)

    error = np.zeros(x.shape)

    data = np.sin(x) + np.cos(y)

    data = IQazimuthal(intensity=data, error=error, qx=x, qy=y)

def test_IQazimuthal_2d_transposed(backend, filename, reference_name, test_iq2d_data):

    """Test plotting IQazimuthal with 2d Qx and Qy arrays"""

    # construct IQazimuthal

    x, y, intensity, error = test_iq2d_data

    data = IQazimuthal(intensity=intensity.T, error=error.T, qx=x.T, qy=y.T)

    # plot

    plot_IQazimuthal(data, filename=filename, backend=backend)

    plt.close()

    fileCheckAndRemove(filename, False)

    # verify

    if reference_name:

        fileCheckAndRemove(filename, remove=False)

        verify_images(filename, reference_name)

    fileCheckAndRemove(filename, remove=True)

# Broken!

@pytest.mark.parametrize(

    "backend, filename",

    "backend, filename, reference_name",

    [

        ("mpl", "test_IQazimuthal_2d_selections.png"),

        ("d3", "test_IQazimuthal_2d_selections.json"),

        ("mpl", "test_IQazimuthal_2d_selections.png", 'tests/unit/references/new_IQazimuthal_2d_selections.png'),

        ("d3", "test_IQazimuthal_2d_selections.json", None),

    ],

    ids=["mpl", "d3"],

)

def test_IQazimuthal_2d_selections(backend, filename):

def test_IQazimuthal_2d_selections(backend, filename, reference_name, test_iq2d_data):

    """Test plotting IQazimuthal with 2d Qx and Qy arrays"""

    x, y = np.meshgrid(

        np.linspace(0.0, 4 * np.pi, 50),

        np.linspace(0.5 * np.pi, 4.5 * np.pi, 50),

        sparse=False,

        indexing="ij",

    )

    error = np.zeros(x.shape)

    data = np.sin(x) + np.cos(y)

    data = IQazimuthal(intensity=data, error=error, qx=x, qy=y)

    # construct IQazimuthal

    x, y, intensity, error = test_iq2d_data

    data = IQazimuthal(intensity=intensity, error=error, qx=x, qy=y)

    # plot

    plot_IQazimuthal(

        data,

        filename=filename,

        backend=backend,

        qmin=0.0,

        qmax=9.0,

        wedges=((-30.0, 30.0),),

    )

    plt.close()

    # verify

    if reference_name:

        fileCheckAndRemove(filename, remove=False)

        verify_images(filename, reference_name)

    fileCheckAndRemove(filename, remove=True)

# Broken

@pytest.mark.parametrize(

    "backend, filename, reference_name",

    [

        ("mpl", "test_IQazimuthal_2d_asymmetric_wedge.png",

         'tests/unit/references/new_IQazimuthal_2d_asymmetric_wedge.png'),

        ("d3", "test_IQazimuthal_2d_asymmetric_wedge.json", None),

    ],

    ids=["mpl", "d3"],

)

def test_IQazimuthal_2d_asymmetric_wedge(backend, filename, reference_name, test_iq2d_data):

    """Test plotting IQazimuthal with 2d Qx and Qy arrays"""