mcutils/xray/utils.py

from __future__ import print_function,division

import logging as log
log.basicConfig(level=log.INFO)

import numpy as np
np.seterr(all='ignore')
import os
import glob
import pathlib
import re
from . import storage as storage

try:
  import matplotlib.pyplot as plt
except ImportError:
  log.warn("Can't import matplotlib !")

_time_regex =      re.compile( "(-?\d+\.?\d*(?:ps|ns|us|ms)?)")
_timeInStr_regex = re.compile("_(-?\d+\.?\d*(?:ps|ns|us|ms)?)")

def getDelayFromString(string) :
  match = _timeInStr_regex_regex.search(string)
  return match and match.group(1) or None

_time_regex = re.compile("(-?\d+\.?\d*)((?:s|fs|ms|ns|ps|us)?)")
def strToTime(delay) :
  _time2value = dict( fs = 1e-15, ps = 1e-12, ns = 1e-9, us = 1e-6, ms = 1e-3, s = 1)

  match = _time_regex.search(delay)
  if match:
    n,t = float(match.group(1)),match.group(2)
    value = _time2value.get(t,1)
    return n*value
  else:
    return None

def timeToStr(delay,fmt="%+.0f"):
  a_delay = abs(delay)
  if a_delay >= 1:
    ret = fmt % delay + "s"
  elif 1e-3 <= a_delay < 1: 
    ret = fmt % (delay*1e3) + "ms"
  elif 1e-6 <= a_delay < 1e-3: 
    ret = fmt % (delay*1e6) + "us"
  elif 1e-9 <= a_delay < 1e-6: 
    ret = fmt % (delay*1e9) + "ns"
  elif 1e-12 <= a_delay < 1e-9: 
    ret = fmt % (delay*1e12) + "ps"
  elif 1e-15 <= a_delay < 1e-12: 
    ret = fmt % (delay*1e12) + "fs"
  elif 1e-18 <= a_delay < 1e-15: 
    ret = fmt % (delay*1e12) + "as"
  return ret

def removeExt(fname):
  """ special remove extension meant to work with compressed files.edf and .edf.gz files """
  if fname[-3:] == ".gz": fname = fname[-3:]
  return os.path.splitext(fname)[0]

def getBasename(fname):
  return removeExt(os.path.basename(fname));

def saveTxt(fname,q,i,e=None,headerv=None,info=None,overwrite=True):
  """ Write data to file 'fname' in text format.
      Inputs:
        q = x vector
        i = 1D or 2D
        e = 1D (discarded when i is 2D)
        info = dictionary (saved as '# key : value') or string
        headerv = vector to be used as header or string
  """
  if os.path.exists(fname) and not overwrite:
    log.warn("File %s exists, returning",fname)
    return
  if isinstance(info,dict):
    header = [ "# %s : %s" %(k,str(v)) for (k,v) in info.items() ]
    header = "\n".join(header); # skip first #, will be added by np
  elif isinstance(info,str):
    header = info
  else:
    header = ""
  if isinstance(headerv,str): header += "\n%s" % headerv
  if i.ndim == 1:
    x = np.vstack( (q,i,e) ) if e is not None else np.vstack( (q,i) )
  if i.ndim == 2:
    x = np.vstack( (q,i,) )
    if headerv is not None:
      headerv = np.concatenate(( (i.shape[1],),headerv))
      x = np.hstack( (headerv[:,np.newaxis],x) )
  np.savetxt(fname,x.T,fmt="%+10.5e",header=header,comments='')

class data_storage(dict):
  """ Storage for 1d integrated info """
  def __init__(self,fileOrDict):
    if isinstance(fileOrDict,dict):
      self.filename = None
      d = fileOrDict
    else:
      assert isinstance(fileOrDict,str)
      if os.path.isdir(fileOrDict): fileOrDict = fileOrDict + "/pyfai_1d.h5"
      self.filename = fileOrDict
      d = storage.read(fileOrDict)

    # allow accessing with .data, .delays, etc.
    for k,v in d.items(): setattr(self,k,v)

    # allow accessing as proper dict
    self.update( **dict(d) )

  def __setitem__(self, key, value):
    setattr(self,key,value)
    super().__setitem__(key, value)

  def __delitem__(self, key):
    delattr(self,key)
    super().__delitem__(key)

  def save(self,fname=None):
    if fname is None: fname = self.filename
    assert fname is not None
    storage.save(fname,dict(self))
 
  #def asdict(self): return dict(self)


def reshapeToBroadcast(what,ref):
  """ expand the 1d array 'what' to allow broadbasting to match 
      multidimentional array 'ref'. The two arrays have to same the same 
      dimensions along the first axis
  """
  assert what.shape[0] == ref.shape[0]
  shape  = [ref.shape[0],] + [1,]*(ref.ndim-1)
  return what.reshape(shape)