예제 #1
0
    def test_integer_levels(self):
        data = self.data + 1
        cmap_params = _determine_cmap_params(data, levels=5, vmin=0, vmax=5, cmap="Blues")
        self.assertEqual(cmap_params["vmin"], cmap_params["levels"][0])
        self.assertEqual(cmap_params["vmax"], cmap_params["levels"][-1])
        self.assertEqual(cmap_params["cmap"].name, "Blues")
        self.assertEqual(cmap_params["extend"], "neither")
        self.assertEqual(cmap_params["cmap"].N, 5)
        self.assertEqual(cmap_params["norm"].N, 6)

        cmap_params = _determine_cmap_params(data, levels=5, vmin=0.5, vmax=1.5)
        self.assertEqual(cmap_params["cmap"].name, "viridis")
        self.assertEqual(cmap_params["extend"], "max")
예제 #2
0
    def test_list_levels(self):
        data = self.data + 1

        orig_levels = [0, 1, 2, 3, 4, 5]
        # vmin and vmax should be ignored if levels are explicitly provided
        cmap_params = _determine_cmap_params(data, levels=orig_levels, vmin=0, vmax=3)
        self.assertEqual(cmap_params["vmin"], 0)
        self.assertEqual(cmap_params["vmax"], 5)
        self.assertEqual(cmap_params["cmap"].N, 5)
        self.assertEqual(cmap_params["norm"].N, 6)

        for wrap_levels in [list, np.array, pd.Index, DataArray]:
            cmap_params = _determine_cmap_params(data, levels=wrap_levels(orig_levels))
            self.assertArrayEqual(cmap_params["levels"], orig_levels)
예제 #3
0
파일: test_plot.py 프로젝트: OXPHOS/xarray
    def test_integer_levels(self):
        data = self.data + 1
        cmap_params = _determine_cmap_params(data, levels=5, vmin=0, vmax=5,
                                             cmap='Blues')
        self.assertEqual(cmap_params['vmin'], cmap_params['levels'][0])
        self.assertEqual(cmap_params['vmax'], cmap_params['levels'][-1])
        self.assertEqual(cmap_params['cmap'].name, 'Blues')
        self.assertEqual(cmap_params['extend'], 'neither')
        self.assertEqual(cmap_params['cmap'].N, 5)
        self.assertEqual(cmap_params['norm'].N, 6)

        cmap_params = _determine_cmap_params(data, levels=5,
                                             vmin=0.5, vmax=1.5)
        self.assertEqual(cmap_params['cmap'].name, 'viridis')
        self.assertEqual(cmap_params['extend'], 'max')
예제 #4
0
    def test_integer_levels(self):
        data = self.data + 1
        cmap_params = _determine_cmap_params(data, levels=5, vmin=0, vmax=5,
                                             cmap='Blues')
        self.assertEqual(cmap_params['vmin'], cmap_params['levels'][0])
        self.assertEqual(cmap_params['vmax'], cmap_params['levels'][-1])
        self.assertEqual(cmap_params['cmap'].name, 'Blues')
        self.assertEqual(cmap_params['extend'], 'neither')
        self.assertEqual(cmap_params['cmap'].N, 5)
        self.assertEqual(cmap_params['norm'].N, 6)

        cmap_params = _determine_cmap_params(data, levels=5,
                                             vmin=0.5, vmax=1.5)
        self.assertEqual(cmap_params['cmap'].name, 'viridis')
        self.assertEqual(cmap_params['extend'], 'max')
예제 #5
0
 def test_center(self):
     cmap_params = _determine_cmap_params(self.data, center=0.5)
     self.assertEqual(cmap_params['vmax'] - 0.5, 0.5 - cmap_params['vmin'])
     self.assertEqual(cmap_params['cmap'], 'RdBu_r')
     self.assertEqual(cmap_params['extend'], 'neither')
     self.assertIsNone(cmap_params['levels'])
     self.assertIsNone(cmap_params['norm'])
예제 #6
0
 def test_center(self):
     cmap_params = _determine_cmap_params(self.data, center=0.5)
     self.assertEqual(cmap_params["vmax"] - 0.5, 0.5 - cmap_params["vmin"])
     self.assertEqual(cmap_params["cmap"], "RdBu_r")
     self.assertEqual(cmap_params["extend"], "neither")
     self.assertIsNone(cmap_params["levels"])
     self.assertIsNone(cmap_params["norm"])
예제 #7
0
 def test_center(self):
     cmap_params = _determine_cmap_params(self.data, center=0.5)
     self.assertEqual(cmap_params['vmax'] - 0.5, 0.5 - cmap_params['vmin'])
     self.assertEqual(cmap_params['cmap'], 'RdBu_r')
     self.assertEqual(cmap_params['extend'], 'neither')
     self.assertIsNone(cmap_params['levels'])
     self.assertIsNone(cmap_params['norm'])
예제 #8
0
    def test_list_levels(self):
        data = self.data + 1

        orig_levels = [0, 1, 2, 3, 4, 5]
        # vmin and vmax should be ignored if levels are explicitly provided
        cmap_params = _determine_cmap_params(data, levels=orig_levels,
                                             vmin=0, vmax=3)
        self.assertEqual(cmap_params['vmin'], 0)
        self.assertEqual(cmap_params['vmax'], 5)
        self.assertEqual(cmap_params['cmap'].N, 5)
        self.assertEqual(cmap_params['norm'].N, 6)

        for wrap_levels in [list, np.array, pd.Index, DataArray]:
            cmap_params = _determine_cmap_params(
                data, levels=wrap_levels(orig_levels))
            self.assertArrayEqual(cmap_params['levels'], orig_levels)
 def test_center(self):
     cmap_params = _determine_cmap_params(self.data, center=0.5)
     assert cmap_params['vmax'] - 0.5 == 0.5 - cmap_params['vmin']
     assert cmap_params['cmap'] == 'RdBu_r'
     assert cmap_params['extend'] == 'neither'
     assert cmap_params['levels'] is None
     assert cmap_params['norm'] is None
예제 #10
0
 def test_robust(self):
     cmap_params = _determine_cmap_params(self.data, robust=True)
     self.assertEqual(cmap_params["vmin"], np.percentile(self.data, 2))
     self.assertEqual(cmap_params["vmax"], np.percentile(self.data, 98))
     self.assertEqual(cmap_params["cmap"].name, "viridis")
     self.assertEqual(cmap_params["extend"], "both")
     self.assertIsNone(cmap_params["levels"])
     self.assertIsNone(cmap_params["norm"])
예제 #11
0
 def test_robust(self):
     cmap_params = _determine_cmap_params(self.data, robust=True)
     self.assertEqual(cmap_params['vmin'], np.percentile(self.data, 2))
     self.assertEqual(cmap_params['vmax'], np.percentile(self.data, 98))
     self.assertEqual(cmap_params['cmap'].name, 'viridis')
     self.assertEqual(cmap_params['extend'], 'both')
     self.assertIsNone(cmap_params['levels'])
     self.assertIsNone(cmap_params['norm'])
예제 #12
0
 def test_robust(self):
     cmap_params = _determine_cmap_params(self.data, robust=True)
     self.assertEqual(cmap_params['vmin'], np.percentile(self.data, 2))
     self.assertEqual(cmap_params['vmax'], np.percentile(self.data, 98))
     self.assertEqual(cmap_params['cmap'].name, 'viridis')
     self.assertEqual(cmap_params['extend'], 'both')
     self.assertIsNone(cmap_params['levels'])
     self.assertIsNone(cmap_params['norm'])
 def test_robust(self):
     cmap_params = _determine_cmap_params(self.data, robust=True)
     assert cmap_params['vmin'] == np.percentile(self.data, 2)
     assert cmap_params['vmax'] == np.percentile(self.data, 98)
     assert cmap_params['cmap'].name == 'viridis'
     assert cmap_params['extend'] == 'both'
     assert cmap_params['levels'] is None
     assert cmap_params['norm'] is None
예제 #14
0
    def test_integer_levels(self):
        data = self.data + 1

        # default is to cover full data range but with no guarantee on Nlevels
        for level in np.arange(2, 10, dtype=int):
            cmap_params = _determine_cmap_params(data, levels=level)
            self.assertEqual(cmap_params['vmin'], cmap_params['levels'][0])
            self.assertEqual(cmap_params['vmax'], cmap_params['levels'][-1])
            self.assertEqual(cmap_params['extend'], 'neither')

        # with min max we are more strict
        cmap_params = _determine_cmap_params(data,
                                             levels=5,
                                             vmin=0,
                                             vmax=5,
                                             cmap='Blues')
        self.assertEqual(cmap_params['vmin'], 0)
        self.assertEqual(cmap_params['vmax'], 5)
        self.assertEqual(cmap_params['vmin'], cmap_params['levels'][0])
        self.assertEqual(cmap_params['vmax'], cmap_params['levels'][-1])
        self.assertEqual(cmap_params['cmap'].name, 'Blues')
        self.assertEqual(cmap_params['extend'], 'neither')
        self.assertEqual(cmap_params['cmap'].N, 4)
        self.assertEqual(cmap_params['norm'].N, 5)

        cmap_params = _determine_cmap_params(data,
                                             levels=5,
                                             vmin=0.5,
                                             vmax=1.5)
        self.assertEqual(cmap_params['cmap'].name, 'viridis')
        self.assertEqual(cmap_params['extend'], 'max')

        cmap_params = _determine_cmap_params(data, levels=5, vmin=1.5)
        self.assertEqual(cmap_params['cmap'].name, 'viridis')
        self.assertEqual(cmap_params['extend'], 'min')

        cmap_params = _determine_cmap_params(data,
                                             levels=5,
                                             vmin=1.3,
                                             vmax=1.5)
        self.assertEqual(cmap_params['cmap'].name, 'viridis')
        self.assertEqual(cmap_params['extend'], 'both')
    def test_integer_levels(self):
        data = self.data + 1

        # default is to cover full data range but with no guarantee on Nlevels
        for level in np.arange(2, 10, dtype=int):
            cmap_params = _determine_cmap_params(data, levels=level)
            assert cmap_params['vmin'] == cmap_params['levels'][0]
            assert cmap_params['vmax'] == cmap_params['levels'][-1]
            assert cmap_params['extend'] == 'neither'

        # with min max we are more strict
        cmap_params = _determine_cmap_params(data,
                                             levels=5,
                                             vmin=0,
                                             vmax=5,
                                             cmap='Blues')
        assert cmap_params['vmin'] == 0
        assert cmap_params['vmax'] == 5
        assert cmap_params['vmin'] == cmap_params['levels'][0]
        assert cmap_params['vmax'] == cmap_params['levels'][-1]
        assert cmap_params['cmap'].name == 'Blues'
        assert cmap_params['extend'] == 'neither'
        assert cmap_params['cmap'].N == 4
        assert cmap_params['norm'].N == 5

        cmap_params = _determine_cmap_params(data,
                                             levels=5,
                                             vmin=0.5,
                                             vmax=1.5)
        assert cmap_params['cmap'].name == 'viridis'
        assert cmap_params['extend'] == 'max'

        cmap_params = _determine_cmap_params(data, levels=5, vmin=1.5)
        assert cmap_params['cmap'].name == 'viridis'
        assert cmap_params['extend'] == 'min'

        cmap_params = _determine_cmap_params(data,
                                             levels=5,
                                             vmin=1.3,
                                             vmax=1.5)
        assert cmap_params['cmap'].name == 'viridis'
        assert cmap_params['extend'] == 'both'
예제 #16
0
    def test_integer_levels(self):
        data = self.data + 1

        # default is to cover full data range but with no guarantee on Nlevels
        for level in np.arange(2, 10, dtype=int):
            cmap_params = _determine_cmap_params(data, levels=level)
            self.assertEqual(cmap_params['vmin'], cmap_params['levels'][0])
            self.assertEqual(cmap_params['vmax'], cmap_params['levels'][-1])
            self.assertEqual(cmap_params['extend'], 'neither')

        # with min max we are more strict
        cmap_params = _determine_cmap_params(data, levels=5, vmin=0, vmax=5,
                                             cmap='Blues')
        self.assertEqual(cmap_params['vmin'], 0)
        self.assertEqual(cmap_params['vmax'], 5)
        self.assertEqual(cmap_params['vmin'], cmap_params['levels'][0])
        self.assertEqual(cmap_params['vmax'], cmap_params['levels'][-1])
        self.assertEqual(cmap_params['cmap'].name, 'Blues')
        self.assertEqual(cmap_params['extend'], 'neither')
        self.assertEqual(cmap_params['cmap'].N, 4)
        self.assertEqual(cmap_params['norm'].N, 5)

        cmap_params = _determine_cmap_params(data, levels=5,
                                             vmin=0.5, vmax=1.5)
        self.assertEqual(cmap_params['cmap'].name, 'viridis')
        self.assertEqual(cmap_params['extend'], 'max')

        cmap_params = _determine_cmap_params(data, levels=5,
                                             vmin=1.5)
        self.assertEqual(cmap_params['cmap'].name, 'viridis')
        self.assertEqual(cmap_params['extend'], 'min')

        cmap_params = _determine_cmap_params(data, levels=5,
                                             vmin=1.3, vmax=1.5)
        self.assertEqual(cmap_params['cmap'].name, 'viridis')
        self.assertEqual(cmap_params['extend'], 'both')
예제 #17
0
def geo_plot(darray, ax=None, method='contourf',
             projection='PlateCarree', grid=False, **kwargs):
    """ Create a global plot of a given variable.

    Parameters:
    -----------
    darray : xray.DataArray
        The darray to be plotted.
    ax : axis
        An existing axis instance, else one will be created.
    method : str
        String to use for looking up name of plotting function via iris
    projection : str or tuple
        Name of the cartopy projection to use and any args
        necessary for initializing it passed as a dictionary;
        see func:`make_geoaxes` for more information
    grid : bool
        Include lat-lon grid overlay
    **kwargs : dict
        Any additional keyword arguments to pass to the plotter,
        including colormap params. If 'vmin' is not in this
        set of optional keyword arguments, the plot colormap will be
        automatically inferred.

    """

    # Set up plotting function
    if method in _PLOTTYPE_ARGS:
        extra_args = _PLOTTYPE_ARGS[method].copy()
    else:
        raise ValueError("Don't know how to deal with '%s' method" % method)
    extra_args.update(**kwargs)

    # Alias a plot function based on the requested method and the
    # datatype being plotted
    plot_func = plt.__dict__[method]

    # `transform` should be the ORIGINAL coordinate system -
    # which is always a simple lat-lon coordinate system in CESM
    # output
    extra_args['transform'] = ccrs.PlateCarree()

    # Was an axis passed to plot on?
    new_axis = ax is None

    if new_axis: # Create a new cartopy axis object for plotting
        if isinstance(projection, (list, tuple)):
            if len(projection) != 2:
                raise ValueError("Expected 'projection' to only have 2 values")
            projection, proj_kwargs = projection[0], projection[1]
        else:
            proj_kwargs = {}

        # hack to look up the name of the projection in the cartopy
        # reference system namespace; makes life a bit easier, so you
        # can just pass a string with the name of the projection wanted.
        proj = ccrs.__dict__[projection](**proj_kwargs)
        ax = plt.axes(projection=proj)
    else: # Set current axis to one passed as argument
        if not hasattr(ax, 'projection'):
            raise ValueError("Expected `ax` to be a GeoAxes instance")
        plt.sca(ax)

    # Setup map
    ax.set_global()
    ax.coastlines()

    try:
        gl = ax.gridlines(crs=extra_args['transform'], draw_labels=True,
                          linewidth=0.5, color='grey', alpha=0.8)
        LON_TICKS = [ -180, -90, 0, 90, 180 ]
        LAT_TICKS = [ -90, -60, -30, 0, 30, 60, 90 ]
        gl.xlabels_top   = False
        gl.ylabels_right = False
        gl.xlines = grid
        gl.ylines = grid
        gl.xlocator = mticker.FixedLocator(LON_TICKS)
        gl.ylocator = mticker.FixedLocator(LAT_TICKS)
        gl.xformatter = LONGITUDE_FORMATTER
        gl.yformatter = LATITUDE_FORMATTER
    except TypeError:
        warnings.warn("Could not label the given map projection.")

    # Infer colormap settings if not provided
    if not ('vmin' in kwargs):
        warnings.warn("Re-inferring color parameters...")
        cmap_kws = _determine_cmap_params(darray.data)
        extra_args.update(cmap_kws)

    gp = plot_func(darray.lon.values, darray.lat.values, darray.data,
                   **extra_args)

    return ax, gp
예제 #18
0
def TS_diagram(od, display=True, FuncAnimation_kwargs=None, **kwargs):
    """
    Animate TS diagrams.

    Parameters
    ----------
    od: OceanDataset
        oceandataset used to plot.
    display: bool
        If True, display the animation.
    FuncAnimation_kwargs: dict
        Keyword arguments from :py:func:`matplotlib.animation.FuncAnimation`
    **kwargs:
        Keyword arguments from :py:func:`oceanspy.plot.TS_diagram`

    Returns
    -------
    anim: matplotlib.animation.FuncAnimation
        Animation object

    See also
    --------
    oceanspy.plot.TS_diagram
    """

    # Check parameters
    _check_instance(
        {
            "od": od,
            "display": display,
            "FuncAnimation_kwargs": FuncAnimation_kwargs
        },
        {
            "od": "oceanspy.OceanDataset",
            "display": "bool",
            "FuncAnimation_kwargs": ["type(None)", "dict"],
        },
    )

    # Handle kwargs
    if FuncAnimation_kwargs is None:
        FuncAnimation_kwargs = {}
    FuncAnimation_kwargs = dict(FuncAnimation_kwargs)

    # Name of the plot_functions
    plot_func = eval("_plot.TS_diagram")

    # First cutout and get time
    cutout_kwargs = kwargs.pop("cutout_kwargs", None)
    if cutout_kwargs is not None:
        od = od.subsample.cutout(**cutout_kwargs)
    time = od._ds["time"]

    # Check Temp and S
    varList = ["Temp", "S"]
    od = _compute._add_missing_variables(od, varList)

    # Fix T and S axes
    Tlim = kwargs.pop("Tlim", None)
    Slim = kwargs.pop("Slim", None)
    if Tlim is None:
        cmap_params = _determine_cmap_params(od._ds["Temp"].values,
                                             center=False)
        Tlim = [cmap_params["vmin"], cmap_params["vmax"]]
    if Slim is None:
        cmap_params = _determine_cmap_params(od._ds["S"].values, center=False)
        Slim = [cmap_params["vmin"], cmap_params["vmax"]]
    kwargs["Tlim"] = Tlim
    kwargs["Slim"] = Slim

    # Fix density
    dens = kwargs.pop("dens", None)
    if dens is None:
        t, s = _xr.broadcast(
            _xr.DataArray(_np.linspace(Tlim[0], Tlim[-1], 100), dims=("t")),
            _xr.DataArray(_np.linspace(Slim[0], Slim[-1], 100), dims=("s")),
        )
        odSigma0 = _ospy.OceanDataset(_xr.Dataset({"Temp": t, "S": s}))
        odSigma0 = odSigma0.set_parameters(od.parameters)
        odSigma0 = odSigma0.compute.potential_density_anomaly()
        odSigma0._ds = odSigma0._ds.set_coords(["Temp", "S"])

        # Freezing point
        paramsList = ["tempFrz0", "dTempFrz_dS"]
        params2use = {
            par: od.parameters[par]
            for par in od.parameters if par in paramsList
        }
        tempFrz0 = params2use["tempFrz0"]
        dTempFrz_dS = params2use["dTempFrz_dS"]
        freez_point = tempFrz0 + odSigma0._ds["S"] * dTempFrz_dS

        # Extract Density
        dens = odSigma0._ds["Sigma0"]
        dens = dens.where(odSigma0._ds["Temp"] > freez_point)
    kwargs["dens"] = dens

    # Fix colorbar
    colorName = kwargs.pop("colorName", None)
    if colorName is not None:

        # Add missing variables (use private)
        _colorName = _rename_aliased(od, colorName)
        od = _compute._add_missing_variables(od, _colorName)

        # Extract color (use public)
        color = od.dataset[colorName]

        # Create colorbar (stolen from xarray)
        cmap_kwargs = kwargs.pop("cmap_kwargs", None)
        if cmap_kwargs is None:
            cmap_kwargs = {}
        cmap_kwargs = dict(cmap_kwargs)
        cmap_kwargs["plot_data"] = color.values
        kwargs["cmap_kwargs"] = _determine_cmap_params(**cmap_kwargs)
    kwargs["colorName"] = colorName

    # Remove ax
    _ax_warning(kwargs)

    # Animation
    anim = _create_animation(od=od,
                             time=time,
                             plot_func=plot_func,
                             func_kwargs=kwargs,
                             display=display,
                             **FuncAnimation_kwargs)
    return anim
예제 #19
0
def TS_diagram(od, display=True, FuncAnimation_kwargs=None,
               **kwargs):
    """
    Animate TS diagrams.

    Parameters
    ----------
    od: OceanDataset
        oceandataset used to plot.
    display: bool
        If True, display the animation.
    FuncAnimation_kwargs: dict
        Keyword arguments from :py:func:`matplotlib.animation.FuncAnimation`
    **kwargs:
        Keyword arguments from :py:func:`oceanspy.plot.TS_diagram`

    Returns
    -------
    anim: matplotlib.animation.FuncAnimation
        Animation object

    See also
    --------
    oceanspy.plot.TS_diagram
    """

    # Check parameters
    _check_instance({'od': od,
                     'display': display,
                     'FuncAnimation_kwargs': FuncAnimation_kwargs},
                    {'od': 'oceanspy.OceanDataset',
                     'display': 'bool',
                     'FuncAnimation_kwargs': ['type(None)', 'dict']})

    # Handle kwargs
    if FuncAnimation_kwargs is None:
        FuncAnimation_kwargs = {}
    FuncAnimation_kwargs = dict(FuncAnimation_kwargs)

    # Name of the plot_functions
    plot_func = eval('_plot.TS_diagram')

    # First cutout and get time
    cutout_kwargs = kwargs.pop('cutout_kwargs', None)
    if cutout_kwargs is not None:
        od = od.subsample.cutout(**cutout_kwargs)
    time = od._ds['time']

    # Check Temp and S
    varList = ['Temp', 'S']
    od = _compute._add_missing_variables(od, varList)

    # Fix T and S axes
    Tlim = kwargs.pop('Tlim', None)
    Slim = kwargs.pop('Slim', None)
    if Tlim is None:
        cmap_params = _determine_cmap_params(od._ds['Temp'].values,
                                             center=False)
        Tlim = [cmap_params['vmin'], cmap_params['vmax']]
    if Slim is None:
        cmap_params = _determine_cmap_params(od._ds['S'].values,
                                             center=False)
        Slim = [cmap_params['vmin'], cmap_params['vmax']]
    kwargs['Tlim'] = Tlim
    kwargs['Slim'] = Slim

    # Fix density
    dens = kwargs.pop('dens', None)
    if dens is None:
        t, s = _xr.broadcast(_xr.DataArray(_np.linspace(Tlim[0],
                                                        Tlim[-1],
                                                        100), dims=('t')),
                             _xr.DataArray(_np.linspace(Slim[0],
                                                        Slim[-1],
                                                        100), dims=('s')))
        odSigma0 = _ospy.OceanDataset(_xr.Dataset({'Temp': t,
                                                   'S': s}))
        odSigma0 = odSigma0.set_parameters(od.parameters)
        odSigma0 = odSigma0.compute.potential_density_anomaly()
        odSigma0._ds = odSigma0._ds.set_coords(['Temp', 'S'])

        # Freezing point
        paramsList = ['tempFrz0', 'dTempFrz_dS']
        params2use = {par: od.parameters[par]
                      for par in od.parameters
                      if par in paramsList}
        tempFrz0 = params2use['tempFrz0']
        dTempFrz_dS = params2use['dTempFrz_dS']
        freez_point = tempFrz0 + odSigma0._ds['S']*dTempFrz_dS

        # Extract Density
        dens = odSigma0._ds['Sigma0']
        dens = dens.where(odSigma0._ds['Temp'] > freez_point)
    kwargs['dens'] = dens

    # Fix colorbar
    colorName = kwargs.pop('colorName', None)
    if colorName is not None:

        # Add missing variables (use private)
        _colorName = _rename_aliased(od, colorName)
        od = _compute._add_missing_variables(od, _colorName)

        # Extract color (use public)
        color = od.dataset[colorName]

        # Create colorbar (stolen from xarray)
        cmap_kwargs = kwargs.pop('cmap_kwargs', None)
        if cmap_kwargs is None:
            cmap_kwargs = {}
        cmap_kwargs = dict(cmap_kwargs)
        cmap_kwargs['plot_data'] = color.values
        kwargs['cmap_kwargs'] = _determine_cmap_params(**cmap_kwargs)
    kwargs['colorName'] = colorName

    # Remove ax
    _ax_warning(kwargs)

    # Animation
    anim = _create_animation(od=od,
                             time=time,
                             plot_func=plot_func,
                             func_kwargs=kwargs,
                             display=display,
                             **FuncAnimation_kwargs)
    return anim
예제 #20
0
def geo_plot(darray,
             ax=None,
             method='contourf',
             projection='PlateCarree',
             grid=False,
             **kwargs):
    """ Create a global plot of a given variable.

    Parameters:
    -----------
    darray : xray.DataArray
        The darray to be plotted.
    ax : axis
        An existing axis instance, else one will be created.
    method : str
        String to use for looking up name of plotting function via iris
    projection : str or tuple
        Name of the cartopy projection to use and any args
        necessary for initializing it passed as a dictionary;
        see func:`make_geoaxes` for more information
    grid : bool
        Include lat-lon grid overlay
    **kwargs : dict
        Any additional keyword arguments to pass to the plotter,
        including colormap params. If 'vmin' is not in this
        set of optional keyword arguments, the plot colormap will be
        automatically inferred.

    """

    # Set up plotting function
    if method in _PLOTTYPE_ARGS:
        extra_args = _PLOTTYPE_ARGS[method].copy()
    else:
        raise ValueError("Don't know how to deal with '%s' method" % method)
    extra_args.update(**kwargs)

    # Alias a plot function based on the requested method and the
    # datatype being plotted
    plot_func = plt.__dict__[method]

    # `transform` should be the ORIGINAL coordinate system -
    # which is always a simple lat-lon coordinate system in CESM
    # output
    extra_args['transform'] = ccrs.PlateCarree()

    # Was an axis passed to plot on?
    new_axis = ax is None

    if new_axis:  # Create a new cartopy axis object for plotting
        if isinstance(projection, (list, tuple)):
            if len(projection) != 2:
                raise ValueError("Expected 'projection' to only have 2 values")
            projection, proj_kwargs = projection[0], projection[1]
        else:
            proj_kwargs = {}

        # hack to look up the name of the projection in the cartopy
        # reference system namespace; makes life a bit easier, so you
        # can just pass a string with the name of the projection wanted.
        proj = ccrs.__dict__[projection](**proj_kwargs)
        ax = plt.axes(projection=proj)
    else:  # Set current axis to one passed as argument
        if not hasattr(ax, 'projection'):
            raise ValueError("Expected `ax` to be a GeoAxes instance")
        plt.sca(ax)

    # Setup map
    ax.set_global()
    ax.coastlines()

    try:
        gl = ax.gridlines(crs=extra_args['transform'],
                          draw_labels=True,
                          linewidth=0.5,
                          color='grey',
                          alpha=0.8)
        LON_TICKS = [-180, -90, 0, 90, 180]
        LAT_TICKS = [-90, -60, -30, 0, 30, 60, 90]
        gl.xlabels_top = False
        gl.ylabels_right = False
        gl.xlines = grid
        gl.ylines = grid
        gl.xlocator = mticker.FixedLocator(LON_TICKS)
        gl.ylocator = mticker.FixedLocator(LAT_TICKS)
        gl.xformatter = LONGITUDE_FORMATTER
        gl.yformatter = LATITUDE_FORMATTER
    except TypeError:
        warnings.warn("Could not label the given map projection.")

    # Infer colormap settings if not provided
    if not ('vmin' in kwargs):
        warnings.warn("Re-inferring color parameters...")
        cmap_kws = _determine_cmap_params(darray.data)
        extra_args.update(cmap_kws)

    gp = plot_func(darray.lon.values, darray.lat.values, darray.data,
                   **extra_args)

    return ax, gp
    # Read colorfile arguments; else, infer the color parameters
    if args.colorfile is not None:
        colorfile = args.colorfile
        try:
            with open(colorfile, 'rb') as f:
                color_data = pickle.load(f)
        except (FileNotFoundError, IOError):
            print("Could not open colorfile '%s'" % args.colorfile)
            sys.exit(1)

        # Warn about variables where color will be freshly inferred
        for v in dataset.variables:
            if v in dataset.dims: continue
            if not (v in color_data):
                warnings.warn("Couldn't find color data for %s" % v)
                color_data[v] = _determine_cmap_params(dataset[v].data)
    else:
        print("Inferring new colormaps")

        color_data = {}
        for v in dataset.variables:
            if v in dataset.dims: continue
            print("   " + v)
            color_data[v] = _determine_cmap_params(dataset[v].data,
                                                   levels=21,
                                                   robust=True,
                                                   extend='both')

    # Save/update the new colorfile
    print("Saving colormap data")
    fn_basename, _ = os.path.splitext(args.nc_file)
예제 #22
0
    def test_divergentcontrol(self):
        neg = self.data - 0.1
        pos = self.data

        # Default with positive data will be a normal cmap
        cmap_params = _determine_cmap_params(pos)
        self.assertEqual(cmap_params["vmin"], 0)
        self.assertEqual(cmap_params["vmax"], 1)
        self.assertEqual(cmap_params["cmap"].name, "viridis")

        # Default with negative data will be a divergent cmap
        cmap_params = _determine_cmap_params(neg)
        self.assertEqual(cmap_params["vmin"], -0.9)
        self.assertEqual(cmap_params["vmax"], 0.9)
        self.assertEqual(cmap_params["cmap"], "RdBu_r")

        # Setting vmin or vmax should prevent this only if center is false
        cmap_params = _determine_cmap_params(neg, vmin=-0.1, center=False)
        self.assertEqual(cmap_params["vmin"], -0.1)
        self.assertEqual(cmap_params["vmax"], 0.9)
        self.assertEqual(cmap_params["cmap"].name, "viridis")
        cmap_params = _determine_cmap_params(neg, vmax=0.5, center=False)
        self.assertEqual(cmap_params["vmin"], -0.1)
        self.assertEqual(cmap_params["vmax"], 0.5)
        self.assertEqual(cmap_params["cmap"].name, "viridis")

        # Setting center=False too
        cmap_params = _determine_cmap_params(neg, center=False)
        self.assertEqual(cmap_params["vmin"], -0.1)
        self.assertEqual(cmap_params["vmax"], 0.9)
        self.assertEqual(cmap_params["cmap"].name, "viridis")

        # However, I should still be able to set center and have a div cmap
        cmap_params = _determine_cmap_params(neg, center=0)
        self.assertEqual(cmap_params["vmin"], -0.9)
        self.assertEqual(cmap_params["vmax"], 0.9)
        self.assertEqual(cmap_params["cmap"], "RdBu_r")

        # Setting vmin or vmax alone will force symetric bounds around center
        cmap_params = _determine_cmap_params(neg, vmin=-0.1)
        self.assertEqual(cmap_params["vmin"], -0.1)
        self.assertEqual(cmap_params["vmax"], 0.1)
        self.assertEqual(cmap_params["cmap"], "RdBu_r")
        cmap_params = _determine_cmap_params(neg, vmax=0.5)
        self.assertEqual(cmap_params["vmin"], -0.5)
        self.assertEqual(cmap_params["vmax"], 0.5)
        self.assertEqual(cmap_params["cmap"], "RdBu_r")
        cmap_params = _determine_cmap_params(neg, vmax=0.6, center=0.1)
        self.assertEqual(cmap_params["vmin"], -0.4)
        self.assertEqual(cmap_params["vmax"], 0.6)
        self.assertEqual(cmap_params["cmap"], "RdBu_r")

        # But this is only true if vmin or vmax are negative
        cmap_params = _determine_cmap_params(pos, vmin=-0.1)
        self.assertEqual(cmap_params["vmin"], -0.1)
        self.assertEqual(cmap_params["vmax"], 0.1)
        self.assertEqual(cmap_params["cmap"], "RdBu_r")
        cmap_params = _determine_cmap_params(pos, vmin=0.1)
        self.assertEqual(cmap_params["vmin"], 0.1)
        self.assertEqual(cmap_params["vmax"], 1)
        self.assertEqual(cmap_params["cmap"].name, "viridis")
        cmap_params = _determine_cmap_params(pos, vmax=0.5)
        self.assertEqual(cmap_params["vmin"], 0)
        self.assertEqual(cmap_params["vmax"], 0.5)
        self.assertEqual(cmap_params["cmap"].name, "viridis")

        # If both vmin and vmax are provided, output is non-divergent
        cmap_params = _determine_cmap_params(neg, vmin=-0.2, vmax=0.6)
        self.assertEqual(cmap_params["vmin"], -0.2)
        self.assertEqual(cmap_params["vmax"], 0.6)
        self.assertEqual(cmap_params["cmap"].name, "viridis")
예제 #23
0
    def test_divergentcontrol(self):
        neg = self.data - 0.1
        pos = self.data

        # Default with positive data will be a normal cmap
        cmap_params = _determine_cmap_params(pos)
        self.assertEqual(cmap_params['vmin'], 0)
        self.assertEqual(cmap_params['vmax'], 1)
        self.assertEqual(cmap_params['cmap'].name, "viridis")

        # Default with negative data will be a divergent cmap
        cmap_params = _determine_cmap_params(neg)
        self.assertEqual(cmap_params['vmin'], -0.9)
        self.assertEqual(cmap_params['vmax'], 0.9)
        self.assertEqual(cmap_params['cmap'], "RdBu_r")

        # Setting vmin or vmax should prevent this only if center is false
        cmap_params = _determine_cmap_params(neg, vmin=-0.1, center=False)
        self.assertEqual(cmap_params['vmin'], -0.1)
        self.assertEqual(cmap_params['vmax'], 0.9)
        self.assertEqual(cmap_params['cmap'].name, "viridis")
        cmap_params = _determine_cmap_params(neg, vmax=0.5, center=False)
        self.assertEqual(cmap_params['vmin'], -0.1)
        self.assertEqual(cmap_params['vmax'], 0.5)
        self.assertEqual(cmap_params['cmap'].name, "viridis")

        # Setting center=False too
        cmap_params = _determine_cmap_params(neg, center=False)
        self.assertEqual(cmap_params['vmin'], -0.1)
        self.assertEqual(cmap_params['vmax'], 0.9)
        self.assertEqual(cmap_params['cmap'].name, "viridis")

        # However, I should still be able to set center and have a div cmap
        cmap_params = _determine_cmap_params(neg, center=0)
        self.assertEqual(cmap_params['vmin'], -0.9)
        self.assertEqual(cmap_params['vmax'], 0.9)
        self.assertEqual(cmap_params['cmap'], "RdBu_r")

        # Setting vmin or vmax alone will force symmetric bounds around center
        cmap_params = _determine_cmap_params(neg, vmin=-0.1)
        self.assertEqual(cmap_params['vmin'], -0.1)
        self.assertEqual(cmap_params['vmax'], 0.1)
        self.assertEqual(cmap_params['cmap'], "RdBu_r")
        cmap_params = _determine_cmap_params(neg, vmax=0.5)
        self.assertEqual(cmap_params['vmin'], -0.5)
        self.assertEqual(cmap_params['vmax'], 0.5)
        self.assertEqual(cmap_params['cmap'], "RdBu_r")
        cmap_params = _determine_cmap_params(neg, vmax=0.6, center=0.1)
        self.assertEqual(cmap_params['vmin'], -0.4)
        self.assertEqual(cmap_params['vmax'], 0.6)
        self.assertEqual(cmap_params['cmap'], "RdBu_r")

        # But this is only true if vmin or vmax are negative
        cmap_params = _determine_cmap_params(pos, vmin=-0.1)
        self.assertEqual(cmap_params['vmin'], -0.1)
        self.assertEqual(cmap_params['vmax'], 0.1)
        self.assertEqual(cmap_params['cmap'], "RdBu_r")
        cmap_params = _determine_cmap_params(pos, vmin=0.1)
        self.assertEqual(cmap_params['vmin'], 0.1)
        self.assertEqual(cmap_params['vmax'], 1)
        self.assertEqual(cmap_params['cmap'].name, "viridis")
        cmap_params = _determine_cmap_params(pos, vmax=0.5)
        self.assertEqual(cmap_params['vmin'], 0)
        self.assertEqual(cmap_params['vmax'], 0.5)
        self.assertEqual(cmap_params['cmap'].name, "viridis")

        # If both vmin and vmax are provided, output is non-divergent
        cmap_params = _determine_cmap_params(neg, vmin=-0.2, vmax=0.6)
        self.assertEqual(cmap_params['vmin'], -0.2)
        self.assertEqual(cmap_params['vmax'], 0.6)
        self.assertEqual(cmap_params['cmap'].name, "viridis")
예제 #24
0
    # Read colorfile arguments; else, infer the color parameters
    if args.colorfile is not None:
        colorfile = args.colorfile
        try:
            with open(colorfile, 'rb') as f:
                color_data = pickle.load(f)
        except (FileNotFoundError, IOError):
            print("Could not open colorfile '%s'" % args.colorfile)
            sys.exit(1)

        # Warn about variables where color will be freshly inferred
        for v in dataset.variables:
            if v in dataset.dims: continue
            if not (v in color_data):
                warnings.warn("Couldn't find color data for %s" % v)
                color_data[v] = _determine_cmap_params(dataset[v].data)
    else:
        print("Inferring new colormaps")

        color_data = {}
        for v in dataset.variables:
            if v in dataset.dims: continue
            print("   " + v)
            color_data[v] = _determine_cmap_params(dataset[v].data,
                                                   levels=21,
                                                   robust=True,
                                                   extend='both')

    # Save/update the new colorfile
    print("Saving colormap data")
    fn_basename, _ = os.path.splitext(args.nc_file)