basic raw volumetric data support

molviewspec/molviewspec/builder.py

@@ -11,6 +11,8 @@
 
 from pydantic import BaseModel, PrivateAttr
 
+# TODO: fix imports prior to PR
+# from .nodes import (
 from molviewspec.nodes import (
     AdditionalProperties,
     ApplySelectionInlineParams,
@@ -36,6 +38,9 @@
     Node,
     ParseFormatT,
     ParseParams,
+    RawVolumeOptionsT,
+    RawVolumeParams,
+    RawVolumeSourceT,
     RepresentationParams,
     RepresentationTypeT,
     SchemaFormatT,
@@ -48,6 +53,10 @@
     TooltipInlineParams,
     TransformParams,
     TransparencyInlineParams,
+    VSVolumeOptionsT,
+    VSVolumeParams,
+    VolumeRepresentationParams,
+    VolumeRepresentationTypeT,
 )
 from molviewspec.utils import get_major_version_tag, make_params
 
@@ -224,6 +233,37 @@ class Parse(_Base):
     """
     Builder step with operations needed after parsing structure data.
     """
+    def raw_volume(
+        self,
+        *,
+        source: RawVolumeSourceT,
+        options: RawVolumeOptionsT | None = None,
+        additional_properties: AdditionalProperties = None,
+    ) -> Volume:
+        """
+        Create a raw volume
+        """
+        params = make_params(RawVolumeParams, locals(), type="raw_volume")
+        node = Node(kind="raw_volume", params=params, additional_properties=additional_properties)
+        self._add_child(node)
+        return Volume(node=node, root=self._root)
+
+    # TODO: consider merging with raw_volume into a single "volume" node
+    # followed by adding another source (possibly with prior renaming
+    # of the "source" param to e.g. "type" or something like this).
+    def vs_volume(
+        self,
+        *,
+        options: VSVolumeOptionsT | None = None,
+        additional_properties: AdditionalProperties = None,
+    ) -> Volume:
+        """
+        Create a volume based on Volume Server (VS) data
+        """
+        params = make_params(VSVolumeParams, locals(), type="vs_volume")
+        node = Node(kind="vs_volume", params=params, additional_properties=additional_properties)
+        self._add_child(node)
+        return Volume(node=node, root=self._root)
 
     def model_structure(
         self,
@@ -315,7 +355,24 @@ def symmetry_mates_structure(
         node = Node(kind="structure", params=params, additional_properties=additional_properties)
         self._add_child(node)
         return Structure(node=node, root=self._root)
-
+
+class Volume(_Base):
+    """
+    Builder step with operations needed after defining the volume to work with.
+    """
+    def volume_representation(
+        self, *, type: VolumeRepresentationTypeT = "isosurface", additional_properties: AdditionalProperties = None
+    ) -> VolumeRepresentation:
+        """
+        Add a representation for this component.
+        :param type: the type of representation, defaults to 'isosurface'
+        :param additional_properties: optional, custom data to attach to this node
+        :return: a builder that handles operations at representation level
+        """
+        params = make_params(VolumeRepresentationParams, locals())
+        node = Node(kind="volume_representation", params=params, additional_properties=additional_properties)
+        self._add_child(node)
+        return VolumeRepresentation(node=node, root=self._root)
 
 class Structure(_Base):
     """
@@ -710,6 +767,46 @@ def color(
         self._add_child(node)
         return self
 
+    def transparency(
+        self, *, transparency: float = 0.8, additional_properties: AdditionalProperties = None
+    ) -> Representation:
+        """
+        Customize the transparency/opacity of this representation.
+        :param transparency: float describing how transparent this representation should be, 0.0: fully opaque, 1.0: fully transparent
+        :param additional_properties: optional, custom data to attach to this node
+        :return: this builder
+        """
+        raise NotImplementedError("'transparency' method is not implemented yet")
+        # params = make_params(TransparencyInlineParams, locals())
+        # node = Node(kind="transparency", params=params, additional_properties=additional_properties)
+        # self._add_child(node)
+        # return self
+
+
+class VolumeRepresentation(_Base):
+    """
+    Builder step with operations relating to particular representations.
+    """
+    def color(
+        self,
+        *,
+        color: ColorT,
+        selector: ComponentSelectorT | ComponentExpression | list[ComponentExpression] = "all",
+        additional_properties: AdditionalProperties = None,
+    ) -> Representation:
+        """
+        Customize the color of this representation.
+        :param color: color using SVG color names or RGB hex code
+        :param selector: optional selector, defaults to applying the color to the whole representation
+        :param additional_properties: optional, custom data to attach to this node
+        :return: this builder
+        """
+        params = make_params(ColorInlineParams, locals())
+        node = Node(kind="color", params=params, additional_properties=additional_properties)
+        self._add_child(node)
+        return self
+
+    # TODO: make it work
     def transparency(
         self, *, transparency: float = 0.8, additional_properties: AdditionalProperties = None
     ) -> Representation:

molviewspec/molviewspec/nodes.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 from datetime import datetime, timezone
-from typing import Any, Literal, Mapping, Optional, Union
+from typing import Any, Literal, Mapping, Optional, TypedDict, Union
 from uuid import uuid4
 
 from pydantic import BaseModel, Field
@@ -29,6 +29,7 @@
     "label_from_uri",
     "line",
     "parse",
+    "raw_volume",
     "representation",
     "sphere",
     "structure",
@@ -37,6 +38,8 @@
     "tooltip_from_uri",
     "transform",
     "transparency",
+    "volume_representation",
+    "vs_volume"
 ]
 
 
@@ -160,9 +163,46 @@ class DownloadParams(BaseModel):
     url: str = Field(description="URL from which to pull structure data.")
 
 
-ParseFormatT = Literal["mmcif", "bcif", "pdb"]
+ParseFormatT = Literal["mmcif", "bcif", "pdb", "map", "vs-density"]
+# RawVolumeSourceT = Literal["map", "omezarr", "ometiff_image", "tiff_stack"]
+RawVolumeSourceT = Literal["map"]
 
+ChannelIdsMapping = dict[str, str]
+
 
+class RawVolumeOptionsT(TypedDict):
+    """
+    Specifies the desired voxel size. Overwrites the automatically determined voxel size (e.g., the one based on the map header or its analogue for other formats).
+    Specifies the mapping of sequential channel IDs to user-defined ones.    
+    """
+    voxel_size: float | None
+    channel_ids_mapping: ChannelIdsMapping | None
+
+class VSVolumeOptionsT(TypedDict):
+    """
+    Specifies the max points of the volume to be loaded.
+    Specifies the time frame index the data for which will be loaded. If not provided, the data for all of the available time frame indices will be loaded.
+    Specifies the channel ID the data for which will be loaded. If not provided, the data for all of the available channel IDs will be loaded.
+    """
+    max_points: int | None
+    time: int | None
+    channel_id: str | None
+
+class RawVolumeParams(BaseModel):
+    """
+    Create a volume from a parsed data resource based on the provided parameters.
+    """
+
+    source: RawVolumeSourceT = Field(description="The type of the raw input file with volumetric data.")
+    options: RawVolumeOptionsT | None = Field(description="Specifies the voxel size and mapping of sequential channel IDs to user-defined channel IDs.")
+
+class VSVolumeParams(BaseModel):
+    """
+    Create a volume from a parsed data resource based on the provided parameters.
+    """
+
+    options: VSVolumeOptionsT | None = Field(description="Optionally specifies the max points, time frame index, and channel ID")
+
 class ParseParams(BaseModel):
     """
     Parse node, describing how to parse downloaded data.
@@ -233,8 +273,9 @@ class ComponentExpression(BaseModel):
     atom_id: Optional[int] = Field(description="Unique atom identifier (`_atom_site.id`)")
     atom_index: Optional[int] = Field(description="0-based atom index in the source file")
 
-
-RepresentationTypeT = Literal["ball_and_stick", "cartoon", "surface"]
+# TODO: "slice"
+VolumeRepresentationTypeT = Literal["isosurface", "direct_volume", "slice"]
+RepresentationTypeT = Literal["ball_and_stick", "cartoon"]
 ColorNamesT = Literal[
     "aliceblue",
     "antiquewhite",
@@ -387,6 +428,14 @@ class ComponentExpression(BaseModel):
 ColorT = Union[ColorNamesT, str]  # str represents hex colors for now
 
 
+# TODO: Segmentation too?
+class VolumeRepresentationParams(BaseModel):
+    """
+    Representation node, describing how to represent a volume.
+    """
+    # TODO: add slice
+    type: VolumeRepresentationTypeT = Field(description="Representation type, i.e. isosurface or direct_volume")
+
 class RepresentationParams(BaseModel):
     """
     Representation node, describing how to represent a component.