nmreval/src/nmreval/io/graceeditor.py

755 lines
25 KiB
Python

from __future__ import annotations
import pathlib
import re
from io import StringIO
from typing import Optional, Tuple
import numpy as np
from numpy import log10
from ..configs import config_paths
from ..lib.logger import logger
class GraceEditor:
_RE_HEADER_START = re.compile(r'# Grace project file')
_RE_HEADER_END = re.compile(r'@timestamp def')
_RE_OBJECT_START = re.compile(r'@with (box|ellipse|string|)')
_RE_REGION_START = re.compile(r'@r(\w+) (on|off)', re.IGNORECASE)
_RE_GRAPH_START = re.compile(r'@g(\w+) (on|off)', re.IGNORECASE)
_RE_SET_START = re.compile(r'@target G(\d+).S(\d+)', re.IGNORECASE)
_RE_SET_END = re.compile(r'^\s*&\s*$')
_RE_COLOR = re.compile(r'@map color (?P<id>\d+) to '
r'\((?P<red>\d+), (?P<green>\d+), (?P<blue>\d+)\),\s+\"(?P<disp>.+)\"',
re.MULTILINE)
_RE_FONT = re.compile(r'@map font (?P<id>\d+) to \"(?P<name>.+?)\",\s+\"(?P<disp>.+)\"')
_RE_PAGE_SIZE = re.compile(r'@page size (\d+), (\d+)')
def __init__(self, filename=None):
self.header = GraceHeader()
self.drawing_objects = []
self.regions = []
self.graphs = []
self.sets = []
_default_file = config_paths().joinpath('Default.agr')
if filename == _default_file:
self.defaults = self
else:
self.defaults = GraceEditor(filename=_default_file)
self.file = filename
if filename is not None:
self.parse(filename)
def _fix_tuda_colors(self):
# 2023-05-11: Default.agr had wrong TUDa colors (4a, 7b, 9b, 9d), so set the values given in colors.py
color_mapping = [
('tuda4a', (175, 204, 80)),
('tuda7b', (245, 163, 0)),
('tuda9b', (230, 0, 26)),
('tuda9d', (156, 28, 38)),
]
for i, line in enumerate(self.header):
m = self._RE_COLOR.match(line)
if m:
for name, right_color in color_mapping:
if m['disp'].lower() == name:
self.header[i] = f'@map color {m["id"]} to {right_color}, "{m["disp"]}"\n'
def __call__(self, filename: str):
self.clear()
self.parse(filename)
return self
def __str__(self):
s = str(self.header)
s += ''.join(map(str, self.regions))
s += ''.join(map(str, self.drawing_objects))
s += ''.join(map(str, self.graphs))
s += ''.join(map(str, self.sets))
return s
def clear(self):
self.header = GraceHeader()
self.drawing_objects = []
self.regions = []
self.graphs = []
self.sets = []
self.graphs = []
def new_graph(self):
if self.file is None:
self.parse(config_paths().joinpath('Default.agr'))
else:
max_idx = 0
for g in self.graphs:
max_idx = max(g.idx, max_idx)
self.graphs.append(self.defaults.graphs[-1].copy())
self.graphs[-1].renumber(max_idx+1)
return self.graphs[-1]
def new_set(self, graph):
s = None
g_idx = -1
for g in self.graphs:
if g.idx == graph:
s = g.new_set()
g_idx = g.idx
break
if s is None:
raise ValueError(f'Graph {graph} was not found.')
self.sets.append(GraceSet(g_idx, s.idx))
self.sets[-1].append(f'@target G{g_idx}.S{s.idx}\n')
self.sets[-1].append('@type xy\n')
self.sets[-1].append('&\n')
return s
def parse(self, filename: str | pathlib.Path):
self.file = pathlib.Path(filename)
# we start always with the header
current_pos = 'header'
with self.file.open('r', errors='replace') as f:
for line_nr, line in enumerate(f):
# lots of states to check
# agr file order: header, drawing object, region, graph, set
if current_pos == 'header':
self.header.append(line)
# already at end of header
if self._RE_HEADER_END.match(line):
current_pos = 'header_end'
elif self._RE_GRAPH_START.match(line):
current_pos = self._make_graph(line)
elif current_pos == 'header_end':
# what comes after the header? region, graph or drawing object?
if self._RE_REGION_START.match(line):
current_pos = 'region'
self.regions.append(GraceRegion())
self.regions[-1].append(line)
elif self._RE_GRAPH_START.match(line):
current_pos = self._make_graph(line)
elif self._RE_OBJECT_START.match(line):
current_pos = 'drawing'
self.drawing_objects.append(GraceDrawing())
self.drawing_objects[-1].append(line)
else:
logger.warn('What happened with this line?', line.rstrip())
elif current_pos == 'drawing':
if self._RE_REGION_START.match(line):
current_pos = 'region'
self.regions.append(GraceRegion())
self.regions[-1].append(line)
else:
m = self._RE_GRAPH_START.match(line)
if m:
current_pos = 'graph'
self.graphs.append(GraceGraph(int(m.group(1))))
self.graphs[-1].append(line)
else:
if self._RE_OBJECT_START.match(line):
self.drawing_objects.append(GraceDrawing())
self.drawing_objects[-1].append(line)
elif current_pos == 'region':
# regions are followed by regions or graphs
if self._RE_GRAPH_START.match(line):
current_pos = self._make_graph(line)
self.graphs[-1].append(line)
else:
if self._RE_REGION_START.match(line):
self.regions.append(GraceRegion())
self.regions[-1].append(line)
elif current_pos == 'graph':
m = self._RE_SET_START.match(line)
if m:
current_pos = 'set'
self.sets.append(GraceSet(int(m.group(1)), int(m.group(2))))
self.sets[-1].append(line)
else:
if self._RE_GRAPH_START.match(line):
current_pos = self._make_graph(line)
self.graphs[-1].append(line)
elif current_pos == 'set':
m = self._RE_SET_END.match(line)
if m:
current_pos = 'set_end'
self.sets[-1].append(line)
elif current_pos == 'set_end':
m = self._RE_SET_START.match(line)
if m:
current_pos = 'set'
self.sets.append(GraceSet(int(m.group(1)), int(m.group(2))))
self.sets[-1].append(line)
else:
if self._RE_GRAPH_START.match(line):
current_pos = self._make_graph(line)
else:
if GraceGraph._RE_SET_START.match(line):
current_pos = 'graph'
else:
raise ValueError(f'Cannot parse line {line}')
self.graphs[-1].append(line)
self._fix_tuda_colors()
def _make_graph(self, line: str):
m = self._RE_GRAPH_START.match(line)
g_idx = int(m.group(1))
if g_idx < len(self.graphs):
# this assumes that graphs are ordered in agr file even if missing, e.g. we read gß, g1, g3, ...
while g_idx != len(self.graphs):
self.graphs.append(GraceGraph(len(self.graphs)))
self.graphs.append(GraceGraph(g_idx))
return 'graph'
@property
def colors(self):
_colors = {}
for line in self.header:
m = self._RE_COLOR.match(line)
if m:
_colors[int(m['id'])] = (m['disp'], (int(m['red']), int(m['green']), int(m['blue'])))
return _colors
def get_color(self, color_num):
color_num = str(color_num)
for line in self.header:
m = self._RE_COLOR.match(line)
if m:
if m['id'] == color_num:
return m['disp'], (int(m['red']), int(m['green']), int(m['blue']))
return
def set_color(self, name: str, rgb: tuple, idx: int = None):
pos = 'before_colors'
cnt = 0
for i, line in enumerate(self.header):
m = self._RE_COLOR.match(line)
if m:
pos = 'in_colors'
cnt += 1
if (int(m['red']), int(m['green']), int(m['blue'])) == rgb:
logger.info(f'color already defined as {m["disp"]}')
return m['id']
elif m['id'] == idx:
self.header[i] = f'@map color {idx} to {rgb}, "{name}"\n'
elif pos == 'in_colors':
# landing here for first mismatch after all colors
if idx is None:
if cnt >= 64:
logger.warn('Maximum numbers of color reached, color is not used by xmgrace')
self.header.insert(i, f'@map color {cnt} to {rgb}, "{name}"\n')
return cnt
else:
self.header.insert(i, f'@map color {idx} to {rgb}, "{name}"\n')
return idx
@property
def fonts(self):
_fonts = {}
for line in self.header:
m = self._RE_FONT.match(line)
if m:
_fonts[m['id']] = m['disp']
return _fonts
@property
def size(self):
for line in self.header:
m = self._RE_PAGE_SIZE.match(line)
if m:
return int(m.group(1)) / 72 * 2.54, int(m.group(2)) / 72 * 2.54
@size.setter
def size(self, value: tuple):
width = value[0] * 72 / 2.54
height = value[1] * 72 / 2.54
for i, line in enumerate(self.header):
m = self._RE_PAGE_SIZE.match(line)
if m:
self.header[i] = f'@page size {width:.0f}, {height:.0f}\n'
break
def convert(self, value, direction='rel'):
page_size = min(self.size)
if direction.startswith('rel'):
return value / page_size
elif direction.startswith('abs'):
return value * page_size
else:
raise ValueError(f'Use `abs`, `absolute`, `rel`, or `relative` to convert, found {direction} instead.')
def get_property(self, g, *args):
if len(args) == 1:
return self.graphs[g].get_property(args[0])
elif len(args) == 2:
return self.graphs[g].set[args[0]].get_property(args[1])
else:
raise TypeError(f'get_property takes two or three arguments ({len(args)+1} given)')
def set_property(self, g, *args):
if len(args) == 2:
self.graphs[g][args[0]] = args[1]
elif len(args) == 3:
self.graphs[g].set[args[0]][args[1]] = args[2]
else:
raise TypeError(f'get_property takes three or four arguments ({len(args)+1} given)')
def dataset(self, g, s):
cmp_to = (int(g), int(s))
for dataset in self.sets:
if dataset.graph_set == cmp_to:
return dataset
return
def write(self, fname):
outfile = pathlib.Path(fname)
with outfile.open('w') as f:
f.write(str(self.header))
f.flush()
for element in [self.regions, self.drawing_objects, self.graphs, self.sets]:
for obj in element:
f.write(str(obj))
f.flush()
class GraceDrawing(list):
def __str__(self):
return ''.join(self)
class GraceHeader(list):
def __str__(self):
return ''.join(self)
class GraceProperties(list):
_RE_ENTRY = re.compile(r'(?!.*(on|off)$)' # ignore lines that end with on or off
r'@\s*(?P<graph>[gs]\d+)?\s*' # @ maybe followed by g0 or s12
r'(?P<key>[\w\s]*)\s+' # key: stops at last space unless comma-separated values
r'(?P<val>(?:\s*[\d\w.+-]+\s*,)*\s*[\\{}\"\w.+\- ]+)', # one value, maybe more with commas
re.IGNORECASE | re.VERBOSE)
_RE_ONOFF = re.compile(r'@\s(?P<graph>[gs]\d+)*\s*(?P<val>[\w\s]*)\s+(on|off)')
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def get_property(self, key: str):
props = []
for line in self:
m = self._RE_ENTRY.match(line)
if m:
k = m.group('key').strip()
if k == key:
props.append(_convert_to_value(m.group('val')))
if len(props) == 1:
return props[0]
else:
return
def set_property(self, **kwargs):
raise NotImplementedError
def _set_property(self, **kwargs):
for i, line in enumerate(self):
m = self._RE_ENTRY.match(line)
if m:
k = m.group('key')
if k in kwargs:
# re.sub is probably better, but I could not work it out for labels with backslashes (\S)
self[i] = line[:m.start('val')] + str(kwargs[k]) + line[m.end('val'):]
kwargs.pop(k)
if not kwargs:
break
return kwargs
def get_all_properties(self):
props = {}
for line in self:
m = self._RE_ENTRY.match(line)
if m:
key = m.group('key')
val = _convert_to_value(m.group('val'))
if key not in props:
props[key] = val
elif isinstance(props[key], tuple):
props[key] = props[key] + (val,)
else:
props[key] = props[key], val
props.pop('with')
return props
def get_onoff(self, key: str) -> Optional[bool]:
onoff_pattern = re.compile(rf'@\s*{key}\s+(?P<val>(on|off))', re.IGNORECASE)
for i, line in enumerate(self):
m = onoff_pattern.match(line)
if m:
return m.group('val') == 'on'
return
def set_onoff(self, key: str, onoff: bool):
raise NotImplementedError
def _set_onoff(self, key , onoff: bool):
bool_str = 'on' if onoff else 'off'
for i, line in enumerate(self):
m = key.match(line)
if m:
self[i] = line[:m.start('val')] + bool_str + '\n'
return True
return False
class GraceGraph(GraceProperties):
_RE_SET_START = re.compile(r'@\s*s(\d+)\s+\w*', re.IGNORECASE)
def __init__(self, idx):
super().__init__()
self.set = []
self.idx = int(idx)
self.__current_pos = 'graph'
self.__current_idx = None
def __str__(self):
return ''.join(self) + ''.join(map(str, self.set))
def copy(self):
_cp = GraceGraph(self.idx)
for line in self:
_cp.append(line)
return _cp
def append(self, line: str):
m = self._RE_SET_START.match(line)
if m:
if self.__current_pos == 'graph' or self.__current_idx != m.group(1):
self.__current_pos = 'set'
self.__current_idx = m.group(1)
while int(self.__current_idx) != len(self.set):
self.set.append(GraceSetProps(len(self.set)))
self.set.append(GraceSetProps(self.__current_idx))
if self.__current_pos == 'graph':
super().append(line)
elif self.__current_pos == 'set':
self.set[-1].append(line)
else:
raise TypeError('Unknown state')
def new_set(self):
max_set_idx = -1
for s in self.set:
if s is None:
continue
max_set_idx = max(max_set_idx, s.idx)
self.set.append(GraceSetProps(max_set_idx+1))
return self.set[-1]
def renumber(self, idx: int):
number_re = re.compile('g' + str(self.idx), re.IGNORECASE)
for i, line in enumerate(self):
m = list(number_re.finditer(line))
if m:
self[i] = line[:m[0].start()+1] + str(idx) + line[m[0].end():]
self.idx = idx
def set_limits(self, x=None, y=None):
for i, line in enumerate(self):
m = self._RE_ENTRY.match(line)
if m and m.group('key') == 'world':
_prev_lims = _convert_to_value(m.group('val'))
if x is None:
x = _prev_lims[0]
if y is None:
y = _prev_lims[1]
self[i] = f'@ world {x[0]}, {y[0]}, {x[1]}, {y[1]}\n'
def set_log(self, x=None, y=None):
kwargs = {}
for ax, label in [(x, 'x'), (y, 'y')]:
if ax is not None:
if ax:
kwargs[label + 'axes scale'] = 'Logarithmic'
kwargs[label + 'axis ticklabel format'] = 'power'
kwargs[label + 'axis ticklabel prec'] = '0'
else:
kwargs[label + 'axes scale'] = 'Normal'
kwargs[label + 'axis ticklabel format'] = 'general'
kwargs[label + 'axis ticklabel prec'] = '5'
for i, line in enumerate(self):
m = self._RE_ENTRY.match(line)
if m and m.group('key') in kwargs:
self[i] = re.sub(m.group('val'), kwargs[m.group('key')], line)
def set_property(self, **kwargs):
remaining_kw = self._set_property(**kwargs)
for k, v in remaining_kw.items():
self.append(f'@ {k} {v}\n')
def set_axis_property(self, axis: str, **kwargs):
ax_kwargs = {}
for k, v in kwargs.items():
if k in ['invert', 'scale']:
ax_kwargs[axis + 'axes ' + k] = v
else:
ax_kwargs[axis + 'axis ' + k] = v
self.set_property(**ax_kwargs)
def get_axis_property(self, axis: str, key: str):
if key in ['invert', 'scale']:
ax_key = axis + 'axes ' + key
else:
ax_key = axis + 'axis ' + key
return self.get_property(ax_key)
def set_axis_onoff(self, axis: str, key: str, onoff: bool):
if key in ['invert', 'scale']:
ax_key = axis + 'axes ' + key
else:
ax_key = axis + 'axis ' + key
return self.set_onoff(ax_key, onoff)
def set_onoff(self, key: str, onoff: bool):
onoff_pattern = re.compile(rf'@\s*{key}\s+(?P<val>(on|off))', re.IGNORECASE)
if not self._set_onoff(onoff_pattern, onoff):
self.append(f'@ {key} {"on" if onoff else "off"}\n')
def view_to_world(self, pos) -> Tuple[float, float]:
view = self.get_property('view')
world = self.get_property('world')
log_x = self.get_property('xaxes scale').lower()
if log_x == 'logarithmic':
x_world = 10**(log10(world[2]/world[0]) / (view[2]-view[0]) * (view[2]-pos[0]) + log10(world[2]))
else:
x_world = (world[2]-world[0]) / (view[2]-view[0]) * (view[2]-pos[0]) + world[2]
log_y = self.get_property('yaxes scale').lower()
if log_y == 'logarithmic':
y_world = 10**(log10(world[3]/world[1]) / (view[3]-view[1]) * (view[3]-pos[1]) + log10(world[3]))
else:
y_world = (world[3] - world[1]) / (view[3]-view[1]) * (view[3]-pos[1]) + world[3]
return x_world, y_world
def world_to_view(self, pos) -> Tuple[float, float]:
view = self.get_property('view')
world = self.get_property('world')
view_pos = pos[:]
log_axes = [self.get_property('xaxes scale'), self.get_property('yaxes scale')]
for i in [0, 1]:
if log_axes[i].lower() == 'logarithmic':
view_pos[i] = (view[i+2]-view[i]) / log10(world[i+2]/world[i]) * log10(pos[i]/world[i]) + view[i]
else:
view_pos[i] = (view[i+2]-view[i]) / (world[i+2]-world[i]) * (pos[i]-world[i]) + view[i]
return view_pos
class GraceSetProps(GraceProperties):
def __init__(self, idx):
super().__init__()
self.idx = int(idx)
def __str__(self):
return ''.join(self)
def set_property(self, **kwargs):
remaining_kw = self._set_property(**kwargs)
for k, v in remaining_kw.items():
self.append(f'@ s{self.idx} {k} {v}\n')
def set_onoff(self, key: str, onoff: bool):
onoff_pattern = re.compile(rf'@\s*s\d{1,2}\s+{key}\s+(?P<val>(on|off))', re.IGNORECASE)
if not self._set_onoff(onoff_pattern, onoff):
self.append(f'@ s{self.idx} {key} {"on" if onoff else "off"}\n')
def set_line(self, **kwargs):
_kwargs = {'line '+k: v for k, v in kwargs.items()}
self.set_property(**_kwargs)
def set_symbol(self, **kwargs):
_kwargs = {'symbol '+k: v for k, v in kwargs.items()}
if 'symbol' in kwargs:
_kwargs['symbol'] = kwargs['symbol']
_kwargs.pop('symbol symbol')
self.set_property(**_kwargs)
class GraceSet(list):
_RE_TYPE = re.compile(r'@type (\w+)')
column_types = {'xy': 2, 'xydx': 3, 'xydy': 3, 'xydxdx': 4, 'xydydy': 4, 'xydxdy': 4, 'xydxdxdydy': 6,
'bar': 2, 'bardy': 3, 'bardydy': 4,
'xyhilo': 5, 'xyz': 3,
'xyr': 3, 'xysize': 3, 'xycolor': 3, 'xycolpat': 4,
'xyvmap': 4, 'xyboxplot': 6}
def __init__(self, g, s):
super().__init__()
self._graph = int(g)
self._set = int(s)
def __str__(self):
return ''.join(self)
@property
def graph(self):
return self._graph
@graph.setter
def graph(self, new_graph: int):
self._graph = new_graph
self[0] = f'@target G{new_graph}.S{self._set}\n'
@property
def set(self):
return self._set
@set.setter
def set(self, new_set: int):
self._set = int(new_set)
self[0] = f'@target G{self._graph}.S{new_set}\n'
@property
def graph_set(self):
return self._graph, self._set
@graph_set.setter
def graph_set(self, new_idx: tuple):
self._graph, self._set = new_idx
self[0] = f'@target G{new_idx[0]}.S{new_idx[1]}\n'
@property
def type(self):
m = GraceSet._RE_TYPE.match(self[1])
if m:
return m.group(1)
@type.setter
def type(self, new_type):
if new_type not in GraceSet.column_types:
raise ValueError(f'Unknown plot type {new_type}.')
self[1] = f'@type {new_type}\n'
@property
def rows(self):
return len(self) - 3
@property
def columns(self):
if self[2] != '&\n':
return len(self[2].split())
else:
return 0
@property
def shape(self):
return self.rows, self.columns
@property
def data(self):
return np.genfromtxt(StringIO(str(self)), skip_header=2, skip_footer=1)
@data.setter
def data(self, value):
if self.columns != 0 and value.shape[1] != self.columns:
raise TypeError(f'Number of columns is not {self.columns} (given {value.shape[1]})')
dtype = self.type
self.clear()
self.append(f'@target G{self._graph}.S{self._set}\n')
self.append(f'@type {dtype}\n')
for i, line in enumerate(value):
self.append(' '.join(line.astype(str)) + '\n')
self.append('&\n')
class GraceRegion(list):
"""
Just another name for a list
"""
def __str__(self):
return ''.join(self)
def _convert_to_value(parse_string):
tuples = parse_string.split(',')
for i, v in enumerate(tuples):
v = v.strip()
if re.match(r'[+-]?\d+(?:.\d+)?(?:[Ee][+-]?\d+)?', v):
try:
tuples[i] = int(v)
except ValueError:
tuples[i] = float(v)
elif re.match(r'\".*\"', v):
tuples[i] = v[1:-1]
else:
tuples[i] = v
if len(tuples) == 1:
return tuples[0]
return tuples
def _convert_to_str(value):
if isinstance(value, (tuple, list)):
return ', '.join(map(str, value))
else:
return str(value)