def testWriteReadWriteBack(self):
# setup our temp filenames
test_outfile = os.path.join(self.tmpdir, "sample_saved.star")
test_outfile2 = os.path.join(self.tmpdir, "sampled_saved2.star")
# create a new StarFile object directly via an OrderedDict of DataFrames
# not by reading a file
data = OrderedDict()
# note that GEMMI requires the names of the fields to start with _
# initialize a key-value set (a set of pairs in GEMMI parlance)
block0 = {
"_key1": "val1",
"_key2": "val2",
"_key3": "val3",
"_key4": "val4"
}
# initialize a single-row loop. we want this to be distinct from a
# set of key-value pairs
block1_dict = {"_field1": 31, "_field2": 32, "_field3": 33}
block1 = DataFrame([block1_dict], columns=block1_dict.keys())
block2_keys = ["_field4", "_field5", "_field6"]
block2_arr = [[f"{x}{y}" for x in range(3)] for y in range(3)]
# initialize a loop data block with a list of lists
block2 = DataFrame(block2_arr, columns=block2_keys)
data["pair"] = block0
data["single_row"] = block1
data["loops"] = block2
# initialize with blocks kwarg
original = StarFile(blocks=data)
original.write(test_outfile)
read_back = StarFile(test_outfile)
# assert that the read-back objects are equal
self.assertEqual(original, read_back)
# write back the second star file object
read_back.write(test_outfile2)
# compare the two .star files line by line
with open(test_outfile) as f_original, open(
test_outfile2) as f_read_back:
lines_original = f_original.readlines()
lines_read_back = f_read_back.readlines()
self.assertEqual(lines_original, lines_read_back)
os.remove(test_outfile)
os.remove(test_outfile2)
def write_star(self, df1, df2, ang, cs, voltage, pixel_size, amp, name,
output_dir):
"""
Writes CTF parameters to starfile.
"""
if not os.path.isdir(output_dir):
os.mkdir(output_dir)
data_block = {}
data_block["_rlnMicrographName"] = name
data_block["_rlnDefocusU"] = df1
data_block["_rlnDefocusV"] = df2
data_block["_rlnDefocusAngle"] = ang
data_block["_rlnSphericalAbberation"] = cs
data_block["_rlnAmplitudeContrast"] = amp
data_block["_rlnVoltage"] = voltage
data_block["_rlnDetectorPixelSize"] = pixel_size
df = DataFrame([data_block])
blocks = OrderedDict()
blocks["root"] = df
star = StarFile(blocks=blocks)
star.write(
os.path.join(output_dir,
os.path.splitext(name)[0]) + ".star")
class StarFileTestCase(TestCase):
def setUp(self):
with importlib_resources.path(tests.saved_test_data,
"sample_data_model.star") as path:
self.starfile = StarFile(path)
# Independent Image object for testing Image source methods
L = 768
self.im = Image(misc.face(gray=True).astype("float64")[:L, :L])
self.img_src = ArrayImageSource(self.im)
# We also want to flex the stack logic.
self.n = 21
im_stack = np.broadcast_to(self.im.data, (self.n, L, L))
# make each image methodically different
im_stack = np.multiply(im_stack, np.arange(self.n)[:, None, None])
self.im_stack = Image(im_stack)
self.img_src_stack = ArrayImageSource(self.im_stack)
# Create a tmpdir object for this test instance
self._tmpdir = tempfile.TemporaryDirectory()
# Get the directory from the name attribute of the instance
self.tmpdir = self._tmpdir.name
def tearDown(self):
# Destroy the tmpdir instance and contents
self._tmpdir.cleanup()
def testLength(self):
# StarFile is an iterable that gives us blocks
# blocks are pandas DataFrames
# We have 6 blocks in our sample starfile.
self.assertEqual(6, len(self.starfile))
def testIteration(self):
# A StarFile can be iterated over, yielding DataFrames for loops
# or dictionaries for pairs
for _, block in self.starfile:
self.assertTrue(
isinstance(block, DataFrame) or isinstance(block, dict))
def testBlockByIndex(self):
# We can use get_block_by_index to retrieve the blocks in
# the OrderedDict by index
# our first block is a set of pairs, represented by a dict
block0 = self.starfile.get_block_by_index(0)
self.assertTrue(isinstance(block0, dict))
self.assertEqual(block0["_rlnReferenceDimensionality"], "3")
# our second block is a loop, represented by a DataFrame
block1 = self.starfile.get_block_by_index(1)
self.assertTrue(isinstance(block1, DataFrame))
self.assertEqual(block1.at[0, "_rlnClassDistribution"], "1.000000")
def testBlockByName(self):
# Indexing a StarFile with a string gives us a block with that name
# ("data_<name>" in starfile).
# the block at index 0 has the name 'model_general'
block0 = self.starfile["model_general"]
# this block is a pair/dict with 22 key value pairs
self.assertEqual(len(block0), 22)
# the block at index 1 has name 'model_classes'
block1 = self.starfile["model_classes"]
# This block is a loop/DF with one row
self.assertEqual(len(block1), 1)
def testData(self):
df = self.starfile["model_class_1"]
self.assertEqual(76, len(df))
self.assertEqual(8, len(df.columns))
# Note that no typecasting of values is performed at io.StarFile level
self.assertEqual(
"0.000000",
df[df["_rlnSpectralIndex"] == "0"].iloc[0]["_rlnResolution"])
def testFileNotFound(self):
with self.assertRaises(FileNotFoundError):
StarFile("badfile.star")
def testReadWriteReadBack(self):
# Save the StarFile object to a .star file
# Read it back for object equality
# Note that __eq__ is supported for the class
# it checks the equality of the underlying OrderedDicts of DataFrames
# using pd.DataFrame.equals()
test_outfile = os.path.join(self.tmpdir, "sample_saved.star")
self.starfile.write(test_outfile)
starfile2 = StarFile(test_outfile)
self.assertEqual(self.starfile, starfile2)
os.remove(test_outfile)
def testWriteReadWriteBack(self):
# setup our temp filenames
test_outfile = os.path.join(self.tmpdir, "sample_saved.star")
test_outfile2 = os.path.join(self.tmpdir, "sampled_saved2.star")
# create a new StarFile object directly via an OrderedDict of DataFrames
# not by reading a file
data = OrderedDict()
# note that GEMMI requires the names of the fields to start with _
# initialize a key-value set (a set of pairs in GEMMI parlance)
block0 = {
"_key1": "val1",
"_key2": "val2",
"_key3": "val3",
"_key4": "val4"
}
# initialize a single-row loop. we want this to be distinct from a
# set of key-value pairs
block1_dict = {"_field1": 31, "_field2": 32, "_field3": 33}
block1 = DataFrame([block1_dict], columns=block1_dict.keys())
block2_keys = ["_field4", "_field5", "_field6"]
block2_arr = [[f"{x}{y}" for x in range(3)] for y in range(3)]
# initialize a loop data block with a list of lists
block2 = DataFrame(block2_arr, columns=block2_keys)
data["pair"] = block0
data["single_row"] = block1
data["loops"] = block2
# initialize with blocks kwarg
original = StarFile(blocks=data)
original.write(test_outfile)
read_back = StarFile(test_outfile)
# assert that the read-back objects are equal
self.assertEqual(original, read_back)
# write back the second star file object
read_back.write(test_outfile2)
# compare the two .star files line by line
with open(test_outfile) as f_original, open(
test_outfile2) as f_read_back:
lines_original = f_original.readlines()
lines_read_back = f_read_back.readlines()
self.assertEqual(lines_original, lines_read_back)
os.remove(test_outfile)
os.remove(test_outfile2)
def testArgsError(self):
with self.assertRaises(StarFileError):
_blocks = OrderedDict()
_blocks[""] = DataFrame(["test", "data"])
with importlib_resources.path(tests.saved_test_data,
"sample_data_model.star") as path:
StarFile(filepath=path, blocks=_blocks)
def testEmptyInit(self):
empty = StarFile()
self.assertTrue(isinstance(empty.blocks, OrderedDict))
self.assertEqual(len(empty.blocks), 0)
def save_metadata(self,
starfile_filepath,
new_mrcs=True,
batch_size=512,
save_mode=None):
"""
Save updated metadata to a STAR file
:param starfile_filepath: Path to STAR file where we want to
save image_source
:param new_mrcs: Whether to save all images to new MRCS files or not.
If True, new file names and pathes need to be created.
:param batch_size: Batch size of images to query from the
`ImageSource` object. Every `batch_size` rows, entries are
written to STAR file.
:param save_mode: Whether to save all images in a single or
multiple files in batch size.
:return: None
"""
df = self._metadata.copy()
# Drop any column that doesn't start with a *single* underscore
df = df.drop(
[
str(col) for col in df.columns
if not col.startswith("_") or col.startswith("__")
],
axis=1,
)
if new_mrcs:
# Create a new column that we will be populating in the loop below
df["_rlnImageName"] = ""
if save_mode == "single":
# Save all images into one single mrc file
fname = os.path.basename(starfile_filepath)
fstem = os.path.splitext(fname)[0]
mrcs_filename = f"{fstem}_{0}_{self.n-1}.mrcs"
# Then set name in dataframe for the StarFile
# Note, here the row_indexer is :, representing all rows in this data frame.
# df.loc will be reponsible for dereferencing and assigning values to df.
# Pandas will assert df.shape[0] == self.n
df.loc[:, "_rlnImageName"] = [
f"{j + 1:06}@{mrcs_filename}" for j in range(self.n)
]
else:
# save all images into multiple mrc files in batch size
for i_start in np.arange(0, self.n, batch_size):
i_end = min(self.n, i_start + batch_size)
num = i_end - i_start
mrcs_filename = (os.path.splitext(
os.path.basename(starfile_filepath))[0] +
f"_{i_start}_{i_end-1}.mrcs")
# Note, here the row_indexer is a slice.
# df.loc will be reponsible for dereferencing and assigning values to df.
# Pandas will assert the lnegth of row_indexer equals num.
row_indexer = df[i_start:i_end].index
df.loc[row_indexer, "_rlnImageName"] = [
"{0:06}@{1}".format(j + 1, mrcs_filename)
for j in range(num)
]
filename_indices = df._rlnImageName.str.split(pat="@",
expand=True)[1].tolist()
# initialize the star file object and save it
odict = OrderedDict()
# since our StarFile only has one block, the convention is to save it with the header "data_", i.e. its name is blank
# if we had a block called "XYZ" it would be saved as "XYZ"
# thus we index the metadata block with ""
odict[""] = df
out_star = StarFile(blocks=odict)
out_star.write(starfile_filepath)
return filename_indices