"""Plotters and formatoption of the psy-maps psyplot plugin."""

# SPDX-FileCopyrightText: 2016-2024 University of Lausanne

# SPDX-FileCopyrightText: 2020-2021 Helmholtz-Zentrum Geesthacht

# SPDX-FileCopyrightText: 2021-2024 Helmholtz-Zentrum hereon GmbH

#

# SPDX-License-Identifier: LGPL-3.0-only

import copy

import re

import warnings

from abc import abstractproperty

from difflib import get_close_matches

from itertools import chain, repeat, starmap

import cartopy

import cartopy.crs as ccrs

import cartopy.feature as cf

import matplotlib as mpl

import matplotlib.ticker as ticker

import numpy as np

import psy_simple.base

import psy_simple.plotters as psyps

import six

from cartopy.mpl.gridliner import Gridliner

from matplotlib.colors import BoundaryNorm

from psy_simple.colors import FixedBoundaryNorm

from psyplot import rcParams

from psyplot.data import InteractiveList, _infer_interval_breaks

from psyplot.docstring import docstrings

from psyplot.plotter import (

BEFOREPLOTTING,

END,

START,

DictFormatoption,

Formatoption,

)

from psy_maps.boxes import lonlatboxes

wrap_proj_types = (

ccrs._RectangularProjection,

ccrs._WarpedRectangularProjection,

ccrs.InterruptedGoodeHomolosine,

ccrs.Mercator,

)

cartopy_version = list(map(int, cartopy.__version__.split(".")[:2]))

@docstrings.get_sections(base="shiftdata")

def shiftdata(lonsin, datain, lon_0):

"""

Shift longitudes (and optionally data) so that they match map projection

region.

Only valid for cylindrical/pseudo-cylindrical global projections and data

on regular lat/lon grids. longitudes and data can be 1-d or 2-d, if 2-d

it is assumed longitudes are 2nd (rightmost) dimension.

Parameters

----------

lonsin

original 1-d or 2-d longitudes.

datain

original 1-d or 2-d data

lon_0

center of map projection region

References

----------

This function is copied and taken from the

:class:`mpl_toolkits.basemap.Basemap` class. The only difference is that

we do not mask values outside the map projection region

"""

lonsin = np.asarray(lonsin)

if lonsin.ndim not in [1, 2]:

raise ValueError("1-d or 2-d longitudes required")

if datain is not None:

# if it's a masked array, leave it alone.

if not np.ma.isMA(datain):

datain = np.asarray(datain)

if datain.ndim not in [1, 2]:

raise ValueError("1-d or 2-d data required")

# 2-d data.

if lonsin.ndim == 2:

raise NotImplementedError(

"Shifting of 2D-data is currently not supported"

)

# 1-d data.

elif lonsin.ndim == 1:

lonsin = np.where(lonsin > lon_0 + 180, lonsin - 360, lonsin)

lonsin = np.where(lonsin < lon_0 - 180, lonsin + 360, lonsin)

londiff = np.abs(lonsin[0:-1] - lonsin[1:])

londiff_sort = np.sort(londiff)

thresh = 360.0 - londiff_sort[-2]

itemindex = len(lonsin) - np.where(londiff >= thresh)[0]

if itemindex:

# check to see if cyclic (wraparound) point included

# if so, remove it.

if np.abs(lonsin[0] - lonsin[-1]) < 1.0e-4:

hascyclic = True

lonsin_save = lonsin.copy()

lonsin = lonsin[1:]

if datain is not None:

datain_save = datain.copy()

datain = datain[1:]

else:

hascyclic = False

lonsin = np.roll(lonsin, itemindex - 1)

if datain is not None:

datain = np.roll(datain, itemindex - 1)

# add cyclic point back at beginning.

if hascyclic:

lonsin_save[1:] = lonsin

lonsin_save[0] = lonsin[-1] - 360.0

lonsin = lonsin_save

if datain is not None:

datain_save[1:] = datain

datain_save[0] = datain[-1]

datain = datain_save

if datain is not None:

return lonsin, datain

else:

return lonsin

def degree_format():

if mpl.rcParams["text.usetex"]:

return r"${%g\/^{\circ}\/%s}$"

else:

return "%g\N{DEGREE SIGN}%s"

def format_lons(x, pos):

fmt_string = degree_format()

x = x - 360 if x > 180 else x

if x == 0:

return fmt_string % (x, "")

return fmt_string % (abs(x), "W" if x < 0 else "E")

def format_lats(x, pos):

fmt_string = degree_format()

if x == 0:

return fmt_string % (abs(x), "")

return fmt_string % (abs(x), "S" if x < 0 else "N")

class Transpose(Formatoption):

"""Transpose the data before plotting

This formatoption can be used in case the data order of the variable

dimensions is (x, y) instead of (y, x)"""

priority = START

group = "axes"

name = "Transpose (switch) x- and y-dimensions"

def update(self, value):

for i, (raw, data) in enumerate(

zip(self.iter_raw_data, self.iter_data)

):

base_var = next(raw.psy.iter_base_variables)

is_unstructured = raw.psy.decoder.is_unstructured(base_var)

if is_unstructured and value:

decoder = copy.copy(data.psy.decoder)

xcoord = data.psy.get_coord("x").name

ycoord = data.psy.get_coord("y").name

decoder.x.add(ycoord)

decoder.y.add(xcoord)

self.set_decoder(decoder, i)

elif is_unstructured:

self.set_decoder(raw.psy.decoder, i)

elif value:

if data.ndim > 2:

new_dims = data.dims[:-2] + data.dims[-2:][::-1]

else:

new_dims = data.dims[::-1]

self.set_data(data.transpose(*new_dims), i)

lon_formatter = ticker.FuncFormatter(format_lons)

lat_formatter = ticker.FuncFormatter(format_lats)

@docstrings.get_sections(

base="ProjectionBase", sections=["Possible types", "See Also"]

)

class ProjectionBase(Formatoption):

"""

Base class for formatoptions that uses cartopy.crs.CRS instances

Possible types

--------------

cartopy.crs.CRS

A cartopy projection instance (e.g. :class:`cartopy.crs.PlateCarree`)

str

A string specifies the projection instance to use. The centered

longitude and latitude are determined by the :attr:`clon` and

:attr:`clat` formatoptions.

Possible strings are (each standing for the specified projection)

=========== =======================================

cf try to decode the CF-conventions

cyl :class:`cartopy.crs.PlateCarree`

robin :class:`cartopy.crs.Robinson`

moll :class:`cartopy.crs.Mollweide`

geo :class:`cartopy.crs.Geostationary`

northpole :class:`cartopy.crs.NorthPolarStereo`

southpole :class:`cartopy.crs.SouthPolarStereo`

ortho :class:`cartopy.crs.Orthographic`

stereo :class:`cartopy.crs.Stereographic`

near :class:`cartopy.crs.NearsidePerspective`

rotated :class:`cartopy.crs.RotatedPole`

=========== =======================================

The special case ``'cf'`` tries to decode the CF-conventions in the

data. If this is not possible, we assume a standard lat-lon projection

(``'cyl'``)

See Also

--------

`Grid-mappings of cf-conventions <http://cfconventions.org/Data/cf-conventions/cf-conventions-1.8/cf-conventions.html#appendix-grid-mappings>`__

"""

projections = {

"cyl": ccrs.PlateCarree,

"robin": ccrs.Robinson,

"moll": ccrs.Mollweide,

"geo": ccrs.Geostationary,

"northpole": ccrs.NorthPolarStereo,

"southpole": ccrs.SouthPolarStereo,

"ortho": ccrs.Orthographic,

"stereo": ccrs.Stereographic,

"near": ccrs.NearsidePerspective,

"rotated": ccrs.RotatedPole,

}

projection_kwargs = dict(

chain(zip(projections.keys(), repeat(["central_longitude"])))

)

projection_kwargs["ortho"] = ["central_longitude", "central_latitude"]

projection_kwargs["stereo"] = ["central_longitude", "central_latitude"]

projection_kwargs["near"] = ["central_longitude", "central_latitude"]

projection_kwargs["rotated"] = ["pole_longitude", "pole_latitude"]

# CRS that are supported to be interpreted by CF-conventions

supported_crs = [

"albers_conical_equal_area",

"azimuthal_equidistant",

"geostationary",

"lambert_azimuthal_equal_area",

"lambert_conformal_conic",

"lambert_cylindrical_equal_area",

"latitude_longitude",

"mercator",

# 'oblique_mercator', # not available for cartopy

"orthographic",

"polar_stereographic",

"rotated_latitude_longitude",

"sinusoidal",

"stereographic",

"transverse_mercator",

# 'vertical_perspective', # not available for cartopy

]

def transform_lonlatbox(self, value):

"""Transform a lon-lat-box to the specific projection"""

value = np.asarray(value)

transformed = self.projection.transform_points(

ccrs.PlateCarree(), value[:2], value[2:]

)[..., :2]

value[:2] = transformed[..., 0]

value[2:] = transformed[..., 1]

if value[0] == value[1] and isinstance(

self.projection, ccrs.PlateCarree

):

value[1] += 360

return value

@property

def cf_projection(self):

data = next(self.iter_data)

if "grid_mapping" in data.attrs:

try:

crs = data[data.attrs["grid_mapping"]]

except KeyError:

try:

crs = data.psy.base[data.attrs["grid_mapping"]]

except KeyError:

warnings.warn(

"Grid mapping variable %(grid_mapping)s specified but "

"impossible to find!" % data.attrs

)

return

if "grid_mapping_name" not in crs.attrs:

warnings.warn(

"Grid mapping variable %(grid_mapping)s specified but "

"without 'grid_mapping_name' attribute!" % data.attrs

)

return

elif crs.attrs["grid_mapping_name"] not in self.supported_crs:

warnings.warn(

"Grid mapping %(grid_mapping_name)s not supported!"

% crs.attrs

)

return

else:

func = getattr(

self, crs.attrs["grid_mapping_name"] + "_from_cf"

)

try:

return func(crs)

except Exception:

self.logger.debug(

"Failed to get CRS from crs variable %s",

crs.name,

exc_info=True,

)

warnings.warn(

f"Failed to get CRS from crs variable {crs.name}"

)

def albers_conical_equal_area_from_cf(self, crs):

kwargs = dict(

central_longitude=crs.longitude_of_central_meridian,

central_latitude=crs.latitude_of_projection_origin,

standard_parallels=crs.standard_parallel,

)

try:

iter(kwargs["standard_parallels"])

except TypeError:

kwargs["standard_parallels"] = [kwargs["standard_parallels"]]

if getattr(crs, "false_easting", None):

kwargs["false_easting"] = crs.false_easting

if getattr(crs, "false_northing", None):

kwargs["false_northing"] = crs.false_northing

return ccrs.AlbersEqualArea(**kwargs)

def azimuthal_equidistant_from_cf(self, crs):

kwargs = dict(

central_longitude=crs.longitude_of_projection_origin,

central_latitude=crs.latitude_of_projection_origin,

)

if getattr(crs, "false_easting", None):

kwargs["false_easting"] = crs.false_easting

if getattr(crs, "false_northing", None):

kwargs["false_northing"] = crs.false_northing

return ccrs.AzimuthalEquidistant(**kwargs)

def geostationary_from_cf(self, crs):

kwargs = dict(

central_longitude=crs.longitude_of_projection_origin,

satellite_height=crs.perspective_point_height,

sweep_axis=crs.sweep_angle_axis,

)

if getattr(crs, "false_easting", None):

kwargs["false_easting"] = crs.false_easting

if getattr(crs, "false_northing", None):

kwargs["false_northing"] = crs.false_northing

return ccrs.Geostationary(**kwargs)

def lambert_azimuthal_equal_area_from_cf(self, crs):

kwargs = dict(

central_longitude=crs.longitude_of_projection_origin,

central_latitude=crs.latitude_of_projection_origin,

)

if getattr(crs, "false_easting", None):

kwargs["false_easting"] = crs.false_easting

if getattr(crs, "false_northing", None):

kwargs["false_northing"] = crs.false_northing

return ccrs.LambertAzimuthalEqualArea(**kwargs)

def lambert_conformal_conic_from_cf(self, crs):

kwargs = dict(

central_longitude=crs.longitude_of_central_meridian,

central_latitude=crs.latitude_of_projection_origin,

standard_parallels=crs.standard_parallel,

)

try:

iter(kwargs["standard_parallels"])

except TypeError:

kwargs["standard_parallels"] = [kwargs["standard_parallels"]]

if getattr(crs, "false_easting", None):

kwargs["false_easting"] = crs.false_easting

if getattr(crs, "false_northing", None):

kwargs["false_northing"] = crs.false_northing

return ccrs.LambertConformal(**kwargs)

def lambert_cylindrical_equal_area_from_cf(self, crs):

kwargs = dict(central_longitude=crs.longitude_of_central_meridian)

return ccrs.LambertCylindrical(**kwargs)

def latitude_longitude_from_cf(self, crs):

return ccrs.PlateCarree()

def mercator_from_cf(self, crs):

kwargs = dict(

central_longitude=crs.longitude_of_projection_origin,

)

if hasattr(crs, "scale_factor_at_projection_origin"):

kwargs["scale_factor"] = crs.scale_factor_at_projection_origin

if getattr(crs, "false_easting", None):

kwargs["false_easting"] = crs.false_easting

if getattr(crs, "false_northing", None):

kwargs["false_northing"] = crs.false_northing

return ccrs.Mercator(**kwargs)

def orthographic_from_cf(self, crs):

kwargs = dict(

central_longitude=crs.longitude_of_projection_origin,

central_latitude=crs.latitude_of_projection_origin,

)

return ccrs.Orthographic(**kwargs)

def polar_stereographic_from_cf(self, crs):

kwargs = dict(

central_longitude=crs.straight_vertical_longitude_from_pole

)

if crs.latitude_of_projection_origin == -90:

return ccrs.SouthPolarStereo(**kwargs)

else:

return ccrs.NorthPolarStereo(**kwargs)

def rotated_latitude_longitude_from_cf(self, crs):

args = [crs.grid_north_pole_longitude, crs.grid_north_pole_latitude]

if hasattr(crs, "north_pole_grid_longitude"):

args.append(crs.north_pole_grid_longitude)

return ccrs.RotatedPole(*args)

def sinusoidal_from_cf(self, crs):

try:

kwargs = dict(

central_longitude=crs.longitude_of_projection_origin,

)

except AttributeError:

kwargs = dict(

central_longitude=crs.longitude_of_central_meridian,

)

if getattr(crs, "false_easting", None):

kwargs["false_easting"] = crs.false_easting

if getattr(crs, "false_northing", None):

kwargs["false_northing"] = crs.false_northing

return ccrs.Sinusoidal(**kwargs)

def stereographic_from_cf(self, crs):

kwargs = dict(

central_latitude=crs.latitude_of_projection_origin,

central_longitude=crs.longitude_of_projection_origin,

scale_factor=crs.scale_factor_at_projection_origin,

)

if getattr(crs, "false_easting", None):

kwargs["false_easting"] = crs.false_easting

if getattr(crs, "false_northing", None):

kwargs["false_northing"] = crs.false_northing

return ccrs.Stereographic(**kwargs)

def transverse_mercator_from_cf(self, crs):

kwargs = dict(

central_longitude=crs.longitude_of_central_meridian,

central_latitude=crs.latitude_of_projection_origin,

scale_factor=crs.scale_factor_at_central_meridian,

)

if getattr(crs, "false_easting", None):

kwargs["false_easting"] = crs.false_easting

if getattr(crs, "false_northing", None):

kwargs["false_northing"] = crs.false_northing

return ccrs.TransverseMercator(**kwargs)

def set_projection(self, value, *args, **kwargs):

if value == "cf":

ret = self.cf_projection

if ret is not None:

return ret

if isinstance(value, ccrs.CRS):

return value

else:

value = "cyl" if value == "cf" else value

return self.projections[value](

**self.get_kwargs(value, *args, **kwargs)

)

def get_kwargs(self, value, clon=None, clat=None):

ret = {}

keys = self.projection_kwargs[value]

if "central_longitude" in keys:

ret["central_longitude"] = self.clon.clon if clon is None else clon

if "central_latitude" in keys:

ret["central_latitude"] = self.clat.clat if clat is None else clat

if "pole_longitude" in keys: # rotated pole

ret["pole_longitude"] = self.clon.clon if clon is None else clon

ret["pole_longitude"] -= 180

if "pole_latitude" in keys: # rotated pole

ret["pole_latitude"] = self.clat.clat if clat is None else clat

self.logger.debug("Setting projection with %s", ret)

return ret

class Projection(ProjectionBase):

"""

Specify the projection for the plot

This formatoption defines the projection of the plot

Possible types

--------------

%(ProjectionBase.possible_types)s

See Also

--------

%(ProjectionBase.see_also)s

Warnings

--------

An update of the projection clears the axes!

"""

# the axes has to be cleared completly if the projection is updated

priority = START

#: an update of this formatoption requires that the axes is cleared

requires_clearing = True

name = "Projection of the plot"

dependencies = ["clon", "clat"]

connections = ["transform", "lonlatbox"]

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

super(Projection, self).__init__(*args, **kwargs)

self.projection = None

@property

def lonlatbox_transformed(self):

"""Transform the lonlatbox according to the projection"""

return self.transform_lonlatbox(self.lonlatbox.lonlatbox)

def initialize_plot(self, value, clear=True):

"""Initialize the plot and set the projection for the axes"""

self.projection = self.set_projection(value)

if self.plotter.cleared:

self.ax.projection = self.projection

self.ax.clear()

def update(self, value):

"""Update the formatoption

Since this formatoption requires clearing, this method does nothing.

Everything is done in the :meth:`initialize_plot` method.

"""

pass

@docstrings.get_sections(base="BoxBase")

class BoxBase(Formatoption):

"""

Abstract base class for specifying a longitude-latitude box

Possible types

--------------

str

A pattern that matches any of the keys in the :attr:`psyplot.rcParams`

``'extents.boxes'`` item (contains user-defined longitude-latitude

boxes) or the :attr:`psyplot.plotter.boxes.lonlatboxes` dictionary

(contains longitude-latitude boxes of different countries and

continents)

[lonmin, lonmax, latmin, latmax]

The surrounding longitude-latitude that shall be used. Values can be

either a float or a string as above

See Also

--------

LonLatBox, MapExtent

"""

def lola_from_pattern(self, s):

"""

Calculate the longitude-latitude box based upon a pattern

This method uses the psyplot.rcParams ``'extents.boxes'`` item to find

longitude that match `s` and takes the surrounding box.

Parameters

----------

s: str

The pattern to use for the keys in the

:attr:`psyplot.plotter.maps.lonlatboxes` dictionary and the

``'extents.boxes'`` item in the :attr:`psyplot.rcParams`

Returns

-------

float: lonmin, lonmax, latmin, latmax or None

The surrounding longitude-latitude box of all items in

``psyplot.rcParams['extents.boxes']`` whose key match `s` if there

was any match. Otherwise None is returned

"""

patt = re.compile(s)

boxes = np.array(

[

box

for key, box in chain(

*map(

six.iteritems,

[lonlatboxes, rcParams["lonlatbox.boxes"]],

)

)

if patt.search(key)

]

)

if len(boxes) == 0:

similar_keys = get_close_matches(s, rcParams["lonlatbox.boxes"])

message = "Did not find any matches for %s!" % s

if similar_keys:

message += " Maybe you mean on of " + ", ".join(similar_keys)

warnings.warn(message, RuntimeWarning)

return

return [

boxes[:, 0].min(),

boxes[:, 1].max(),

boxes[:, 2].min(),

boxes[:, 3].max(),

]

docstrings.keep_types("BoxBase.possible_types", "str", "str")

class CenterLon(BoxBase):

"""

Set the center longitude of the plot

Parameters

----------

None

Let the :attr:`lonlatbox` formatoption determine the center

float

Specifiy the center manually

%(BoxBase.possible_types.str)s

"""

priority = START

name = "Longitude of the center of the plot"

requires_clearing = True

dependencies = ["lonlatbox"]

def update(self, value):

self.lon_mean = float(np.mean(self.lonlatbox.lonlatbox[:2]))

if value is not None:

if isinstance(value, six.string_types):

box = self.lola_from_pattern(value)

if box is not None:

self.clon = float(np.mean(box[:2]))

else:

value = None

else:

self.clon = value

if value is None:

lonmin, lonmax = self.lonlatbox.lonlatbox[:2]

if lonmax - lonmin > 350:

self.clon = 0.0

else:

self.clon = self.lon_mean

class CenterLat(BoxBase):

"""

Set the center latitude of the plot

Parameters

----------

None

Let the :attr:`lonlatbox` formatoption determine the center

float

Specifiy the center manually

%(BoxBase.possible_types.str)s

"""

priority = START

name = "Latitude of the center of the plot"

requires_clearing = True

dependencies = ["lonlatbox"]

def update(self, value):

self.lat_mean = float(np.mean(self.lonlatbox.lonlatbox[2:]))

if value is not None:

if isinstance(value, six.string_types):

box = self.lola_from_pattern(value)

if box is not None:

self.clat = float(np.mean(box[2:]))

else:

value = None

else:

self.clat = value

if value is None:

latmin, latmax = self.lonlatbox.lonlatbox[2:]

if latmax - latmin > 170:

self.clat = 0

else:

self.clat = self.lat_mean

@docstrings.get_sections(base="LonLatBox")

class LonLatBox(BoxBase):

"""

Set the longitude-latitude box of the data shown

This formatoption extracts the data that matches the specified box.

Possible types

--------------

None

Use the full data

%(BoxBase.possible_types)s

Notes

-----

- For only specifying the region of the plot, see the :attr:`map_extent`

formatoption

- If the coordinates are two-dimensional (e.g. for a circumpolar grid),

than the data is not extracted but values outside the specified

longitude-latitude box are set to NaN

See Also

--------

map_extent"""

priority = START

name = "Longitude-Latitude box of the data"

requires_clearing = True

dependencies = ["transform", "transpose"]

@property

def lonlatbox_transformed(self):

return self.transform.transform_lonlatbox(self.lonlatbox)

def data_dependent(self, data, set_data=True):

if isinstance(data, InteractiveList):

data = data[0]

decoder = data.psy.decoder

lon, lat = self._get_lola(data, decoder)

new_lonlatbox = self.calc_lonlatbox(

lon, lat, decoder.is_unstructured(data)

)

update = self.data_lonlatbox != new_lonlatbox

if not update and set_data and self.value is not None:

self.update(self.value)

elif update:

self.logger.debug(

"Reinitializing because lonlatbox of new data %s does not "

"match the old one %s",

new_lonlatbox,

self.data_lonlatbox,

)

return update

def update(self, value):

if isinstance(self.data, InteractiveList):

for i, arr in enumerate(self.data):

decoder = self.get_decoder(i)

self.data[i] = self.update_array(

value,

arr,

decoder,

next(six.itervalues(self.raw_data[i].psy.base_variables)),

)

else:

arr = self.data

arr = self.update_array(

value,

arr,

self.decoder,

next(six.itervalues(self.raw_data.psy.base_variables)),

)

self.data = arr

def _get_lola(self, data, decoder):

"""Get longitude and latitde informations from the given data array"""

lon_da = decoder.get_x(data, data.coords)

lat_da = decoder.get_y(data, data.coords)

lon, lat = self.to_degree(

lon_da.attrs.get("units"), lon_da.values, lat_da.values

)

data_shape = (

data.shape[-2:]

if not decoder.is_unstructured(data)

else (data.shape[-1],)

)

is_combined = data_shape != data.shape

if lon.shape == data_shape and lat.shape == data_shape:

vals = data.values[0 if is_combined else slice(None)]

lon = np.where(np.isnan(vals), np.nan, lon)

lat = np.where(np.isnan(vals), np.nan, lat)

return lon, lat

def update_array(self, value, data, decoder, base_var=None):

"""Update the given `data` array"""

lon, lat = self._get_lola(data, decoder)

self.data_lonlatbox = self.calc_lonlatbox(

lon, lat, decoder.is_unstructured(data)

)

if value is None:

self.lonlatbox = self.data_lonlatbox

return data

else:

if isinstance(value, six.string_types):

value = self.lola_from_pattern(value)

if value is None:

self.lonlatbox = self.data_lonlatbox

return data

value = list(value)

for i, v in enumerate(value):

if isinstance(v, six.string_types):

value[i] = self.lola_from_pattern(v)[i]

is_unstructured = decoder.is_unstructured(

base_var if base_var is not None else data

)

is_rectilinear = lon.ndim == 1 and not is_unstructured

shift = isinstance(self.transform.projection, ccrs.PlateCarree)

if is_rectilinear and shift:

data = data.copy(True)

lon, data.values = self.shiftdata(

lon, data.values, np.mean(value[:2])

)

lon_name = decoder.get_x(data, data.coords).name

data[lon_name] = (

data[lon_name].dims,

lon,

data[lon_name].attrs,

)

elif is_unstructured and shift:

# make sure that we are inside the map extent

ret = self.transform.projection.transform_points(

self.transform.projection, lon, lat

)

lon = ret[..., 0]

lat = ret[..., 1]

self.lonlatbox = value

value = self.lonlatbox_transformed

with warnings.catch_warnings():

warnings.filterwarnings(

"ignore", "invalid value encountered", RuntimeWarning

)

if is_rectilinear and (lat[1:] < lat[:-1]).all():

lat_values = value[3:1:-1]

else:

lat_values = value[2:]

if is_rectilinear:

kwargs = dict(

zip(

data.dims[-2:],

starmap(slice, [lat_values, value[:2]]),

)

)

ret = data.sel(**kwargs)

else:

ret = self.mask_outside(

data.copy(True),

lon,

lat,

*value,

is_unstructured=is_unstructured,

decoder=decoder,

)

lon, lat = self._get_lola(ret, decoder)

self.data_lonlatbox = self.calc_lonlatbox(

lon, lat, is_unstructured

)

ret.psy.decoder = decoder

decoder.clear_cache()

return ret

def to_degree(self, units=None, *args):

"""Converts arrays with radian units to degree

Parameters

----------

units: str

if ``'radian'``, the arrays in ``*args`` will be converted

``*args``

numpy arrays

Returns

-------

list of np.ndarray

returns the arrays provided with ``*args``

Notes

-----

if `units` is ``'radian'``, a copy of the array will be returned"""

args = list(args)

if units == "radian" and isinstance(

self.transform.projection, ccrs.PlateCarree

):

for i, array in enumerate(args):

args[i] = array * 180.0 / np.pi

return args

def mask_outside(

self,

data,

lon,

lat,

lonmin,

lonmax,

latmin,

latmax,

is_unstructured=False,

decoder=None,

):

data.values = data.values.copy()

ndim = 2 if not is_unstructured else 1

if lonmax < lonmin:

lonmax += 360

lon[lon < 0] += 360

mask = np.any(

[lon < lonmin, lon > lonmax, lat < latmin, lat > latmax], axis=0

)

decoder = decoder or data.psy.decoder

if data.ndim > ndim:

if is_unstructured and decoder.supports_spatial_slicing:

data = data.psy[:, np.where(~mask)[0]]

else:

for i, arr in enumerate(data.values):

arr[mask] = np.nan

data.values[i, :] = arr

else:

if is_unstructured and decoder.supports_spatial_slicing:

data = data.psy[np.where(~mask)[0]]

else:

data.values[mask] = np.nan

return data

def calc_lonlatbox(self, lon, lat, is_unstructured=False):

if isinstance(self.transform.projection, ccrs.PlateCarree):

with warnings.catch_warnings():

warnings.filterwarnings(

"ignore", "invalid value encountered", RuntimeWarning

)

lon = lon[np.all([lon >= -180, lon <= 360], axis=0)]

lat = lat[np.all([lat >= -90, lat <= 90], axis=0)]

return [lon.min(), lon.max(), lat.min(), lat.max()]

else:

if lon.ndim == 1 and not is_unstructured:

lon, lat = np.meshgrid(lon, lat)

points = ccrs.PlateCarree().transform_points(

self.transform.projection, lon, lat

)

lon = points[..., 0]

lat = points[..., 1]

lon_0 = self.transform.projection.proj4_params.get("lon_0")

if lon_0 is not None:

lon = np.where(lon > lon_0 + 180, lon - 360, lon)

lon = np.where(lon < lon_0 - 180, lon + 360, lon)

return [lon.min(), lon.max(), lat.min(), lat.max()]

def shiftdata(self, lonsin, datain, lon_0):

"""

Shift the data such that it matches the region we want to show

Parameters

----------

%(shiftdata.parameters)s

Notes

-----

`datain` can also be multiple fields stored in a three-dimensional

array. Then we shift all fields along the first dimension

"""

# shiftdata does not work properly if we do not give the bounds of

# the array

lonsin = lonsin.copy()

if lonsin.ndim == 1:

get_centers = True

lonsin = _infer_interval_breaks(lonsin)

else:

get_centers = False

if datain.ndim == 2:

lonsin, datain = shiftdata(lonsin, datain, lon_0)

elif datain.ndim == 3:

lon_save = lonsin.copy()

for i, a in enumerate(datain):

lonsin, datain[i] = shiftdata(lon_save, a, lon_0)

if get_centers:

lonsin = np.mean([lonsin[1:], lonsin[:-1]], axis=0)

return lonsin, datain