update calculate_limits.py to produce nicer plots

1 files changed, 343 insertions, 15 deletions

diff --git a/utils/calculate_limits.py b/utils/calculate_limits.py
index 464051a..3737850 100755
--- a/utils/calculate_limits.py
+++ b/utils/calculate_limits.py
@@ -14,18 +14,177 @@
 # You should have received a copy of the GNU General Public License along with
 # this program. If not, see <https://www.gnu.org/licenses/>.
 
+from __future__ import print_function, division
 import numpy as np
 from scipy.integrate import quad
 from scipy.stats import expon
 from scipy.special import spherical_jn, erf
 from numpy import pi
-from functools import lru_cache
+from collections import namedtuple
+from functools import update_wrapper
+from threading import RLock
+
+# Backport of lru_cache from http://code.activestate.com/recipes/578078-py26-and-py30-backport-of-python-33s-lru-cache/
+_CacheInfo = namedtuple("CacheInfo", ["hits", "misses", "maxsize", "currsize"])
+
+class _HashedSeq(list):
+    __slots__ = 'hashvalue'
+
+    def __init__(self, tup, hash=hash):
+        self[:] = tup
+        self.hashvalue = hash(tup)
+
+    def __hash__(self):
+        return self.hashvalue
+
+def _make_key(args, kwds, typed,
+             kwd_mark = (object(),),
+             fasttypes = {int, str, frozenset, type(None)},
+             sorted=sorted, tuple=tuple, type=type, len=len):
+    'Make a cache key from optionally typed positional and keyword arguments'
+    key = args
+    if kwds:
+        sorted_items = sorted(kwds.items())
+        key += kwd_mark
+        for item in sorted_items:
+            key += item
+    if typed:
+        key += tuple(type(v) for v in args)
+        if kwds:
+            key += tuple(type(v) for k, v in sorted_items)
+    elif len(key) == 1 and type(key[0]) in fasttypes:
+        return key[0]
+    return _HashedSeq(key)
+
+def lru_cache(maxsize=100, typed=False):
+    """Least-recently-used cache decorator.
+
+    If *maxsize* is set to None, the LRU features are disabled and the cache
+    can grow without bound.
+
+    If *typed* is True, arguments of different types will be cached separately.
+    For example, f(3.0) and f(3) will be treated as distinct calls with
+    distinct results.
+
+    Arguments to the cached function must be hashable.
+
+    View the cache statistics named tuple (hits, misses, maxsize, currsize) with
+    f.cache_info().  Clear the cache and statistics with f.cache_clear().
+    Access the underlying function with f.__wrapped__.
+
+    See:  http://en.wikipedia.org/wiki/Cache_algorithms#Least_Recently_Used
 
-# on retina screens, the default plots are way too small
-# by using Qt5 and setting QT_AUTO_SCREEN_SCALE_FACTOR=1
-# Qt5 will scale everything using the dpi in ~/.Xresources
-import matplotlib
-matplotlib.use("Qt5Agg")
+    """
+
+    # Users should only access the lru_cache through its public API:
+    #       cache_info, cache_clear, and f.__wrapped__
+    # The internals of the lru_cache are encapsulated for thread safety and
+    # to allow the implementation to change (including a possible C version).
+
+    def decorating_function(user_function):
+
+        cache = dict()
+        stats = [0, 0]                  # make statistics updateable non-locally
+        HITS, MISSES = 0, 1             # names for the stats fields
+        make_key = _make_key
+        cache_get = cache.get           # bound method to lookup key or return None
+        _len = len                      # localize the global len() function
+        lock = RLock()                  # because linkedlist updates aren't threadsafe
+        root = []                       # root of the circular doubly linked list
+        root[:] = [root, root, None, None]      # initialize by pointing to self
+        nonlocal_root = [root]                  # make updateable non-locally
+        PREV, NEXT, KEY, RESULT = 0, 1, 2, 3    # names for the link fields
+
+        if maxsize == 0:
+
+            def wrapper(*args, **kwds):
+                # no caching, just do a statistics update after a successful call
+                result = user_function(*args, **kwds)
+                stats[MISSES] += 1
+                return result
+
+        elif maxsize is None:
+
+            def wrapper(*args, **kwds):
+                # simple caching without ordering or size limit
+                key = make_key(args, kwds, typed)
+                result = cache_get(key, root)   # root used here as a unique not-found sentinel
+                if result is not root:
+                    stats[HITS] += 1
+                    return result
+                result = user_function(*args, **kwds)
+                cache[key] = result
+                stats[MISSES] += 1
+                return result
+
+        else:
+
+            def wrapper(*args, **kwds):
+                # size limited caching that tracks accesses by recency
+                key = make_key(args, kwds, typed) if kwds or typed else args
+                with lock:
+                    link = cache_get(key)
+                    if link is not None:
+                        # record recent use of the key by moving it to the front of the list
+                        root, = nonlocal_root
+                        link_prev, link_next, key, result = link
+                        link_prev[NEXT] = link_next
+                        link_next[PREV] = link_prev
+                        last = root[PREV]
+                        last[NEXT] = root[PREV] = link
+                        link[PREV] = last
+                        link[NEXT] = root
+                        stats[HITS] += 1
+                        return result
+                result = user_function(*args, **kwds)
+                with lock:
+                    root, = nonlocal_root
+                    if key in cache:
+                        # getting here means that this same key was added to the
+                        # cache while the lock was released.  since the link
+                        # update is already done, we need only return the
+                        # computed result and update the count of misses.
+                        pass
+                    elif _len(cache) >= maxsize:
+                        # use the old root to store the new key and result
+                        oldroot = root
+                        oldroot[KEY] = key
+                        oldroot[RESULT] = result
+                        # empty the oldest link and make it the new root
+                        root = nonlocal_root[0] = oldroot[NEXT]
+                        oldkey = root[KEY]
+                        oldvalue = root[RESULT]
+                        root[KEY] = root[RESULT] = None
+                        # now update the cache dictionary for the new links
+                        del cache[oldkey]
+                        cache[key] = oldroot
+                    else:
+                        # put result in a new link at the front of the list
+                        last = root[PREV]
+                        link = [last, root, key, result]
+                        last[NEXT] = root[PREV] = cache[key] = link
+                    stats[MISSES] += 1
+                return result
+
+        def cache_info():
+            """Report cache statistics"""
+            with lock:
+                return _CacheInfo(stats[HITS], stats[MISSES], maxsize, len(cache))
+
+        def cache_clear():
+            """Clear the cache and cache statistics"""
+            with lock:
+                cache.clear()
+                root = nonlocal_root[0]
+                root[:] = [root, root, None, None]
+                stats[:] = [0, 0]
+
+        wrapper.__wrapped__ = user_function
+        wrapper.cache_info = cache_info
+        wrapper.cache_clear = cache_clear
+        return update_wrapper(wrapper, user_function)
+
+    return decorating_function
 
 # speed of light (exact)
 SPEED_OF_LIGHT = 299792458 # m/s
@@ -328,8 +487,157 @@ def get_event_rate_xenon(m, cs, A, threshold):
 
     return rate
 
+# Taken from https://raw.githubusercontent.com/mwaskom/seaborn/c73055b2a9d9830c6fbbace07127c370389d04dd/seaborn/utils.py
+def despine(fig=None, ax=None, top=True, right=True, left=False,
+            bottom=False, offset=None, trim=False):
+    """Remove the top and right spines from plot(s).
+
+    fig : matplotlib figure, optional
+        Figure to despine all axes of, default uses current figure.
+    ax : matplotlib axes, optional
+        Specific axes object to despine.
+    top, right, left, bottom : boolean, optional
+        If True, remove that spine.
+    offset : int or dict, optional
+        Absolute distance, in points, spines should be moved away
+        from the axes (negative values move spines inward). A single value
+        applies to all spines; a dict can be used to set offset values per
+        side.
+    trim : bool, optional
+        If True, limit spines to the smallest and largest major tick
+        on each non-despined axis.
+
+    Returns
+    -------
+    None
+
+    """
+    # Get references to the axes we want
+    if fig is None and ax is None:
+        axes = plt.gcf().axes
+    elif fig is not None:
+        axes = fig.axes
+    elif ax is not None:
+        axes = [ax]
+
+    for ax_i in axes:
+        for side in ["top", "right", "left", "bottom"]:
+            # Toggle the spine objects
+            is_visible = not locals()[side]
+            ax_i.spines[side].set_visible(is_visible)
+            if offset is not None and is_visible:
+                try:
+                    val = offset.get(side, 0)
+                except AttributeError:
+                    val = offset
+                _set_spine_position(ax_i.spines[side], ('outward', val))
+
+        # Potentially move the ticks
+        if left and not right:
+            maj_on = any(
+                t.tick1line.get_visible()
+                for t in ax_i.yaxis.majorTicks
+            )
+            min_on = any(
+                t.tick1line.get_visible()
+                for t in ax_i.yaxis.minorTicks
+            )
+            ax_i.yaxis.set_ticks_position("right")
+            for t in ax_i.yaxis.majorTicks:
+                t.tick2line.set_visible(maj_on)
+            for t in ax_i.yaxis.minorTicks:
+                t.tick2line.set_visible(min_on)
+
+        if bottom and not top:
+            maj_on = any(
+                t.tick1line.get_visible()
+                for t in ax_i.xaxis.majorTicks
+            )
+            min_on = any(
+                t.tick1line.get_visible()
+                for t in ax_i.xaxis.minorTicks
+            )
+            ax_i.xaxis.set_ticks_position("top")
+            for t in ax_i.xaxis.majorTicks:
+                t.tick2line.set_visible(maj_on)
+            for t in ax_i.xaxis.minorTicks:
+                t.tick2line.set_visible(min_on)
+
+        if trim:
+            # clip off the parts of the spines that extend past major ticks
+            xticks = ax_i.get_xticks()
+            if xticks.size:
+                firsttick = np.compress(xticks >= min(ax_i.get_xlim()),
+                                        xticks)[0]
+                lasttick = np.compress(xticks <= max(ax_i.get_xlim()),
+                                       xticks)[-1]
+                ax_i.spines['bottom'].set_bounds(firsttick, lasttick)
+                ax_i.spines['top'].set_bounds(firsttick, lasttick)
+                newticks = xticks.compress(xticks <= lasttick)
+                newticks = newticks.compress(newticks >= firsttick)
+                ax_i.set_xticks(newticks)
+
+            yticks = ax_i.get_yticks()
+            if yticks.size:
+                firsttick = np.compress(yticks >= min(ax_i.get_ylim()),
+                                        yticks)[0]
+                lasttick = np.compress(yticks <= max(ax_i.get_ylim()),
+                                       yticks)[-1]
+                ax_i.spines['left'].set_bounds(firsttick, lasttick)
+                ax_i.spines['right'].set_bounds(firsttick, lasttick)
+                newticks = yticks.compress(yticks <= lasttick)
+                newticks = newticks.compress(newticks >= firsttick)
+                ax_i.set_yticks(newticks)
+
 if __name__ == '__main__':
-    import matplotlib.pyplot as plt
+    import argparse
+
+    parser = argparse.ArgumentParser("plot fit results")
+    parser.add_argument("--save", action="store_true", default=False, help="save plots")
+    args = parser.parse_args()
+
+    if args.save:
+        # default \textwidth for a fullpage article in Latex is 16.50764 cm.
+        # You can figure this out by compiling the following TeX document:
+        #
+        #    \documentclass{article}
+        #    \usepackage{fullpage}
+        #    \usepackage{layouts}
+        #    \begin{document}
+        #    textwidth in cm: \printinunitsof{cm}\prntlen{\textwidth}
+        #    \end{document}
+
+        width = 16.50764
+        width /= 2.54 # cm -> inches
+        # According to this page:
+        # http://www-personal.umich.edu/~jpboyd/eng403_chap2_tuftegospel.pdf,
+        # Tufte suggests an aspect ratio of 1.5 - 1.6.
+        height = width/1.5
+        FIGSIZE = (width,height)
+
+        import matplotlib.pyplot as plt
+
+        font = {'family':'serif', 'serif': ['computer modern roman']}
+        plt.rc('font',**font)
+
+        plt.rc('text', usetex=True)
+    else:
+        # on retina screens, the default plots are way too small
+        # by using Qt5 and setting QT_AUTO_SCREEN_SCALE_FACTOR=1
+        # Qt5 will scale everything using the dpi in ~/.Xresources
+        import matplotlib
+        matplotlib.use("Qt5Agg")
+
+        import matplotlib.pyplot as plt
+
+        # Default figure size. Currently set to my monitor width and height so that
+        # things are properly formatted
+        FIGSIZE = (13.78,7.48)
+
+        # Make the defalt font bigger
+        plt.rc('font', size=22)
+
+    plt.rcParams['figure.figsize'] = FIGSIZE
 
     ls = np.logspace(-1,8,1000)
 
@@ -341,7 +649,7 @@ if __name__ == '__main__':
 
     colors = plt.rcParams['axes.prop_cycle'].by_key()['color']
 
-    plt.figure()
+    fig = plt.figure(1)
     plt.subplot(111)
     plt.plot(ls, rate/np.max(rate),color=colors[0])
     plt.xlabel("Mediator Decay Length (m)")
@@ -351,8 +659,9 @@ if __name__ == '__main__':
     plt.axvline(x=R, ls='--', color=colors[3], label="Earth radius")
     plt.gca().set_xscale('log')
     plt.gca().set_xlim((ls[0], ls[-1]))
-    plt.title("Event Rate of Self-Destructing Dark Matter in SNO")
-    plt.legend()
+    plt.legend(framealpha=1.0)
+    despine(fig,trim=True)
+    plt.tight_layout()
 
     # threshold is ~5 keVr
     xenon_100t_threshold = 5e-6
@@ -394,7 +703,7 @@ if __name__ == '__main__':
     # from arxiv:1705.06655
     xenon_100t_livetime = 34.2 # days
 
-    plt.figure()
+    fig = plt.figure(2)
     plt.subplot(111)
     plt.gca().set_xscale('log')
     plt.gca().set_yscale('log')
@@ -414,12 +723,12 @@ if __name__ == '__main__':
     plt.clabel(CS6, inline=1, fmt=r"SNO ($\mathrm{L}_V$ = 1000 km)", fontsize=10)
     plt.xlabel(r"$m_\chi$ (GeV)")
     plt.ylabel(r"WIMP-nucleon scattering cross section ($\mathrm{cm}^2$)")
-    plt.title("Self-Destructing Dark Matter Limits")
+    plt.tight_layout()
 
     x = np.linspace(0,300e-6,1000)
     
     # reproducing Figure 2.1 in Tom Caldwell's thesis
-    plt.figure()
+    fig = plt.figure(3)
     plt.plot(x*1e6, list(map(lambda x: get_nuclear_form_factor(XENON_MASS, x)**2,x)), label="Xe")
     plt.plot(x*1e6, list(map(lambda x: get_nuclear_form_factor(NEON_MASS, x)**2,x)), label="Ne")
     plt.plot(x*1e6, list(map(lambda x: get_nuclear_form_factor(ARGON_MASS, x)**2,x)), label="Ar")
@@ -430,5 +739,24 @@ if __name__ == '__main__':
     plt.legend()
     plt.gca().set_ylim((1e-5,1))
     plt.gca().set_xlim((0,300))
-    plt.title("Nuclear Form Factors for various elements")
-    plt.show()
+    plt.tight_layout()
+
+    if args.save:
+        fig = plt.figure(1)
+        plt.savefig("sddm_event_rate.pdf")
+        plt.savefig("sddm_event_rate.eps")
+        fig = plt.figure(2)
+        plt.savefig("sno_sensitivity.pdf")
+        plt.savefig("sno_sensitivity.eps")
+        fig = plt.figure(3)
+        plt.savefig("nuclear_form_factors.pdf")
+        plt.savefig("nuclear_form_factors.eps")
+    else:
+        plt.figure(1)
+        plt.title("Event Rate of Self-Destructing Dark Matter in SNO")
+        plt.figure(2)
+        plt.title("Self-Destructing Dark Matter Limits")
+        plt.figure(3)
+        plt.title("Nuclear Form Factors for various elements")
+
+        plt.show()