Skip to content

Commit 2bb1800

Browse files
Speed up create_input when dates are already np.datetime64
Simplify line of code formatting ruff
1 parent a1f0057 commit 2bb1800

File tree

1 file changed

+8
-9
lines changed

1 file changed

+8
-9
lines changed

pymsis/msis.py

Lines changed: 8 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -377,15 +377,14 @@ def create_input(
377377
(ndates == nlons == nlats == nalts), then the shape is (ndates,).
378378
"""
379379
# Turn everything into arrays
380-
dates_arr: npt.NDArray[np.datetime64] = np.atleast_1d(
381-
np.array(dates, dtype=np.datetime64)
382-
)
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)
380+
dates_arr: npt.NDArray[np.datetime64] = np.atleast_1d(dates).astype(np.datetime64)
381+
dates_arr_y = dates_arr.astype("datetime64[Y]")
382+
dates_arr_d = dates_arr.astype("datetime64[D]")
383+
dates_arr_s = dates_arr.astype("datetime64[s]")
384+
385+
# dyear is DOY 1-366
386+
dyear: npt.NDArray[float] = (dates_arr_d - dates_arr_y).astype(float) + 1
387+
dseconds: npt.NDArray[float] = (dates_arr_s - dates_arr_d).astype(float)
389388
# TODO: Make it a continuous day of year?
390389
# The new code mentions it should be and accepts float, but the
391390
# regression tests indicate it should still be integer DOY

0 commit comments

Comments
 (0)