Changes to _eval_step_func (#2895)

* changes to _eval_step_func * Update simulation.py * double to single quotes * incorrect if statement * reverted _eval_step_func, rm get_name_args, edited get_num_args * Update python/simulation.py * Update python/simulation.py * Update python/simulation.py * if statement readability update * added space from pre-check --------- Co-authored-by: Steven G. Johnson <[email protected]>
NanoComp · Sep 5, 2024 · 1f77440 · 1f77440
1 parent bfd134f
commit 1f77440
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/python/simulation.py b/python/simulation.py
@@ -3,6 +3,7 @@
 """
 
 import functools
+import inspect
 import math
 import numbers
 import os
@@ -93,11 +94,12 @@ def fix_dft_args(args, i):
 
 
 def get_num_args(func):
-    return (
-        2
-        if isinstance(func, Harminv) or isinstance(func, PadeDFT)
-        else func.__code__.co_argcount
-    )
+    if isinstance(func, Harminv) or isinstance(func, PadeDFT):
+        return 2
+    elif inspect.ismethod(func):
+        return func.__code__.co_argcount - 1  # remove 'self' from count
+    else:
+        return func.__code__.co_argcount
 
 
 def vec(*args):