def save_to_igor_itx(file_path: str,
xs: List[np.ndarray],
datas: List[np.ndarray],
names: List[str],
ys: Optional[List[np.ndarray]] = None,
x_labels: Optional[Union[str, List[str]]] = None,
y_labels: Optional[Union[str, List[str]]] = None):
"""Save data to a .itx file which can be dropped into Igor"""
def check_axis_linear(arr: np.ndarray, axis: str, name: str,
current_waves: list) -> bool:
if arr.shape[-1] > 1 and not np.all(
np.isclose(np.diff(arr),
np.diff(arr)[0])):
logger.warning(
f"{file_path}: Igor doesn't support a non-linear {axis}-axis. Saving as separate wave"
)
axis_wave = IgorWave(arr, name=name + f'_{axis}')
current_waves.append(axis_wave)
return False
else:
return True
if x_labels is None or isinstance(x_labels, str):
x_labels = [x_labels] * len(datas)
if y_labels is None or isinstance(y_labels, str):
y_labels = [y_labels] * len(datas)
if ys is None:
ys = [None] * len(datas)
assert all([
len(datas) == len(list_) for list_ in [xs, names, x_labels, y_labels]
])
waves = []
for x, y, data, name, x_label, y_label in zip(xs, ys, datas, names,
x_labels, y_labels):
wave = IgorWave(data, name=name)
if x is not None:
if check_axis_linear(x, 'x', name, waves):
wave.set_dimscale('x',
x[0],
np.mean(np.diff(x)),
units=x_label)
if y is not None:
if check_axis_linear(y, 'y', name, waves):
wave.set_dimscale('y',
y[0],
np.mean(np.diff(y)),
units=y_label)
elif y_label is not None:
wave.set_datascale(y_label)
waves.append(wave)
with open(file_path, 'w') as fp:
for wave in waves:
wave.save_itx(
fp, image=True
) # Image = True hopefully makes np and igor match in x/y
def test_invalid_name(self):
name = '\'invalid_name\'' # wave cannot contain quotation marks
array = np.random.randint(0, 100, 10, dtype=np.int32)
wave = IgorWave(array)
self.assertRaises(validator.InvalidNameError,
wave.rename,
name,
on_errors='raise')
wave.rename(name, on_errors='fix')
self.assertEqual(wave.name, name.replace('\'', '_'))
def test_dimscale(self):
array = np.random.randint(0, 100, 10, dtype=np.int32)
wave = IgorWave(array)
wave.set_dimscale('x', 0, 0.01, 's')
with TemporaryFile('wb') as fp:
wave.save(fp)
with TemporaryFile('wt') as fp:
wave.save_itx(fp)
def test_pint(self):
import pint
ureg = pint.UnitRegistry()
Q = ureg.Quantity
a = np.arange(0.0, 10.0, 1.0)
with self.subTest('short units'):
wave = IgorWave(Q(a, 's'))
bunits = wave._wave_header.dataUnits
expected = 's'.encode(ENCODING)
self.assertEqual(bunits, expected)
with self.subTest('long units'):
q = Q(a, 'kg m / s**2')
wave = IgorWave(q)
bunits = wave._extended_data_units
expected = '{:~}'.format(q.units).encode(ENCODING)
self.assertEqual(bunits, expected)
def test_datetime(self):
array = np.arange(np.datetime64('2019-01-01'),
np.datetime64('2019-12-31'), np.timedelta64(1, 'D'))
wave = IgorWave(array)
self.assertIs(wave._check_array().dtype.type, np.float64)
with TemporaryFile('wb') as bin, TemporaryFile('wt') as text:
wave.save(bin)
wave.save_itx(text)
def check_axis_linear(arr: np.ndarray, axis: str, name: str,
current_waves: list) -> bool:
if arr.shape[-1] > 1 and not np.all(
np.isclose(np.diff(arr),
np.diff(arr)[0])):
logger.warning(
f"{file_path}: Igor doesn't support a non-linear {axis}-axis. Saving as separate wave"
)
axis_wave = IgorWave(arr, name=name + f'_{axis}')
current_waves.append(axis_wave)
return False
else:
return True
def dict_to_ibw(dict_like, prefix, on_errors='fix'):
"""
:param dict_like: dictionary-like object ({name: array-like, ...})
:param prefix: file name prefix (each wave is saved as <prefix>_<column>.ibw)
:param on_errors: behavior when invalid name is given. 'fix': fix errors. 'raise': raise an exception.
:return: dictionary of generated waves ({name: wave, ...})
"""
wavelist = [
IgorWave(array, name, on_errors)
for (name, array) in dict_like.items()
]
for wave in wavelist:
wave.save(str(prefix) + '_%s.ibw' % wave.name)
return {wave.name: wave for wave in wavelist}
def test_data_units(self):
wavename = 'wave0'
unitslist = ('', 'sec', 'second')
descs = ('empty units', 'short units', 'long units')
for desc, units in zip(descs, unitslist):
with self.subTest(desc):
com = 'X SetScale d,.*"%s",.*\'%s\'' % (units, wavename)
array = np.random.randint(0, 100, 10, dtype=np.int32)
wave = IgorWave(array)
wave.set_datascale(units)
with TemporaryFile('wb') as fp:
wave.save(fp)
with TemporaryFile('w+t') as fp:
wave.save_itx(fp)
fp.seek(0)
content = fp.read()
self.assertRegex(content, com)
def test_array_type(self):
valid_types = (np.bool_, np.float16, np.int32, np.uint32, np.int64,
np.uint64, np.float32, np.float64, np.complex128)
invalid_types = (object, str)
for vt in valid_types:
with self.subTest('type: %r' % vt):
wave = IgorWave(np.random.randint(0, 100, 10).astype(vt))
with TemporaryFile('wb') as fp:
wave.save(fp)
with TemporaryFile('wt') as fp:
wave.save_itx(fp)
for it in invalid_types:
with self.subTest('type: %r' % it):
wave = IgorWave(np.random.randint(0, 100, 10).astype(it))
with TemporaryFile('wb') as fp:
self.assertRaises(TypeError, wave.save, fp)
with TemporaryFile('wt') as fp:
self.assertRaises(TypeError, wave.save_itx, fp)
def make_dos_waves(path_list):
for rt in path_list:
print(rt)
fl = os.listdir(rt)
fl = [f for f in fl if '.dos' in f]
fl.sort()
fm = []
head = ''
for f in fl:
if not f.split('.')[0] == head:
head = f.split('.')[0]
fm.append([])
fm[len(fm) - 1].append(f)
print(fm)
for fl in fm:
ws = {}
head = fl[0].split('.')[0]
for f in fl:
sp = f.split('eV')[1]
if sp == 'up':
r = 1
elif sp == 'dn':
r = -1
en = 0
wp = rt + head + '_dos.itx'
path = rt + f
dos = load_dos(path)
for ky in dos.keys():
ws[ky + sp] = dos[ky] * r
print(ws.keys())
el = ws['ENERGYup']
e_s = el[0]
e_n = el.shape[0]
e_e = el[e_n - 1]
e_st = (e_e - e_s) / (e_n - 1)
with open(wp, 'w') as fp:
for ky in ws.keys():
if 'ENERGY' in ky:
pass
else:
wn = head + '_' + ky.replace(':', '_')
wn = wn.replace('-', '_')
wave = IgorWave(ws[ky], name=wn)
wave.set_dimscale('x', e_s, e_st, 'eV')
wave.save_itx(fp)
print(len(fm))
def fig_to_igor_itx(f: go.Figure, filepath: str):
d = data_from_plotly_fig(f)
waves = []
for k in d:
if not k.endswith('_x') and not k.endswith('_y'):
wave = IgorWave(d[k], name=k)
wave.set_datascale(f.layout.yaxis.title.text)
for dim in ['x', 'y']:
if f'{k}_{dim}' in d:
dim_arr = d[f'{k}_{dim}']
wave.set_dimscale('x',
dim_arr[0],
np.mean(np.diff(dim_arr)),
units=f.layout.xaxis.title.text)
waves.append(wave)
with open(filepath, 'w') as fp:
for wave in waves:
wave.save_itx(
fp, image=True
) # Image = True hopefully makes np and igor match in x/y
def dict_to_itx(dict_like, path_or_buf, on_errors='fix'):
"""
:param dict_like: dictionary-like object ({name: array-like, ...})
:param path_or_buf: filename or file-like object
:param on_errors: behavior when invalid name is given. 'fix': fix errors. 'raise': raise an exception.
:return: dictionary of generated waves ({name: wave, ...})
"""
fp = path_or_buf if hasattr(path_or_buf, 'write') else open(
path_or_buf, 'w', encoding=igorwriter.ENCODING)
wavelist = [
IgorWave(array, name, on_errors)
for (name, array) in dict_like.items()
]
try:
for wave in wavelist:
wave.save_itx(fp)
finally:
if fp is not path_or_buf:
fp.close()
return {wave.name: wave for wave in wavelist}
def make_3Dband_ibw(npypath, savepath, name, e_s, e_e):
ch = np.load(npypath)
print('npy data loaded')
dims = len(ch.shape) - 1
ky_s = -1
ky_e = 1
kz_s = -1
kz_e = 1
sp.call(['mkdir', '-p', savepath])
if dims == 2:
ch = np.transpose(ch, (1, 2, 0))
print(ch.shape)
ch = np.concatenate([np.flip(ch, 1), np.delete(ch, 0, 1)], 1)
ch = np.concatenate([np.flip(ch, 2), np.delete(ch, 0, 2)], 2)
elif dims == 3:
ch = np.transpose(ch, (2, 3, 1, 0))
ch = np.concatenate([np.flip(ch, 1), np.delete(ch, 0, 1)], 1)
ch = np.concatenate([np.flip(ch, 2), np.delete(ch, 0, 2)], 2)
ch = np.concatenate([np.flip(ch, 3), np.delete(ch, 0, 3)], 3)
kz_n = ch.shape[3]
kz_st = (kz_e - kz_s) / (kz_n - 1)
print('concatenate end')
ky_n = ch.shape[1]
ky_st = (ky_e - ky_s) / (ky_n - 1)
for ba in range(ch.shape[0]):
out = savepath + name + '_Band' + str(ba + 1) + '.ibw'
if (np.amax(ch[ba]) >= e_s and np.amin(ch[ba]) <= e_e):
wave = IgorWave(ch[ba], name=name + '_Band' + str(ba + 1))
wave.set_dimscale('x', ky_s, ky_st, '2pi/a')
wave.set_dimscale('y', ky_s, ky_st, '2pi/a')
if dims == 3:
wave.set_dimscale('z', kz_s, kz_st, '2pi/a')
wave.save(out)
print('save : ' + out)
def test_file_io(self):
array = np.random.randint(0, 100, 10, dtype=np.int32)
wave = IgorWave(array)
with TemporaryDirectory() as datadir:
itx = datadir + '/wave0.itx'
ibw = datadir + '/wave0.ibw'
wave.save_itx(itx)
with open(itx, 'w') as fp:
wave.save_itx(fp)
wave.save(ibw)
with self.subTest('save in a new binary file'):
with open(ibw, 'wb') as fp:
wave.save(fp)
with self.subTest('append to a binary should fail'):
with open(ibw, 'wb') as fp:
fp.write(b'something')
self.assertRaises(ValueError, wave.save, fp)
with open(ibw, 'ab') as fp:
self.assertRaises(ValueError, wave.save, fp)
with self.subTest('save in a truncated binary file'):
with open(ibw, 'wb') as fp:
wave.save(fp)