Skip to content

Commit 9b5b23b

Browse files
Speed up create_input when dates are already np.datetime64
1 parent a1f0057 commit 9b5b23b

File tree

1 file changed

+8
-7
lines changed

1 file changed

+8
-7
lines changed

pymsis/msis.py

Lines changed: 8 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -378,14 +378,15 @@ def create_input(
378378
"""
379379
# Turn everything into arrays
380380
dates_arr: npt.NDArray[np.datetime64] = np.atleast_1d(
381-
np.array(dates, dtype=np.datetime64)
381+
np.asarray(dates, dtype=np.datetime64)
382382
)
383-
dyear: npt.NDArray[np.datetime64] = (
384-
dates_arr.astype("datetime64[D]") - dates_arr.astype("datetime64[Y]")
385-
).astype(float) + 1 # DOY 1-366
386-
dseconds: npt.NDArray[np.datetime64] = (
387-
dates_arr.astype("datetime64[s]") - dates_arr.astype("datetime64[D]")
388-
).astype(float)
383+
dates_arr_y = dates_arr.astype("datetime64[Y]")
384+
dates_arr_d = dates_arr.astype("datetime64[D]")
385+
dates_arr_s = dates_arr.astype("datetime64[s]")
386+
387+
# dyear is DOY 1-366
388+
dyear: npt.NDArray[float] = (dates_arr_d - dates_arr_y).astype(float) + 1
389+
dseconds: npt.NDArray[float] = (dates_arr_s - dates_arr_d).astype(float)
389390
# TODO: Make it a continuous day of year?
390391
# The new code mentions it should be and accepts float, but the
391392
# regression tests indicate it should still be integer DOY

0 commit comments

Comments
 (0)