Python
Collapsed Code
1
def optimize_params_window(sif_w, ogvi_w, ndwi_w, lst_w,
2
                           param_ini, param_bounds, min_obs):
3
    """
4
    Wrapper function to optimize SIF parameters for a single window.
5
    Designed to be used with apply_ufunc or iteration.
6
    Assumes input arrays (sif_w, etc.) are 2D numpy arrays (window data).
7
    """
8
    # Flatten the window arrays and filter out NaNs
9
    # Important: Filter consistently across all variables
10
    mask = ~np.isnan(sif_w) & ~np.isnan(ogvi_w) & ~np.isnan(ndwi_w) & ~np.isnan(lst_w)
11
    n_valid = np.sum(mask)
12
13
    if n_valid < min_obs:
14
        # Not enough valid observations in the window
15
        return np.full(len(param_ini), np.nan, dtype=np.float32)
16
    
17
18
    sif_obs_f = sif_w[mask]
19
    vi_f = ogvi_w[mask]
20
    et_f = ndwi_w[mask] # Using NDWI as proxy for water stress/ET effect
21
    lst_f = lst_w[mask]
22
23
    # Define bounds for L-BFGS-B
24
    bounds_scipy = list(zip(*param_bounds)) # [(min1, max1), (min2, max2), ...]
25
26
   
27
    result = minimize(cost_function,
28
                              x0=param_ini,
29
                              args=(vi_f, et_f, lst_f, sif_obs_f),
30
                              method='L-BFGS-B',
31
                              bounds=bounds_scipy,
32
                              options={'maxiter': 1000, 'ftol': 1e-7, 'gtol': 1e-5}) # Adjust options as needed
33
34
    optimized_params = np.array(np.clip(result.x, param_bounds[0], param_bounds[1]))
35
    #print(optimized_params.dtype)
36
    #print(optimized_params.size)
37
    #print(optimized_params.shape)
38
    return optimized_params
Python

Python
Collapsed Code
1
# Define optimization parameters and bounds
2
param_ini = np.array([1.0, 2.0, -295.0, 10.0])
3
param_min = np.array([0.5, 0.1, -310.0, 1.0])
4
param_max = np.array([1.5, 5.0, -290.0, 50.0])
5
param_bounds = np.array([param_min, param_max])
6

7
min_obs_optim = 21 # Minimum valid observations in window
8
window_size_optim = 5 # Must match Julia 'window_edge'
9

10
# using apply_ufunc
11
parameters_cube = xr.apply_ufunc(
12
    optimize_params_window, # Function to apply
13
    # Input arrays:
14
    sif_cube_low_july['SIF_743'].rolling(lat=window_size_lat, lon=window_size_lon, center=True).construct(lat = 'lat_roll', lon = 'lon_roll'),
15
    ogvi_cube_low_july['OGVI'].rolling(lat=window_size_lat, lon=window_size_lon, center=True).construct(lat = 'lat_roll', lon = 'lon_roll'),
16
    ndwi_cube_low_july['NDWI'].rolling(lat=window_size_lat, lon=window_size_lon, center=True).construct(lat = 'lat_roll', lon = 'lon_roll'),
17
    lst_cube_low_july['LST'].rolling(lat=window_size_lat, lon=window_size_lon, center=True).construct(lat = 'lat_roll', lon = 'lon_roll'),
18
    # Keyword arguments for the function:
19
    kwargs={'param_ini': param_ini, 'param_bounds': param_bounds, 'min_obs': min_obs_optim},
20
    # Define input core dimensions (the window dimensions):
21
    input_core_dims=[['lat_roll', 'lon_roll'], ['lat_roll', 'lon_roll'], ['lat_roll', 'lon_roll'], ['lat_roll', 'lon_roll']],
22
    # Define output core dimensions (the parameters dimension):
23
    output_core_dims=[['parameters']],
24
    dask_gufunc_kwargs={'output_sizes': {'parameters': 6}},
25
    # Specify the parameters dimension coordinates:
26
    dask='parallelized', # Enable Dask parallelization
27
    output_dtypes=[np.float64], # Specify output data type
28
    vectorize = True,
29
    # Add a new dimension 'parameters' to the output
30
    exclude_dims=set(('lat_roll', 'lon_roll')), # Exclude window dims from output
31
).rename('optimized_parameters')
32

33
# Assign coordinates to the new 'parameters' dimension
34

35
parameters_cube['parameters'] = ["b1", "b2", "b3", "b4", "b5", "b6"]
36
parameters_cube = parameters_cube.chunk({'lat': 80, 'lon': 80, 'parameters': 6}) # Rechunk
37

Python

Python
Collapsed Code
1
parameters_cube_high = parameters_cube.interp_like(
2
    lst_cube_high_july, # Use one of the high-res grids as template
3
    method="linear",    # Options: 'linear', 'nearest'
4
    #kwargs={"fill_value": "extrapolate"} # Optional: handle edges
5
    # Or fill_value=np.nan if extrapolation is not desired
6
   # kwargs = {'fill_value':np.nan},
7
).chunk({'lat': 400, 'lon': 400, 'parameters': -1})
Python

Python
Collapsed Code
1
# Define the wrapper function for downscaling within a window
2
def sif_downscaling_window(ogvi_hw, lst_hw, ndwi_hw, params_hw):
3
    """
4
    Calculates downscaled SIF for a central pixel based on window means.
5
    Inputs are 3D numpy arrays (window_x, window_y, [parameters]).
6
    """
7
   
8
9
    # Calculate mean parameters within the window (ignore NaNs)
10
    # Need to handle case where all params in window are NaN
11
    mean_params = np.nanmean(params_hw, axis=(0,1))
12
13
    if np.all(np.isnan(mean_params)):
14
        return np.array([np.nan]) # Cannot calculate if no valid parameters in window)
15
16
    # Calculate mean predictors within the window (ignore NaNs)
17
    mean_ogvi = np.nanmean(ogvi_hw)
18
    mean_ndwi = np.nanmean(ndwi_hw)
19
    mean_lst = np.nanmean(lst_hw)
20
21
    # Check if any mean predictor is NaN (means all values in window were NaN)
22
    if np.isnan(mean_ogvi) or np.isnan(mean_lst):
23
        return np.array([np.nan])
24
25
    # Apply SIF model using mean values
26
    sif_ds = sif_model(mean_ogvi, mean_lst, mean_params)
27
    return np.array([sif_ds])
28
    
29
## applying the function to predict SIF
30

31
window_size_downscale = 3
32

33
sif_cube_high = xr.apply_ufunc(
34
    sif_downscaling_window,
35
    # Input arrays with rolling windows constructed
36
    ogvi_cube_high_july['OGVI'].rolling(lat=window_size_downscale, lon=window_size_downscale, center=True).construct(lat = 'lat_roll', lon = 'lon_roll'),
37
    ndwi_aligned.rolling({x_dim: window_size_downscale, y_dim: window_size_downscale}, center=True).construct(x_dim+'_w', y_dim+'_w'),
38
    lst_cube_high_july['LST'].rolling(lat=window_size_downscale, lon=window_size_downscale, center=True).construct(lat = 'lat_roll', lon = 'lon_roll'),
39
    parameters_cube_high['optimized_parameters'].rolling(lat=window_size_downscale, lon=window_size_downscale, center=True).construct(lat = 'lat_roll', lon = 'lon_roll').chunk({'parameters':-1}),
40
    # Input core dimensions now include the window dims and the parameter dim for the last input
41
    input_core_dims=[['lat_roll', 'lon_roll'],
42
                     ['lat_roll', 'lon_roll'],
43
                     ['lat_roll', 'lon_roll', 'parameters']],
44
    # Output is scalar for each pixel (no core dims)
45
    output_core_dims=[['SIF_downscaled']],
46
    dask_gufunc_kwargs={'output_sizes': {'SIF_downscaled': 1}},
47
    dask='parallelized',
48
    output_dtypes=[np.float64],
49
    vectorize = True,
50
    exclude_dims=set(('lat_roll', 'lon_roll')), # Exclude window dims from output
51
).rename('SIF_downscaled')
52

53
sif_cube_high = sif_cube_high.to_dataset(name = 'SIF')
54

55
sif_cube_high.attrs['description'] = f'Downscaled SIF using Duveiller&Cescatti method, {window_size_downscale}x{window_size_downscale} window smoothing'
56

57

Python