Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Changes to _eval_step_func #2895

Merged
merged 11 commits into from
Sep 5, 2024
20 changes: 15 additions & 5 deletions python/simulation.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,6 +100,14 @@ def get_num_args(func):
)


def get_name_args(func):
return (
("sim", "todo")
if isinstance(func, Harminv) or isinstance(func, PadeDFT)
else func.__code__.co_varnames[: func.__code__.co_argcount]
)
kombatEldridge marked this conversation as resolved.
Show resolved Hide resolved


def vec(*args):
try:
# Check for vec(x, [y, [z]])
Expand Down Expand Up @@ -4996,14 +5004,16 @@ def _combine(sim, todo):

def _eval_step_func(sim, func, todo):
num_args = get_num_args(func)
name_args = get_name_args(func)
self_count = int("self" in name_args)
kombatEldridge marked this conversation as resolved.
Show resolved Hide resolved

if num_args != 1 and num_args != 2:
raise ValueError(f"Step function '{func.__name__}'' requires 1 or 2 arguments")
elif num_args == 1:
if num_args not in {1 + self_count, 2 + self_count}:
raise ValueError(f"Step function '{func.__name__}' requires 1 or 2 arguments")
elif num_args == 2 + self_count:
func(sim, todo)
elif num_args == 1 + self_count:
if todo == "step":
func(sim)
elif num_args == 2:
func(sim, todo)


def _when_true_funcs(cond, *step_funcs):
Expand Down