fieldset.to_windowed_arrays is slow on CROCO tutorial - not loading time-invariant variables

[erikvansebille]

What version of Parcels are you running?

v4

Is your feature request related to a problem?

While working on using fieldset.to_windowed_arrays() in the documentation (#2748), I realised that the new method is performant in all tutorials except for the tutorial_CROCO_3D tutorial.

On my computer, the pset.execute in this tutorial ran in ~5 seconds with ds.load(), but more than 3 minutes with fieldset.to_windowed_arrays()

A bit of digging showed that this was because the variables that are time-independent (like lon and lat, but also h and Zeta) were not converted to numpy so stayed in dask; and then indexing on them was very slow

Describe the solution you'd like

By default, also convert every variable that is independent of time into a numpy array. These are time-constant fields, so we can assume that they are much smaller than the time-varying fields.

Describe alternatives you've considered

Leave it to users to pick which variables are turned to windowed_arrays, which are turned to numpy and which are left as dask. But that will be a steep learning curve

Additional context

Here's a bit of code that shows the difference in speed

def test_speed_on_CROCO_data():
    import parcels.tutorial
    ds_fields = parcels.tutorial.open_dataset("CROCOidealized_data/data")

    fields = {
        "U": ds_fields["u"],
        "V": ds_fields["v"],
        "W": ds_fields["w"],
        "h": ds_fields["h"],
        "zeta": ds_fields["zeta"],
        "Cs_w": ds_fields["Cs_w"],
    }

    ds_fset = parcels.convert.croco_to_sgrid(fields=fields, coords=ds_fields)

    ### REMOVE THIS CODE FOR SLOW TESTING
    for v in ["lon", "lat", "h", "zeta", "Cs_w"]:
        ds_fset[v].load()
    ###

    fieldset = parcels.FieldSet.from_sgrid_conventions(ds_fset)
    fieldset.add_context("hc", ds_fields.hc.item())
    fieldset = fieldset.to_windowed_arrays()

    for name, var in ds_fset.variables.items():
        data = var.data
        if isinstance(data, da.Array):
            backend = "dask"
        elif isinstance(data, np.ndarray):
            backend = "numpy"
        else:
            backend = type(data).__name__
        print(f"{name}: {backend}")

    X, Z = np.meshgrid(
        [40e3, 80e3, 120e3],
        [100, -10, -130, -250, -400, -850, -1400, -1550],
    )
    Y = np.ones(X.size) * 98000


    def DeleteParticle(particles, fieldset):
        any_error = particles.state >= 50
        particles[any_error].state = parcels.StatusCode.Delete


    pset = parcels.ParticleSet(fieldset=fieldset, pclass=parcels.Particle, x=X, y=Y, z=Z)

    outputfile = parcels.ParticleFile(
        path="croco_particles3D.parquet",
        outputdt=np.timedelta64(5000, "s"),
        mode="w",
    )

    pset.execute(
        [parcels.kernels.AdvectionRK4_3D_CROCO, DeleteParticle],
        runtime=np.timedelta64(50_000, "s"),
        dt=np.timedelta64(100, "s"),
        output_file=outputfile,
    )