def test_checkplot_pickle_update():
'''
Tests if a checkplot pickle can be made, read back, and updated.
'''
outpath = os.path.join(os.path.dirname(LCPATH), 'test-checkplot.pkl')
lcd, msg = hatlc.read_and_filter_sqlitecurve(LCPATH)
gls = periodbase.pgen_lsp(lcd['rjd'], lcd['aep_000'], lcd['aie_000'])
assert isinstance(gls, dict)
assert_allclose(gls['bestperiod'], 1.54289477)
pdm = periodbase.stellingwerf_pdm(lcd['rjd'], lcd['aep_000'],
lcd['aie_000'])
assert isinstance(pdm, dict)
assert_allclose(pdm['bestperiod'], 3.08578956)
# test write
cpf = checkplot.checkplot_pickle([gls, pdm],
lcd['rjd'],
lcd['aep_000'],
lcd['aie_000'],
outfile=outpath,
objectinfo=lcd['objectinfo'])
assert os.path.exists(outpath)
# test read back
cpd = _read_checkplot_picklefile(cpf)
assert isinstance(cpd, dict)
cpdkeys = set(list(cpd.keys()))
testset = {
'0-gls', '1-pdm', 'comments', 'externalplots', 'finderchart',
'magseries', 'neighbors', 'normmingap', 'normto', 'objectid',
'objectinfo', 'pfmethods', 'sigclip', 'signals', 'status', 'varinfo'
}
assert (testset - cpdkeys) == set()
assert_allclose(cpd['0-gls']['bestperiod'], 1.54289477)
assert_allclose(cpd['1-pdm']['bestperiod'], 3.08578956)
# test update write to pickle
cpd['comments'] = ('this is a test of the checkplot pickle '
'update mechanism. this is only a test.')
cpfupdated = checkplot_pickle_update(cpf, cpd)
cpdupdated = _read_checkplot_picklefile(cpfupdated)
assert cpdupdated['comments'] == cpd['comments']
def checkplot_infokey_worker(task):
'''This gets the required keys from the requested file.
Parameters
----------
task : tuple
Task is a two element tuple::
- task[0] is the dict to work on
- task[1] is a list of lists of str indicating all the key address to
extract items from the dict for
Returns
-------
list
This is a list of all of the items at the requested key addresses.
'''
cpf, keys = task
cpd = _read_checkplot_picklefile(cpf)
resultkeys = []
for k in keys:
try:
resultkeys.append(_dict_get(cpd, k))
except Exception as e:
resultkeys.append(np.nan)
return resultkeys
def test_checkplot_with_multiple_same_pfmethods():
'''
This tests running the same period-finder for different period ranges.
'''
outpath = os.path.join(os.path.dirname(LCPATH), 'test-checkplot.pkl')
lcd, msg = hatlc.read_and_filter_sqlitecurve(LCPATH)
gls_1 = periodbase.pgen_lsp(lcd['rjd'],
lcd['aep_000'],
lcd['aie_000'],
startp=0.01,
endp=0.1)
gls_2 = periodbase.pgen_lsp(lcd['rjd'],
lcd['aep_000'],
lcd['aie_000'],
startp=0.1,
endp=300.0)
pdm_1 = periodbase.stellingwerf_pdm(lcd['rjd'],
lcd['aep_000'],
lcd['aie_000'],
startp=0.01,
endp=0.1)
pdm_2 = periodbase.stellingwerf_pdm(lcd['rjd'],
lcd['aep_000'],
lcd['aie_000'],
startp=0.1,
endp=300.0)
assert isinstance(gls_1, dict)
assert isinstance(gls_2, dict)
assert isinstance(pdm_1, dict)
assert isinstance(pdm_2, dict)
cpf = checkplot.checkplot_pickle([gls_1, gls_2, pdm_1, pdm_2],
lcd['rjd'],
lcd['aep_000'],
lcd['aie_000'],
outfile=outpath,
objectinfo=lcd['objectinfo'])
assert os.path.exists(cpf)
assert os.path.abspath(cpf) == os.path.abspath(outpath)
cpd = _read_checkplot_picklefile(cpf)
pfmethods = list(cpd['pfmethods'])
assert len(pfmethods) == 4
assert '0-gls' in pfmethods
assert '1-gls' in pfmethods
assert '2-pdm' in pfmethods
assert '3-pdm' in pfmethods
def xmatch_cplist_external_catalogs(cplist,
xmatchpkl,
xmatchradiusarcsec=2.0,
updateexisting=True,
resultstodir=None):
'''This xmatches external catalogs to a collection of checkplots.
Parameters
----------
cplist : list of str
This is the list of checkplot pickle files to process.
xmatchpkl : str
The filename of a pickle prepared beforehand with the
`checkplot.pkl_xmatch.load_xmatch_external_catalogs` function,
containing collected external catalogs to cross-match the objects in the
input `cplist` against.
xmatchradiusarcsec : float
The match radius to use for the cross-match in arcseconds.
updateexisting : bool
If this is True, will only update the `xmatch` dict in each checkplot
pickle with any new cross-matches to the external catalogs. If False,
will overwrite the `xmatch` dict with results from the current run.
resultstodir : str or None
If this is provided, then it must be a directory to write the resulting
checkplots to after xmatch is done. This can be used to keep the
original checkplots in pristine condition for some reason.
Returns
-------
dict
Returns a dict with keys = input checkplot pickle filenames and vals =
xmatch status dict for each checkplot pickle.
'''
# load the external catalog
with open(xmatchpkl, 'rb') as infd:
xmd = pickle.load(infd)
# match each object. this is fairly fast, so this is not parallelized at the
# moment
status_dict = {}
for cpf in cplist:
cpd = _read_checkplot_picklefile(cpf)
try:
# match in place
xmatch_external_catalogs(cpd,
xmd,
xmatchradiusarcsec=xmatchradiusarcsec,
updatexmatch=updateexisting)
for xmi in cpd['xmatch']:
if cpd['xmatch'][xmi]['found']:
LOGINFO('checkplot %s: %s matched to %s, '
'match dist: %s arcsec' %
(os.path.basename(cpf), cpd['objectid'],
cpd['xmatch'][xmi]['name'],
cpd['xmatch'][xmi]['distarcsec']))
if not resultstodir:
outcpf = _write_checkplot_picklefile(cpd, outfile=cpf)
else:
xcpf = os.path.join(resultstodir, os.path.basename(cpf))
outcpf = _write_checkplot_picklefile(cpd, outfile=xcpf)
status_dict[cpf] = outcpf
except Exception:
LOGEXCEPTION('failed to match objects for %s' % cpf)
status_dict[cpf] = None
return status_dict
def add_cmd_to_checkplot(cpx,
cmdpkl,
require_cmd_magcolor=True,
save_cmd_pngs=False):
'''This adds CMD figures to a checkplot dict or pickle.
Looks up the CMDs in `cmdpkl`, adds the object from `cpx` as a gold(-ish)
star in the plot, and then saves the figure to a base64 encoded PNG, which
can then be read and used by the `checkplotserver`.
Parameters
----------
cpx : str or dict
This is the input checkplot pickle or dict to add the CMD to.
cmdpkl : str or dict
The CMD pickle generated by the `colormagdiagram_cplist` or
`colormagdiagram_cpdir` functions above, or the dict produced by reading
this pickle in.
require_cmd_magcolor : bool
If this is True, a CMD plot will not be made if the color and mag keys
required by the CMD are not present or are nan in this checkplot's
objectinfo dict.
save_cmd_png : bool
If this is True, then will save the CMD plots that were generated and
added back to the checkplotdict as PNGs to the same directory as
`cpx`. If `cpx` is a dict, will save them to the current working
directory.
Returns
-------
str or dict
If `cpx` was a str filename of checkplot pickle, this will return that
filename to indicate that the CMD was added to the file. If `cpx` was a
checkplotdict, this will return the checkplotdict with a new key called
'colormagdiagram' containing the base64 encoded PNG binary streams of
all CMDs generated.
'''
# get the checkplot
if isinstance(cpx, str) and os.path.exists(cpx):
cpdict = _read_checkplot_picklefile(cpx)
elif isinstance(cpx, dict):
cpdict = cpx
else:
LOGERROR('unknown type of checkplot provided as the cpx arg')
return None
# get the CMD
if isinstance(cmdpkl, str) and os.path.exists(cmdpkl):
with open(cmdpkl, 'rb') as infd:
cmd = pickle.load(infd)
elif isinstance(cmdpkl, dict):
cmd = cmdpkl
cpdict['colormagdiagram'] = {}
# get the mags and colors from the CMD dict
cplist_mags = cmd['mags']
cplist_colors = cmd['colors']
# now make the CMD plots for each color-mag combination in the CMD
for c1, c2, ym, ind in zip(cmd['color_mag1'],
cmd['color_mag2'], cmd['yaxis_mag'],
range(len(cmd['color_mag1']))):
# get these from the checkplot for this object
if (c1 in cpdict['objectinfo']
and cpdict['objectinfo'][c1] is not None):
c1mag = cpdict['objectinfo'][c1]
else:
c1mag = np.nan
if (c2 in cpdict['objectinfo']
and cpdict['objectinfo'][c2] is not None):
c2mag = cpdict['objectinfo'][c2]
else:
c2mag = np.nan
if (ym in cpdict['objectinfo']
and cpdict['objectinfo'][ym] is not None):
ymmag = cpdict['objectinfo'][ym]
else:
ymmag = np.nan
if (require_cmd_magcolor
and not (np.isfinite(c1mag) and np.isfinite(c2mag)
and np.isfinite(ymmag))):
LOGWARNING("required color: %s-%s or mag: %s are not "
"in this checkplot's objectinfo dict "
"(objectid: %s), skipping CMD..." %
(c1, c2, ym, cpdict['objectid']))
continue
# make the CMD for this color-mag combination
try:
thiscmd_title = r'%s-%s/%s' % (CMD_LABELS[c1], CMD_LABELS[c2],
CMD_LABELS[ym])
# make the scatter plot
fig = plt.figure(figsize=(10, 8))
plt.plot(cplist_colors[:, ind],
cplist_mags[:, ind],
rasterized=True,
marker='o',
linestyle='none',
mew=0,
ms=3)
# put this object on the plot
plt.plot([c1mag - c2mag], [ymmag],
ms=20,
color='#b0ff05',
marker='*',
mew=0)
plt.xlabel(r'$%s - %s$' % (CMD_LABELS[c1], CMD_LABELS[c2]))
plt.ylabel(r'$%s$' % CMD_LABELS[ym])
plt.title('%s - $%s$ CMD' % (cpdict['objectid'], thiscmd_title))
plt.gca().invert_yaxis()
# now save the figure to StrIO and put it back in the checkplot
cmdpng = StrIO()
plt.savefig(cmdpng,
bbox_inches='tight',
pad_inches=0.0,
format='png')
cmdpng.seek(0)
cmdb64 = base64.b64encode(cmdpng.read())
cmdpng.close()
plt.close('all')
plt.gcf().clear()
cpdict['colormagdiagram']['%s-%s/%s' % (c1, c2, ym)] = cmdb64
# if we're supposed to export to PNG, do so
if save_cmd_pngs:
if isinstance(cpx, str):
outpng = os.path.join(
os.path.dirname(cpx), 'cmd-%s-%s-%s.%s.png' %
(cpdict['objectid'], c1, c2, ym))
else:
outpng = 'cmd-%s-%s-%s.%s.png' % (cpdict['objectid'], c1,
c2, ym)
_base64_to_file(cmdb64, outpng)
except Exception:
LOGEXCEPTION('CMD for %s-%s/%s does not exist in %s, skipping...' %
(c1, c2, ym, cmdpkl))
continue
#
# end of making CMDs
#
if isinstance(cpx, str):
cpf = _write_checkplot_picklefile(cpdict, outfile=cpx, protocol=4)
return cpf
elif isinstance(cpx, dict):
return cpdict
def colormagdiagram_cplist(cplist,
outpkl,
color_mag1=('gaiamag', 'sdssg'),
color_mag2=('kmag', 'kmag'),
yaxis_mag=('gaia_absmag', 'rpmj')):
'''This makes color-mag diagrams for all checkplot pickles in the provided
list.
Can make an arbitrary number of CMDs given lists of x-axis colors and y-axis
mags to use.
Parameters
----------
cplist : list of str
This is the list of checkplot pickles to process.
outpkl : str
The filename of the output pickle that will contain the color-mag
information for all objects in the checkplots specified in `cplist`.
color_mag1 : list of str
This a list of the keys in each checkplot's `objectinfo` dict that will
be used as color_1 in the equation::
x-axis color = color_mag1 - color_mag2
color_mag2 : list of str
This a list of the keys in each checkplot's `objectinfo` dict that will
be used as color_2 in the equation::
x-axis color = color_mag1 - color_mag2
yaxis_mag : list of str
This is a list of the keys in each checkplot's `objectinfo` dict that
will be used as the (absolute) magnitude y-axis of the color-mag
diagrams.
Returns
-------
str
The path to the generated CMD pickle file for the collection of objects
in the input checkplot list.
Notes
-----
This can make many CMDs in one go. For example, the default kwargs for
`color_mag`, `color_mag2`, and `yaxis_mag` result in two CMDs generated and
written to the output pickle file:
- CMD1 -> gaiamag - kmag on the x-axis vs gaia_absmag on the y-axis
- CMD2 -> sdssg - kmag on the x-axis vs rpmj (J reduced PM) on the y-axis
'''
# first, we'll collect all of the info
cplist_objectids = []
cplist_mags = []
cplist_colors = []
for cpf in cplist:
cpd = _read_checkplot_picklefile(cpf)
cplist_objectids.append(cpd['objectid'])
thiscp_mags = []
thiscp_colors = []
for cm1, cm2, ym in zip(color_mag1, color_mag2, yaxis_mag):
if (ym in cpd['objectinfo'] and cpd['objectinfo'][ym] is not None):
thiscp_mags.append(cpd['objectinfo'][ym])
else:
thiscp_mags.append(np.nan)
if (cm1 in cpd['objectinfo'] and cpd['objectinfo'][cm1] is not None
and cm2 in cpd['objectinfo']
and cpd['objectinfo'][cm2] is not None):
thiscp_colors.append(cpd['objectinfo'][cm1] -
cpd['objectinfo'][cm2])
else:
thiscp_colors.append(np.nan)
cplist_mags.append(thiscp_mags)
cplist_colors.append(thiscp_colors)
# convert these to arrays
cplist_objectids = np.array(cplist_objectids)
cplist_mags = np.array(cplist_mags)
cplist_colors = np.array(cplist_colors)
# prepare the outdict
cmddict = {
'objectids': cplist_objectids,
'mags': cplist_mags,
'colors': cplist_colors,
'color_mag1': color_mag1,
'color_mag2': color_mag2,
'yaxis_mag': yaxis_mag
}
# save the pickled figure and dict for fast retrieval later
with open(outpkl, 'wb') as outfd:
pickle.dump(cmddict, outfd, pickle.HIGHEST_PROTOCOL)
plt.close('all')
return cmddict