add continue feature for depletion

[lewisgross1296]

Description

This PR closes #2871. To summarize the motivation, I'd like depletion restarts to be a bit easier for users. If you set up a depletion simulation that doesn't end up getting run all the way (think hitting max wall time on an HPC or some other interrupt), the burden then falls to the user to correctly determine what time steps were not run and to create a new script that picks up where the other one left off. In my own experience, it is easy to make a mistake and re-run a simulation that is not what I desired.

All we do here is add a continue_timesteps flag (default to False) that adds a new logic block to the Integrator class. It would be useful in the above case and requires users to provide the Integrator with a set of timesteps and one of power/power_density/source_rate that matches what already exists in a prev_results passed to the Opperator.

With this flag, the depletion restart python script can now be exactly the same as the initial script, just with a continue_timesteps = True flag and a prev_results object loaded from a depletion_statepoint.h5 passed to the Operator. It could contain more timesteps at the end too, but the only requirement is that the data matches what is in the Operator provided to the Integrator for all the existing steps in the results.

This PR does not eliminate any of the existing capability of the depletion API, so all the old use cases/syntax for restarts remain. The flag is optional and defaults to False so no one will need to update past scripts with this change.

Checklist

• I have performed a self-review of my own code
• I have run clang-format (version 15) on any C++ source files (if applicable)
• I have followed the style guidelines for Python source files (if applicable)
• I have made corresponding changes to the documentation (if applicable)
• I have added tests that prove my fix is effective or that my feature works (if applicable)

[lewisgross1296]

Looking for input on this line (636)

[dean-krueger]

Just one fairly minor comment.

[connoramoreno]

Re-suggesting my comment from before since it was outdated by the commit of Edgar's suggestion. Might need to apply proper formatting for OpenMC, if applicable (should be black-compliant)

[lewisgross1296]

Looks like it hit the ValueError on the test_continue that was supposed to use a correct continue run. Is this a case where the action-tmate would help us debug?

            if continue_timesteps:
                completed_timesteps = operator.prev_res.get_times()
                completed_source_rates = operator.prev_res.get_source_rates()
                num_previous_steps_run = len(completed_timesteps)
                if (
                    np.array_equal(completed_timesteps, timesteps[: num_previous_steps_run]) and
                    np.array_equal(completed_source_rates, source_rates[: num_previous_steps_run])
                ):
                    seconds = seconds[num_previous_steps_run:]
                    source_rates = source_rates[num_previous_steps_run:]
                else:
>                   raise ValueError(
                        "You are attempting to continue a run in which the previous results "
                        "do not have the same initial steps as those provided to the "
                        "Integrator. Please make sure you are using the correct timesteps, "
                        "powers or power densities, and previous results file."
                    )
E                   ValueError: You are attempting to continue a run in which the previous results do not have the same initial steps as those provided to the Integrator. Please make sure you are using the correct timesteps, powers or power densities, and previous results file.

[pshriwise]

A few things to address here:

• We'll want to make sure that the call to Results.get_times() provides the same unis as the time steps passed to the Integrator class here.
• The times coming from the Results object are points in time, not dt, so we'll need to take the diff of that array to make the comparison to the time steps passed for the continuation run.
• Similar to the get_times situation, the array coming from the Results.get_source_rates() will include more values than the number of time steps and will need to be modified to make a valid comparison to the continuation source rates.

[pshriwise]

Let's wrap this string or use a subset that is sufficient to detect the error message.

[lewisgross1296]

Let me know if the wrap looks good. For the subset method, do you just use an r-string and glob? I saw some other examples, so if you'd prefer, I can do something like this

    with pytest.raises(
        ValueError,
        match = "You are attempting to continue a run in which the previous *"
    ):
        bundle.solver(
            operator, [0.75, 0.5, 0.75], [1.0, 1.0, 1.0], continue_timesteps=True
        ).integrate()

[pshriwise]

Let's make this a docstring for the new test

[pshriwise]

Looks like it hit the ValueError on the test_continue that was supposed to use a correct continue run. Is this a case where the action-tmate would help us debug?

            if continue_timesteps:
                completed_timesteps = operator.prev_res.get_times()
                completed_source_rates = operator.prev_res.get_source_rates()
                num_previous_steps_run = len(completed_timesteps)
                if (
                    np.array_equal(completed_timesteps, timesteps[: num_previous_steps_run]) and
                    np.array_equal(completed_source_rates, source_rates[: num_previous_steps_run])
                ):
                    seconds = seconds[num_previous_steps_run:]
                    source_rates = source_rates[num_previous_steps_run:]
                else:
>                   raise ValueError(
                        "You are attempting to continue a run in which the previous results "
                        "do not have the same initial steps as those provided to the "
                        "Integrator. Please make sure you are using the correct timesteps, "
                        "powers or power densities, and previous results file."
                    )
E                   ValueError: You are attempting to continue a run in which the previous results do not have the same initial steps as those provided to the Integrator. Please make sure you are using the correct timesteps, powers or power densities, and previous results file.

Possibly yes, but this is an error that can be reproduced locally easily enough.

[lewisgross1296]

Turn out some of the tests were triggering the ValueError for one of the two checks in the if statement. As suggested by @pshriwise, I broke them up into two blocks: one to check consistent timesteps and one to check consistent source rates.

Here's my local version of the code

            # validate existing depletion steps are consistent with those passed to operator
            if continue_timesteps:
                completed_times = operator.prev_res.get_times(time_units=timestep_units)
                completed_timesteps = completed_times[1:] - completed_times[:-1] # convert absolute t to dt
                completed_source_rates = operator.prev_res.get_source_rates()
                num_previous_steps_run = len(completed_timesteps)
                if (np.array_equal(completed_timesteps, timesteps[:num_previous_steps_run])):
                    seconds = seconds[num_previous_steps_run:]
                else:
                    print("timestep check failed")
                    print(f"timesteps given to continue run: {timesteps}")
                    print(f"found completed timesteps: {completed_timesteps}")
                    raise ValueError(
                        "You are attempting to continue a run in which the previous timesteps "
                        "do not have the same initial timesteps as those provided to the "
                        "Integrator. Please make sure you are using the correct timesteps."
                    )
                if(np.array_equal(completed_source_rates, source_rates[:num_previous_steps_run])):
                    source_rates = source_rates[num_previous_steps_run:]
                else:
                    print("source rate check failed")
                    print(f"SRs given to coninute run: {source_rates}")
                    print(f"completed SRs: {completed_source_rates}")
                    raise ValueError(
                        "You are attempting to continue a run in which the previous results "
                        "do not have the same initial source rates, powers, or power densities "
                        "as those provided to the Integrator. Please make sure you are using "
                        "the correct powers, power densities, or source rates and previous results file."
                    )

I've added some print statements to get an idea of what the important variables look like during tests but I'm not totally sure why some tests are failing despite printing values that seem to be the same?

test_deplete_continue.py::test_continue[predictor] [openmc.deplete] t=0.0 s, dt=0.75 s, source=1.0
[openmc.deplete] t=0.75 (final operator evaluation)
source rate check failed
SRs given to coninute run: [1.0, 1.0]
completed SRs: [1. 1.]
FAILED
test_deplete_continue.py::test_continue[cecm] [openmc.deplete] t=0.0 s, dt=0.75 s, source=1.0
[openmc.deplete] t=0.75 (final operator evaluation)
source rate check failed
SRs given to coninute run: [1.0, 1.0]
completed SRs: [1. 1.]
FAILED
test_deplete_continue.py::test_continue[celi] [openmc.deplete] t=0.0 s, dt=0.75 s, source=1.0
[openmc.deplete] t=0.75 (final operator evaluation)
source rate check failed
SRs given to coninute run: [1.0, 1.0]
completed SRs: [1. 1.]
FAILED
test_deplete_continue.py::test_continue[cf4] [openmc.deplete] t=0.0 s, dt=0.75 s, source=1.0
[openmc.deplete] t=0.75 (final operator evaluation)
source rate check failed
SRs given to coninute run: [1.0, 1.0]
completed SRs: [1. 1.]
FAILED
test_deplete_continue.py::test_continue[epc_rk4] [openmc.deplete] t=0.0 s, dt=0.75 s, source=1.0
[openmc.deplete] t=0.75 (final operator evaluation)
source rate check failed
SRs given to coninute run: [1.0, 1.0]
completed SRs: [1. 1.]
FAILED
test_deplete_continue.py::test_continue[leqi] [openmc.deplete] t=0.0 s, dt=0.75 s, source=1.0
[openmc.deplete] t=0.75 (final operator evaluation)
source rate check failed
SRs given to coninute run: [1.0, 1.0]
completed SRs: [1. 1.]
FAILED
test_deplete_continue.py::test_continue[si_leqi] [openmc.deplete] t=0.0 s, dt=0.75 s, source=1.0
source rate check failed
SRs given to coninute run: [1.0, 1.0]
completed SRs: [1. 1.]
FAILED
test_deplete_continue.py::test_continue[si_celi] [openmc.deplete] t=0.0 s, dt=0.75 s, source=1.0
source rate check failed
SRs given to coninute run: [1.0, 1.0]
completed SRs: [1. 1.]
FAILED
test_deplete_continue.py::test_mismatched_initial_times[predictor] [openmc.deplete] t=0.0 s, dt=0.75 s, source=1.0
[openmc.deplete] t=0.75 s, dt=0.75 s, source=1.0
[openmc.deplete] t=1.5 (final operator evaluation)
timestep check failed
timesteps given to continue run: [0.75, 0.5, 0.75]
found completed timesteps: [0.75 0.75]

I tested the numpy.array_equal function for 1. and 1.0

lgross@ulam:~/openmc (auto_continue_depletion) $ python
Python 3.11.11 (main, Jan 13 2025, 12:51:04) [GCC 13.3.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import numpy
>>> print(numpy.array_equal([1.],[1.0]))
True

I know we discussed some schemes might have more data in them than the Integrator has but this doesn't seem to be happening in the above? Maybe I'm misinterpreting

[lewisgross1296]

As discussed today, I printed the non-sliced versions of the data above which clears up why the arrays appeared the same. E.g. we should actually do this to test

                if(np.array_equal(completed_source_rates, source_rates[:num_previous_steps_run])):
                    source_rates = source_rates[num_previous_steps_run:]
                else:
                    print("source rate check failed")
                    print(f"Initial SRs given to coninute run: {source_rates[:num_previous_steps_run]}")
                    print(f"completed SRs: {completed_source_rates}")

When I print what's actually compared in the if statements, we see

test_deplete_continue.py::test_continue[predictor] [openmc.deplete] t=0.0 s, dt=0.75 s, source=1.0
[openmc.deplete] t=0.75 (final operator evaluation)
source rate check failed
SRs (sliced) given to coninute run: [1.0]
completed SRs: [1. 1.]
FAILED
test_deplete_continue.py::test_continue[cecm] [openmc.deplete] t=0.0 s, dt=0.75 s, source=1.0
[openmc.deplete] t=0.75 (final operator evaluation)
source rate check failed
SRs (sliced) given to coninute run: [1.0]
completed SRs: [1. 1.]
FAILED
test_deplete_continue.py::test_continue[celi] [openmc.deplete] t=0.0 s, dt=0.75 s, source=1.0
[openmc.deplete] t=0.75 (final operator evaluation)
source rate check failed
SRs (sliced) given to coninute run: [1.0]
completed SRs: [1. 1.]
FAILED
test_deplete_continue.py::test_continue[cf4] [openmc.deplete] t=0.0 s, dt=0.75 s, source=1.0
[openmc.deplete] t=0.75 (final operator evaluation)
source rate check failed
SRs (sliced) given to coninute run: [1.0]
completed SRs: [1. 1.]
FAILED
test_deplete_continue.py::test_continue[epc_rk4] [openmc.deplete] t=0.0 s, dt=0.75 s, source=1.0
[openmc.deplete] t=0.75 (final operator evaluation)
source rate check failed
SRs (sliced) given to coninute run: [1.0]
completed SRs: [1. 1.]
FAILED
test_deplete_continue.py::test_continue[leqi] [openmc.deplete] t=0.0 s, dt=0.75 s, source=1.0
[openmc.deplete] t=0.75 (final operator evaluation)
source rate check failed
SRs (sliced) given to coninute run: [1.0]
completed SRs: [1. 1.]
FAILED
test_deplete_continue.py::test_continue[si_leqi] [openmc.deplete] t=0.0 s, dt=0.75 s, source=1.0
source rate check failed
SRs (sliced) given to coninute run: [1.0]
completed SRs: [1. 1.]
FAILED
test_deplete_continue.py::test_continue[si_celi] [openmc.deplete] t=0.0 s, dt=0.75 s, source=1.0
source rate check failed
SRs (sliced) given to coninute run: [1.0]
completed SRs: [1. 1.]
FAILED

This is due to the fact that certain integrators with sub-steps write out more source rates than are initially provided. I will work on looking into how to capture this case and test appropriately.