Loading Checkpoints After Initialization + Reset()

[nzolman]

I had a use-case where I initialized a medium mesh cylinder FlowEnv without a checkpoint (just the default flow), and then later used env.flow.load_checkpoint. The env.render() function worked fine, and showed my checkpointed flow. However, I received an error when I next used env.reset(). Everything seems to work as expected if I use the checkpoint "restart" key in the flow config. But it seems like this is worth pointing out, because intuitively you would expect to be able to load a flow and use env.reset() as long as the meshes were identical.

EDIT: Using latest version of hydrogym on main branch. (38826f4)
Python 3.10.6
firedrake version: 0.13.0+5289.gee720d61

Code:

from hydrogym import FlowEnv
import hydrogym.firedrake as hgym

flow_env_config = {
    "flow": hgym.Cylinder,
    "solver": hgym.IPCS,
    "solver_config":{'dt': 0.001},
    "flow_config": {'mesh': 'medium'}
}

env = FlowEnv(flow_env_config)

# Renders Correctly
env.render()

path_to_ckpt = '/path/to/checkpoint.ckpt'
env.flow.load_checkpoint(path_to_ckpt)

# Renders Correctly
env.render()

# errors
env.reset()

Error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In [4], line 27
     24 env.render()
     26 # errors
---> 27 env.reset()

File /home/firedrake/hydrogym/hydrogym/core.py:374, in FlowEnv.reset(self)
    372 def reset(self) -> Union[PDEBase.ObsType, Tuple[PDEBase.ObsType, dict]]:
    373     self.iter = 0
--> 374     self.flow.reset(q0=self.q0)
    375     self.solver.reset()
    377     return self.flow.get_observations()

File /home/firedrake/hydrogym/hydrogym/core.py:102, in PDEBase.reset(self, q0)
     94 """Reset the PDE to an initial state
     95 
     96 Args:
   (...)
     99         Defaults to None.
    100 """
    101 if q0 is not None:
--> 102     self.set_state(q0)
    103 self.reset_controls()

File /home/firedrake/hydrogym/hydrogym/firedrake/flow.py:91, in FlowConfig.set_state(self, q)
     85 def set_state(self, q: fd.Function):
     86     """Set the current state fields
     87 
     88     Args:
     89         q (fd.Function): State to be assigned
     90     """
---> 91     self.q.assign(q)

File PETSc/Log.pyx:115, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func()

File PETSc/Log.pyx:116, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func()

File /home/firedrake/firedrake/src/firedrake/firedrake/adjoint/function.py:117, in FunctionMixin._ad_annotate_assign.<locals>.wrapper(self, other, *args, **kwargs)
    114     tape.add_block(block)
    116 with stop_annotating():
--> 117     ret = assign(self, other, *args, **kwargs)
    119 if annotate:
    120     block.add_output(self.create_block_variable())

File <decorator-gen-48>:2, in assign(self, expr, subset)

File /home/firedrake/firedrake/src/firedrake/firedrake/utils.py:74, in known_pyop2_safe.<locals>.wrapper(f, *args, **kwargs)
     72 opts["type_check"] = safe
     73 try:
---> 74     return f(*args, **kwargs)
     75 finally:
     76     opts["type_check"] = check

File /home/firedrake/firedrake/src/firedrake/firedrake/function.py:412, in Function.assign(self, expr, subset)
    409     return self
    411 from firedrake import assemble_expressions
--> 412 assemble_expressions.evaluate_expression(
    413     assemble_expressions.Assign(self, expr), subset)
    414 return self

File PETSc/Log.pyx:115, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func()

File PETSc/Log.pyx:116, in petsc4py.PETSc.Log.EventDecorator.decorator.wrapped_func()

File <decorator-gen-38>:2, in evaluate_expression(expr, subset)

File /home/firedrake/firedrake/src/firedrake/firedrake/utils.py:71, in known_pyop2_safe.<locals>.wrapper(f, *args, **kwargs)
     69 safe = parameters["type_check_safe_par_loops"]
     70 if check == safe:
---> 71     return f(*args, **kwargs)
     72 opts["type_check"] = safe
     73 try:

File /home/firedrake/firedrake/src/firedrake/firedrake/assemble_expressions.py:520, in evaluate_expression(expr, subset)
    516         except ReferenceError:
    517             # TODO: Is there a situation where some of the kernels
    518             # succeed and others don't?
    519             pass
--> 520 arguments = expr.par_loop_args
    521 if cache is not None:
    522     cache[slow_key] = arguments

File /home/firedrake/firedrake/src/PyOP2/pyop2/utils.py:62, in cached_property.__get__(self, obj, cls)
     60 if obj is None:
     61     return self
---> 62 obj.__dict__[self.__name__] = result = self.fget(obj)
     63 return result

File /home/firedrake/firedrake/src/firedrake/firedrake/assemble_expressions.py:327, in Assign.par_loop_args(self)
    325 result = []
    326 grouping = OrderedDict()
--> 327 for e in self.split:
    328     grouping.setdefault(e.lvalue.node_set, []).append(e)
    329 for iterset, exprs in grouping.items():

File /home/firedrake/firedrake/src/PyOP2/pyop2/utils.py:62, in cached_property.__get__(self, obj, cls)
     60 if obj is None:
     61     return self
---> 62 obj.__dict__[self.__name__] = result = self.fget(obj)
     63 return result

File /home/firedrake/firedrake/src/firedrake/firedrake/assemble_expressions.py:270, in Assign.split(self)
    267 if indices == set([None]):
    268     if len((set(spaces) | {V}) - {None}) != 1:
    269         # Check that there were no unindexed coefficients
--> 270         raise ValueError("Saw indexed coefficients in rvalue, "
    271                          "perhaps you meant to index the lvalue with .sub(...)")
    272     rvalues = self.splitter(self.rvalue, False)
    273     return tuple(type(self)(lvalue, rvalue)
    274                  for lvalue, rvalue in zip(self.lvalue.split(), rvalues))

ValueError: Saw indexed coefficients in rvalue, perhaps you meant to index the lvalue with .sub(...)

[ludgerpaehler]

Do I understand it correctly, that the above is already enough to reproduce the bug? :)

(Also, sorry for not getting on top of this earlier)

[nzolman]

@ludgerpaehler—should be! I've updated to include my python and firedrake versions. The only additional assumption should be that there is already some saved mesh (also a medium mesh from the cylinder) created from env.flow.save_checkpoint.