def load(fname, sme):
"""
Load the SME Structure from disk
Parameters
----------
fname : str
file to load
sme : SME_Structure
empty sme structure with default values set
Returns
-------
sme : SME_Structure
loaded sme structure
"""
try:
ff = FlexFile.read(fname)
sme = from_flex(ff, sme)
ff.close()
return sme
except Exception as ex:
logger.error(ex)
try:
sme = load_v1(fname, sme)
except:
raise ex
return sme
def write(self, filename, format=None):
os.makedirs(dirname(filename), exist_ok=True)
if format is None:
format = self.__class__.determine_filetype_based_on_filename(
filename)
if format == "flex":
ff = FlexFile(extensions={"spectrum": self.__flex_save__()})
ff.write(filename)
elif format == "fits":
hdu = self.get_fits_hdu()
hdu.writeto(filename, overwrite=True)
elif format == "npz":
data = self.to_dict()
np.savez(filename, **data)
else:
raise ValueError(
f"Format not understood expected one of ('flex', 'fits') but got {format}"
)
def to_flex(sme):
header = {}
extensions = {}
for name in sme._names:
value = sme[name]
if isinstance(value, IPersist):
extensions[name] = value._save()
elif value is not None:
header[name] = value
ff = FlexFile(header, extensions)
return ff
def test_read_write(tmp_fname):
file = FlexFile()
file.header["Hello"] = "World"
file.write(tmp_fname)
del file
f2 = FlexFile.read(tmp_fname)
assert isinstance(f2.header, dict)
assert f2.header["Hello"] == "World"
def test_header_numpy_arrays(tmp_fname):
file = FlexFile()
file.header["Hello"] = np.ones(3)
file.write(tmp_fname)
del file
f2 = FlexFile.read(tmp_fname)
assert isinstance(f2.header, dict)
assert np.all(np.asarray(f2.header["Hello"]) == 1)
def test_empty_data_json(tmp_fname):
file = FlexFile()
ext = BinaryDataExtension(data=[])
file.extensions["img"] = ext
file.to_json(tmp_fname)
del file
f2 = FlexFile.from_json(tmp_fname)
assert f2["img"].data.size == 0
def test_empty_data(tmp_fname):
file = FlexFile()
ext = BinaryDataExtension(data=[])
file.extensions["img"] = ext
file.write(tmp_fname)
del file
f2 = FlexFile.read(tmp_fname)
assert f2["img"].data.size == 0
def test_json_multi(tmp_fname):
file = FlexFile()
ext = MultipleDataExtension(data={"bla": np.zeros(10)})
file.extensions["img"] = ext
file.to_json(tmp_fname)
del file
f2 = FlexFile.from_json(tmp_fname)
assert f2["img"].data["bla"].size == 10
assert np.all(f2["img"].data["bla"] == 0)
def test_read_write(tmp_fname):
file = FlexFile()
ext = BinaryDataExtension(data=np.zeros(10))
file.extensions["img"] = ext
file.write(tmp_fname)
del file
f2 = FlexFile.read(tmp_fname)
assert f2["img"].data.size == 10
assert np.all(f2["img"].data == 0)
def test_json(tableclass, tmp_fname, table):
file = FlexFile()
ext = tableclass(data=table)
file.extensions["tab"] = ext
file.to_json(tmp_fname)
del file
f2 = FlexFile.from_json(tmp_fname)
assert f2["tab"].data.size == 10 * 2
assert np.all(f2["tab"].data == 1)
assert len(f2["tab"].data.columns) == 2
assert "A" in f2["tab"].data.columns
assert "B" in f2["tab"].data.columns
def read(cls, filename, format=None):
if format is None:
format = cls.determine_filetype_based_on_filename(filename)
if format == "flex":
ff = FlexFile.read(filename)
spec = ff["spectrum"]
return spec
elif format == "fits":
hdulist = fits.open(filename)
spec = cls.read_fits_hdu(hdulist[1])
return spec
elif format == "npz":
data = np.load(filename, allow_pickle=True)
data = {k: v[()] for k, v in data.items()}
spec = cls.from_dict(data)
return spec
else:
raise ValueError(
f"Format not understood expected one of ('flex', 'fits') but got {format}"
)
import numpy as np
import pandas as pd
from flex.flex import FlexFile
from flex.extensions.bindata import BinaryDataExtension, MultipleDataExtension
from flex.extensions.tabledata import TableExtension, JSONTableExtension
file = FlexFile()
ext_bin = BinaryDataExtension(data=np.linspace(0, 1, 4 * 3 * 2).reshape((4, 3,
2)))
ext = MultipleDataExtension()
ext.data["data"] = np.arange(12, dtype="i2").reshape((3, 2, 2))
tab = JSONTableExtension()
tab.data = pd.DataFrame(np.arange(20).reshape((10, 2)), columns=["A", "B"])
file.header["bla"] = np.float(3e10)
file.header["blub"] = np.nan
file.header["blrub"] = float("inf")
file.header["blurb"] = np.array([1, 2, 3])
file.header["blurb2"] = [1, 2, 3]
file.header["int"] = 1
file.header["np_int"] = np.int(1)
file.header["str"] = "Hello"
file.header["np_str"] = np.str("World")
# FITS can handle either long header keys, or long entries
# But not both at the same time...
file.header["really_long_entry"] = "bla"
file.header["really_l"] = "bla" * 100
file.header["really_long_text_with_HIERARCH"] = "bla" * 100
def save(fname_in, fname_out):
ff = FlexFile.from_json(fname_in)
sme = SME_Structure()
sme = from_flex(ff, sme)
sme.save(fname_out)
def load(self, filename=None):
if filename is None:
filename = self.savefilename
ff = FlexFile.read(filename)
return ff.header
def save(self, data, filename=None):
if filename is None:
filename = self.savefilename
ff = FlexFile(header=data)
ff.write(filename)
