Skip to content

Commit 30ad179

Browse files
Add MRC file writing support (#147)
* Fix #108 * Add MRC file writing support - Implement MRC.write() method - Add Grid.export() integration for MRC format - Preserve header information and handle edge cases * Added Test * update CHANGELOG
1 parent 394161b commit 30ad179

File tree

5 files changed

+313
-5
lines changed

5 files changed

+313
-5
lines changed

AUTHORS

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,3 +28,4 @@ Contributors:
2828
* Olivier Languin-Cattoën <ollyfutur>
2929
* Andrés Montoya <conradolandia> (logo)
3030
* Rich Waldo <plethorachutney>
31+
* Pradyumn Prasad <Pradyumn-cloud>

CHANGELOG

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,7 @@ The rules for this file:
1313
* accompany each entry with github issue/PR number (Issue #xyz)
1414

1515
------------------------------------------------------------------------------
16-
??/??/???? IAlibay, ollyfutur, conradolandia, orbeckst, PlethoraChutney
16+
??/??/???? IAlibay, ollyfutur, conradolandia, orbeckst, PlethoraChutney, Pradyumn-cloud
1717
* 1.1.0
1818

1919
Changes
@@ -30,6 +30,7 @@ The rules for this file:
3030
(PR #142)
3131
* `Grid` now allows forcing MRC/CCP4 maps to be read as volumes even when
3232
the header indicates they are stacks of 2D images. (#150, PR #149)
33+
* Added MRC file writing support (Issue #108)
3334

3435
Fixes
3536

gridData/core.py

Lines changed: 39 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -212,6 +212,7 @@ def __init__(self, grid=None, edges=None, origin=None, delta=None,
212212
'PKL': self._export_python,
213213
'PICKLE': self._export_python, # compatibility
214214
'PYTHON': self._export_python, # compatibility
215+
'MRC': self._export_mrc,
215216
}
216217
self._loaders = {
217218
'CCP4': self._load_mrc,
@@ -600,6 +601,8 @@ def export(self, filename, file_format=None, type=None, typequote='"'):
600601
601602
dx
602603
:mod:`OpenDX`
604+
mrc
605+
:mod:`mrc` MRC/CCP4 format
603606
pickle
604607
pickle (use :meth:`Grid.load` to restore); :meth:`Grid.save`
605608
is simpler than ``export(format='python')``.
@@ -609,7 +612,7 @@ def export(self, filename, file_format=None, type=None, typequote='"'):
609612
filename : str
610613
name of the output file
611614
612-
file_format : {'dx', 'pickle', None} (optional)
615+
file_format : {'dx', 'pickle', 'mrc', None} (optional)
613616
output file format, the default is "dx"
614617
615618
type : str (optional)
@@ -686,6 +689,41 @@ def _export_dx(self, filename, type=None, typequote='"', **kwargs):
686689
if ext == '.gz':
687690
filename = root + ext
688691
dx.write(filename)
692+
693+
def _export_mrc(self, filename, **kwargs):
694+
"""Export the density grid to an MRC/CCP4 file.
695+
696+
The MRC2014 file format is used via the mrcfile library.
697+
698+
Parameters
699+
----------
700+
filename : str
701+
Output filename
702+
**kwargs
703+
Additional keyword arguments (currently ignored)
704+
705+
Notes
706+
-----
707+
* Only orthorhombic unit cells are supported
708+
* If the Grid was loaded from an MRC file, the original header
709+
information (including axis ordering) is preserved
710+
* For new grids, standard ordering (mapc=1, mapr=2, maps=3) is used
711+
712+
.. versionadded:: 1.1.0
713+
"""
714+
# Create MRC object and populate with Grid data
715+
mrc_file = mrc.MRC()
716+
mrc_file.array = self.grid
717+
mrc_file.delta = numpy.diag(self.delta)
718+
mrc_file.origin = self.origin
719+
mrc_file.rank = 3
720+
721+
# Transfer header if it exists (preserves axis ordering and other metadata)
722+
if hasattr(self, '_mrc_header'):
723+
mrc_file.header = self._mrc_header
724+
725+
# Write to file
726+
mrc_file.write(filename)
689727

690728
def save(self, filename):
691729
"""Save a grid object to `filename` and add ".pickle" extension.

gridData/mrc.py

Lines changed: 104 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ class MRC(object):
8585
-----
8686
* Only volumetric (3D) densities are read.
8787
* Only orthorhombic unitcells supported (other raise :exc:`ValueError`)
88-
* Only reading is currently supported.
88+
* Reading and writing are supported.
8989
9090
9191
.. versionadded:: 0.7.0
@@ -160,5 +160,106 @@ def histogramdd(self):
160160
"""Return array data as (edges,grid), i.e. a numpy nD histogram."""
161161
return (self.array, self.edges)
162162

163-
164-
163+
def write(self, filename):
164+
"""Write grid data to MRC/CCP4 file format.
165+
166+
Parameters
167+
----------
168+
filename : str
169+
Output filename for the MRC file
170+
171+
Notes
172+
-----
173+
The data array should be in xyz order (axis 0=X, axis 1=Y, axis 2=Z).
174+
175+
If the MRC object was created by reading an existing file, the original
176+
header information (including mapc, mapr, maps ordering) is preserved.
177+
Otherwise, standard ordering (mapc=1, mapr=2, maps=3) is used.
178+
179+
.. versionadded:: 1.1.0
180+
"""
181+
if filename is not None:
182+
self.filename = filename
183+
184+
# Preserve header if it exists, otherwise use defaults
185+
if hasattr(self, 'header'):
186+
# File was read - preserve original ordering
187+
h = self.header
188+
axes_order = np.hstack([h.mapc, h.mapr, h.maps])
189+
mapc, mapr, maps = int(h.mapc), int(h.mapr), int(h.maps)
190+
else:
191+
# New file - use standard ordering
192+
axes_order = np.array([1, 2, 3])
193+
mapc, mapr, maps = 1, 2, 3
194+
h = None
195+
196+
# Reverse the transformation done in read()
197+
transpose_order = np.argsort(axes_order[::-1])
198+
inverse_transpose_order = np.argsort(transpose_order)
199+
200+
# Transform our xyz array back to the file's native ordering
201+
data_for_file = np.transpose(self.array, axes=inverse_transpose_order)
202+
203+
# Ensure proper data type (float32 is standard for mode 2)
204+
data_for_file = data_for_file.astype(np.float32)
205+
206+
# Create new MRC file
207+
with mrcfile.new(filename, overwrite=True) as mrc:
208+
mrc.set_data(data_for_file)
209+
210+
# Set voxel size from delta (diagonal elements)
211+
voxel_size = np.diag(self.delta).astype(np.float32)
212+
mrc.voxel_size = tuple(voxel_size)
213+
214+
# Set map ordering
215+
mrc.header.mapc = mapc
216+
mrc.header.mapr = mapr
217+
mrc.header.maps = maps
218+
219+
# Handle nstart and origin
220+
if h is not None:
221+
# Preserve original header values
222+
nxstart = int(h.nxstart)
223+
nystart = int(h.nystart)
224+
nzstart = int(h.nzstart)
225+
header_origin_xyz = np.array([h.origin.x, h.origin.y, h.origin.z], dtype=np.float32)
226+
227+
mrc.header.mx = int(h.mx)
228+
mrc.header.my = int(h.my)
229+
mrc.header.mz = int(h.mz)
230+
231+
# Preserve cell dimensions
232+
if hasattr(h, 'cella'):
233+
mrc.header.cella.x = float(h.cella.x)
234+
mrc.header.cella.y = float(h.cella.y)
235+
mrc.header.cella.z = float(h.cella.z)
236+
if hasattr(h, 'cellb'):
237+
mrc.header.cellb.alpha = float(h.cellb.alpha)
238+
mrc.header.cellb.beta = float(h.cellb.beta)
239+
mrc.header.cellb.gamma = float(h.cellb.gamma)
240+
# Copy space group if available
241+
if hasattr(h, 'ispg'):
242+
mrc.header.ispg = int(h.ispg)
243+
else:
244+
# For new files, calculate nstart from origin
245+
if np.any(voxel_size <= 0):
246+
raise ValueError(f"Voxel size must be positive, got {voxel_size}")
247+
248+
# Set header.origin = 0 and encode everything in nstart
249+
header_origin_xyz = np.zeros(3, dtype=np.float32)
250+
nxstart = int(np.round(self.origin[0] / voxel_size[0]))
251+
nystart = int(np.round(self.origin[1] / voxel_size[1]))
252+
nzstart = int(np.round(self.origin[2] / voxel_size[2]))
253+
254+
# Set the start positions
255+
mrc.header.nxstart = nxstart
256+
mrc.header.nystart = nystart
257+
mrc.header.nzstart = nzstart
258+
259+
# Set explicit origin
260+
mrc.header.origin.x = float(header_origin_xyz[0])
261+
mrc.header.origin.y = float(header_origin_xyz[1])
262+
mrc.header.origin.z = float(header_origin_xyz[2])
263+
264+
# Update statistics only
265+
mrc.update_header_stats()

gridData/tests/test_mrc.py

Lines changed: 167 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -142,3 +142,170 @@ def test_origin(self, grid, ccp4data):
142142

143143
def test_data(self, grid, ccp4data):
144144
assert_allclose(grid.grid, ccp4data.array)
145+
146+
class TestMRCWrite:
147+
"""Tests for MRC write functionality"""
148+
149+
def test_mrc_write_roundtrip(self, ccp4data, tmpdir):
150+
"""Test writing and reading back preserves data"""
151+
outfile = str(tmpdir / "roundtrip.mrc")
152+
153+
# Write the file
154+
ccp4data.write(outfile)
155+
156+
# Read it back
157+
mrc_read = mrc.MRC(outfile)
158+
159+
# Check data matches
160+
assert_allclose(mrc_read.array, ccp4data.array)
161+
assert_allclose(mrc_read.origin, ccp4data.origin, rtol=1e-5, atol=1e-3)
162+
assert_allclose(mrc_read.delta, ccp4data.delta, rtol=1e-5)
163+
164+
def test_mrc_write_header_preserved(self, ccp4data, tmpdir):
165+
"""Test that header fields are preserved"""
166+
outfile = str(tmpdir / "header.mrc")
167+
168+
ccp4data.write(outfile)
169+
mrc_read = mrc.MRC(outfile)
170+
171+
# Check axis ordering preserved
172+
assert mrc_read.header.mapc == ccp4data.header.mapc
173+
assert mrc_read.header.mapr == ccp4data.header.mapr
174+
assert mrc_read.header.maps == ccp4data.header.maps
175+
176+
# Check offsets preserved
177+
assert mrc_read.header.nxstart == ccp4data.header.nxstart
178+
assert mrc_read.header.nystart == ccp4data.header.nystart
179+
assert mrc_read.header.nzstart == ccp4data.header.nzstart
180+
181+
def test_mrc_write_new_file(self, tmpdir):
182+
"""Test creating new MRC file from scratch"""
183+
outfile = str(tmpdir / "new.mrc")
184+
185+
# Create new MRC object
186+
mrc_new = mrc.MRC()
187+
mrc_new.array = np.arange(24).reshape(2, 3, 4).astype(np.float32)
188+
mrc_new.delta = np.diag([1.0, 2.0, 3.0])
189+
mrc_new.origin = np.array([5.0, 10.0, 15.0])
190+
mrc_new.rank = 3
191+
192+
# Write and read back
193+
mrc_new.write(outfile)
194+
mrc_read = mrc.MRC(outfile)
195+
196+
# Verify
197+
assert_allclose(mrc_read.array, mrc_new.array, rtol=1e-5)
198+
assert_allclose(mrc_read.origin, mrc_new.origin, rtol=1e-4)
199+
assert_allclose(np.diag(mrc_read.delta), np.diag(mrc_new.delta), rtol=1e-5)
200+
201+
def test_mrc_write_zero_voxel_raises(self, tmpdir):
202+
"""Test that zero voxel size raises ValueError"""
203+
outfile = str(tmpdir / "invalid.mrc")
204+
205+
mrc_obj = mrc.MRC()
206+
mrc_obj.array = np.ones((2, 2, 2), dtype=np.float32)
207+
mrc_obj.delta = np.diag([0.0, 1.0, 1.0])
208+
mrc_obj.origin = np.array([0.0, 0.0, 0.0])
209+
mrc_obj.rank = 3
210+
211+
with pytest.raises(ValueError, match="Voxel size must be positive"):
212+
mrc_obj.write(outfile)
213+
214+
215+
class TestGridMRCWrite:
216+
"""Tests for Grid.export() with MRC format"""
217+
218+
def test_grid_export_mrc(self, tmpdir):
219+
"""Test Grid.export() with file_format='mrc'"""
220+
outfile = str(tmpdir / "grid.mrc")
221+
222+
# Create simple grid
223+
data = np.arange(60).reshape(3, 4, 5).astype(np.float32)
224+
g = Grid(data, origin=[0, 0, 0], delta=[1.0, 1.0, 1.0])
225+
226+
# Export and read back
227+
g.export(outfile, file_format='mrc')
228+
g_read = Grid(outfile)
229+
230+
# Verify
231+
assert_allclose(g_read.grid, g.grid, rtol=1e-5)
232+
assert_allclose(g_read.origin, g.origin, rtol=1e-4)
233+
assert_allclose(g_read.delta, g.delta, rtol=1e-5)
234+
235+
def test_grid_export_mrc_roundtrip(self, tmpdir):
236+
"""Test MRC → Grid → export → Grid preserves data"""
237+
outfile = str(tmpdir / "roundtrip_grid.mrc")
238+
239+
# Load original
240+
g_orig = Grid(datafiles.CCP4_1JZV)
241+
242+
# Export and reload
243+
g_orig.export(outfile, file_format='mrc')
244+
g_read = Grid(outfile)
245+
246+
# Verify
247+
assert_allclose(g_read.grid, g_orig.grid, rtol=1e-5)
248+
assert_allclose(g_read.origin, g_orig.origin, rtol=1e-4)
249+
assert_allclose(g_read.delta, g_orig.delta, rtol=1e-5)
250+
assert_equal(g_read.grid.shape, g_orig.grid.shape)
251+
252+
def test_grid_export_mrc_preserves_header(self, tmpdir):
253+
"""Test that Grid preserves MRC header through export"""
254+
outfile = str(tmpdir / "header_grid.mrc")
255+
256+
g_orig = Grid(datafiles.CCP4_1JZV)
257+
orig_mapc = g_orig._mrc_header.mapc
258+
orig_mapr = g_orig._mrc_header.mapr
259+
orig_maps = g_orig._mrc_header.maps
260+
261+
# Export and check
262+
g_orig.export(outfile, file_format='mrc')
263+
g_read = Grid(outfile)
264+
265+
assert g_read._mrc_header.mapc == orig_mapc
266+
assert g_read._mrc_header.mapr == orig_mapr
267+
assert g_read._mrc_header.maps == orig_maps
268+
269+
def test_mrc_write_4x4x4_with_header(self, tmpdir):
270+
"""Test writing 4x4x4 MRC file with custom header values."""
271+
272+
# Create 4x4x4 data
273+
data = np.arange(64, dtype=np.float32).reshape((4, 4, 4))
274+
outfile = str(tmpdir / "test_with_header.mrc")
275+
276+
# Create and write MRC
277+
m = mrc.MRC()
278+
m.array = data
279+
m.delta = np.diag([1.5, 2.0, 2.5])
280+
m.origin = np.array([10.0, 20.0, 30.0])
281+
m.rank = 3
282+
m.write(outfile)
283+
284+
# Read back and verify
285+
m_read = mrc.MRC(outfile)
286+
assert_allclose(m_read.array, data)
287+
assert_allclose(np.diag(m_read.delta), [1.5, 2.0, 2.5])
288+
assert_allclose(m_read.origin, [10.0, 20.0, 30.0], rtol=1e-4, atol=1.0)
289+
290+
291+
def test_mrc_write_4x4x4_without_header(self, tmpdir):
292+
"""Test writing 4x4x4 MRC file with default header."""
293+
294+
# Create 4x4x4 random data
295+
np.random.seed(42)
296+
data = np.random.rand(4, 4, 4).astype(np.float32)
297+
outfile = str(tmpdir / "test_without_header.mrc")
298+
299+
# Create and write MRC
300+
m = mrc.MRC()
301+
m.array = data
302+
m.delta = np.diag([1.0, 1.0, 1.0])
303+
m.origin = np.array([0.0, 0.0, 0.0])
304+
m.rank = 3
305+
m.write(outfile)
306+
307+
# Read back and verify
308+
m_read = mrc.MRC(outfile)
309+
assert_allclose(m_read.array, data)
310+
assert_allclose(np.diag(m_read.delta), [1.0, 1.0, 1.0])
311+
assert_allclose(m_read.origin, [0.0, 0.0, 0.0], rtol=1e-4, atol=1.0)

0 commit comments

Comments
 (0)