cross-batch runs not working

[jeanbraun]

Hello @benbovy,

I am not sure if this is the right place (or in xsimlab) to raise this issue, but since updating to the new version of fastscapelib-fortran (and fastscape), I cannot run cross-batch jobs, i.e., containing two batch dimensions that I "stack" and "unstack". I keep getting the following message: ValueError: conflicting multi-index level name VARNAME with dimension VARNAME

I attach a notebook containing a very simple example where the issue arises.

Many thanks for your help Jean (@jeanbraun)

` import numpy as np import xsimlab as xs import xarray as xr

from fastscape.models import basic_model

ds_in = xs.create_setup( model=basic_model, clocks={'time':np.linspace(0,10e7,101)}, input_vars={ 'gridshape': [101,101], 'gridlength': [10e3,10e3], 'boundarystatus': 'fixed_value', 'upliftrate': ('U', [1e-4,1e-3]), 'splk_coef': ('K', [1e-6,1e-5]), 'splarea_exp': 0.4, 'splslope_exp': 1, 'spl__tol_rel': 0.0001, 'spltol_abs': 0.0001, 'splmax_iter': 100, 'diffusion__diffusivity': 1e-3, 'init_topographyseed': None, }, output_vars={'topography__elevation': 'time'} )

import zarr zgroup = zarr.group("Test.zarr", overwrite=True)

with basic_model, xs.monitoring.ProgressBar(): ds_out = ds_in.stack(batch=['U','K']).xsimlab.run(check_dims='transpose',store=zgroup, batch_dim='batch', parallel=True, scheduler='processes').unstack('batch') `

[benbovy]

Hi Jean,

Thanks for the report.

It looks like xarray-simlab needs to be updated in order to be fully compatible with the last version(s) of Xarray. Which version of Xarray do you use (i.e., output of xr.__versions__)? I could reproduce the issue with version 2022.12.0. I guess Xarray has been updated as well when updating your environment.

In the meantime, for your example here is a workaround (working with Xarray version 2022.12.0):

with basic_model, xs.monitoring.ProgressBar():
    ds_out = (
        ds_in
        .stack(batch=['U','K'])
        .drop_indexes(["batch", "U", "K"])    # explicitly drop the multi-index 
        .drop_vars("batch")                   # remove the batch dimension coordinate
        .xsimlab.run(
            check_dims='transpose',
            store=zgroup,
            batch_dim='batch',
            parallel=True,
            scheduler='processes',
        )
        .set_index(batch=["U", "K"])          # explicitly set the multi-index
        .unstack("batch")
    )

Alternatively, you could define dimensions that corresponds to the parameter variables (so that you can use those parameter values as Xarray coordinates!):

ds_in = xs.create_setup(
    input_vars={
        'uplift__rate': ('uplift__rate', [1e-4,1e-3]),
        'spl__k_coef': ('spl__k_coef', [1e-6,1e-5]),
        ...
    },
)

And then

with basic_model, xs.monitoring.ProgressBar():
    ds_out = (
        ds_in
        .stack(batch=['uplift__rate', 'spl__k_coef'], create_index=False)    # do not create an index
        .xsimlab.run(
            check_dims='transpose',
            store=zgroup,
            batch_dim='batch',
            parallel=True,
            scheduler='processes',
        )
        .set_index(batch=["uplift__rate", "spl__k_coef"])           # explicitly set the multi-index
        .unstack("batch")
    )

[jeanbraun]

Hello @benbovy . Thanks. I have version 2022.12.0 too. I'll use your workaround. Can I get it to work if I downgrade array? To which version?

[benbovy]

It should work with version 2022.3.0.

[jeanbraun]

Thanks

[benbovy]

Re-opening as we should make it work with last versions of Xarray.

[benbovy]

Closing as this should be fixed upstream: https://github.com/xarray-contrib/xarray-simlab/issues/193