feat(nice-lite): add abinitio3D volume viewer

Author: djmaity

nice/nice_lite/data_structures/batchjob.py

@@ -25,7 +25,11 @@
 from .class_selection import ClassSelectionError, SIMPLEProjectFileReader
 from .simple import SIMPLEBatch, SIMPLEProjFile, SIMPLEProject
 from .job import Job
-from .mrc import read_mrc_stack_info, render_mrc_particle_png
+from .mrc import (
+    read_mrc_stack_info,
+    read_mrc_volume_info,
+    render_mrc_particle_png,
+)
 from .movie import render_movie_webp
 from .workspace import Workspace
 
@@ -660,6 +664,89 @@ def get_artifact_summary(self, max_previews=12):
             "images": images,
         }
 
+    def get_volume_outputs(self):
+        """Return project-declared, owned ab initio 3D density volumes."""
+        if self.prog != "abinitio3D" or self.status != "finished":
+            return []
+
+        job_dir = self.get_safe_job_dir()
+        result_project = self.get_result_project_path()
+        if job_dir is None or result_project is None:
+            return []
+        try:
+            records = SIMPLEProjectFileReader(result_project).read_records("out")
+        except (ClassSelectionError, OSError, OverflowError, struct.error):
+            return []
+
+        outputs = []
+        seen_paths = set()
+        for record in records:
+            if not isinstance(record, dict) or record.get("imgkind") != "vol":
+                continue
+            declared_path = record.get("vol")
+            if not isinstance(declared_path, str) or not declared_path.strip():
+                continue
+            declared_path = declared_path.strip()
+            if not os.path.isabs(declared_path):
+                declared_path = os.path.join(
+                    os.path.dirname(result_project),
+                    declared_path,
+                )
+            resolved_path = os.path.realpath(declared_path)
+            volume_name = os.path.basename(resolved_path)
+            safe_path = self._safe_job_file(volume_name, job_dir)
+            if (
+                safe_path is None
+                or safe_path != resolved_path
+                or safe_path in seen_paths
+            ):
+                continue
+
+            info = read_mrc_volume_info(safe_path)
+            if info is None:
+                continue
+            state = record.get("state")
+            if (
+                isinstance(state, bool)
+                or not isinstance(state, (int, float))
+                or not math.isfinite(state)
+                or state <= 0
+                or not float(state).is_integer()
+            ):
+                state = None
+            else:
+                state = int(state)
+            population = record.get("pop")
+            if (
+                isinstance(population, bool)
+                or not isinstance(population, (int, float))
+                or not math.isfinite(population)
+                or population < 0
+            ):
+                population = None
+            elif float(population).is_integer():
+                population = int(population)
+
+            seen_paths.add(safe_path)
+            outputs.append({
+                "path": safe_path,
+                "name": volume_name,
+                "state": state,
+                "population": population,
+                "width": info.width,
+                "height": info.height,
+                "depth": info.depth,
+                "voxel_size": info.voxel_size,
+                "minimum": info.minimum,
+                "maximum": info.maximum,
+            })
+
+        return sorted(outputs, key=lambda output: (
+            output["state"] is None,
+            output["state"] or 0,
+            output["name"],
+        ))
+
     def get_particle_stack_page(self, page=1, page_size=40):
         """Return one page of addressable images from owned output stacks.
 

nice/nice_lite/data_structures/mrc.py

@@ -29,6 +29,36 @@ class MRCStackInfo:
     data_offset: int
 
 
+@dataclass(frozen=True)
+class MRCVolumeInfo:
+    """Validated metadata for one three-dimensional MRC density map."""
+
+    width: int
+    height: int
+    depth: int
+    mode: int
+    data_offset: int
+    voxel_size: tuple
+    minimum: float
+    maximum: float
+
+
+@dataclass(frozen=True)
+class MRCVolumePayload:
+    """A bounded, normalized volume ready for a browser 3D texture."""
+
+    data: bytes
+    width: int
+    height: int
+    depth: int
+    source_width: int
+    source_height: int
+    source_depth: int
+    voxel_size: tuple
+    minimum: float
+    maximum: float
+
+
 def read_mrc_stack_info(path):
     """Return validated MRC stack metadata without reading particle pixels."""
     try:
@@ -73,6 +103,127 @@ def read_mrc_stack_info(path):
     )
 
 
+def read_mrc_volume_info(path):
+    """Return validated 3D-map metadata without reading density voxels."""
+    stack_info = read_mrc_stack_info(path)
+    if stack_info is None or stack_info.count <= 1:
+        return None
+
+    try:
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", RuntimeWarning)
+            with mrcfile.open(
+                path,
+                mode="r",
+                permissive=True,
+                header_only=True,
+            ) as volume:
+                voxel_size = tuple(
+                    float(axis)
+                    for axis in (
+                        volume.voxel_size.x,
+                        volume.voxel_size.y,
+                        volume.voxel_size.z,
+                    )
+                )
+                minimum = float(volume.header.dmin)
+                maximum = float(volume.header.dmax)
+    except (OSError, OverflowError, TypeError, ValueError):
+        return None
+
+    if not all(np.isfinite(axis) and axis > 0.0 for axis in voxel_size):
+        voxel_size = (1.0, 1.0, 1.0)
+    if not np.isfinite(minimum) or not np.isfinite(maximum):
+        minimum = maximum = 0.0
+
+    return MRCVolumeInfo(
+        width=stack_info.width,
+        height=stack_info.height,
+        depth=stack_info.count,
+        mode=stack_info.mode,
+        data_offset=stack_info.data_offset,
+        voxel_size=voxel_size,
+        minimum=minimum,
+        maximum=maximum,
+    )
+
+
+def _sample_volume_axis(length, max_dimension):
+    output_length = min(length, max_dimension)
+    return np.minimum(
+        length - 1,
+        ((np.arange(output_length) + 0.5) * length / output_length).astype(np.intp),
+    )
+
+
+def read_mrc_volume_payload(path, max_dimension=128):
+    """Read a bounded 3D texture, normalized to the source density range."""
+    if (
+        not isinstance(max_dimension, int)
+        or isinstance(max_dimension, bool)
+        or not 8 <= max_dimension <= 256
+    ):
+        return None
+
+    info = read_mrc_volume_info(path)
+    if info is None:
+        return None
+
+    source_slices = _sample_volume_axis(info.depth, max_dimension)
+    source_rows = _sample_volume_axis(info.height, max_dimension)
+    source_columns = _sample_volume_axis(info.width, max_dimension)
+    try:
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore", RuntimeWarning)
+            with mrcfile.mmap(path, mode="r", permissive=True) as volume:
+                values = volume.data
+                if (
+                    values is None
+                    or values.ndim != 3
+                    or values.shape != (info.depth, info.height, info.width)
+                ):
+                    return None
+                sampled = np.asarray(
+                    values[np.ix_(source_slices, source_rows, source_columns)],
+                    dtype=np.float32,
+                )
+    except (IndexError, OSError, OverflowError, TypeError, ValueError):
+        return None
+
+    finite_mask = np.isfinite(sampled)
+    if not finite_mask.any():
+        return None
+
+    minimum = info.minimum
+    maximum = info.maximum
+    if maximum <= minimum:
+        finite_values = sampled[finite_mask]
+        minimum = float(np.min(finite_values))
+        maximum = float(np.max(finite_values))
+
+    dynamic_range = maximum - minimum
+    if dynamic_range <= 0.0:
+        texture = np.zeros(sampled.shape, dtype=np.uint8)
+    else:
+        safe_values = np.where(finite_mask, sampled, minimum)
+        texture = np.rint(
+            np.clip(255.0 * (safe_values - minimum) / dynamic_range, 0.0, 255.0)
+        ).astype(np.uint8)
+
+    return MRCVolumePayload(
+        data=np.ascontiguousarray(texture).tobytes(order="C"),
+        width=texture.shape[2],
+        height=texture.shape[1],
+        depth=texture.shape[0],
+        source_width=info.width,
+        source_height=info.height,
+        source_depth=info.depth,
+        voxel_size=info.voxel_size,
+        minimum=minimum,
+        maximum=maximum,
+    )
+
+
 def _encode_grayscale_image(pixels, image_format):
     """Encode a two-dimensional uint8 array as an in-memory image."""
     if (