Skip to content

Fixed potential memory leak in local_solar_time #2356

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Mar 7, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 27 additions & 25 deletions esmvalcore/preprocessor/_time.py
Original file line number Diff line number Diff line change
Expand Up @@ -1436,9 +1436,9 @@ def _get_time_index_and_mask(

def _transform_to_lst_eager(
data: np.ndarray,
*,
time_index: np.ndarray,
mask: np.ndarray,
*,
time_dim: int,
lon_dim: int,
**__,
Expand Down Expand Up @@ -1479,9 +1479,9 @@ def _transform_to_lst_eager(

def _transform_to_lst_lazy(
data: da.core.Array,
*,
time_index: np.ndarray,
mask: np.ndarray,
*,
time_dim: int,
lon_dim: int,
output_dtypes: DTypeLike,
Expand All @@ -1505,28 +1505,25 @@ def _transform_to_lst_lazy(
`mask` is 2D with shape (time, lon) that will be applied to the final data.

"""
_transform_chunk_to_lst = partial(
_transform_to_lst_eager,
time_index=time_index,
mask=mask,
time_dim=-2, # this is ensured by da.apply_gufunc
lon_dim=-1, # this is ensured by da.apply_gufunc
)
new_data = da.apply_gufunc(
_transform_chunk_to_lst,
'(t,y)->(t,y)',
_transform_to_lst_eager,
'(t,x),(t,x),(t,x)->(t,x)',
data,
axes=[(time_dim, lon_dim), (time_dim, lon_dim)],
time_index,
mask,
axes=[(time_dim, lon_dim), (0, 1), (0, 1), (time_dim, lon_dim)],
output_dtypes=output_dtypes,
time_dim=-2, # this is ensured by da.apply_gufunc
lon_dim=-1, # this is ensured by da.apply_gufunc
)
return new_data


def _transform_arr_to_lst(
data: np.ndarray | da.core.Array,
*,
time_index: np.ndarray,
mask: np.ndarray,
*,
time_dim: int,
lon_dim: int,
output_dtypes: DTypeLike,
Expand All @@ -1545,8 +1542,8 @@ def _transform_arr_to_lst(
func = _transform_to_lst_lazy # type: ignore
new_data = func(
data, # type: ignore
time_index=time_index,
mask=mask,
time_index,
mask,
time_dim=time_dim,
lon_dim=lon_dim,
output_dtypes=output_dtypes,
Expand All @@ -1571,13 +1568,10 @@ def _transform_cube_to_lst(cube: Cube) -> Cube:

# Transform cube data
(time_index, mask) = _get_time_index_and_mask(time_coord, lon_coord)
_transform_arr = partial(
_transform_arr_to_lst,
time_index=time_index,
mask=mask,
)
cube.data = _transform_arr(
cube.data = _transform_arr_to_lst(
cube.core_data(),
time_index,
mask,
time_dim=time_dim,
lon_dim=lon_dim,
output_dtypes=cube.dtype,
Expand All @@ -1589,15 +1583,19 @@ def _transform_cube_to_lst(cube: Cube) -> Cube:
if time_dim in dims and lon_dim in dims:
time_dim_ = dims.index(time_dim)
lon_dim_ = dims.index(lon_dim)
coord.points = _transform_arr(
coord.points = _transform_arr_to_lst(
coord.core_points(),
time_index,
mask,
time_dim=time_dim_,
lon_dim=lon_dim_,
output_dtypes=coord.dtype,
)
if coord.has_bounds():
coord.bounds = _transform_arr(
coord.bounds = _transform_arr_to_lst(
coord.core_bounds(),
time_index,
mask,
time_dim=time_dim_,
lon_dim=lon_dim_,
output_dtypes=coord.bounds_dtype,
Expand All @@ -1609,8 +1607,10 @@ def _transform_cube_to_lst(cube: Cube) -> Cube:
if time_dim in dims and lon_dim in dims:
time_dim_ = dims.index(time_dim)
lon_dim_ = dims.index(lon_dim)
cell_measure.data = _transform_arr(
cell_measure.data = _transform_arr_to_lst(
cell_measure.core_data(),
time_index,
mask,
time_dim=time_dim_,
lon_dim=lon_dim_,
output_dtypes=cell_measure.dtype,
Expand All @@ -1622,8 +1622,10 @@ def _transform_cube_to_lst(cube: Cube) -> Cube:
if time_dim in dims and lon_dim in dims:
time_dim_ = dims.index(time_dim)
lon_dim_ = dims.index(lon_dim)
anc_var.data = _transform_arr(
anc_var.data = _transform_arr_to_lst(
anc_var.core_data(),
time_index,
mask,
time_dim=time_dim_,
lon_dim=lon_dim_,
output_dtypes=anc_var.dtype,
Expand Down