outer_steps = 1  # total of 4 days
timedelta = np.timedelta64(1, 'h') * inner_steps
times = (np.arange(outer_steps) * inner_steps)  # time axis in hours

# initialize model state
data_dict, forcings = model.data_from_xarray(eval_era5.head(time=1))
inputs, input_forcings = pytree_utils.slice_along_axis((data_dict, forcings), axis=0, idx=0)
state = model.encode(
    inputs, forcings=input_forcings, rng_key=None
)

# make forecast
final_state, predictions = model.unroll(
    state,
    forcings,
    steps=outer_steps,
    timedelta=timedelta,
    start_with_input=False,
)
predictions_ds = model.data_to_xarray(predictions, times=times)

print(np.isnan(predictions_ds["u_component_of_wind"]).sum())

import matplotlib.pyplot as plt
for i in range(len(predictions_ds.level)):
    print(predictions_ds.level[i])
    array = predictions_ds["u_component_of_wind"][0,i,:]
    plt.imshow(array, cmap='viridis', aspect='auto')
    plt.colorbar()
    plt.title('2D Array Visualization')
    plt.xlabel('X-axis')
    plt.ylabel('Y-axis')
    plt.show()