ucndata 5.0.5__py3-none-any.whl

ucndata/__init__.py

@@ -0,0 +1,55 @@
+# init
+# Derek Fujimoto
+# June 2024
+
+from .ucnrun import ucnrun
+from .ucncycle import ucncycle
+from .ucnperiod import ucnperiod
+from .ucnbase import ucnbase
+from .ttreeslow import ttreeslow
+from .applylist import applylist
+from .exceptions import *
+from .datetime import to_datetime, from_datetime
+
+# path to the directory which contains the root files
+DATADIR = "/data3/ucn/root_files"
+
+# detector tree names
+DET_NAMES = {'He3':{'hits':         'UCNHits_He3',
+                    'hits_orig':    'UCNHits_He3', # as saved in the root file
+                    'charge':       'He3_Charge',
+                    'rate':         'He3_Rate',
+                    'transitions':  'RunTransitions_He3',
+                    'hitsseq':      'hitsinsequence_he3',
+                    'hitsseqcumul': 'hitsinsequencecumul_he3',
+                    },
+            'Li6':{'hits':         'UCNHits_Li6',
+                   'hits_orig':    'UCNHits_Li-6', # as saved in the root file
+                   'charge':       'Li6_Charge',
+                   'rate':         'Li6_Rate',
+                   'transitions':  'RunTransitions_Li6',
+                   'hitsseq':      'hitsinsequence_li6',
+                   'hitsseqcumul': 'hitsinsequencecumul_li6',
+                   },
+            }
+
+# needed slow control trees: for checking data quality
+SLOW_TREES = ('BeamlineEpics', 'SequencerTree')
+
+# EPICS trees to group into a single slow control tree
+EPICS_TREES = [ 'BeamlineEpics',        'UCN2EpPha5Last',   'UCN2EpicsPhase2B',
+                'UCN2EpPha5Oth',        'UCN2EpicsPhase3',  'UCN2EpPha5Pre',
+                'UCN2EpicsPressures',   'UCN2EpPha5Tmp',    'UCN2EpicsTemperature',
+                'UCN2Epics',            'UCN2Pur',          'UCN2EpicsOthers',
+                'UCN2EpLD2' 
+              ]
+
+# data thresholds for checking data
+DATA_CHECK_THRESH = {'beam_min_current': 0.1, # uA
+                     'pileup_cnt_per_ms': 10, # if larger than this, then pileup and delete
+                     'pileup_within_first_s': 1, # time frame for pileup in each period
+                    }
+
+# precise cycle defaults
+DEFAULT_CYCLE_TIMES_PRECISE = { 'hw_channel': 10,
+                                'detector':   'Li6'}

ucndata/applylist.py

@@ -0,0 +1,205 @@
+# Array which syntatic sugar to apply functions to each element contained
+# Derek Fujimoto
+# Oct 2024
+
+import numpy as np
+
+class applylist(list):
+    """A list object with the following enhancements:
+
+    * An `apply` function: like in pandas, `apply` takes a function handle and applies it to every element in the list. This acts recursively, if the list has a depth of more than 1
+    * Element access: accessing attributes, if not an attribute of the `applylist`, instead try to fetch attributes of the contained objects. The same for functions. This also works recursively.
+    * Numpy-like array slicing: slicing with a `np.ndarray` first converts this object to an array, then does the slice, then converts back. This allows slicing on arrays of indices for re-ordering or booleans for selection based on criteria
+    * Arithmetic works element-wise as in numpy arrays
+
+    Examples:
+
+        >>> # Slicing
+        >>> x = applylist(range(10))
+        >>> print(x[3:7])
+        [3, 4, 5, 6]
+
+        >>> # Element access
+        >>> x = applylist([ucnrun(1846), ucnrun(1847)])
+        >>> print(x.run_number)
+        [1846, 1847]
+        >>> print(x.beam_current_uA.mean())
+        [np.float64(0.16612837637441483), np.float64(0.18927602913972205)]
+
+        >>> # Arithmetic and comparisons
+        >>> x = applylist([1,2,3])
+        >>> print(x*2)
+        [np.int64(2), np.int64(4), np.int64(6)]
+        >>> print(x>2)
+        [np.False_, np.False_, np.True_]
+    """
+
+    def __call__(self, *args, **kwargs):
+
+        # call function of sub-elements
+        return applylist([x.__call__(*args, **kwargs) for x in self])
+
+    # numpy-style comparisons
+    def __eq__(self, other):    return applylist(np.array(self).__eq__(other))
+    def __ne__(self, other):    return applylist(np.array(self).__ne__(other))
+    def __lt__(self, other):    return applylist(np.array(self).__lt__(other))
+    def __gt__(self, other):    return applylist(np.array(self).__gt__(other))
+    def __le__(self, other):    return applylist(np.array(self).__le__(other))
+    def __ge__(self, other):    return applylist(np.array(self).__ge__(other))
+
+    # numpy-style unary operators
+    def __pos__(self):          return applylist(np.array(self).__pos__())
+    def __neg__(self):          return applylist(np.array(self).__neg__())
+    def __abs__(self):          return applylist(np.array(self).__abs__())
+    def __invert__(self):       return applylist(np.array(self).__invert__())
+    def __round__(self):        return applylist(np.array(self).__round__())
+    def __floor__(self):        return applylist(np.array(self).__floor__())
+    def __ceil__(self):         return applylist(np.array(self).__ceil__())
+    def __trunc__(self):        return applylist(np.array(self).__trunc__())
+
+    # numpy-style arithmetic operators
+    def __add__(self, other):   return applylist(np.array(self).__add__(other))
+    def __sub__(self, other):   return applylist(np.array(self).__sub__(other))
+    def __mul__(self, other):   return applylist(np.array(self).__mul__(other))
+    def __floordiv__(self, other):   return applylist(np.array(self).__floordiv__(other))
+    def __div__(self, other):   return applylist(np.array(self).__div__(other))
+    def __truediv__(self, other):   return applylist(np.array(self).__truediv__(other))
+    def __mod__(self, other):   return applylist(np.array(self).__mod__(other))
+    def __divmod__(self, other):   return applylist(np.array(self).__divmod__(other))
+    def __pow__(self, other):   return applylist(np.array(self).__pow__(other))
+    def __and__(self, other):   return applylist(np.array(self).__and__(other))
+    def __or__(self, other):   return applylist(np.array(self).__or__(other))
+    def __xor__(self, other):   return applylist(np.array(self).__xor__(other))
+
+    def __getattr__(self, name):
+
+        # retain default behaviour
+        try:
+            return super().__getattr__(name)
+
+        # enhanced behaviour
+        except AttributeError as err:
+
+            # check for scalar numpy datatypes, all contents assumed to be of the same type
+            if isinstance(self[0], (int, float, np.integer, np.floating)) or '__array_' in name:
+                raise err from None
+
+            # try to get the contents of subarray
+            return applylist([getattr(x, name) for x in self])
+
+    def __getitem__(self, key):
+
+        # numpy-like slicing
+        if isinstance(key, (np.ndarray, tuple, list, applylist)):
+            key = np.asarray(key)
+
+            # boolean array
+            if key.dtype == bool:
+                return applylist([i for i, k in zip(self, key) if k])
+            
+            # index array
+            elif key.dtype == int:
+                return applylist([self[k] for k in key])
+
+        # default behaviour
+        else:
+            x = super().__getitem__(key)
+            if isinstance(x, list):
+                return applylist(x)
+            return x 
+
+    def astype(self, typecast):
+        """Convert all elements in-place to the given type.
+
+        Args:
+            typecast (type): type to cast each element to (e.g. int, float, str)
+
+        Returns:
+            None: modifies the list in-place
+
+        Example:
+
+            >>> x = applylist(np.arange(5))
+            >>> print(x)
+            [np.int64(0), np.int64(1), np.int64(2), np.int64(3), np.int64(4)]
+            >>> x.astype(int)
+            >>> print(x)
+            [0, 1, 2, 3, 4]
+        """
+        for i in range(len(self)):
+            self[i] = typecast(self[i])
+
+    def apply(self, fn, inplace=False):
+        """Apply function to each element contained, similar to pandas functionality
+
+        Args:
+            fn (function handle): function to apply to each element
+            inplace (bool): if `False` return a copy, else act in-place
+
+        Returns:
+            applylist|None: new `applylist` with results if `inplace` is `False`, else `None`
+
+        Examples:
+
+            >>> # With return value
+            >>> x = applylist(np.arange(5))
+            >>> y = x.apply(lambda a: a**2)
+            >>> print(x)
+            [np.int64(0), np.int64(1), np.int64(2), np.int64(3), np.int64(4)]
+            >>> print(y)
+            [np.int64(0), np.int64(1), np.int64(4), np.int64(9), np.int64(16)]
+
+            >>> # Inplace
+            >>> x = applylist(np.arange(5))
+            >>> x.apply(lambda a: a**2, inplace=True)
+            >>> print(x)
+            [np.int64(0), np.int64(1), np.int64(4), np.int64(9), np.int64(16)]
+        """
+
+        # decide if inplace
+        if inplace: copy = self
+        else:       copy = self.copy()
+
+        # do the operation
+        for i in range(len(copy)):
+            if isinstance(copy[i], applylist):
+                copy[i] = copy[i].apply(fn)
+            else:
+                copy[i] = fn(copy[i])
+
+        # return the copy
+        if not inplace: return copy
+
+    def copy(self):
+        """Return a shallow copy of this applylist.
+
+        Returns:
+            applylist: a new `applylist` with the same elements
+
+        Example:
+
+            >>> x = applylist([1, 2, 3])
+            >>> y = x.copy()
+            >>> y[0] = 99
+            >>> print(x)
+            [1, 2, 3]
+            >>> print(y)
+            [99, 2, 3]
+        """
+        return applylist(super().copy())
+
+    def transpose(self):
+        """Transpose a 2D applylist by converting to np.array and back.
+
+        Returns:
+            applylist: transposed list, where rows and columns are swapped
+
+        Example:
+
+            >>> x = applylist([[1,2,3], [4,5,6]])
+            >>> print(x)
+            [[1, 2, 3], [4, 5, 6]]
+            >>> print(x.transpose())
+            [array([1, 4]), array([2, 5]), array([3, 6])]
+        """
+        return applylist(np.array(self).transpose())