Coverage for /opt/anaconda/lib/python3.7/site-packages/llama/utils.py : 60%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# (c) Stefan Countryman, September 2, 2016
3"""
4Utility functions for LLAMA with *no dependencies* on other LLAMA tooling.
5"""
7import sys
8import warnings
9import tempfile
10import tarfile
11import logging
12import gzip
13import os
14import subprocess
15import importlib
16from pathlib import Path
17from pkgutil import iter_modules
18from types import FunctionType
19from stat import S_ISREG, ST_MTIME, ST_MODE
20from subprocess import check_call
21from collections import namedtuple
22from textwrap import wrap
23from pygments import highlight
24from pygments.lexers import get_lexer_by_name
25from pygments.formatters import Terminal256Formatter
26from pygments.util import ClassNotFound
27import six
28from six import string_types
29import functools
30from llama.classes import (
31 COLOR,
32 optional_env_var,
33 GenerationError, # gets imported from utils elsewhere TODO move to classes
34 AbstractFileHandler,
35)
37_MAINPROC = [True]
40def label_worker_proc():
41 """
42 Label the current process as a worker process (not the main process). You
43 should only call this as the initializer to an executor.
44 """
45 _MAINPROC[0] = False
48def is_main_process():
49 """
50 Return whether this is the main process (as opposed to a worker process).
51 """
52 return _MAINPROC[0]
55def color(fg=None, bg=None):
56 """
57 Return full-color escape codes for changing foreground and background in
58 terminals with RGB color support.
60 Parameters
61 ----------
62 fg : str or list, optional
63 The foreground (text) color. A 3-item list containing integer values
64 for red, green, and blue *or* a hexadecimal representation of the color
65 prefixed by an octothorpe (like in CSS). Background will not be set if
66 not specified.
67 bg : str or list, optional
68 The background color. Same type as ``fg``.
70 Returns
71 -------
72 esc : str
73 The teriminal escape code to use.
75 Raises
76 ------
77 ValueError
78 If ``fg`` or ``bg`` are not as specified.
79 """
80 if fg is None and bg is None:
81 return ''
82 args = [fg, bg]
83 bases = ['38;2;', '48;2;']
84 fmt_part = '{};{};{}'
85 parts = []
86 for i, a in enumerate(args):
87 if a is not None:
88 if not hasattr(a, '__len__'):
89 raise ValueError("Arguments must be string, tuple, or None. "
90 f"Got fg = {fg} and bg = {bg}")
91 if len(a) == 7:
92 a = args[i] = [int(a[1+2*j:3+2*j], base=16) for j in range(3)]
93 elif len(a) != 3:
94 raise ValueError(f"Arg {i} should be hex or (r, g, b)")
95 parts.append(bases[i]+fmt_part.format(*a))
96 return '\033['+';'.join(parts)+'m'
99class ColorFormatter(logging.Formatter):
100 """A formatter that colors log output for easier reading."""
102 COLORS = {
103 'WARNING': COLOR.MAGENTA,
104 'INFO': COLOR.GREEN,
105 'DEBUG': COLOR.BLUE,
106 'CRITICAL': COLOR.RED,
107 'ERROR': COLOR.YELLOW,
108 }
110 def format(self, record):
111 levelname = record.levelname
112 if levelname in self.COLORS:
113 lname = record.levelname
114 record.levelname = (COLOR.BOLD + self.COLORS[lname] +
115 f"{lname:<8}" + COLOR.CLEAR + COLOR.UNDERLINE)
116 return logging.Formatter.format(self, record)
119LOGGER = logging.getLogger(__name__)
120# use this formatter throughout the library
121LOG_FORMATTER = ColorFormatter(
122 COLOR.UNDERLINE + (
123 '|%(levelname)s|%(asctime)-19s|%(processName)16s:%(process)-6s|'
124 '%(threadName)-14s|%(name)-18s|%(filename)14s:%(lineno)-4s|'
125 '%(funcName)14s|'
126 ) + COLOR.CLEAR + ' %(message)s',
127 datefmt='%Y-%m-%d %H:%M:%S'
128)
129# define some other important directories and make sure they exist
130DEFAULT_OUTPUT_DIR = os.path.join(
131 os.environ.get(
132 'XDG_DATA_HOME',
133 os.path.join(os.path.expanduser('~'), '.local', 'share')
134 ),
135 'llama',
136)
137OUTPUT_DIR = optional_env_var(
138 ["LLAMA_OUTPUT_DIR"],
139 "Specify where LLAMA data, triggers, and results should be saved locally.",
140 [DEFAULT_OUTPUT_DIR],
141)[0]
142DATADIR = os.path.join(OUTPUT_DIR, 'inputs')
143TEST_DATA = os.path.join(DATADIR, 'tests')
144TEST_EVENT_FILES = os.path.join(TEST_DATA, 'event_files')
145DEFAULT_CACHEDIR = os.path.join(
146 os.environ.get(
147 'XDG_CACHE_HOME',
148 os.path.join(os.path.expanduser('~'), '.cache')
149 ),
150 'llama',
151)
152CACHEDIR = optional_env_var(
153 ["LLAMA_CACHEDIR"],
154 "Specify where LLAMA data should be cached.",
155 [DEFAULT_CACHEDIR],
156)[0]
157OBJECT_DIR = os.path.join(CACHEDIR, 'objects')
158SINGLE_DETECTOR_TRIGGER_DIR = os.path.join(OUTPUT_DIR,
159 'single_detector_triggers')
160DEFAULT_RUN_DIR = os.path.join(OUTPUT_DIR, 'current_run/')
161PAST_RUNS_DIR = os.path.join(OUTPUT_DIR, 'past_runs')
162LOGDIR = os.path.join(OUTPUT_DIR, 'logs')
163LOGFILE = os.path.join(LOGDIR, 'llama.log')
164for default_dir in (OUTPUT_DIR, DATADIR, DEFAULT_RUN_DIR, PAST_RUNS_DIR,
165 OBJECT_DIR, CACHEDIR, LOGDIR, SINGLE_DETECTOR_TRIGGER_DIR):
166 if not os.path.isdir(default_dir):
167 try:
168 os.makedirs(default_dir)
169 except OSError as err:
170 LOGGER.error("Could not create default directory %s. Error: %s",
171 default_dir, err)
172 warnings.warn(("Could not create default directory expected by "
173 "llama: {} Error: {}").format(default_dir, err))
175TB_LEXER = get_lexer_by_name('py3tb') # lexer
176try:
177 TB_FORMATTER = Terminal256Formatter(style='native')
178except ClassNotFound:
179 LOGGER.debug("Can't load 'native' pygments style, using default.")
180 TB_FORMATTER = Terminal256Formatter()
181SIDEREAL_DAY = 0.997269566
182MAX_GRID_SCALE_RAD = 1e-7
183# graphviz .dot graph file format for visualization functions
184# https://www.graphviz.org/doc/info/shapes.html
185EDGEFMT = (r'"{depnum}" -> "{num}" ['
186 r'arrowtail="none", '
187 r'color="{color}", '
188 # r'arrowhead="dot"'
189 r'];')
190DOTFMT = """digraph "{name}" {{
191charset="utf-8"
192splines=ortho
193bgcolor={bgcolor}
194margin=0
195rankdir=LR
196{nodes}
197{edges}
198}}
199"""
200GRAPH_EXTENSIONS = (
201 "png",
202 "pdf",
203 "svg",
204)
207def setup_logger(logfile, loglevel='DEBUG'):
208 """Set up a logger for ``llama`` and all submodules that logs to both
209 ``logfile`` and to STDOUT. You will still need to actually make a logger
210 for whatever module you are calling this from.
212 Parameters
213 ----------
214 logfile : str
215 The logfile to write all output to. If the path equals or resolves to
216 ``/dev/null``, no logfile will be configured.
217 loglevel : int, str, or NoneType, optional
218 The loglevel to pass to ``setLevel`` for the ``StreamHandler`` writing
219 to the terminal; the ``logfile`` handler always writes at maximum
220 verbosity (``DEBUG``).
221 """
222 logger_llama = logging.getLogger('llama')
223 logger_llama.setLevel(logging.DEBUG)
224 # set up stream handler for logger and format it properly
225 if loglevel is not None:
226 log_stream = logging.StreamHandler()
227 log_stream.setFormatter(LOG_FORMATTER)
228 log_stream.setLevel(loglevel)
229 logger_llama.addHandler(log_stream)
230 # set up file handler for logger and format it properly
231 log = Path(logfile)
232 if log == '/dev/null' or log.resolve() == Path('/dev/null').resolve():
233 return
234 log_file = logging.FileHandler(logfile, 'a')
235 log_file.setFormatter(LOG_FORMATTER)
236 logger_llama.addHandler(log_file)
239# TODO update this so that it dynamically downloads test files from some
240# testing branch as needed to $XDG_DATA_HOME.
241def get_test_file(filename, eventid):
242 """Get the path to an example file as would be generated by LLAMA. These
243 are stored in a single directory (equivalent to an event directory) and can
244 be used for unit, integration, and doc tests.
246 Parameters
247 ----------
248 filename : str
249 The name of the file as it would appear in an event directory with no
250 path information before it.
251 eventid : str
252 The name of the event directory we will be using for this unit test.
253 Since different test cases apply to different file handlers, you'll
254 need to make sure to avoid specifying an ``eventid`` that this
255 filehandler is not meant to be tested on (whatever the reason).
257 Returns
258 -------
259 test_filepath : str
260 The full path to the file to be used for testing.
261 """
262 test_filepath = os.path.join(TEST_EVENT_FILES, eventid, filename)
263 assert os.path.isfile(test_filepath)
264 return test_filepath
267def bytes2str(bytes_or_str):
268 """Convert an object that could be a bytes object (in python3) to a str
269 object. If it is already a str object, leave it alone. Used for
270 python2/python3 compatibility."""
271 if not isinstance(bytes_or_str, str):
272 return bytes_or_str.decode("utf-8")
273 return bytes_or_str
276try:
277 from llama.version import version as __version__
278except ModuleNotFoundError:
279 LOGGER.critical("Could not determine version; llama.version not "
280 "defined. Your installation is broken. If you are a "
281 "developer, install the library with `pip install .` (or, "
282 "if you know what you're doing, with `pip install -e .`, "
283 "though this will break automatic version detection and "
284 "necessitate a reinstall when you want to set "
285 "llama.version to the correct value).")
286 raise
289def memoize(method=False):
290 """A decorator that takes a method and memoizes it by storing the
291 value in a class or instance variable. Useful when the same arguments are
292 likely to keep recurring and the computation time is long. Data is
293 stored in the function (for regular functions and staticmethods), the class
294 (for classmethods) or in the instance (for properties and methods), this
295 approach can't work for staticmethods nor bare
296 functions. For this to work, all arguments (except ``self`` or ``cls``)
297 must be hashable.
299 NOTE: all arguments must be hashable (with the exception of the class or
300 instance when ``method`` is ``True``; see below).
302 Parameters
303 ----------
304 method : bool, default=False
305 If true, treat this function like a method, so that the first argument
306 is a class or instance, and store the cache with that class or instance
307 instead of storing it with the function. This lets you use function
308 memoization on properties, classmethods, and instance methods even if
309 the class or instance is not hashable.
310 """
311 # pylint: disable=missing-docstring
312 def decorator(func):
313 #TODO remove this to reinstate memoization, or else permanently remove
314 # memoization later.
315 return func
316 if not method:
317 func_cache_closure = dict()
318 # pylint: disable=missing-docstring
319 @functools.wraps(func)
320 def wrapper(*args, **kwargs):
321 if method:
322 if not hasattr(args[0], '_cache'):
323 setattr(args[0], '_cache', dict())
324 # have a cache for all functions in this class/instance
325 cache = args[0]._cache
326 # this is the cache just for this function
327 if func not in cache:
328 cache[func] = dict()
329 func_cache = cache[func]
330 hash_args = args[1:]
331 else:
332 func_cache = func_cache_closure
333 hash_args = args
334 # force argument key to be immutable and hence hashable; fails if
335 # this wouldn't work anyway (e.g. one of ``args`` is a list)
336 key = str(
337 hash((
338 tuple(hash_args),
339 tuple((k, kwargs[k]) for k in sorted(kwargs))
340 ))
341 )
342 if key not in func_cache:
343 val = func(*args, **kwargs)
344 try:
345 func_cache[key] = val
346 except TypeError as err:
347 err.args = ((err.args[0] + (" Can't save result to HDF5: "
348 "{}").format(val),) +
349 err.args[1:])
350 raise err
351 return func_cache[key]
353 wrapper.__doc__ = func.__doc__
354 return wrapper
356 return decorator
359def write_gzip(infile, outfilename):
360 """See whether a file is actually a valid gzip file. If not, zip it. Either
361 way, write the result to ``outfilename``.
363 Parameters
364 ----------
365 infile : file-like
366 File-like object to be compressed to a gzip file.
367 outfilename : str
368 Path to the output compressed file. Will be overwritten if it already
369 exists.
371 Returns
372 -------
373 success : bool
374 ``True`` if it was already a valid gzip file and ``False`` if we were
375 forced to compress it.
377 Examples
378 --------
379 Compress a string in a file-like ``BytesIO`` wrapper to some temporary
380 ``.gz`` file:
381 >>> from io import BytesIO
382 >>> import tempfile
383 >>> with tempfile.NamedTemporaryFile(suffix='.gz', delete=False) as f:
384 ... outfilename = f.name
385 >>> infile = BytesIO('this is obviously not a valid gzip file'.encode())
386 >>> assert not write_gzip(infile, outfilename)
387 >>> with open(outfilename, 'rb') as actuallyzippedfile:
388 ... assert write_gzip(actuallyzippedfile, outfilename)
389 """
390 with open(outfilename, 'wb') as outfile:
391 outfile.write(infile.read())
392 try:
393 with gzip.open(outfilename) as gzfile:
394 gzfile.read(1) # read 1 byte if this is really a gz file
395 return True
396 except IOError:
397 infile.seek(0)
398 with gzip.open(outfilename, 'wb') as gzfile:
399 gzfile.write(infile.read())
400 return False
403# "ra" is a perfectly clear contraction for "Right Ascension".
404# pylint: disable=invalid-name
405def rotate_angs2vec(ra, dec, yrot, zrot, degrees=True):
406 """Take a list of angular sky positions (points on the sphere) and rotate
407 them first through an angle ``yrot`` about the y-axis (right-handed)
408 followed by an angle ``zrot`` about the z-axis. Return the new points as
409 (x, y, z) coordinates of the rotated directions as vectors on the unit
410 sphere.
412 Parameters
413 ----------
414 ra : array-like
415 A list or ``numpy.ndarray`` of Right Ascension values for each point on
416 the sphere.
417 dec : array-like
418 A list or ``numpy.ndarray`` of Declination values for each point on the
419 sphere. Must have same length as ``ra``.
420 yrot : float
421 The angle through which to do a right-handed rotation about the
422 positive y-axis. This rotation is applied first.
423 zrot : float
424 The angle through which to do a right-handed rotation about the
425 positive z-axis. This rotation is applied second.
426 degrees : bool, optional
427 (DEFAULT: True) If true, interpret all input angles as being measured
428 in degrees. Otherwise, interpret them as being measured in radians.
430 Returns
431 -------
432 x : array
433 The x-coordinates of the rotated input vectors [cartesian].
434 y : array
435 The y-coordinates of the rotated input vectors [cartesian].
436 z : array
437 The z-coordinates of the rotated input vectors [cartesian].
439 Examples
440 --------
441 >>> import numpy as np
442 >>> def eql(a, b, prec=1e-15): # compares floating point arrays
443 ... return all(prec > np.array(a).flatten() - np.array(b).flatten())
444 >>> eql(rotate_angs2vec(0, 0, 0, 0), [1, 0, 0])
445 True
446 >>> eql(rotate_angs2vec(270, 0, 0, 0), [0, -1, 0])
447 True
448 >>> eql(rotate_angs2vec(270, 0, 0, 90), [1, 0, 0])
449 True
450 >>> eql(rotate_angs2vec(180, 0, 90, 90), [0, 0, 1])
451 True
452 >>> eql(rotate_angs2vec(0, 90, 90, 90), [0, 1, 0])
453 True
454 >>> eql(rotate_angs2vec(np.pi/2, 0, 0, np.pi, degrees=False), [0, -1, 0])
455 True
456 >>> eql(rotate_angs2vec(np.pi, 0, np.pi/2, 0, degrees=False), [0, 0, 1])
457 True
458 """
459 import numpy as np
460 twopi = 2 * np.pi
461 if degrees:
462 ra = np.radians(ra).reshape((-1,)) % twopi
463 dec = np.radians(dec).reshape((-1,)) % twopi
464 yrot = np.radians(yrot) % twopi
465 zrot = np.radians(zrot) % twopi
466 else:
467 ra = np.array(ra).reshape((-1,)) % twopi
468 dec = np.array(dec).reshape((-1,)) % twopi
469 yrot = np.array(yrot) % twopi
470 zrot = np.array(zrot) % twopi
471 # create a rotation matrix to move everything into place
472 cosy = np.cos(yrot)
473 siny = np.sin(yrot)
474 cosz = np.cos(zrot)
475 sinz = np.sin(zrot)
476 rotation_matrix = np.array([[cosz*cosy, -sinz, cosz*siny],
477 [sinz*cosy, cosz, sinz*siny],
478 [-siny, 0, cosy]])
479 # get the x, y, z coordinates of the points to be rotated
480 cosd = np.cos(dec)
481 xyz = np.array([cosd*np.cos(ra), cosd*np.sin(ra), np.sin(dec)])
482 # in python3.5+, @ is the matrix product operator in numpy; we can replace
483 # the matrix product function call with that infix operator once we drop
484 # support for python 3.4-.
485 prods = np.matmul(rotation_matrix, xyz)
486 return tuple([np.array(prods[i, :]).flatten() for i in range(3)])
489# pylint: disable=C0103
490def rotate_angs2angs(ra, dec, yrot, zrot, degrees=True):
491 """Take a list of angular sky positions (points on the sphere) and rotate
492 them first through an angle ``yrot`` about the y-axis (right-handed)
493 followed by an angle ``zrot`` about the z-axis. Return the new points as
494 (ra, dec) coordinates of the rotated vectors.
496 Parameters
497 ----------
498 ra : array-like
499 A list or ``numpy.ndarray`` of Right Ascension values for each point on
500 the sphere.
501 dec : array-like
502 A list or ``numpy.ndarray`` of Declination values for each point on the
503 sphere. Must have same length as ``ra``.
504 yrot : float
505 The angle through which to do a right-handed rotation about the
506 positive y-axis. This rotation is applied first.
507 zrot : float
508 The angle through which to do a right-handed rotation about the
509 positive z-axis. This rotation is applied second.
510 degrees : bool, optional
511 (DEFAULT: True) If true, interpret all input angles as being measured
512 in degrees and return values in degrees. Otherwise, interpret them as
513 being measured in radians and return values in radians.
515 Returns
516 -------
517 outra : array
518 The Right Ascensions of the rotated input vectors in the specified
519 units.
520 outdec : array
521 The Declinations of the rotated input vectors in the specified units.
523 Examples
524 --------
525 >>> import numpy as np
526 >>> def eql(a, b, prec=1e-13, mod=360): # compares floating point arrays
527 ... a = np.array(a).flatten()
528 ... b = np.array(b).flatten()
529 ... return all(prec > (a%mod - b%mod))
530 >>> ra = range(360)
531 >>> dec = np.zeros(360)
532 >>> eql(rotate_angs2angs(ra, dec, 0, 0), [ra, dec])
533 True
534 """
535 import numpy as np
536 x, y, z = rotate_angs2vec(ra, dec, yrot, zrot, degrees)
537 r = np.sqrt(x**2 + y**2)
538 outra = np.arctan2(y, x)
539 outdec = np.arctan2(z, r)
540 if degrees:
541 outra = np.degrees(outra)
542 outdec = np.degrees(outdec)
543 return (outra, outdec)
546# Notes on defining a circular sky region (unused for now):
547# mask = cosd(sigma)-sind(cDEC)*sind(dec) <= cosd(cDEC)*cosd(dec).*cosd(ra -
548# cRA);
549# at circle:
550# cosd(sigma)-sind(cDEC)*sind(dec) = cosd(cDEC)*cosd(dec).*cosd(ra - cRA);
551# cosd(ra - cRA) = (cosd(sigma)-sind(cDEC)*sind(dec)) / (cosd(cDEC)*cosd(dec))
552# ra - cRA = acosd((cosd(sigma)-sind(cDEC)*sind(dec)) / (cosd(cDEC)*cosd(dec)))
553# ra = acosd((cosd(sigma)-sind(cDEC)*sind(dec)) / (cosd(cDEC)*cosd(dec))) + cRA
554# ramax = +ra % 360
555# ramin = -ra % 360
558# pylint: disable=C0103
559def get_grid(ra, dec, radius, pixels=100, degrees=True):
560 """Get a list of pixel positions arranged in a rectangular grid and filling
561 a square with sizes of length ``radius`` centered at Right Ascension ``ra``
562 and Declination ``dec`` with ``pixels`` as the width and height in pixels
563 of the square grid.
565 This technically only works properly if the radius is very small, so a
566 ValueError will be raised for radii larger than ``MAX_GRID_SCALE_RAD``.
567 For larger sky areas, HEALPix pixelizations should be used to assure equal
568 area.
570 Parameters
571 ----------
572 ra : float
573 The Right Ascension of the center of the grid.
574 dec : float
575 The Declination of the center of the grid.
576 radius : float
577 The radius of the circle. Must be larger than ``MAX_GRID_SCALE_RAD``
578 (in radians) to keep pixel areas nearly constant.
579 pixels : float, optional
580 The diameter, in pixels, of the circle along the cardinal directions of
581 the grid. Higher values imply higher resolutions for the grid.
582 degrees : bool, optional
583 If true, interpret all input and returned angles as being measured in
584 degrees. Otherwise, interpret them all as being measured in radians.
586 Returns
587 -------
588 ras : array
589 Right Ascension values of the grid pixels in the specified units.
590 decs : array
591 Declination values of the grid pixels in the specified units.
592 area : float
593 the (approximate) solid-angle per-pixel in the specified units.
595 Raises
596 ------
597 ValueError
598 Raised when the ``radius`` is too large (larger than
599 ``MAX_GRID_SCALE_RAD``, see note above).
600 """
601 import numpy as np
602 area = (float(radius) / pixels)**2
603 if degrees:
604 ra = np.radians(ra)
605 dec = np.radians(dec)
606 radius = np.radians(radius)
607 if radius > MAX_GRID_SCALE_RAD:
608 raise ValueError(
609 'radius must be less than {}'.format(MAX_GRID_SCALE_RAD)
610 )
611 grid_vals = np.linspace(-radius, radius, pixels).reshape((1, -1))
612 ras = (grid_vals * np.ones((pixels, 1))).flatten()
613 decs = (grid_vals.transpose() * np.ones((1, pixels))).flatten()
614 ras, decs = rotate_angs2angs(ras, decs, -dec, ra, degrees=False)
615 if degrees:
616 ras = np.degrees(ras)
617 decs = np.degrees(decs)
618 return (ras, decs, area)
620# convert between zenith/azumuth and right ascension/declination units, where
621# zenith/aziumuth are measured at ICECUBE's position (south pole)
624def color_logger(loglevel=None, outfile=sys.stdout):
625 """Return a function ``log`` that prints formatted strings to either
626 ``sys.stdout`` or with the logger using a certain log level (if
627 ``loglevel`` is specified). If ``loglevel`` is specified, then ``outfile``
628 is ignored."""
630 def log(*args, **kwargs):
631 """Print formatted strings to STDOUT, unless some log level is set,
632 in which case log it. Optionall, specify a color from ``COLOR`` as
633 ``col``."""
634 msg = ''.join([format(a) for a in args])
635 if 'col' in kwargs:
636 msg = getattr(COLOR, kwargs['col']) + msg + COLOR.CLEAR
637 if loglevel is None:
638 print(msg, file=outfile)
639 else:
640 LOGGER.__getattribute__(loglevel)(msg)
642 return log
645def archive_figs(plots, outfile, exts=('pdf', 'png'), fname_list=None,
646 close_figs=False):
647 """Take a list of figures and save them with filenames that look like
648 ``{fname}.{suffix}`` to a ``.tar.gz`` archive named ``outfile``. The
649 ``list_index`` will be the zero-padded index of the given plot in the
650 ``plots`` list. Default filenames are generated if ``fname_list`` is not
651 given. For instance, with no ``fname_list`` specified, the third plot in
652 ``plots`` would have files named named ``plt_003.pdf`` and ``plt_003.png``
653 saved to ``outfile`` (assuming default ``kwargs`` values).
655 Parameters
656 ----------
657 plots : list
658 A list of ``matplotlib`` figures to save to file.
659 outfile : str
660 Filename of the output archive. Should end in ``.tar.gz`` since this will
661 be a zipped tarball of figures.
662 exts : list, optional
663 List of filename extensions to use. Each extension becomes the
664 ``suffix`` in the above format string. Allows multiple image formats to
665 be saved for each plot.
666 fname_list : list, optional
667 A list of file names to use for each figure (sans file extensions).
668 *NOTE:* If not provided, some default filenames are generated instead.
669 These follow the filename pattern ``plt_{list_index}.{suffix}``.
670 close_figs : bool, optional
671 If ``True``, clear the figures in ``plots`` after plotting is finished.
672 This will help reduce memory usage if the plots are not going to be
673 reused.
674 """
675 if fname_list is None:
676 fname_list = ['plt_{:03d}.{}'.format(i, ext)
677 for ext in exts
678 for i in range(len(plots))]
679 else:
680 fname_list = [f + '.' + e for e in exts for f in fname_list]
681 # the default arg value below is due to a scoping issue with lambda
682 # closures in python; it explicitly establishes the scope of ``p`` within
683 # the lambda rather than outside of it. for details, see:
684 # stackoverflow.com/questions/2295290
685 plot_writers = [lambda fname, plt=p: plt.savefig(fname) for p in plots]
686 write_to_zip(fname_list, plot_writers * len(exts), outfile,
687 suffix=sum([['.' + ext]*len(plots) for ext in exts], []))
688 if close_figs:
689 for fig in plots:
690 fig.clf()
693def write_to_zip(fname_list, write_functions, outfilename, suffix=''):
694 """Create a new ``.tar.gz`` tarfile archive names ``outfilename`` and fill
695 it with filenames given in ``fname_list``. ``write_functions`` should
696 be a list of functions that takes an input filename as its only argument
697 and writes data to that filename. It should have the same length as
698 ``fname_list``, since each element of ``write_functions`` will be used to
699 generate the corresponding file in ``fname_list``. If the
700 ``write_functions`` expect some sort of specific file extension, that can
701 be provided by manually setting ``suffix`` (DEFAULT: '') to the appropriate
702 extension. If each file requires its own specific filename extension, then
703 ``suffix`` can be specified as a list of suffixes with length matching
704 ``fname_list`` and ``write_functions``.
706 Examples
707 --------
708 >>> import tempfile, tarfile
709 >>> def foo(file):
710 ... with open(file, 'w') as f:
711 ... f.write('foo!')
712 >>> write_functions = [foo] * 3
713 >>> fname_list = ['bar.txt', 'baz.txt', 'quux.txt']
714 >>> outfile = tempfile.NamedTemporaryFile(delete=False)
715 >>> outfilename = outfile.name
716 >>> outfile.close()
717 >>> write_to_zip(fname_list, write_functions, outfilename, suffix='.txt')
718 >>> tar = tarfile.open(outfilename)
719 >>> [f.name for f in tar.getmembers()]
720 ['bar.txt', 'baz.txt', 'quux.txt']
721 >>> os.remove(outfilename)
722 """
723 if not isinstance(suffix, list):
724 suffix = len(fname_list) * [suffix]
725 if not len(fname_list) == len(write_functions):
726 raise ValueError('fname_list and write_functions must have same len')
727 with tarfile.open(outfilename, 'w|gz') as tar:
728 for i, fname in enumerate(fname_list):
729 tmp = tempfile.NamedTemporaryFile(suffix=suffix[i], delete=False)
730 tmpname = tmp.name
731 tmp.close()
732 write_functions[i](tmpname)
733 tar.add(tmpname, arcname=fname)
734 os.remove(tmpname)
737# some common time conversions
738def mjd2utc(mjd):
739 """Convert MJD time to UTC string"""
740 import astropy.time
741 return astropy.time.Time(mjd, format='mjd').isot
744def mjd2gps(mjd):
745 """Convert MJD time to GPS seconds"""
746 import astropy.time
747 return astropy.time.Time(mjd, format='mjd').gps
750def utc2mjd(utc):
751 """Convert UTC time string to MJD time"""
752 import astropy.time
753 return astropy.time.Time(utc).mjd
756def utc2gps(utc):
757 """Convert UTC time string to GPS seconds"""
758 import astropy.time
759 return astropy.time.Time(utc).gps
762def gps2mjd(gps):
763 """Convert GPS seconds to MJD time"""
764 import astropy.time
765 return astropy.time.Time(gps, format='gps', scale='utc').mjd
768def gps2utc(gps):
769 """Convert GPS seconds to UTC string"""
770 import astropy.time
771 return astropy.time.Time(gps, format='gps', scale='utc').isot
774def get_voevent_param(voeventpath, param_name):
775 """Get a 'What' parameter from a VOEvent XML file and return it as a string
776 (no clever type inference is performed). For example, get the FAR
777 (False Alarm Rate) for event G298048 (first BNS detection), which is
778 included as a standard test case in 'llama/data/tests/voevents':
780 Examples
781 --------
782 >>> get_voevent_param(get_test_file('lvc_gcn.xml',
783 ... 'MS181101ab-2-Initial'), 'FAR')
784 '9.11069936486e-14'
786 Note that the type of the returned data will (probably) be a unicode
787 string, so be sure to convert as necessary."""
788 if not os.path.isfile(voeventpath):
789 raise IOError(voeventpath + ': file does not exist.')
790 import untangle
791 event = untangle.parse(voeventpath)
792 params = event.voe_VOEvent.What.Param
793 # return filter(lambda a: a['name'] == param_name, params)[0]['value']
794 return [p for p in params if p['name'] == param_name][0]['value']
797def get_voevent_notice_time(voeventpath):
798 """Get a unicode string with the date and time of the notification rather
799 than the event itself. Is in ISO format."""
800 if not os.path.isfile(voeventpath):
801 raise IOError(voeventpath + ': file does not exist.')
802 import untangle
803 event = untangle.parse(voeventpath)
804 return event.voe_VOEvent.Who.Date.cdata
807def get_voevent_time(voeventpath):
808 """Get a unicode string with the ISO date and time of the event:
810 Examples
811 --------
812 >>> get_voevent_time(get_test_file('lvc_gcn.xml',
813 ... 'MS181101ab-2-Initial'))
814 '2018-11-01T22:22:46.654437'
815 """
816 if not os.path.isfile(voeventpath):
817 raise IOError(voeventpath + ': file does not exist.')
818 import untangle
819 event = untangle.parse(voeventpath)
820 ol = event.voe_VOEvent.WhereWhen.ObsDataLocation.ObservationLocation
821 return ol.AstroCoords.Time.TimeInstant.ISOTime.cdata
824def getVOEventGPSSeconds(voeventpath):
825 """Get the GPS time of the event described in this VOEvent file
826 to the nearest second."""
827 return int(utc2gps(get_voevent_time(voeventpath)))
830def getVOEventGPSNanoseconds(voeventpath):
831 """Get the number of nanoseconds past the start of the current
832 GPS second of the event described in this VOEvent file. For example,
833 if the current GPS time is 123456789.033, this method will return
834 33000000."""
835 return (get_voevent_gps_time(voeventpath) -
836 getVOEventGPSSeconds(voeventpath))
839def get_voevent_gps_time(voeventpath):
840 """Get the GPS time of the event described in this VOEvent file.
841 Maximum precision is in nanoseconds, though in practice it will
842 not be this high."""
843 return utc2gps(get_voevent_time(voeventpath))
846def get_voevent_role(voeventpath):
847 """Get the role of this VOEvent, test or observation."""
848 import untangle
849 event = untangle.parse(voeventpath)
850 return event.voe_VOEvent['role']
853def send_email(subject, recipients, body='', attachments=()):
854 """Send an email using the default configured email address. No return
855 value.
857 Parameters
858 ----------
859 recipients : list
860 A list of strings representing email addresses of the recipients.
861 body : str or file-like, optional
862 Send contents of this string or file-like object.
863 attachments : str or list, optional
864 A single filepath or a list of filepaths to files that should also be
865 attached to this email.
866 """
867 if not isinstance(recipients, list):
868 raise ValueError('``recipients`` must be a list of strings.')
869 if isinstance(body, six.string_types):
870 tmp = tempfile.TemporaryFile(mode='w')
871 tmp.write(body)
872 tmp.seek(0)
873 body = tmp
874 # add attachments, if specified
875 if isinstance(attachments, six.string_types):
876 attachments = [attachments]
877 mail_cmd = ['mail', '-s', subject]
878 for attachment in attachments:
879 mail_cmd.append('-a')
880 mail_cmd.append(attachment)
881 mail_cmd.append(','.join(recipients))
882 LOGGER.warning('running send_email command: %s\n', mail_cmd)
883 # try sending the email, making sure to explicitly close the file
884 # descriptor if it was opened within this function
885 proc = subprocess.Popen(mail_cmd, stdin=body, stdout=subprocess.PIPE,
886 stderr=subprocess.PIPE)
887 res, err = proc.communicate()
888 body.close()
889 LOGGER.debug('STDOUT:\n%s\n', res)
890 LOGGER.debug('STDERR:\n%s\n', err)
891 if proc.returncode != 0:
892 raise Exception(
893 """
894 MAIL FAILED.
896 subject: {}
897 dumping stdout: {}
898 dumping stderr: {}
899 """.format(subject, res, err)
900 )
903def parameter_factory(classname, description, **kwargs):
904 r"""Return a class that can be used as a namespace for a bunch of
905 parameters related to a specific search. Very similar to defining a
906 ``namedtuple`` except that each parameter is defined as a property with a
907 docstring, making these objects suitable for interactive use without having
908 to open the source code to see parameter descriptions. Parameter names are
909 provided as keyword argument names (see below).
911 The main point of this factory is that it provides a compact way of
912 defining a new hashable class with descriptive docstrings; these features
913 are good for classes representing collections of parameters (hence the
914 function name).
916 The reason to implement this as a factory function rather than an abstract
917 superclass is that the returned class has static attributes with docstrings
918 available in a way that ipython can parse, making for easier interactive
919 use of the built-in documentation.
921 Parameters
922 ----------
923 classname : str
924 The name of the new class.
925 description : str
926 A description of what sort of analysis these parameters are to be used
927 for. Becomes part of the docstring of the new class (along with the
928 descriptions included in ``parameters``).
929 kwargs : str or function
930 The names and descriptions (or, for derived values, formulas) of
931 parameters to be stored. The names of the arguments become the
932 parameter names; these will be attributes of the returned class. The
933 values of the arguments
934 can be strings serving as descriptions of the parameter (for
935 INDEPENDENT parameters provided by the user, in which case the
936 descriptive string becomes the docstring for the parameter in the
937 returned class) or functions that take ``self`` as
938 their only argument where ``self`` is an instance of the new parameter
939 class and some of the attributes of said instance are
940 used to calculate the value of this (DEPENDENT) parameter; in these
941 cases, the docstring of the function is reused for that parameter.
942 These dependent parameters become properties of the new class and are
943 calculated dynamically from the class's other values (possibly
944 including other dependent parameters as long as there are no circular
945 method calls).
947 Returns
948 -------
949 new_class : type
950 a class whose name is ``classname``, whose properties are all the
951 parameter names, whose docstrings are their descriptions, whose
952 docstring is the class ``description`` combined with the names and
953 descriptions in the ``kwargs`` parameters, and whose ``__init__``
954 signature requires that the parameters be passed as kwargs for the sake
955 of explicitness in instantiation.
957 Examples
958 --------
959 >>> def buz(self):
960 ... "bar+baz"
961 ... return self.bar + self.baz
962 >>> def boz(self):
963 ... "double buz"
964 ... return 2*self.buz
965 >>> Foo = parameter_factory(
966 ... "Foo",
967 ... "A silly example set of parameters.",
968 ... bar='a fake param',
969 ... baz='another fake param',
970 ... buz=buz,
971 ... boz=boz,
972 ... )
973 >>> Foo.bar.__doc__
974 '(independent)\na fake param'
975 >>> Foo.baz.__doc__
976 '(independent)\nanother fake param'
977 >>> Foo.buz.__doc__
978 '(dependent)\nbar+baz'
979 >>> Foo.boz.__doc__
980 '(dependent)\ndouble buz'
981 >>> quux = Foo(bar=1, baz=2)
982 >>> quux.bar
983 1
984 >>> quux.baz
985 2
986 >>> quux.buz
987 3
988 >>> quux.boz
989 6
990 """
991 docstr = description + ("\n\nAll properties listed are REQUIRED and must "
992 "be passed as kwargs.\n\nProperties\n----------\n")
993 parameters = kwargs
994 newclassdict = dict()
996 # create getters and docstrings for all parameters and a class docstring
997 # that mentions all parameters together
998 for parameter_name, value in parameters.items():
999 if callable(value):
1000 docstr += parameter_name + " : (dependent)\n"
1001 def newmethod(self, func=value):
1002 return func(self)
1003 doc = ("No description given." if value.__doc__ is None else
1004 value.__doc__)
1005 newmethod.__doc__ = "(dependent)\n" + doc
1006 docstr += ' {}\n'.format(('\n ').join(wrap(doc, width=75)))
1007 else:
1008 docstr += parameter_name + " : (independent)\n"
1009 def newmethod(self, pname=parameter_name):
1010 return getattr(self, '_'+pname)
1011 newmethod.__doc__ = "(independent)\n" + value
1012 docstr += ' {}\n'.format(('\n ').join(wrap(value, width=75)))
1013 # indent the description of each parameter
1014 newmethod.__name__ = parameter_name
1015 newclassdict[parameter_name] = property(newmethod)
1016 newclassdict['__doc__'] = docstr
1018 def init(self, **kwargs):
1019 for parameter_name, value in kwargs.items():
1020 if (not callable(value) and parameter_name in parameters and
1021 not callable(parameters[parameter_name])):
1022 setattr(self, '_'+parameter_name, value)
1024 newclassdict['__init__'] = init
1026 def _repr(self):
1027 return "{}({})".format(
1028 type(self).__name__,
1029 ", ".join(
1030 "{}={}".format(n, getattr(self, n))
1031 for n in sorted(parameters) if not callable(parameters[n])
1032 )
1033 )
1035 newclassdict['__repr__'] = _repr
1036 newclassdict['__str__'] = _repr
1038 def _iter(self):
1039 return ((p, getattr(self, p)) for p in sorted(parameters))
1041 newclassdict['__iter__'] = _iter
1043 def _hash(self):
1044 return hash(
1045 tuple(
1046 (n, getattr(self, n))
1047 for n in sorted(parameters) if not callable(parameters[n])
1048 )
1049 )
1051 newclassdict['__hash__'] = _hash
1053 def _eq(self, other):
1054 return hash(self) == hash(other) and type(self) == type(other)
1056 newclassdict['__eq__'] = _eq
1057 return type(classname, (object,), newclassdict)
1060def vectypes(func, types):
1061 """Implementation of ``vecstr`` and ``veccls``."""
1062 def wrapper(fh, query):
1063 if any(isinstance(query, t) for t in types):
1064 return func(fh, query)
1065 return any(func(fh, q) for q in query)
1066 return wrapper
1069def vecstr(func):
1070 """Given a function taking a ``FileHandler`` instance ``fh`` and a
1071 ``query`` argument and returning a ``bool``, return a function with the
1072 same signature that first checks whether ``query`` is a ``str`` before
1073 calling ``func``. If ``query`` is a ``str``, return ``func(fh, query)``,
1074 and if it is not, return ``any(func(fh, q) for q in query)``, i.e. check
1075 whether ``func`` is true for any of the values given in ``query``. Note
1076 that this means ``query`` can either be a check against a ``str`` or
1077 against several possible values of ``str`` contained in an iterable. This
1078 function is an implementation detail used for downselection checks."""
1079 return vectypes(func, string_types)
1082def veccls(func):
1083 """Like ``vecstr`` but for class arguments instead of string arguments."""
1084 return vectypes(func, [type])
1087def vecfh(func):
1088 """Like ``vecstr`` but for ``FileHandler`` instances."""
1089 return vectypes(func, [AbstractFileHandler])
1092def plot_graphviz(dot, outfile):
1093 """Plot a ``.dot`` graph (graph here used in the mathematical sense) and save
1094 it to ``outfile``. Calls graphviz ``dot`` executable, used for plotting
1095 ``.dot`` graphs, so graphviz must be installed on the system.
1097 Parameters
1098 ----------
1099 dot : str
1100 Graphviz ``.dot`` format graph (the same contents you would expect in a
1101 ``.dot`` file).
1102 outfile : str
1103 Output file to save graph to. File type is inferred from ``outfile``'s
1104 extension. This file type must be one of the ones recognized in
1105 ``GRAPH_EXTENSIONS``.
1107 Raises
1108 ------
1109 subprocess.CalledProcessError
1110 If the ``dot`` process fails to generate the output file. This
1111 exception object will have the nonzero return code in it.
1112 """
1113 if outfile is not None:
1114 if outfile.endswith('.dot'):
1115 with open(outfile, "w") as out:
1116 out.write(dot)
1117 return
1118 for ext in GRAPH_EXTENSIONS:
1119 if outfile.lower().endswith(ext):
1120 with tempfile.NamedTemporaryFile("w") as tempout:
1121 tempout.write(dot)
1122 tempout.flush()
1123 ext = outfile.split(".")[-1]
1124 check_call(['dot', '-T'+ext, tempout.name, '-o', outfile])
1125 return
1126 raise ValueError(("Unrecognized file extension {}, pick from "
1127 "{}").format(outfile, GRAPH_EXTENSIONS))
1130def sort_files_by_modification_time(paths):
1131 """Take the list of ``path`` strings sorted by ascending modification
1132 time. This is the last time the contents were modified, or, for a new file,
1133 its creation time."""
1134 entries = ((os.stat(path), path) for path in paths)
1135 regular_files = ((stat[ST_MTIME], path)
1136 for stat, path in entries if S_ISREG(stat[ST_MODE]))
1137 sorted_files = sorted(regular_files, key=lambda f: f[0])
1138 return [f[1] for f in sorted_files]
1141def find_in_submodules(modulename: str, predicate: FunctionType):
1142 """Return a dict of fully-qualified variable names mapping to all objects
1143 in ``modulename`` and its submodules for which ``predicate`` is True."""
1144 print("searching in ", modulename)
1145 matches = dict()
1146 locs = importlib.util.find_spec(modulename).submodule_search_locations
1147 for _importer, subname, ispkg in iter_modules(locs):
1148 fullname = f"{modulename}.{subname}"
1149 mod = importlib.import_module(fullname)
1150 for varname in dir(mod):
1151 item = getattr(mod, varname)
1152 if predicate(item):
1153 matches[f"{fullname}.{varname}"] = item
1154 if ispkg:
1155 matches.update(find_in_submodules(fullname, predicate))
1156 return matches
1159RemoteFileCacherTuple = namedtuple("RemoteFileCacherTuple",
1160 ("url", "localpath"))
1163class RemoteFileCacher(RemoteFileCacherTuple):
1164 """A simple loader for remotely-cached data accessible on public URLs. Use
1165 this to download cloud-based resources and locally cache them in {objdir}.
1166 Files will only be downloaded as-needed to save space and bandwidth.
1168 Parameters
1169 ----------
1170 url : str
1171 Specify the URL of the remote resource.
1172 localpath : str, optional
1173 The (optional) local path at which to cache this resource. By default,
1174 will just be ``{objdir}/filename`` where ``filename`` is actually
1175 taken from the remote URL filename.
1176 """
1178 __doc__ = __doc__.format(objdir=OBJECT_DIR)
1180 def __new__(cls, url, localpath=None):
1181 if localpath is None:
1182 from urllib.parse import urlparse
1183 localpath = Path(OBJECT_DIR) / Path(urlparse(url).path).parts[-1]
1184 else:
1185 localpath = Path(localpath)
1186 return RemoteFileCacherTuple.__new__(cls, url, localpath)
1188 def get(self, query=None):
1189 """If the file is not available locally, download it and store it at
1190 ``localpath`` (do nothing if present). Return ``localpath``. Optionally
1191 append a query string ``query``."""
1192 if not self.localpath.exists():
1193 if query is not None:
1194 url = self.url+"?"+query
1195 else:
1196 url = self.url
1197 LOGGER.info("File not cached locally, downloading %s -> %s",
1198 url, self.localpath)
1199 from llama.com.dl import download
1200 download(url, str(self.localpath.absolute()))
1201 return self.localpath
1204# https://stackoverflow.com/a/1094933/3601493
1205def sizeof_fmt(num, suffix='B'):
1206 """Format the size of a file in a human-readable way. ``num`` is the file
1207 size, presumed to be in bytes. Specify ``suffix='b'`` to indicate that the
1208 unit of ``num`` is bits."""
1209 for unit in ['', 'Ki', 'Mi', 'Gi', 'Ti', 'Pi', 'Ei', 'Zi']:
1210 if abs(num) < 1024.0:
1211 return "%3.1f%s%s" % (num, unit, suffix)
1212 num /= 1024.0
1213 return "%.1f%s%s" % (num, 'Yi', suffix)
1216def tbhighlight(tb):
1217 """
1218 Syntax-highlight a traceback for a 256-color terminal.
1220 Parameters
1221 ----------
1222 tb : str
1223 The traceback string (as provided by ``traceback.format_exc()``)
1225 Returns
1226 -------
1227 highlighted_tb : str
1228 The same traceback string with syntax highlighting applied.
1229 """
1230 return highlight(tb, TB_LEXER, TB_FORMATTER)