From 1fec618582fa86b1b1df89c234840c0a47f7b9d3 Mon Sep 17 00:00:00 2001
From: Jameson Graef Rollins <jameson.rollins@ligo.org>
Date: Thu, 11 Feb 2021 12:11:16 -0800
Subject: [PATCH] allow specifying budget class to load
By default the budget class loaded will have the same name as the budget
module. This adds the ability to specify which budget class to load,
allowing for loading sub-budgets directly.
Hooks added to the CLI as well.
---
gwinc/__init__.py | 18 ++++++++++++------
gwinc/__main__.py | 5 ++++-
2 files changed, 16 insertions(+), 7 deletions(-)
diff --git a/gwinc/__init__.py b/gwinc/__init__.py
index 530aaf8d..7613c93e 100644
--- a/gwinc/__init__.py
+++ b/gwinc/__init__.py
@@ -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:
diff --git a/gwinc/__main__.py b/gwinc/__main__.py
index b70f8956..f74293b1 100644
--- a/gwinc/__main__.py
+++ b/gwinc/__main__.py
@@ -105,6 +105,9 @@ group.add_argument(
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")
@@ -137,7 +140,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
--
GitLab