[WIP] feat: Write model files (#24)

* Create a scattered_point model

* Writing works

* Write docstring for test

* Remove not needed files

* Adjust objectsize before writing

* Use pydantic defaults

* Further simplify models

* Better naming and new test

* Update src/imodmodel/models.py

* Update src/imodmodel/writers.py

* Remove API prototype

* Add a roundtrip test

* Update tests/test_model_api.py

---------

Co-authored-by: alisterburt <[email protected]>

src/imodmodel/models.py

@@ -22,33 +22,33 @@ class GeneralStorage(BaseModel):
 
 class ModelHeader(BaseModel):
     """https://bio3d.colorado.edu/imod/doc/binspec.html"""
-    name: str
-    xmax: int
-    ymax: int
-    zmax: int
-    objsize: int
-    flags: int
-    drawmode: int
-    mousemode: int
-    blacklevel: int
-    whitelevel: int
-    xoffset: float
-    yoffset: float
-    zoffset: float
-    xscale: float
-    yscale: float
-    zscale: float
-    object: int
-    contour: int
-    point: int
-    res: int
-    thresh: int
-    pixelsize: float
-    units: int
-    csum: int
-    alpha: float
-    beta: float
-    gamma: float
+    name: str = 'IMOD-NewModel'
+    xmax: int = 0
+    ymax: int = 0
+    zmax: int = 0
+    objsize: int = 0
+    flags: int = 402653704
+    drawmode: int = 1
+    mousemode: int = 2
+    blacklevel: int = 0
+    whitelevel: int = 255
+    xoffset: float = 0.0
+    yoffset: float = 0.0
+    zoffset: float = 0.0
+    xscale: float = 1.0
+    yscale: float = 1.0
+    zscale: float = 1.0
+    object: int = 0
+    contour: int = 0
+    point: int = -1
+    res: int = 3
+    thresh: int = 128
+    pixelsize: float = 1.0
+    units: int = 0
+    csum: int = 0
+    alpha: float = 0.0
+    beta: float = 0.0
+    gamma: float = 0.0
 
     @field_validator('name', mode="before")
     @classmethod
@@ -59,26 +59,26 @@ def decode_null_terminated_byte_string(cls, value: bytes):
 
 class ObjectHeader(BaseModel):
     """https://bio3d.colorado.edu/imod/doc/binspec.html"""
-    name: str
-    extra_data: List[int]
-    contsize: int
-    flags: int
-    axis: int
-    drawmode: int
-    red: float
-    green: float
-    blue: float
-    pdrawsize: int
-    symbol: int
-    symsize: int
-    linewidth2: int
-    linewidth: int
-    linesty: int
-    symflags: int
-    sympad: int
-    trans: int
-    meshsize: int
-    surfsize: int
+    name: str = ''
+    extra_data: List[int] = [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]
+    contsize: int = 1
+    flags: int = 402653704
+    axis: int = 0
+    drawmode: int = 1
+    red: float = 0.0
+    green: float = 1.0
+    blue: float = 0.0
+    pdrawsize: int = 2
+    symbol: int = 1
+    symsize: int = 3
+    linewidth2: int = 1
+    linewidth: int = 1
+    linesty: int = 0
+    symflags: int = 0
+    sympad: int = 0
+    trans: int = 0
+    meshsize: int = 0
+    surfsize: int = 0
 
     @field_validator('name', mode="before")
     @classmethod
@@ -174,19 +174,19 @@ def face_values(self) -> Optional[np.ndarray]:
 
 class IMAT(BaseModel):
     """https://bio3d.colorado.edu/imod/doc/binspec.html"""
-    ambient: int
-    diffuse: int
-    specular: int
-    shininess: int
-    fillred: int
-    fillgreen: int
-    fillblue: int
-    quality: int
-    mat2: int
-    valblack: int
-    valwhite: int
-    matflags2: int
-    mat3b3: int
+    ambient: int = 102
+    diffuse: int = 255
+    specular: int = 127
+    shininess: int = 4
+    fillred: int = 0
+    fillgreen: int = 0
+    fillblue: int = 0
+    quality: int = 0
+    mat2: int = 0
+    valblack: int = 0
+    valwhite: int = 255
+    matflags2: int = 0 
+    mat3b3: int = 0
 
 class MINX(BaseModel):
     """https://bio3d.colorado.edu/imod/doc/binspec.html"""
@@ -237,9 +237,10 @@ class SLAN(BaseModel):
     center: Tuple[float,float,float]
     label: str
 
+
 class Object(BaseModel):
     """https://bio3d.colorado.edu/imod/doc/binspec.html"""
-    header: ObjectHeader
+    header: ObjectHeader = ObjectHeader()
     contours: List[Contour] = []
     meshes: List[Mesh] = []
     extra: List[GeneralStorage] = []
@@ -251,11 +252,11 @@ class ImodModel(BaseModel):
 
     https://bio3d.colorado.edu/imod/doc/binspec.html
     """
-    id: ID
-    header: ModelHeader
-    objects: List[Object]
+    id: ID = ID(IMOD_file_id='IMOD', version_id='V1.2')
+    header: ModelHeader = ModelHeader()
+    objects: List[Object] = []
     slicer_angles: List[SLAN] = []
-    minx: Optional[MINX]
+    minx: Optional[MINX] = None
     extra: List[GeneralStorage] = []
 
     @classmethod
@@ -264,4 +265,10 @@ def from_file(cls, filename: os.PathLike):
         from .parsers import parse_model
         with open(filename, 'rb') as file:
             return parse_model(file)
-
+
+    def to_file(self, filename: os.PathLike):
+        """Write an IMOD model to disk."""
+        from .writers import write_model
+        self.header.objsize = len(self.objects)
+        with open(filename, 'wb') as file:
+            write_model(file, self)

src/imodmodel/writers.py

@@ -0,0 +1,145 @@
+
+from struct import Struct
+from typing import BinaryIO, List, Union
+
+import numpy as np
+
+from .models import (
+    ID,
+    IMAT,
+    SLAN,
+    Contour,
+    ContourHeader,
+    GeneralStorage,
+    Mesh,
+    MeshHeader,
+    MINX,
+    ImodModel,
+    ModelHeader,
+    Object,
+    ObjectHeader,
+)
+from .binary_specification import ModFileSpecification
+
+
+
+def _write_to_format_str(file: BinaryIO, format_str: str, data: Union[tuple, list]):
+    # Convert data to bytes
+    data = [s.encode("utf-8") if isinstance(s, str) else s for s in data]
+    struct = Struct(format_str)
+    file.write(struct.pack(*data))
+
+
+def _write_to_specification(file: BinaryIO, specification: dict, data: dict):
+    format_str = f">{''.join(specification.values())}"
+    data_1d = []
+    for key, value in specification.items():
+        if value[0].isdigit() and value[-1] in "iIlLqQfd":
+            data_1d.extend(data[key])
+        else:
+            data_1d.append(data[key])
+    _write_to_format_str(file, format_str, data_1d)
+
+
+def _write_id(file: BinaryIO, id: ID):
+    _write_to_specification(file, ModFileSpecification.ID, id.dict())
+
+
+def _write_model_header(file: BinaryIO, header: ModelHeader):
+    _write_to_specification(file, ModFileSpecification.MODEL_HEADER, header.dict())
+
+
+def _write_object_header(file: BinaryIO, header: ObjectHeader):
+    _write_to_specification(file, ModFileSpecification.OBJECT_HEADER, header.dict())
+
+
+def _write_contour_header(file: BinaryIO, header: ContourHeader):
+    _write_to_specification(file, ModFileSpecification.CONTOUR_HEADER, header.dict())
+
+
+def _write_mesh_header(file: BinaryIO, header: MeshHeader):
+    _write_to_specification(file, ModFileSpecification.MESH_HEADER, header.dict())
+
+
+def _write_imat(file: BinaryIO, imat: IMAT):
+    _write_chunk_size(file, 48)
+    _write_to_specification(file, ModFileSpecification.IMAT, imat.dict())
+
+
+def _write_minx(file: BinaryIO, minx: MINX):
+    _write_chunk_size(file, 72)
+    _write_to_specification(file, ModFileSpecification.MINX, minx.dict())
+
+
+def _write_slicer_angle(file: BinaryIO, slicer_angle: SLAN):
+    _write_chunk_size(file, 48)
+    _write_to_specification(file, ModFileSpecification.SLAN, slicer_angle.dict())
+
+
+def _write_general_storage(file: BinaryIO, storages: List[GeneralStorage]):
+    size = len(storages) * 12
+    _write_chunk_size(file, size)
+    for storage in storages:
+        _write_to_format_str(file, '>hh', (storage.type, storage.flags))
+        _write_to_format_str(file, '>i' if isinstance(storage.index, int) else '>f', (storage.index,))
+        _write_to_format_str(file, '>i' if isinstance(storage.value, int) else '>f', (storage.value,))
+
+
+def _write_chunk_size(file: BinaryIO, size: int):
+    file.write(size.to_bytes(4, byteorder="big"))
+
+
+def _write_control_sequence(file: BinaryIO, sequence: str):
+    file.write(sequence.encode("utf-8"))
+
+
+def _write_contour(file: BinaryIO, contour: Contour):
+    _write_contour_header(file, contour.header)
+    points = contour.points.flatten()
+    _write_to_format_str(file, f">{'f' * len(points)}", points)
+    if contour.extra:
+        _write_general_storage(file, contour.extra)
+
+
+def _write_mesh(file: BinaryIO, mesh: Mesh):
+    _write_mesh_header(file, mesh.header)
+    vertices = mesh.raw_vertices.flatten()
+    indices = mesh.raw_indices.flatten()
+    _write_to_format_str(file, f">{'f' * len(vertices)}", vertices)
+    _write_to_format_str(file, f">{'i' * len(indices)}", indices)
+    if mesh.extra:
+        _write_general_storage(file, mesh.extra)
+
+
+def _write_object(file: BinaryIO, obj: Object):
+    _write_object_header(file, obj.header)
+    for contour in obj.contours:
+        _write_control_sequence(file, "CONT")
+        _write_contour(file, contour)
+    for mesh in obj.meshes:
+        _write_control_sequence(file, "MESH")
+        _write_mesh(file, mesh)
+    if obj.imat:
+        _write_control_sequence(file, "IMAT")
+        _write_imat(file, obj.imat)
+    if obj.extra:
+        _write_control_sequence(file, "OBST")
+        _write_general_storage(file, obj.extra)
+
+
+def write_model(file: BinaryIO, model: ImodModel):
+    _write_id(file, model.id)
+    _write_model_header(file, model.header)
+    for obj in model.objects:
+        _write_control_sequence(file, "OBJT")
+        _write_object(file, obj)
+    for slicer_angle in model.slicer_angles:
+        _write_control_sequence(file, "SLAN")
+        _write_slicer_angle(file, slicer_angle)
+    if model.minx:
+        _write_control_sequence(file, "MINX")
+        _write_minx(file, model.minx)
+    if model.extra:
+        _write_control_sequence(file, "MOST")
+        _write_general_storage(file, model.extra)
+    _write_control_sequence(file, "IEOF")

tests/test_model_api.py

@@ -77,5 +77,16 @@ def test_read_minx(meshed_contour_model_file):
     assert model.minx.ctrans == pytest.approx((-2228.0, 2228.0, 681.099976), abs=1e-6)
     assert model.minx.crot == pytest.approx((0.0, 0.0, 0.0), abs=1e-6)
 
+def test_read_write_read_roundtrip(two_contour_model_file, tmp_path):
+    """Check that reading and writing a model file results in the same data."""
+    import numpy as np
 
-
+    model = ImodModel.from_file(two_contour_model_file)
+    model.to_file(tmp_path / "test_model.imod")
+    model2 = ImodModel.from_file(tmp_path / "test_model.imod")
+    assert model.header == model2.header
+    assert model.objects[0].header == model2.objects[0].header
+    assert model.objects[0].contours[0].header == model2.objects[0].contours[0].header
+    assert np.allclose(model.objects[0].contours[0].points, model2.objects[0].contours[0].points)
+    assert model.objects[0].contours[1].header == model2.objects[0].contours[1].header
+    assert np.allclose(model.objects[0].contours[1].points, model2.objects[0].contours[1].points)