Cannot run advect_model (new user)

[lalligagger]

Hi, I've just started taking a look into xarray-simlab and would love to try using for some modeling I'm doing! Currently I'm not able to get all the way through the example model, after doing a pip install.

I have defined my processes/ models as in the "Creating Models" example:

@xs.process
class UniformGrid1D:
    """Create a 1-dimensional, equally spaced grid."""

    spacing = xs.variable(description="uniform spacing", static=True)
    length = xs.variable(description="total length", static=True)
    x = xs.index(dims="x")

    def initialize(self):
        self.x = np.arange(0, self.length, self.spacing)

@xs.process
class ProfileU:
    """Compute the evolution of the profile of quantity `u`."""

    u_vars = xs.group("u_vars")
    u = xs.variable(
        dims="x", intent="inout", description="quantity u", attrs={"units": "m"}
    )

    def run_step(self):
        self._delta_u = sum((v for v in self.u_vars))

    def finalize_step(self):
        self.u += self._delta_u

@xs.process
class AdvectionLax:
    """Advection using finite difference (Lax method) on
    a fixed grid with periodic boundary conditions.

    """

    v = xs.variable(dims=[(), "x"], description="velocity")
    grid_spacing = xs.foreign(UniformGrid1D, "spacing")
    u = xs.foreign(ProfileU, "u")
    u_advected = xs.variable(dims="x", intent="out", groups="u_vars")

@xs.runtime(args="step_delta")
def run_step(self, dt):
    factor = self.v / (2 * self.grid_spacing)

    u_left = np.roll(self.u, 1)
    u_right = np.roll(self.u, -1)
    u_1 = 0.5 * (u_right + u_left) - factor * dt * (u_right - u_left)

    self.u_advected = u_1 - self.u

@xs.process
class InitUGauss:
    """Initialize `u` profile using a Gaussian pulse."""

    loc = xs.variable(description="location of initial pulse", static=True)
    scale = xs.variable(description="scale of initial pulse", static=True)
    x = xs.foreign(UniformGrid1D, "x")
    u = xs.foreign(ProfileU, "u", intent="out")

    def initialize(self):
        self.u = np.exp(-1 / self.scale ** 2 * (self.x - self.loc) ** 2)

advect_model = xs.Model(
    {
        "grid": UniformGrid1D,
        "profile": ProfileU,
        "init": InitUGauss,
        "advect": AdvectionLax,
    }
)

here is the ds_in I generated using the autogeneration and the "Setup and Run" example values:

# %create_setup advect_model --default --verbose

ds_in = xs.create_setup(
    model=advect_model,
    clocks={
        'time': np.linspace(0., 1., 101),
        'otime': [0, 0.5, 1],     
    },
    master_clock='time',
    input_vars={
        # uniform spacing
        'grid__spacing': 0.01,
        # total length
        'grid__length': 1.5,
        # location of initial pulse
        'init__loc': 0.3,
        # scale of initial pulse
        'init__scale': 0.1,
        # velocity
        'advect__v': 1.,
    },
    output_vars={
        'profile__u': 'otime'
    }
)

On run, I get key errors on ('advect', 'u_advected')....

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-15-41c902d1e936> in <module>
      2 
      3 with advect_model:
----> 4     out_ds2 = (ds_in.xsimlab.update_vars(input_vars=in_vars).xsimlab.run())

~/.virtualenvs/stac/lib/python3.9/site-packages/xsimlab/xr_accessor.py in run(self, model, batch_dim, check_dims, validate, store, encoding, decoding, hooks, parallel, scheduler, safe_mode)
    821         )
    822 
--> 823         driver.run_model()
    824 
    825         return driver.get_results()

~/.virtualenvs/stac/lib/python3.9/site-packages/xsimlab/drivers.py in run_model(self)
    470 
    471         if self.batch_dim is None:
--> 472             _run(
    473                 ds_in,
    474                 self.model,

~/.virtualenvs/stac/lib/python3.9/site-packages/xsimlab/drivers.py in _run(dataset, model, store, hooks, validate, batch, batch_size, parallel, scheduler)
    347             in_vars = _get_input_vars(ds_step, model)
    348             model.update_state(in_vars, validate=validate_inputs, ignore_static=False)
--> 349             signal = model.execute("run_step", rt_context, **execute_kwargs)
    350 
    351             if signal == RuntimeSignal.BREAK:

~/.virtualenvs/stac/lib/python3.9/site-packages/xsimlab/model.py in execute(self, stage, runtime_context, hooks, validate, parallel, scheduler)
   1021         else:
   1022             for p_obj in self._processes.values():
-> 1023                 _, (_, signal_process) = self._execute_process(p_obj, *execute_args)
   1024 
   1025                 if signal_process == RuntimeSignal.BREAK:

~/.virtualenvs/stac/lib/python3.9/site-packages/xsimlab/model.py in _execute_process(self, p_obj, stage, runtime_context, hooks, validate, state)
    840             return p_name, ({}, signal_pre)
    841 
--> 842         state_out, signal_out = executor.execute(
    843             p_obj, stage, runtime_context, state=state
    844         )

~/.virtualenvs/stac/lib/python3.9/site-packages/xsimlab/process.py in execute(self, p_obj, stage, runtime_context, state)
    528             return {}, RuntimeSignal.NONE
    529         else:
--> 530             signal_out = executor.execute(p_obj, runtime_context, state=state)
    531 
    532             skeys = [p_obj.__xsimlab_state_keys__[k] for k in self.out_vars]

~/.virtualenvs/stac/lib/python3.9/site-packages/xsimlab/process.py in execute(self, p_obj, runtime_context, state)
    392         args = [runtime_context[k] for k in self.args]
    393 
--> 394         signal = self.meth(p_obj, *args)
    395 
    396         if signal is None:

<ipython-input-2-1a1f078aa29b> in run_step(self)
     20 
     21     def run_step(self):
---> 22         self._delta_u = sum((v for v in self.u_vars))
     23 
     24     def finalize_step(self):

<ipython-input-2-1a1f078aa29b> in <genexpr>(.0)
     20 
     21     def run_step(self):
---> 22         self._delta_u = sum((v for v in self.u_vars))
     23 
     24     def finalize_step(self):

~/.virtualenvs/stac/lib/python3.9/site-packages/xsimlab/process.py in getter_state_or_on_demand(self)
    275 
    276         for key in state_keys:
--> 277             yield self.__xsimlab_state__[key]
    278 
    279         for key in od_keys:

KeyError: ('advect', 'u_advected')

[lalligagger]

• attempting to run model with ds_out = ds_in.xsimlab.run(model=advect_model)

[lalligagger]

This worked when using the advect model here: https://github.com/xarray-contrib/xarray-simlab/blob/master/doc/scripts/advection_model.py

Looking forward to using this more and maybe helping out once I get up the learning curve!