diff --git a/environment.yml b/environment.yml index dc35dc6b1..2e9519a3c 100644 --- a/environment.yml +++ b/environment.yml @@ -28,5 +28,6 @@ dependencies: - sphinx-design - networkx - pyarrow + - cftime - pip: - git+https://github.com/ultraplot/UltraTheme.git diff --git a/ultraplot/internals/rcsetup.py b/ultraplot/internals/rcsetup.py index e719d1f96..143fa2378 100644 --- a/ultraplot/internals/rcsetup.py +++ b/ultraplot/internals/rcsetup.py @@ -264,6 +264,37 @@ def _validate_belongs(value): # noqa: E306 return _validate_belongs +_CFTIME_RESOLUTIONS = ( + "SECONDLY", + "MINUTELY", + "HOURLY", + "DAILY", + "MONTHLY", + "YEARLY", +) + + +def _validate_cftime_resolution_format(units: dict) -> dict: + if not isinstance(units, dict): + raise ValueError("Cftime units expects a dict") + + for resolution, format_ in units.items(): + unit = _validate_cftime_resolution(resolution) + + # Delegate format parsing to cftime + _rc_ultraplot_default["cftime.time_resolution_format"].update(units) + return _rc_ultraplot_default["cftime.time_resolution_format"] + + +def _validate_cftime_resolution(unit: str) -> str: + if not isinstance(unit, str): + raise TypeError("Time unit cftime is expecting str") + if unit in _CFTIME_RESOLUTIONS: + return unit + msg = f"Unit not understood. Got {unit} expected one of: {_CFTIME_RESOLUTIONS}" + raise ValueError(msg) + + def _validate_cmap(subtype): """ Validate the colormap or cycle. Possibly skip name registration check @@ -997,6 +1028,33 @@ def copy(self): _validate_fontweight, "Font weight for column labels on the bottom of the figure.", ), + "cftime.time_unit": ( + "days since 2000-01-01", + _validate_string, + "Time unit for non-Gregorian calendars.", + ), + "cftime.resolution": ( + "DAILY", + _validate_cftime_resolution, + "Default time resolution for non-Gregorian calendars.", + ), + "cftime.time_resolution_format": ( + { + "SECONDLY": "%S", + "MINUTELY": "%M", + "HOURLY": "%H", + "DAILY": "%d", + "MONTHLY": "%m", + "YEARLY": "%Y", + }, + _validate_cftime_resolution_format, + "Dict used for formatting non-Gregorian calendars.", + ), + "cftime.max_display_ticks": ( + 7, + _validate_int, + "Number of ticks to display for cftime units.", + ), # Coastlines "coast": (False, _validate_bool, "Toggles coastline lines on and off."), "coast.alpha": ( diff --git a/ultraplot/tests/test_tickers.py b/ultraplot/tests/test_tickers.py new file mode 100644 index 000000000..61791fa05 --- /dev/null +++ b/ultraplot/tests/test_tickers.py @@ -0,0 +1,785 @@ +import pytest, numpy as np, xarray as xr, ultraplot as uplt, cftime +from ultraplot.ticker import AutoCFDatetimeLocator +from unittest.mock import patch +import importlib +import cartopy.crs as ccrs + + +@pytest.mark.mpl_image_compare +def test_datetime_calendars_comparison(): + # Time axis centered at mid-month + # Standard calendar + time1 = xr.date_range("2000-01", periods=120, freq="MS") + time2 = xr.date_range("2000-02", periods=120, freq="MS") + time = time1 + 0.5 * (time2 - time1) + # Non-standard calendar (uses cftime) + time1 = xr.date_range("2000-01", periods=120, freq="MS", calendar="noleap") + time2 = xr.date_range("2000-02", periods=120, freq="MS", calendar="noleap") + time_noleap = time1 + 0.5 * (time2 - time1) + + da = xr.DataArray( + data=np.sin(np.arange(0.0, 2 * np.pi, np.pi / 60.0)), + dims=("time",), + coords={ + "time": time, + }, + attrs={"long_name": "low freq signal", "units": "normalized"}, + ) + + da_noleap = xr.DataArray( + data=np.sin(2.0 * np.arange(0.0, 2 * np.pi, np.pi / 60.0)), + dims=("time",), + coords={ + "time": time_noleap, + }, + attrs={"long_name": "high freq signal", "units": "normalized"}, + ) + + fig, axs = uplt.subplots(ncols=2) + axs.format(title=("Standard calendar", "Non-standard calendar")) + axs[0].plot(da) + axs[1].plot(da_noleap) + + return fig + + +@pytest.mark.parametrize( + "calendar", + [ + "standard", + "gregorian", + "proleptic_gregorian", + "julian", + "all_leap", + "360_day", + "365_day", + "366_day", + ], +) +def test_datetime_calendars(calendar): + time = xr.date_range("2000-01-01", periods=365 * 10, calendar=calendar) + da = xr.DataArray( + data=np.sin(np.linspace(0, 2 * np.pi, 365 * 10)), + dims=("time",), + coords={"time": time}, + ) + fig, ax = uplt.subplots() + ax.plot(da) + ax.format(title=f"Calendar: {calendar}") + + +@pytest.mark.mpl_image_compare +def test_datetime_short_range(): + time = xr.date_range("2000-01-01", periods=10, calendar="standard") + da = xr.DataArray( + data=np.sin(np.linspace(0, 2 * np.pi, 10)), + dims=("time",), + coords={"time": time}, + ) + fig, ax = uplt.subplots() + ax.plot(da) + ax.format(title="Short time range (days)") + return fig + + +@pytest.mark.mpl_image_compare +def test_datetime_long_range(): + time = xr.date_range( + "2000-01-01", periods=365 * 200, calendar="standard" + ) # 200 years + da = xr.DataArray( + data=np.sin(np.linspace(0, 2 * np.pi, 365 * 200)), + dims=("time",), + coords={"time": time}, + ) + fig, ax = uplt.subplots() + ax.plot(da) + ax.format(title="Long time range (centuries)") + return fig + + +def test_datetime_explicit_formatter(): + time = xr.date_range("2000-01-01", periods=365 * 2, calendar="noleap") + da = xr.DataArray( + data=np.sin(np.linspace(0, 2 * np.pi, 365 * 2)), + dims=("time",), + coords={"time": time}, + ) + fig, ax = uplt.subplots() + + formatter = uplt.ticker.CFDatetimeFormatter("%b %Y", calendar="noleap") + locator = uplt.ticker.AutoCFDatetimeLocator(calendar="noleap") + ax.xaxis.set_major_locator(locator) + ax.xaxis.set_major_formatter(formatter) + ax.plot(da) + + fig.canvas.draw() + + labels = [label.get_text() for label in ax.get_xticklabels()] + assert len(labels) > 1 + # check first label + import cftime + + cftime.datetime.strptime(labels[1], "%b %Y", calendar="noleap") + + +@pytest.mark.parametrize( + "date1, date2, num1, num2, expected_resolution, expected_n", + [ + ( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + cftime.datetime(2020, 1, 1, calendar="gregorian"), + 0, + 7305, + "YEARLY", + 20, + ), + ( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + cftime.datetime(2001, 1, 1, calendar="gregorian"), + 0, + 365, + "MONTHLY", + 12, + ), + ( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + cftime.datetime(2000, 1, 10, calendar="gregorian"), + 0, + 9, + "DAILY", + 9, + ), + ( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + cftime.datetime(2000, 1, 1, 12, calendar="gregorian"), + 0, + 0.5, + "HOURLY", + 12, + ), + ( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + cftime.datetime(2000, 1, 1, 0, 0, 10, calendar="gregorian"), + 0, + 1.1574074074074073e-4, + "SECONDLY", + 10, + ), + ], +) +def test_compute_resolution(date1, date2, num1, num2, expected_resolution, expected_n): + locator = AutoCFDatetimeLocator() + resolution, n = locator.compute_resolution(num1, num2, date1, date2) + assert resolution == expected_resolution + assert np.allclose(n, expected_n) + + +@pytest.mark.parametrize( + "date1, date2, num1, num2", + [ + ( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + cftime.datetime(2020, 1, 1, calendar="gregorian"), + 0, + 7305, + ), + ( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + cftime.datetime(2001, 1, 1, calendar="gregorian"), + 0, + 365, + ), + ( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + cftime.datetime(2000, 1, 10, calendar="gregorian"), + 0, + 9, + ), + ( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + cftime.datetime(2000, 1, 1, 12, calendar="gregorian"), + 0, + 0.5, + ), + ( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + cftime.datetime(2000, 1, 1, 0, 0, 10, calendar="gregorian"), + 0, + 1.1574074074074073e-4, + ), + # Additional test cases to cover the while loop + ( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + cftime.datetime(2000, 12, 31, calendar="gregorian"), + 0, + 365, + ), + ( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + cftime.datetime(2000, 1, 2, calendar="gregorian"), + 0, + 1, + ), + ], +) +def test_tick_values(date1, date2, num1, num2): + locator = AutoCFDatetimeLocator() + locator.compute_resolution(num1, num2, date1, date2) + ticks = locator.tick_values(num1, num2) + assert len(ticks) > 0 + assert all( + isinstance( + cftime.num2date(t, locator.date_unit, calendar=locator.calendar), + cftime.datetime, + ) + for t in ticks + ) + + +def test_datetime_maxticks(): + time = xr.date_range("2000-01-01", periods=365 * 20, calendar="noleap") + da = xr.DataArray( + data=np.sin(np.linspace(0, 2 * np.pi, 365 * 20)), + dims=("time",), + coords={"time": time}, + ) + fig, ax = uplt.subplots() + ax.plot(da) + locator = ax.xaxis.get_major_locator() + locator.set_params(maxticks=6) + fig.canvas.draw() + assert len(ax.get_xticks()) <= 6 + + +@pytest.mark.parametrize("module_name", ["cftime", "cartopy.crs"]) +def test_missing_modules(module_name): + """Test fallback behavior when modules are missing.""" + with patch.dict("sys.modules", {module_name: None}): + # Reload the ultraplot.ticker module to apply the mocked sys.modules + import ultraplot.ticker + + importlib.reload(ultraplot.ticker) + + if module_name == "cftime": + from ultraplot.ticker import cftime + + assert cftime is None + elif module_name == "ccrs": + from ultraplot.ticker import ( + ccrs, + LatitudeFormatter, + LongitudeFormatter, + _PlateCarreeFormatter, + ) + + assert ccrs is None + assert LatitudeFormatter is object + assert LongitudeFormatter is object + assert _PlateCarreeFormatter is object + + +def test_index_locator(): + from ultraplot.ticker import IndexLocator + + # Initialize with default values + locator = IndexLocator() + assert locator._base == 1 + assert locator._offset == 0 + + # Update parameters + locator.set_params(base=2, offset=1) + assert locator._base == 2 + assert locator._offset == 1 + + +def test_default_precision_zerotrim(): + from ultraplot.ticker import _default_precision_zerotrim + + # Case 1: Default behavior + precision, zerotrim = _default_precision_zerotrim() + assert precision == 6 # Default when zerotrim is True + assert zerotrim is True + + # Case 2: Explicit precision and zerotrim + precision, zerotrim = _default_precision_zerotrim(precision=3, zerotrim=False) + assert precision == 3 + assert zerotrim is False + + +def test_index_locator_tick_values(): + from ultraplot.ticker import IndexLocator + + locator = IndexLocator(base=2, offset=1) + ticks = locator.tick_values(0, 10) + assert np.array_equal(ticks, [1, 3, 5, 7, 9]) + + +def test_discrete_locator_call(): + from ultraplot.ticker import DiscreteLocator + + locator = DiscreteLocator(locs=[0, 1, 2, 3, 4]) + ticks = locator() + assert np.array_equal(ticks, [0, 1, 2, 3, 4]) + + +def test_discrete_locator_set_params(): + from ultraplot.ticker import DiscreteLocator + + locator = DiscreteLocator(locs=[0, 1, 2, 3, 4]) + locator.set_params(steps=[1, 2], nbins=3, minor=True, min_n_ticks=2) + + assert np.array_equal(locator._steps, [1, 2, 10]) + assert locator._nbins == 3 + assert locator._minor is True + assert locator._min_n_ticks == 2 + + +def test_discrete_locator_tick_values(): + from ultraplot.ticker import DiscreteLocator + + # Create a locator with specific locations + locator = DiscreteLocator(locs=[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]) + locator.set_params(steps=[1, 2], nbins=5, minor=False, min_n_ticks=2) + + ticks = locator.tick_values(None, None) + + assert np.array_equal(ticks, list(range(10))) + + +@pytest.mark.parametrize( + "value, string, expected", + [ + (1e-10, "0.0", True), # Case 1: Small number close to zero + (1000, "1000", False), # Case 2: Large number + ], +) +def test_auto_formatter_fix_small_number(value, string, expected): + from ultraplot.ticker import AutoFormatter + + formatter = AutoFormatter() + result = formatter._fix_small_number(value, string) + if expected: + assert result != string + else: + assert result == string + + +@pytest.mark.parametrize( + "start_date, end_date", + [ + ( + cftime.date2num( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + "days since 2000-01-01", + ), + cftime.date2num( + cftime.datetime(2020, 1, 1, calendar="gregorian"), + "days since 2000-01-01", + ), + ), + # Case 2: Monthly resolution + ( + cftime.date2num( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + "days since 2000-01-01", + ), + cftime.date2num( + cftime.datetime(2000, 12, 31, calendar="gregorian"), + "days since 2000-01-01", + ), + ), + ], +) +def test_auto_datetime_locator_tick_values(start_date, end_date): + from ultraplot.ticker import AutoCFDatetimeLocator + + locator = AutoCFDatetimeLocator(calendar="gregorian") + import cftime + + ticks = locator.tick_values(start_date, end_date) + assert len(ticks) > 0 # Ensure ticks are generated + + +@pytest.mark.parametrize( + "start_date, end_date, calendar, expected_exception, expected_resolution", + [ + ( + cftime.date2num( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + "days since 2000-01-01", + ), + cftime.date2num( + cftime.datetime(2020, 1, 1, calendar="gregorian"), + "days since 2000-01-01", + ), + "gregorian", + None, + "YEARS", + ), # Case 1: Valid yearly resolution + ( + cftime.date2num( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + "days since 2000-01-01", + ), + cftime.date2num( + cftime.datetime(2000, 12, 31, calendar="gregorian"), + "days since 2000-01-01", + ), + "gregorian", + None, + "MONTHS", + ), # Case 2: Valid monthly resolution + ( + cftime.date2num( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + "days since 2000-01-01", + ), + cftime.date2num( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + "days since 2000-01-01", + ), + "gregorian", + None, + "DAYS", + ), # Case 3: Empty range + ( + cftime.date2num( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + "days since 2000-01-01", + ), + cftime.date2num( + cftime.datetime(2000, 12, 31, calendar="gregorian"), + "days since 2000-01-01", + ), + "gregorian", + None, + "MONTHS", + ), # Case 4: Months data range + ( + cftime.date2num( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + "days since 2000-01-01", + ), + cftime.date2num( + cftime.datetime(2000, 1, 1, calendar="gregorian"), + "days since 2000-01-01", + ), + "gregorian", + None, + None, + ), # Case 5: Empty range (valid calendar) + ], +) +def test_auto_datetime_locator_tick_values( + start_date, + end_date, + calendar, + expected_exception, + expected_resolution, +): + from ultraplot.ticker import AutoCFDatetimeLocator + import cftime + + locator = AutoCFDatetimeLocator(calendar=calendar) + resolution = expected_resolution + if expected_exception == ValueError: + with pytest.raises( + ValueError, match="Incorrectly formatted CF date-time unit_string" + ): + cftime.date2num( + cftime.datetime(2000, 1, 1, calendar="gregorian"), "invalid unit" + ) + if expected_exception: + with pytest.raises(expected_exception): + locator.tick_values(start_date, end_date) + else: + ticks = locator.tick_values(start_date, end_date) + assert len(ticks) > 0 + + # Verify that the ticks are at the correct resolution + if expected_resolution == "YEARLY": + assert all( + cftime.num2date(t, locator.date_unit, calendar=locator.calendar).month + == 1 + for t in ticks + ) + assert all( + cftime.num2date(t, locator.date_unit, calendar=locator.calendar).day + == 1 + for t in ticks + ) + elif expected_resolution == "MONTHLY": + assert all( + cftime.num2date(t, locator.date_unit, calendar=locator.calendar).day + == 1 + for t in ticks + ) + elif expected_resolution == "DAILY": + assert all( + cftime.num2date(t, locator.date_unit, calendar=locator.calendar).hour + == 0 + for t in ticks + ) + elif expected_resolution == "HOURLY": + assert all( + cftime.num2date(t, locator.date_unit, calendar=locator.calendar).minute + == 0 + for t in ticks + ) + elif expected_resolution == "MINUTELY": + assert all( + cftime.num2date(t, locator.date_unit, calendar=locator.calendar).second + == 0 + for t in ticks + ) + + +@pytest.mark.parametrize( + "vmin, vmax, expected_ticks", + [ + (5, 5, []), # Case 1: Empty range + (-5, 5, [1, 3]), # Case 2: Negative range + ], +) +def test_index_locator_tick_values_edge_cases(vmin, vmax, expected_ticks): + from ultraplot.ticker import IndexLocator + + locator = IndexLocator(base=2, offset=1) + ticks = locator.tick_values(vmin, vmax) + print(f"vmin: {vmin}, vmax: {vmax}, ticks: {ticks}") + assert np.array_equal(ticks, expected_ticks) + + +@pytest.mark.parametrize( + "locs, steps, nbins, expected_length, expected_ticks", + [ + ([], None, None, 0, []), # Case 1: Empty locs + ([5], None, None, 1, [5]), # Case 2: Single loc + (np.arange(0, 100, 1), [1, 2], 10, 100, None), # Case 3: Large range with steps + ], +) +def test_discrete_locator_tick_values_edge_cases( + locs, steps, nbins, expected_length, expected_ticks +): + from ultraplot.ticker import DiscreteLocator + + locator = DiscreteLocator(locs=locs) + if steps and nbins: + locator.set_params(steps=steps, nbins=nbins) + ticks = locator.tick_values(None, None) + + assert len(ticks) == expected_length # Check the number of ticks + if expected_ticks is not None: + assert np.array_equal(ticks, expected_ticks) # Check the tick values + + +@pytest.mark.parametrize( + "steps, expected", + [ + ((0, 2000000), [1, 2, 3, 6, 10]), # large range + ((0, 2), [1, 1.5, 2, 2.5, 3, 5, 10]), # small range + ], +) +def test_degree_locator_guess_steps(steps, expected): + from ultraplot.ticker import DegreeLocator + + locator = DegreeLocator() + locator._guess_steps(*steps) + assert np.array_equal(locator._steps, expected) + + +import pytest + + +@pytest.mark.parametrize( + "formatter_args, value, expected", + [ + ({"precision": 2, "zerotrim": True}, 12345, r"$1.23{\times}10^{4}$"), + ({"precision": 2, "zerotrim": True}, 0.12345, r"$1.23{\times}10^{−1}$"), + ({"precision": 2, "zerotrim": True}, 0, r"$0$"), + ({"precision": 2, "zerotrim": True}, 1.2, r"$1.2$"), + ({"precision": 2, "zerotrim": False}, 1.2, r"$1.20$"), + ], +) +def test_sci_formatter(formatter_args, value, expected): + from ultraplot.ticker import SciFormatter + + formatter = SciFormatter(**formatter_args) + assert formatter(value) == expected + + +@pytest.mark.parametrize( + "formatter_args, value, expected", + [ + ({"sigfig": 3}, 12345, "12300"), + ({"sigfig": 3}, 123.45, "123"), + ({"sigfig": 3}, 0.12345, "0.123"), + ({"sigfig": 3}, 0.0012345, "0.00123"), + ({"sigfig": 2, "base": 5}, 87, "85"), + ({"sigfig": 2, "base": 5}, 8.7, "8.5"), + ], +) +def test_sig_fig_formatter(formatter_args, value, expected): + from ultraplot.ticker import SigFigFormatter + + formatter = SigFigFormatter(**formatter_args) + assert formatter(value) == expected + + +@pytest.mark.parametrize( + "formatter_args, value, expected", + [ + ( + {"symbol": r"$\\pi$", "number": 3.141592653589793}, + 3.141592653589793 / 2, + r"$\\pi$/2", + ), + ( + {"symbol": r"$\\pi$", "number": 3.141592653589793}, + 3.141592653589793, + r"$\\pi$", + ), + ( + {"symbol": r"$\\pi$", "number": 3.141592653589793}, + 2 * 3.141592653589793, + r"2$\\pi$", + ), + ({"symbol": r"$\\pi$", "number": 3.141592653589793}, 0, "0"), + ({}, 0.5, "1/2"), + ({}, 1, "1"), + ({}, 1.5, "3/2"), + ], +) +def test_frac_formatter(formatter_args, value, expected): + from ultraplot.ticker import FracFormatter + + formatter = FracFormatter(**formatter_args) + assert formatter(value) == expected + + +def test_frac_formatter_unicode_minus(): + from ultraplot.ticker import FracFormatter + from ultraplot.config import rc + import numpy as np + + formatter = FracFormatter(symbol=r"$\\pi$", number=np.pi) + with rc.context({"axes.unicode_minus": True}): + assert formatter(-np.pi / 2) == r"−$\\pi$/2" + + +@pytest.mark.parametrize( + "fmt, calendar, dt_args, expected", + [ + ("%Y-%m-%d", "noleap", (2001, 2, 28), "2001-02-28"), + ], +) +def test_cfdatetime_formatter_direct_call(fmt, calendar, dt_args, expected): + from ultraplot.ticker import CFDatetimeFormatter + import cftime + + formatter = CFDatetimeFormatter(fmt, calendar=calendar) + dt = cftime.datetime(*dt_args, calendar=calendar) + assert formatter(dt) == expected + + +@pytest.mark.parametrize( + "start_date_str, end_date_str, calendar, resolution", + [ + ("2000-01-01 00:00:00", "2000-01-01 12:00:00", "standard", "HOURLY"), + ("2000-01-01 00:00:00", "2000-01-01 00:10:00", "standard", "MINUTELY"), + ("2000-01-01 00:00:00", "2000-01-01 00:00:10", "standard", "SECONDLY"), + ], +) +def test_autocftime_locator_subdaily( + start_date_str, end_date_str, calendar, resolution +): + from ultraplot.ticker import AutoCFDatetimeLocator + import cftime + + locator = AutoCFDatetimeLocator(calendar=calendar) + units = locator.date_unit + + start_dt = cftime.datetime.strptime( + start_date_str, "%Y-%m-%d %H:%M:%S", calendar=calendar + ) + end_dt = cftime.datetime.strptime( + end_date_str, "%Y-%m-%d %H:%M:%S", calendar=calendar + ) + + start_num = cftime.date2num(start_dt, units, calendar=calendar) + end_num = cftime.date2num(end_dt, units, calendar=calendar) + + res, _ = locator.compute_resolution(start_num, end_num, start_dt, end_dt) + assert res == resolution + + ticks = locator.tick_values(start_num, end_num) + assert len(ticks) > 0 + + +def test_autocftime_locator_safe_helpers(): + from ultraplot.ticker import AutoCFDatetimeLocator + import cftime + + # Test _safe_num2date with invalid value + locator_gregorian = AutoCFDatetimeLocator(calendar="gregorian") + with pytest.raises(OverflowError): + locator_gregorian._safe_num2date(1e30) + + # Test _safe_create_datetime with invalid date + locator_noleap = AutoCFDatetimeLocator(calendar="noleap") + assert locator_noleap._safe_create_datetime(2001, 2, 29) is None + + +def test_autocftime_locator_safe_daily_locator(): + from ultraplot.ticker import AutoCFDatetimeLocator + + locator = AutoCFDatetimeLocator() + assert locator._safe_daily_locator(0, 10) is not None + + +def test_latitude_locator(): + from ultraplot.ticker import LatitudeLocator + import numpy as np + + locator = LatitudeLocator() + ticks = np.array(locator.tick_values(-100, 100)) + assert np.all(ticks >= -90) + assert np.all(ticks <= 90) + + +def test_cftime_converter(): + from ultraplot.ticker import CFTimeConverter, cftime + from ultraplot.config import rc + import numpy as np + + converter = CFTimeConverter() + + # test default_units + assert converter.default_units(np.array([1, 2]), None) is None + assert converter.default_units([], None) is None + assert converter.default_units([1, 2], None) is None + with pytest.raises(ValueError): + converter.default_units( + [ + cftime.datetime(2000, 1, 1, calendar="gregorian"), + cftime.datetime(2000, 1, 1, calendar="noleap"), + ], + None, + ) + + with pytest.raises(ValueError): + converter.default_units(cftime.datetime(2000, 1, 1, calendar=""), None) + + # test convert + assert converter.convert(1, None, None) == 1 + assert np.array_equal( + converter.convert(np.array([1, 2]), None, None), np.array([1, 2]) + ) + assert converter.convert([1, 2], None, None) == [1, 2] + assert converter.convert([], None, None) == [] + + # test axisinfo with no unit + with rc.context({"cftime.time_unit": "days since 2000-01-01"}): + info = converter.axisinfo(None, None) + assert isinstance(info.majloc, uplt.ticker.AutoCFDatetimeLocator) diff --git a/ultraplot/ticker.py b/ultraplot/ticker.py index bc3f7b9a9..7e22010db 100644 --- a/ultraplot/ticker.py +++ b/ultraplot/ticker.py @@ -7,9 +7,18 @@ from fractions import Fraction import matplotlib.axis as maxis +import matplotlib.dates as mdates import matplotlib.ticker as mticker +import matplotlib.transforms as mtransforms +import matplotlib.units as munits +from datetime import datetime, timedelta import numpy as np +try: + import cftime +except ModuleNotFoundError: + cftime = None + from .config import rc from .internals import ic # noqa: F401 from .internals import _not_none, context, docstring @@ -37,6 +46,9 @@ "SciFormatter", "SigFigFormatter", "FracFormatter", + "CFDatetimeFormatter", + "AutoCFDatetimeFormatter", + "AutoCFDatetimeLocator", "DegreeFormatter", "LongitudeFormatter", "LatitudeFormatter", @@ -414,19 +426,24 @@ def __call__(self, x, pos=None): x, tail = self._neg_pos_format(x, self._negpos, wraprange=self._wraprange) # Default string formatting + print(x, pos) string = super().__call__(x, pos) + print(string) # Fix issue where non-zero string is formatted as zero string = self._fix_small_number(x, string) + print(string) # Custom string formatting string = self._minus_format(string) + print(string) if self._zerotrim: string = self._trim_trailing_zeros(string, self._get_decimal_point()) # Prefix and suffix string = self._add_prefix_suffix(string, self._prefix, self._suffix) string = string + tail # add negative-positive indicator + print(string) return string def get_offset(self): @@ -822,6 +839,487 @@ def __call__(self, x, pos=None): # noqa: U100 return string +class CFDatetimeFormatter(mticker.Formatter): + """ + Format dates using `cftime.datetime.strftime` format strings. + """ + + def __init__(self, fmt, calendar="standard", units="days since 2000-01-01"): + """ + Parameters + ---------- + fmt : str + The `strftime` format string. + calendar : str, default: 'standard' + The calendar for interpreting numeric tick values. + units : str, default: 'days since 2000-01-01' + The time units for interpreting numeric tick values. + """ + if cftime is None: + raise ModuleNotFoundError("cftime is required for CFDatetimeFormatter.") + self._format = fmt + self._calendar = calendar + self._units = units + + def __call__(self, x, pos=None): + if isinstance(x, cftime.datetime): + dt = x + else: + dt = cftime.num2date(x, self._units, calendar=self._calendar) + return dt.strftime(self._format) + + +class AutoCFDatetimeFormatter(mticker.Formatter): + """Automatic formatter for `cftime.datetime` data.""" + + def __init__(self, locator, calendar, time_units=None): + self.locator = locator + self.calendar = calendar + self.time_units = time_units or rc["cftime.time_unit"] + + def pick_format(self, resolution): + return rc["cftime.time_resolution_format"][resolution] + + def __call__(self, x, pos=0): + format_string = self.pick_format(self.locator.resolution) + if isinstance(x, cftime.datetime): + dt = x + else: + dt = cftime.num2date(x, self.time_units, calendar=self.calendar) + return dt.strftime(format_string) + + +class AutoCFDatetimeLocator(mticker.Locator): + """Determines tick locations when plotting `cftime.datetime` data.""" + + if cftime: + real_world_calendars = cftime._cftime._calendars + else: + real_world_calendars = () + + def __init__(self, maxticks=None, calendar="standard", date_unit=None, minticks=3): + super().__init__() + self.minticks = minticks + # These are thresholds for *determining resolution*, NOT directly for MaxNLocator tick count + self.maxticks = { + "YEARLY": 1, + "MONTHLY": 12, + "DAILY": 8, + "HOURLY": 11, + "MINUTELY": 1, + "SECONDLY": 11, # Added for completeness, though not a resolution threshold + } + if maxticks is not None: + if isinstance(maxticks, dict): + self.maxticks.update(maxticks) + else: + # If a single value is provided for maxticks, apply it to all resolution *thresholds* + self.maxticks = dict.fromkeys(self.maxticks.keys(), maxticks) + + self.calendar = calendar + self.date_unit = date_unit or rc["cftime.time_unit"] + if not self.date_unit.lower().startswith("days since"): + emsg = "The date unit must be days since for a NetCDF time locator." + raise ValueError(emsg) + self.resolution = rc["cftime.resolution"] + self._cached_resolution = {} + + self._max_display_ticks = rc["cftime.max_display_ticks"] + + def set_params(self, maxticks=None, minticks=None, max_display_ticks=None): + """Set the parameters for the locator.""" + if maxticks is not None: + if isinstance(maxticks, dict): + self.maxticks.update(maxticks) + else: + self.maxticks = dict.fromkeys(self.maxticks.keys(), maxticks) + if minticks is not None: + self.minticks = minticks + if max_display_ticks is not None: + self._max_display_ticks = max_display_ticks + + def compute_resolution(self, num1, num2, date1, date2): + """Returns the resolution of the dates. + Also updates self.calendar from date1 for consistency. + """ + if isinstance(date1, cftime.datetime): + self.calendar = date1.calendar + # Assuming self.date_unit (e.g., "days since 0001-01-01") is a universal epoch. + + num_days = float(np.abs(num1 - num2)) + + num_years = num_days / 365.0 + num_months = num_days / 30.0 + num_hours = num_days * 24.0 + num_minutes = num_days * 60.0 * 24.0 + + if num_years > self.maxticks["YEARLY"]: + resolution = "YEARLY" + n = abs(date1.year - date2.year) + elif num_months > self.maxticks["MONTHLY"]: + resolution = "MONTHLY" + n = abs((date2.year - date1.year) * 12 + (date2.month - date1.month)) + elif num_days > self.maxticks["DAILY"]: + resolution = "DAILY" + n = num_days + elif num_hours > self.maxticks["HOURLY"]: + resolution = "HOURLY" + n = num_hours + elif num_minutes > self.maxticks["MINUTELY"]: + resolution = "MINUTELY" + n = num_minutes + else: + resolution = "SECONDLY" + n = num_days * 24 * 3600 + + self.resolution = resolution + return resolution, n + + def __call__(self): + vmin, vmax = self.axis.get_view_interval() + return self.tick_values(vmin, vmax) + + def tick_values(self, vmin, vmax): + vmin, vmax = mtransforms.nonsingular(vmin, vmax, expander=1e-7, tiny=1e-13) + result = self._safe_num2date(vmin, vmax) + lower, upper = result if result else (None, None) + if lower is None or upper is None: + return np.array([]) # Return empty array if conversion fails + + resolution, _ = self.compute_resolution(vmin, vmax, lower, upper) + + ticks_cftime = [] + + if resolution == "YEARLY": + years_start = lower.year + years_end = upper.year + # Use MaxNLocator to find "nice" years to tick + year_candidates = mticker.MaxNLocator( + self._max_display_ticks, integer=True, prune="both" + ).tick_values(years_start, years_end) + for year in year_candidates: + if ( + year == 0 and self.calendar in self.real_world_calendars + ): # Skip year 0 for real-world calendars + continue + try: + dt = cftime.datetime(int(year), 1, 1, calendar=self.calendar) + ticks_cftime.append(dt) + except ( + ValueError + ): # Catch potential errors for invalid dates in specific calendars + pass + + elif resolution == "MONTHLY": + # Generate all first-of-month dates between lower and upper bounds + all_months_cftime = [] + current_date = self._safe_create_datetime(lower.year, lower.month, 1) + end_date_limit = self._safe_create_datetime(upper.year, upper.month, 1) + if current_date is None or end_date_limit is None: + return np.array([]) # Return empty array if date creation fails + + while current_date <= end_date_limit: + all_months_cftime.append(current_date) + # Increment month + if current_date.month == 12: + current_date = cftime.datetime( + current_date.year + 1, 1, 1, calendar=self.calendar + ) + else: + current_date = cftime.datetime( + current_date.year, + current_date.month + 1, + 1, + calendar=self.calendar, + ) + + # Select a reasonable number of these using a stride + num_all_months = len(all_months_cftime) + if num_all_months == 0: + pass + elif num_all_months <= self._max_display_ticks: + ticks_cftime = all_months_cftime + else: + stride = max(1, num_all_months // self._max_display_ticks) + ticks_cftime = all_months_cftime[::stride] + # Ensure first and last are included if not already + if ( + all_months_cftime + and ticks_cftime + and ticks_cftime[0] != all_months_cftime[0] + ): + ticks_cftime.insert(0, all_months_cftime[0]) + if ( + all_months_cftime + and ticks_cftime + and ticks_cftime[-1] != all_months_cftime[-1] + ): + ticks_cftime.append(all_months_cftime[-1]) + + elif resolution == "DAILY": + # MaxNLocator_days works directly on the numeric values (days since epoch) + try: + days_numeric = self._safe_daily_locator(vmin, vmax) + if days_numeric is None: + return np.array([]) # Return empty array if locator fails + for dt_num in days_numeric: + tick = self._safe_num2date(dt_num) + if tick is not None: + ticks_cftime.append(tick) + except ValueError: + return np.array([]) # Return empty array if locator fails + + elif resolution == "HOURLY": + current_dt_cftime = self._safe_create_datetime( + lower.year, lower.month, lower.day, lower.hour + ) + if current_dt_cftime is None: + return np.array([]) # Return empty array if date creation fails + total_hours = (vmax - vmin) * 24.0 + hour_step_candidates = [1, 2, 3, 4, 6, 8, 12] + hour_step = 1 + + if total_hours > 0: + hour_step = hour_step_candidates[ + np.argmin( + np.abs( + np.array(hour_step_candidates) + - total_hours / self._max_display_ticks + ) + ) + ] + if hour_step == 0: + hour_step = 1 + + while ( + cftime.date2num(current_dt_cftime, self.date_unit, self.calendar) > vmin + ): + current_dt_cftime += timedelta(hours=-hour_step) + if ( + cftime.date2num(current_dt_cftime, self.date_unit, self.calendar) + < vmin - (vmax - vmin) * 2 + ): + current_dt_cftime = cftime.datetime( + lower.year, + lower.month, + lower.day, + lower.hour, + 0, + 0, + calendar=self.calendar, + ) # Reset + break # Break if we overshot significantly + + while ( + cftime.date2num(current_dt_cftime, self.date_unit, self.calendar) + <= vmax + (vmax - vmin) * 0.01 + ): # Small buffer to include last tick + ticks_cftime.append(current_dt_cftime) + current_dt_cftime += timedelta(hours=hour_step) + if ( + len(ticks_cftime) > 2 * self._max_display_ticks and hour_step != 0 + ): # Safety break to prevent too many ticks + break + + elif resolution == "MINUTELY": + try: + current_dt_cftime = cftime.datetime( + lower.year, + lower.month, + lower.day, + lower.hour, + lower.minute, + 0, + calendar=self.calendar, + ) + except ValueError: + return np.array([]) # Return empty array if date creation fails + total_minutes = (vmax - vmin) * 24.0 * 60.0 + minute_step_candidates = [1, 2, 5, 10, 15, 20, 30] + minute_step = 1 + + if total_minutes > 0: + minute_step = minute_step_candidates[ + np.argmin( + np.abs( + np.array(minute_step_candidates) + - total_minutes / self._max_display_ticks + ) + ) + ] + if minute_step == 0: + minute_step = 1 + + while ( + cftime.date2num(current_dt_cftime, self.date_unit, self.calendar) > vmin + ): + current_dt_cftime += timedelta(minutes=-minute_step) + if ( + cftime.date2num(current_dt_cftime, self.date_unit, self.calendar) + < vmin - (vmax - vmin) * 2 + ): + try: + current_dt_cftime = cftime.datetime( + lower.year, + lower.month, + lower.day, + lower.hour, + lower.minute, + 0, + calendar=self.calendar, + ) + except ValueError: + return np.array([]) # Return empty array if date creation fails + break + + while ( + cftime.date2num(current_dt_cftime, self.date_unit, self.calendar) + <= vmax + (vmax - vmin) * 0.01 + ): + ticks_cftime.append(current_dt_cftime) + current_dt_cftime += timedelta(minutes=minute_step) + if len(ticks_cftime) > 2 * self._max_display_ticks and minute_step != 0: + break + + elif resolution == "SECONDLY": + current_dt_cftime = self._safe_create_datetime( + lower.year, + lower.month, + lower.day, + lower.hour, + lower.minute, + lower.second, + ) + if current_dt_cftime is None: + return np.array([]) # Return empty array if date creation fails + total_seconds = (vmax - vmin) * 24.0 * 60.0 * 60.0 + second_step_candidates = [1, 2, 5, 10, 15, 20, 30] + second_step = 1 + + if total_seconds > 0: + second_step = second_step_candidates[ + np.argmin( + np.abs( + np.array(second_step_candidates) + - total_seconds / self._max_display_ticks + ) + ) + ] + if second_step == 0: + second_step = 1 + + while ( + cftime.date2num(current_dt_cftime, self.date_unit, self.calendar) > vmin + ): + current_dt_cftime += timedelta(seconds=-second_step) + if ( + cftime.date2num(current_dt_cftime, self.date_unit, self.calendar) + < vmin - (vmax - vmin) * 2 + ): + current_dt_cftime = cftime.datetime( + lower.year, + lower.month, + lower.day, + lower.hour, + lower.minute, + lower.second, + calendar=self.calendar, + ) + break + + while ( + cftime.date2num(current_dt_cftime, self.date_unit, self.calendar) + <= vmax + (vmax - vmin) * 0.01 + ): + ticks_cftime.append(current_dt_cftime) + current_dt_cftime += timedelta( + seconds=second_step + ) # Uses standard library timedelta + if len(ticks_cftime) > 2 * self._max_display_ticks and second_step != 0: + break + else: + emsg = f"Resolution {resolution} not implemented yet." + raise ValueError(emsg) + + if self.calendar in self.real_world_calendars: + # Filters out year 0 for calendars where it's not applicable + ticks_cftime = [t for t in ticks_cftime if t.year != 0] + + # Convert the generated cftime.datetime objects back to numeric values + ticks_numeric = [] + for t in ticks_cftime: + try: + num_val = cftime.date2num(t, self.date_unit, calendar=self.calendar) + if vmin <= num_val <= vmax: # Final filter for actual view interval + ticks_numeric.append(num_val) + except ValueError: # Handle potential issues with invalid cftime dates + pass + + # Fallback: if no ticks are found (e.g., very short range or specific calendar issues), + # ensure at least two end points if vmin <= vmax + if not ticks_numeric and vmin <= vmax: + return np.array([vmin, vmax]) + elif not ticks_numeric: # If vmin > vmax or other edge cases + return np.array([]) + + return np.unique(np.array(ticks_numeric)) + + def _safe_num2date(self, value, vmax=None): + """ + Safely converts numeric values to cftime.datetime objects. + + If a single value is provided, it converts and returns a single datetime object. + If both value (vmin) and vmax are provided, it converts and returns a tuple of + datetime objects (lower, upper). + + This helper is used to handle cases where the conversion might fail + due to invalid inputs or calendar-specific constraints. If the conversion + fails, it returns None or a tuple of Nones. + """ + try: + if vmax is not None: + lower = cftime.num2date(value, self.date_unit, calendar=self.calendar) + upper = cftime.num2date(vmax, self.date_unit, calendar=self.calendar) + return lower, upper + else: + return cftime.num2date(value, self.date_unit, calendar=self.calendar) + except ValueError: + return (None, None) if vmax is not None else None + + def _safe_create_datetime(self, year, month=1, day=1, hour=0, minute=0, second=0): + """ + Safely creates a cftime.datetime object with the given date and time components. + + This helper is used to handle cases where creating a datetime object might fail + due to invalid inputs (e.g., invalid dates in specific calendars). If the creation + fails, it returns None. + """ + try: + return cftime.datetime( + year, month, day, hour, minute, second, calendar=self.calendar + ) + except ValueError: + return None + + def _safe_daily_locator(self, vmin, vmax): + """ + Safely generates daily tick values using MaxNLocator. + + This helper is used to handle cases where the locator might fail + due to invalid input ranges or other issues. If the locator fails, + it returns None. + """ + try: + return mticker.MaxNLocator( + self._max_display_ticks, + integer=True, + steps=[1, 2, 4, 7, 10], + prune="both", + ).tick_values(vmin, vmax) + except ValueError: + return None + + class _CartopyFormatter(object): """ Mixin class for cartopy formatters. @@ -897,3 +1395,125 @@ def __init__(self, *args, **kwargs): %(ticker.dms)s """ super().__init__(*args, **kwargs) + + +class CFTimeConverter(mdates.DateConverter): + """ + Converter for cftime.datetime data. + """ + + @staticmethod + def axisinfo(unit, axis): + """Returns the :class:`~matplotlib.units.AxisInfo` for *unit*.""" + if cftime is None: + raise ModuleNotFoundError("cftime is required for CFTimeConverter.") + + if unit is None: + calendar, date_unit, date_type = ( + "standard", + rc["cftime.time_unit"], + getattr(cftime, "DatetimeProlepticGregorian", None), + ) + else: + calendar, date_unit, date_type = unit + + majloc = AutoCFDatetimeLocator(calendar=calendar, date_unit=date_unit) + majfmt = AutoCFDatetimeFormatter( + majloc, calendar=calendar, time_units=date_unit + ) + + year = datetime.now().year + if date_type is not None: + datemin = date_type(year - 1, 1, 1) + datemax = date_type(year, 1, 1) + else: + # Fallback if date_type is None + datemin = cftime.datetime(year - 1, 1, 1, calendar="standard") + datemax = cftime.datetime(year, 1, 1, calendar="standard") + + return munits.AxisInfo( + majloc=majloc, + majfmt=majfmt, + label="", + default_limits=(datemin, datemax), + ) + + @classmethod + def default_units(cls, x, axis): + """Computes some units for the given data point.""" + if isinstance(x, np.ndarray) and x.dtype != object: + return None # It's already numeric + + if hasattr(x, "__iter__") and not isinstance(x, str): + if isinstance(x, np.ndarray): + x = x.reshape(-1) + + first_value = next( + (v for v in x if v is not None and v is not np.ma.masked), None + ) + if first_value is None: + return None + + if not isinstance(first_value, cftime.datetime): + return None + + # Check all calendars are the same + calendar = first_value.calendar + if any( + getattr(v, "calendar", None) != calendar + for v in x + if v is not None and v is not np.ma.masked + ): + raise ValueError("Calendar units are not all equal.") + + date_type = type(first_value) + else: + # scalar + if not isinstance(x, cftime.datetime): + return None + calendar = x.calendar + date_type = type(x) + + if not calendar: + raise ValueError( + "A calendar must be defined to plot dates using a cftime axis." + ) + + return calendar, rc["cftime.time_unit"], date_type + + @classmethod + def convert(cls, value, unit, axis): + """Converts value with :py:func:`cftime.date2num`.""" + if cftime is None: + raise ModuleNotFoundError("cftime is required for CFTimeConverter.") + + if isinstance(value, (int, float, np.number)): + return value + + if isinstance(value, np.ndarray) and value.dtype != object: + return value + + # Get calendar from unit, or from data + if unit: + calendar, date_unit, _ = unit + else: + if hasattr(value, "__iter__") and not isinstance(value, str): + first_value = next( + (v for v in value if v is not None and v is not np.ma.masked), None + ) + else: + first_value = value + + if first_value is None or not isinstance(first_value, cftime.datetime): + return value # Cannot convert + + calendar = first_value.calendar + date_unit = rc["cftime.time_unit"] + + result = cftime.date2num(value, date_unit, calendar=calendar) + return result + + +if cftime is not None: + if cftime.datetime not in munits.registry: + munits.registry[cftime.datetime] = CFTimeConverter()