Skip to content

Instantly share code, notes, and snippets.

@mmechtley
Last active April 13, 2022 05:31
Show Gist options
  • Save mmechtley/5633015 to your computer and use it in GitHub Desktop.
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 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))
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