Last active
April 13, 2022 05:31
-
-
Save mmechtley/5633015 to your computer and use it in GitHub Desktop.
Python Metaprogramming: An example of how to use Abstract Syntax Trees (ast module) in Python to turn bare expressions (not bound to any target variable names) in an end-user configuration file into elements of a list.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
""" | |
This file is made by the user, to set up model components for a Markov Chain | |
Monte Carlo galaxy simulator. | |
For simplicity, I didn't want to force the user to bind everything to variables, | |
or follow some particular naming scheme, or remember to add them to a list, etc. | |
YAML or such seemed like overkill, since python keyword arguments already make it | |
pretty human-readable | |
""" | |
from psfMC.ModelComponents import Sky, PSF, Sersic | |
from psfMC.distributions import Normal, Uniform | |
# First component (modeling the brightness of the background sky) | |
Sky(adu=Normal(mu=0, tau=100)) | |
# Point source component | |
PSF(xy=Uniform(lower=(60, 60), upper=(70, 70)), | |
mag=Uniform(lower=18, upper=24)) | |
# Sersic profile, modeling a galaxy under the point source | |
Sersic(xy=Uniform(lower=(60, 60), upper=(70, 70)), | |
mag=Uniform(lower=22, upper=27.5), | |
reff=Uniform(lower=1.0, upper=12.0), | |
index=Uniform(lower=0.5, upper=8), | |
axis_ratio=Uniform(lower=0.25, upper=1.0), | |
angle=Uniform(lower=0, upper=360)) | |
# A second sersic profile, modeling a fainter nearby galaxy | |
Sersic(xy=Uniform(lower=(42, 82), upper=(50, 90)), | |
mag=Uniform(lower=22, upper=27.5), | |
reff=Uniform(lower=1.0, upper=4.0), | |
index=Uniform(lower=0.5, upper=8.0), | |
axis_ratio=Uniform(lower=0.7, upper=1.0), | |
angle=Uniform(lower=0, upper=360)) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from ast import * | |
_comps_name = 'components' | |
class ExprsToAssigns(NodeTransformer): | |
""" | |
Walks the Abstract Syntax Tree from the parsed model file, and transforms | |
bare expressions into augmented assignments that are tacked on to the end | |
of the components list, ie: | |
Sky(...) | |
becomes: | |
components += [Sky(...)] | |
""" | |
def visit_Expr(self, node): | |
return copy_location(AugAssign( | |
target=Name(id=_comps_name, ctx=Store()), | |
op=Add(), | |
value=List(elts=[node.value], ctx=Load()) | |
), node) | |
def component_list_from_file(filename): | |
""" | |
Read in a file containing model components and return them as a list | |
""" | |
with open(filename) as f: | |
model_tree = parse(f.read()) | |
# Inject distribution and component imports. Put them at the beginning | |
# so user imports may override them. level=1 means relative import, e.g.: | |
# from .ModelComponents import * | |
increment_lineno(model_tree, n=3) | |
comps = ImportFrom(module='ModelComponents', | |
names=[alias(name='*', asname=None)], level=1) | |
dists = ImportFrom(module='distributions', | |
names=[alias(name='*', asname=None)], level=1) | |
model_tree.body.insert(0, comps) | |
model_tree.body.insert(1, dists) | |
# Insert a statement creating an empty list called components | |
comp_list_node = Assign(targets=[Name(id=_comps_name, ctx=Store())], | |
value=List(elts=[], ctx=Load())) | |
model_tree.body.insert(2, comp_list_node) | |
# Transform bare components expressions into list append statements | |
model_tree = ExprsToAssigns().visit(model_tree) | |
fix_missing_locations(model_tree) | |
exec(compile(model_tree, filename, mode='exec')) | |
return locals()[_comps_name] |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment