Added initial code for density maps

environment.yml

@@ -6,6 +6,7 @@ dependencies:
   - codecov
   - coverage
   - gemmi
+  - conda-forge::jupyter
   - mrcfile
   - numba
   - numpy
@@ -20,4 +21,3 @@ dependencies:
   - pip :
       - git+https://github.com/compSPI/simSPI.git
       - starfile
-

ioSPI/density_maps.py

@@ -0,0 +1,42 @@
+"""Read density maps."""
+
+import math
+
+import mrcfile
+import numpy as np
+
+
+def read_density_map_from_mrc(path, make_cubic=False):
+    """Return density map from an input .mrc file.
+
+    Parameters
+    ----------
+    path : str
+        File name for .mrc file to convert to density map
+    make_cubic : bool
+        If true, pads the map to reach a cubic array
+    """
+    map_mrc = mrcfile.open(path)
+    data = map_mrc.data
+    if make_cubic:
+        cube_length = max(data.shape)
+        cubic_data = np.pad(
+            data,
+            (
+                (
+                    math.floor((cube_length - data.shape[0]) / 2.0),
+                    math.ceil((cube_length - data.shape[0]) / 2),
+                ),
+                (
+                    math.floor((cube_length - data.shape[1]) / 2.0),
+                    math.ceil((cube_length - data.shape[1]) / 2),
+                ),
+                (
+                    math.floor((cube_length - data.shape[2]) / 2.0),
+                    math.ceil((cube_length - data.shape[2]) / 2),
+                ),
+            ),
+            "minimum",
+        )
+        data = cubic_data
+    return data

tests/test_density_maps.py

@@ -0,0 +1,41 @@
+"""Unit tests for density maps helper functions."""
+
+import os
+import tempfile
+
+import mrcfile
+import numpy as np
+
+from ioSPI import density_maps
+
+
+def test_read_density_map_from_mrc():
+    """Test read_density_map_from_mrc util function with a random mrc file."""
+    tmp_mrc = tempfile.NamedTemporaryFile(delete=False, suffix=".mrc")
+    tmp_mrc.close()
+    data = np.random.uniform(0, 1, (5, 5, 5)).astype(dtype="float32")
+
+    try:
+        with mrcfile.new(tmp_mrc.name, overwrite=True) as mrc:
+            mrc.set_data(data)
+        out_data = density_maps.read_density_map_from_mrc(tmp_mrc.name)
+        assert (out_data == data).all()
+    finally:
+        os.unlink(tmp_mrc.name)
+
+
+def test_read_density_map_from_mrc_cubic():
+    """Test read_density_map_from_mrc util function make_cubic function."""
+    tmp_mrc = tempfile.NamedTemporaryFile(delete=False, suffix=".mrc")
+    tmp_mrc.close()
+    data = np.random.uniform(0, 1, (4, 5, 6)).astype(dtype="float32")
+
+    try:
+        with mrcfile.new(tmp_mrc.name, overwrite=True) as mrc:
+            mrc.set_data(data)
+        out_data = density_maps.read_density_map_from_mrc(tmp_mrc.name, make_cubic=True)
+        assert out_data.shape[0] == 6
+        assert out_data.shape[1] == 6
+        assert out_data.shape[2] == 6
+    finally:
+        os.unlink(tmp_mrc.name)

tests/test_notebooks.py

@@ -0,0 +1,54 @@
+"""Unit tests for the notebooks."""
+
+import glob
+import os
+import subprocess
+import tempfile
+
+import pytest
+
+NOTEBOOKS_DIR = "notebooks"
+NOTEBOOKS_TO_SKIP = os.path.join(NOTEBOOKS_DIR, "download_and_upload_with_osf.ipynb")
+
+
+def _exec_notebook(path):
+    """Execute notebook at path.
+
+    Parameters
+    ----------
+    path : str
+        Relative path of the notebook.
+        E.g. notebooks/particle_metadata.ipynb
+    """
+    file_name = tempfile.NamedTemporaryFile(suffix=".ipynb").name
+    args = [
+        "jupyter",
+        "nbconvert",
+        "--to",
+        "notebook",
+        "--execute",
+        "--ExecutePreprocessor.timeout=1000",
+        "--ExecutePreprocessor.kernel_name=python3",
+        "--output",
+        file_name,
+        path,
+    ]
+    subprocess.check_call(args)
+
+
+paths = sorted(glob.glob(f"{NOTEBOOKS_DIR}/*.ipynb"))
+
+
+@pytest.mark.parametrize("path", paths)
+def test_notebook(path):
+    """Test the notebook at path by executing it.
+
+    Parameters
+    ----------
+    path : str
+        Relative path of the notebooks.
+        E.g. notebooks/particle_metadata.ipynb
+    """
+    if path in NOTEBOOKS_TO_SKIP:
+        pytest.skip()
+    _exec_notebook(path)