From fee9e86dbde95a3efcb8e307ec2a8017d7341c7e Mon Sep 17 00:00:00 2001 From: wpbonelli Date: Tue, 28 Oct 2025 16:20:52 -0400 Subject: [PATCH] adapt time utilities from imod-python --- docs/examples/quickstart.py | 4 +- flopy4/mf6/component.py | 2 +- flopy4/mf6/gwf/chd.py | 2 +- flopy4/mf6/gwf/drn.py | 2 +- flopy4/mf6/gwf/rch.py | 2 +- flopy4/mf6/gwf/wel.py | 2 +- flopy4/mf6/simulation.py | 10 +- flopy4/mf6/tdis.py | 44 ++++- flopy4/mf6/utils/cbc_reader.py | 2 +- flopy4/mf6/utils/{grid_utils.py => grid.py} | 0 flopy4/mf6/utils/heads_reader.py | 2 +- flopy4/mf6/utils/time.py | 99 +++++++++++ test/test_codec.py | 10 +- test/test_component.py | 34 ++-- test/test_time_utils.py | 180 ++++++++++++++++++++ 15 files changed, 357 insertions(+), 38 deletions(-) rename flopy4/mf6/utils/{grid_utils.py => grid.py} (100%) create mode 100644 flopy4/mf6/utils/time.py create mode 100644 test/test_time_utils.py diff --git a/docs/examples/quickstart.py b/docs/examples/quickstart.py index 5740a481..ad2e5228 100644 --- a/docs/examples/quickstart.py +++ b/docs/examples/quickstart.py @@ -2,16 +2,16 @@ import matplotlib.pyplot as plt import numpy as np -from flopy.discretization.modeltime import ModelTime from flopy.discretization.structuredgrid import StructuredGrid from flopy4.mf6.gwf import Chd, Gwf, Ic, Npf, Oc from flopy4.mf6.ims import Ims from flopy4.mf6.simulation import Simulation +from flopy4.mf6.utils.time import Time name = "quickstart" workspace = Path(__file__).parent / name -time = ModelTime(perlen=[1.0], nstp=[1]) +time = Time(perlen=[1.0], nstp=[1]) grid = StructuredGrid( nlay=1, nrow=10, diff --git a/flopy4/mf6/component.py b/flopy4/mf6/component.py index 09ef9c1b..13c3cebc 100644 --- a/flopy4/mf6/component.py +++ b/flopy4/mf6/component.py @@ -11,7 +11,7 @@ from flopy4.mf6.constants import MF6 from flopy4.mf6.spec import field, fields_dict, to_field -from flopy4.mf6.utils.grid_utils import update_maxbound +from flopy4.mf6.utils.grid import update_maxbound from flopy4.uio import IO, Loader, Writer COMPONENTS = {} diff --git a/flopy4/mf6/gwf/chd.py b/flopy4/mf6/gwf/chd.py index 365163f0..4ca77608 100644 --- a/flopy4/mf6/gwf/chd.py +++ b/flopy4/mf6/gwf/chd.py @@ -10,7 +10,7 @@ from flopy4.mf6.converter import structure_array from flopy4.mf6.package import Package from flopy4.mf6.spec import array, field, path -from flopy4.mf6.utils.grid_utils import update_maxbound +from flopy4.mf6.utils.grid import update_maxbound from flopy4.utils import to_path diff --git a/flopy4/mf6/gwf/drn.py b/flopy4/mf6/gwf/drn.py index 222cd999..62b9eed0 100644 --- a/flopy4/mf6/gwf/drn.py +++ b/flopy4/mf6/gwf/drn.py @@ -10,7 +10,7 @@ from flopy4.mf6.converter import structure_array from flopy4.mf6.package import Package from flopy4.mf6.spec import array, field, path -from flopy4.mf6.utils.grid_utils import update_maxbound +from flopy4.mf6.utils.grid import update_maxbound from flopy4.utils import to_path diff --git a/flopy4/mf6/gwf/rch.py b/flopy4/mf6/gwf/rch.py index 03ed0e76..88379d80 100644 --- a/flopy4/mf6/gwf/rch.py +++ b/flopy4/mf6/gwf/rch.py @@ -10,7 +10,7 @@ from flopy4.mf6.converter import structure_array from flopy4.mf6.package import Package from flopy4.mf6.spec import array, field, path -from flopy4.mf6.utils.grid_utils import update_maxbound +from flopy4.mf6.utils.grid import update_maxbound from flopy4.utils import to_path diff --git a/flopy4/mf6/gwf/wel.py b/flopy4/mf6/gwf/wel.py index e2b97049..289c65e5 100644 --- a/flopy4/mf6/gwf/wel.py +++ b/flopy4/mf6/gwf/wel.py @@ -10,7 +10,7 @@ from flopy4.mf6.converter import structure_array from flopy4.mf6.package import Package from flopy4.mf6.spec import array, field, path -from flopy4.mf6.utils.grid_utils import update_maxbound +from flopy4.mf6.utils.grid import update_maxbound from flopy4.utils import to_path diff --git a/flopy4/mf6/simulation.py b/flopy4/mf6/simulation.py index 7c0df905..86a58498 100644 --- a/flopy4/mf6/simulation.py +++ b/flopy4/mf6/simulation.py @@ -1,7 +1,6 @@ from os import PathLike from warnings import warn -from flopy.discretization.modeltime import ModelTime from modflow_devtools.misc import cd, run_cmd from xattree import xattree @@ -11,14 +10,15 @@ from flopy4.mf6.solution import Solution from flopy4.mf6.spec import field from flopy4.mf6.tdis import Tdis +from flopy4.mf6.utils.time import Time def convert_time(value): - if isinstance(value, ModelTime): + if isinstance(value, Time): return Tdis.from_time(value) if isinstance(value, Tdis): return value - raise TypeError(f"Expected ModelTime or Tdis, got {type(value)}") + raise TypeError(f"Expected Time or Tdis, got {type(value)}") @xattree @@ -41,8 +41,8 @@ def __attrs_post_init__(self): self.filename = "mfsim.nam" @property - def time(self) -> ModelTime: - """Return the simulation time discretization.""" + def time(self) -> Time: + """Return a `Time` object describing the simulation's time discretization.""" return self.tdis.to_time() def run(self, exe: str | PathLike = "mf6", verbose: bool = False) -> None: diff --git a/flopy4/mf6/tdis.py b/flopy4/mf6/tdis.py index 2c3ac4de..5e362ca4 100644 --- a/flopy4/mf6/tdis.py +++ b/flopy4/mf6/tdis.py @@ -3,13 +3,13 @@ import numpy as np from attrs import Converter, define -from flopy.discretization.modeltime import ModelTime -from numpy.typing import NDArray +from numpy.typing import ArrayLike, NDArray from xattree import ROOT, xattree from flopy4.mf6.converter import structure_array from flopy4.mf6.package import Package from flopy4.mf6.spec import array, dim, field +from flopy4.mf6.utils.time import Time @xattree @@ -42,9 +42,9 @@ class PeriodData: converter=Converter(structure_array, takes_self=True, takes_field=True), ) - def to_time(self) -> ModelTime: - """Convert to a `ModelTime` object.""" - return ModelTime( + def to_time(self) -> Time: + """Convert to a `Time` object.""" + return Time( nper=self.nper, time_units=self.time_units, start_date_time=self.start_date_time, @@ -54,8 +54,8 @@ def to_time(self) -> ModelTime: ) @classmethod - def from_time(cls, time: ModelTime) -> "Tdis": - """Create a time discretization from a `ModelTime`.""" + def from_time(cls, time: Time) -> "Tdis": + """Create a time discretization from a `Time` object.""" return cls( nper=time.nper, time_units=None if time.time_units in [None, "unknown"] else time.time_units, @@ -64,3 +64,33 @@ def from_time(cls, time: ModelTime) -> "Tdis": nstp=time.nstp, tsmult=time.tsmult, ) + + @classmethod + def from_timestamps( + cls, + timestamps: ArrayLike, + nstp: Optional[ArrayLike] = None, + tsmult: Optional[ArrayLike] = None, + ) -> "Tdis": + """ + Create a time discretization from timestamps. + + Parameters + ---------- + timestamps : sequence of datetime-likes + Stress period start times + nstp : int or sequence of int, optional + Number of timesteps per stress period. If scalar, applied to all periods. + If None, defaults to 1 for all periods. + tsmult : float or sequence of float, optional + Timestep multiplier per stress period. If scalar, applied to all periods. + If None, defaults to 1.0 for all periods. + + Returns + ------- + Tdis + Time discretization object + """ + + time = Time.from_timestamps(timestamps, nstp=nstp, tsmult=tsmult) + return cls.from_time(time) diff --git a/flopy4/mf6/utils/cbc_reader.py b/flopy4/mf6/utils/cbc_reader.py index 14fc7b9a..aac33acb 100644 --- a/flopy4/mf6/utils/cbc_reader.py +++ b/flopy4/mf6/utils/cbc_reader.py @@ -13,7 +13,7 @@ from flopy.discretization import StructuredGrid from flopy4.adapters import StructuredGridWrapper -from flopy4.mf6.utils.grid_utils import get_coords +from flopy4.mf6.utils.grid import get_coords @define diff --git a/flopy4/mf6/utils/grid_utils.py b/flopy4/mf6/utils/grid.py similarity index 100% rename from flopy4/mf6/utils/grid_utils.py rename to flopy4/mf6/utils/grid.py diff --git a/flopy4/mf6/utils/heads_reader.py b/flopy4/mf6/utils/heads_reader.py index c9274056..83712996 100644 --- a/flopy4/mf6/utils/heads_reader.py +++ b/flopy4/mf6/utils/heads_reader.py @@ -8,7 +8,7 @@ import xarray as xr from flopy.discretization import StructuredGrid -from .grid_utils import get_coords +from .grid import get_coords def open_hds( diff --git a/flopy4/mf6/utils/time.py b/flopy4/mf6/utils/time.py new file mode 100644 index 00000000..10588237 --- /dev/null +++ b/flopy4/mf6/utils/time.py @@ -0,0 +1,99 @@ +import numpy as np +import pandas as pd +from flopy.discretization.modeltime import ModelTime +from numpy.typing import ArrayLike + + +class Time(ModelTime): + """Extend flopy3's ModelTime""" + + @classmethod + def from_timestamps( + cls, + timestamps: ArrayLike, + nstp: ArrayLike | None = None, + tsmult: ArrayLike | None = None, + ) -> "Time": + """ + Create a time discretization from timestamps. + + Parameters + ---------- + timestamps : sequence of datetime-likes + Stress period start times + nstp : int or sequence of int, optional + Number of timesteps per stress period. If scalar, applied to all periods. + If None, defaults to 1 for all periods. + tsmult : float or sequence of float, optional + Timestep multiplier per stress period. If scalar, applied to all periods. + If None, defaults to 1.0 for all periods. + + Returns + ------- + Time + Time discretization object + """ + + unique_times = np.unique(np.hstack(pd.to_datetime(timestamps))) # np.unique also sorts + if len(unique_times) < 2: + raise ValueError("Need at least two timestamps to create time discretization") + + timedeltas = to_timedeltas(unique_times) + perlen = np.array( + [delta.total_seconds() / 86400.0 for delta in timedeltas], dtype=np.float64 + ) + nper = len(perlen) + + if nstp is None: + nstp_array = np.ones(nper, dtype=np.int64) + elif np.isscalar(nstp): + nstp_array = np.full(nper, nstp, dtype=np.int64) + else: + nstp_array = np.array(nstp, dtype=np.int64) + if len(nstp_array) != nper: + raise ValueError( + f"nstp length ({len(nstp_array)}) must match number of periods ({nper})" + ) + + if tsmult is None: + tsmult_array = np.ones(nper, dtype=np.float64) + elif np.isscalar(tsmult): + tsmult_array = np.full(nper, tsmult, dtype=np.float64) + else: + tsmult_array = np.array(tsmult, dtype=np.float64) + if len(tsmult_array) != nper: + raise ValueError( + f"tsmult length ({len(tsmult_array)}) must match number of periods ({nper})" + ) + + return cls( + perlen=perlen, + nstp=nstp_array, + tsmult=tsmult_array, + time_units="days", + start_datetime=pd.to_datetime(unique_times[0]).to_pydatetime(), + ) + + +def to_timedeltas( + timestamps: ArrayLike, +) -> list[pd.Timedelta]: + """ + Converts a sequence of datetime-like objects to a list of timedelta + objects representing the durations between consecutive time points. + + Parameters + ---------- + timestamps : sequence of datetime-likes + A sequence of datetime-like objects representing time points. + + Returns + ------- + timedeltas : list of pd.Timedelta + A list of durations between consecutive time points. + """ + if len(np.atleast_1d(timestamps)) < 2: + return [] + + timestamps = pd.to_datetime(timestamps) + return [pd.Timedelta(end - start) for start, end in zip(timestamps[:-1], timestamps[1:])] # type: ignore diff --git a/test/test_codec.py b/test/test_codec.py index cb641240..cdb06222 100644 --- a/test/test_codec.py +++ b/test/test_codec.py @@ -118,11 +118,10 @@ def test_dumps_dis(): def test_dumps_tdis(): - from flopy.discretization.modeltime import ModelTime - from flopy4.mf6.tdis import Tdis + from flopy4.mf6.utils.time import Time - tdis = Tdis.from_time(ModelTime(perlen=[1.0, 2.0], nstp=[1, 2])) + tdis = Tdis.from_time(Time(perlen=[1.0, 2.0], nstp=[1, 2])) tdis.time_units = "days" dumped = dumps(COMPONENT_CONVERTER.unstructure(tdis)) @@ -404,11 +403,10 @@ def test_dumps_gwf(): def test_dumps_simulation(): - from flopy.discretization.modeltime import ModelTime - from flopy4.mf6.gwf import Dis, Gwf, Ic, Npf, Oc from flopy4.mf6.simulation import Simulation from flopy4.mf6.tdis import Tdis + from flopy4.mf6.utils.time import Time # Create model components dis = Dis(nlay=1, nrow=5, ncol=5, delr=100.0, delc=100.0) @@ -427,7 +425,7 @@ def test_dumps_simulation(): ) # Create time discretization - time = ModelTime(perlen=[1.0], nstp=[1]) + time = Time(perlen=[1.0], nstp=[1]) tdis = Tdis.from_time(time) # Create simulation diff --git a/test/test_component.py b/test/test_component.py index 0be0d31a..e05ee69b 100644 --- a/test/test_component.py +++ b/test/test_component.py @@ -1,9 +1,9 @@ from pathlib import Path import numpy as np +import pandas as pd import pytest from flopy.discretization import StructuredGrid -from flopy.discretization.modeltime import ModelTime from xarray import DataTree from flopy4.mf6.component import COMPONENTS @@ -12,6 +12,7 @@ from flopy4.mf6.ims import Ims from flopy4.mf6.simulation import Simulation from flopy4.mf6.tdis import Tdis +from flopy4.mf6.utils.time import Time def test_registry(): @@ -28,7 +29,7 @@ def test_init_empty_sim(): def test_init_gwf_explicit_dims(): - time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0]) + time = Time(perlen=[1.0], nstp=[1], tsmult=[1.0]) grid = StructuredGrid(nlay=1, nrow=2, ncol=2) dims = { "nper": time.nper, @@ -67,7 +68,7 @@ def test_init_gwf_explicit_dims(): @pytest.mark.skip(reason="TODO") def test_init_gwf_from_grid_context(): - time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0]) + time = Time(perlen=[1.0], nstp=[1], tsmult=[1.0]) grid = StructuredGrid(nlay=1, nrow=2, ncol=2) # TODO maybe a dumb idea, but we could put the # time and grid in a context manager? then you @@ -154,7 +155,7 @@ def test_init_gwf_top_down_misaligned(): def test_init_sim_explicit_dims(): - time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0]) + time = Time(perlen=[1.0], nstp=[1], tsmult=[1.0]) grid = StructuredGrid(nlay=1, nrow=10, ncol=10) dims = { "nlay": grid.nlay, @@ -211,7 +212,7 @@ def test_init_sim_explicit_dims(): def test_init_big_sim(): # if size over threshold, arrays should be sparse - time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0]) + time = Time(perlen=[1.0], nstp=[1], tsmult=[1.0]) grid = StructuredGrid(nlay=1, nrow=100, ncol=100) sim = Simulation(tdis=time) gwf = Gwf(parent=sim, dis=grid) @@ -280,7 +281,7 @@ def test_ims_dfn(): def test_gwf_chd01(function_tmpdir): sim_name = "chd01" gwf_name = "gwf_chd01" - time = ModelTime(perlen=[5.0], nstp=[1], tsmult=[1.0], time_units="days") + time = Time(perlen=[5.0], nstp=[1], tsmult=[1.0], time_units="days") ims = Ims( filename="sln1.ims", @@ -363,7 +364,7 @@ def test_gwf_chd01(function_tmpdir): def test_quickstart(function_tmpdir): sim_name = "quickstart" gwf_name = "mymodel" - time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0]) + time = Time(perlen=[1.0], nstp=[1], tsmult=[1.0]) ims = Ims(models=[gwf_name]) dis = Dis( nlay=1, @@ -397,7 +398,7 @@ def test_quickstart(function_tmpdir): def test_write_ascii(function_tmpdir): sim_name = "sim" gwf_name = "gwf" - time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0]) + time = Time(perlen=[1.0], nstp=[1], tsmult=[1.0]) ims = Ims(models=[gwf_name]) dis = Dis( nlay=1, @@ -435,7 +436,7 @@ def test_write_ascii(function_tmpdir): def test_to_dict_fields(): - time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0]) + time = Time(perlen=[1.0], nstp=[1], tsmult=[1.0]) grid = StructuredGrid(nlay=1, nrow=10, ncol=10) dims = { "nlay": grid.nlay, @@ -463,7 +464,7 @@ def test_to_dict_fields(): def test_to_dict_blocks(): - time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0]) + time = Time(perlen=[1.0], nstp=[1], tsmult=[1.0]) grid = StructuredGrid(nlay=1, nrow=10, ncol=10) dims = { "nlay": grid.nlay, @@ -515,7 +516,7 @@ def test_to_dict_on_component(): def test_to_dict_on_context(): - time = ModelTime(perlen=[1.0], nstp=[1], tsmult=[1.0]) + time = Time(perlen=[1.0], nstp=[1], tsmult=[1.0]) ims = Ims(models=["gwf"]) sim = Simulation(tdis=time, solutions={"ims": ims}) @@ -541,3 +542,14 @@ def test_to_dict_with_strict_excludes_fields_without_block_metadata(): assert "nrow" in result assert "ncol" in result assert "nodes" not in result + + +def test_tdis_from_timestamps(): + tdis = Tdis.from_timestamps(["2020-01-01", "2020-01-05", "2020-01-15"], nstp=5, tsmult=1.2) + + assert tdis.nper == 2 + assert tdis.time_units == "days" + assert tdis.start_date_time == pd.Timestamp("2020-01-01").to_pydatetime() + np.testing.assert_array_equal(tdis.perlen, [4.0, 10.0]) + np.testing.assert_array_equal(tdis.nstp, [5, 5]) + np.testing.assert_array_equal(tdis.tsmult, [1.2, 1.2]) diff --git a/test/test_time_utils.py b/test/test_time_utils.py new file mode 100644 index 00000000..b646fb50 --- /dev/null +++ b/test/test_time_utils.py @@ -0,0 +1,180 @@ +from datetime import timedelta + +import numpy as np +import pandas as pd +import pytest + +from flopy4.mf6.utils.time import Time, to_timedeltas + + +def test_to_deltas_regular_intervals(): + # daily + times = pd.to_datetime(["2020-01-01", "2020-01-02", "2020-01-03", "2020-01-04"]) + durations = to_timedeltas(times.values) + assert len(durations) == 3 + assert durations == [timedelta(days=1), timedelta(days=1), timedelta(days=1)] + + # hourly + times = pd.to_datetime( + ["2020-01-01 00:00:00", "2020-01-01 01:00:00", "2020-01-01 02:00:00", "2020-01-01 03:00:00"] + ) + durations = to_timedeltas(times.values) + assert len(durations) == 3 + assert durations == [pd.Timedelta(hours=1), pd.Timedelta(hours=1), pd.Timedelta(hours=1)] + + +def test_to_deltas_varying_intervals(): + times = pd.to_datetime(["2020-01-01", "2020-01-02", "2020-01-05", "2020-01-10"]) + durations = to_timedeltas(times.values) + assert len(durations) == 3 + assert durations == [pd.Timedelta(days=1), pd.Timedelta(days=3), pd.Timedelta(days=5)] + + times = pd.to_datetime( + [ + "2020-01-01 00:00:00", + "2020-01-01 12:30:45", + "2020-01-02 00:00:00", + ] + ) + durations = to_timedeltas(times.values) + assert len(durations) == 2 + expected_first = pd.Timedelta(hours=12, minutes=30, seconds=45) + expected_second = pd.Timedelta(hours=11, minutes=29, seconds=15) + assert durations == [expected_first, expected_second] + + +def test_to_deltas_cross_month_boundaries(): + times = pd.to_datetime( + [ + "2020-01-31", + "2020-02-01", + "2020-02-29", + "2020-03-01", + ] + ) + durations = to_timedeltas(times.values) + assert len(durations) == 3 + assert durations == [pd.Timedelta(days=1), pd.Timedelta(days=28), pd.Timedelta(days=1)] + + +def test_leap_year_vs_regular_year(): + # Leap year + times_leap = pd.to_datetime(["2020-02-01", "2020-03-01"]) + durations_leap = to_timedeltas(times_leap.values) + + # Regular year + times_regular = pd.to_datetime(["2021-02-01", "2021-03-01"]) + durations_regular = to_timedeltas(times_regular.values) + + assert durations_leap[0] == pd.Timedelta(days=29) + assert durations_regular[0] == pd.Timedelta(days=28) + + +def test_to_deltas_empty_with_less_than_2_items(): + times = pd.to_datetime(["2020-01-01"]) + assert not any(to_timedeltas(times.values)) + + times = pd.to_datetime([]) + assert not any(to_timedeltas(times.values)) + + +def test_time_from_timestamps_basic(): + times = pd.to_datetime(["2020-01-01", "2020-01-05", "2020-01-15", "2020-02-01"]) + time = Time.from_timestamps(times) + + assert time.nper == 3 + assert time.time_units == "days" + assert time.start_datetime == pd.Timestamp("2020-01-01").to_pydatetime() + np.testing.assert_array_equal(time.perlen, [4.0, 10.0, 17.0]) + np.testing.assert_array_equal(time.nstp, [1, 1, 1]) + np.testing.assert_array_equal(time.tsmult, [1.0, 1.0, 1.0]) + + +def test_time_from_timestamps_with_duplicates(): + times = np.array(["2020-01-01", "2020-01-05", "2020-01-01", "2020-01-10", "2020-01-05"]) + time = Time.from_timestamps(times) + + assert time.nper == 2 + np.testing.assert_array_equal(time.perlen, [4.0, 5.0]) + + +def test_time_from_timestamps_single_timestamp(): + times = pd.to_datetime(["2020-01-01"]) + + with pytest.raises(ValueError, match="at least two timestamps"): + Time.from_timestamps(times) + + +def test_time_from_timestamps_realistic_scenario(): + times = np.hstack( + [ + pd.date_range("2020-01-01", periods=12, freq="MS"), + pd.to_datetime(["2020-03-15", "2020-06-15", "2020-09-15"]), + ] + ) + time = Time.from_timestamps(times) + + assert time.nper == 14 # 12 months + 3 mid-month splits - 1 + assert time.time_units == "days" + assert time.start_datetime.year == 2020 + assert time.start_datetime.month == 1 + assert time.start_datetime.day == 1 + + +def test_time_from_timestamps_with_nstp_scalar(): + times = pd.to_datetime(["2020-01-01", "2020-01-05", "2020-01-15"]) + time = Time.from_timestamps(times, nstp=10) + + assert time.nper == 2 + np.testing.assert_array_equal(time.nstp, [10, 10]) + np.testing.assert_array_equal(time.tsmult, [1.0, 1.0]) + + +def test_time_from_timestamps_with_nstp_array(): + times = pd.to_datetime(["2020-01-01", "2020-01-05", "2020-01-15", "2020-02-01"]) + time = Time.from_timestamps(times, nstp=[5, 10, 20]) + + assert time.nper == 3 + np.testing.assert_array_equal(time.nstp, [5, 10, 20]) + np.testing.assert_array_equal(time.tsmult, [1.0, 1.0, 1.0]) + + +def test_time_from_timestamps_with_tsmult_scalar(): + times = pd.to_datetime(["2020-01-01", "2020-01-05", "2020-01-15"]) + time = Time.from_timestamps(times, nstp=10, tsmult=1.5) + + assert time.nper == 2 + np.testing.assert_array_equal(time.nstp, [10, 10]) + np.testing.assert_array_equal(time.tsmult, [1.5, 1.5]) + + +def test_time_from_timestamps_with_tsmult_array(): + times = pd.to_datetime(["2020-01-01", "2020-01-05", "2020-01-15", "2020-02-01"]) + time = Time.from_timestamps(times, tsmult=[1.0, 1.2, 1.5]) + + assert time.nper == 3 + np.testing.assert_array_equal(time.nstp, [1, 1, 1]) + np.testing.assert_array_equal(time.tsmult, [1.0, 1.2, 1.5]) + + +def test_time_from_timestamps_with_both_arrays(): + times = pd.to_datetime(["2020-01-01", "2020-01-05", "2020-01-15", "2020-02-01"]) + time = Time.from_timestamps(times, nstp=[5, 10, 20], tsmult=[1.0, 1.2, 1.5]) + + assert time.nper == 3 + np.testing.assert_array_equal(time.nstp, [5, 10, 20]) + np.testing.assert_array_equal(time.tsmult, [1.0, 1.2, 1.5]) + + +def test_time_from_timestamps_nstp_length_mismatch(): + times = pd.to_datetime(["2020-01-01", "2020-01-05", "2020-01-15"]) + + with pytest.raises(ValueError, match="nstp length"): + Time.from_timestamps(times, nstp=[5, 10, 20]) + + +def test_time_from_timestamps_tsmult_length_mismatch(): + timestamps = pd.to_datetime(["2020-01-01", "2020-01-05", "2020-01-15"]) + + with pytest.raises(ValueError, match="tsmult length"): + Time.from_timestamps(timestamps, tsmult=[1.0, 1.2, 1.5])