id like to adapt this code to work with a convolutional neural network. i want the network to intake a 5 by 5grid of data where these values are:     z_flat,              # Shape: (35520, 3)
    q_flat,
    t_flat,
    u_flat,
    v_flat,
    w_flat,
    msl_flat[:, None],
    lsm_flat[:, None],
    t2m_flat[:, None],
    u10_flat[:, None],
    v10_flat[:, None], so that is one input to the network, the other are three more features:     year_flat[:, None],
    lat_flat[:, None],
    lon_flat[:, None]. please adapt my code so that it works like i decrivbed. import numpy as np

def datetime64_to_year_fraction(dt64):
    # Convert to datetime64[D] to get just the date part
    dt = dt64.astype('datetime64[D]')

    # Extract year, month, and day using numpy's functions
    years = dt.astype('datetime64[Y]').astype(int) + 1970  # Adjusting for epoch
    months = (dt.astype('datetime64[M]').astype(int) % 12) + 1  # Months range from 1 to 12
    days = (dt.astype('datetime64[D]').astype(int) % 31) + 1  # Days range from 1 to 31

    # Calculate month and day portions
    MONTH_PORTION = (months - 1) / 12
    DAY_PORTION = (days - 1) / (365.25)  # Using average days in a year

    return MONTH_PORTION + DAY_PORTIONimport random
import xarray as xr
from tqdm import tqdm

list_of_arrays = []

# Select the required data points using the mask
dataset = training_dataset_2013_to_2023_france_spain

# Extract the required variables using vectorized operations
year_array = datetime64_to_year_fraction(dataset.time.values)
lat_array = dataset.latitude.values
lon_array = dataset.longitude.values

z_array = dataset.z.sel(isobaricInhPa=[1000, 925, 850]).values
q_array = dataset.q.sel(isobaricInhPa=[1000, 925, 850]).values
t_array = dataset.t.sel(isobaricInhPa=[1000, 925, 850]).values
u_array = dataset.u.sel(isobaricInhPa=[1000, 925, 850]).values
v_array = dataset.v.sel(isobaricInhPa=[1000, 925, 850]).values
w_array = dataset.w.sel(isobaricInhPa=[1000, 925, 850]).values

msl_array = dataset.msl.values
lsm_array = dataset.lsm.values
t2m_array = dataset.t2m.values
u10_array = dataset.u10.values
v10_array = dataset.v10.values
u100_array = dataset.u100.values
v100_array = dataset.v100.values

print(year_array.shape)
print(lat_array.shape)
print(lon_array.shape)
print(z_array.shape)
print(q_array.shape)
print(t_array.shape)
print(u_array.shape)
print(v_array.shape)
print(w_array.shape)
print(msl_array.shape)
print(lsm_array.shape)
print(t2m_array.shape)
print(u10_array.shape)
print(v10_array.shape)
print(u100_array.shape)
print(v100_array.shape)import numpy as np

time_dimension = dataset.time.values.shape[0]

# Create meshgrids for time, latitude, and longitude to align with your data dimensions
year_grid = np.broadcast_to(year_array[:, None, None], (time_dimension, v10_array.shape[1], v10_array.shape[2]))
lat_grid = np.broadcast_to(lat_array[None, :, None], (time_dimension, v10_array.shape[1], v10_array.shape[2]))
lon_grid = np.broadcast_to(lon_array[None, None, :], (time_dimension, v10_array.shape[1], v10_array.shape[2]))

# Flatten the grids to create 1D arrays
year_flat = year_grid.flatten()
year_grid = [].copy()
lat_flat = lat_grid.flatten()
lat_grid = [].copy()
lon_flat = lon_grid.flatten()
lon_grid = [].copy()

# Function to transpose and flatten variables with the isobaric dimension
def transpose_and_flatten(var_array):
    # Transpose to move the isobaric dimension to the last axis
    var_array_t = np.transpose(var_array, (0, 2, 3, 1))  # Shape: (32, 30, 37, 3)
    # Flatten to 2D array with shape (35520, 3)
    var_flat = var_array_t.reshape(-1, 3)
    return var_flat

# Flatten the variables with the isobaric dimension
print("Flattening Z")
z_flat = transpose_and_flatten(z_array)
z_array = [].copy()
print("Flattening Q")
q_flat = transpose_and_flatten(q_array)
q_array = [].copy()
print("Flattening T")
t_flat = transpose_and_flatten(t_array)
t_array = [].copy()
print("Flattening U")
u_flat = transpose_and_flatten(u_array)
u_array = [].copy()
print("Flattening V")
v_flat = transpose_and_flatten(v_array)
v_array = [].copy()
print("Flattening W")
w_flat = transpose_and_flatten(w_array)
w_array = [].copy()

# Flatten the single-level variables
print("Flattening MSL")
msl_flat = msl_array.flatten()
msl_array = [].copy()
print("Flattening LSM")
lsm_flat = lsm_array.flatten()
lsm_array = [].copy()
print("Flattening T2M")
t2m_flat = t2m_array.flatten()
t2m_array = [].copy()
print("Flattening U10")
u10_flat = u10_array.flatten()
u10_array = [].copy()
print("Flattening V10")
v10_flat = v10_array.flatten()
v10_array = [].copy()
print("Flattening U100")
u100_flat = u100_array.flatten()
u100_array = [].copy()
print("Flattening V100")
v100_flat = v100_array.flatten()
v100_array = [].copy()

# **Now, use `time_seconds` instead of `time_flat`**
variables_list = [
    year_flat[:, None],
    lat_flat[:, None],
    lon_flat[:, None],
    z_flat,              # Shape: (35520, 3)
    q_flat,
    t_flat,
    u_flat,
    v_flat,
    w_flat,
    msl_flat[:, None],
    lsm_flat[:, None],
    t2m_flat[:, None],
    u10_flat[:, None],
    v10_flat[:, None],
    u100_flat[:, None],
    v100_flat[:, None],
]