forked from IPKM/nmreval
141 lines
4.2 KiB
Python
141 lines
4.2 KiB
Python
import pickle
|
|
from collections import OrderedDict
|
|
from typing import Tuple, Union
|
|
|
|
from ..data.nmr import FID, Spectrum
|
|
from ..data.points import Points
|
|
from ..fit.result import FitResult, FitResultCreator
|
|
from .read_old_nmr import HAS_BSDDB3, _read_file_v1
|
|
from ..lib.colors import Colors
|
|
from ..lib.lines import LineStyle
|
|
from ..lib.symbols import SymbolStyle
|
|
|
|
|
|
class NMRReader:
|
|
def __init__(self, fname: str = None, version: str = '0'):
|
|
self.filename = fname
|
|
|
|
self._set_version(version)
|
|
|
|
def __call__(self, fname, version: str = None):
|
|
self.filename = fname
|
|
self._set_version(version)
|
|
|
|
return self
|
|
|
|
def _set_version(self, vers: str):
|
|
if vers is None:
|
|
return
|
|
|
|
if vers == '-1':
|
|
self.version = 1
|
|
else:
|
|
self.version = 2
|
|
|
|
def make_data(self, fname: str = None) -> Union[Tuple[dict, dict], Tuple[dict]]:
|
|
if fname is None:
|
|
fname = self.filename
|
|
|
|
if fname:
|
|
if self.version == 1:
|
|
return self._make_old(fname),
|
|
else:
|
|
return self._make_new(fname)
|
|
|
|
@staticmethod
|
|
def _make_new(fname) -> Tuple[dict, dict]:
|
|
with open(fname, 'rb') as fp:
|
|
# remove magic
|
|
_ = fp.read(16)
|
|
states = pickle.load(fp)
|
|
|
|
datalist = OrderedDict()
|
|
_dtypes = {'pts': Points, 'fit': FitResult, 'fid': FID}
|
|
|
|
for s in states['sets']:
|
|
set_id = s.pop('id')
|
|
dtype = _dtypes[s.pop('mode')]
|
|
data = dtype.set_state(s.pop('data'))
|
|
datalist[set_id] = (data, s)
|
|
|
|
return datalist, states['graphs']
|
|
|
|
def _make_old(self, fname: str) -> dict:
|
|
if not HAS_BSDDB3:
|
|
raise IOError('Legacy .nmr cannot be read without bsddb3')
|
|
|
|
datadic = _read_file_v1(fname)
|
|
datalist = OrderedDict()
|
|
|
|
for w_key, window in datadic.items():
|
|
order = window['data']['order']
|
|
for data_id in order:
|
|
if data_id in ['selected', 'order']:
|
|
continue
|
|
|
|
properties = window['data'][data_id]
|
|
|
|
data_properties = properties['_data']
|
|
|
|
meta = data_properties['meta']
|
|
meta['name'] = meta['value']
|
|
meta.pop('value')
|
|
|
|
mode = data_properties['mode'].lower()
|
|
if mode == 'points':
|
|
dtype = Points
|
|
elif mode == 'spectrum':
|
|
dtype = Spectrum
|
|
else:
|
|
dtype = FID
|
|
|
|
data = dtype(x=data_properties['x'], y=data_properties['y'],
|
|
y_err=data_properties['y_err'],
|
|
**meta)
|
|
|
|
data.mask = data_properties['mask'].data
|
|
datalist[data_id] = data, self.plot_opts(properties['dataplot_real'], properties['dataplot_imag'])
|
|
|
|
_fit = properties['_fit']
|
|
if _fit is not None:
|
|
ff = FitResultCreator.make_from_session(data_properties['x'], data_properties['y'], data_id, _fit)
|
|
datalist[data_id+'-fit'] = ff, self.plot_opts(properties['fitplot'], {})
|
|
|
|
return datalist
|
|
|
|
@staticmethod
|
|
def plot_opts(real_dic: dict, imag_dic: dict) -> dict:
|
|
opts = {}
|
|
|
|
pen = real_dic['pen']
|
|
symsize = real_dic['symbolSize']
|
|
symbol = SymbolStyle.from_str(real_dic['symbol'])
|
|
symcolor = None
|
|
|
|
try:
|
|
symcolor = Colors(real_dic['symbolBrush'][:3])
|
|
except ValueError:
|
|
pass
|
|
|
|
if pen is None:
|
|
ls = LineStyle.No
|
|
lw = 1.
|
|
lc = symcolor
|
|
else:
|
|
# pen is not None
|
|
lc = Colors(real_dic['pen'][0][:3])
|
|
if symcolor is None:
|
|
symcolor = lc
|
|
|
|
lw = real_dic['pen'][1]
|
|
ls = LineStyle.Solid
|
|
|
|
opts['real'] = {'symbol': symbol, 'size': symsize, 'color': symcolor}, \
|
|
{'style': ls, 'width': lw, 'color': lc}
|
|
|
|
if imag_dic:
|
|
opts['imag'] = {'symbol': symbol, 'size': symsize, 'color': symcolor}, \
|
|
{'style': LineStyle.Dashed, 'width': lw, 'color': lc}
|
|
|
|
return opts
|