Commit f488e64e authored by Jameson Rollins's avatar Jameson Rollins

Merge branch 'load-sub-budget' into 'master'

sub-budget plotting

Closes #70 and #82

See merge request !122
parents d23fb2ad 1c6a785a
Pipeline #193696 passed with stages
in 2 minutes and 24 seconds
......@@ -74,20 +74,25 @@ def load_module(name_or_path):
return mod, path
def load_budget(name_or_path, freq=None):
def load_budget(name_or_path, freq=None, bname=None):
"""Load GWINC Budget
Accepts either the name of a built-in canonical budget (see
gwinc.IFOS), the path to a budget package (directory) or module
(ending in .py), or the path to an IFO Struct definition file (see
gwinc.Struct).
gwinc.Struct). If an IFO Struct is specified, the base "aLIGO"
budget definition will be used.
If `bname` is specified the Budget class with that name will be
loaded from the budget module. Otherwise, the Budget class with
the same name as the budget module will be load.
If the budget is a package directory which includes an 'ifo.yaml'
file the ifo Struct will be loaded from that file and assigned to
the budget.ifo attribute. If a Struct definition file is provided
the base aLIGO budget definition will be assumed.
Returns an instantiated Budget object. If a frequency array or
Returns the instantiated Budget object. If a frequency array or
frequency specification string (see `freq_from_spec()`) is
provided, the budget will be instantiated with the provided array.
If a frequency array is not provided and the Budget class
......@@ -103,7 +108,7 @@ def load_budget(name_or_path, freq=None):
if os.path.exists(name_or_path):
path = name_or_path.rstrip('/')
bname, ext = os.path.splitext(os.path.basename(path))
base, ext = os.path.splitext(os.path.basename(path))
if ext in Struct.STRUCT_EXT:
logger.info("loading struct {}...".format(path))
......@@ -112,6 +117,7 @@ def load_budget(name_or_path, freq=None):
modname = 'gwinc.ifo.aLIGO'
else:
bname = bname or base
modname = path
else:
......@@ -120,10 +126,10 @@ def load_budget(name_or_path, freq=None):
name_or_path,
IFOS,
))
bname = name_or_path
bname = bname or name_or_path
modname = 'gwinc.ifo.'+name_or_path
logger.info("loading module {}...".format(modname))
logger.info(f"loading budget {modname}.{bname}...")
mod, modpath = load_module(modname)
Budget = getattr(mod, bname)
if freq is None:
......
......@@ -102,12 +102,21 @@ group.add_argument(
group.add_argument(
'--diff', '-d', metavar='IFO',
help="show difference table between IFO and another IFO description (name or path) and exit (budget not calculated)")
group.add_argument(
'--list', '-l', action='store_true',
help="list all elements of Budget (budget not calculated)")
parser.add_argument(
'--no-plot', '-np', action='store_false', dest='plot',
help="suppress plotting")
parser.add_argument(
'--bname', '-b',
help="name of top-level Budget class to load (defaults to IFO name)")
parser.add_argument(
'IFO',
help="IFO name or path")
parser.add_argument(
'subbudget', metavar='SUBBUDGET', nargs='?',
help="subbudget to plot; can be nested (e.g. 'Thermal.Substrate')")
def main():
......@@ -137,7 +146,7 @@ def main():
except IndexError:
parser.error(f"Improper frequency specification: {args.freq}")
try:
budget = load_budget(args.IFO, freq=freq)
budget = load_budget(args.IFO, freq=freq, bname=args.bname)
except RuntimeError as e:
parser.exit(2, f"Error: {e}\n")
name = budget.name
......@@ -178,6 +187,18 @@ def main():
ov = repr(p[2])
print(fmt.format(k, v, ov))
return
if args.list:
for i in budget.walk():
name = '.'.join([n.__class__.__name__ for n in i])
type = i[-1].__class__.__bases__[0].__name__
print(f'{name} ({type})')
return
if args.subbudget:
try:
budget[args.subbudget]
except KeyError:
parser.exit(3, f"Error: Unknown budget item '{args.subbudget}'.\n")
out_data_files = set()
out_plot_files = set()
......@@ -239,11 +260,6 @@ def main():
logger.info("calculating budget...")
trace = budget.run(freq=freq)
if args.title:
plot_style['title'] = args.title
else:
plot_style['title'] = "GWINC Noise Budget: {}".format(name)
if args.range:
logger.info("calculating inspiral ranges...")
metrics, H = inspiral_range.all_ranges(freq, trace.psd, **RANGE_PARAMS)
......@@ -263,6 +279,15 @@ def main():
else:
subtitle = None
if args.subbudget:
trace = trace[args.subbudget]
name += f': {args.subbudget}'
if args.title:
plot_style['title'] = args.title
else:
plot_style['title'] = "GWINC Noise Budget: {}".format(name)
##########
# interactive
......
......@@ -418,13 +418,33 @@ class Budget(Noise):
return name
def __getitem__(self, name):
"""Get a (possibly nested) sub-BudgetItem.
"""
try:
name, rest = name.split('.', 1)
except ValueError:
rest = None
try:
return self._noise_objs[name]
o = self._noise_objs[name]
except KeyError:
try:
return self._cal_objs[name]
o = self._cal_objs[name]
except KeyError:
raise KeyError("unknown noise or cal name '{}".format(name))
if rest:
return o[rest]
else:
return o
def get(self, key, default=None):
"""Get a (possibly nested) sub-BudgetItem.
"""
try:
return self[key]
except KeyError:
return default
def keys(self):
"""Iterate over budget noise names."""
......
......@@ -61,8 +61,23 @@ class BudgetTrace:
raise AttributeError
def __getitem__(self, name):
"""get budget trace by name"""
return self._bdict[name]
"""get budget trace by name
"""
try:
name, rest = name.split('.', 1)
return self._bdict[name][rest]
except ValueError:
return self._bdict[name]
def get(self, key, default=None):
"""get a (possibly nested) Trace item.
"""
try:
return self[key]
except KeyError:
return default
def items(self):
"""iterator of budget (name, trace) tuples"""
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment