def _dump_dist_idxs_cmd(options):
"""
Return the post_setup hook function for 'openmdao dump_idxs'.
Parameters
----------
options : argparse Namespace
Command line options.
Returns
-------
function
The hook function.
"""
if options.outfile is None:
out = sys.stdout
else:
out = open(options.outfile, 'w')
def _dumpdist(prob):
dump_dist_idxs(prob, vec_name=options.vecname, stream=out)
exit()
hooks._register_hook('final_setup', 'Problem', post=_dumpdist)
return _dumpdist
def _view_dyn_shapes_cmd(options, user_args):
"""
Return the post_setup hook function for 'openmdao view_dyn_shapes'.
Parameters
----------
options : argparse Namespace
Command line options.
user_args : list of str
Args to be passed to the user script.
"""
def _view_shape_graph(model):
view_dyn_shapes(model,
outfile=options.outfile,
show=not options.no_display,
title=options.title)
exit()
def _set_dyn_hook(prob):
# we can't wait until the end of Problem.setup because we'll die in _setup_sizes
# if there were any unresolved dynamic shapes, so put the hook immediately after
# _setup_dynamic_shapes. inst_id is None here because no system's pathname will
# have been set at the time this hook is triggered.
hooks._register_hook('_setup_dynamic_shapes',
class_name='Group',
inst_id=None,
post=_view_shape_graph)
hooks._setup_hooks(prob.model)
# register the hooks
hooks._register_hook('setup', 'Problem', pre=_set_dyn_hook)
ignore_errors(True)
_load_and_exec(options.file[0], user_args)
def _n2_cmd(options, user_args):
"""
Process command line args and call n2 on the specified file.
Parameters
----------
options : argparse Namespace
Command line options.
user_args : list of str
Command line options after '--' (if any). Passed to user script.
"""
filename = _to_filename(options.file[0])
if filename.endswith('.py'):
# the file is a python script, run as a post_setup hook
def _noraise(prob):
prob.model._raise_connection_errors = False
if options.use_declare_partial_info:
warn_deprecation("'--use_declare_partial_info' is now the"
" default and the option is ignored.")
def _viewmod(prob):
n2(prob, outfile=options.outfile, show_browser=not options.no_browser,
title=options.title, embeddable=options.embeddable)
hooks._register_hook('setup', 'Problem', pre=_noraise, ncalls=1)
hooks._register_hook('final_setup', 'Problem', post=_viewmod, exit=True)
ignore_errors(True)
_load_and_exec(options.file[0], user_args)
else:
# assume the file is a recording, run standalone
n2(filename, outfile=options.outfile, title=options.title,
show_browser=not options.no_browser, embeddable=options.embeddable)
def _cite_cmd(options, user_args):
"""
Run the `openmdao cite` command.
Parameters
----------
options : argparse Namespace
Command line options.
user_args : list of str
Args to be passed to the user script.
"""
if options.outfile is None:
out = sys.stdout
else:
out = open(options.outfile, 'w')
if not options.classes:
options.classes = None
def _cite(prob):
if not MPI or MPI.COMM_WORLD.rank == 0:
print_citations(prob, classes=options.classes, out_stream=out)
exit()
hooks._register_hook('setup', 'Problem', post=_cite)
ignore_errors(True)
_load_and_exec(options.file[0], user_args)
def _timing_cmd(options, user_args):
"""
Implement the 'openmdao timing' command.
Parameters
----------
options : argparse Namespace
Command line options.
user_args : list of str
Args to be passed to the user script.
"""
if not options.funcs:
options.funcs = _default_timer_methods.copy()
filename = _to_filename(options.file[0])
if filename.endswith('.py'):
hooks._register_hook('setup',
'Problem',
pre=partial(_set_timer_setup_hook, options))
# register an atexit function to write out all of the timing data
atexit.register(partial(_postprocess, options))
with timing_context(not options.use_context):
_load_and_exec(options.file[0], user_args)
else: # assume file is a pickle file
if options.use_context:
issue_warning(
f"Since given file '{options.file[0]}' is not a python script, the "
"'--use_context' option is ignored.")
_show_view(options.file[0], options)
def _cite_cmd(options):
"""
Return the post setup hook function for `openmdao cite`.
Parameters
----------
options : argparse Namespace
Command line options.
Returns
-------
function
The hook function.
"""
if options.outfile is None:
out = sys.stdout
else:
out = open(options.outfile, 'w')
if not options.classes:
options.classes = None
def _cite(prob):
if not MPI or MPI.COMM_WORLD.rank == 0:
print_citations(prob, classes=options.classes, out_stream=out)
exit()
hooks._register_hook('setup', 'Problem', post=_cite)
return _cite
def _n2_cmd(options, user_args):
"""
Process command line args and call n2 on the specified file.
Parameters
----------
options : argparse Namespace
Command line options.
user_args : list of str
Command line options after '--' (if any). Passed to user script.
"""
filename = options.file[0]
if filename.endswith('.py'):
# the file is a python script, run as a post_setup hook
def _viewmod(prob):
n2(prob, outfile=options.outfile, show_browser=not options.no_browser,
title=options.title, embeddable=options.embeddable,
use_declare_partial_info=options.use_declare_partial_info)
exit() # could make this command line selectable later
options.func = lambda options: _viewmod
hooks._register_hook('final_setup', 'Problem', pre=_viewmod)
_simple_exec(options, user_args)
else:
# assume the file is a recording, run standalone
n2(filename, outfile=options.outfile, title=options.title,
show_browser=not options.no_browser, embeddable=options.embeddable,
use_declare_partial_info=options.use_declare_partial_info)
def _view_connections_cmd(options):
"""
Return the post_setup hook function for 'openmdao view_connections'.
Parameters
----------
options : argparse Namespace
Command line options.
Returns
-------
function
The hook function.
"""
def _viewconns(prob):
if options.title:
title = options.title
else:
title = "Connections for %s" % os.path.basename(options.file[0])
view_connections(prob, outfile=options.outfile, show_browser=not options.no_browser,
show_values=options.show_values, title=title)
exit()
# register the hook
if options.show_values:
funcname = 'final_setup'
else:
funcname = 'setup'
hooks._register_hook(funcname, class_name='Problem', inst_id=options.problem, post=_viewconns)
return _viewconns
def test_ncalls(self):
hooks._register_hook('final_setup',
'Problem',
pre=make_hook('pre_final'),
post=make_hook('post_final'),
ncalls=2)
hooks._register_hook('final_setup',
'Problem',
pre=make_hook('pre_final2'),
post=make_hook('post_final2'))
prob = self.build_model()
prob.run_model()
prob.run_model()
prob.run_model()
self.assertEqual(prob.calls, [
'pre_final',
'pre_final2',
'post_final',
'post_final2',
'pre_final',
'pre_final2',
'post_final',
'post_final2',
'pre_final2',
'post_final2',
])
def _view_connections_cmd(options, user_args):
"""
Return the post_setup hook function for 'openmdao view_connections'.
Parameters
----------
options : argparse Namespace
Command line options.
user_args : list of str
Args to be passed to the user script.
"""
def _viewconns(prob):
if options.title:
title = options.title
else:
title = "Connections for %s" % os.path.basename(options.file[0])
view_connections(prob,
outfile=options.outfile,
show_browser=not options.no_browser,
show_values=options.show_values,
title=title)
exit()
# register the hook
if options.show_values:
funcname = 'final_setup'
else:
funcname = 'setup'
hooks._register_hook(funcname,
class_name='Problem',
inst_id=options.problem,
post=_viewconns)
ignore_errors(True)
_load_and_exec(options.file[0], user_args)
def _set_dyn_hook(prob):
# we can't wait until the end of Problem.setup because we'll die in _setup_sizes
# if there were any unresolved dynamic shapes, so put the hook immediately after
# _setup_dynamic_shapes. inst_id is None here because no system's pathname will
# have been set at the time this hook is triggered.
hooks._register_hook('_setup_dynamic_shapes', class_name='Group', inst_id=None,
post=_view_shape_graph)
hooks._setup_hooks(prob.model)
def register_report(name,
func,
desc,
class_name,
method,
pre_or_post,
report_filename,
inst_id=None):
"""
Register a report with the reporting system.
Parameters
----------
name : str
Name of report. Report names must be unique across all reports.
func : function
A function to do the reporting. Expects the first argument to be an instance of class_name.
desc : str
A description of the report.
class_name : str
The name of the class owning the method where the report will be run.
method : str
In which method of class_name should this be run.
pre_or_post : str
Valid values are 'pre' and 'post'. Indicates when to run the report in the method.
report_filename : str
Name of file to use when saving the report.
inst_id : str or None
Either the instance ID of an OpenMDAO object (e.g. Problem, Driver) or None.
If None, then this report will be run for all objects of type class_name.
"""
global _reports_registry
report = _Report(func, desc, class_name, inst_id, method, pre_or_post,
report_filename)
if name in _reports_registry:
raise ValueError(f"A report with the name {name} already exists")
_reports_registry[name] = report
if pre_or_post == 'pre':
_register_hook(method,
class_name,
pre=func,
inst_id=inst_id,
report_filename=report_filename)
elif pre_or_post == 'post':
_register_hook(method,
class_name,
post=func,
inst_id=inst_id,
report_filename=report_filename)
else:
raise ValueError(
f"The argument 'pre_or_post' can only have values of 'pre' or 'post', but {pre_or_post}"
" was given")
return
def _xdsm_cmd(options, user_args):
"""
Process command line args and call xdsm on the specified file.
Parameters
----------
options : argparse Namespace
Command line options.
user_args : list of str
Command line options after '--' (if any). Passed to user script.
"""
filename = _to_filename(options.file[0])
kwargs = {}
for name in ['box_stacking', 'box_width', 'box_lines', 'numbered_comps', 'number_alignment']:
val = getattr(options, name)
if val is not None:
kwargs[name] = val
if filename.endswith('.py'):
# the file is a python script, run as a post_setup hook
def _xdsm(prob):
write_xdsm(prob, filename=options.outfile, model_path=options.model_path,
recurse=options.recurse,
include_external_outputs=not options.no_extern_outputs,
out_format=options.format,
include_solver=options.include_solver, subs=_CHAR_SUBS,
show_browser=not options.no_browser, show_parallel=not options.no_parallel,
add_process_conns=not options.no_process_conns,
output_side=options.output_side,
legend=options.legend,
class_names=options.class_names,
equations=options.equations,
include_indepvarcomps=not options.no_indepvarcomps,
**kwargs)
exit()
hooks._register_hook('setup', 'Problem', post=_xdsm)
_load_and_exec(options.file[0], user_args)
else:
# assume the file is a recording, run standalone
write_xdsm(options.file[0], filename=options.outfile, model_path=options.model_path,
recurse=options.recurse,
include_external_outputs=not options.no_extern_outputs,
out_format=options.format,
include_solver=options.include_solver, subs=_CHAR_SUBS,
show_browser=not options.no_browser, show_parallel=not options.no_parallel,
add_process_conns=not options.no_process_conns, output_side=options.output_side,
legend=options.legend,
class_names=options.class_names,
equations=options.equations,
**kwargs)
def _scaling_cmd(options, user_args):
"""
Return the post_setup hook function for 'openmdao driver_scaling'.
Parameters
----------
options : argparse Namespace
Command line options.
user_args : list of str
Args to be passed to the user script.
"""
def _set_flag(problem):
global _run_driver_called
_run_driver_called = True
def _scaling_check(problem):
if _run_driver_called:
# If run_driver has been called, we know no more user changes are coming.
_scaling(problem)
def _scaling(problem):
hooks._unregister_hook('final_setup',
'Problem') # avoid recursive loop
hooks._unregister_hook('run_driver', 'Problem')
driver = problem.driver
if options.title:
title = options.title
else:
title = "Driver scaling for %s" % os.path.basename(options.file[0])
view_driver_scaling(driver,
outfile=options.outfile,
show_browser=not options.no_browser,
title=title,
jac=not options.nojac)
exit()
# register the hooks
hooks._register_hook('final_setup',
class_name='Problem',
inst_id=options.problem,
post=_scaling_check)
hooks._register_hook('run_driver',
class_name='Problem',
inst_id=options.problem,
pre=_set_flag)
# register an atexit function to check if scaling report was triggered during the script
import atexit
atexit.register(_exitfunc)
ignore_errors(True)
_load_and_exec(options.file[0], user_args)
def _set_timer_setup_hook(options, problem):
# This just sets a hook into the top level system of the model after we know it exists.
# Note that this means that no timings can happen until AFTER _setup_procs is done.
global _timing_managers
inst_id = problem._get_inst_id()
if inst_id not in _timing_managers:
_timing_managers[inst_id] = TimingManager(options)
hooks._register_hook('_setup_procs',
'System',
inst_id='',
post=partial(_setup_timers, options))
hooks._setup_hooks(problem.model)
def upload_vars_init_cmd(self, py_filename, options):
def upload_vars_init(prob):
self.upload_vars_init(prob, options)
sys.exit()
d = os.path.dirname(py_filename)
run_analysis_filename = os.path.join(d, "run_analysis.py")
if not os.path.exists(run_analysis_filename):
error(
f"Can not get analysis init: script {run_analysis_filename} not found."
)
hooks.use_hooks = True
hooks._register_hook("final_setup", "Problem", post=upload_vars_init)
_load_and_exec(run_analysis_filename, [])
def _config_summary_cmd(options, user_args):
"""
Return the post_setup hook function for 'openmdao summary'.
Parameters
----------
options : argparse Namespace
Command line options.
user_args : list of str
Args to be passed to the user script.
"""
hooks._register_hook('final_setup', 'Problem', post=config_summary, exit=True)
ignore_errors(True)
_load_and_exec(options.file[0], user_args)
def _tree_cmd(options):
"""
Return the post_setup hook function for 'openmdao tree'.
Parameters
----------
options : argparse Namespace
Command line options.
Returns
-------
function
The hook function.
"""
if options.outfile is None:
out = sys.stdout
else:
out = open(options.outfile, 'w')
if options.attrs or options.vecvars:
filt = _get_tree_filter(options.attrs, options.vecvars)
else:
filt = None
def _tree(prob):
tree(prob,
show_colors=options.show_colors,
show_sizes=options.show_sizes,
show_approx=options.show_approx,
filter=filt,
max_depth=options.depth,
rank=options.rank,
stream=out)
exit()
# register the hook
if options.vecvars or options.show_sizes or options.show_approx:
funcname = 'final_setup'
else:
funcname = 'setup'
hooks._register_hook(funcname,
class_name='Problem',
inst_id=options.problem,
post=_tree)
return _tree
def _check_config_cmd(options, user_args):
"""
Return the post_setup hook function for 'openmdao check'.
Parameters
----------
options : argparse Namespace
Command line options.
user_args : list of str
Args to be passed to the user script.
Returns
-------
function
The post-setup hook function.
"""
def _check_config(prob):
if not MPI or MPI.COMM_WORLD.rank == 0:
if options.outfile is None:
logger = get_logger('check_config',
out_stream='stdout',
out_file=None,
use_format=True)
else:
logger = get_logger('check_config',
out_file=options.outfile,
use_format=True)
if not options.checks:
options.checks = sorted(_default_checks)
elif 'all' in options.checks:
options.checks = sorted(_all_checks)
prob.check_config(logger, options.checks)
exit()
# register the hook
_register_hook('final_setup',
class_name='Problem',
inst_id=options.problem,
post=_check_config)
ignore_errors(True)
_load_and_exec(options.file[0], user_args)
def _tree_cmd(options, user_args):
"""
Return the post_setup hook function for 'openmdao tree'.
Parameters
----------
options : argparse Namespace
Command line options.
user_args : list of str
Args to be passed to the user script.
"""
if options.outfile is None:
out = sys.stdout
else:
out = open(options.outfile, 'w')
if options.attrs or options.vecvars:
filt = _get_tree_filter(options.attrs, options.vecvars)
else:
filt = None
def _tree(prob):
tree(prob,
show_colors=options.show_colors,
show_sizes=options.show_sizes,
show_approx=options.show_approx,
filter=filt,
max_depth=options.depth,
rank=options.rank,
stream=out)
exit()
# register the hook
if options.vecvars or options.show_sizes or options.show_approx:
funcname = 'final_setup'
else:
funcname = 'setup'
hooks._register_hook(funcname,
class_name='Problem',
inst_id=options.problem,
post=_tree)
ignore_errors(True)
_load_and_exec(options.file[0], user_args)
def push_mda_cmd(self, py_filename, options):
def push_mda(prob):
name = options["--name"]
pbname = prob.model.__class__.__name__
if name and pbname != name:
info("Analysis %s skipped" % pbname)
# do not exit seeking for another problem (ie analysis)
else:
options["--pyfilename"] = py_filename
xdsm = self.push_mda(prob, options)
if options.get("--xdsm"): # show command
# required to interrupt pb execution
raise AnalysisPushedException(xdsm=xdsm)
else:
sys.exit()
hooks.use_hooks = True
hooks._register_hook("final_setup", "Problem", post=push_mda)
_load_and_exec(py_filename, [])
return push_mda
def _config_summary_cmd(options):
"""
Return the post_setup hook function for 'openmdao summary'.
Parameters
----------
options : argparse Namespace
Command line options.
Returns
-------
function
The hook function.
"""
def summary(prob):
config_summary(prob)
sys.exit(0)
hooks._register_hook('final_setup', 'Problem', post=summary)
return summary
def test_exit(self):
hooks._register_hook('final_setup',
'Problem',
pre=make_hook('pre_final'),
post=make_hook('post_final'))
hooks._register_hook('final_setup',
'Problem',
pre=make_hook('pre_final2'),
post=make_hook('post_final2'),
exit=True)
prob = self.build_model()
try:
prob.run_model()
prob.run_model()
prob.run_model()
except SystemExit:
self.assertEqual(
prob.calls,
['pre_final', 'pre_final2', 'post_final', 'post_final2'])
else:
self.fail("sys.exit() was not called")
def test_problem_hooks_kwargs(self):
x0 = 33.0
y0 = 44.0
def set_prob_vars_hook_func(prob, **kwargs):
if 'x0' in kwargs:
prob['p1.x'] = kwargs['x0']
if 'y0' in kwargs:
prob['p2.y'] = kwargs['y0']
hooks._register_hook('final_setup',
'Problem',
pre=set_prob_vars_hook_func,
x0=x0,
y0=y0)
prob = self.build_model()
prob.run_model()
self.assertEqual(prob['comp.x'], x0)
self.assertEqual(prob['comp.y'], y0)
def _dist_idxs_exec(options, user_args):
"""
Return the post_setup hook function for 'openmdao dump_idxs'.
Parameters
----------
options : argparse Namespace
Command line options.
user_args : list of str
Args to be passed to the user script.
"""
if options.outfile is None:
out = sys.stdout
else:
out = open(options.outfile, 'w')
def _dumpdist(prob):
dump_dist_idxs(prob, full=options.full, stream=out)
exit()
_register_hook('final_setup', 'Problem', post=_dumpdist)
ignore_errors(True)
_load_and_exec(options.file[0], user_args)
def test_multiwrap(self):
pre_final = make_hook('pre_final')
post_final = make_hook('post_final')
hooks._register_hook('final_setup',
'Problem',
pre=pre_final,
post=post_final)
hooks._register_hook('final_setup',
'Problem',
pre=make_hook('pre_final2'),
post=make_hook('post_final2'))
prob = self.build_model()
prob.run_model()
prob.run_model()
prob.run_model()
self.assertEqual(prob.calls, [
'pre_final',
'pre_final2',
'post_final',
'post_final2',
'pre_final',
'pre_final2',
'post_final',
'post_final2',
'pre_final',
'pre_final2',
'post_final',
'post_final2',
])
hooks._unregister_hook('final_setup',
'Problem',
pre=pre_final,
post=False)
prob.calls = []
prob.run_model()
prob.run_model()
prob.run_model()
self.assertEqual(prob.calls, [
'pre_final2',
'post_final',
'post_final2',
'pre_final2',
'post_final',
'post_final2',
'pre_final2',
'post_final',
'post_final2',
])
hooks._unregister_hook('final_setup', 'Problem', pre=True, post=False)
prob.calls = []
prob.run_model()
prob.run_model()
prob.run_model()
self.assertEqual(prob.calls, [
'post_final',
'post_final2',
'post_final',
'post_final2',
'post_final',
'post_final2',
])
def _meta_model_cmd(options, user_args):
"""
Return the post_setup hook function for 'openmdao meta_model'.
Parameters
----------
options : argparse Namespace
Command line options.
user_args : list of str
Args to be passed to the user script.
"""
def _view_metamodel(prob):
if bokeh is None:
print(
"bokeh must be installed to view a MetaModel. Use the command:\n",
" pip install bokeh")
exit()
hooks._unregister_hook('final_setup', 'Problem')
mm_types = (MetaModelStructuredComp, MetaModelUnStructuredComp)
pathname = options.pathname
port_number = options.port_number
resolution = options.resolution
browser = options.browser
if pathname:
comp = prob.model._get_subsystem(pathname)
if comp and isinstance(comp, mm_types):
view_metamodel(comp, resolution, port_number, browser)
exit()
else:
comp = None
metamodels = {
mm.pathname: mm
for mm in prob.model.system_iter(include_self=True, typ=mm_types)
}
mm_names = list(metamodels.keys())
mm_count = len(mm_names)
if mm_count == 0:
print("No Metamodel components found in model.")
elif mm_count == 1 and not pathname:
comp = metamodels[mm_names[0]]
view_metamodel(comp, resolution, port_number, browser)
else:
try_str = "Try one of the following: {}.".format(mm_names)
if not pathname:
print("\nMetamodel not specified. {}".format(try_str))
elif not comp:
print("\nMetamodel '{}' not found.\n {}".format(
pathname, try_str))
else:
print("\n'{}' is not a Metamodel.\n {}".format(
pathname, try_str))
exit()
hooks._register_hook('final_setup', 'Problem', post=_view_metamodel)
_load_and_exec(options.file[0], user_args)
def test_problem_hooks(self):
def make_hook(name):
def hook_func(prob):
prob.calls.append(name)
return hook_func
hooks.use_hooks = True
hooks._register_hook('setup', 'Problem', pre=make_hook('pre_setup'), post=make_hook('post_setup'))
hooks._register_hook('final_setup', 'Problem', pre=make_hook('pre_final'), post=make_hook('post_final'))
hooks._register_hook('run_model', 'Problem', pre=make_hook('pre_run_model'), post=make_hook('post_run_model'))
try:
prob = om.Problem()
prob.calls = []
model = prob.model
model.add_subsystem('p1', om.IndepVarComp('x', 3.0))
model.add_subsystem('p2', om.IndepVarComp('y', -4.0))
model.add_subsystem('comp', om.ExecComp("f_xy=2.0*x+3.0*y"))
model.connect('p1.x', 'comp.x')
model.connect('p2.y', 'comp.y')
prob.setup()
prob.run_model()
prob.run_model()
prob.run_model()
self.assertEqual(prob.calls, ['pre_setup', 'post_setup',
'pre_run_model', 'pre_final', 'post_final', 'post_run_model',
'pre_run_model', 'pre_final', 'post_final', 'post_run_model',
'pre_run_model', 'pre_final', 'post_final', 'post_run_model',
])
np.testing.assert_allclose(prob['comp.f_xy'], -6.0)
hooks._unregister_hook('setup', 'Problem', pre=False)
hooks._unregister_hook('final_setup', 'Problem')
hooks._unregister_hook('run_model', 'Problem', post=False)
prob.calls = []
prob.setup()
prob.run_model()
prob.run_model()
prob.run_model()
self.assertEqual(prob.calls, ['pre_setup', 'post_run_model', 'post_run_model', 'post_run_model'])
hooks._unregister_hook('setup', 'Problem')
msg = "No hook found for method 'final_setup' for class 'Problem' and instance 'None'."
# already removed final_setup hooks earlier, so expect a warning here
with assert_warning(UserWarning, msg):
hooks._unregister_hook('final_setup', 'Problem')
hooks._unregister_hook('run_model', 'Problem')
prob.calls = []
prob.setup()
prob.run_model()
prob.run_model()
self.assertEqual(prob.calls, [])
self.assertEqual(len(hooks._hooks), 0) # should be no hooks left
finally:
hooks.use_hooks = False
hooks._reset_all_hooks()
def test_multiwrap_mixed_inst_None(self):
pre_final = make_hook('pre_final')
post_final = make_hook('post_final')
hooks._register_hook('final_setup',
'Problem',
inst_id='problem1',
ncalls=2,
pre=make_hook('pre_final1'),
post=make_hook('post_final1'))
hooks._register_hook('final_setup',
'Problem',
inst_id='problem2',
pre=make_hook('pre_final2'),
post=make_hook('post_final2'))
hooks._register_hook('final_setup',
'Problem',
pre=pre_final,
post=post_final)
probs = [self.build_model(f"problem{i+1}") for i in range(2)]
for prob in probs:
prob.calls = []
for i in range(3):
prob.run_model()
self.assertEqual(probs[0].calls, [
'pre_final1', 'pre_final', 'post_final1', 'post_final',
'pre_final1', 'pre_final', 'post_final1', 'post_final',
'pre_final', 'post_final'
])
self.assertEqual(probs[1].calls, [
'pre_final2', 'pre_final', 'post_final2', 'post_final',
'pre_final2', 'pre_final', 'post_final2', 'post_final',
'pre_final2', 'pre_final', 'post_final2', 'post_final'
])
hooks._unregister_hook('final_setup',
'Problem',
pre=pre_final,
post=False)
for prob in probs:
prob.calls = []
for i in range(3):
prob.run_model()
self.assertEqual(probs[0].calls,
['post_final', 'post_final', 'post_final'])
self.assertEqual(probs[1].calls, [
'pre_final2', 'post_final2', 'post_final', 'pre_final2',
'post_final2', 'post_final', 'pre_final2', 'post_final2',
'post_final'
])
hooks._unregister_hook('final_setup', 'Problem', pre=True, post=False)
for prob in probs:
prob.calls = []
for i in range(3):
prob.run_model()
self.assertEqual(probs[0].calls,
['post_final', 'post_final', 'post_final'])
self.assertEqual(probs[1].calls, [
'post_final2', 'post_final', 'post_final2', 'post_final',
'post_final2', 'post_final'
])
def test_problem_hooks(self):
hooks._register_hook('setup',
'Problem',
pre=make_hook('pre_setup'),
post=make_hook('post_setup'))
hooks._register_hook('final_setup',
'Problem',
pre=make_hook('pre_final'),
post=make_hook('post_final'))
hooks._register_hook('run_model',
'Problem',
pre=make_hook('pre_run_model'),
post=make_hook('post_run_model'))
prob = self.build_model()
prob.run_model()
prob.run_model()
prob.run_model()
self.assertEqual(prob.calls, [
'pre_setup',
'post_setup',
'pre_run_model',
'pre_final',
'post_final',
'post_run_model',
'pre_run_model',
'pre_final',
'post_final',
'post_run_model',
'pre_run_model',
'pre_final',
'post_final',
'post_run_model',
])
np.testing.assert_allclose(prob['comp.f_xy'], -6.0)
hooks._unregister_hook('setup', 'Problem', pre=False)
hooks._unregister_hook('final_setup', 'Problem')
hooks._unregister_hook('run_model', 'Problem', post=False)
prob.calls = []
prob.setup()
prob.run_model()
prob.run_model()
prob.run_model()
self.assertEqual(prob.calls, [
'pre_setup', 'post_run_model', 'post_run_model', 'post_run_model'
])
hooks._unregister_hook('setup', 'Problem')
msg = "No hook found for method 'final_setup' for class 'Problem' and instance 'None'."
# already removed final_setup hooks earlier, so expect a warning here
with assert_warning(UserWarning, msg):
hooks._unregister_hook('final_setup', 'Problem')
hooks._unregister_hook('run_model', 'Problem')
prob.calls = []
prob.setup()
prob.run_model()
prob.run_model()
self.assertEqual(prob.calls, [])
self.assertEqual(len(hooks._hooks), 0) # should be no hooks left
